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

flyte/_internal/resolvers/_task_module.py

@@ -1,13 +1,11 @@
-import inspect
-import os
 import pathlib
-import sys
 from typing import Tuple
 
+from flyte._module import extract_obj_module
 from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
 
 
-def extract_task_module(task: TaskTemplate, /, source_dir: pathlib.Path | None = None) -> Tuple[str, str]:
+def extract_task_module(task: TaskTemplate, /, source_dir: pathlib.Path) -> Tuple[str, str]:
     """
     Extract the task module from the task template.
 
@@ -15,40 +13,9 @@ def extract_task_module(task: TaskTemplate, /, source_dir: pathlib.Path | None =
     :param source_dir: The source directory to use for relative paths.
     :return: A tuple containing the entity name, module
     """
-    entity_name = task.name
     if isinstance(task, AsyncFunctionTaskTemplate):
-        entity_module = inspect.getmodule(task.func)
-        if entity_module is None:
-            raise ValueError(f"Task {entity_name} has no module.")
-
-        fp = entity_module.__file__
-        if fp is None:
-            raise ValueError(f"Task {entity_name} has no module.")
-
-        file_path = pathlib.Path(fp)
-        # Get the relative path to the current directory
-        # Will raise ValueError if the file is not in the source directory
-        relative_path = file_path.relative_to(str(source_dir))
-
-        if relative_path == pathlib.Path("."):
-            entity_module_name = entity_module.__name__
-        else:
-            # Replace file separators with dots and remove the '.py' extension
-            dotted_path = os.path.splitext(str(relative_path))[0].replace(os.sep, ".")
-            entity_module_name = dotted_path
-
         entity_name = task.func.__name__
+        entity_module_name = extract_obj_module(task.func, source_dir)
+        return entity_name, entity_module_name
     else:
-        raise NotImplementedError(f"Task module {entity_name} not implemented.")
-
-    if entity_module_name == "__main__":
-        """
-        This case is for the case in which the task is run from the main module.
-        """
-        fp = sys.modules["__main__"].__file__
-        if fp is None:
-            raise ValueError(f"Task {entity_name} has no module.")
-        main_path = pathlib.Path(fp)
-        entity_module_name = main_path.stem
-
-    return entity_name, entity_module_name
+        raise NotImplementedError(f"Task module {task.name} not implemented.")

flyte/_internal/resolvers/default.py

@@ -1,6 +1,6 @@
 import importlib
 from pathlib import Path
-from typing import List, Optional
+from typing import List
 
 from flyte._internal.resolvers._task_module import extract_task_module
 from flyte._internal.resolvers.common import Resolver
@@ -23,6 +23,6 @@ class DefaultTaskResolver(Resolver):
         task_def = getattr(task_module, task_name)
         return task_def
 
-    def loader_args(self, task: TaskTemplate, root_dir: Optional[Path] = None) -> List[str]:  # type:ignore
+    def loader_args(self, task: TaskTemplate, root_dir: Path) -> List[str]:  # type:ignore
         t, m = extract_task_module(task, root_dir)
         return ["mod", m, "instance", t]

flyte/_internal/runtime/convert.py

@@ -8,27 +8,33 @@ from dataclasses import dataclass
 from types import NoneType
 from typing import Any, Dict, List, Tuple, Union, get_args
 
-from flyteidl.core import execution_pb2, interface_pb2, literals_pb2
+from flyteidl2.core import execution_pb2, interface_pb2, literals_pb2
+from flyteidl2.task import common_pb2, task_definition_pb2
 
 import flyte.errors
 import flyte.storage as storage
-from flyte._protos.workflow import common_pb2, run_definition_pb2, task_definition_pb2
+from flyte._context import ctx
 from flyte.models import ActionID, NativeInterface, TaskContext
 from flyte.types import TypeEngine, TypeTransformerFailedError
 
 
 @dataclass(frozen=True)
 class Inputs:
-    proto_inputs: run_definition_pb2.Inputs
+    proto_inputs: common_pb2.Inputs
 
     @classmethod
     def empty(cls) -> "Inputs":
-        return cls(proto_inputs=run_definition_pb2.Inputs())
+        return cls(proto_inputs=common_pb2.Inputs())
+
+    @property
+    def context(self) -> Dict[str, str]:
+        """Get the context as a dictionary."""
+        return {kv.key: kv.value for kv in self.proto_inputs.context}
 
 
 @dataclass(frozen=True)
 class Outputs:
-    proto_outputs: run_definition_pb2.Outputs
+    proto_outputs: common_pb2.Outputs
 
 
 @dataclass
@@ -102,15 +108,30 @@ def is_optional_type(tp) -> bool:
     return NoneType in get_args(tp)  # fastest check
 
 
-async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwargs) -> Inputs:
+async def convert_from_native_to_inputs(
+    interface: NativeInterface, *args, custom_context: Dict[str, str] | None = None, **kwargs
+) -> Inputs:
     kwargs = interface.convert_to_kwargs(*args, **kwargs)
 
     missing = [key for key in interface.required_inputs() if key not in kwargs]
     if missing:
         raise ValueError(f"Missing required inputs: {', '.join(missing)}")
 
+    # Read custom_context from TaskContext if available (inside task execution)
+    # Otherwise use the passed parameter (for remote run initiation)
+    context_kvs = None
+    tctx = ctx()
+    if tctx and tctx.custom_context:
+        # Inside a task - read from TaskContext
+        context_to_use = tctx.custom_context
+        context_kvs = [literals_pb2.KeyValuePair(key=k, value=v) for k, v in context_to_use.items()]
+    elif custom_context:
+        # Remote run initiation
+        context_kvs = [literals_pb2.KeyValuePair(key=k, value=v) for k, v in custom_context.items()]
+
     if len(interface.inputs) == 0:
-        return Inputs.empty()
+        # Handle context even for empty inputs
+        return Inputs(proto_inputs=common_pb2.Inputs(context=context_kvs))
 
     # fill in defaults if missing
     type_hints: Dict[str, type] = {}
@@ -122,13 +143,14 @@ async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwa
             (default_value is not None and default_value is not inspect.Signature.empty)
             or (default_value is None and is_optional_type(input_type))
             or input_type is None
+            or input_type is type(None)
         ):
             if default_value == NativeInterface.has_default:
                 if interface._remote_defaults is None or input_name not in interface._remote_defaults:
                     raise ValueError(f"Input '{input_name}' has a default value but it is not set in the interface.")
                 already_converted_kwargs[input_name] = interface._remote_defaults[input_name]
-            elif input_type is None:
-                # If the type is None, we assume it's a placeholder for no type
+            elif input_type is None or input_type is type(None):
+                # If the type is 'None' or 'class<None>', we assume it's a placeholder for no type
                 kwargs[input_name] = None
                 type_hints[input_name] = NoneType
             else:
@@ -144,12 +166,12 @@ async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwa
         for k, v in already_converted_kwargs.items():
             copied_literals[k] = v
         literal_map = literals_pb2.LiteralMap(literals=copied_literals)
+
     # Make sure we the interface, not literal_map or kwargs, because those may have a different order
     return Inputs(
-        proto_inputs=run_definition_pb2.Inputs(
-            literals=[
-                run_definition_pb2.NamedLiteral(name=k, value=literal_map.literals[k]) for k in interface.inputs.keys()
-            ]
+        proto_inputs=common_pb2.Inputs(
+            literals=[common_pb2.NamedLiteral(name=k, value=literal_map.literals[k]) for k in interface.inputs.keys()],
+            context=context_kvs,
         )
     )
 
@@ -191,11 +213,11 @@ async def convert_from_native_to_outputs(o: Any, interface: NativeInterface, tas
     for (output_name, python_type), v in zip(interface.outputs.items(), o):
         try:
             lit = await TypeEngine.to_literal(v, python_type, TypeEngine.to_literal_type(python_type))
-            named.append(run_definition_pb2.NamedLiteral(name=output_name, value=lit))
+            named.append(common_pb2.NamedLiteral(name=output_name, value=lit))
         except TypeTransformerFailedError as e:
             raise flyte.errors.RuntimeDataValidationError(output_name, e, task_name)
 
-    return Outputs(proto_outputs=run_definition_pb2.Outputs(literals=named))
+    return Outputs(proto_outputs=common_pb2.Outputs(literals=named))
 
 
 async def convert_outputs_to_native(interface: NativeInterface, outputs: Outputs) -> Union[Any, Tuple[Any, ...]]:
@@ -222,7 +244,7 @@ def convert_error_to_native(err: execution_pb2.ExecutionError | Exception | Erro
     if isinstance(err, Error):
         err = err.err
 
-    user_code, server_code = _clean_error_code(err.code)
+    user_code, _server_code = _clean_error_code(err.code)
     match err.kind:
         case execution_pb2.ExecutionError.UNKNOWN:
             return flyte.errors.RuntimeUnknownError(code=user_code, message=err.message, worker=err.worker)
@@ -351,7 +373,7 @@ def generate_inputs_repr_for_literal(literal: literals_pb2.Literal) -> bytes:
     return literal.SerializeToString(deterministic=True)
 
 
-def generate_inputs_hash_for_named_literals(inputs: list[run_definition_pb2.NamedLiteral]) -> str:
+def generate_inputs_hash_for_named_literals(inputs: list[common_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
@@ -375,7 +397,7 @@ def generate_inputs_hash_for_named_literals(inputs: list[run_definition_pb2.Name
     return hash_data(combined_bytes)
 
 
-def generate_inputs_hash_from_proto(inputs: run_definition_pb2.Inputs) -> str:
+def generate_inputs_hash_from_proto(inputs: common_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.
@@ -404,7 +426,7 @@ def generate_cache_key_hash(
     task_interface: interface_pb2.TypedInterface,
     cache_version: str,
     ignored_input_vars: List[str],
-    proto_inputs: run_definition_pb2.Inputs,
+    proto_inputs: common_pb2.Inputs,
 ) -> str:
     """
     Generate a cache key hash based on the inputs hash, task name, task interface, and cache version.
@@ -420,7 +442,7 @@ 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 = common_pb2.Inputs(literals=filtered)
         final_inputs = generate_inputs_hash_from_proto(final)
     else:
         final_inputs = inputs_hash

flyte/_internal/runtime/entrypoints.py

@@ -1,4 +1,6 @@
 import importlib
+import os
+import traceback
 from typing import List, Optional, Tuple, Type
 
 import flyte.errors
@@ -10,6 +12,7 @@ from flyte._logging import log, logger
 from flyte._task import TaskTemplate
 from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
 
+from ..._utils import adjust_sys_path
 from .convert import Error, Inputs, Outputs
 from .taskrunner import (
     convert_and_run,
@@ -72,7 +75,26 @@ def load_task(resolver: str, *resolver_args: str) -> TaskTemplate:
     """
     resolver_class = load_class(resolver)
     resolver_instance = resolver_class()
-    return resolver_instance.load_task(resolver_args)
+    try:
+        return resolver_instance.load_task(resolver_args)
+    except ModuleNotFoundError as e:
+        cwd = os.getcwd()
+        files = []
+        try:
+            for root, dirs, filenames in os.walk(cwd):
+                for name in dirs + filenames:
+                    rel_path = os.path.relpath(os.path.join(root, name), cwd)
+                    files.append(rel_path)
+        except Exception as list_err:
+            files = [f"(Failed to list directory: {list_err})"]
+
+        msg = (
+            "\n\nFull traceback:\n" + "".join(traceback.format_exc()) + f"\n[ImportError Diagnostics]\n"
+            f"Module '{e.name}' not found in either the Python virtual environment or the current working directory.\n"
+            f"Current working directory: {cwd}\n"
+            f"Files found under current directory:\n" + "\n".join(f"  - {f}" for f in files)
+        )
+        raise ModuleNotFoundError(msg) from e
 
 
 def load_pkl_task(code_bundle: CodeBundle) -> TaskTemplate:
@@ -100,6 +122,7 @@ async def download_code_bundle(code_bundle: CodeBundle) -> CodeBundle:
     :param code_bundle: The code bundle to download.
     :return: The code bundle with the downloaded path.
     """
+    adjust_sys_path()
     logger.debug(f"Downloading {code_bundle}")
     downloaded_path = await download_bundle(code_bundle)
     return code_bundle.with_downloaded_path(downloaded_path)

flyte/_internal/runtime/io.py

@@ -5,10 +5,12 @@ It uses the storage module to handle the actual uploading and downloading of fil
 TODO: Convert to use streaming apis
 """
 
-from flyteidl.core import errors_pb2, execution_pb2
+from flyteidl.core import errors_pb2
+from flyteidl2.core import execution_pb2
+from flyteidl2.task import common_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
 
@@ -69,7 +71,7 @@ async def upload_outputs(outputs: Outputs, output_path: str, max_bytes: int = -1
     await storage.put_stream(data_iterable=outputs.proto_outputs.SerializeToString(), to_path=output_uri)
 
 
-async def upload_error(err: execution_pb2.ExecutionError, output_prefix: str):
+async def upload_error(err: execution_pb2.ExecutionError, output_prefix: str) -> str:
     """
     :param err: execution_pb2.ExecutionError
     :param output_prefix: The output prefix of the remote uri.
@@ -86,17 +88,18 @@ async def upload_error(err: execution_pb2.ExecutionError, output_prefix: str):
         )
     )
     error_uri = error_path(output_prefix)
-    await storage.put_stream(data_iterable=error_document.SerializeToString(), to_path=error_uri)
+    return await storage.put_stream(data_iterable=error_document.SerializeToString(), to_path=error_uri)
 
 
 # ------------------------------- 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()
+    lm = common_pb2.Inputs()
 
     if max_bytes == -1:
         proto_str = b"".join([c async for c in storage.get_stream(path=path)])
@@ -115,6 +118,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)
 
 
@@ -125,7 +138,7 @@ async def load_outputs(path: str, max_bytes: int = -1) -> Outputs:
                       If -1, reads the entire file.
     :return: Outputs object
     """
-    lm = run_definition_pb2.Outputs()
+    lm = common_pb2.Outputs()
 
     if max_bytes == -1:
         proto_str = b"".join([c async for c in storage.get_stream(path=path)])
@@ -157,7 +170,7 @@ async def load_error(path: str) -> execution_pb2.ExecutionError:
     err.ParseFromString(proto_str)
 
     if err.error is not None:
-        user_code, server_code = _clean_error_code(err.error.code)
+        user_code, _server_code = _clean_error_code(err.error.code)
         return execution_pb2.ExecutionError(
             code=user_code,
             message=err.error.message,

flyte/_internal/runtime/resources_serde.py

@@ -1,8 +1,8 @@
-from typing import List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
-from flyteidl.core import tasks_pb2
+from flyteidl2.core import tasks_pb2
 
-from flyte._resources import CPUBaseType, Resources
+from flyte._resources import CPUBaseType, DeviceClass, Resources
 
 ACCELERATOR_DEVICE_MAP = {
     "A100": "nvidia-tesla-a100",
@@ -24,6 +24,14 @@ ACCELERATOR_DEVICE_MAP = {
     "V6E": "tpu-v6e-slice",
 }
 
+_DeviceClassToProto: Dict[DeviceClass, "tasks_pb2.GPUAccelerator.DeviceClass"] = {
+    "GPU": tasks_pb2.GPUAccelerator.NVIDIA_GPU,
+    "TPU": tasks_pb2.GPUAccelerator.GOOGLE_TPU,
+    "NEURON": tasks_pb2.GPUAccelerator.AMAZON_NEURON,
+    "AMD_GPU": tasks_pb2.GPUAccelerator.AMD_GPU,
+    "HABANA_GAUDI": tasks_pb2.GPUAccelerator.HABANA_GAUDI,
+}
+
 
 def _get_cpu_resource_entry(cpu: CPUBaseType) -> tasks_pb2.Resources.ResourceEntry:
     return tasks_pb2.Resources.ResourceEntry(
@@ -54,11 +62,17 @@ def _get_gpu_extended_resource_entry(resources: Resources) -> Optional[tasks_pb2
     device = resources.get_device()
     if device is None:
         return None
-    if device.device not in ACCELERATOR_DEVICE_MAP:
-        raise ValueError(f"GPU of type {device.device} unknown, cannot map to device name")
+
+    device_class = _DeviceClassToProto.get(device.device_class, tasks_pb2.GPUAccelerator.NVIDIA_GPU)
+    if device.device is None:
+        raise RuntimeError("Device type must be specified for GPU string.")
+    else:
+        device_type = device.device
+    device_type = ACCELERATOR_DEVICE_MAP.get(device_type, device_type)
     return tasks_pb2.GPUAccelerator(
-        device=ACCELERATOR_DEVICE_MAP[device.device],
+        device=device_type,
         partition_size=device.partition if device.partition else None,
+        device_class=device_class,
     )
 
 

flyte/_internal/runtime/reuse.py

@@ -2,7 +2,7 @@ import hashlib
 import typing
 from venv import logger
 
-from flyteidl.core import tasks_pb2
+from flyteidl2.core import tasks_pb2
 
 import flyte.errors
 from flyte import ReusePolicy

flyte/_internal/runtime/rusty.py

@@ -1,5 +1,4 @@
 import asyncio
-import sys
 import time
 from typing import Any, List, Tuple
 
@@ -11,7 +10,8 @@ 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._utils import adjust_sys_path
+from flyte.models import ActionID, Checkpoints, CodeBundle, PathRewrite, RawDataPath
 
 
 async def download_tgz(destination: str, version: str, tgz: str) -> CodeBundle:
@@ -23,7 +23,7 @@ async def download_tgz(destination: str, version: str, tgz: str) -> CodeBundle:
     :return: The CodeBundle object.
     """
     logger.info(f"[rusty] Downloading tgz code bundle from {tgz} to {destination} with version {version}")
-    sys.path.insert(0, ".")
+    adjust_sys_path()
 
     code_bundle = CodeBundle(
         tgz=tgz,
@@ -42,7 +42,7 @@ async def download_load_pkl(destination: str, version: str, pkl: str) -> Tuple[C
     :return: The CodeBundle object.
     """
     logger.info(f"[rusty] Downloading pkl code bundle from {pkl} to {destination} with version {version}")
-    sys.path.insert(0, ".")
+    adjust_sys_path()
 
     code_bundle = CodeBundle(
         pkl=pkl,
@@ -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),

flyte/_internal/runtime/task_serde.py

@@ -4,19 +4,18 @@ It includes a Resolver interface for loading tasks, and functions to load classe
 """
 
 import copy
-import sys
 import typing
 from datetime import timedelta
 from typing import Optional, cast
 
-from flyteidl.core import identifier_pb2, literals_pb2, security_pb2, tasks_pb2
+from flyteidl2.core import identifier_pb2, literals_pb2, security_pb2, tasks_pb2
+from flyteidl2.task import common_pb2, environment_pb2, task_definition_pb2
 from google.protobuf import duration_pb2, wrappers_pb2
 
 import flyte.errors
 from flyte._cache.cache import VersionParameters, cache_from_request
 from flyte._logging import logger
 from flyte._pod import _PRIMARY_CONTAINER_NAME_FIELD, PodTemplate
-from flyte._protos.workflow import common_pb2, environment_pb2, task_definition_pb2
 from flyte._secret import SecretRequest, secrets_from_request
 from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
 from flyte.models import CodeBundle, SerializationContext
@@ -120,7 +119,7 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
         version=serialize_context.version,
     )
 
-    # TODO Add support for SQL, extra_config, custom
+    # TODO Add support for extra_config, custom
     extra_config: typing.Dict[str, str] = {}
 
     if task.pod_template and not isinstance(task.pod_template, str):
@@ -133,7 +132,7 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
 
     custom = task.custom_config(serialize_context)
 
-    sql = None
+    sql = task.sql(serialize_context)
 
     # -------------- CACHE HANDLING ----------------------
     task_cache = cache_from_request(task.cache)
@@ -171,8 +170,9 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
             retries=get_proto_retry_strategy(task.retries),
             timeout=get_proto_timeout(task.timeout),
             pod_template_name=(task.pod_template if task.pod_template and isinstance(task.pod_template, str) else None),
-            interruptible=task.interruptable,
+            interruptible=task.interruptible,
             generates_deck=wrappers_pb2.BoolValue(value=task.report),
+            debuggable=task.debuggable,
         ),
         interface=transform_native_to_typed_interface(task.native_interface),
         custom=custom if len(custom) > 0 else None,
@@ -209,34 +209,40 @@ def _get_urun_container(
         else None
     )
     resources = get_proto_resources(task_template.resources)
-    # pr: under what conditions should this return None?
+
     if isinstance(task_template.image, str):
         raise flyte.errors.RuntimeSystemError("BadConfig", "Image is not a valid image")
-    image_id = task_template.image.identifier
+
+    env_name = task_template.parent_env_name
+    if env_name is None:
+        raise flyte.errors.RuntimeSystemError("BadConfig", f"Task {task_template.name} has no parent environment name")
+
     if not serialize_context.image_cache:
         # This computes the image uri, computing hashes as necessary so can fail if done remotely.
         img_uri = task_template.image.uri
-    elif serialize_context.image_cache and image_id not in serialize_context.image_cache.image_lookup:
-        img_uri = task_template.image.uri
-
-        logger.warning(
-            f"Image {task_template.image} not found in the image cache: {serialize_context.image_cache.image_lookup}."
+    elif serialize_context.image_cache and env_name not in serialize_context.image_cache.image_lookup:
+        raise flyte.errors.RuntimeUserError(
+            "MissingEnvironment",
+            f"Environment '{env_name}' not found in image cache.\n\n"
+            "💡 To fix this:\n"
+            "  1. If your parent environment calls a task in another environment,"
+            " declare that dependency using 'depends_on=[...]'.\n"
+            "     Example:\n"
+            "         env1 = flyte.TaskEnvironment(\n"
+            "             name='outer',\n"
+            "             image=flyte.Image.from_debian_base().with_pip_packages('requests'),\n"
+            "             depends_on=[env2, env3],\n"
+            "         )\n"
+            "  2. If you're using os.getenv() to set the environment name,"
+            " make sure the runtime environment has the same environment variable defined.\n"
+            "     Example:\n"
+            "         env = flyte.TaskEnvironment(\n"
+            '             name=os.getenv("my-name"),\n'
+            '             env_vars={"my-name": os.getenv("my-name")},\n'
+            "         )\n",
         )
     else:
-        python_version_str = "{}.{}".format(sys.version_info.major, sys.version_info.minor)
-        version_lookup = serialize_context.image_cache.image_lookup[image_id]
-        if python_version_str in version_lookup:
-            img_uri = version_lookup[python_version_str]
-        elif version_lookup:
-            # Fallback: try to get any available version
-            fallback_py_version, img_uri = next(iter(version_lookup.items()))
-            logger.warning(
-                f"Image {task_template.image} for python version {python_version_str} "
-                f"not found in the image cache: {serialize_context.image_cache.image_lookup}.\n"
-                f"Fall back using image {img_uri} for python version {fallback_py_version} ."
-            )
-        else:
-            img_uri = task_template.image.uri
+        img_uri = serialize_context.image_cache.image_lookup[env_name]
 
     return tasks_pb2.Container(
         image=img_uri,

flyte/_internal/runtime/taskrunner.py

@@ -129,6 +129,14 @@ async def convert_and_run(
     in a context tree.
     """
     ctx = internal_ctx()
+
+    # Load inputs first to get context
+    if input_path:
+        inputs = await load_inputs(input_path, path_rewrite_config=raw_data_path.path_rewrite)
+
+    # Extract context from inputs
+    custom_context = inputs.context if inputs else {}
+
     tctx = TaskContext(
         action=action,
         checkpoints=checkpoints,
@@ -142,9 +150,10 @@ async def convert_and_run(
         report=flyte.report.Report(name=action.name),
         mode="remote" if not ctx.data.task_context else ctx.data.task_context.mode,
         interactive_mode=interactive_mode,
+        custom_context=custom_context,
     )
+
     with ctx.replace_task_context(tctx):
-        inputs = await load_inputs(input_path) if input_path else inputs
         inputs_kwargs = await convert_inputs_to_native(inputs, task.native_interface)
         out, err = await run_task(tctx=tctx, controller=controller, task=task, inputs=inputs_kwargs)
         if err is not None: