Skip to main content

Accessing Runtime Metadata with TaskContext

When writing task logic, tasks often need awareness of their execution environment—such as obtaining the active execution ID for tagging logs, determining a unique storage prefix for offloading intermediate artifacts, inspecting retry checkpoints, or generating dynamic HTML dashboards. Hardcoding paths or environment flags breaks portability across local tests and remote cluster runs. The flyte.ctx() function provides access to an active TaskContext instance containing execution metadata, storage paths, reporting handles, and environment state.

Accessing the Active TaskContext​

To access runtime metadata inside any task execution, call flyte.ctx(). If called during task execution, it returns the current TaskContext instance stored in Flyte's async context variables; if called outside an active task run (such as during top-level script import), it returns None.

import flyte

@flyte.task
def process_data(batch_id: int) -> str:
    ctx = flyte.ctx()
    if ctx is None:
        return f"Executing outside task context: {batch_id}"

    # Check environment mode and task metadata
    mode = ctx.mode
    is_cluster = ctx.is_in_cluster()
    action_name = ctx.action.name
    version = ctx.version

    return f"Action {action_name} (v{version}) running in mode '{mode}' (cluster={is_cluster})"

Context Attributes and Execution Mode​

The TaskContext dataclass in models.py holds core execution attributes:

  • action: An ActionID instance representing the task execution identifier.
  • version: A string containing the task version hash.
  • mode: A literal indicating the execution mode: "local", "remote", or "hybrid".
  • interactive_mode: A boolean indicating whether the task is running interactively.
  • output_path: The destination URI prefix for task outputs.
  • run_base_dir: The root storage directory for the entire run.
  • raw_data_path: A RawDataPath instance for generating non-colliding storage URIs.
  • checkpoints: A Checkpoints instance (or None) holding paths for state persistence.
  • code_bundle: A CodeBundle instance (or None) holding code package metadata.
  • report: A Report instance for writing HTML visualizations.
  • data: A dictionary for custom contextual key-value pairs, accessible via ctx["key"].

The helper method is_in_cluster() evaluates whether the execution is remote:

def is_in_cluster(self):
    return self.mode == "remote"

Because TaskContext is defined with @dataclass(frozen=True, kw_only=True), its fields are immutable. To produce a modified context, use ctx.replace(**kwargs). Calling replace(data=...) shallowly merges the new dictionary into the existing data dictionary:

# Create an updated context instance with merged data
new_ctx = ctx.replace(data={"custom_metric": 42})
metric = new_ctx["custom_metric"]

Action Identifiers (ActionID)​

The ctx.action attribute holds an ActionID dataclass identifying the task within a Flyte run hierarchy:

@dataclass(frozen=True, kw_only=True)
class ActionID:
    name: str
    run_name: str | None = None
    project: str | None = None
    domain: str | None = None
    org: str | None = None

When initialized without an explicit run_name, ActionID.__post_init__ defaults run_name to match name.

Inside tasks, logging formatters and observability tools use ctx.action to associate logs with specific execution runs:

import flyte

@flyte.task
def logged_step(item: str) -> None:
    ctx = flyte.ctx()
    if ctx:
        action = ctx.action
        print(f"[{action.project}/{action.domain}][{action.run_name}][{action.name}] Processing {item}")

Deterministic Sub-Action IDs​

When launching child actions, mapped tasks, or sub-runs, flyte-sdk generates deterministic action identifiers using ActionID.new_sub_action_from:

sub_action = ctx.action.new_sub_action_from(
    task_call_seq=0,
    task_hash="task-abc1234",
    input_hash="input-xyz5678",
    group="data-prep",
)

Internally, this method formats the components as "{self.name}-{input_hash}-{task_hash}-{task_call_seq}-{group}", computes an MD5 digest, and encodes the digest using base-36 to produce a unique sub-action name.

Raw Data Paths and Storage Artifacts (RawDataPath)​

Offloaded files, binary blobs, and partitions require storage paths that do not collide across concurrent actions. The ctx.raw_data_path attribute holds a RawDataPath instance providing the root prefix and URI generation utilities.

import flyte

@flyte.task
def export_artifacts(data: bytes) -> str:
    ctx = flyte.ctx()
    if ctx is None:
        raise RuntimeError("No active TaskContext")

    # Generate a unique path with a 128-bit random UUID prefix
    target_path = ctx.raw_data_path.get_random_remote_path("export.parquet")
    
    # Store or offload the file to target_path
    return target_path

RawDataPath.get_random_remote_path(file_name=None) handles both local filesystem paths and remote object store URIs (such as s3://, gs://, or file://):

  1. It generates a 128-bit random hex UUID via UUID(int=random.getrandbits(128)).hex.
  2. For local paths (file protocol), it constructs the child directory, creates parent directories if a file_name is given, and touches the file.
  3. For remote storage, it uses fsspec.utils.get_protocol and fsspec.filesystem(protocol).sep to join the base prefix, the random UUID directory, and the optional file_name.

For local development or custom folder structures, RawDataPath.from_local_folder(local_folder) constructs a path from a string or pathlib.Path, creating temporary directories via tempfile.mkdtemp() if no path is supplied.

Checkpoint Management (Checkpoints)​

For fault-tolerant or iterative workloads, ctx.checkpoints provides the Checkpoints dataclass:

@dataclass(frozen=True)
class Checkpoints:
    prev_checkpoint_path: str | None
    checkpoint_path: str | None
  • prev_checkpoint_path: Points to the storage URI containing the checkpoint from a prior retry or execution attempt.
  • checkpoint_path: Points to the writable storage location for the current execution attempt.
import os
import flyte

@flyte.task
def iterative_training(epochs: int) -> int:
    ctx = flyte.ctx()
    start_epoch = 0

    if ctx and ctx.checkpoints:
        if ctx.checkpoints.prev_checkpoint_path and os.path.exists(ctx.checkpoints.prev_checkpoint_path):
            # Restore state from previous attempt
            with open(ctx.checkpoints.prev_checkpoint_path, "r") as f:
                start_epoch = int(f.read().strip())

        # Train and write new checkpoint
        for epoch in range(start_epoch, epochs):
            # ... training step ...
            if ctx.checkpoints.checkpoint_path:
                with open(ctx.checkpoints.checkpoint_path, "w") as f:
                    f.write(str(epoch + 1))

    return epochs

Code Bundles (CodeBundle)​

The ctx.code_bundle field exposes packaged task artifacts via the CodeBundle dataclass:

@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

CodeBundle computes the code version using a hash of the code files and requires at least one archive artifact (tgz or pkl), raising a ValueError if both are None. During remote execution, downloaded_path contains the local filesystem path where the code archive was unpacked.

Emitting Visualizations with Reports​

Every TaskContext carries a Report object in ctx.report. Flyte's reporting system allows tasks to generate custom HTML visual decks and dashboards organized into tabs.

You can interact with reports directly using the module-level functions in flyte.report:

import flyte
import flyte.report

@flyte.task(report=True)
def analyze_dataset() -> None:
    # Append HTML content to the default "main" tab
    flyte.report.log("<h2>Dataset Summary</h2><p>Processed 10,000 records successfully.</p>")

    # Add content to a dedicated custom tab
    custom_tab = flyte.report.get_tab("Metrics", create_if_missing=True)
    custom_tab.log("<p>Accuracy: 98.5%</p>")

    # Replace entire tab content if needed
    custom_tab.replace("<p>Accuracy updated: 99.1%</p>")

Report Flushing and Lifecycle​

Inside flyte.report:

  • log(content: str, do_flush: bool = False): Appends an HTML snippet into the active tab. Setting do_flush=True immediately uploads the report.
  • get_tab(name: str, create_if_missing: bool = True): Retrieves or initializes a tab.
  • replace(content: str, do_flush: bool = False): Clears existing logs in the main tab and replaces them with the new HTML snippet.
  • flush(): Renders all tabs into the template and uploads the HTML payload using flyte.storage.put_stream to io.report_path(ctx.output_path).

When a task decorator specifies @env.task(report=True), the flyte-sdk task runner automatically calls flyte.report.flush() upon task completion. If report=False (the default), call flyte.report.flush() explicitly within the task to write the HTML report to storage.