Skip to main content

The Serialization Process

When a task in flyte-sdk is prepared for remote execution, it undergoes a serialization process that packages the task along with its necessary contextual information. This process relies on two core classes: SerializationContext and CodeBundle. These classes ensure that tasks are consistently versioned, correctly located, and properly configured for execution in a distributed environment.

SerializationContext​

The SerializationContext class serves as a central repository for all contextual information required during the serialization of a task. It is a dataclass that encapsulates metadata such as versioning details, project and domain information, and paths for input and output data. This context is crucial for providing a consistent environment, whether the task is being deployed or executed at runtime.

Key attributes of SerializationContext include:

  • version (str): The specific version identifier for the task.
  • project (str | None): The project to which the task belongs.
  • domain (str | None): The domain within the project.
  • org (str | None): The organization associated with the task.
  • code_bundle (Optional[CodeBundle]): An optional CodeBundle instance, which packages the task's code.
  • input_path (str): A templated path, defaulting to {{.input}}, indicating where task inputs will be located.
  • output_path (str): A templated path, defaulting to {{.outputPrefix}}, indicating where task outputs should be stored.
  • interpreter_path (str): The path to the Python interpreter, defaulting to /opt/venv/bin/python.
  • image_cache (ImageCache | None): An optional cache for Docker images.
  • root_dir (Optional[pathlib.Path]): The root directory for the task.

The input_path and output_path attributes utilize templating, allowing the Flyte engine to dynamically resolve these paths at runtime based on the execution environment. The SerializationContext also provides a method to determine the task's entrypoint:

import os
from typing import Optional

class SerializationContext:
    # ... (other attributes and methods)

    def get_entrypoint_path(self, interpreter_path: Optional[str] = None) -> str:
        """
        Get the entrypoint path for the task. This is used to determine the entrypoint for the task execution.
        :param interpreter_path: The path to the interpreter (python)
        """
        if interpreter_path is None:
            interpreter_path = self.interpreter_path
        return os.path.join(os.path.dirname(interpreter_path), "runtime.py")

SerializationContext instances are created at various stages of the task lifecycle. For example, during deployment, it gathers comprehensive project and versioning details:

# From src/flyte/_deploy.py
sc = SerializationContext(
    project=cfg.project,
    domain=cfg.domain,
    org=cfg.org,
    code_bundle=code_bundle,
    version=version,
    image_cache=image_cache,
    root_dir=cfg.root_dir,
)

When running a task locally, a SerializationContext is created with relevant local execution details:

# From src/flyte/_run.py
s_ctx = SerializationContext(
    code_bundle=code_bundle,
    version=version,
    image_cache=image_cache,
    root_dir=cfg.root_dir,
)

For remote task execution, the context is tailored to the action ID and existing runtime information:

# From src/flyte/_internal/controllers/remote/_controller.py
new_serialization_context = SerializationContext(
    project=current_action_id.project,
    domain=current_action_id.domain,
    org=current_action_id.org,
    code_bundle=code_bundle,
    version=tctx.version,
    image_cache=tctx.compiled_image_cache,
)

It is important to note that SerializationContext is only available during the serialization phase, which occurs either during deployment or at runtime when a task is being prepared for execution.

CodeBundle​

The CodeBundle class is a frozen dataclass that represents the packaged code for a task. It is a critical component of the SerializationContext, providing the means to transfer and inflate task code in the execution environment. The immutability of CodeBundle (frozen=True) ensures that its contents remain consistent once created.

Key attributes of CodeBundle include:

  • computed_version (str): A version identifier for the code bundle, typically derived from a hash of the code itself.
  • destination (str): The target path where the code bundle will be inflated (unpacked), defaulting to ..
  • tgz (str | None): An optional path to a .tgz archive containing the code.
  • pkl (str | None): An optional path to a .pkl (pickle) file, which might contain serialized Python objects.
  • downloaded_path (pathlib.Path | None): The path where the code bundle has been downloaded and inflated at runtime. This attribute is only populated during runtime after the bundle has been processed.

CodeBundle enforces that at least one of tgz or pkl must be provided upon instantiation, ensuring that there is always a source for the code. This validation is performed in its __post_init__ method:

import rich.repr
from dataclasses import dataclass, field, replace
import pathlib

@rich.repr.auto
@dataclass(frozen=True, kw_only=True)
class CodeBundle:
    # ... (other attributes)

    def __post_init__(self):
        if self.tgz is None and self.pkl is None:
            raise ValueError("Either tgz or pkl must be provided")

    def with_downloaded_path(self, path: pathlib.Path) -> CodeBundle:
        """
        Create a new CodeBundle with the given downloaded path.
        """
        return replace(self, downloaded_path=path)

At runtime, particularly within the Flyte agent, a CodeBundle is constructed from command-line arguments to prepare the execution environment:

# From src/flyte/_bin/runtime.py
bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)

The with_downloaded_path method allows for the creation of a new CodeBundle instance with the downloaded_path attribute set, reflecting the state of the bundle after it has been retrieved and unpacked in the execution environment.

Interplay and Usage​

The SerializationContext and CodeBundle work in tandem to facilitate the serialization of tasks. The SerializationContext often contains a CodeBundle instance, providing the necessary code alongside other contextual metadata. This combined information is then used by internal Flyte components to translate tasks into their wire format for execution. For instance, the translate_task_to_wire function in _internal/runtime/task_serde.py explicitly accepts a SerializationContext:

# From src/flyte/_internal/runtime/task_serde.py
import typing
from typing import Optional

# Assuming TaskTemplate, common_pb2.NamedParameter, and task_definition_pb2.TaskSpec are defined elsewhere
def translate_task_to_wire(
    task: "TaskTemplate", # Use forward references for types not explicitly imported in this snippet
    serialization_context: "SerializationContext",
    default_inputs: Optional[typing.List["common_pb2.NamedParameter"]] = None,
) -> "task_definition_pb2.TaskSpec":
    pass

This demonstrates how the SerializationContext provides the comprehensive environment needed to transform a task definition into a format suitable for the Flyte platform.

Gotchas and Considerations​

  • Path Templating: Be aware that SerializationContext.input_path and SerializationContext.output_path use templated strings ({{.input}}, {{.outputPrefix}}) that are resolved by the Flyte engine at runtime.
  • CodeBundle Requirements: A CodeBundle must always be initialized with either a tgz or pkl path, but not necessarily both. This ensures that the code source is always specified.
  • Runtime Availability: The CodeBundle.downloaded_path attribute is only populated and available during runtime after the code bundle has been successfully downloaded and inflated into the execution environment.
  • Versioning: The version attribute in SerializationContext is explicitly required when deploying or running a task, highlighting its importance for tracking and reproducibility. Future enhancements, as indicated by TODOs in _deploy.py, aim to incorporate image cache digests and code bundle digests for more robust versioning.
  • Testing: It's worth noting that dedicated unit test files for models.py and related modules are not present within the repository, which might imply that testing for these core components is integrated into broader functional or integration tests.
  • Runtime Environment Variables: For successful runtime execution, certain environment variables are expected to be set, such as FLYTE_INTERNAL_TASK_PROJECT, FLYTE_INTERNAL_TASK_DOMAIN, _U_ORG_NAME, RUN_NAME, and ACTION_NAME. These variables can be used to resolve templated values like run_name and name in _bin/runtime.py.

Configuration​

Several environment variables influence the serialization and runtime behavior, providing configuration points for project details, run names, and internal settings:

  • ACTION_NAME: Specifies the name of the action.
  • RUN_NAME: Defines the name of the run.
  • FLYTE_INTERNAL_TASK_PROJECT: Sets the project for the internal task.
  • FLYTE_INTERNAL_TASK_DOMAIN: Sets the domain for the internal task.
  • _U_ORG_NAME: Specifies the organization name.
  • _U_EP_OVERRIDE: Allows overriding the endpoint.
  • _U_RUN_BASE: Defines the base directory for run outputs.
  • _F_E_VS: Enables VSCode debugging.
  • _UNION_EAGER_API_KEY: A temporary mechanism for API key authentication.