Merge branch 'master' of github.com:flyteorg/flytekit into Add_user_specified_field

Author: arbaobao

flytekit/clis/sdk_in_container/run.py

@@ -160,7 +160,7 @@ class RunLevelParams(PyFlyteParams):
     )
     poll_interval: int = make_click_option_field(
         click.Option(
-            param_decls=["-i", "--poll-interval", "poll_interval"],
+            param_decls=["--poll-interval", "poll_interval"],
             required=False,
             type=int,
             default=None,

flytekit/core/base_task.py

@@ -131,6 +131,8 @@ class TaskMetadata(object):
         pod_template_name (Optional[str]): The name of an existing PodTemplate resource in the cluster which will be used for this task.
         generates_deck (bool): Indicates whether the task will generate a Deck URI.
         is_eager (bool): Indicates whether the task should be treated as eager.
+        labels (Optional[dict[str, str]]): Labels to be applied to the task resource.
+        annotations (Optional[dict[str, str]]): Annotations to be applied to the task resource.
     """
 
     cache: bool = False
@@ -144,6 +146,8 @@ class TaskMetadata(object):
     pod_template_name: Optional[str] = None
     generates_deck: bool = False
     is_eager: bool = False
+    labels: Optional[dict[str, str]] = None
+    annotations: Optional[dict[str, str]] = None
 
     def __post_init__(self):
         if self.timeout:
@@ -185,6 +189,10 @@ def to_taskmetadata_model(self) -> _task_model.TaskMetadata:
             pod_template_name=self.pod_template_name,
             cache_ignore_input_vars=self.cache_ignore_input_vars,
             is_eager=self.is_eager,
+            k8s_object_metadata=_task_model.K8sObjectMetadata(
+                labels=self.labels,
+                annotations=self.annotations,
+            ),
         )
 
 

flytekit/core/python_auto_container.py

@@ -100,6 +100,11 @@ def __init__(
         kwargs["metadata"] = kwargs["metadata"] if "metadata" in kwargs else TaskMetadata()
         kwargs["metadata"].pod_template_name = pod_template_name
 
+        if kwargs.get("labels") is not None:
+            kwargs["metadata"].labels = kwargs["labels"]
+        if kwargs.get("annotations") is not None:
+            kwargs["metadata"].annotations = kwargs["annotations"]
+
         self._container_image = container_image
         # TODO(katrogan): Implement resource overrides
 

flytekit/core/task.py

@@ -128,6 +128,8 @@ def task(
     pickle_untyped: bool = ...,
     shared_memory: Optional[Union[L[True], str]] = None,
     resources: Optional[Resources] = ...,
+    labels: Optional[dict[str, str]] = ...,
+    annotations: Optional[dict[str, str]] = ...,
     **kwargs,
 ) -> Callable[[Callable[..., FuncOut]], PythonFunctionTask[T]]: ...
 
@@ -167,6 +169,8 @@ def task(
     pickle_untyped: bool = ...,
     shared_memory: Optional[Union[L[True], str]] = ...,
     resources: Optional[Resources] = ...,
+    labels: Optional[dict[str, str]] = ...,
+    annotations: Optional[dict[str, str]] = ...,
     **kwargs,
 ) -> Union[Callable[P, FuncOut], PythonFunctionTask[T]]: ...
 
@@ -211,6 +215,8 @@ def task(
     pickle_untyped: bool = False,
     shared_memory: Optional[Union[L[True], str]] = None,
     resources: Optional[Resources] = None,
+    labels: Optional[dict[str, str]] = None,
+    annotations: Optional[dict[str, str]] = None,
     **kwargs,
 ) -> Union[
     Callable[P, FuncOut],
@@ -348,6 +354,8 @@ def launch_dynamically():
         first value is the request and the second value is the limit. If the value is a single value, then both the
         requests and limit is set to that value. For example, the `Resource(cpu=("1", "2"), mem="1Gi")` will set the cpu
         request to 1, cpu limit to 2, and mem request to 1Gi.
+    :param labels: Labels to be applied to the task resource.
+    :param annotations: Annotations to be applied to the task resource.
     """
     # Maintain backwards compatibility with the old cache parameters, while cleaning up the task function definition.
     cache_serialize = kwargs.pop("cache_serialize", None)
@@ -445,6 +453,8 @@ def wrapper(fn: Callable[P, FuncOut]) -> PythonFunctionTask[T]:
             pickle_untyped=pickle_untyped,
             shared_memory=shared_memory,
             resources=resources,
+            labels=labels,
+            annotations=annotations,
         )
         update_wrapper(task_instance, decorated_fn)
         return task_instance

flytekit/core/type_engine.py

@@ -660,10 +660,8 @@ def get_literal_type(self, t: Type[T]) -> LiteralType:
                 python_type = hints.get(name, field.type)
                 literal_type[name] = TypeEngine.to_literal_type(python_type)
             except Exception as e:
-                logger.warning(
-                    "Field {} of type {} cannot be converted to a literal type. Error: {}".format(
-                        field.name, field.type, e
-                    )
+                logger.debug(
+                    f"Field {field.name} of type {field.type} cannot be converted to a literal type. Error: {e}"
                 )
 
         # This is for attribute access in FlytePropeller.
@@ -1771,6 +1769,19 @@ def _type_essence(x: LiteralType) -> LiteralType:
 
 
 def _are_types_castable(upstream: LiteralType, downstream: LiteralType) -> bool:
+    if upstream.union_type is not None:
+        # for each upstream variant, there must be a compatible type downstream
+        for v in upstream.union_type.variants:
+            if not _are_types_castable(v, downstream):
+                return False
+        return True
+
+    if downstream.union_type is not None:
+        # there must be a compatible downstream type
+        for v in downstream.union_type.variants:
+            if _are_types_castable(upstream, v):
+                return True
+
     if upstream.collection_type is not None:
         if downstream.collection_type is None:
             return False
@@ -1816,19 +1827,6 @@ def _are_types_castable(upstream: LiteralType, downstream: LiteralType) -> bool:
 
         return True
 
-    if upstream.union_type is not None:
-        # for each upstream variant, there must be a compatible type downstream
-        for v in upstream.union_type.variants:
-            if not _are_types_castable(v, downstream):
-                return False
-        return True
-
-    if downstream.union_type is not None:
-        # there must be a compatible downstream type
-        for v in downstream.union_type.variants:
-            if _are_types_castable(upstream, v):
-                return True
-
     if upstream.enum_type is not None:
         # enums are castable to string
         if downstream.simple == SimpleType.STRING:
@@ -2115,7 +2113,7 @@ async def dict_to_generic_literal(
                     ),
                     metadata={"format": "pickle"},
                 )
-            raise TypeTransformerFailedError(f"Cannot convert `{v}` to Flyte Literal.\n" f"Error Message: {e}")
+            raise TypeTransformerFailedError(f"Cannot convert `{v}` to Flyte Literal.\nError Message: {e}")
 
     @staticmethod
     async def dict_to_binary_literal(
@@ -2141,7 +2139,7 @@ async def dict_to_binary_literal(
                     ),
                     metadata={"format": "pickle"},
                 )
-            raise TypeTransformerFailedError(f"Cannot convert `{v}` to Flyte Literal.\n" f"Error Message: {e}")
+            raise TypeTransformerFailedError(f"Cannot convert `{v}` to Flyte Literal.\nError Message: {e}")
 
     @staticmethod
     def is_pickle(python_type: Type[dict]) -> bool:

flytekit/extras/accelerators.py

@@ -106,6 +106,10 @@ def to_flyte_idl(self) -> tasks_pb2.GPUAccelerator:
 #: `NVIDIA L4 Tensor Core GPU https://www.nvidia.com/en-us/data-center/l4
 L4_VWS = GPUAccelerator("nvidia-l4-vws")
 
+#: use this constant to specify that the task should run on an
+#: `NVIDIA L40S Tensor Core GPU https://www.nvidia.com/en-us/data-center/l40s
+L40S = GPUAccelerator("nvidia-l40s")
+
 #: use this constant to specify that the task should run on an
 #: `NVIDIA Tesla K80 GPU https://www.nvidia.com/en-gb/data-center/tesla-k80
 K80 = GPUAccelerator("nvidia-tesla-k80")
@@ -130,6 +134,14 @@ def to_flyte_idl(self) -> tasks_pb2.GPUAccelerator:
 #: `NVIDIA Tesla V100 GPU https://images.nvidia.com/content/technologies/volta/pdf/tesla-volta-v100-datasheet-letter-fnl-web.pdf
 V100 = GPUAccelerator("nvidia-tesla-v100")
 
+#: use this constant to specify that the task should run on an
+#: `NVIDIA H100 GPU https://www.nvidia.com/en-us/data-center/h100
+H100 = GPUAccelerator("nvidia-h100")
+
+#: use this constant to specify that the task should run on an
+#: `NVIDIA H200 GPU https://www.nvidia.com/en-us/data-center/h200
+H200 = GPUAccelerator("nvidia-h200")
+
 
 class MultiInstanceGPUAccelerator(BaseAccelerator):
     """

flytekit/image_spec/default_builder.py

@@ -29,7 +29,8 @@
     --mount=type=bind,target=uv.lock,src=uv.lock \
     --mount=type=bind,target=pyproject.toml,src=pyproject.toml \
     $PIP_SECRET_MOUNT \
-    uv sync $PIP_INSTALL_ARGS
+    uv sync $PIP_INSTALL_ARGS && \
+    chown -R flytekit /root/.venv
 WORKDIR /
 
 # Update PATH and UV_PYTHON to point to the venv created by uv sync
@@ -54,12 +55,12 @@
     --mount=type=bind,target=poetry.lock,src=poetry.lock \
     --mount=type=bind,target=pyproject.toml,src=pyproject.toml \
     $PIP_SECRET_MOUNT \
-    poetry install $PIP_INSTALL_ARGS
-
+    poetry install $PIP_INSTALL_ARGS && \
+    chown -R flytekit /root/.venv
 WORKDIR /
 
 # Update PATH and UV_PYTHON to point to venv
-ENV PATH="/root/.venv/bin:$$PATH" \
+ENV PATH="/root/.venv/bin:$$PATH"  \
     UV_PYTHON=/root/.venv/bin/python
 """
 )
@@ -81,14 +82,19 @@
     $APT_PACKAGES
 """)
 
+# make sure that micromamba python installation is owned by flytekit user
 MICROMAMBA_INSTALL_COMMAND_TEMPLATE = Template("""\
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/opt/micromamba/pkgs,\
 id=micromamba \
     --mount=from=micromamba,source=/usr/bin/micromamba,target=/usr/bin/micromamba \
     micromamba config set use_lockfiles False && \
-    micromamba create -n runtime --root-prefix /opt/micromamba \
+    ( micromamba create -n runtime --root-prefix /opt/micromamba \
+    -c conda-forge $CONDA_CHANNELS \
+    python=$PYTHON_VERSION $CONDA_PACKAGES \
+    || micromamba install -n runtime --root-prefix /opt/micromamba \
     -c conda-forge $CONDA_CHANNELS \
-    python=$PYTHON_VERSION $CONDA_PACKAGES
+    python=$PYTHON_VERSION $CONDA_PACKAGES ) && \
+    chown -R flytekit /opt/micromamba
 """)
 
 DOCKER_FILE_TEMPLATE = Template("""\
@@ -98,6 +104,7 @@
 
 FROM $BASE_IMAGE
 
+WORKDIR /
 USER root
 $APT_INSTALL_COMMAND
 RUN --mount=from=micromamba,source=/etc/ssl/certs/ca-certificates.crt,target=/tmp/ca-certificates.crt \
@@ -118,7 +125,7 @@
     SSL_CERT_DIR=/etc/ssl/certs \
     $ENV
 
-$UV_PYTHON_INSTALL_COMMAND
+$PYTHON_INSTALL_COMMAND
 
 # Adds nvidia just in case it exists
 ENV PATH="$$PATH:/usr/local/nvidia/bin:/usr/local/cuda/bin" \
@@ -336,7 +343,7 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path):
         )
         raise ValueError(msg)
 
-    uv_python_install_command = prepare_python_install(image_spec, tmp_dir)
+    python_install_command = prepare_python_install(image_spec, tmp_dir)
     env_dict = {"PYTHONPATH": "/root"}
 
     if image_spec.env:
@@ -422,11 +429,11 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path):
     _f_img_id_env = f"{_F_IMG_ID}={image_spec.id}"
 
     docker_content = DOCKER_FILE_TEMPLATE.substitute(
-        UV_PYTHON_INSTALL_COMMAND=uv_python_install_command,
-        APT_INSTALL_COMMAND=apt_install_command,
         INSTALL_PYTHON_TEMPLATE=python_install_template.template,
         EXTRA_PATH=python_install_template.extra_path,
         PYTHON_EXEC=python_install_template.python_exec,
+        APT_INSTALL_COMMAND=apt_install_command,
+        PYTHON_INSTALL_COMMAND=python_install_command,
         BASE_IMAGE=base_image,
         ENV=env,
         _F_IMG_ID_ENV=_f_img_id_env,

flytekit/models/core/workflow.py

@@ -423,7 +423,7 @@ def to_flyte_idl(self) -> _core_workflow.ArrayNode:
             execution_mode=self._execution_mode,
             is_original_sub_node_interface=BoolValue(value=self._is_original_sub_node_interface),
             data_mode=self._data_mode,
-            bound_inputs=self._bound_inputs,
+            bound_inputs=sorted(self._bound_inputs),
         )
 
     @classmethod

flytekit/models/task.py

@@ -185,6 +185,7 @@ def __init__(
         cache_ignore_input_vars,
         is_eager: bool = False,
         generates_deck: bool = False,
+        k8s_object_metadata: typing.Optional["K8sObjectMetadata"] = None,
     ):
         """
         Information needed at runtime to determine behavior such as whether or not outputs are discoverable, timeouts,
@@ -208,6 +209,7 @@ def __init__(
         :param pod_template_name: The name of the existing PodTemplate resource which will be used in this task.
         :param cache_ignore_input_vars: Input variables that should not be included when calculating hash for cache.
         :param is_eager:
+        :param metadata: Kubernetes metadata for the task.
         """
         self._discoverable = discoverable
         self._runtime = runtime
@@ -221,6 +223,7 @@ def __init__(
         self._cache_ignore_input_vars = cache_ignore_input_vars
         self._is_eager = is_eager
         self._generates_deck = generates_deck
+        self._k8s_object_metadata = k8s_object_metadata
 
     @property
     def is_eager(self):
@@ -318,6 +321,14 @@ def cache_ignore_input_vars(self):
         """
         return self._cache_ignore_input_vars
 
+    @property
+    def k8s_object_metadata(self) -> typing.Optional["K8sObjectMetadata"]:
+        """
+        Kubernetes metadata for the task.
+        :rtype: K8sObjectMetadata
+        """
+        return self._k8s_object_metadata
+
     def to_flyte_idl(self):
         """
         :rtype: flyteidl.admin.task_pb2.TaskMetadata
@@ -334,6 +345,7 @@ def to_flyte_idl(self):
             cache_ignore_input_vars=self.cache_ignore_input_vars,
             is_eager=self.is_eager,
             generates_deck=BoolValue(value=self.generates_deck),
+            metadata=self.k8s_object_metadata.to_flyte_idl() if self.k8s_object_metadata else None,
         )
         if self.timeout:
             tm.timeout.FromTimedelta(self.timeout)
@@ -358,6 +370,9 @@ def from_flyte_idl(cls, pb2_object: _core_task.TaskMetadata):
             cache_ignore_input_vars=pb2_object.cache_ignore_input_vars,
             is_eager=pb2_object.is_eager,
             generates_deck=pb2_object.generates_deck.value if pb2_object.HasField("generates_deck") else False,
+            k8s_object_metadata=K8sObjectMetadata.from_flyte_idl(pb2_object.metadata)
+            if pb2_object.HasField("metadata")
+            else None,
         )
 
 
@@ -985,7 +1000,11 @@ def from_flyte_idl(cls, pb2_object):
 
 
 class K8sObjectMetadata(_common.FlyteIdlEntity):
-    def __init__(self, labels: typing.Dict[str, str] = None, annotations: typing.Dict[str, str] = None):
+    def __init__(
+        self,
+        labels: typing.Optional[typing.Dict[str, str]] = None,
+        annotations: typing.Optional[typing.Dict[str, str]] = None,
+    ):
         """
         This defines additional metadata for building a kubernetes pod.
         """