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

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)
@@ -308,15 +330,82 @@ def generate_inputs_hash(serialized_inputs: str | bytes) -> str:
     return hash_data(serialized_inputs)
 
 
-def generate_inputs_hash_from_proto(inputs: run_definition_pb2.Inputs) -> str:
+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[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
+    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: 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.
     :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:
@@ -337,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.
@@ -353,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)
@@ -142,6 +165,7 @@ async def load_and_run_task(
     code_bundle: CodeBundle | None = None,
     input_path: str | None = None,
     image_cache: ImageCache | None = None,
+    interactive_mode: bool = False,
 ):
     """
     This method is invoked from the runtime/CLI and is used to run a task. This creates the context tree,
@@ -159,6 +183,7 @@ async def load_and_run_task(
     :param code_bundle: The code bundle to use for the task.
     :param input_path: The input path to use for the task.
     :param image_cache: Mappings of Image identifiers to image URIs.
+    :param interactive_mode: Whether to run the task in interactive mode.
     """
     task = await _download_and_load_task(code_bundle, resolver, resolver_args)
 
@@ -175,4 +200,5 @@ async def load_and_run_task(
         code_bundle=code_bundle,
         input_path=input_path,
         image_cache=image_cache,
+        interactive_mode=interactive_mode,
     )

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

@@ -8,14 +8,14 @@ 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
@@ -54,7 +54,7 @@ def translate_task_to_wire(
     return task_definition_pb2.TaskSpec(
         task_template=tt,
         default_inputs=default_inputs,
-        short_name=task.friendly_name[:_MAX_TASK_SHORT_NAME_LENGTH],
+        short_name=task.short_name[:_MAX_TASK_SHORT_NAME_LENGTH],
         environment=env,
     )
 
@@ -119,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):
@@ -132,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)
@@ -145,7 +145,6 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
             logger.debug(f"Detected pkl bundle for task {task.name}, using computed version as cache version")
             cache_version = serialize_context.code_bundle.computed_version
         else:
-            version_parameters = None
             if isinstance(task, AsyncFunctionTaskTemplate):
                 version_parameters = VersionParameters(func=task.func, image=task.image)
             else:
@@ -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,25 +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
-        from flyte._version import __version__
-
-        logger.warning(
-            f"Image {task_template.image} not found in the image cache: {serialize_context.image_cache.image_lookup}.\n"
-            f"This typically occurs when the Flyte SDK version (`{__version__}`) used in the task environment "
-            f"differs from the version used to compile or deploy it.\n"
-            f"Ensure both environments use the same Flyte SDK version to avoid inconsistencies in image resolution."
+    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:
-        img_uri = serialize_context.image_cache.image_lookup[image_id]
+        img_uri = serialize_context.image_cache.image_lookup[env_name]
 
     return tasks_pb2.Container(
         image=img_uri,