flyte 0.2.0b22__py3-none-any.whl → 0.2.0b24__py3-none-any.whl

flyte/__init__.py

@@ -1,19 +1,5 @@
 """
-Flyte SDK for authoring Compound AI applications, services and workflows.
-
-## Environments
-
-TaskEnvironment class to define a new environment for a set of tasks.
-
-Example usage:
-
-```python
-env = flyte.TaskEnvironment(name="my_env", image="my_image", resources=Resources(cpu="1", memory="1Gi"))
-
-@env.task
-async def my_task():
-    pass
-```
+Flyte SDK for authoring compound AI applications, services and workflows.
 """
 
 __all__ = [

flyte/_bin/runtime.py

@@ -32,7 +32,12 @@ ENABLE_REF_TASKS = "_REF_TASKS"  # This is a temporary flag to enable reference
 _UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
 
 
-@click.command("a0")
+@click.group()
+def _pass_through():
+    pass
+
+
+@_pass_through.command("a0")
 @click.option("--inputs", "-i", required=True)
 @click.option("--outputs-path", "-o", required=True)
 @click.option("--version", "-v", required=True)
@@ -158,3 +163,7 @@ def main(
         await controller.stop()
 
     asyncio.run(_run_and_stop())
+
+
+if __name__ == "__main__":
+    _pass_through()

flyte/_environment.py

@@ -49,6 +49,13 @@ class Environment:
         """
         Add a dependency to the environment.
         """
+        for e in env:
+            if not isinstance(e, Environment):
+                raise TypeError(f"Expected Environment, got {type(e)}")
+            if e.name == self.name:
+                raise ValueError("Cannot add self as a dependency")
+            if e in self.depends_on:
+                continue
         self.depends_on.extend(env)
 
     def clone_with(

flyte/_image.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import base64
 import hashlib
 import sys
+import typing
 from abc import abstractmethod
 from dataclasses import asdict, dataclass, field
 from functools import cached_property
@@ -10,6 +11,7 @@ from pathlib import Path
 from typing import Callable, ClassVar, Dict, List, Literal, Optional, Tuple, TypeVar, Union
 
 import rich.repr
+from packaging.version import Version
 
 # Supported Python versions
 PYTHON_3_10 = (3, 10)
@@ -352,21 +354,29 @@ class Image:
             layer.validate()
 
     @classmethod
-    def _get_default_image_for(cls, python_version: Tuple[int, int], flyte_version: Optional[str] = None) -> Image:
+    def _get_default_image_for(
+        cls, python_version: Tuple[int, int], flyte_version: Optional[str] = None, install_flyte: bool = True
+    ) -> Image:
         # Would love a way to move this outside of this class (but still needs to be accessible via Image.auto())
         # this default image definition may need to be updated once there is a released pypi version
         from flyte._version import __version__
 
-        dev_mode = (cls._is_editable_install() or (__version__ and "dev" in __version__)) and not flyte_version
-        if flyte_version is None:
-            flyte_version = __version__.replace("+", "-")
-        preset_tag = flyte_version if flyte_version.startswith("v") else f"v{flyte_version}"
-        preset_tag = f"py{python_version[0]}.{python_version[1]}-{preset_tag}"
+        dev_mode = (
+            (cls._is_editable_install() or (__version__ and "dev" in __version__))
+            and not flyte_version
+            and install_flyte
+        )
+        if install_flyte is False:
+            preset_tag = f"py{python_version[0]}.{python_version[1]}"
+        else:
+            if flyte_version is None:
+                flyte_version = __version__.replace("+", "-")
+            suffix = flyte_version if flyte_version.startswith("v") else f"v{flyte_version}"
+            preset_tag = f"py{python_version[0]}.{python_version[1]}-{suffix}"
         image = Image._new(
             base_image=f"python:{python_version[0]}.{python_version[1]}-slim-bookworm",
             registry=_BASE_REGISTRY,
             name=_DEFAULT_IMAGE_NAME,
-            # _tag=preset_tag,
             platform=("linux/amd64", "linux/arm64"),
         )
         labels_and_user = _DockerLines(
@@ -389,11 +399,15 @@ class Image:
         )
         image = image.with_apt_packages("build-essential", "ca-certificates")
 
-        # Add in flyte library
-        if dev_mode:
-            image = image._with_local_v2()
-        else:
-            image = image.with_pip_packages(f"flyte=={flyte_version}")
+        if install_flyte:
+            if dev_mode:
+                image = image.with_local_v2()
+            else:
+                flyte_version = typing.cast(str, flyte_version)
+                if Version(flyte_version).is_devrelease or Version(flyte_version).is_prerelease:
+                    image = image.with_pip_packages(f"flyte=={flyte_version}", pre=True)
+                else:
+                    image = image.with_pip_packages(f"flyte=={flyte_version}")
         object.__setattr__(image, "_tag", preset_tag)
         # Set this to auto for all auto images because the meaning of "auto" can change (based on logic inside
         # _get_default_image_for, acts differently in a running task container) so let's make sure it stays auto.
@@ -413,6 +427,7 @@ class Image:
         cls,
         python_version: Optional[Tuple[int, int]] = None,
         flyte_version: Optional[str] = None,
+        install_flyte: bool = True,
         registry: Optional[str] = None,
         name: Optional[str] = None,
     ) -> Image:
@@ -422,6 +437,7 @@ class Image:
 
         :param python_version: If not specified, will use the current Python version
         :param flyte_version: Union version to use
+        :param install_flyte: If True, will install the flyte library in the image
         :param registry: Registry to use for the image
         :param name: Name of the image if you want to override the default name
 
@@ -430,7 +446,9 @@ class Image:
         if python_version is None:
             python_version = _detect_python_version()
 
-        base_image = cls._get_default_image_for(python_version=python_version, flyte_version=flyte_version)
+        base_image = cls._get_default_image_for(
+            python_version=python_version, flyte_version=flyte_version, install_flyte=install_flyte
+        )
 
         if registry and name:
             return base_image.clone(registry=registry, name=name)
@@ -534,6 +552,7 @@ class Image:
         self,
         registry: Optional[str] = None,
         name: Optional[str] = None,
+        python_version: Optional[Tuple[int, int]] = None,
         addl_layer: Optional[Layer] = None,
     ) -> Image:
         """
@@ -541,17 +560,16 @@ class Image:
 
         :param registry: Registry to use for the image
         :param name: Name of the image
+        :param python_version: Python version for the image, if not specified, will use the current Python version
         :param addl_layer: Additional layer to add to the image. This will be added to the end of the layers.
-        :param addl_layer: Additional layer to add to the image. This will be added on top of the existing layers.
 
         :return:
         """
         registry = registry if registry else self.registry
         name = name if name else self.name
-        if addl_layer and (not name or not registry):
+        if addl_layer and (not name):
             raise ValueError(
-                f"Cannot add additional layer {addl_layer} to an image"
-                f" without a registry and name. Please first clone()."
+                f"Cannot add additional layer {addl_layer} to an image without name. Please first clone()."
             )
         new_layers = (*self._layers, addl_layer) if addl_layer else self._layers
         img = Image._new(
@@ -560,7 +578,7 @@ class Image:
             registry=registry,
             name=name,
             platform=self.platform,
-            python_version=self.python_version,
+            python_version=python_version or self.python_version,
             _layers=new_layers,
         )
 
@@ -710,7 +728,7 @@ class Image:
         :param dst: destination folder in the image
         :return: Image
         """
-        new_image = self.clone(addl_layer=CopyConfig(path_type=1, src=src, dst=dst))
+        new_image = self.clone(addl_layer=CopyConfig(path_type=1, src=src, dst=dst, _compute_identifier=lambda x: dst))
         return new_image
 
     def with_source_file(self, src: Path, dst: str = ".") -> Image:
@@ -722,7 +740,7 @@ class Image:
         :param dst: destination folder in the image
         :return: Image
         """
-        new_image = self.clone(addl_layer=CopyConfig(path_type=0, src=src, dst=dst))
+        new_image = self.clone(addl_layer=CopyConfig(path_type=0, src=src, dst=dst, _compute_identifier=lambda x: dst))
         return new_image
 
     def with_uv_project(
@@ -774,7 +792,7 @@ class Image:
         new_image = self.clone(addl_layer=Commands(commands=new_commands))
         return new_image
 
-    def _with_local_v2(self) -> Image:
+    def with_local_v2(self) -> Image:
         """
         Use this method to create a new image with the local v2 builder
         This will override any existing builder

flyte/_internal/controllers/remote/_action.py

@@ -1,12 +1,15 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import Any, Literal
 
 from flyteidl.core import execution_pb2
 
 from flyte._protos.workflow import run_definition_pb2, state_service_pb2, task_definition_pb2
 from flyte.models import GroupData
 
+ActionType = Literal["task", "trace"]
+
 
 @dataclass
 class Action:
@@ -17,6 +20,7 @@ class Action:
 
     action_id: run_definition_pb2.ActionIdentifier
     parent_action_name: str
+    type: ActionType = "task"  # type of action, task or trace
     friendly_name: str | None = None
     group: GroupData | None = None
     task: task_definition_pb2.TaskSpec | None = None
@@ -144,3 +148,28 @@ class Action:
             err=obj.error if obj.HasField("error") else None,
             realized_outputs_uri=obj.output_uri,
         )
+
+    @classmethod
+    def from_trace(
+        cls,
+        parent_action_name: str,
+        action_id: run_definition_pb2.ActionIdentifier,
+        friendly_name: str,
+        group_data: GroupData | None,
+        trace_spec: Any,  # TODO
+        inputs_uri: str,
+        outputs_uri: str,
+    ) -> Action:
+        """
+        This creates a new action for tracing purposes. It is used to track the execution of a trace.
+        """
+        return cls(
+            action_id=action_id,
+            parent_action_name=parent_action_name,
+            type="trace",
+            friendly_name=friendly_name,
+            group=group_data,
+            inputs_uri=inputs_uri,
+            realized_outputs_uri=outputs_uri,
+            phase=run_definition_pb2.Phase.PHASE_SUCCEEDED,
+        )

flyte/_internal/controllers/remote/_controller.py

@@ -174,13 +174,12 @@ class RemoteController(Controller):
         )
 
         task_spec = translate_task_to_wire(_task, new_serialization_context)
-        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
-
-        inputs_hash = convert.generate_inputs_hash(serialized_inputs)
+        inputs_hash = convert.generate_inputs_hash_from_proto(inputs.proto_inputs)
         sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
             tctx, task_spec, inputs_hash, _task_call_seq
         )
 
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
         inputs_uri = io.inputs_path(sub_action_output_path)
         await upload_inputs_with_retry(serialized_inputs, inputs_uri)
 
@@ -391,10 +390,12 @@ class RemoteController(Controller):
         if tctx is None:
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
 
+        current_action_id = tctx.action
         current_output_path = tctx.output_path
         sub_run_output_path = storage.join(current_output_path, info.action.name)
 
         if info.interface.has_outputs():
+            outputs_file_path: str = ""
             if info.output:
                 outputs = await convert.convert_from_native_to_outputs(info.output, info.interface)
                 outputs_file_path = io.outputs_path(sub_run_output_path)
@@ -405,6 +406,33 @@ class RemoteController(Controller):
                 await io.upload_error(err.err, error_path)
             else:
                 raise flyte.errors.RuntimeSystemError("BadTraceInfo", "Trace info does not have output or error")
+            trace_action = Action.from_trace(
+                parent_action_name=current_action_id.name,
+                action_id=run_definition_pb2.ActionIdentifier(
+                    name=info.action.name,
+                    run=run_definition_pb2.RunIdentifier(
+                        name=current_action_id.run_name,
+                        project=current_action_id.project,
+                        domain=current_action_id.domain,
+                        org=current_action_id.org,
+                    ),
+                ),
+                inputs_uri=info.inputs_path,
+                outputs_uri=outputs_file_path,
+                trace_spec="",
+                friendly_name=info.name,
+                group_data=tctx.group_data,
+            )
+            try:
+                logger.info(
+                    f"Submitting Trace action Run:[{trace_action.run_name}, Parent:[{trace_action.parent_action_name}],"
+                    f" Trace fn:[{info.name}], action:[{info.action.name}]"
+                )
+                await self.submit_action(trace_action)
+                logger.info(f"Trace Action for [{info.name}] action id: {info.action.name}, completed!")
+            except asyncio.CancelledError:
+                # If the action is cancelled, we need to cancel the action on the server as well
+                raise
 
     async def submit_task_ref(self, _task: task_definition_pb2.TaskDetails, *args, **kwargs) -> Any:
         ctx = internal_ctx()
@@ -420,12 +448,12 @@ class RemoteController(Controller):
             _task.spec.task_template.interface, default_inputs=_task.spec.default_inputs
         )
         inputs = await convert.convert_from_native_to_inputs(native_interface, *args, **kwargs)
-        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
-        inputs_hash = convert.generate_inputs_hash(serialized_inputs)
+        inputs_hash = convert.generate_inputs_hash_from_proto(inputs.proto_inputs)
         sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
             tctx, task_name, inputs_hash, invoke_seq_num
         )
 
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
         inputs_uri = io.inputs_path(sub_action_output_path)
         await upload_inputs_with_retry(serialized_inputs, inputs_uri)
         # cache key - task name, task signature, inputs, cache version

flyte/_internal/controllers/remote/_core.py

@@ -309,29 +309,43 @@ class Controller:
         """
         Attempt to launch an action.
         """
-        if not action.is_started() and action.task is not None:
-            logger.debug(f"Attempting to launch action: {action.name}")
-            try:
+        if not action.is_started():
+            task: queue_service_pb2.TaskAction | None = None
+            trace: queue_service_pb2.TraceAction | None = None
+            if action.type == "task":
+                if action.task is None:
+                    raise flyte.errors.RuntimeSystemError(
+                        "NoTaskSpec", "Task Spec not found, cannot launch Task Action."
+                    )
                 cache_key = None
                 logger.info(f"Action {action.name} has cache version {action.cache_key}")
                 if action.cache_key:
                     cache_key = StringValue(value=action.cache_key)
 
+                task = queue_service_pb2.TaskAction(
+                    id=task_definition_pb2.TaskIdentifier(
+                        version=action.task.task_template.id.version,
+                        org=action.task.task_template.id.org,
+                        project=action.task.task_template.id.project,
+                        domain=action.task.task_template.id.domain,
+                        name=action.task.task_template.id.name,
+                    ),
+                    spec=action.task,
+                    cache_key=cache_key,
+                )
+            else:
+                trace = queue_service_pb2.TraceAction(
+                    name=action.friendly_name,
+                )
+
+            logger.debug(f"Attempting to launch action: {action.name}")
+            try:
                 await self._queue_service.EnqueueAction(
                     queue_service_pb2.EnqueueActionRequest(
                         action_id=action.action_id,
                         parent_action_name=action.parent_action_name,
-                        task=queue_service_pb2.TaskAction(
-                            id=task_definition_pb2.TaskIdentifier(
-                                version=action.task.task_template.id.version,
-                                org=action.task.task_template.id.org,
-                                project=action.task.task_template.id.project,
-                                domain=action.task.task_template.id.domain,
-                                name=action.task.task_template.id.name,
-                            ),
-                            spec=action.task,
-                            cache_key=cache_key,
-                        ),
+                        task=task,
+                        trace=trace,
                         input_uri=action.inputs_uri,
                         run_output_base=action.run_output_base,
                         group=action.group.name if action.group else None,

flyte/_internal/imagebuild/__init__.py

@@ -1,15 +1,3 @@
-import asyncio
-from typing import List
-
-from flyte._image import Image
-from flyte._internal.imagebuild.docker_builder import DockerImageBuilder
 from flyte._internal.imagebuild.image_builder import ImageBuildEngine
-from flyte._internal.imagebuild.remote_builder import RemoteImageBuilder
-
-__all__ = ["DockerImageBuilder", "ImageBuildEngine", "RemoteImageBuilder"]
-
 
-async def build(images: List[Image]) -> List[str]:
-    builder = DockerImageBuilder()
-    ts = [asyncio.create_task(builder.build_image(image)) for image in images]
-    return list(await asyncio.gather(*ts))
+__all__ = ["ImageBuildEngine"]

flyte/_internal/imagebuild/docker_builder.py

@@ -72,7 +72,7 @@ RUN --mount=type=cache,sharing=locked,mode=0777,target=/var/cache/apt,id=apt \
 
 UV_PYTHON_INSTALL_COMMAND = Template("""\
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
-    uv pip install --prerelease=allow $PIP_INSTALL_ARGS
+    uv pip install $PIP_INSTALL_ARGS
 """)
 
 # uv pip install --python /root/env/bin/python

flyte/_internal/runtime/convert.py

@@ -155,9 +155,18 @@ async def convert_from_native_to_outputs(o: Any, interface: NativeInterface, tas
     if not isinstance(o, tuple):
         o = (o,)
 
-    assert len(interface.outputs) == len(interface.outputs), (
-        f"Received {len(o)} outputs but interface has {len(interface.outputs)}"
-    )
+    if len(interface.outputs) == 0:
+        if len(o) != 0:
+            if len(o) == 1 and o[0] is not None:
+                raise flyte.errors.RuntimeDataValidationError(
+                    "o0",
+                    f"Expected no outputs but got {o},did you miss a return type annotation?",
+                    task_name,
+                )
+    else:
+        assert len(o) == len(interface.outputs), (
+            f"Received {len(o)} outputs but return annotation has {len(interface.outputs)} outputs specified. "
+        )
     named = []
     for (output_name, python_type), v in zip(interface.outputs.items(), o):
         try:
@@ -279,6 +288,29 @@ def generate_inputs_hash(serialized_inputs: str | bytes) -> str:
     return hash_data(serialized_inputs)
 
 
+def generate_inputs_hash_from_proto(inputs: run_definition_pb2.Inputs) -> str:
+    """
+    Generate a hash for the inputs. This is used to uniquely identify the inputs for a task.
+    :param inputs: The inputs to hash.
+    :return: A hexadecimal string representation of the hash.
+    """
+    if not inputs:
+        return ""
+    return generate_inputs_hash(inputs.SerializeToString(deterministic=True))
+
+
+def generate_interface_hash(task_interface: interface_pb2.TypedInterface) -> str:
+    """
+    Generate a hash for the task interface. This is used to uniquely identify the task interface.
+    :param task_interface: The interface of the task.
+    :return: A hexadecimal string representation of the hash.
+    """
+    if not task_interface:
+        return ""
+    serialized_interface = task_interface.SerializeToString(deterministic=True)
+    return hash_data(serialized_interface)
+
+
 def generate_cache_key_hash(
     task_name: str,
     inputs_hash: str,
@@ -302,10 +334,13 @@ def generate_cache_key_hash(
     if ignored_input_vars:
         filtered = [named_lit for named_lit in proto_inputs.literals if named_lit.name not in ignored_input_vars]
         final = run_definition_pb2.Inputs(literals=filtered)
-        final_inputs = final.SerializeToString(deterministic=True)
+        final_inputs = generate_inputs_hash_from_proto(final)
     else:
         final_inputs = inputs_hash
-    data = f"{final_inputs}{task_name}{task_interface.SerializeToString(deterministic=True)}{cache_version}"
+
+    interface_hash = generate_interface_hash(task_interface)
+
+    data = f"{final_inputs}{task_name}{interface_hash}{cache_version}"
     return hash_data(data)
 
 

flyte/_internal/runtime/task_serde.py

@@ -143,7 +143,7 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
 
     # TODO Add support for SQL, extra_config, custom
     extra_config: typing.Dict[str, str] = {}
-    custom = {}  # type: ignore
+    custom = task.custom_config(serialize_context)
 
     sql = None
     if task.pod_template and not isinstance(task.pod_template, str):
@@ -220,9 +220,16 @@ def _get_urun_container(
     if not serialize_context.image_cache:
         # This computes the image uri, computing hashes as necessary so can fail if done remotely.
         img_uri = task_template.image.uri
-    elif image_id not in serialize_context.image_cache.image_lookup:
+    elif serialize_context.image_cache and image_id not in serialize_context.image_cache.image_lookup:
         img_uri = task_template.image.uri
-        logger.warning(f"Image uri {img_uri} not found in image_cache {serialize_context.image_cache.image_lookup}")
+        from flyte._version import __version__
+
+        logger.warning(
+            f"Image {task_template.image} not found in the image cache: {serialize_context.image_cache.image_lookup}.\n"
+            f"This typically occurs when the Flyte SDK version (`{__version__}`) used in the task environment "
+            f"differs from the version used to compile or deploy it.\n"
+            f"Ensure both environments use the same Flyte SDK version to avoid inconsistencies in image resolution."
+        )
     else:
         img_uri = serialize_context.image_cache.image_lookup[image_id]
 

flyte/_pod.py

@@ -2,7 +2,8 @@ from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Dict, Optional
 
 if TYPE_CHECKING:
-    from kubernetes.client import V1PodSpec
+    from flyteidl.core.tasks_pb2 import K8sPod
+    from kubernetes.client import ApiClient, V1PodSpec
 
 
 _PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
@@ -17,3 +18,12 @@ class PodTemplate(object):
     primary_container_name: str = _PRIMARY_CONTAINER_DEFAULT_NAME
     labels: Optional[Dict[str, str]] = None
     annotations: Optional[Dict[str, str]] = None
+
+    def to_k8s_pod(self) -> "K8sPod":
+        from flyteidl.core.tasks_pb2 import K8sObjectMetadata, K8sPod
+
+        return K8sPod(
+            metadata=K8sObjectMetadata(labels=self.labels, annotations=self.annotations),
+            pod_spec=ApiClient().sanitize_for_serialization(self.pod_spec),
+            primary_container_name=self.primary_container_name,
+        )

flyte/_resources.py

@@ -1,8 +1,15 @@
-from dataclasses import dataclass
-from typing import Literal, Optional, Tuple, Union, get_args
+from dataclasses import dataclass, fields
+from typing import TYPE_CHECKING, Literal, Optional, Tuple, Union, get_args
 
 import rich.repr
 
+from flyte._pod import _PRIMARY_CONTAINER_DEFAULT_NAME
+
+if TYPE_CHECKING:
+    from kubernetes.client import V1PodSpec
+
+PRIMARY_CONTAINER_DEFAULT_NAME = "primary"
+
 GPUType = Literal["T4", "A100", "A100 80G", "H100", "L4", "L40s"]
 GPUQuantity = Literal[1, 2, 3, 4, 5, 6, 7, 8]
 A100Parts = Literal["1g.5gb", "2g.10gb", "3g.20gb", "4g.20gb", "7g.40gb"]
@@ -224,3 +231,61 @@ class Resources:
         if self.shm == "auto":
             return ""
         return self.shm
+
+
+def _check_resource_is_singular(resource: Resources):
+    """
+    Raise a value error if the resource has a tuple.
+    """
+    for field in fields(resource):
+        value = getattr(resource, field.name)
+        if isinstance(value, (tuple, list)):
+            raise ValueError(f"{value} can not be a list or tuple")
+    return resource
+
+
+def pod_spec_from_resources(
+    primary_container_name: str = _PRIMARY_CONTAINER_DEFAULT_NAME,
+    requests: Optional[Resources] = None,
+    limits: Optional[Resources] = None,
+    k8s_gpu_resource_key: str = "nvidia.com/gpu",
+) -> "V1PodSpec":
+    from kubernetes.client import V1Container, V1PodSpec, V1ResourceRequirements
+
+    def _construct_k8s_pods_resources(resources: Optional[Resources], k8s_gpu_resource_key: str):
+        if resources is None:
+            return None
+
+        resources_map = {
+            "cpu": "cpu",
+            "memory": "memory",
+            "gpu": k8s_gpu_resource_key,
+            "ephemeral_storage": "ephemeral-storage",
+        }
+
+        k8s_pod_resources = {}
+
+        _check_resource_is_singular(resources)
+        for resource in fields(resources):
+            resource_value = getattr(resources, resource.name)
+            if resource_value is not None:
+                k8s_pod_resources[resources_map[resource.name]] = resource_value
+
+        return k8s_pod_resources
+
+    requests = _construct_k8s_pods_resources(resources=requests, k8s_gpu_resource_key=k8s_gpu_resource_key)
+    limits = _construct_k8s_pods_resources(resources=limits, k8s_gpu_resource_key=k8s_gpu_resource_key)
+    requests = requests or limits
+    limits = limits or requests
+
+    return V1PodSpec(
+        containers=[
+            V1Container(
+                name=primary_container_name,
+                resources=V1ResourceRequirements(
+                    requests=requests,
+                    limits=limits,
+                ),
+            )
+        ]
+    )

flyte/_run.py

@@ -6,13 +6,9 @@ import uuid
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
 
 import flyte.errors
-from flyte.errors import InitializationError
-from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath, SerializationContext, TaskContext
-from flyte.syncify import syncify
-
-from ._context import contextual_run, internal_ctx
-from ._environment import Environment
-from ._initialize import (
+from flyte._context import contextual_run, internal_ctx
+from flyte._environment import Environment
+from flyte._initialize import (
     _get_init_config,
     get_client,
     get_common_config,
@@ -20,9 +16,19 @@ from ._initialize import (
     requires_initialization,
     requires_storage,
 )
-from ._logging import logger
-from ._task import P, R, TaskTemplate
-from ._tools import ipython_check
+from flyte._logging import logger
+from flyte._task import P, R, TaskTemplate
+from flyte._tools import ipython_check
+from flyte.errors import InitializationError
+from flyte.models import (
+    ActionID,
+    Checkpoints,
+    CodeBundle,
+    RawDataPath,
+    SerializationContext,
+    TaskContext,
+)
+from flyte.syncify import syncify
 
 if TYPE_CHECKING:
     from flyte.remote import Run
@@ -132,7 +138,9 @@ class _Runner:
 
             if self._interactive_mode:
                 code_bundle = await build_pkl_bundle(
-                    obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
+                    obj,
+                    upload_to_controlplane=not self._dry_run,
+                    copy_bundle_to=self._copy_bundle_to,
                 )
             else:
                 if self._copy_files != "none":
@@ -253,7 +261,8 @@ class _Runner:
                     pb2=run_definition_pb2.Run(
                         action=run_definition_pb2.Action(
                             id=run_definition_pb2.ActionIdentifier(
-                                name="a0", run=run_definition_pb2.RunIdentifier(name="dry-run")
+                                name="a0",
+                                run=run_definition_pb2.RunIdentifier(name="dry-run"),
                             )
                         )
                     )
@@ -286,7 +295,7 @@ class _Runner:
         if obj.parent_env is None:
             raise ValueError("Task is not attached to an environment. Please attach the task to an environment.")
 
-        image_cache = build_images.aio(cast(Environment, obj.parent_env()))
+        image_cache = await build_images.aio(cast(Environment, obj.parent_env()))
 
         code_bundle = None
         if self._name is not None:
@@ -296,7 +305,9 @@ class _Runner:
         if not code_bundle:
             if self._interactive_mode:
                 code_bundle = await build_pkl_bundle(
-                    obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
+                    obj,
+                    upload_to_controlplane=not self._dry_run,
+                    copy_bundle_to=self._copy_bundle_to,
                 )
             else:
                 if self._copy_files != "none":
@@ -381,7 +392,8 @@ class _Runner:
         tctx = TaskContext(
             action=action,
             checkpoints=Checkpoints(
-                prev_checkpoint_path=internal_ctx().raw_data.path, checkpoint_path=internal_ctx().raw_data.path
+                prev_checkpoint_path=internal_ctx().raw_data.path,
+                checkpoint_path=internal_ctx().raw_data.path,
             ),
             code_bundle=None,
             output_path=self._metadata_path,
@@ -403,7 +415,10 @@ class _Runner:
 
     @syncify
     async def run(
-        self, task: TaskTemplate[P, Union[R, Run]] | LazyEntity, *args: P.args, **kwargs: P.kwargs
+        self,
+        task: TaskTemplate[P, Union[R, Run]] | LazyEntity,
+        *args: P.args,
+        **kwargs: P.kwargs,
     ) -> Union[R, Run]:
         """
         Run an async `@env.task` or `TaskTemplate` instance. The existing async context will be used.
@@ -430,6 +445,10 @@ class _Runner:
 
         if isinstance(task, LazyEntity) and self._mode != "remote":
             raise ValueError("Remote task can only be run in remote mode.")
+
+        if not isinstance(task, TaskTemplate) and not isinstance(task, LazyEntity):
+            raise TypeError("On Flyte tasks can be run, not generic functions or methods.")
+
         if self._mode == "remote":
             return await self._run_remote(task, *args, **kwargs)
         task = cast(TaskTemplate, task)

flyte/_task.py

@@ -361,6 +361,7 @@ class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
     """
 
     func: FunctionTypes
+    plugin_config: Optional[Any] = None  # This is used to pass plugin specific configuration
 
     def __post_init__(self):
         super().__post_init__()

flyte/_task_environment.py

@@ -60,8 +60,9 @@ class TaskEnvironment(Environment):
     :param reusable: Reuse policy for the environment, if set, a python process may be reused for multiple tasks.
     """
 
-    cache: Union[CacheRequest] = "auto"
+    cache: Union[CacheRequest] = "disable"
     reusable: ReusePolicy | None = None
+    plugin_config: Optional[Any] = None
     # TODO Shall we make this union of string or env? This way we can lookup the env by module/file:name
     # TODO also we could add list of files that are used by this environment
 
@@ -133,6 +134,8 @@ class TaskEnvironment(Environment):
         used.
         :param report: Optional Whether to generate the html report for the task, defaults to False.
         """
+        from ._task import P, R
+
         if self.reusable is not None:
             if pod_template is not None:
                 raise ValueError("Cannot set pod_template when environment is reusable.")
@@ -141,7 +144,22 @@ class TaskEnvironment(Environment):
             friendly_name = name or func.__name__
             task_name = self.name + "." + func.__name__
 
-            tmpl: AsyncFunctionTaskTemplate = AsyncFunctionTaskTemplate(
+            if self.plugin_config is not None:
+                from flyte.extend import TaskPluginRegistry
+
+                task_template_class: type[AsyncFunctionTaskTemplate[P, R]] | None = TaskPluginRegistry.find(
+                    config_type=type(self.plugin_config)
+                )
+                if task_template_class is None:
+                    raise ValueError(
+                        f"No task plugin found for config type {type(self.plugin_config)}. "
+                        f"Please register a plugin using flyte.extend.TaskPluginRegistry.register() api."
+                    )
+            else:
+                task_template_class = AsyncFunctionTaskTemplate
+
+            task_template_class = cast(type[AsyncFunctionTaskTemplate[P, R]], task_template_class)
+            tmpl = task_template_class(
                 func=func,
                 name=task_name,
                 image=self.image,
@@ -158,6 +176,7 @@ class TaskEnvironment(Environment):
                 interface=NativeInterface.from_callable(func),
                 report=report,
                 friendly_name=friendly_name,
+                plugin_config=self.plugin_config,
             )
             self._tasks[task_name] = tmpl
             return tmpl

flyte/_task_plugins.py

@@ -0,0 +1,45 @@
+import typing
+from typing import Type
+
+import rich.repr
+
+if typing.TYPE_CHECKING:
+    from ._task import AsyncFunctionTaskTemplate
+
+T = typing.TypeVar("T", bound="AsyncFunctionTaskTemplate")
+
+
+class _Registry:
+    """
+    A registry for task plugins.
+    """
+
+    def __init__(self):
+        self._plugins: typing.Dict[Type, Type[T]] = {}
+
+    def register(self, config_type: Type, plugin: Type[T]):
+        """
+        Register a plugin.
+        """
+        self._plugins[config_type] = plugin
+
+    def find(self, config_type: Type) -> typing.Optional[Type[T]]:
+        """
+        Get a plugin by name.
+        """
+        return self._plugins.get(config_type)
+
+    def list_plugins(self):
+        """
+        List all registered plugins.
+        """
+        return list(self._plugins.keys())
+
+    def __rich_repr__(self) -> "rich.repr.Result":
+        yield from (("Name", i) for i in self.list_plugins())
+
+    def __repr__(self):
+        return f"TaskPluginRegistry(plugins={self.list_plugins()})"
+
+
+TaskPluginRegistry = _Registry()

flyte/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.0b22'
-__version_tuple__ = version_tuple = (0, 2, 0, 'b22')
+__version__ = version = '0.2.0b24'
+__version_tuple__ = version_tuple = (0, 2, 0, 'b24')

flyte/cli/_common.py

@@ -151,7 +151,9 @@ class InvokeBaseMixin:
                 # If the user has requested verbose output, print the full traceback
                 console = Console()
                 console.print(Traceback.from_exception(type(e), e, e.__traceback__))
-            raise click.ClickException(f"Error invoking command: {e}") from e
+                exit(1)
+            else:
+                raise click.ClickException(f"Error invoking command: {e}") from e
 
 
 class CommandBase(InvokeBaseMixin, click.RichCommand):

flyte/errors.py

@@ -157,9 +157,9 @@ class RuntimeDataValidationError(RuntimeUserError):
     This error is raised when the user tries to access a resource that does not exist or is invalid.
     """
 
-    def __init__(self, var: str, e: Exception, task_name: str = ""):
+    def __init__(self, var: str, e: Exception | str, task_name: str = ""):
         super().__init__(
-            "DataValiationError", f"In task {task_name} variable {var}, failed to serialize/deserialize because {e}"
+            "DataValiationError", f"In task {task_name} variable {var}, failed to serialize/deserialize because of {e}"
         )
 
 

flyte/extend.py

@@ -0,0 +1,12 @@
+from ._initialize import is_initialized
+from ._resources import PRIMARY_CONTAINER_DEFAULT_NAME, pod_spec_from_resources
+from ._task import AsyncFunctionTaskTemplate
+from ._task_plugins import TaskPluginRegistry
+
+__all__ = [
+    "PRIMARY_CONTAINER_DEFAULT_NAME",
+    "AsyncFunctionTaskTemplate",
+    "TaskPluginRegistry",
+    "is_initialized",
+    "pod_spec_from_resources",
+]

flyte/types/_type_engine.py

@@ -1315,7 +1315,8 @@ class TypeEngine(typing.Generic[T]):
             try:
                 return transformer.guess_python_type(flyte_type)
             except ValueError:
-                logger.debug(f"Skipping transformer {transformer.name} for {flyte_type}")
+                # Skipping transformer
+                continue
 
         # Because the dataclass transformer is handled explicitly in the get_transformer code, we have to handle it
         # separately here too.
@@ -1438,6 +1439,19 @@ def _type_essence(x: types_pb2.LiteralType) -> types_pb2.LiteralType:
 
 
 def _are_types_castable(upstream: types_pb2.LiteralType, downstream: types_pb2.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.HasField("collection_type"):
         if not downstream.HasField("collection_type"):
             return False
@@ -1483,7 +1497,7 @@ def _are_types_castable(upstream: types_pb2.LiteralType, downstream: types_pb2.L
 
         return True
 
-    if upstream.HasField("union_type"):
+    if upstream.HasField("union_type") and 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):

flyte-0.2.0b24.data/scripts/runtime.py

@@ -0,0 +1,169 @@
+"""
+Flyte runtime module, this is the entrypoint script for the Flyte runtime.
+
+Caution: Startup time for this module is very important, as it is the entrypoint for the Flyte runtime.
+Refrain from importing any modules here. If you need to import any modules, do it inside the main function.
+"""
+
+import asyncio
+import os
+import sys
+from typing import Any, List
+
+import click
+
+# Todo: work with pvditt to make these the names
+# ACTION_NAME = "_U_ACTION_NAME"
+# RUN_NAME = "_U_RUN_NAME"
+# PROJECT_NAME = "_U_PROJECT_NAME"
+# DOMAIN_NAME = "_U_DOMAIN_NAME"
+# ORG_NAME = "_U_ORG_NAME"
+
+ACTION_NAME = "ACTION_NAME"
+RUN_NAME = "RUN_NAME"
+PROJECT_NAME = "FLYTE_INTERNAL_TASK_PROJECT"
+DOMAIN_NAME = "FLYTE_INTERNAL_TASK_DOMAIN"
+ORG_NAME = "_U_ORG_NAME"
+ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
+RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
+ENABLE_REF_TASKS = "_REF_TASKS"  # This is a temporary flag to enable reference tasks in the runtime.
+
+# TODO: Remove this after proper auth is implemented
+_UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
+
+
+@click.group()
+def _pass_through():
+    pass
+
+
+@_pass_through.command("a0")
+@click.option("--inputs", "-i", required=True)
+@click.option("--outputs-path", "-o", required=True)
+@click.option("--version", "-v", required=True)
+@click.option("--run-base-dir", envvar=RUN_OUTPUT_BASE_DIR, required=True)
+@click.option("--raw-data-path", "-r", required=False)
+@click.option("--checkpoint-path", "-c", required=False)
+@click.option("--prev-checkpoint", "-p", required=False)
+@click.option("--name", envvar=ACTION_NAME, required=False)
+@click.option("--run-name", envvar=RUN_NAME, required=False)
+@click.option("--project", envvar=PROJECT_NAME, required=False)
+@click.option("--domain", envvar=DOMAIN_NAME, required=False)
+@click.option("--org", envvar=ORG_NAME, required=False)
+@click.option("--image-cache", required=False)
+@click.option("--tgz", required=False)
+@click.option("--pkl", required=False)
+@click.option("--dest", required=False)
+@click.option("--resolver", required=False)
+@click.argument(
+    "resolver-args",
+    type=click.UNPROCESSED,
+    nargs=-1,
+)
+def main(
+    run_name: str,
+    name: str,
+    project: str,
+    domain: str,
+    org: str,
+    image_cache: str,
+    version: str,
+    inputs: str,
+    run_base_dir: str,
+    outputs_path: str,
+    raw_data_path: str,
+    checkpoint_path: str,
+    prev_checkpoint: str,
+    tgz: str,
+    pkl: str,
+    dest: str,
+    resolver: str,
+    resolver_args: List[str],
+):
+    sys.path.insert(0, ".")
+
+    import flyte
+    import flyte._utils as utils
+    from flyte._initialize import init
+    from flyte._internal.controllers import create_controller
+    from flyte._internal.imagebuild.image_builder import ImageCache
+    from flyte._internal.runtime.entrypoints import load_and_run_task
+    from flyte._logging import logger
+    from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
+
+    logger.info(f"Initializing flyte runtime - version {flyte.__version__}")
+
+    assert org, "Org is required for now"
+    assert project, "Project is required"
+    assert domain, "Domain is required"
+    assert run_name, f"Run name is required {run_name}"
+    assert name, f"Action name is required {name}"
+
+    if run_name.startswith("{{"):
+        run_name = os.getenv("RUN_NAME", "")
+    if name.startswith("{{"):
+        name = os.getenv("ACTION_NAME", "")
+
+    # Figure out how to connect
+    # This detection of api key is a hack for now.
+    controller_kwargs: dict[str, Any] = {"insecure": False}
+    if api_key := os.getenv(_UNION_EAGER_API_KEY_ENV_VAR):
+        logger.info("Using api key from environment")
+        controller_kwargs["api_key"] = api_key
+    else:
+        ep = os.environ.get(ENDPOINT_OVERRIDE, "host.docker.internal:8090")
+        controller_kwargs["endpoint"] = ep
+        if "localhost" in ep or "docker" in ep:
+            controller_kwargs["insecure"] = True
+        logger.debug(f"Using controller endpoint: {ep} with kwargs: {controller_kwargs}")
+
+    bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
+    enable_ref_tasks = os.getenv(ENABLE_REF_TASKS, "false").lower() in ("true", "1", "yes")
+    # We init regular client here so that reference tasks can work
+    # Current reference tasks will not work with remote controller, because we create 2 different
+    # channels on different threads and this is not supported by grpcio or the auth system. It ends up leading
+    # File "src/python/grpcio/grpc/_cython/_cygrpc/aio/completion_queue.pyx.pxi", line 147,
+    # in grpc._cython.cygrpc.PollerCompletionQueue._handle_events
+    # BlockingIOError: [Errno 11] Resource temporarily unavailable
+    # TODO solution is to use a single channel for both controller and reference tasks, but this requires a refactor
+    if enable_ref_tasks:
+        logger.warning(
+            "Reference tasks are enabled. This will initialize client and you will see a BlockIOError. "
+            "This is harmless, but a nuisance. We are working on a fix."
+        )
+        init(org=org, project=project, domain=domain, **controller_kwargs)
+    else:
+        init()
+    # Controller is created with the same kwargs as init, so that it can be used to run tasks
+    controller = create_controller(ct="remote", **controller_kwargs)
+
+    ic = ImageCache.from_transport(image_cache) if image_cache else None
+
+    # Create a coroutine to load the task and run it
+    task_coroutine = load_and_run_task(
+        resolver=resolver,
+        resolver_args=resolver_args,
+        action=ActionID(name=name, run_name=run_name, project=project, domain=domain, org=org),
+        raw_data_path=RawDataPath(path=raw_data_path),
+        checkpoints=Checkpoints(checkpoint_path, prev_checkpoint),
+        code_bundle=bundle,
+        input_path=inputs,
+        output_path=outputs_path,
+        run_base_dir=run_base_dir,
+        version=version,
+        controller=controller,
+        image_cache=ic,
+    )
+    # Create a coroutine to watch for errors
+    controller_failure = controller.watch_for_errors()
+
+    # Run both coroutines concurrently and wait for first to finish and cancel the other
+    async def _run_and_stop():
+        await utils.run_coros(controller_failure, task_coroutine)
+        await controller.stop()
+
+    asyncio.run(_run_and_stop())
+
+
+if __name__ == "__main__":
+    _pass_through()

{flyte-0.2.0b22.dist-info → flyte-0.2.0b24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flyte
-Version: 0.2.0b22
+Version: 0.2.0b24
 Summary: Add your description here
 Author-email: Ketan Umare <kumare3@users.noreply.github.com>
 Requires-Python: >=3.10

{flyte-0.2.0b22.dist-info → flyte-0.2.0b24.dist-info}/RECORD

@@ -1,35 +1,37 @@
-flyte/__init__.py,sha256=VjIdD4AyX7mYNoWDSqZ6tmNKo14ZFwfToKeb49SB60U,1733
+flyte/__init__.py,sha256=jWSynBJyJ0WuSjE8HID2MT386ZNlRzx7LzCHW1y_XNw,1468
 flyte/_build.py,sha256=MkgfLAPeL56YeVrGRNZUCZgbwzlEzVP3wLbl5Qru4yk,578
 flyte/_context.py,sha256=K0-TCt-_pHOoE5Xni87_8uIe2vCBOhfNQEtjGT4Hu4k,5239
 flyte/_deploy.py,sha256=HU2ksZUvHj77DrJm7MryN0n26DUJqo4h7-op4fUTBUg,9593
 flyte/_doc.py,sha256=_OPCf3t_git6UT7kSJISFaWO9cfNzJhhoe6JjVdyCJo,706
 flyte/_docstring.py,sha256=SsG0Ab_YMAwy2ABJlEo3eBKlyC3kwPdnDJ1FIms-ZBQ,1127
-flyte/_environment.py,sha256=dmIFmFLRIklOEYL9gsP2IH3-MYcjHYyyOlqlcf3E6_A,2924
+flyte/_environment.py,sha256=oKVXLBX0ky2eE_wjBdzvQGI_2LiT2Nbx58ur7GMt50c,3231
 flyte/_excepthook.py,sha256=nXts84rzEg6-7RtFarbKzOsRZTQR4plnbWVIFMAEprs,1310
 flyte/_group.py,sha256=7o1j16sZyUmYB50mOiq1ui4TBAKhRpDqLakV8Ya1kw4,803
 flyte/_hash.py,sha256=Of_Zl_DzzzF2jp4ZsLm-3o-xJFCCJ8_GubmLI1htx78,504
-flyte/_image.py,sha256=6d_bxy2KOn0_wj895iPOxXKN5uzygUeT2fOhxG95dg4,30633
+flyte/_image.py,sha256=ksG_BuDqQyUrrgkoI-I2FWvHvj0Zt8h1x9FAO2DryXI,31439
 flyte/_initialize.py,sha256=xKl_LYMluRt21wWqa6RTKuLo0_DCbSaTfUk27_brtNk,18232
 flyte/_interface.py,sha256=1B9zIwFDjiVp_3l_mk8EpA4g3Re-6DUBEBi9z9vDvPs,3504
 flyte/_logging.py,sha256=_yNo-Nx2yzh0MLoZGbnIYHGKei4wtQmSGM0lE30Ev7w,3662
 flyte/_map.py,sha256=efPd8O-JKUg1OY3_MzU3KGbhsGYDVRNBwWr0ceNIXhQ,7444
-flyte/_pod.py,sha256=lNaQuWX22QG6Xji7-8GpuKUkqCmVFaRxOVU-eUEa-Vk,637
-flyte/_resources.py,sha256=UOLyEVhdxolvrHhddiBbYdJuE1RkM_l7xeS9G1abe6M,7583
+flyte/_pod.py,sha256=--72b0c6IkOEbBwZPLmgl-ll-j7ECfG-kh75LzBnNN8,1068
+flyte/_resources.py,sha256=L2JuvQDlMo1JLJeUmJPRwtWbunhR2xJEhFgQW5yc72c,9690
 flyte/_retry.py,sha256=rfLv0MvWxzPByKESTglEmjPsytEAKiIvvmzlJxXwsfE,941
 flyte/_reusable_environment.py,sha256=P4FBATVKAYcIKpdFN98sI8acPyKy8eIGx6V0kUb9YdM,1289
-flyte/_run.py,sha256=HDObXGfZKBafebuh1OR5e1NWLUSbBDcceIz3kX2GMss,22828
+flyte/_run.py,sha256=ZdpnWNkMTWLA7QZLzczaBdcgWfzv6DMMUAe5OQdv_Ik,23227
 flyte/_secret.py,sha256=SqIHs6mi8hEkIIBZe3bI9jJsPt65Mt6dV5uh9_op1ME,2392
-flyte/_task.py,sha256=-3VLSROj8g3-ZWzV272Ww7mt5yIeA753kmpNyr75I58,18087
-flyte/_task_environment.py,sha256=B493b54xpnNFNtHek3cmT4XuNiPNVf2gx9Yylyb1ncw,6905
+flyte/_task.py,sha256=uXKhOn5Avz4TmUUlDbJ5f-XnOr1C4ps-lB8aP5celKE,18181
+flyte/_task_environment.py,sha256=u2JZCtP1A3C_xCILQGFkfRdPgsyjDSORcCznZFkUZ2E,7779
+flyte/_task_plugins.py,sha256=9MH3nFPOH_e8_92BT4sFk4oyAnj6GJFvaPYWaraX7yE,1037
 flyte/_timeout.py,sha256=zx5sFcbYmjJAJbZWSGzzX-BpC9HC7Jfs35T7vVhKwkk,1571
 flyte/_tools.py,sha256=JewkQZBR_M85tS6QY8e4xXue75jbOE48nID4ZHnc9jY,632
 flyte/_trace.py,sha256=C788bgoSc3st8kE8Cae2xegnLx2CT6uuRKKfaDrDUys,5122
-flyte/_version.py,sha256=EDrKK4XJyvTSZe8-ZzLNFdqQWtheoqh-sXUajAe6Zug,521
-flyte/errors.py,sha256=skXcdexLisFZFcTnUmMvMmuh4Ty96oJkyLKaipzkyeI,4954
+flyte/_version.py,sha256=FFwCuNAGsDblVk5R0UiFXYSEp_KeGkD1OMVoRLzyNbc,521
+flyte/errors.py,sha256=IUXHuEgszis9ZyhBUOSSlZx7VPQk0BF4N875yj1FBKo,4963
+flyte/extend.py,sha256=GB4ZedGzKa30vYWRVPOdxEeK62xnUVFY4z2tD6H9eEw,376
 flyte/models.py,sha256=my7Vxo-NK6gHGahyqtHr42wYjPGw0nl2SGBBoSb6hIc,14733
 flyte/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 flyte/_bin/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-flyte/_bin/runtime.py,sha256=_ADnpWY4nVvjaujfcx2pBpOskAvnm0WtXBf1uPnZIAQ,6230
+flyte/_bin/runtime.py,sha256=kiRtbRxfFy2zQn6T2-Qfnb3BYopt7FIXBjSCBO63UL4,6334
 flyte/_cache/__init__.py,sha256=zhdO5UuHQRdzn8GHmSN40nrxfAmI4ihDRuHZM11U84Y,305
 flyte/_cache/cache.py,sha256=ErhWzzJdEjTIuEF4f-r6IBgko-3Al9iUs1Eq4O42TUE,5021
 flyte/_cache/defaults.py,sha256=gzJZW0QJPUfd2OPnGpv3tzIfwPtgFjAKoie3NP1P97U,217
@@ -44,14 +46,14 @@ flyte/_internal/controllers/__init__.py,sha256=5CBnS9lb1VFMzZuRXUiaPhlN3G9qh7Aq9
 flyte/_internal/controllers/_local_controller.py,sha256=Wpgtd50C_ovIHpQSZC6asQc7iKyKIraEf-MAHCwcNJI,7124
 flyte/_internal/controllers/_trace.py,sha256=biI-lXSIe3gXuWI-KT6T-jTtojQCQ7BLOHTCG3J6MQc,1145
 flyte/_internal/controllers/remote/__init__.py,sha256=9_azH1eHLqY6VULpDugXi7Kf1kK1ODqEnsQ_3wM6IqU,1919
-flyte/_internal/controllers/remote/_action.py,sha256=aGmTCciHkx9YSU6gPWWFsgjIan-mgelHfmZM4SzZJS4,5121
+flyte/_internal/controllers/remote/_action.py,sha256=tB3GfcB4-PcqaupeoFKxqkodFuNCaruHbXHFGkPaTFM,6040
 flyte/_internal/controllers/remote/_client.py,sha256=HPbzbfaWZVv5wpOvKNtFXR6COiZDwd1cUJQqi60A7oU,1421
-flyte/_internal/controllers/remote/_controller.py,sha256=G4gEUl-2f4FqACz3IK7sK0-w93kkJndzkDpbISqo_4w,21247
-flyte/_internal/controllers/remote/_core.py,sha256=UKjiL3dGidXcR1dPzJswn67g38HQmVXV7n0zoX_8AZY,18422
+flyte/_internal/controllers/remote/_controller.py,sha256=hTy1cfCVBUxBtQg4ibTzL6RUiJyVfO2xTmMU7fZzXNY,22692
+flyte/_internal/controllers/remote/_core.py,sha256=FJe1ZAWu_0w5SQUFgAHY4Hjjf_AjuI8sYu_orFkPFeU,18889
 flyte/_internal/controllers/remote/_informer.py,sha256=StiPcQLLW0h36uEBhKsupMY79EeFCKA3QQzvv2IyvRo,14188
 flyte/_internal/controllers/remote/_service_protocol.py,sha256=B9qbIg6DiGeac-iSccLmX_AL2xUgX4ezNUOiAbSy4V0,1357
-flyte/_internal/imagebuild/__init__.py,sha256=4fPz5RN3QB1DGw5O2sZo7x71F3EWhyE1jmCCqy4iezk,573
-flyte/_internal/imagebuild/docker_builder.py,sha256=LIfqtLQ4NVxkjgWlASiIi_xqvxhztlUDTVt4P7McLDo,15946
+flyte/_internal/imagebuild/__init__.py,sha256=dwXdJ1jMhw9RF8itF7jkPLanvX1yCviSns7hE5eoIts,102
+flyte/_internal/imagebuild/docker_builder.py,sha256=4-AJLX9CgPmVblzV-caIn3uvHLIlBMntOUKUmC-ptM8,15927
 flyte/_internal/imagebuild/image_builder.py,sha256=dXBXl62qcPabus6dR3eP8P9mBGNhpZHZ2Xm12AymKkk,11150
 flyte/_internal/imagebuild/remote_builder.py,sha256=wm-D3xhdrk8OSQxeyI0nfMpw3P1gW7mqDKAOKOnwY6o,10089
 flyte/_internal/resolvers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -59,11 +61,11 @@ flyte/_internal/resolvers/_task_module.py,sha256=jwy1QYygUK7xmpCZLt1SPTfJCkfox3C
 flyte/_internal/resolvers/common.py,sha256=ADQLRoyGsJ4vuUkitffMGrMKKjy0vpk6X53g4FuKDLc,993
 flyte/_internal/resolvers/default.py,sha256=nX4DHUYod1nRvEsl_vSgutQVEdExu2xL8pRkyi4VWbY,981
 flyte/_internal/runtime/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-flyte/_internal/runtime/convert.py,sha256=rvKXEDQxFLyAslSJeOrqeG1kkXaYxHWxkBMloaZNCPc,13795
+flyte/_internal/runtime/convert.py,sha256=W0nMLZU3GkTGtIg5ZAz0ioIk3aitbaK2M2qbSe26VDQ,15079
 flyte/_internal/runtime/entrypoints.py,sha256=Kyi19i7LYk7YM3ZV_Y4FXGt5Pc1tIftGkIDohopblyY,5127
 flyte/_internal/runtime/io.py,sha256=Lgdy4iPjlKjUO-V_AkoPZff6lywaFjZUG-PErRukmx4,4248
 flyte/_internal/runtime/resources_serde.py,sha256=TObMVsSjVcQhcY8-nY81pbvrz7TP-adDD5xV-LqAaxM,4813
-flyte/_internal/runtime/task_serde.py,sha256=wjYZSr0aXCvo6egWLlTrrRgYbx-V6L1clY2jstJNQlk,14039
+flyte/_internal/runtime/task_serde.py,sha256=x24z2U_StqeavhQ5U5ljrHLMlSOjz-jOrLPliqOKrxM,14479
 flyte/_internal/runtime/taskrunner.py,sha256=rHWS4t5qgZnzGdGrs0_O0sSs_PVGoE1CNPDb-fTwwmo,7332
 flyte/_internal/runtime/types_serde.py,sha256=EjRh9Yypx9-20XXQprtNgp766LeQVRoYWtY6XPGMZQg,1813
 flyte/_protos/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -153,7 +155,7 @@ flyte/_utils/uv_script_parser.py,sha256=PxqD8lSMi6xv0uDd1s8LKB2IPZr4ttZJCUweqlyM
 flyte/cli/__init__.py,sha256=aeCcumeP9xD_5aCmaRYUPCe2QRJSGCaxcUbTZ3co768,341
 flyte/cli/_abort.py,sha256=Ty-63Gtd2PUn6lCuL5AaasfBoPu7TDSU5EQKVbkF4qw,661
 flyte/cli/_build.py,sha256=SBgybTVWOZ22VBHFL8CVFB_oo34lF9wvlwNirYFFyk0,3543
-flyte/cli/_common.py,sha256=CFBmdc0c59XqWb7PuZfLr7Z55HXnlqKaxDCq1sgY9MU,11115
+flyte/cli/_common.py,sha256=prQPs8_qSNCvhBPgHa6AY8c9OwutWw3js9jd5vdL_kk,11161
 flyte/cli/_create.py,sha256=8U4kePYmfW02j54184IurNYiTGvjdr3J-ERQM9goHdg,4715
 flyte/cli/_delete.py,sha256=VTmXv09PBjkdtyl23mbSjIQQlN7Y1AI_bO0GkHP-f9E,546
 flyte/cli/_deploy.py,sha256=Zxm7vn1zrqmll73CJTiVGYJI95P-XBI1AlhOlmbmkD0,4635
@@ -219,10 +221,11 @@ flyte/types/_interface.py,sha256=5y9EC5r897xz03Hh0ltF8QVGKMfMfAznws-hKSEO4Go,167
 flyte/types/_pickle.py,sha256=PjdR66OTDMZ3OYq6GvM_Ua0cIo5t2XQaIjmpJ9xo4Ys,4050
 flyte/types/_renderer.py,sha256=ygcCo5l60lHufyQISFddZfWwLlQ8kJAKxUT_XnR_6dY,4818
 flyte/types/_string_literals.py,sha256=NlG1xV8RSA-sZ-n-IFQCAsdB6jXJOAKkHWtnopxVVDk,4231
-flyte/types/_type_engine.py,sha256=CCjpqXNX2BMza2cKq42hJXwabWy8GWsimsgiGZJ_00M,96583
+flyte/types/_type_engine.py,sha256=qIrXXBFqv4CydzLW-9iBeTKK2Xo08DHtEYbF9iYDkbU,97087
 flyte/types/_utils.py,sha256=pbts9E1_2LTdLygAY0UYTLYJ8AsN3BZyviSXvrtcutc,2626
-flyte-0.2.0b22.dist-info/METADATA,sha256=RAb1kOFofWEX_kiGOtmRH3HNRfNOu_1iUznOooQsrUI,5857
-flyte-0.2.0b22.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-flyte-0.2.0b22.dist-info/entry_points.txt,sha256=MIq2z5dBurdCJfpXfMKzgBv7sJOakKRYxr8G0cMiTrg,75
-flyte-0.2.0b22.dist-info/top_level.txt,sha256=7dkyFbikvA12LEZEqawx8oDG1CMod6hTliPj7iWzgYo,6
-flyte-0.2.0b22.dist-info/RECORD,,
+flyte-0.2.0b24.data/scripts/runtime.py,sha256=kiRtbRxfFy2zQn6T2-Qfnb3BYopt7FIXBjSCBO63UL4,6334
+flyte-0.2.0b24.dist-info/METADATA,sha256=fi8nuO9osztTEnuZkk9tVSxMWKP3sC2xqgIZVJ_aW44,5857
+flyte-0.2.0b24.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+flyte-0.2.0b24.dist-info/entry_points.txt,sha256=MIq2z5dBurdCJfpXfMKzgBv7sJOakKRYxr8G0cMiTrg,75
+flyte-0.2.0b24.dist-info/top_level.txt,sha256=7dkyFbikvA12LEZEqawx8oDG1CMod6hTliPj7iWzgYo,6
+flyte-0.2.0b24.dist-info/RECORD,,