Managing Data Paths

When executing tasks, flyte-sdk separates task metadata (action inputs, outputs, and status descriptors) from raw user payloads (files, dataframes, and binary blobs). It manages these storage locations through two primary classes: RawDataPath, which provisions directory prefixes and allocates collision-free paths, and TaskContext, which provides execution metadata and data locations at runtime.

Managing Storage Locations with RawDataPath

RawDataPath encapsulates the root storage prefix for task inputs, intermediate artifacts, and outputs. It handles path generation across local filesystems and remote object storage (such as AWS S3, Google Cloud Storage, or Azure Blob Storage).

Creating Local Raw Data Paths

To initialize a raw data directory on the local filesystem, use RawDataPath.from_local_folder():

import pathlib
from flyte.models import RawDataPath

# 1. Automatically create a temporary directory using tempfile.mkdtemp()
temp_raw_path = RawDataPath.from_local_folder()
print(temp_raw_path.path)  # e.g., /tmp/tmp_abc123

# 2. Pass a pathlib.Path object (creates directory recursively on disk if missing)
local_dir = pathlib.Path("/tmp/flyte/experiments")
pathlib_raw_path = RawDataPath.from_local_folder(local_dir)
print(pathlib_raw_path.path)  # /tmp/flyte/experiments

# 3. Pass a raw string (stores the path string directly without disk verification)
string_raw_path = RawDataPath.from_local_folder("/tmp/flyte/custom")
print(string_raw_path.path)  # /tmp/flyte/custom

Generating Unique Remote and Local Paths

RawDataPath.get_random_remote_path() appends a 128-bit UUID hex string to the base path. It detects whether the protocol is local (file:// or local paths) or remote (s3://, gs://, abfs://) via fsspec:

from flyte.models import RawDataPath

# Remote object store prefix
remote_rd = RawDataPath(path="s3://my-flyte-bucket/data-prefix")

# Generate a random remote directory URI
remote_dir = remote_rd.get_random_remote_path()
# Output: s3://my-flyte-bucket/data-prefix/4a6d61f1c7504a3e9c60e45cf56f272a

# Generate a random remote file URI
remote_file = remote_rd.get_random_remote_path(file_name="results.parquet")
# Output: s3://my-flyte-bucket/data-prefix/4a6d61f1c7504a3e9c60e45cf56f272a/results.parquet

# Local filesystem prefix
local_rd = RawDataPath(path="/tmp/flyte/storage")

# Generate a local path with a file name
# Note: When file_name is specified on local filesystems, parent directories are created and the file is touched
local_file = local_rd.get_random_remote_path(file_name="output.csv")
# Output: /tmp/flyte/storage/e2b4f98129a04a8b835e5cc4de62f01d/output.csv

Accessing TaskContext in Tasks

TaskContext holds all execution metadata for a running task. Inside a task, you can inspect runtime data, check cluster execution status, and access the active RawDataPath via flyte.ctx().

import flyte
from flyte.io import File

env = flyte.TaskEnvironment("data-pipeline")

@env.task
async def generate_dataset() -> File:
    current_ctx = flyte.ctx()
    if current_ctx is None:
        raise RuntimeError("TaskContext is only accessible inside active task executions.")

    # Check execution environment
    if current_ctx.is_in_cluster():
        print(f"Running remotely in cluster mode. Version: {current_ctx.version}")
    else:
        print(f"Running in mode: {current_ctx.mode}")

    # Inspect context paths
    print(f"Action ID: {current_ctx.action.name}")
    print(f"Run Base Directory: {current_ctx.run_base_dir}")
    print(f"Output Path: {current_ctx.output_path}")

    # Allocate a destination path for arbitrary file uploads
    destination_uri = current_ctx.raw_data_path.get_random_remote_path(file_name="data.txt")
    
    local_file = "/tmp/local_data.txt"
    with open(local_file, "w") as f:
        f.write("Processed content")

    # Upload using flyte.storage
    remote_path = await flyte.storage.put(from_path=local_file, to_path=destination_uri)
    return File(path=remote_path)

Context Immutability and Custom Data

TaskContext is a frozen dataclass (@dataclass(frozen=True)). To update parameters or attach custom execution attributes, use replace():

# TaskContext must be updated via replace() rather than direct assignment
updated_ctx = current_ctx.replace(
    data={"pipeline_id": "exp_102", "batch_size": 64}
)

# Custom data dictionary can be indexed directly on TaskContext
batch_size = updated_ctx["batch_size"]  # 64

Integration with High-Level Storage and I/O

Flyte's high-level types (File, DataFrame, and flyte.storage) automatically delegate path allocation to the context's RawDataPath whenever explicit destination paths are omitted.

Transparent Allocation in flyte.storage.put

When to_path is omitted in flyte.storage.put(), flyte-sdk queries ctx.raw_data to generate a random destination path, preserving the file name if recursive=False:

import flyte

@env.task
async def save_artifact() -> str:
    # Omitting to_path automatically generates a unique path under ctx.raw_data_path
    uploaded_uri = await flyte.storage.put(from_path="local_artifact.json")
    return uploaded_uri

Allocating URIs with File.new_remote()

You can initialize an unpopulated File pointer bound to a random remote path in the active raw data prefix:

from flyte.io import File
import flyte

@env.task
async def create_remote_file() -> File:
    # Allocates a new random URI using ctx.raw_data.get_random_remote_path()
    remote_file = File.new_remote()
    print(f"Allocated remote URI: {remote_file.path}")
    return remote_file

Configuring Custom Data Paths

By default, local runs provision a temporary local folder, while remote runs allocate directories under run_base_dir/rd/<random_id>. You can override the base directory and raw data destination using flyte.with_runcontext().

Custom Local Run Paths

import flyte

env = flyte.TaskEnvironment("path-config")

@env.task
async def process_task() -> int:
    return 42

# Configure a specific local folder for task outputs and artifacts
runner = flyte.with_runcontext(
    mode="local",
    raw_data_path="/var/tmp/flyte/custom_raw_data",
)
await runner.run(process_task)

Custom Remote Storage Prefix

For hybrid or remote task execution against cloud object storage:

import flyte

# Configure run_base_dir and raw_data_path for remote execution
runner = flyte.with_runcontext(
    mode="remote",
    run_base_dir="s3://my-bucket/flyte-runs/run-001",
    raw_data_path="s3://my-bucket/flyte-runs/run-001/custom_raw_data",
)
await runner.run(process_task)

Troubleshooting and Gotchas

Local Filesystem Creation Differences in RawDataPath.from_local_folder

RawDataPath.from_local_folder behaves differently depending on the input type:

• pathlib.Path: Automatically calls mkdir(parents=True, exist_ok=True) on disk.
• None: Creates a new temporary folder on disk via tempfile.mkdtemp().
• str: Initializes the dataclass with the string path directly without creating the directory on disk. If passing a string path that does not yet exist, create it manually or wrap it in pathlib.Path(...).

Local File Touching in get_random_remote_path

When targeting local storage (file:// or standard file paths):

• If file_name is provided (e.g., get_random_remote_path(file_name="data.csv")), RawDataPath creates the parent directory and touches the file (local_path.touch()) on disk.
• If file_name is None (e.g., get_random_remote_path()), no file or folder is touched on disk, leaving directory creation to downstream writers.

Accessing flyte.ctx() Outside Task Execution

Calling flyte.ctx() outside of an active task execution returns None. When writing utility functions that may run both inside tasks and as standalone scripts, check whether flyte.ctx() is None before dereferencing attributes like ctx.raw_data_path.

Missing run_base_dir in Remote Execution

Hybrid and remote task executions require a root storage prefix. If run_base_dir is not configured in remote mode, flyte-sdk raises ValueError: Raw data path is required when running task. Ensure run_base_dir or raw_data_path is passed via flyte.with_runcontext(...) when launching remote runs.