The Task Execution Context

When running tasks across local development environments, containerized remote clusters, or nested execution pipelines, task code often needs access to runtime metadata—such as the execution run identifier, deterministic storage paths for artifacts, code versions, checkpoint locations, or execution mode. Hardcoding paths or attempting to pass runtime identifiers down through function arguments creates brittle code and breaks task reusability.

In flyte-sdk, runtime state is encapsulated in the TaskContext object and made available to tasks through flyte.ctx().

Accessing the Current Task Context

Inside any running Flyte task, call flyte.ctx() to retrieve the current TaskContext instance. If your code is called outside of an active execution (such as a plain Python unit test or script invocation without flyte.run()), flyte.ctx() returns None.

import flyte

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

    action = ctx.action
    print(f"Executing action '{action.name}' in run '{action.run_name}'")
    print(f"Task version: {ctx.version}, Mode: {ctx.mode}")
    print(f"Output directory: {ctx.output_path}")

    return f"Processed batch {batch_id} in {action.name}"

Context Storage and Propagation

Internally, flyte-sdk manages runtime state using Python's contextvars.ContextVar. In _context.py, the global context variable root_context_var stores the active Context container:

# _context.py
root_context_var = contextvars.ContextVar("root", default=Context(data=ContextData()))

def ctx() -> Optional[TaskContext]:
    """Retrieve the current task context from the context variable."""
    return internal_ctx().data.task_context

def internal_ctx() -> Context:
    """Retrieve the current context from the context variable."""
    return root_context_var.get()

When task runners (such as _internal/runtime/taskrunner.py or entrypoint scripts in _bin/runtime.py) initialize a task, they instantiate a TaskContext and bind it into the execution tree using ctx.replace_task_context(tctx).

---

TaskContext Structure and Core Attributes

The TaskContext dataclass in models.py is defined as a frozen, keyword-only dataclass (@dataclass(frozen=True, kw_only=True)):

Attribute	Type	Description
action	ActionID	Identifiers for the current task action, run, project, domain, and org.
version	str	The version string of the executing task.
raw_data_path	RawDataPath	Root directory and helper for allocating offloaded data storage paths.
input_path	str | None	Storage path where task input parameters are staged.
output_path	str	Storage path where task output results are written.
run_base_dir	str	Base working directory for the parent run.
report	Report	Report collector for logging deck/report outputs during task execution.
group_data	GroupData | None	Group metadata if running inside a flyte.group(...) block.
checkpoints	Checkpoints | None	Paths for loading and saving recovery state across retries or iterations.
code_bundle	CodeBundle | None	Code packaging metadata, including version hash and inflated file locations.
compiled_image_cache	ImageCache | None	Cache reference for container images built or resolved during runtime.
data	Dict[str, Any]	Custom dictionary for storing arbitrary execution-scoped context data.
mode	Literal["local", "remote", "hybrid"]	The execution environment mode (defaults to "remote").
interactive_mode	bool	Flag indicating whether the task is running in interactive mode.

---

Action Hierarchy and Identification (ActionID)

Every task execution corresponds to an action within a run, tracked by flyte.models.ActionID:

# models.py
@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

During initialization, ActionID.__post_init__ ensures that if run_name is omitted, it defaults to the action name.

Action Naming and Sub-Actions

1. Random Action Creation: ActionID.create_random() generates a randomized UUID name for both name and run_name.
2. Sub-Actions: When a task spawns child operations or nested workflows, new_sub_action(name=None) creates a child ActionID preserving project, domain, org, and run_name while updating name.
3. Deterministic Sub-Actions: In runtime controllers (_internal/runtime/convert.py), sub-actions are created deterministically based on hashes of inputs, task template specs, call sequences, and group names:

sub_action_id = current_action_id.new_sub_action_from(
    task_hash=task_hash,
    input_hash=inputs_hash,
    group=tctx.group_data.name if tctx.group_data else None,
    task_call_seq=invoke_seq,
)

new_sub_action_from computes an MD5 digest over "{name}-{input_hash}-{task_hash}-{task_call_seq}[-{group}]" and encodes it using base36.

Context Enrichment in Logging

flyte-sdk uses ActionID to tag all log records emitted during execution. In _logging.py, ContextFilter inspects flyte.ctx() and prepends the run and action names:

# _logging.py
class ContextFilter(logging.Filter):
    def filter(self, record):
        from flyte._context import ctx

        c = ctx()
        if c:
            action = c.action
            record.msg = f"[{action.run_name}][{action.name}] {record.msg}"
        return True

---

Offloaded Storage and Data Paths (RawDataPath)

Tasks handling large files, datasets, or dataframes require dedicated storage locations offloaded from standard metadata stores. TaskContext.raw_data_path is a RawDataPath instance pointing to the root storage directory:

# models.py
@dataclass(frozen=True, kw_only=True)
class RawDataPath:
    path: str

Path Generation Behavior

RawDataPath.get_random_remote_path(file_name=None) generates storage target URIs using fsspec:

• Local Storage (file:// or local paths): Resolves the path, generates a 128-bit hex UUID subfolder, and if file_name is supplied, creates parent directories and touches the file locally.
• Remote Storage (S3, GCS, ABFS): Inspects the protocol with fsspec.utils.get_protocol(file_prefix), strips trailing separators, and joins the UUID and optional file_name using the filesystem's separator fs.sep.

Internal IO mechanisms like flyte.io.File and dataframe encoders call this method to stage files:

import flyte

@flyte.task
def export_custom_artifact() -> str:
    ctx = flyte.ctx()
    # Allocate a unique location under the task's raw data path
    destination_uri = ctx.raw_data_path.get_random_remote_path(file_name="metrics.json")
    
    # Write directly to destination_uri or return the URI
    return destination_uri

Locally, RawDataPath.from_local_folder(local_folder) handles initialization: passing a pathlib.Path ensures the directory exists, passing None creates a new directory via tempfile.mkdtemp(), and strings are wrapped directly.

---

Checkpointing State (Checkpoints)

For fault-tolerant, iterative, or resumable tasks, TaskContext.checkpoints holds paths pointing to previous and current checkpoint locations:

# models.py
@dataclass(frozen=True)
class Checkpoints:
    prev_checkpoint_path: str | None
    checkpoint_path: str | None

When a task fails and is retried, or when an iterative task executes across cycles, prev_checkpoint_path allows the task to read the latest saved state, while checkpoint_path specifies where the new state should be written.

---

Code Packaging and Versioning (CodeBundle)

The task context tracks the deployed code artifacts and version identifiers:

• TaskContext.version: The string version identifier for the running task. Remote task definitions configured with auto_version="current" inspect flyte.ctx().version at runtime to inherit the active version.
• TaskContext.code_bundle: A CodeBundle dataclass that describes packaging:

# models.py
@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")

    def with_downloaded_path(self, path: pathlib.Path) -> CodeBundle:
        return replace(self, downloaded_path=path)

CodeBundle requires either tgz (archive file path) or pkl (serialized package path) upon creation. During container bootstrapping, with_downloaded_path updates the bundle once unpacked to local disk.

---

Grouping Tasks (GroupData and flyte.group)

Nested tasks and map operations can be organized into logical groups using the flyte.group context manager in _group.py:

import flyte

@flyte.task
def sub_step(val: int) -> int:
    return val * 2

@flyte.task
def workflow_task(x: int) -> int:
    with flyte.group("preprocessing_stage"):
        res = sub_step(x)
    return res

Under the hood, flyte.group(name) reads the current task context, creates a GroupData(name=name) object, and updates the active context:

# _group.py
@contextmanager
def group(name: str):
    ctx = internal_ctx()
    if ctx.data.task_context is None:
        yield
        return
    tctx = ctx.data.task_context
    new_tctx = tctx.replace(group_data=GroupData(name))
    with ctx.replace_task_context(new_tctx):
        yield

Any sub-actions generated within this block include the group name in their deterministic action ID hashing.

---

Execution Environment and Immutability

Checking Cluster Execution

Tasks can check whether they are running in a remote cluster environment using is_in_cluster():

ctx = flyte.ctx()
if ctx and ctx.is_in_cluster():
    print("Running remotely in cluster pod")
else:
    print(f"Running locally or hybrid (mode={ctx.mode if ctx else 'none'})")

is_in_cluster() returns True when mode == "remote".

Immutability with replace()

Because TaskContext is a frozen dataclass, direct attribute assignment (ctx.version = "v2") raises a FrozenInstanceError. To create modified contexts, call TaskContext.replace(**kwargs):

# models.py
def replace(self, **kwargs) -> TaskContext:
    if "data" in kwargs:
        rec_data = kwargs.pop("data")
        if rec_data is None:
            return replace(self, **kwargs)
        data = {}
        if self.data is not None:
            data = self.data.copy()
        data.update(rec_data)
        kwargs.update({"data": data})
    return replace(self, **kwargs)

When passing custom keys to data, replace() automatically creates a copy of the existing dictionary and merges new entries rather than overwriting the whole dictionary. Task dictionary values can also be retrieved directly using indexing (ctx["custom_key"]).