flyte 2.0.0b17__py3-none-any.whl → 2.0.0b18__py3-none-any.whl

flyte/_bin/runtime.py

@@ -101,6 +101,7 @@ def main(
     from flyte._logging import logger
     from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
 
+    logger.warning(f"Flyte runtime started for action {name} with run name {run_name}")
     logger.info("Registering faulthandler for SIGUSR1")
     faulthandler.register(signal.SIGUSR1)
 
@@ -168,6 +169,7 @@ def main(
         await controller.stop()
 
     asyncio.run(_run_and_stop())
+    logger.warning(f"Flyte runtime completed for action {name} with run name {run_name}")
 
 
 if __name__ == "__main__":

flyte/_debug/vscode.py

@@ -8,6 +8,7 @@ import subprocess
 import sys
 import tarfile
 import time
+from pathlib import Path
 from typing import List
 
 import aiofiles
@@ -173,7 +174,7 @@ def prepare_launch_json(ctx: click.Context, pid: int):
     Generate the launch.json and settings.json for users to easily launch interactive debugging and task resumption.
     """
 
-    virtual_venv = os.getenv("VIRTUAL_ENV")
+    virtual_venv = os.getenv("VIRTUAL_ENV", str(Path(sys.executable).parent.parent))
     if virtual_venv is None:
         raise RuntimeError("VIRTUAL_ENV is not found in environment variables.")
 
@@ -258,11 +259,12 @@ async def _start_vscode_server(ctx: click.Context):
     await asyncio.gather(download_tgz(ctx.params["dest"], ctx.params["version"], ctx.params["tgz"]), download_vscode())
     child_process = multiprocessing.Process(
         target=lambda cmd: asyncio.run(asyncio.run(execute_command(cmd))),
-        kwargs={"cmd": f"code-server --bind-addr 0.0.0.0:8080 --disable-workspace-trust --auth none {os.getcwd()}"},
+        kwargs={"cmd": f"code-server --bind-addr 0.0.0.0:6060 --disable-workspace-trust --auth none {os.getcwd()}"},
     )
     child_process.start()
     if child_process.pid is None:
         raise RuntimeError("Failed to start vscode server.")
+
     prepare_launch_json(ctx, child_process.pid)
 
     start_time = time.time()

flyte/_deploy.py

@@ -153,7 +153,7 @@ async def _build_images(deployment: DeploymentPlan) -> ImageCache:
 
         elif env.image == "auto" and "auto" not in image_identifier_map:
             auto_image = Image.from_debian_base()
-            image_identifier_map["auto"] = auto_image.uri
+            images.append(_build_image_bg(env_name, auto_image))
     final_images = await asyncio.gather(*images)
 
     for env_name, image_uri in final_images:
@@ -161,6 +161,8 @@ async def _build_images(deployment: DeploymentPlan) -> ImageCache:
         env = deployment.envs[env_name]
         if isinstance(env.image, Image):
             image_identifier_map[env.image.identifier] = image_uri
+        elif env.image == "auto":
+            image_identifier_map["auto"] = image_uri
 
     return ImageCache(image_lookup=image_identifier_map)
 

flyte/_environment.py

@@ -2,16 +2,14 @@ from __future__ import annotations
 
 import re
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Union
 
 import rich.repr
 
 from ._image import Image
+from ._pod import PodTemplate
 from ._resources import Resources
-from ._secret import SecretRequest
-
-if TYPE_CHECKING:
-    from kubernetes.client import V1PodTemplate
+from ._secret import Secret, SecretRequest
 
 # Global registry to track all Environment instances in load order
 _ENVIRONMENT_REGISTRY: List[Environment] = []
@@ -44,7 +42,7 @@ class Environment:
 
     name: str
     depends_on: List[Environment] = field(default_factory=list)
-    pod_template: Optional[Union[str, "V1PodTemplate"]] = None
+    pod_template: Optional[Union[str, PodTemplate]] = None
     description: Optional[str] = None
     secrets: Optional[SecretRequest] = None
     env_vars: Optional[Dict[str, str]] = None
@@ -54,6 +52,17 @@ class Environment:
     def __post_init__(self):
         if not is_snake_or_kebab_with_numbers(self.name):
             raise ValueError(f"Environment name '{self.name}' must be in snake_case or kebab-case format.")
+        if not isinstance(self.image, (Image, str)):
+            raise TypeError(f"Expected image to be of type str or Image, got {type(self.image)}")
+        if self.secrets and not isinstance(self.secrets, (str, Secret, List)):
+            raise TypeError(f"Expected secrets to be of type SecretRequest, got {type(self.secrets)}")
+        for dep in self.depends_on:
+            if not isinstance(dep, Environment):
+                raise TypeError(f"Expected depends_on to be of type List[Environment], got {type(dep)}")
+        if self.resources is not None and not isinstance(self.resources, Resources):
+            raise TypeError(f"Expected resources to be of type Resources, got {type(self.resources)}")
+        if self.env_vars is not None and not isinstance(self.env_vars, dict):
+            raise TypeError(f"Expected env_vars to be of type Dict[str, str], got {type(self.env_vars)}")
         # Automatically register this environment instance in load order
         _ENVIRONMENT_REGISTRY.append(self)
 

flyte/_hash.py

@@ -1,4 +1,4 @@
-from typing import Callable, Generic, TypeVar
+from typing import TypeVar
 
 T = TypeVar("T")
 
@@ -6,18 +6,3 @@ T = TypeVar("T")
 class HashOnReferenceMixin(object):
     def __hash__(self):
         return hash(id(self))
-
-
-class HashMethod(Generic[T]):
-    """
-    Flyte-specific object used to wrap the hash function for a specific type
-    """
-
-    def __init__(self, function: Callable[[T], str]):
-        self._function = function
-
-    def calculate(self, obj: T) -> str:
-        """
-        Calculate hash for `obj`.
-        """
-        return self._function(obj)

flyte/_image.py

@@ -887,7 +887,12 @@ class Image:
         Use this method to create a new image with the specified uv.lock file layered on top of the current image
         Must have a corresponding pyproject.toml file in the same directory
         Cannot be used in conjunction with conda
-        In the Union builders, using this will change the virtual env to /root/.venv
+
+        By default, this method copies the entire project into the image,
+         including files such as pyproject.toml, uv.lock, and the src/ directory.
+
+        If you prefer not to install the current project, you can pass the extra argument --no-install-project.
+         In this case, the image builder will only copy pyproject.toml and uv.lock into the image.
 
         :param pyproject_file: path to the pyproject.toml file, needs to have a corresponding uv.lock file
         :param uvlock: path to the uv.lock file, if not specified, will use the default uv.lock file in the same

flyte/_initialize.py

@@ -248,23 +248,22 @@ async def init_from_config(
     import flyte.config as config
 
     cfg: config.Config
-    if path_or_config is None or isinstance(path_or_config, str):
-        # If a string is passed, treat it as a path to the config file
-        if root_dir and path_or_config:
-            cfg = config.auto(str(root_dir / path_or_config))
-        elif path_or_config:
-            if not Path(path_or_config).exists():
-                raise InitializationError(
-                    "ConfigFileNotFoundError",
-                    "user",
-                    f"Configuration file '{path_or_config}' does not exist., current working directory is {Path.cwd()}",
-                )
-            cfg = config.auto(path_or_config)
+    if path_or_config is None:
+        # If no path is provided, use the default config file
+        cfg = config.auto()
+    elif isinstance(path_or_config, str):
+        if root_dir:
+            cfg_path = str(root_dir / path_or_config)
         else:
-            # If no path is provided, use the default config file
-            cfg = config.auto()
+            cfg_path = path_or_config
+        if not Path(cfg_path).exists():
+            raise InitializationError(
+                "ConfigFileNotFoundError",
+                "user",
+                f"Configuration file '{cfg_path}' does not exist., current working directory is {Path.cwd()}",
+            )
+        cfg = config.auto(cfg_path)
     else:
-        # If a Config object is passed, use it directly
         cfg = path_or_config
 
     logger.debug(f"Flyte config initialized as {cfg}")

flyte/_internal/controllers/remote/_controller.py

@@ -413,8 +413,7 @@ class RemoteController(Controller):
             else:
                 logger.warning(f"Action {prev_action.action_id.name} failed, but no error was found, re-running trace!")
         elif prev_action.realized_outputs_uri is not None:
-            outputs_file_path = io.outputs_path(prev_action.realized_outputs_uri)
-            o = await io.load_outputs(outputs_file_path, max_bytes=MAX_TRACE_BYTES)
+            o = await io.load_outputs(prev_action.realized_outputs_uri, max_bytes=MAX_TRACE_BYTES)
             outputs = await convert.convert_outputs_to_native(_interface, o)
             return (
                 TraceInfo(func_name, sub_action_id, _interface, inputs_uri, output=outputs),
@@ -439,15 +438,13 @@ class RemoteController(Controller):
         outputs_file_path: str = ""
 
         if info.interface.has_outputs():
-            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)
-                await io.upload_outputs(outputs, sub_run_output_path, max_bytes=MAX_TRACE_BYTES)
-            elif info.error:
+            if info.error:
                 err = convert.convert_from_native_to_error(info.error)
                 await io.upload_error(err.err, sub_run_output_path)
             else:
-                raise flyte.errors.RuntimeSystemError("BadTraceInfo", "Trace info does not have output or error")
+                outputs = await convert.convert_from_native_to_outputs(info.output, info.interface)
+                outputs_file_path = io.outputs_path(sub_run_output_path)
+                await io.upload_outputs(outputs, sub_run_output_path, max_bytes=MAX_TRACE_BYTES)
 
         typed_interface = transform_native_to_typed_interface(info.interface)
 

flyte/_internal/controllers/remote/_core.py

@@ -158,8 +158,8 @@ class Controller:
         self._thread.start()
 
         # Wait for the thread to be ready
-        logger.info("Waiting for controller thread to be ready...")
         if not self._thread_ready.wait(timeout=self._thread_wait_timeout):
+            logger.warning("Controller thread did not finish within timeout")
             raise TimeoutError("Controller thread failed to start in time")
 
         if self._get_exception():

flyte/_internal/imagebuild/docker_builder.py

@@ -44,19 +44,19 @@ _F_IMG_ID = "_F_IMG_ID"
 FLYTE_DOCKER_BUILDER_CACHE_FROM = "FLYTE_DOCKER_BUILDER_CACHE_FROM"
 FLYTE_DOCKER_BUILDER_CACHE_TO = "FLYTE_DOCKER_BUILDER_CACHE_TO"
 
-UV_LOCK_INSTALL_TEMPLATE = Template("""\
-WORKDIR /root
+UV_LOCK_WITHOUT_PROJECT_INSTALL_TEMPLATE = Template("""\
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
-    --mount=type=bind,target=uv.lock,src=uv.lock \
-    --mount=type=bind,target=pyproject.toml,src=pyproject.toml \
+    --mount=type=bind,target=uv.lock,src=$UV_LOCK_PATH \
+    --mount=type=bind,target=pyproject.toml,src=$PYPROJECT_PATH \
     $SECRET_MOUNT \
-    uv sync $PIP_INSTALL_ARGS
-WORKDIR /
+    uv sync --active $PIP_INSTALL_ARGS
+""")
 
-# Update PATH and UV_PYTHON to point to the venv created by uv sync
-ENV PATH="/root/.venv/bin:$$PATH" \
-    VIRTUALENV=/root/.venv \
-    UV_PYTHON=/root/.venv/bin/python
+UV_LOCK_INSTALL_TEMPLATE = Template("""\
+COPY $PYPROJECT_PATH $PYPROJECT_PATH
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
+    $SECRET_MOUNT \
+    uv sync --active $PIP_INSTALL_ARGS --project $PYPROJECT_PATH
 """)
 
 UV_PACKAGE_INSTALL_COMMAND_TEMPLATE = Template("""\
@@ -90,7 +90,7 @@ RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
 # new template
 DOCKER_FILE_UV_BASE_TEMPLATE = Template("""\
 # syntax=docker/dockerfile:1.10
-FROM ghcr.io/astral-sh/uv:0.6.12 AS uv
+FROM ghcr.io/astral-sh/uv:0.8.13 AS uv
 FROM $BASE_IMAGE
 
 USER root
@@ -177,6 +177,7 @@ class PythonWheelHandler:
                 "/dist",
                 "--no-deps",
                 "--no-index",
+                "--reinstall",
                 layer.package_name,
             ],
         ]
@@ -230,20 +231,27 @@ class AptPackagesHandler:
 class UVProjectHandler:
     @staticmethod
     async def handle(layer: UVProject, context_path: Path, dockerfile: str) -> str:
-        # copy the two files
-        shutil.copy(layer.pyproject, context_path)
-        shutil.copy(layer.uvlock, context_path)
-
-        # --locked: Assert that the `uv.lock` will remain unchanged
-        # --no-dev: Omit the development dependency group
-        # --no-install-project: Do not install the current project
-        additional_pip_install_args = ["--locked", "--no-dev", "--no-install-project"]
         secret_mounts = _get_secret_mounts_layer(layer.secret_mounts)
-        delta = UV_LOCK_INSTALL_TEMPLATE.substitute(
-            PIP_INSTALL_ARGS=" ".join(additional_pip_install_args), SECRET_MOUNT=secret_mounts
-        )
-        dockerfile += delta
+        if layer.extra_index_urls and "--no-install-project" in layer.extra_index_urls:
+            # Only Copy pyproject.yaml and uv.lock.
+            pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
+            uvlock_dst = copy_files_to_context(layer.uvlock, context_path)
+            delta = UV_LOCK_WITHOUT_PROJECT_INSTALL_TEMPLATE.substitute(
+                UV_LOCK_PATH=uvlock_dst.relative_to(context_path),
+                PYPROJECT_PATH=pyproject_dst.relative_to(context_path),
+                PIP_INSTALL_ARGS=" ".join(layer.get_pip_install_args()),
+                SECRET_MOUNT=secret_mounts,
+            )
+        else:
+            # Copy the entire project.
+            pyproject_dst = copy_files_to_context(layer.pyproject.parent, context_path)
+            delta = UV_LOCK_INSTALL_TEMPLATE.substitute(
+                PYPROJECT_PATH=pyproject_dst.relative_to(context_path),
+                PIP_INSTALL_ARGS=" ".join(layer.get_pip_install_args()),
+                SECRET_MOUNT=secret_mounts,
+            )
 
+        dockerfile += delta
         return dockerfile
 
 

flyte/_internal/imagebuild/remote_builder.py

@@ -1,5 +1,7 @@
+import gzip
 import os
 import shutil
+import tarfile
 import tempfile
 import typing
 from datetime import datetime, timezone
@@ -7,11 +9,13 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Optional, Tuple, cast
 from uuid import uuid4
 
+import aiofiles
 import click
 
 import flyte
 import flyte.errors
 from flyte import Image, remote
+from flyte._code_bundle._utils import tar_strip_file_attributes
 from flyte._image import (
     AptPackages,
     Architecture,
@@ -130,7 +134,7 @@ class RemoteImageBuilder(ImageBuilder):
                     auto_version="latest",
                 )
                 await flyte.with_runcontext(project=IMAGE_TASK_PROJECT, domain=IMAGE_TASK_DOMAIN).run.aio(
-                    entity, spec=spec, context=context, target_image=image_name
+                    entity, target_image=image_name
                 )
             except Exception as e:
                 # Ignore the error if optimize is not enabled in the backend.
@@ -160,17 +164,32 @@ 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
-        archive = Path(shutil.make_archive(str(tmp_path / "context"), "xztar", context_path))
-        st = archive.stat()
-        if st.st_size > 5 * 1024 * 1024:
+        tar_path = tmp_path / "context.tar"
+        with tarfile.open(tar_path, "w", dereference=False) as tar:
+            files: typing.List[str] = os.listdir(context_path)
+            for ws_file in files:
+                tar.add(
+                    os.path.join(context_path, ws_file),
+                    recursive=True,
+                    arcname=ws_file,
+                    filter=tar_strip_file_attributes,
+                )
+        context_dst = Path(f"{tar_path!s}.gz")
+        with gzip.GzipFile(filename=context_dst, mode="wb", mtime=0) as gzipped:
+            async with aiofiles.open(tar_path, "rb") as tar_file:
+                content = await tar_file.read()
+                gzipped.write(content)
+
+        context_size = tar_path.stat().st_size
+        if context_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. "
+                    f"Context size is {context_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)
+        _, context_url = await remote.upload_file.aio(context_dst)
     else:
         context_url = ""
 
@@ -248,12 +267,20 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
             for line in layer.pyproject.read_text().splitlines():
                 if "tool.uv.index" in line:
                     raise ValueError("External sources are not supported in pyproject.toml")
-            shutil.copy2(layer.pyproject, context_path / layer.pyproject.name)
+
+            if layer.extra_index_urls and "--no-install-project" in layer.extra_index_urls:
+                # Copy pyproject itself
+                pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
+            else:
+                # Copy the entire project
+                pyproject_dst = copy_files_to_context(layer.pyproject.parent, context_path)
 
             uv_layer = image_definition_pb2.Layer(
                 uv_project=image_definition_pb2.UVProject(
-                    pyproject=str(layer.pyproject.name),
-                    uvlock=str(layer.uvlock.name),
+                    pyproject=str(pyproject_dst.relative_to(context_path)),
+                    uvlock=str(copy_files_to_context(layer.uvlock, context_path).relative_to(context_path)),
+                    options=pip_options,
+                    secret_mounts=secret_mounts,
                 )
             )
             layers.append(uv_layer)
@@ -303,7 +330,7 @@ def _get_fully_qualified_image_name(outputs: ActionOutputs) -> str:
 
 def _get_build_secrets_from_image(image: Image) -> Optional[typing.List[Secret]]:
     secrets = []
-    DEFAULT_SECRET_DIR = Path("etc/flyte/secrets")
+    DEFAULT_SECRET_DIR = Path("/etc/flyte/secrets")
     for layer in image._layers:
         if isinstance(layer, (PipOption, Commands, AptPackages)) and layer.secret_mounts is not None:
             for secret_mount in layer.secret_mounts:

flyte/_internal/imagebuild/utils.py

@@ -23,7 +23,8 @@ def copy_files_to_context(src: Path, context_path: Path) -> Path:
         dst_path = context_path / src
     dst_path.parent.mkdir(parents=True, exist_ok=True)
     if src.is_dir():
-        shutil.copytree(src, dst_path, dirs_exist_ok=True)
+        # TODO: Add support dockerignore
+        shutil.copytree(src, dst_path, dirs_exist_ok=True, ignore=shutil.ignore_patterns(".idea", ".venv"))
     else:
         shutil.copy(src, dst_path)
     return dst_path

flyte/_internal/runtime/convert.py

@@ -308,15 +308,82 @@ def generate_inputs_hash(serialized_inputs: str | bytes) -> str:
     return hash_data(serialized_inputs)
 
 
+def generate_inputs_repr_for_literal(literal: literals_pb2.Literal) -> bytes:
+    """
+    Generate a byte representation for a single literal that is meant to be hashed as part of the cache key
+    computation for an Action. This function should just serialize the literal deterministically, but will
+    use an existing hash value if present in the Literal.  This is trivial, except we need to handle nested literals
+    (inside collections and maps), that may have the hash property set.
+
+    :param literal: The literal to get a hashable representation for.
+    :return: byte representation of the literal that can be fed into a hash function.
+    """
+    # If the literal has a hash value, use that instead of serializing the full literal
+    if literal.hash:
+        return literal.hash.encode("utf-8")
+
+    if literal.HasField("collection"):
+        buf = bytearray()
+        for nested_literal in literal.collection.literals:
+            if nested_literal.hash:
+                buf += nested_literal.hash.encode("utf-8")
+            else:
+                buf += generate_inputs_repr_for_literal(nested_literal)
+
+        b = bytes(buf)
+        return b
+
+    elif literal.HasField("map"):
+        buf = bytearray()
+        # Sort keys to ensure deterministic ordering
+        for key in sorted(literal.map.literals.keys()):
+            nested_literal = literal.map.literals[key]
+            buf += key.encode("utf-8")
+            if nested_literal.hash:
+                buf += nested_literal.hash.encode("utf-8")
+            else:
+                buf += generate_inputs_repr_for_literal(nested_literal)
+
+        b = bytes(buf)
+        return b
+
+    # For all other cases (scalars, etc.), just serialize the literal normally
+    return literal.SerializeToString(deterministic=True)
+
+
+def generate_inputs_hash_for_named_literals(inputs: list[run_definition_pb2.NamedLiteral]) -> str:
+    """
+    Generate a hash for the inputs using the new literal representation approach that respects
+    hash values already present in literals. This is used to uniquely identify the inputs for a task
+    when some literals may have precomputed hash values.
+
+    :param inputs: List of NamedLiteral inputs to hash.
+    :return: A base64-encoded string representation of the hash.
+    """
+    if not inputs:
+        return ""
+
+    # Build the byte representation by concatenating each literal's representation
+    combined_bytes = b""
+    for named_literal in inputs:
+        # Add the name to ensure order matters
+        name_bytes = named_literal.name.encode("utf-8")
+        literal_bytes = generate_inputs_repr_for_literal(named_literal.value)
+        # Combine name and literal bytes with a separator to avoid collisions
+        combined_bytes += name_bytes + b":" + literal_bytes + b";"
+
+    return hash_data(combined_bytes)
+
+
 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:
+    if not inputs or not inputs.literals:
         return ""
-    return generate_inputs_hash(inputs.SerializeToString(deterministic=True))
+    return generate_inputs_hash_for_named_literals(list(inputs.literals))
 
 
 def generate_interface_hash(task_interface: interface_pb2.TypedInterface) -> str:

flyte/_internal/runtime/taskrunner.py

@@ -4,6 +4,7 @@ invoked within a context tree.
 """
 
 import pathlib
+import time
 from typing import Any, Dict, List, Optional, Tuple
 
 import flyte.report
@@ -172,6 +173,8 @@ async def extract_download_run_upload(
     This method is invoked from the CLI (urun) and is used to run a task. This assumes that the context tree
     has already been created, and the task has been loaded. It also handles the loading of the task.
     """
+    t = time.time()
+    logger.warning(f"Task {action.name} started at {t}")
     outputs, err = await convert_and_run(
         task=task,
         input_path=input_path,
@@ -194,4 +197,4 @@ async def extract_download_run_upload(
         logger.info(f"Task {task.name} completed successfully, no outputs")
         return
     await upload_outputs(outputs, output_path) if output_path else None
-    logger.info(f"Task {task.name} completed successfully, uploaded outputs to {output_path}")
+    logger.warning(f"Task {task.name} completed successfully, uploaded outputs to {output_path} in {time.time() - t}s")