flyte 2.0.0b3__py3-none-any.whl → 2.0.0b5__py3-none-any.whl

flyte/_internal/imagebuild/docker_builder.py

@@ -11,6 +11,7 @@ from typing import ClassVar, Optional, Protocol, cast
 import aiofiles
 import click
 
+from flyte import Secret
 from flyte._image import (
     AptPackages,
     Commands,
@@ -19,12 +20,15 @@ from flyte._image import (
     Env,
     Image,
     Layer,
+    PipOption,
     PipPackages,
     PythonWheels,
     Requirements,
     UVProject,
+    UVScript,
     WorkDir,
     _DockerLines,
+    _ensure_tuple,
 )
 from flyte._internal.imagebuild.image_builder import (
     DockerAPIImageChecker,
@@ -44,6 +48,7 @@ WORKDIR /root
 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 \
+    $SECRET_MOUNT \
     uv sync $PIP_INSTALL_ARGS
 WORKDIR /
 
@@ -56,31 +61,35 @@ ENV PATH="/root/.venv/bin:$$PATH" \
 UV_PACKAGE_INSTALL_COMMAND_TEMPLATE = Template("""\
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
     --mount=type=bind,target=requirements_uv.txt,src=requirements_uv.txt \
-    uv pip install --prerelease=allow --python $$UV_PYTHON $PIP_INSTALL_ARGS
+    $SECRET_MOUNT \
+    uv pip install --python $$UV_PYTHON $PIP_INSTALL_ARGS
 """)
 
 UV_WHEEL_INSTALL_COMMAND_TEMPLATE = Template("""\
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=wheel \
     --mount=source=/dist,target=/dist,type=bind \
-    uv pip install --prerelease=allow --python $$UV_PYTHON $PIP_INSTALL_ARGS
+    $SECRET_MOUNT \
+    uv pip install --python $$UV_PYTHON $PIP_INSTALL_ARGS
 """)
 
 APT_INSTALL_COMMAND_TEMPLATE = Template("""\
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/var/cache/apt,id=apt \
+    $SECRET_MOUNT \
     apt-get update && apt-get install -y --no-install-recommends \
     $APT_PACKAGES
 """)
 
 UV_PYTHON_INSTALL_COMMAND = Template("""\
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
+    $SECRET_MOUNT \
     uv pip install $PIP_INSTALL_ARGS
 """)
 
 # uv pip install --python /root/env/bin/python
 # new template
 DOCKER_FILE_UV_BASE_TEMPLATE = Template("""\
-#syntax=docker/dockerfile:1.5
-FROM ghcr.io/astral-sh/uv:0.6.12 as uv
+# syntax=docker/dockerfile:1.10
+FROM ghcr.io/astral-sh/uv:0.6.12 AS uv
 FROM $BASE_IMAGE
 
 USER root
@@ -139,8 +148,10 @@ class PipAndRequirementsHandler:
 
         pip_install_args = layer.get_pip_install_args()
         pip_install_args.extend(["--requirement", "requirements_uv.txt"])
-
-        delta = UV_PACKAGE_INSTALL_COMMAND_TEMPLATE.substitute(PIP_INSTALL_ARGS=" ".join(pip_install_args))
+        secret_mounts = _get_secret_mounts_layer(layer.secret_mounts)
+        delta = UV_PACKAGE_INSTALL_COMMAND_TEMPLATE.substitute(
+            PIP_INSTALL_ARGS=" ".join(pip_install_args), SECRET_MOUNT=secret_mounts
+        )
         dockerfile += delta
 
         return dockerfile
@@ -152,8 +163,10 @@ class PythonWheelHandler:
         shutil.copytree(layer.wheel_dir, context_path / "dist", dirs_exist_ok=True)
         pip_install_args = layer.get_pip_install_args()
         pip_install_args.extend(["/dist/*.whl"])
-
-        delta = UV_WHEEL_INSTALL_COMMAND_TEMPLATE.substitute(PIP_INSTALL_ARGS=" ".join(pip_install_args))
+        secret_mounts = _get_secret_mounts_layer(layer.secret_mounts)
+        delta = UV_WHEEL_INSTALL_COMMAND_TEMPLATE.substitute(
+            PIP_INSTALL_ARGS=" ".join(pip_install_args), SECRET_MOUNT=secret_mounts
+        )
         dockerfile += delta
 
         return dockerfile
@@ -181,9 +194,10 @@ class EnvHandler:
 
 class AptPackagesHandler:
     @staticmethod
-    async def handle(layer: AptPackages, context_path: Path, dockerfile: str) -> str:
+    async def handle(layer: AptPackages, _: Path, dockerfile: str) -> str:
         packages = layer.packages
-        delta = APT_INSTALL_COMMAND_TEMPLATE.substitute(APT_PACKAGES=" ".join(packages))
+        secret_mounts = _get_secret_mounts_layer(layer.secret_mounts)
+        delta = APT_INSTALL_COMMAND_TEMPLATE.substitute(APT_PACKAGES=" ".join(packages), SECRET_MOUNT=secret_mounts)
         dockerfile += delta
 
         return dockerfile
@@ -200,7 +214,10 @@ class UVProjectHandler:
         # --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"]
-        delta = UV_LOCK_INSTALL_TEMPLATE.substitute(PIP_INSTALL_ARGS=" ".join(additional_pip_install_args))
+        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
 
         return dockerfile
@@ -250,19 +267,82 @@ class CommandsHandler:
 
 class WorkDirHandler:
     @staticmethod
-    async def handle(layer: WorkDir, context_path: Path, dockerfile: str) -> str:
+    async def handle(layer: WorkDir, _: Path, dockerfile: str) -> str:
         # cd to the workdir
         dockerfile += f"\nWORKDIR {layer.workdir}\n"
 
         return dockerfile
 
 
+def _get_secret_commands(layers: typing.Tuple[Layer, ...]) -> typing.List[str]:
+    commands = []
+
+    def _get_secret_command(secret: str | Secret) -> typing.List[str]:
+        secret_id = hash(secret)
+        if isinstance(secret, str):
+            if not os.path.exists(secret):
+                raise FileNotFoundError(f"Secret file '{secret}' not found")
+            return ["--secret", f"id={secret_id},src={secret}"]
+        secret_env_key = "_".join([k.upper() for k in filter(None, (secret.group, secret.key))])
+        secret_env = os.getenv(secret_env_key)
+        if secret_env:
+            return ["--secret", f"id={secret_id},env={secret_env}"]
+        secret_file_name = "_".join(list(filter(None, (secret.group, secret.key))))
+        secret_file_path = f"/etc/secrets/{secret_file_name}"
+        if not os.path.exists(secret_file_path):
+            raise FileNotFoundError(f"Secret not found in Env Var {secret_env_key} or file {secret_file_path}")
+        return ["--secret", f"id={secret_id},src={secret_file_path}"]
+
+    for layer in layers:
+        if isinstance(layer, (PipOption, AptPackages)):
+            if layer.secret_mounts:
+                for secret_mount in layer.secret_mounts:
+                    commands.extend(_get_secret_command(secret_mount))
+    return commands
+
+
+def _get_secret_mounts_layer(secrets: typing.Tuple[str | Secret, ...] | None) -> str:
+    if secrets is None:
+        return ""
+    secret_mounts_layer = ""
+    for secret in secrets:
+        secret_id = hash(secret)
+        if isinstance(secret, str):
+            secret_mounts_layer += f"--mount=type=secret,id={secret_id},target=/run/secrets/{os.path.basename(secret)}"
+        elif isinstance(secret, Secret):
+            if secret.mount:
+                secret_mounts_layer += f"--mount=type=secret,id={secret_id},target={secret.mount}"
+            elif secret.as_env_var:
+                secret_mounts_layer += f"--mount=type=secret,id={secret_id},env={secret.as_env_var}"
+            else:
+                secret_file_name = "_".join(list(filter(None, (secret.group, secret.key))))
+                secret_mounts_layer += f"--mount=type=secret,id={secret_id},src=/run/secrets/{secret_file_name}"
+
+    return secret_mounts_layer
+
+
 async def _process_layer(layer: Layer, context_path: Path, dockerfile: str) -> str:
     match layer:
         case PythonWheels():
             # Handle Python wheels
             dockerfile = await PythonWheelHandler.handle(layer, context_path, dockerfile)
 
+        case UVScript():
+            # Handle UV script
+            from flyte._utils import parse_uv_script_file
+
+            header = parse_uv_script_file(layer.script)
+            if header.dependencies:
+                pip = PipPackages(
+                    packages=_ensure_tuple(header.dependencies) if header.dependencies else None,
+                    secret_mounts=layer.secret_mounts,
+                    index_url=layer.index_url,
+                    extra_args=layer.extra_args,
+                    pre=layer.pre,
+                    extra_index_urls=layer.extra_index_urls,
+                )
+                dockerfile = await PipAndRequirementsHandler.handle(pip, context_path, dockerfile)
+
         case Requirements() | PipPackages():
             # Handle pip packages and requirements
             dockerfile = await PipAndRequirementsHandler.handle(layer, context_path, dockerfile)
@@ -357,6 +437,8 @@ class DockerImageBuilder(ImageBuilder):
         else:
             command.append("--load")
 
+        command.extend(_get_secret_commands(layers=image._layers))
+
         concat_command = " ".join(command)
         logger.debug(f"Build command: {concat_command}")
         click.secho(f"Run command: {concat_command} ", fg="blue")
@@ -464,6 +546,8 @@ class DockerImageBuilder(ImageBuilder):
                 command.append("--push")
             else:
                 command.append("--load")
+
+            command.extend(_get_secret_commands(layers=image._layers))
             command.append(tmp_dir)
 
             concat_command = " ".join(command)

flyte/_internal/imagebuild/remote_builder.py

@@ -23,6 +23,7 @@ from flyte._image import (
     PythonWheels,
     Requirements,
     UVProject,
+    UVScript,
 )
 from flyte._internal.imagebuild.image_builder import ImageBuilder, ImageChecker
 from flyte._logging import logger
@@ -196,10 +197,19 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
                 )
             )
             layers.append(requirements_layer)
-        elif isinstance(layer, PipPackages):
+        elif isinstance(layer, PipPackages) or isinstance(layer, UVScript):
+            if isinstance(layer, UVScript):
+                from flyte._utils import parse_uv_script_file
+
+                header = parse_uv_script_file(layer.script)
+                if not header.dependencies:
+                    continue
+                packages: typing.Iterable[str] = header.dependencies
+            else:
+                packages = layer.packages or []
             pip_layer = image_definition_pb2.Layer(
                 pip_packages=image_definition_pb2.PipPackages(
-                    packages=layer.packages,
+                    packages=packages,
                     options=image_definition_pb2.PipOptions(
                         index_url=layer.index_url,
                         extra_index_urls=layer.extra_index_urls,

flyte/_internal/runtime/convert.py

@@ -105,11 +105,9 @@ def is_optional_type(tp) -> bool:
 async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwargs) -> Inputs:
     kwargs = interface.convert_to_kwargs(*args, **kwargs)
 
-    if len(kwargs) < interface.num_required_inputs():
-        raise ValueError(
-            f"Received {len(kwargs)} inputs but interface has {interface.num_required_inputs()} required inputs. "
-            f"Please provide all required inputs. Inputs received: {kwargs}, interface: {interface}"
-        )
+    missing = [key for key in interface.required_inputs() if key not in kwargs]
+    if missing:
+        raise ValueError(f"Missing required inputs: {', '.join(missing)}")
 
     if len(interface.inputs) == 0:
         return Inputs.empty()

flyte/_run.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 import pathlib
 import uuid
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
 
 import flyte.errors
@@ -36,10 +37,26 @@ if TYPE_CHECKING:
     from flyte.remote._task import LazyEntity
 
     from ._code_bundle import CopyFiles
+    from ._internal.imagebuild.image_builder import ImageCache
 
 Mode = Literal["local", "remote", "hybrid"]
 
 
+@dataclass(frozen=True)
+class _CacheKey:
+    obj_id: int
+    dry_run: bool
+
+
+@dataclass(frozen=True)
+class _CacheValue:
+    code_bundle: CodeBundle | None
+    image_cache: Optional[ImageCache]
+
+
+_RUN_CACHE: Dict[_CacheKey, _CacheValue] = {}
+
+
 async def _get_code_bundle_for_run(name: str) -> CodeBundle | None:
     """
     Get the code bundle for the run with the given name.
@@ -78,6 +95,7 @@ class _Runner:
         annotations: Dict[str, str] | None = None,
         interruptible: bool = False,
         log_level: int | None = None,
+        disable_run_cache: bool = False,
     ):
         init_config = _get_init_config()
         client = init_config.client if init_config else None
@@ -104,6 +122,7 @@ class _Runner:
         self._annotations = annotations
         self._interruptible = interruptible
         self._log_level = log_level
+        self._disable_run_cache = disable_run_cache
 
     @requires_initialization
     async def _run_remote(self, obj: TaskTemplate[P, R] | LazyEntity, *args: P.args, **kwargs: P.kwargs) -> Run:
@@ -135,24 +154,36 @@ 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 = await build_images.aio(cast(Environment, obj.parent_env()))
-
-            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,
-                )
+            if (
+                not self._disable_run_cache
+                and _RUN_CACHE.get(_CacheKey(obj_id=id(obj), dry_run=self._dry_run)) is not None
+            ):
+                cached_value = _RUN_CACHE[_CacheKey(obj_id=id(obj), dry_run=self._dry_run)]
+                code_bundle = cached_value.code_bundle
+                image_cache = cached_value.image_cache
             else:
-                if self._copy_files != "none":
-                    code_bundle = await build_code_bundle(
-                        from_dir=cfg.root_dir,
-                        dryrun=self._dry_run,
+                image_cache = await build_images.aio(cast(Environment, obj.parent_env()))
+
+                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,
-                        copy_style=self._copy_files,
                     )
                 else:
-                    code_bundle = None
+                    if self._copy_files != "none":
+                        code_bundle = await build_code_bundle(
+                            from_dir=cfg.root_dir,
+                            dryrun=self._dry_run,
+                            copy_bundle_to=self._copy_bundle_to,
+                            copy_style=self._copy_files,
+                        )
+                    else:
+                        code_bundle = None
+            if not self._disable_run_cache:
+                _RUN_CACHE[_CacheKey(obj_id=id(obj), dry_run=self._dry_run)] = _CacheValue(
+                    code_bundle=code_bundle, image_cache=image_cache
+                )
 
             version = self._version or (
                 code_bundle.computed_version if code_bundle and code_bundle.computed_version else None
@@ -516,6 +547,7 @@ def with_runcontext(
     annotations: Dict[str, str] | None = None,
     interruptible: bool = False,
     log_level: int | None = None,
+    disable_run_cache: bool = False,
 ) -> _Runner:
     """
     Launch a new run with the given parameters as the context.
@@ -556,6 +588,7 @@ def with_runcontext(
     :param interruptible: Optional If true, the run can be interrupted by the user.
     :param log_level: Optional Log level to set for the run. If not provided, it will be set to the default log level
         set using `flyte.init()`
+    :param disable_run_cache: Optional If true, the run cache will be disabled. This is useful for testing purposes.
 
     :return: runner
     """
@@ -580,6 +613,7 @@ def with_runcontext(
         project=project,
         domain=domain,
         log_level=log_level,
+        disable_run_cache=disable_run_cache,
     )
 
 

flyte/_secret.py

@@ -7,8 +7,9 @@ from typing import List, Optional, Union
 @dataclass
 class Secret:
     """
-    Secrets are used to inject sensitive information into tasks. Secrets can be mounted as environment variables or
-    files. The secret key is the name of the secret in the secret store. The group is optional and maybe used with some
+    Secrets are used to inject sensitive information into tasks or image build context.
+    Secrets can be mounted as environment variables or files.
+     The secret key is the name of the secret in the secret store. The group is optional and maybe used with some
     secret stores to organize secrets. The secret_mount is used to specify how the secret should be mounted. If the
     secret_mount is set to "env" the secret will be mounted as an environment variable. If the secret_mount is set to
     "file" the secret will be mounted as a file. The as_env_var is an optional parameter that can be used to specify the

flyte/_task.py

@@ -150,6 +150,10 @@ class TaskTemplate(Generic[P, R]):
         self.__dict__.update(state)
         self.parent_env = None
 
+    @property
+    def source_file(self) -> Optional[str]:
+        return None
+
     async def pre(self, *args, **kwargs) -> Dict[str, Any]:
         """
         This is the preexecute function that will be
@@ -395,6 +399,15 @@ class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
         if not iscoroutinefunction(self.func):
             self._call_as_synchronous = True
 
+    @property
+    def source_file(self) -> Optional[str]:
+        """
+        Returns the source file of the function, if available. This is useful for debugging and tracing.
+        """
+        if hasattr(self.func, "__code__") and self.func.__code__:
+            return self.func.__code__.co_filename
+        return None
+
     def forward(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
         # In local execution, we want to just call the function. Note we're not awaiting anything here.
         # If the function was a coroutine function, the coroutine is returned and the await that the caller has

flyte/_utils/__init__.py

@@ -9,6 +9,7 @@ from .coro_management import run_coros
 from .file_handling import filehash_update, update_hasher_for_source
 from .helpers import get_cwd_editable_install
 from .lazy_module import lazy_module
+from .module_loader import load_python_modules
 from .org_discovery import hostname_from_url, org_from_endpoint, sanitize_endpoint
 from .uv_script_parser import parse_uv_script_file
 
@@ -18,6 +19,7 @@ __all__ = [
     "get_cwd_editable_install",
     "hostname_from_url",
     "lazy_module",
+    "load_python_modules",
     "org_from_endpoint",
     "parse_uv_script_file",
     "run_coros",

flyte/_utils/module_loader.py

@@ -0,0 +1,89 @@
+import glob
+import importlib.util
+import os
+import sys
+from pathlib import Path
+from typing import List, Tuple
+
+from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn, TimeRemainingColumn
+
+import flyte.errors
+
+
+def load_python_modules(path: Path, recursive: bool = False) -> Tuple[List[str], List[Tuple[Path, str]]]:
+    """
+    Load all Python modules from a path and return list of loaded module names.
+
+    :param path: File or directory path
+    :param recursive: If True, load modules recursively from subdirectories
+    :return: List of loaded module names, and list of file paths that failed to load
+    """
+    loaded_modules = []
+    failed_paths = []
+
+    if path.is_file() and path.suffix == ".py":
+        # Single file case
+        module_name = _load_module_from_file(path)
+        if module_name:
+            loaded_modules.append(module_name)
+
+    elif path.is_dir():
+        # Directory case
+        pattern = "**/*.py" if recursive else "*.py"
+        python_files = glob.glob(str(path / pattern), recursive=recursive)
+
+        with Progress(
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            "[progress.percentage]{task.percentage:>3.0f}%",
+            TimeElapsedColumn(),
+            TimeRemainingColumn(),
+            TextColumn("• {task.fields[current_file]}"),
+        ) as progress:
+            task = progress.add_task(f"Loading {len(python_files)} files", total=len(python_files), current_file="")
+            for file_path in python_files:
+                p = Path(file_path)
+                progress.update(task, advance=1, current_file=p.name)
+                # Skip __init__.py files
+                if p.name == "__init__.py":
+                    continue
+
+                try:
+                    module_name = _load_module_from_file(p)
+                    if module_name:
+                        loaded_modules.append(module_name)
+                except flyte.errors.ModuleLoadError as e:
+                    failed_paths.append((p, str(e)))
+
+            progress.update(task, advance=1, current_file="[green]Done[/green]")
+
+    return loaded_modules, failed_paths
+
+
+def _load_module_from_file(file_path: Path) -> str | None:
+    """
+    Load a Python module from a file path.
+
+    :param file_path: Path to the Python file
+    :return: Module name if successfully loaded, None otherwise
+    """
+    try:
+        # Use the file stem as module name
+        module_name = file_path.stem
+
+        # Load the module specification
+        spec = importlib.util.spec_from_file_location(module_name, file_path)
+        if spec is None or spec.loader is None:
+            return None
+
+        # Create and execute the module
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+        module_path = os.path.dirname(os.path.abspath(file_path))
+        sys.path.append(module_path)
+        spec.loader.exec_module(module)
+
+        return module_name
+
+    except Exception as e:
+        raise flyte.errors.ModuleLoadError(f"Failed to load module from {file_path}: {e}") from e

flyte/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.0.0b3'
-__version_tuple__ = version_tuple = (2, 0, 0, 'b3')
+__version__ = version = '2.0.0b5'
+__version_tuple__ = version_tuple = (2, 0, 0, 'b5')

flyte/cli/_common.py

@@ -265,7 +265,7 @@ class ObjectsPerFileGroup(GroupBase):
 
         spec = importlib.util.spec_from_file_location(module_name, self.filename)
         if spec is None or spec.loader is None:
-            raise click.ClickException(f"Could not load module {module_name} from {self.filename}")
+            raise click.ClickException(f"Could not load module {module_name} from path [{self.filename}]")
 
         module = importlib.util.module_from_spec(spec)
         sys.modules[module_name] = module
@@ -314,6 +314,10 @@ class FileGroup(GroupBase):
         if self._files is None:
             directory = self._dir or Path(".").absolute()
             self._files = [os.fspath(p) for p in directory.glob("*.py") if p.name != "__init__.py"]
+            if not self._files:
+                self._files = [os.fspath(".")] + [
+                    os.fspath(p.name) for p in directory.iterdir() if not p.name.startswith(("_", ".")) and p.is_dir()
+                ]
         return self._files
 
     def list_commands(self, ctx):