The Runtime `TaskContext`

When code needs the current run ID, task output prefix, or checkpoint locations, read the runtime context with flyte.ctx() inside the task:

import flyte

@flyte.task
def inspect_runtime() -> str:
    task_context = flyte.ctx()
    if task_context is None:
        return "not running as a task"

    return (
        f"run={task_context.action.run_name} "
        f"action={task_context.action.name} "
        f"raw_data={task_context.raw_data_path.path}"
    )

ctx is re-exported by flyte/__init__.py from _context.py. It returns the current TaskContext while execution is installed, and returns None when the callable is being used as ordinary Python rather than through a task run. The context is installed around local, remote, and runtime dispatch; task code should therefore guard the result before dereferencing it.

Read execution-specific fields

TaskContext is a frozen, keyword-only dataclass in models.py. Its required fields are action, version, raw_data_path, output_path, run_base_dir, and report. input_path, group_data, checkpoints, code_bundle, and compiled_image_cache may be absent; data, mode, and interactive_mode have defaults.

Identify the current action

The action field is an ActionID. It contains the action name, its enclosing run_name, and optional project, domain, and org values:

import flyte

@flyte.task
def current_action() -> str:
    task_context = flyte.ctx()
    if task_context is None:
        return "no task context"

    action = task_context.action
    return f"{action.project}:{action.domain}:{action.run_name}:{action.name}"

If run_name is omitted when an ActionID is constructed, ActionID.__post_init__ sets it to name. ActionID.create_random() generates a UUID-based name and sets both the action and run names to that value. The runtime's remote path constructs an action with the supplied action name, run name, project, domain, and organization; local execution either uses the configured name or calls create_random().

Locate raw data and run outputs

raw_data_path is a RawDataPath wrapper. Its .path is the raw-data prefix used by task IO helpers, while output_path and run_base_dir identify the task output and run-level base locations:

import flyte

@flyte.task
def runtime_paths() -> dict[str, str]:
    task_context = flyte.ctx()
    if task_context is None:
        return {}

    return {
        "raw_data": task_context.raw_data_path.path,
        "output": task_context.output_path,
        "run_base": task_context.run_base_dir,
    }

File.new_remote() is a concrete consumer of this context. It calls internal_ctx().raw_data.get_random_remote_path() and uses the result as the new file's path. Consequently, code running inside a task can allocate a destination through the public File API without choosing the raw-data prefix itself:

from flyte.io import File

@flyte.task
def create_file() -> File[str]:
    file = File.new_remote()
    return file

The exact file-writing pattern embedded in io/_file.py is:

@flyte.task
def write_file() -> File[DataFrame]:
    df = pd.DataFrame(...)
    file = File.new_remote()
    async with file.open("wb") as f:
        df.to_csv(f)
    return file

The File example is an embedded source docstring and assumes the DataFrame and pandas names used by that example are available in the surrounding task. new_remote() requires an initialized Flyte context; its path allocation comes from the current raw-data prefix.

Compare local and remote context construction

Remote execution in _run.py derives the raw-data prefix from the run base directory and a random ID, then creates checkpoint prefixes beneath it. The context is installed while run_task executes:

output_path = self._run_base_dir
raw_data_path = f"{output_path}/rd/{random_id}"
raw_data_path_obj = RawDataPath(path=raw_data_path)
checkpoint_path = f"{raw_data_path}/checkpoint"
prev_checkpoint = f"{raw_data_path}/prev_checkpoint"
checkpoints = Checkpoints(checkpoint_path, prev_checkpoint)

async def _run_task() -> Tuple[Any, Optional[Exception]]:
    ctx = internal_ctx()
    tctx = TaskContext(
        action=action,
        checkpoints=checkpoints,
        code_bundle=code_bundle,
        output_path=output_path,
        version=version if version else "na",
        raw_data_path=raw_data_path_obj,
        compiled_image_cache=image_cache,
        run_base_dir=self._run_base_dir,
        report=flyte.report.Report(name=action.name),
    )
    async with ctx.replace_task_context(tctx):
        return await run_task(tctx=tctx, controller=controller, task=obj, inputs=inputs)

The remote constructor does not explicitly set mode, so it uses TaskContext's default, "remote". A local run uses the current context's raw-data path for both checkpoint fields, sets code_bundle and compiled_image_cache to None, and explicitly sets mode="local":

tctx = TaskContext(
    action=action,
    checkpoints=Checkpoints(
        prev_checkpoint_path=internal_ctx().raw_data.path,
        checkpoint_path=internal_ctx().raw_data.path,
    ),
    code_bundle=None,
    output_path=self._metadata_path,
    run_base_dir=self._metadata_path,
    version="na",
    raw_data_path=internal_ctx().raw_data,
    compiled_image_cache=None,
    report=Report(name=action.name),
    mode="local",
)
with ctx.replace_task_context(tctx):
    ...

Runtime dispatch in _internal/runtime/taskrunner.py preserves an existing context's mode for nested execution; otherwise it uses "remote". is_in_cluster() is narrower than “has remote-related work”: it returns True only when mode == "remote", so it returns False for both local and hybrid.

Update context data without mutating the context

Because TaskContext is frozen, use its replace() method instead of assigning fields. It returns a new context. A data mapping is treated specially: supplied keys are copied into the existing data mapping, and new values overwrite keys with the same name. Bracket access reads from that mapping and returns None for a missing key:

import flyte

@flyte.task
def use_task_data() -> str:
    task_context = flyte.ctx()
    if task_context is None:
        return "no task context"

    updated = task_context.replace(data={"request_id": "abc-123"})
    return updated["request_id"]

replace(data=None) leaves the existing mapping unchanged while still applying any other keyword replacements. The nested dictionary itself is copied when a non-None mapping is supplied; identity fields such as action remain unchanged unless explicitly replaced.

Scope grouping with group()

Use flyte.group() around nested task submissions when they should carry a group name:

import flyte

@flyte.task
async def my_task():
    with flyte.group("my_group"):
        t1(x, y)

The implementation reads the current internal context. If no task context is installed, group(name) yields without changing execution. Inside a task it creates GroupData(name), calls TaskContext.replace(group_data=...), and installs that copied context only for the with block. GroupData contains only its name.

Understand child action IDs and output paths

Nested task submission derives child identity from the parent TaskContext. generate_sub_action_id_and_output_path() in _internal/runtime/convert.py takes the parent action, a task specification or task name, an input hash, and an invocation sequence. For a task specification it hashes deterministic serialization; it then passes the task hash, input hash, sequence, and optional group name to ActionID.new_sub_action_from():

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]:
    current_action_id = tctx.action
    current_output_path = tctx.run_base_dir
    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_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

new_sub_action_from() concatenates the parent action name, input hash, task hash, invocation sequence, and optional group, hashes that string with hashlib.md5, and base36-encodes the digest. It returns a copied ActionID, preserving the parent run and organization fields. The MD5 value is used here for deterministic naming. The child output prefix is run_base_dir joined with the generated child action name.

Allocate raw-data paths

RawDataPath.from_local_folder() has three distinct branches:

from pathlib import Path

from flyte.models import RawDataPath

path_from_path = RawDataPath.from_local_folder(Path("/tmp/flyte-data"))
temporary_path = RawDataPath.from_local_folder()
path_from_string = RawDataPath.from_local_folder("/tmp/already-configured")

A pathlib.Path is created recursively. Passing None creates a temporary directory with tempfile.mkdtemp(). A string is only wrapped; the string-path directory is not created by this method. Any other value raises ValueError.

get_random_remote_path(file_name=None) appends a random UUID-derived component. For a file protocol it returns an absolute local path. With a filename, it creates the parent directory and touches the file; without a filename, it only returns the generated path. For other protocols, it obtains the fsspec filesystem and joins the prefix, random component, and optional filename with that filesystem's separator; it does not create a remote object.

The local run() path installs a RawDataPath before constructing its local TaskContext:

with internal_ctx().new_raw_data_path(
    raw_data_path=RawDataPath.from_local_folder(local_folder=self._raw_data_path)
):
    return await self._run_local(task, *args, **kwargs)

At the lower level, internal_ctx().raw_data first checks the task context's raw_data_path, then a context-level raw-data path. If neither is present it raises ValueError.

Use checkpoints and code-bundle metadata

checkpoints is either None or a frozen Checkpoints pair with prev_checkpoint_path and checkpoint_path. The class performs no validation or path normalization. Prefer keyword arguments when constructing one so the two fields cannot be confused:

from flyte.models import Checkpoints

checkpoints = Checkpoints(
    prev_checkpoint_path="s3://bucket/run/previous",
    checkpoint_path="s3://bucket/run/current",
)

The remote _run.py call shown above uses positional arguments as Checkpoints(checkpoint_path, prev_checkpoint), where the class's first field is prev_checkpoint_path; inspect that call carefully when interpreting runtime values. Local execution deliberately assigns the same raw-data path to both fields.

code_bundle is a CodeBundle when runtime task code was packaged. Construction requires tgz or pkl:

from flyte.models import CodeBundle

bundle = CodeBundle(
    computed_version="code-hash",
    destination="/tmp/flyte-bundle",
    pkl="s3://bucket/task.pkl",
)

CodeBundle.with_downloaded_path() returns a copy with a pathlib.Path in downloaded_path; it does not mutate the frozen bundle. Runtime entrypoints call download_bundle(), then attach that path:

async def download_code_bundle(code_bundle: CodeBundle) -> CodeBundle:
    logger.debug(f"Downloading {code_bundle}")
    downloaded_path = await download_bundle(code_bundle)
    return code_bundle.with_downloaded_path(downloaded_path)

download_bundle() requires destination to already be a directory. A tgz bundle is downloaded there and extracted with the system tar command; an existing downloaded archive is reused. A pkl bundle is downloaded directly. Missing both sources raises ValueError, and a failed tar extraction raises RuntimeError. For pickle loading, the runtime reads the downloaded path as a gzip-compressed cloudpickle file.

Troubleshoot context access

Symptom	Check
flyte.ctx() is None	The callable is outside a task context. Guard the result, or execute it through a Flyte run.
internal_ctx().raw_data raises ValueError	Neither TaskContext.raw_data_path nor a context-level raw-data path was installed.
is_in_cluster() is False in hybrid mode	The method returns True only for exactly mode == "remote"; hybrid mode is not treated as remote by this method.
Local string raw-data folder is missing	RawDataPath.from_local_folder("...") does not create the directory. Use a pathlib.Path or None branch when creation is required.
Checkpoint fields appear reversed	Checkpoints declares prev_checkpoint_path first. Use keyword arguments; the remote _run.py construction currently passes positional values in the opposite-looking order.
CodeBundle construction fails	Supply at least one of tgz or pkl. Also ensure destination already exists as a directory before downloading.
Context changes leak across asynchronous work	_context.py documents the context as not generally coroutine-safe and assumes one thread. Use contextual_run for a new context tree and pair context entry and exit correctly.

The runtime CLI populates action and run metadata from its command-line/environment inputs, including run and action names, project, domain, organization, and run base directory. In normal task code, prefer flyte.ctx() over reconstructing those values: the installed TaskContext is the object that the runner, controllers, IO helpers, and runtime entrypoints pass through execution.