Skip to main content

Understanding Task Serialization

Task serialization in flyte-sdk is the process of converting Python task definitions into a language-neutral wire format (Protobuf) that the Flyte backend can understand and execute. This process ensures that the task's code, dependencies, and execution environment are correctly packaged and described for remote execution.

The two primary components driving this process are the SerializationContext, which manages deployment-time metadata, and the CodeBundle, which handles the packaging of the source code itself.

The SerializationContext​

The SerializationContext class (defined in models.py) is a container for information required during the serialization of a task. It acts as a bridge between the deployment environment and the final serialized task template.

Key attributes include:

  • project, domain, org: Identifiers that define where the task belongs in the Flyte ecosystem.
  • version: The specific version string for the task.
  • code_bundle: A reference to the CodeBundle containing the task's source code.
  • image_cache: A mapping of image identifiers to their final URIs, ensuring the correct container images are used.
  • input_path and output_path: Templates (e.g., {{.input}}) that Flyte uses to inject data locations at runtime.

The context is typically instantiated during the deployment process in _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,
)

Packaging Code with CodeBundle​

The CodeBundle class (defined in models.py) represents the physical packaging of the task's source code. It supports two primary formats:

  1. TGZ: A compressed tarball of the source directory.
  2. PKL: A pickled representation of the task, often used in interactive or notebook-based deployments.

A CodeBundle requires either a tgz or pkl path to be valid. It also includes a computed_version, which is a hash of the code content. This hash is frequently used as the default version for the entire task deployment.

@dataclass(frozen=True, kw_only=True)
class CodeBundle:
    computed_version: str
    destination: str = "."
    tgz: str | None = None
    pkl: str | None = None
    downloaded_path: pathlib.Path | None = None

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

The Serialization Pipeline​

The serialization process is triggered during deployment (via apply in _deploy.py) and flows through translate_task_to_wire to get_proto_task in _internal/runtime/task_serde.py.

1. Identity Resolution​

get_proto_task uses the SerializationContext to populate the Identifier for the task, combining the project, domain, and version with the task's name:

task_id = identifier_pb2.Identifier(
    resource_type=identifier_pb2.ResourceType.TASK,
    project=serialize_context.project,
    domain=serialize_context.domain,
    org=serialize_context.org,
    name=task.name,
    version=serialize_context.version,
)

2. Container Configuration​

The _get_urun_container function resolves the container image URI using the image_cache in the context. If the image is not found in the cache, flyte-sdk logs a warning, which often indicates a version mismatch between the deployment and execution environments.

3. Argument Generation​

The container_args method in AsyncFunctionTaskTemplate (in _task.py) generates the command-line arguments that the Flyte container will use at runtime. This includes the paths to the code bundle and the version:

if serialize_context.code_bundle:
    if serialize_context.code_bundle.tgz:
        args = [*args, *["--tgz", f"{serialize_context.code_bundle.tgz}"]]
    elif serialize_context.code_bundle.pkl:
        args = [*args, *["--pkl", f"{serialize_context.code_bundle.pkl}"]]
    args = [*args, *["--dest", f"{serialize_context.code_bundle.destination or '.'}"]]

Runtime Deserialization​

When the task executes on a remote cluster, the process is reversed. The extract_code_bundle function in task_serde.py parses the container's command-line arguments to reconstruct the CodeBundle object. This allows the runtime to locate, download, and inflate the source code before executing the task logic.

def extract_code_bundle(task_spec: task_definition_pb2.TaskSpec) -> Optional[CodeBundle]:
    container = task_spec.task_template.container
    if container and container.args:
        # ... logic to parse --pkl, --tgz, --dest, and --version ...
        if pkl_path or tgz_path:
            return CodeBundle(
                destination=dest_path,
                tgz=tgz_path,
                pkl=pkl_path,
                computed_version=version,
            )

Implementation Constraints​

  • Versioning: While the CodeBundle version is based on a code hash, the overall deployment version in _deploy.py does not currently account for changes in the Docker image or other task metadata.
  • Copy Styles: If copy_style is set to "none" during deployment, a manual version must be provided to the apply function, or a DeploymentError will be raised.
  • Interpreter Path: The SerializationContext defaults the interpreter_path to /opt/venv/bin/python, assuming a standard virtual environment within the execution container.