flyte 0.2.0b23__py3-none-any.whl → 0.2.0b25__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
@@ -353,16 +354,25 @@ 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,
@@ -389,14 +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:
-            if Version(flyte_version).is_devrelease or Version(flyte_version).is_prerelease:
-                image = image.with_pip_packages(f"flyte=={flyte_version}", pre=True)
+        if install_flyte:
+            if dev_mode:
+                image = image.with_local_v2()
             else:
-                image = image.with_pip_packages(f"flyte=={flyte_version}")
+                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.
@@ -416,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:
@@ -425,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
 
@@ -433,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)
@@ -537,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:
         """
@@ -544,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(
@@ -563,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,
         )
 
@@ -713,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:
@@ -725,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(
@@ -777,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

@@ -456,6 +456,10 @@ class DockerImageBuilder(ImageBuilder):
             else:
                 click.secho(f"Run command: {concat_command} ", fg="blue")
 
-            await asyncio.to_thread(subprocess.run, command, check=True)
+            try:
+                await asyncio.to_thread(subprocess.run, command, check=True)
+            except subprocess.CalledProcessError as e:
+                logger.error(f"Failed to build image: {e}")
+                raise RuntimeError(f"Failed to build image: {e}")
 
             return image.uri

flyte/_internal/imagebuild/remote_builder.py

@@ -138,9 +138,17 @@ async def _validate_configuration(image: Image) -> Tuple[str, Optional[str]]:
 
     if any(context_path.iterdir()):
         # If there are files in the context directory, upload it
-        _, context_url = await remote.upload_file.aio(
-            Path(shutil.make_archive(str(tmp_path / "context"), "xztar", context_path))
-        )
+        archive = Path(shutil.make_archive(str(tmp_path / "context"), "xztar", context_path))
+        st = archive.stat()
+        if st.st_size > 5 * 1024 * 1024:
+            logger.warning(
+                click.style(
+                    f"Context size is {st.st_size / (1024 * 1024):.2f} MB, which is larger than 5 MB. "
+                    "Upload and build speed will be impacted.",
+                    fg="yellow",
+                )
+            )
+        _, context_url = await remote.upload_file.aio(archive)
     else:
         context_url = ""
 

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/entrypoints.py

@@ -1,4 +1,5 @@
-from typing import List, Optional, Tuple
+import importlib
+from typing import List, Optional, Tuple, Type
 
 import flyte.errors
 from flyte._code_bundle import download_bundle
@@ -10,7 +11,6 @@ from flyte._task import TaskTemplate
 from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
 
 from .convert import Error, Inputs, Outputs
-from .task_serde import load_task
 from .taskrunner import (
     convert_and_run,
     extract_download_run_upload,
@@ -51,25 +51,68 @@ async def direct_dispatch(
     )
 
 
+def load_class(qualified_name) -> Type:
+    """
+    Load a class from a qualified name. The qualified name should be in the format 'module.ClassName'.
+    :param qualified_name: The qualified name of the class to load.
+    :return: The class object.
+    """
+    module_name, class_name = qualified_name.rsplit(".", 1)  # Split module and class
+    module = importlib.import_module(module_name)  # Import the module
+    return getattr(module, class_name)  # Retrieve the class
+
+
+def load_task(resolver: str, *resolver_args: str) -> TaskTemplate:
+    """
+    Load a task from a resolver. This is a placeholder function.
+
+    :param resolver: The resolver to use to load the task.
+    :param resolver_args: Arguments to pass to the resolver.
+    :return: The loaded task.
+    """
+    resolver_class = load_class(resolver)
+    resolver_instance = resolver_class()
+    return resolver_instance.load_task(resolver_args)
+
+
+def load_pkl_task(code_bundle: CodeBundle) -> TaskTemplate:
+    """
+    Loads a task from a pickled code bundle.
+    :param code_bundle: The code bundle to load the task from.
+    :return: The loaded task template.
+    """
+    logger.debug(f"Loading task from pkl: {code_bundle.downloaded_path}")
+    try:
+        import gzip
+
+        import cloudpickle
+
+        with gzip.open(str(code_bundle.downloaded_path), "rb") as f:
+            return cloudpickle.load(f)
+    except Exception as e:
+        logger.exception(f"Failed to load pickled task from {code_bundle.downloaded_path}. Reason: {e!s}")
+        raise
+
+
+async def download_code_bundle(code_bundle: CodeBundle) -> CodeBundle:
+    """
+    Downloads the code bundle if it is not already downloaded.
+    :param code_bundle: The code bundle to download.
+    :return: The code bundle with the downloaded path.
+    """
+    logger.debug(f"Downloading {code_bundle}")
+    downloaded_path = await download_bundle(code_bundle)
+    return code_bundle.with_downloaded_path(downloaded_path)
+
+
 async def _download_and_load_task(
     code_bundle: CodeBundle | None, resolver: str | None = None, resolver_args: List[str] | None = None
 ) -> TaskTemplate:
     if code_bundle and (code_bundle.tgz or code_bundle.pkl):
         logger.debug(f"Downloading {code_bundle}")
-        downloaded_path = await download_bundle(code_bundle)
-        code_bundle = code_bundle.with_downloaded_path(downloaded_path)
+        code_bundle = await download_code_bundle(code_bundle)
         if code_bundle.pkl:
-            try:
-                logger.debug(f"Loading task from pkl: {code_bundle.downloaded_path}")
-                import gzip
-
-                import cloudpickle
-
-                with gzip.open(str(code_bundle.downloaded_path), "rb") as f:
-                    return cloudpickle.load(f)
-            except Exception as e:
-                logger.exception(f"Failed to load pickled task from {code_bundle.downloaded_path}. Reason: {e!s}")
-                raise
+            return load_pkl_task(code_bundle)
 
         if not resolver or not resolver_args:
             raise flyte.errors.RuntimeSystemError(