flyte 2.0.0b22__py3-none-any.whl → 2.0.0b23__py3-none-any.whl

flyte/_initialize.py

@@ -182,6 +182,28 @@ async def init(
     """
     from flyte._utils import get_cwd_editable_install, org_from_endpoint, sanitize_endpoint
 
+    if endpoint or api_key:
+        if project is None:
+            raise ValueError(
+                "Project must be provided to initialize the client. "
+                "Please set 'project' in the 'task' section of your config file, "
+                "or pass it directly to flyte.init(project='your-project-name')."
+            )
+
+        if domain is None:
+            raise ValueError(
+                "Domain must be provided to initialize the client. "
+                "Please set 'domain' in the 'task' section of your config file, "
+                "or pass it directly to flyte.init(domain='your-domain-name')."
+            )
+
+        if org is None and org_from_endpoint(endpoint) is None:
+            raise ValueError(
+                "Organization must be provided to initialize the client. "
+                "Please set 'org' in the 'task' section of your config file, "
+                "or pass it directly to flyte.init(org='your-org-name')."
+            )
+
     _initialize_logger(log_level=log_level)
 
     global _init_config  # noqa: PLW0603
@@ -278,6 +300,7 @@ async def init_from_config(
     _initialize_logger(log_level=log_level)
 
     logger.info(f"Flyte config initialized as {cfg}", extra={"highlighter": ReprHighlighter()})
+
     await init.aio(
         org=cfg.task.org,
         project=cfg.task.project,

flyte/_internal/controllers/_local_controller.py

@@ -2,16 +2,20 @@ import asyncio
 import atexit
 import concurrent.futures
 import os
+import pathlib
 import threading
 from typing import Any, Callable, Tuple, TypeVar
 
 import flyte.errors
+from flyte._cache.cache import VersionParameters, cache_from_request
+from flyte._cache.local_cache import LocalTaskCache
 from flyte._context import internal_ctx
 from flyte._internal.controllers import TraceInfo
 from flyte._internal.runtime import convert
 from flyte._internal.runtime.entrypoints import direct_dispatch
+from flyte._internal.runtime.types_serde import transform_native_to_typed_interface
 from flyte._logging import log, logger
-from flyte._task import TaskTemplate
+from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
 from flyte._utils.helpers import _selector_policy
 from flyte.models import ActionID, NativeInterface
 from flyte.remote._task import TaskDetails
@@ -81,31 +85,67 @@ class LocalController:
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
 
         inputs = await convert.convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
-        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+        inputs_hash = convert.generate_inputs_hash_from_proto(inputs.proto_inputs)
+        task_interface = transform_native_to_typed_interface(_task.interface)
 
         sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
-            tctx, _task.name, serialized_inputs, 0
+            tctx, _task.name, inputs_hash, 0
         )
         sub_action_raw_data_path = tctx.raw_data_path
-
-        out, err = await direct_dispatch(
-            _task,
-            controller=self,
-            action=sub_action_id,
-            raw_data_path=sub_action_raw_data_path,
-            inputs=inputs,
-            version=tctx.version,
-            checkpoints=tctx.checkpoints,
-            code_bundle=tctx.code_bundle,
-            output_path=sub_action_output_path,
-            run_base_dir=tctx.run_base_dir,
+        # Make sure the output path exists
+        pathlib.Path(sub_action_output_path).mkdir(parents=True, exist_ok=True)
+        pathlib.Path(sub_action_raw_data_path.path).mkdir(parents=True, exist_ok=True)
+
+        task_cache = cache_from_request(_task.cache)
+        cache_enabled = task_cache.is_enabled()
+        if isinstance(_task, AsyncFunctionTaskTemplate):
+            version_parameters = VersionParameters(func=_task.func, image=_task.image)
+        else:
+            version_parameters = VersionParameters(func=None, image=_task.image)
+        cache_version = task_cache.get_version(version_parameters)
+        cache_key = convert.generate_cache_key_hash(
+            _task.name,
+            inputs_hash,
+            task_interface,
+            cache_version,
+            list(task_cache.get_ignored_inputs()),
+            inputs.proto_inputs,
         )
-        if err:
-            exc = convert.convert_error_to_native(err)
-            if exc:
-                raise exc
-            else:
-                raise flyte.errors.RuntimeSystemError("BadError", "Unknown error")
+
+        out = None
+        # We only get output from cache if the cache behavior is set to auto
+        if task_cache.behavior == "auto":
+            out = await LocalTaskCache.get(cache_key)
+            if out is not None:
+                logger.info(
+                    f"Cache hit for task '{_task.name}' (version: {cache_version}), getting result from cache..."
+                )
+
+        if out is None:
+            out, err = await direct_dispatch(
+                _task,
+                controller=self,
+                action=sub_action_id,
+                raw_data_path=sub_action_raw_data_path,
+                inputs=inputs,
+                version=cache_version,
+                checkpoints=tctx.checkpoints,
+                code_bundle=tctx.code_bundle,
+                output_path=sub_action_output_path,
+                run_base_dir=tctx.run_base_dir,
+            )
+
+            if err:
+                exc = convert.convert_error_to_native(err)
+                if exc:
+                    raise exc
+                else:
+                    raise flyte.errors.RuntimeSystemError("BadError", "Unknown error")
+
+            # store into cache
+            if cache_enabled and out is not None:
+                await LocalTaskCache.set(cache_key, out)
+
         if _task.native_interface.outputs:
             if out is None:
                 raise flyte.errors.RuntimeSystemError("BadOutput", "Task output not captured.")
@@ -129,7 +169,7 @@ class LocalController:
         pass
 
     async def stop(self):
-        pass
+        await LocalTaskCache.close()
 
     async def watch_for_errors(self):
         pass
@@ -151,11 +191,11 @@ class LocalController:
             converted_inputs = await convert.convert_from_native_to_inputs(_interface, *args, **kwargs)
             assert converted_inputs
 
-        serialized_inputs = converted_inputs.proto_inputs.SerializeToString(deterministic=True)
+        inputs_hash = convert.generate_inputs_hash_from_proto(converted_inputs.proto_inputs)
         action_id, action_output_path = convert.generate_sub_action_id_and_output_path(
             tctx,
             _func.__name__,
-            serialized_inputs,
+            inputs_hash,
             0,
         )
         assert action_output_path

flyte/_internal/controllers/remote/_action.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Literal
+from typing import Literal, Optional
 
 from flyteidl.core import execution_pb2, interface_pb2
 from google.protobuf import timestamp_pb2
@@ -39,6 +39,7 @@ class Action:
     phase: run_definition_pb2.Phase | None = None
     started: bool = False
     retries: int = 0
+    queue: Optional[str] = None  # The queue to which this action was submitted.
     client_err: Exception | None = None  # This error is set when something goes wrong in the controller.
     cache_key: str | None = None  # None means no caching, otherwise it is the version of the cache.
 
@@ -122,6 +123,7 @@ class Action:
         inputs_uri: str,
         run_output_base: str,
         cache_key: str | None = None,
+        queue: Optional[str] = None,
     ) -> Action:
         return cls(
             action_id=sub_action_id,
@@ -132,6 +134,7 @@ class Action:
             inputs_uri=inputs_uri,
             run_output_base=run_output_base,
             cache_key=cache_key,
+            queue=queue,
         )
 
     @classmethod

flyte/_internal/controllers/remote/_controller.py

@@ -126,7 +126,7 @@ class RemoteController(Controller):
             workers=workers,
             max_system_retries=max_system_retries,
         )
-        default_parent_concurrency = int(os.getenv("_F_P_CNC", "100"))
+        default_parent_concurrency = int(os.getenv("_F_P_CNC", "1000"))
         self._default_parent_concurrency = default_parent_concurrency
         self._parent_action_semaphore: DefaultDict[str, asyncio.Semaphore] = defaultdict(
             lambda: asyncio.Semaphore(default_parent_concurrency)
@@ -238,6 +238,7 @@ class RemoteController(Controller):
             inputs_uri=inputs_uri,
             run_output_base=tctx.run_base_dir,
             cache_key=cache_key,
+            queue=_task.queue,
         )
 
         try:
@@ -377,9 +378,10 @@ class RemoteController(Controller):
         invoke_seq_num = self.generate_task_call_sequence(_func, current_action_id)
         inputs = await convert.convert_from_native_to_inputs(_interface, *args, **kwargs)
         serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+        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, func_name, serialized_inputs, invoke_seq_num
+            tctx, func_name, inputs_hash, invoke_seq_num
         )
 
         inputs_uri = io.inputs_path(sub_action_output_path)
@@ -539,6 +541,7 @@ class RemoteController(Controller):
             inputs_uri=inputs_uri,
             run_output_base=tctx.run_base_dir,
             cache_key=cache_key,
+            queue=None,
         )
 
         try:

flyte/_internal/controllers/remote/_core.py

@@ -118,13 +118,14 @@ class Controller:
                     raise RuntimeError("Failure event not initialized")
                 self._failure_event.set()
         except asyncio.CancelledError:
-            pass
+            raise
 
     async def _bg_watch_for_errors(self):
         if self._failure_event is None:
             raise RuntimeError("Failure event not initialized")
         await self._failure_event.wait()
         logger.warning(f"Failure event received: {self._failure_event}, cleaning up informers and exiting.")
+        self._running = False
 
     async def watch_for_errors(self):
         """Watch for errors in the background thread"""
@@ -351,6 +352,7 @@ class Controller:
                         ),
                         spec=action.task,
                         cache_key=cache_key,
+                        cluster=action.queue,
                     )
                 elif action.type == "trace":
                     trace = action.trace
@@ -440,10 +442,11 @@ class Controller:
                 logger.warning(f"[{worker_id}] Retrying action {action.name} after backoff")
                 await self._shared_queue.put(action)
             except Exception as e:
-                logger.error(f"[{worker_id}] Error in controller loop: {e}")
+                logger.error(f"[{worker_id}] Error in controller loop for {action.name}: {e}")
                 err = flyte.errors.RuntimeSystemError(
                     code=type(e).__name__,
-                    message=f"Controller failed, system retries {action.retries} crossed threshold {self._max_retries}",
+                    message=f"Controller failed, system retries {action.retries} / {self._max_retries} "
+                    f"crossed threshold, for action {action.name}: {e}",
                     worker=worker_id,
                 )
                 err.__cause__ = e

flyte/_internal/controllers/remote/_informer.py

@@ -270,7 +270,7 @@ class Informer:
             logger.warning("Informer already running")
             return cast(asyncio.Task, self._watch_task)
         self._running = True
-        self._watch_task = asyncio.create_task(self.watch())
+        self._watch_task = asyncio.create_task(self.watch(), name=f"InformerWatch-{self.parent_action_name}")
         await self.wait_for_cache_sync(timeout=timeout)
         return self._watch_task
 

flyte/_internal/imagebuild/docker_builder.py

@@ -22,6 +22,7 @@ from flyte._image import (
     Layer,
     PipOption,
     PipPackages,
+    PoetryProject,
     PythonWheels,
     Requirements,
     UVProject,
@@ -46,44 +47,71 @@ FLYTE_DOCKER_BUILDER_CACHE_TO = "FLYTE_DOCKER_BUILDER_CACHE_TO"
 
 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_PATH \
-    --mount=type=bind,target=pyproject.toml,src=$PYPROJECT_PATH \
-    $SECRET_MOUNT \
-    uv sync --active $PIP_INSTALL_ARGS
+   --mount=type=bind,target=uv.lock,src=$UV_LOCK_PATH \
+   --mount=type=bind,target=pyproject.toml,src=$PYPROJECT_PATH \
+   $SECRET_MOUNT \
+   uv sync --active --inexact $PIP_INSTALL_ARGS
 """)
 
 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
+   --mount=type=bind,target=/root/.flyte/$PYPROJECT_PATH,src=$PYPROJECT_PATH,rw \
+   $SECRET_MOUNT \
+   uv sync --active --inexact --no-editable $PIP_INSTALL_ARGS --project /root/.flyte/$PYPROJECT_PATH
+""")
+
+POETRY_LOCK_WITHOUT_PROJECT_INSTALL_TEMPLATE = Template("""\
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
+   uv pip install poetry
+
+ENV POETRY_CACHE_DIR=/tmp/poetry_cache \
+   POETRY_VIRTUALENVS_IN_PROJECT=true
+
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/tmp/poetry_cache,id=poetry \
+   --mount=type=bind,target=poetry.lock,src=$POETRY_LOCK_PATH \
+   --mount=type=bind,target=pyproject.toml,src=$PYPROJECT_PATH \
+   $SECRET_MOUNT \
+   poetry install $POETRY_INSTALL_ARGS
+""")
+
+POETRY_LOCK_INSTALL_TEMPLATE = Template("""\
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
+   uv pip install poetry
+
+ENV POETRY_CACHE_DIR=/tmp/poetry_cache \
+   POETRY_VIRTUALENVS_IN_PROJECT=true
+
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/tmp/poetry_cache,id=poetry \
+   --mount=type=bind,target=/root/.flyte/$PYPROJECT_PATH,src=$PYPROJECT_PATH,rw \
+   $SECRET_MOUNT \
+   poetry install $POETRY_INSTALL_ARGS
 """)
 
 UV_PACKAGE_INSTALL_COMMAND_TEMPLATE = Template("""\
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
-    $REQUIREMENTS_MOUNT \
-    $SECRET_MOUNT \
-    uv pip install --python $$UV_PYTHON $PIP_INSTALL_ARGS
+   $REQUIREMENTS_MOUNT \
+   $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 \
-    $SECRET_MOUNT \
-    uv pip install --python $$UV_PYTHON $PIP_INSTALL_ARGS
+   --mount=source=/dist,target=/dist,type=bind \
+   $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
+   $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
+   $SECRET_MOUNT \
+   uv pip install $PIP_INSTALL_ARGS
 """)
 
 # uv pip install --python /root/env/bin/python
@@ -93,24 +121,29 @@ DOCKER_FILE_UV_BASE_TEMPLATE = Template("""\
 FROM ghcr.io/astral-sh/uv:0.8.13 AS uv
 FROM $BASE_IMAGE
 
+
 USER root
 
+
 # Copy in uv so that later commands don't have to mount it in
 COPY --from=uv /uv /usr/bin/uv
 
+
 # Configure default envs
 ENV UV_COMPILE_BYTECODE=1 \
-    UV_LINK_MODE=copy \
-    VIRTUALENV=/opt/venv \
-    UV_PYTHON=/opt/venv/bin/python \
-    PATH="/opt/venv/bin:$$PATH"
+   UV_LINK_MODE=copy \
+   VIRTUALENV=/opt/venv \
+   UV_PYTHON=/opt/venv/bin/python \
+   PATH="/opt/venv/bin:$$PATH"
+
 
 # Create a virtualenv with the user specified python version
 RUN uv venv $$VIRTUALENV --python=$PYTHON_VERSION
 
+
 # Adds nvidia just in case it exists
 ENV PATH="$$PATH:/usr/local/nvidia/bin:/usr/local/cuda/bin" \
-    LD_LIBRARY_PATH="/usr/local/nvidia/lib64"
+   LD_LIBRARY_PATH="/usr/local/nvidia/lib64"
 """)
 
 # This gets added on to the end of the dockerfile
@@ -258,6 +291,32 @@ class UVProjectHandler:
         return dockerfile
 
 
+class PoetryProjectHandler:
+    @staticmethod
+    async def handel(layer: PoetryProject, context_path: Path, dockerfile: str) -> str:
+        secret_mounts = _get_secret_mounts_layer(layer.secret_mounts)
+        if layer.extra_args and "--no-root" in layer.extra_args:
+            # Only Copy pyproject.yaml and poetry.lock.
+            pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
+            poetry_lock_dst = copy_files_to_context(layer.poetry_lock, context_path)
+            delta = POETRY_LOCK_WITHOUT_PROJECT_INSTALL_TEMPLATE.substitute(
+                POETRY_LOCK_PATH=poetry_lock_dst.relative_to(context_path),
+                PYPROJECT_PATH=pyproject_dst.relative_to(context_path),
+                POETRY_INSTALL_ARGS=layer.extra_args or "",
+                SECRET_MOUNT=secret_mounts,
+            )
+        else:
+            # Copy the entire project.
+            pyproject_dst = copy_files_to_context(layer.pyproject.parent, context_path)
+            delta = POETRY_LOCK_INSTALL_TEMPLATE.substitute(
+                PYPROJECT_PATH=pyproject_dst.relative_to(context_path),
+                POETRY_INSTALL_ARGS=layer.extra_args or "",
+                SECRET_MOUNT=secret_mounts,
+            )
+        dockerfile += delta
+        return dockerfile
+
+
 class DockerIgnoreHandler:
     @staticmethod
     async def handle(layer: DockerIgnore, context_path: Path, _: str):
@@ -332,8 +391,9 @@ class CommandsHandler:
     @staticmethod
     async def handle(layer: Commands, _: Path, dockerfile: str) -> str:
         # Append raw commands to the dockerfile
+        secret_mounts = _get_secret_mounts_layer(layer.secret_mounts)
         for command in layer.commands:
-            dockerfile += f"\nRUN {command}\n"
+            dockerfile += f"\nRUN {secret_mounts} {command}\n"
 
         return dockerfile
 
@@ -355,9 +415,8 @@ def _get_secret_commands(layers: typing.Tuple[Layer, ...]) -> typing.List[str]:
             secret = Secret(key=secret)
         secret_id = hash(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}"]
+        if os.getenv(secret_env_key):
+            return ["--secret", f"id={secret_id},env={secret_env_key}"]
         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):
@@ -365,7 +424,7 @@ def _get_secret_commands(layers: typing.Tuple[Layer, ...]) -> typing.List[str]:
         return ["--secret", f"id={secret_id},src={secret_file_path}"]
 
     for layer in layers:
-        if isinstance(layer, (PipOption, AptPackages)):
+        if isinstance(layer, (PipOption, AptPackages, Commands)):
             if layer.secret_mounts:
                 for secret_mount in layer.secret_mounts:
                     commands.extend(_get_secret_command(secret_mount))
@@ -426,6 +485,10 @@ async def _process_layer(
             # Handle UV project
             dockerfile = await UVProjectHandler.handle(layer, context_path, dockerfile)
 
+        case PoetryProject():
+            # Handle Poetry project
+            dockerfile = await PoetryProjectHandler.handel(layer, context_path, dockerfile)
+
         case CopyConfig():
             # Handle local files and folders
             dockerfile = await CopyConfigHandler.handle(layer, context_path, dockerfile, docker_ignore_file_path)
@@ -572,6 +635,7 @@ class DockerImageBuilder(ImageBuilder):
           - start from the base image
           - use python to create a default venv and export variables
 
+
           Then for the layers
           - for each layer
             - find the appropriate layer handler

flyte/_internal/imagebuild/image_builder.py

@@ -135,11 +135,6 @@ class ImageBuildEngine:
 
     ImageBuilderType = typing.Literal["local", "remote"]
 
-    _SEEN_IMAGES: typing.ClassVar[typing.Dict[str, str]] = {
-        # Set default for the auto container. See Image._identifier_override for more info.
-        "auto": Image.from_debian_base().uri,
-    }
-
     @staticmethod
     @alru_cache
     async def image_exists(image: Image) -> Optional[str]:
@@ -235,7 +230,7 @@ class ImageBuildEngine:
 
 
 class ImageCache(BaseModel):
-    image_lookup: Dict[str, Dict[str, str]]
+    image_lookup: Dict[str, str]
     serialized_form: str | None = None
 
     @property
@@ -273,11 +268,10 @@ class ImageCache(BaseModel):
         """
         tuples = []
         for k, v in self.image_lookup.items():
-            for py_version, image_uri in v.items():
-                tuples.append(
-                    [
-                        ("Name", f"{k} (py{py_version})"),
-                        ("image", image_uri),
-                    ]
-                )
+            tuples.append(
+                [
+                    ("Name", k),
+                    ("image", v),
+                ]
+            )
         return tuples

flyte/_internal/imagebuild/remote_builder.py

@@ -182,6 +182,11 @@ async def _validate_configuration(image: Image) -> Tuple[str, Optional[str]]:
 def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2.ImageSpec":
     from flyte._protos.imagebuilder import definition_pb2 as image_definition_pb2
 
+    if image.dockerfile is not None:
+        raise flyte.errors.ImageBuildError(
+            "Custom Dockerfile is not supported with remote image builder.You can use local image builder instead."
+        )
+
     layers = []
     for layer in image._layers:
         secret_mounts = None
@@ -251,7 +256,7 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
                 if "tool.uv.index" in line:
                     raise ValueError("External sources are not supported in pyproject.toml")
 
-            if layer.extra_index_urls and "--no-install-project" in layer.extra_index_urls:
+            if layer.extra_args and "--no-install-project" in layer.extra_args:
                 # Copy pyproject itself
                 pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
             else:

flyte/_internal/runtime/io.py

@@ -9,6 +9,7 @@ from flyteidl.core import errors_pb2, execution_pb2
 
 import flyte.storage as storage
 from flyte._protos.workflow import run_definition_pb2
+from flyte.models import PathRewrite
 
 from .convert import Inputs, Outputs, _clean_error_code
 
@@ -90,10 +91,11 @@ async def upload_error(err: execution_pb2.ExecutionError, output_prefix: str):
 
 
 # ------------------------------- DOWNLOAD Methods ------------------------------- #
-async def load_inputs(path: str, max_bytes: int = -1) -> Inputs:
+async def load_inputs(path: str, max_bytes: int = -1, path_rewrite_config: PathRewrite | None = None) -> Inputs:
     """
     :param path: Input file to be downloaded
     :param max_bytes: Maximum number of bytes to read from the input file. Default is -1, which means no limit.
+    :param path_rewrite_config: If provided, rewrites paths in the input blobs according to the configuration.
     :return: Inputs object
     """
     lm = run_definition_pb2.Inputs()
@@ -115,6 +117,16 @@ async def load_inputs(path: str, max_bytes: int = -1) -> Inputs:
         proto_str = b"".join(proto_bytes)
 
     lm.ParseFromString(proto_str)
+
+    if path_rewrite_config is not None:
+        for inp in lm.literals:
+            if inp.value.HasField("scalar") and inp.value.scalar.HasField("blob"):
+                scalar_blob = inp.value.scalar.blob
+                if scalar_blob.uri.startswith(path_rewrite_config.old_prefix):
+                    scalar_blob.uri = scalar_blob.uri.replace(
+                        path_rewrite_config.old_prefix, path_rewrite_config.new_prefix, 1
+                    )
+
     return Inputs(proto_inputs=lm)
 
 

flyte/_internal/runtime/rusty.py

@@ -11,7 +11,7 @@ from flyte._internal.runtime.entrypoints import download_code_bundle, load_pkl_t
 from flyte._internal.runtime.taskrunner import extract_download_run_upload
 from flyte._logging import logger
 from flyte._task import TaskTemplate
-from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
+from flyte.models import ActionID, Checkpoints, CodeBundle, PathRewrite, RawDataPath
 
 
 async def download_tgz(destination: str, version: str, tgz: str) -> CodeBundle:
@@ -115,6 +115,7 @@ async def run_task(
     prev_checkpoint: str | None = None,
     code_bundle: CodeBundle | None = None,
     input_path: str | None = None,
+    path_rewrite_cfg: str | None = None,
 ):
     """
     Runs the task with the provided parameters.
@@ -134,6 +135,7 @@ async def run_task(
     :param controller: The controller to use for the task.
     :param code_bundle: Optional code bundle for the task.
     :param input_path: Optional input path for the task.
+    :param path_rewrite_cfg: Optional path rewrite configuration.
     :return: The loaded task template.
     """
     start_time = time.time()
@@ -144,6 +146,19 @@ async def run_task(
         f" at {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))}"
     )
 
+    path_rewrite = PathRewrite.from_str(path_rewrite_cfg) if path_rewrite_cfg else None
+    if path_rewrite:
+        import flyte.storage as storage
+
+        if not await storage.exists(path_rewrite.new_prefix):
+            logger.error(
+                f"[rusty] Path rewrite failed for path {path_rewrite.new_prefix}, "
+                f"not found, reverting to original path {path_rewrite.old_prefix}"
+            )
+            path_rewrite = None
+        else:
+            logger.info(f"[rusty] Using path rewrite: {path_rewrite}")
+
     try:
         await contextual_run(
             extract_download_run_upload,
@@ -151,7 +166,7 @@ async def run_task(
             action=ActionID(name=name, org=org, project=project, domain=domain, run_name=run_name),
             version=version,
             controller=controller,
-            raw_data_path=RawDataPath(path=raw_data_path),
+            raw_data_path=RawDataPath(path=raw_data_path, path_rewrite=path_rewrite),
             output_path=output_path,
             run_base_dir=run_base_dir,
             checkpoints=Checkpoints(prev_checkpoint_path=prev_checkpoint, checkpoint_path=checkpoint_path),