flyte 0.0.1b3__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/_internal/controllers/remote/_informer.py

@@ -38,11 +38,11 @@ class ActionCache:
         """
         Add an action to the cache if it doesn't exist. This is invoked by the watch.
         """
-        logger.info(f"Observing phase {run_definition_pb2.Phase.Name(state.phase)} for {state.action_id.name}")
+        logger.debug(f"Observing phase {run_definition_pb2.Phase.Name(state.phase)} for {state.action_id.name}")
         if state.output_uri:
-            logger.info(f"Output URI: {state.output_uri}")
+            logger.debug(f"Output URI: {state.output_uri}")
         else:
-            logger.info(f"{state.action_id.name} has no output URI")
+            logger.warning(f"{state.action_id.name} has no output URI")
         if state.phase == run_definition_pb2.Phase.PHASE_FAILED:
             logger.error(
                 f"Action {state.action_id.name} failed with error (msg):"
@@ -235,7 +235,7 @@ class Informer:
                     await self._shared_queue.put(node)
                     # hack to work in the absence of sentinel
             except asyncio.CancelledError:
-                logger.warning(f"Watch cancelled: {self.name}")
+                logger.info(f"Watch cancelled: {self.name}")
                 return
             except asyncio.TimeoutError as e:
                 logger.error(f"Watch timeout: {self.name}", exc_info=e)

flyte/_internal/controllers/remote/_service_protocol.py

@@ -28,12 +28,12 @@ class QueueService(Protocol):
     ) -> queue_service_pb2.EnqueueActionResponse:
         """Enqueue a task"""
 
-    async def AbortQueuedAction(
-        self,
-        req: queue_service_pb2.AbortQueuedActionRequest,
-        **kwargs,
-    ) -> queue_service_pb2.AbortQueuedActionResponse:
-        """Dequeue a task"""
+    # async def AbortQueuedAction(
+    #     self,
+    #     req: queue_service_pb2.AbortQueuedActionRequest,
+    #     **kwargs,
+    # ) -> queue_service_pb2.AbortQueuedActionResponse:
+    #     """Dequeue a task"""
 
 
 class ClientSet(Protocol):

flyte/_internal/imagebuild/docker_builder.py

@@ -1,4 +1,5 @@
 import asyncio
+import os
 import shutil
 import subprocess
 import tempfile
@@ -25,6 +26,8 @@ from flyte._image import (
 from flyte._logging import logger
 
 _F_IMG_ID = "_F_IMG_ID"
+FLYTE_DOCKER_BUILDER_CACHE_FROM = "FLYTE_DOCKER_BUILDER_CACHE_FROM"
+FLYTE_DOCKER_BUILDER_CACHE_TO = "FLYTE_DOCKER_BUILDER_CACHE_TO"
 
 UV_LOCK_INSTALL_TEMPLATE = Template("""\
 WORKDIR /root
@@ -193,7 +196,7 @@ class CopyConfigHandler:
             shutil.copy(abs_path, dest_path)
         elif layer.context_source.is_dir():
             # Copy the entire directory
-            shutil.copytree(abs_path, dest_path)
+            shutil.copytree(abs_path, dest_path, dirs_exist_ok=True)
         else:
             raise ValueError(f"Source path is neither file nor directory: {layer.context_source}")
 
@@ -396,6 +399,14 @@ class DockerImageBuilder:
                 "--push" if push else "--load",
             ]
 
+            cache_from = os.getenv(FLYTE_DOCKER_BUILDER_CACHE_FROM)
+            cache_to = os.getenv(FLYTE_DOCKER_BUILDER_CACHE_TO)
+            if cache_from and cache_to:
+                command[3:3] = [
+                    f"--cache-from={cache_from}",
+                    f"--cache-to={cache_to}",
+                ]
+
             if image.registry and push:
                 command.append("--push")
             command.append(tmp_dir)

flyte/_internal/imagebuild/image_builder.py

@@ -32,13 +32,15 @@ class DockerAPIImageChecker(ImageChecker):
     """
 
     @classmethod
-    async def image_exists(cls, repository: str, tag: str, arch: Tuple[Architecture, ...] = ("linux/amd64",)) -> bool:
+    async def image_exists(
+        cls,
+        repository: str,
+        tag: str,
+        arch: Tuple[Architecture, ...] = ("linux/amd64",)
+    ) -> bool:
         import httpx
 
-        if "/" in repository:
-            if not repository.startswith("library/"):
-                raise ValueError("This checker only works with Docker Hub")
-        else:
+        if "/" not in repository:
             repository = f"library/{repository}"
 
         auth_url = "https://auth.docker.io/token"
@@ -50,6 +52,7 @@ class DockerAPIImageChecker(ImageChecker):
             auth_response = await client.get(auth_url, params={"service": service, "scope": scope})
             if auth_response.status_code != 200:
                 raise Exception(f"Failed to get auth token: {auth_response.status_code}")
+
             token = auth_response.json()["token"]
 
             manifest_url = f"https://registry-1.docker.io/v2/{repository}/manifests/{tag}"
@@ -60,18 +63,20 @@ class DockerAPIImageChecker(ImageChecker):
                     "application/vnd.docker.distribution.manifest.list.v2+json"
                 ),
             }
-            manifest_response = await client.get(manifest_url, headers=headers)
 
+            manifest_response = await client.get(manifest_url, headers=headers)
             if manifest_response.status_code != 200:
-                raise Exception(f"Failed to get manifest: {manifest_response.status_code}")
+                logger.warning(f"Image not found: {repository}:{tag} (HTTP {manifest_response.status_code})")
+                return False
+
             manifest_list = manifest_response.json()["manifests"]
-            architectures = [f"{x['platform']['os']}/{x['platform']['architecture']}" for x in manifest_list]
+            architectures = [f"{m['platform']['os']}/{m['platform']['architecture']}" for m in manifest_list]
 
-            if set(architectures) >= set(arch):
-                logger.debug(f"Image {repository}:{tag} found for architecture(s) {arch}, has {architectures}")
+            if set(arch).issubset(set(architectures)):
+                logger.debug(f"Image {repository}:{tag} found with arch {architectures}")
                 return True
             else:
-                logger.debug(f"Image {repository}:{tag} not found for architecture(s) {arch}, only has {architectures}")
+                logger.debug(f"Image {repository}:{tag} has {architectures}, but missing {arch}")
                 return False
 
 

flyte/_internal/runtime/convert.py

@@ -1,16 +1,20 @@
 from __future__ import annotations
 
+import asyncio
+import base64
+import hashlib
+import inspect
 from dataclasses import dataclass
-from typing import Any, Dict, Tuple, Union
+from types import NoneType
+from typing import Any, Dict, List, Tuple, Union, get_args
 
-from flyteidl.core import execution_pb2, literals_pb2
+from flyteidl.core import execution_pb2, interface_pb2, literals_pb2
 
 import flyte.errors
 import flyte.storage as storage
-from flyte._datastructures import ActionID, NativeInterface, TaskContext
-from flyte._internal.controllers import pbhash
-from flyte._protos.workflow import run_definition_pb2
-from flyte.types import TypeEngine
+from flyte._protos.workflow import common_pb2, run_definition_pb2, task_definition_pb2
+from flyte.models import ActionID, NativeInterface, TaskContext
+from flyte.types import TypeEngine, TypeTransformerFailedError
 
 
 @dataclass(frozen=True)
@@ -56,24 +60,97 @@ async def convert_inputs_to_native(inputs: Inputs, python_interface: NativeInter
     return native_vals
 
 
+async def convert_upload_default_inputs(interface: NativeInterface) -> List[common_pb2.NamedParameter]:
+    """
+    Converts the default inputs of a NativeInterface to a list of NamedParameters for upload.
+    This is used to upload default inputs to the Flyte backend.
+    """
+    if not interface.inputs:
+        return []
+
+    vars = []
+    literal_coros = []
+    for input_name, (input_type, default_value) in interface.inputs.items():
+        if default_value is not inspect.Parameter.empty:
+            lt = TypeEngine.to_literal_type(input_type)
+            literal_coros.append(TypeEngine.to_literal(default_value, input_type, lt))
+            vars.append((input_name, lt))
+
+    literals: List[literals_pb2.Literal] = await asyncio.gather(*literal_coros)
+    named_params = []
+    for (name, lt), literal in zip(vars, literals):
+        param = interface_pb2.Parameter(
+            var=interface_pb2.Variable(
+                type=lt,
+            ),
+            default=literal,
+        )
+        named_params.append(
+            common_pb2.NamedParameter(
+                name=name,
+                parameter=param,
+            ),
+        )
+    return named_params
+
+
+def is_optional_type(tp) -> bool:
+    """
+    True if the *annotation* `tp` is equivalent to Optional[…].
+    Works for Optional[T], Union[T, None], and T | None.
+    """
+    return NoneType in get_args(tp)  # fastest check
+
+
 async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwargs) -> Inputs:
     kwargs = interface.convert_to_kwargs(*args, **kwargs)
-    if len(kwargs) == 0:
-        return Inputs.empty()
-    if len(kwargs) < len(interface.inputs):
+
+    if len(kwargs) < interface.num_required_inputs():
         raise ValueError(
-            f"Received {len(kwargs)} inputs but interface has {len(interface.inputs)}. "
-            f"Please provide all required inputs."
+            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}"
         )
-    literal_map = await TypeEngine.dict_to_literal_map(kwargs, interface.get_input_types())
+
+    if len(interface.inputs) == 0:
+        return Inputs.empty()
+
+    # fill in defaults if missing
+    type_hints: Dict[str, type] = {}
+    already_converted_kwargs: Dict[str, literals_pb2.Literal] = {}
+    for input_name, (input_type, default_value) in interface.inputs.items():
+        if input_name in kwargs:
+            type_hints[input_name] = input_type
+        elif (default_value is not None and default_value is not inspect.Signature.empty) or (
+            default_value is None and is_optional_type(input_type)
+        ):
+            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]
+            else:
+                kwargs[input_name] = default_value
+                type_hints[input_name] = input_type
+
+    literal_map = await TypeEngine.dict_to_literal_map(kwargs, type_hints)
+    if len(already_converted_kwargs) > 0:
+        copied_literals: Dict[str, literals_pb2.Literal] = {}
+        for k, v in literal_map.literals.items():
+            copied_literals[k] = v
+        # Add the already converted kwargs to the literal map
+        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=v) for k, v in literal_map.literals.items()]
+            literals=[
+                run_definition_pb2.NamedLiteral(name=k, value=literal_map.literals[k]) for k in interface.inputs.keys()
+            ]
         )
     )
 
 
-async def convert_from_native_to_outputs(o: Any, interface: NativeInterface) -> Outputs:
+async def convert_from_native_to_outputs(o: Any, interface: NativeInterface, task_name: str = "") -> Outputs:
     # Always make it a tuple even if it's just one item to simplify logic below
     if not isinstance(o, tuple):
         o = (o,)
@@ -83,8 +160,11 @@ async def convert_from_native_to_outputs(o: Any, interface: NativeInterface) ->
     )
     named = []
     for (output_name, python_type), v in zip(interface.outputs.items(), o):
-        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))
+        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))
+        except TypeTransformerFailedError as e:
+            raise flyte.errors.RuntimeDataValidationError(output_name, e, task_name)
 
     return Outputs(proto_outputs=run_definition_pb2.Outputs(literals=named))
 
@@ -179,21 +259,84 @@ def convert_from_native_to_error(err: BaseException) -> Error:
         )
 
 
-def generate_sub_action_id_and_output_path(tctx: TaskContext, task_name: str, inputs: Inputs) -> Tuple[ActionID, str]:
+def hash_data(data: Union[str, bytes]) -> str:
+    """
+    Generate a hash for the given data. If the data is a string, it will be encoded to bytes before hashing.
+    :param data: The data to hash, can be a string or bytes.
+    :return: A hexadecimal string representation of the hash.
+    """
+    if isinstance(data, str):
+        data = data.encode("utf-8")
+    digest = hashlib.sha256(data).digest()
+    return base64.b64encode(digest).decode("utf-8")
+
+
+def generate_inputs_hash(serialized_inputs: str | bytes) -> str:
+    """
+    Generate a hash for the inputs. This is used to uniquely identify the inputs for a task.
+    :return: A hexadecimal string representation of the hash.
+    """
+    return hash_data(serialized_inputs)
+
+
+def generate_cache_key_hash(
+    task_name: str,
+    inputs_hash: str,
+    task_interface: interface_pb2.TypedInterface,
+    cache_version: str,
+    ignored_input_vars: List[str],
+    proto_inputs: run_definition_pb2.Inputs,
+) -> str:
+    """
+    Generate a cache key hash based on the inputs hash, task name, task interface, and cache version.
+    This is used to uniquely identify the cache key for a task.
+
+    :param task_name: The name of the task.
+    :param inputs_hash: The hash of the inputs.
+    :param task_interface: The interface of the task.
+    :param cache_version: The version of the cache.
+    :param ignored_input_vars: A list of input variable names to ignore when generating the cache key.
+    :param proto_inputs: The proto inputs for the task, only used if there are ignored inputs.
+    :return: A hexadecimal string representation of the 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_inputs = final.SerializeToString(deterministic=True)
+    else:
+        final_inputs = inputs_hash
+    data = f"{final_inputs}{task_name}{task_interface.SerializeToString(deterministic=True)}{cache_version}"
+    return hash_data(data)
+
+
+def generate_sub_action_id_and_output_path(
+    tctx: TaskContext,
+    task_spec_or_name: task_definition_pb2.TaskSpec | str,
+    inputs_hash: str,
+    invoke_seq: int,
+) -> Tuple[ActionID, str]:
     """
     Generate a sub-action ID and output path based on the current task context, task name, and inputs.
+
+    action name = current action name + task name + input hash + group name (if available)
     :param tctx:
-    :param task_name:
-    :param inputs:
+    :param task_spec_or_name: task specification or task name. Task name is only used in case of trace actions.
+    :param inputs_hash: Consistent hash string of the inputs
+    :param invoke_seq: The sequence number of the invocation, used to differentiate between multiple invocations.
     :return:
     """
     current_action_id = tctx.action
     current_output_path = tctx.run_base_dir
-    inputs_hash = pbhash.compute_hash_string(inputs.proto_inputs)
+    if isinstance(task_spec_or_name, task_definition_pb2.TaskSpec):
+        task_spec_or_name.task_template.interface
+        task_hash = hash_data(task_spec_or_name.SerializeToString(deterministic=True))
+    else:
+        task_hash = task_spec_or_name
     sub_action_id = current_action_id.new_sub_action_from(
-        task_name=task_name,
+        task_hash=task_hash,
         input_hash=inputs_hash,
         group=tctx.group_data.name if tctx.group_data else None,
+        task_call_seq=invoke_seq,
     )
     sub_run_output_path = storage.join(current_output_path, sub_action_id.name)
     return sub_action_id, sub_run_output_path

flyte/_internal/runtime/entrypoints.py

@@ -3,11 +3,11 @@ from typing import List, Optional, Tuple
 import flyte.errors
 from flyte._code_bundle import download_bundle
 from flyte._context import contextual_run
-from flyte._datastructures import ActionID, Checkpoints, CodeBundle, RawDataPath
 from flyte._internal import Controller
 from flyte._internal.imagebuild.image_builder import ImageCache
 from flyte._logging import log, logger
 from flyte._task import TaskTemplate
+from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
 
 from .convert import Error, Inputs, Outputs
 from .task_serde import load_task

flyte/_internal/runtime/io.py

@@ -21,11 +21,11 @@ _OUTPUTS_FILE_NAME = "outputs.pb"
 _CHECKPOINT_FILE_NAME = "_flytecheckpoints"
 _ERROR_FILE_NAME = "error.pb"
 _REPORT_FILE_NAME = "report.html"
-_PKL_FILE_NAME = "code_bundle.pkl.gz"
+_PKL_EXT = ".pkl.gz"
 
 
-def pkl_path(base_path: str) -> str:
-    return storage.join(base_path, _PKL_FILE_NAME)
+def pkl_path(base_path: str, pkl_name: str) -> str:
+    return storage.join(base_path, f"{pkl_name}{_PKL_EXT}")
 
 
 def inputs_path(base_path: str) -> str:

flyte/_internal/runtime/task_serde.py

@@ -3,26 +3,32 @@ This module provides functionality to serialize and deserialize tasks to and fro
 It includes a Resolver interface for loading tasks, and functions to load classes and tasks.
 """
 
+import copy
 import importlib
+import typing
 from datetime import timedelta
-from typing import Optional, Type
+from typing import Optional, Type, cast
 
 from flyteidl.core import identifier_pb2, literals_pb2, security_pb2, tasks_pb2
 from google.protobuf import duration_pb2, wrappers_pb2
 
 import flyte.errors
 from flyte._cache.cache import VersionParameters, cache_from_request
-from flyte._datastructures import SerializationContext
 from flyte._logging import logger
-from flyte._protos.workflow import task_definition_pb2
+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
 
 from ..._retry import RetryStrategy
 from ..._timeout import TimeoutType, timeout_from_request
 from .resources_serde import get_proto_extended_resources, get_proto_resources
 from .types_serde import transform_native_to_typed_interface
 
+_MAX_ENV_NAME_LENGTH = 63  # Maximum length for environment names
+_MAX_TASK_SHORT_NAME_LENGTH = 63  # Maximum length for task short names
+
 
 def load_class(qualified_name) -> Type:
     """
@@ -49,17 +55,31 @@ def load_task(resolver: str, *resolver_args: str) -> TaskTemplate:
 
 
 def translate_task_to_wire(
-    task: TaskTemplate, serialization_context: SerializationContext
+    task: TaskTemplate,
+    serialization_context: SerializationContext,
+    default_inputs: Optional[typing.List[common_pb2.NamedParameter]] = None,
 ) -> task_definition_pb2.TaskSpec:
     """
     Translate a task to a wire format. This is a placeholder function.
 
     :param task: The task to translate.
     :param serialization_context: The serialization context to use for the translation.
+    :param default_inputs: Optional list of default inputs for the task.
 
     :return: The translated task.
     """
-    return get_proto_task(task, serialization_context)
+    tt = get_proto_task(task, serialization_context)
+    env: environment_pb2.Environment | None = None
+    if task.parent_env and task.parent_env():
+        _env = task.parent_env()
+        if _env:
+            env = environment_pb2.Environment(name=_env.name[:_MAX_ENV_NAME_LENGTH])
+    return task_definition_pb2.TaskSpec(
+        task_template=tt,
+        default_inputs=default_inputs,
+        short_name=task.friendly_name[:_MAX_TASK_SHORT_NAME_LENGTH],
+        environment=env,
+    )
 
 
 def get_security_context(secrets: Optional[SecretRequest]) -> Optional[security_pb2.SecurityContext]:
@@ -108,7 +128,7 @@ def get_proto_timeout(timeout: TimeoutType | None) -> Optional[duration_pb2.Dura
     return duration_pb2.Duration(seconds=max_runtime_timeout.seconds)
 
 
-def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext) -> task_definition_pb2.TaskSpec:
+def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext) -> tasks_pb2.TaskTemplate:
     task_id = identifier_pb2.Identifier(
         resource_type=identifier_pb2.ResourceType.TASK,
         project=serialize_context.project,
@@ -121,17 +141,18 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
     # if task.parent_env is None:
     # raise ValueError(f"Task {task.name} must have a parent environment")
 
-    #
-    # This pod will be incorrect when doing fast serialize
-    #
-    container = _get_urun_container(serialize_context, task)
+    # TODO Add support for SQL, extra_config, custom
+    extra_config: typing.Dict[str, str] = {}
+    custom = {}  # type: ignore
 
-    # TODO Add support for SQL, Pod, extra_config, custom
-    pod = None
     sql = None
-    # pod = task.get_k8s_pod(serialize_context)
-    extra_config = {}  # type: ignore
-    custom = {}  # type: ignore
+    if task.pod_template and not isinstance(task.pod_template, str):
+        container = None
+        pod = _get_k8s_pod(_get_urun_container(serialize_context, task), task.pod_template)
+        extra_config[_PRIMARY_CONTAINER_NAME_FIELD] = task.pod_template.primary_container_name
+    else:
+        container = _get_urun_container(serialize_context, task)
+        pod = None
 
     # -------------- CACHE HANDLING ----------------------
     task_cache = cache_from_request(task.cache)
@@ -154,20 +175,24 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
     else:
         logger.debug(f"Cache disabled for task {task.name}")
 
-    tt = tasks_pb2.TaskTemplate(
+    return tasks_pb2.TaskTemplate(
         id=task_id,
         type=task.task_type,
         metadata=tasks_pb2.TaskMetadata(
             discoverable=cache_enabled,
             discovery_version=cache_version,
             cache_serializable=task_cache.serialize,
-            cache_ignore_input_vars=task_cache.ignored_inputs,
-            runtime=tasks_pb2.RuntimeMetadata(),
+            cache_ignore_input_vars=task_cache.get_ignored_inputs() if cache_enabled else None,
+            runtime=tasks_pb2.RuntimeMetadata(
+                version=flyte.version(),
+                type=tasks_pb2.RuntimeMetadata.RuntimeType.FLYTE_SDK,
+                flavor="python",
+            ),
             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,
-            generates_deck=wrappers_pb2.BoolValue(value=False),  # TODO add support for reports
+            generates_deck=wrappers_pb2.BoolValue(value=task.report),
         ),
         interface=transform_native_to_typed_interface(task.native_interface),
         custom=custom,
@@ -179,7 +204,6 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
         sql=sql,
         extended_resources=get_proto_extended_resources(task.resources),
     )
-    return task_definition_pb2.TaskSpec(task_template=tt)
 
 
 def _get_urun_container(
@@ -208,3 +232,99 @@ def _get_urun_container(
         data_config=task_template.data_loading_config(serialize_context),
         config=task_template.config(serialize_context),
     )
+
+
+def _sanitize_resource_name(resource: tasks_pb2.Resources.ResourceEntry) -> str:
+    return tasks_pb2.Resources.ResourceName.Name(resource.name).lower().replace("_", "-")
+
+
+def _get_k8s_pod(primary_container: tasks_pb2.Container, pod_template: PodTemplate) -> Optional[tasks_pb2.K8sPod]:
+    """
+    Get the K8sPod representation of the task template.
+    :param task: The task to convert.
+    :return: The K8sPod representation of the task template.
+    """
+    from kubernetes.client import ApiClient, V1PodSpec
+    from kubernetes.client.models import V1EnvVar, V1ResourceRequirements
+
+    pod_template = copy.deepcopy(pod_template)
+    containers = cast(V1PodSpec, pod_template.pod_spec).containers
+    primary_exists = False
+
+    for container in containers:
+        if container.name == pod_template.primary_container_name:
+            primary_exists = True
+            break
+
+    if not primary_exists:
+        raise ValueError(
+            "No primary container defined in the pod spec."
+            f" You must define a primary container with the name '{pod_template.primary_container_name}'."
+        )
+    final_containers = []
+
+    for container in containers:
+        # We overwrite the primary container attributes with the values given to ContainerTask.
+        # The attributes include: image, command, args, resource, and env (env is unioned)
+
+        if container.name == pod_template.primary_container_name:
+            if container.image is None:
+                # Copy the image from primary_container only if the image is not specified in the pod spec.
+                container.image = primary_container.image
+
+            container.command = list(primary_container.command)
+            container.args = list(primary_container.args)
+
+            limits, requests = {}, {}
+            for resource in primary_container.resources.limits:
+                limits[_sanitize_resource_name(resource)] = resource.value
+            for resource in primary_container.resources.requests:
+                requests[_sanitize_resource_name(resource)] = resource.value
+
+            resource_requirements = V1ResourceRequirements(limits=limits, requests=requests)
+            if len(limits) > 0 or len(requests) > 0:
+                # Important! Only copy over resource requirements if they are non-empty.
+                container.resources = resource_requirements
+
+            if primary_container.env is not None:
+                container.env = [V1EnvVar(name=e.key, value=e.value) for e in primary_container.env] + (
+                    container.env or []
+                )
+
+        final_containers.append(container)
+
+    cast(V1PodSpec, pod_template.pod_spec).containers = final_containers
+    pod_spec = ApiClient().sanitize_for_serialization(pod_template.pod_spec)
+
+    metadata = tasks_pb2.K8sObjectMetadata(labels=pod_template.labels, annotations=pod_template.annotations)
+    return tasks_pb2.K8sPod(pod_spec=pod_spec, metadata=metadata)
+
+
+def extract_code_bundle(task_spec: task_definition_pb2.TaskSpec) -> Optional[CodeBundle]:
+    """
+    Extract the code bundle from the task spec.
+    :param task_spec: The task spec to extract the code bundle from.
+    :return: The extracted code bundle or None if not present.
+    """
+    container = task_spec.task_template.container
+    if container and container.args:
+        pkl_path = None
+        tgz_path = None
+        dest_path: str = "."
+        version = ""
+        for i, v in enumerate(container.args):
+            if v == "--pkl":
+                # Extract the code bundle path from the argument
+                pkl_path = container.args[i + 1] if i + 1 < len(container.args) else None
+            elif v == "--tgz":
+                # Extract the code bundle path from the argument
+                tgz_path = container.args[i + 1] if i + 1 < len(container.args) else None
+            elif v == "--dest":
+                # Extract the destination path from the argument
+                dest_path = container.args[i + 1] if i + 1 < len(container.args) else "."
+            elif v == "--version":
+                # Extract the version from the argument
+                version = container.args[i + 1] if i + 1 < len(container.args) else ""
+        if pkl_path or tgz_path:
+            return CodeBundle(destination=dest_path, tgz=tgz_path, pkl=pkl_path, computed_version=version)
+    return None