# Build tasks
> This bundle contains all pages in the Build tasks section.
> Source: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming ===

# Build tasks

> **📝 Note**
>
> An LLM-optimized bundle of this entire section is available at [`section.md`](section.md).
> This single file contains all pages in this section, optimized for AI coding agent context.

This section covers the essential programming patterns and techniques for developing robust Flyte workflows. Once you understand the basics of task configuration, these guides will help you build sophisticated, production-ready data pipelines and machine learning workflows.

## What you'll learn

The task programming section covers key patterns for building effective Flyte workflows:

**Data handling and types**

- **Tasks > Build tasks > Files and directories**: Work with large datasets using Flyte's efficient file and directory types that automatically handle data upload, storage, and transfer between tasks.

- **Tasks > Build tasks > DataFrames**: Pass DataFrames between tasks without downloading data into memory, with support for Pandas, Polars, PyArrow, Dask, and other DataFrame backends.
- **Tasks > Build tasks > Data classes and structures**: Use Python data classes and Pydantic models as task inputs and outputs to create well-structured, type-safe workflows.
- **Tasks > Build tasks > Custom context**: Use custom context to pass metadata through your task execution hierarchy without adding parameters to every task.

**Execution patterns**

- **Tasks > Build tasks > Fanout**: Scale your workflows by running many tasks in parallel, perfect for processing large datasets or running hyperparameter sweeps.
- **Tasks > Build tasks > Mapping over inputs**: Apply the same task to every item of a list with `flyte.map`: in-order results, error handling, concurrency limits, and partials.
- **Tasks > Build tasks > Consuming a message queue**: Pull messages from an external queue such as AWS SQS and fan out processing across a pool of reusable containers.
- **Tasks > Build tasks > Controlling parallel execution**: Limit concurrent task executions using semaphores or `flyte.map` concurrency for rate-limited APIs, GPU quotas, and resource-constrained workflows.
- **Tasks > Build tasks > Streaming map-reduce**: Process fanout results as they complete with `asyncio.as_completed`, reducing in batches incrementally instead of waiting for every task to finish.
- **Tasks > Build tasks > Task dependencies and ordering**: Replicate DAG-like behavior (sequencing, fan-out, fan-in, and fine-grained dependency-driven scheduling) using `asyncio` in Flyte 2's implicit dependency model.
- **Tasks > Build tasks > Structured concurrency with anyio**: Use `anyio` task groups as a top-level structured-concurrency alternative to raw `asyncio`, with automatic sibling cancellation when one task fails.
- **Tasks > Build tasks > External conditions**: Pause a task until an external signal arrives: a human approval, a callback from an external service, or a value supplied at runtime.
- **Tasks > Build tasks > Grouping actions**: Organize related task executions into logical groups for better visualization and management in the UI.
- **Tasks > Build tasks > Run a bioinformatics tool**: Run arbitrary containers in any language without the Flyte SDK installed, using Flyte's copilot sidecar for data flow.
- **Tasks > Build tasks > Remote tasks**: Use previously deployed tasks without importing their code or dependencies, enabling team collaboration and task reuse.
- [**Pod templates**](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/pod-templates/page.md): Extend tasks with Kubernetes pod templates to add sidecars, volume mounts, and advanced Kubernetes configurations.
- **Tasks > Build tasks > Abort and cancel actions**: Stop in-progress actions automatically, programmatically, or manually via the CLI and UI.
- **Tasks > Build tasks > Regular async function (not a task)**: Advanced patterns like task forwarding and other specialized task execution techniques.
- **Tasks > Build tasks > Higher-order functions**: Write reusable functions that take tasks as arguments — fallback, memory-scaling retry, circuit breaker, and batch map-reduce wrappers built on Flyte's dynamic execution.

**Development and debugging**

- **Tasks > Build tasks > Notebooks**: Write and iterate on workflows directly in Jupyter notebooks for interactive development and experimentation.
- **Tasks > Build tasks > Test business logic directly**: Test your Flyte tasks using direct invocation for business logic or `flyte.run()` for Flyte-specific features.
- **Tasks > Build tasks > Links**: Add clickable URLs to tasks in the Flyte UI, connecting them to external tools like experiment trackers and monitoring dashboards.
- **Tasks > Build tasks > Reports**: Generate custom HTML reports during task execution to display progress, results, and visualizations in the UI.
- **Tasks > Build tasks > Traces**: Add fine-grained observability to helper functions within your tasks for better debugging and resumption capabilities.
- **Tasks > Build tasks > Intra-task checkpoints**: Save in-progress state within a task (such as a training loop) so retries resume from the last checkpoint instead of starting over.
- **Tasks > Build tasks > Error handling**: Implement robust error recovery strategies, including automatic resource scaling and graceful failure handling.

## When to use these patterns

These programming patterns become essential as your workflows grow in complexity:

- Use **fanout** when you need to process multiple items concurrently or run parameter sweeps.
- Use **mapping over inputs** to apply the same task to every item of a list, and **controlling parallel execution** when you need to limit how many run at the same time.
- Apply **streaming map-reduce** when map tasks have uneven durations or you want to reduce results in batches as they complete, rather than waiting for the entire fanout to finish.
- Implement **error handling** for production workflows that need to recover from infrastructure failures.
- Apply **grouping** to organize complex workflows with many task executions.
- Use **files and directories** when working with large datasets that don't fit in memory.
- Use **DataFrames** to efficiently pass tabular data between tasks across different processing engines.
- Choose **container tasks** when you need to run code in non-Python languages, use legacy containers, or execute AI-generated code in sandboxes.
- Use **remote tasks** to reuse tasks deployed by other teams without managing their dependencies.
- Apply **pod templates** when you need advanced Kubernetes features like sidecars or specialized storage configurations.
- Use **traces** to debug non-deterministic operations like API calls or ML inference.
- Use **intra-task checkpoints** to make long-running training loops resumable across retries, preemptions, and interruptions.
- Use **links** to connect tasks to external tools like Weights & Biases, Grafana, or custom dashboards directly from the Flyte UI.
- Create **reports** to monitor long-running workflows and share results with stakeholders.
- Use **custom context** when you need lightweight, cross-cutting metadata to flow through your task hierarchy without becoming part of the task's logical inputs.
- Write **unit tests** to validate your task logic and ensure type transformations work correctly before deployment.
- Use **abort and cancel** to stop unnecessary actions when conditions change, such as early convergence in HPO or manual intervention.
- Use **external conditions** to insert approval gates or data collection checkpoints into automated workflows.
- Apply **higher-order functions** to factor recurring orchestration logic — retry-on-OOM, fallback, circuit breaking, batching — into reusable wrappers that work with any task.

Each guide includes practical examples and best practices to help you implement these patterns effectively in your own workflows.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/files-and-directories ===

# Files and directories

Flyte provides the `flyte.io.File` and
`flyte.io.Dir` types to represent files and directories, respectively.
Together with [`flyte.io.DataFrame`](./dataframes) they constitute the *offloaded data types* - unlike [materialized types](./dataclasses-and-structures) like data classes, these pass references rather than full data content.

A variable of an offloaded type does not contain its actual data, but rather a reference to the data.
The actual data is stored in the internal blob store of your Union/Flyte instance.
When a variable of an offloaded type is first created, its data is uploaded to the blob store.
It can then be passed from task to task as a reference.
The actual data is only downloaded from the blob stored when the task needs to access it, for example, when the task calls `open()` on a `File` or `Dir` object.

This allows Flyte to efficiently handle large files and directories without needing to transfer the data unnecessarily.
Even very large data objects like video files and DNA datasets can be passed efficiently between tasks.

For the full picture of what gets stored in the bucket versus what stays in the control plane database, see [Where your data lives](https://www.union.ai/docs/latest/flyte/user-guide/get-started/core-concepts/where-data-lives).

The `File` and `Dir` classes provide both synchronous and asynchronous methods to interact with the data, so you can use them from either kind of task. See **Tasks > Build tasks > Files and directories > Synchronous and asynchronous APIs** for the full method pairing.

> [!NOTE]
> Because `File` and `Dir` are passed by reference, a downstream cached task does not get a cache hit on identical content stored at a new path. To cache on content, attach a hash at production time - see [Content-based caching for DataFrames, files, and directories](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/caching).

## Example usage

The examples below show the basic use-cases of uploading files and directories created locally, and using them as inputs to a task.

```
import asyncio
import tempfile
from pathlib import Path

import flyte
from flyte.io import Dir, File

env = flyte.TaskEnvironment(name="files-and-folders")

@env.task
async def write_file(name: str) -> File:

    # Create a file and write some content to it
    with open("test.txt", "w") as f:
        f.write(f"hello world {name}")

    # Upload the file using flyte
    uploaded_file_obj = await File.from_local("test.txt")
    return uploaded_file_obj
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/file_and_dir.py*

The upload happens when the [`File.from_local`](https://www.union.ai/docs/latest/flyte/api-reference/flyte-sdk/flyte.io/file) command is called.
Because the upload would otherwise block execution, `File.from_local` is implemented as an `async` function.
The Flyte SDK frequently uses this class constructor pattern, so you will see it with other types as well.

This is a slightly more complicated task that calls the task above to produce `File` objects.

These are assembled into a directory and the `Dir` object is returned, also via invoking `from_local`.

```
@env.task
async def write_and_check_files() -> Dir:
    coros = []
    for name in ["Alice", "Bob", "Eve"]:
        coros.append(write_file(name=name))

    vals = await asyncio.gather(*coros)
    temp_dir = tempfile.mkdtemp()
    for file in vals:
        async with file.open("rb") as fh:
            contents = await fh.read()
            # Convert bytes to string
            contents_str = contents.decode('utf-8') if isinstance(contents, bytes) else str(contents)
            print(f"File {file.path} contents: {contents_str}")
            new_file = Path(temp_dir) / file.name
            with open(new_file, "w") as out:  # noqa: ASYNC230
                out.write(contents_str)
    print(f"Files written to {temp_dir}")

    # walk the directory and ls
    for path in Path(temp_dir).iterdir():
        print(f"File: {path.name}")

    my_dir = await Dir.from_local(temp_dir)
    return my_dir
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/file_and_dir.py*

Finally, these tasks show how to use an offloaded type as an input.
Helper functions like `walk` and `open` have been added to the objects
and do what you might expect.

```
@env.task
async def check_dir(my_dir: Dir):
    print(f"Dir {my_dir.path} contents:")
    async for file in my_dir.walk():
        print(f"File: {file.name}")
        async with file.open("rb") as fh:
            contents = await fh.read()
            # Convert bytes to string
            contents_str = contents.decode('utf-8') if isinstance(contents, bytes) else str(contents)
            print(f"Contents: {contents_str}")

@env.task
async def create_and_check_dir():
    my_dir = await write_and_check_files()
    await check_dir(my_dir=my_dir)

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(create_and_check_dir)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/file_and_dir.py*

## Synchronous and asynchronous APIs

Every I/O operation on `File` and `Dir` comes in two forms, so you can use the offloaded types from both asynchronous and synchronous tasks:

- In an **asynchronous task** (`async def`), use the coroutine methods: `await` the upload, download, and existence calls, use `async with file.open(...)` to stream, and `async for` to walk a `Dir`. This is the pattern shown in **Tasks > Build tasks > Files and directories > Example usage** above.
- In a **synchronous task** (plain `def`), use the `_sync` variants: `File.from_local_sync()`, `file.open_sync()`, `dir.walk_sync()`, and so on. These block until the operation completes.

The two forms are otherwise equivalent — pick the one that matches how your task is defined. A few constructors do no I/O and so have a single form that is used unchanged from either kind of task: `File.new_remote()`, `File.from_existing_remote()`, `Dir.new_remote()`, and `Dir.from_existing_remote()`.

### File methods

| Asynchronous (in `async def` tasks) | Synchronous (in `def` tasks) | Purpose |
|---|---|---|
| `await File.from_local(path)` | `File.from_local_sync(path)` | Upload a local file to the blob store |
| `File.new_remote()` | `File.new_remote()` | Allocate a new remote file to stream into |
| `File.from_existing_remote(uri)` | `File.from_existing_remote(uri)` | Reference a file that already exists remotely |
| `async with file.open(mode) as fh` | `with file.open_sync(mode) as fh` | Open the file as a stream for reading or writing |
| `await file.download()` | `file.download_sync()` | Download the file to the local filesystem |
| `await file.exists()` | `file.exists_sync()` | Check whether the file exists |

### Dir methods

| Asynchronous (in `async def` tasks) | Synchronous (in `def` tasks) | Purpose |
|---|---|---|
| `await Dir.from_local(path)` | `Dir.from_local_sync(path)` | Upload a local directory to the blob store |
| `Dir.new_remote()` | `Dir.new_remote()` | Allocate a new remote directory to stream into |
| `Dir.from_existing_remote(uri)` | `Dir.from_existing_remote(uri)` | Reference a directory that already exists remotely |
| `async for f in dir.walk()` | `for f in dir.walk_sync()` | Iterate over the files in the directory |
| `await dir.list_files()` | `dir.list_files_sync()` | List the files in the directory (non-recursive) |
| `await dir.get_file(name)` | `dir.get_file_sync(name)` | Get a single file from the directory by name |
| `await dir.download()` | `dir.download_sync()` | Download the whole directory to the local filesystem |
| `await dir.exists()` | `dir.exists_sync()` | Check whether the directory exists |

> [!NOTE]
> `walk_sync()` additionally accepts a `file_pattern` glob (for example `file_pattern="*.txt"`) to filter the files it yields. Both forms accept `recursive` and `max_depth`.

### Synchronous example

The **Tasks > Build tasks > Files and directories > Example usage** uses `await`, `async with`, and `async for`. The same kind of workflow written with the synchronous API uses plain `def` tasks and the `_sync` method names:

```python
import flyte
from flyte.io import File

env = flyte.TaskEnvironment(name="sync-file")

@env.task
def write_file(content: str) -> File:
    # Allocate a new remote file and stream content into it
    f = File.new_remote()
    with f.open_sync("wb") as fh:
        fh.write(content.encode("utf-8"))
    return f

@env.task
def read_file(f: File) -> str:
    # Open the file for reading without downloading the whole object
    with f.open_sync("rb") as fh:
        return fh.read().decode("utf-8")

@env.task
def main() -> str:
    f = write_file(content="hello world")
    return read_file(f)
```

Directories work the same way — use `Dir.from_local_sync()` to upload and `walk_sync()` to iterate:

```python
import os
import tempfile

import flyte
from flyte.io import Dir

env = flyte.TaskEnvironment(name="sync-dir")

@env.task
def upload_dir() -> Dir:
    with tempfile.TemporaryDirectory() as tmp:
        for i in range(3):
            with open(os.path.join(tmp, f"file{i}.txt"), "w") as fh:
                fh.write(f"content {i}")
        # Upload the directory to the blob store
        return Dir.from_local_sync(tmp)

@env.task
def read_dir(d: Dir) -> int:
    count = 0
    # Walk and read every file, all synchronously
    for file in d.walk_sync(recursive=True):
        with file.open_sync("rb") as fh:
            print(f"{file.name}: {fh.read().decode('utf-8')}")
        count += 1
    return count

@env.task
def main() -> int:
    d = upload_dir()
    return read_dir(d)
```

## JSONL files

Flyte provides typed JSON Lines (JSONL) I/O through the `flyteplugins-jsonl` plugin, which extends `File` and `Dir` with the `JsonlFile` and `JsonlDir` types, adding streaming record-level read/write, optional zstd compression, and automatic shard rotation for large datasets.

See the [JSONL integration](https://www.union.ai/docs/latest/flyte/user-guide/integrations/jsonl/_index) guide for installation and usage.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/dataclasses-and-structures ===

# Data classes and structures

Dataclasses and Pydantic models are fully supported in Flyte as **materialized data types**:
Structured data where the full content is serialized and passed between tasks.
Use these as you would normally, passing them as inputs and outputs of tasks.

Unlike **offloaded types** like [`DataFrame`s](./dataframes), [`File`s and `Dir`s](./files-and-directories), data class and Pydantic model data is fully serialized, stored, and deserialized between tasks.
This makes them ideal for configuration objects, metadata, and smaller structured data where all fields should be serializable.

## Example: Combining dataclasses and Pydantic models

This example demonstrates how data classes and Pydantic models work together as materialized data types, showing nested structures and batch processing patterns:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "pydantic",
# ]
# main = "main"
# params = ""
# ///

import asyncio
from dataclasses import dataclass
from typing import List

from pydantic import BaseModel
import flyte

env = flyte.TaskEnvironment(name="ex-mixed-structures")

@dataclass
class InferenceRequest:
    feature_a: float
    feature_b: float

@dataclass
class BatchRequest:
    requests: List[InferenceRequest]
    batch_id: str = "default"

class PredictionSummary(BaseModel):
    predictions: List[float]
    average: float
    count: int
    batch_id: str

@env.task
async def predict_one(request: InferenceRequest) -> float:
    """
    A dummy linear model: prediction = 2 * feature_a + 3 * feature_b + bias(=1.0)
    """
    return 2.0 * request.feature_a + 3.0 * request.feature_b + 1.0

@env.task
async def process_batch(batch: BatchRequest) -> PredictionSummary:
    """
    Processes a batch of inference requests and returns summary statistics.
    """
    # Process all requests concurrently
    tasks = [predict_one(request=req) for req in batch.requests]
    predictions = await asyncio.gather(*tasks)

    # Calculate statistics
    average = sum(predictions) / len(predictions) if predictions else 0.0

    return PredictionSummary(
        predictions=predictions,
        average=average,
        count=len(predictions),
        batch_id=batch.batch_id
    )

@env.task
async def summarize_results(summary: PredictionSummary) -> str:
    """
    Creates a text summary from the prediction results.
    """
    return (
        f"Batch {summary.batch_id}: "
        f"Processed {summary.count} predictions, "
        f"average value: {summary.average:.2f}"
    )

@env.task
async def main() -> str:
    batch = BatchRequest(
        requests=[
            InferenceRequest(feature_a=1.0, feature_b=2.0),
            InferenceRequest(feature_a=3.0, feature_b=4.0),
            InferenceRequest(feature_a=5.0, feature_b=6.0),
        ],
        batch_id="demo_batch_001"
    )
    summary = await process_batch(batch)
    result = await summarize_results(summary)
    return result

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataclasses-and-structures/example.py*

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/dataframes ===

# DataFrames

By default, return values in Python are materialized - meaning the actual data is downloaded and loaded into memory. This applies to simple types like integers, as well as more complex types like DataFrames.

To avoid downloading large datasets into memory, Flyte V2 exposes [`flyte.io.dataframe`](https://www.union.ai/docs/latest/flyte/api-reference/flyte-sdk/flyte.io/dataframe): a thin,  uniform wrapper type for DataFrame-style objects that allows you to pass a reference to the data, rather than the fully materialized contents.

The `flyte.io.DataFrame` type provides serialization support for common engines like `pandas`, `polars`, `pyarrow`, `dask`, etc.; enabling you to move data between different DataFrame backends.

DataFrame contents are written to the data plane object store and passed between tasks by reference. For the full map of what goes in the bucket versus what stays in the control plane database, see [Where your data lives](https://www.union.ai/docs/latest/flyte/user-guide/get-started/core-concepts/where-data-lives).

> [!NOTE]
> Because a DataFrame is passed by reference, a downstream cached task does not get a cache hit on identical content stored at a new path. To cache on content, attach a hash with `flyte.io.HashFunction` - see [Content-based caching for DataFrames, files, and directories](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/caching).

## Setting up the environment and sample data

For our example we will start by setting up our task environment with the required dependencies and create some sample data.

```
from typing import Annotated

import numpy as np
import pandas as pd
import flyte
import flyte.io

env = flyte.TaskEnvironment(
    "dataframe_usage",
    image= flyte.Image.from_debian_base().with_pip_packages("pandas", "pyarrow", "numpy"),
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)

BASIC_EMPLOYEE_DATA = {
    "employee_id": range(1001, 1009),
    "name": ["Alice", "Bob", "Charlie", "Diana", "Ethan", "Fiona", "George", "Hannah"],
    "department": ["HR", "Engineering", "Engineering", "Marketing", "Finance", "Finance", "HR", "Engineering"],
    "hire_date": pd.to_datetime(
        ["2018-01-15", "2019-03-22", "2020-07-10", "2017-11-01", "2021-06-05", "2018-09-13", "2022-01-07", "2020-12-30"]
    ),
}

ADDL_EMPLOYEE_DATA = {
    "employee_id": range(1001, 1009),
    "salary": [55000, 75000, 72000, 50000, 68000, 70000, np.nan, 80000],
    "bonus_pct": [0.05, 0.10, 0.07, 0.04, np.nan, 0.08, 0.03, 0.09],
    "full_time": [True, True, True, False, True, True, False, True],
    "projects": [
        ["Recruiting", "Onboarding"],
        ["Platform", "API"],
        ["API", "Data Pipeline"],
        ["SEO", "Ads"],
        ["Budget", "Forecasting"],
        ["Auditing"],
        [],
        ["Platform", "Security", "Data Pipeline"],
    ],
}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/dataframes.py*

## Create a raw DataFrame

Now, let's create a task that returns a native Pandas DataFrame:

```
@env.task
async def create_raw_dataframe() -> pd.DataFrame:
    return pd.DataFrame(BASIC_EMPLOYEE_DATA)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/dataframes.py*

This is the most basic use-case of how to pass DataFrames (of all kinds, not just Pandas).
We simply create the DataFrame as normal, and return it.

Because the task has been declared to return a supported native DataFrame type (in this case `pandas.DataFrame` Flyte will automatically detect it, serialize it correctly and upload it at task completion enabling it to be passed transparently to the next task.

Flyte supports auto-serialization for the following DataFrame types:

* `pandas.DataFrame`
* `pyarrow.Table`
* `dask.dataframe.DataFrame`
* `polars.DataFrame`
* `flyte.io.DataFrame` (see below)

## Create a flyte.io.DataFrame

Alternatively you can also create a `flyte.io.DataFrame` object directly from a native object with the `from_df` method:

```
@env.task
async def create_flyte_dataframe() -> Annotated[flyte.io.DataFrame, "parquet"]:
    pd_df = pd.DataFrame(ADDL_EMPLOYEE_DATA)
    fdf = flyte.io.DataFrame.from_df(pd_df)
    return fdf
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/dataframes.py*

The `flyte.io.DataFrame` class creates a thin wrapper around objects of any standard DataFrame type. It serves as a generic "any DataFrame type" (a concept that Python itself does not currently offer).

As with native DataFrame types, Flyte will automatically serialize and upload the data at task completion.

The advantage of the unified `flyte.io.DataFrame` wrapper is that you can be explicit about the storage format that makes sense for your use case, by using an `Annotated` type where the second argument encodes format or other lightweight hints. For example, here we specify that the DataFrame should be stored as Parquet:

## Automatically convert between types

You can use Flyte to automatically download and convert the DataFrame between types when needed:

```
@env.task
async def join_data(raw_dataframe: pd.DataFrame, flyte_dataframe: pd.DataFrame) -> flyte.io.DataFrame:
    joined_df = raw_dataframe.merge(flyte_dataframe, on="employee_id", how="inner")
    return flyte.io.DataFrame.from_df(joined_df)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/dataframes.py*

This task takes two DataFrames as input. We'll pass one raw Pandas DataFrame, and one `flyte.io.DataFrame`.
Flyte automatically converts the `flyte.io.DataFrame` to a Pandas DataFrame (since we declared that as the input type) before passing it to the task.
The actual download and conversion happens only when we access the data, in this case, when we do the merge.

## Downloading DataFrames

When a task receives a `flyte.io.DataFrame`, you can request a concrete backend representation. For example, to download as a pandas DataFrame:

```
@env.task
async def download_data(joined_df: flyte.io.DataFrame):
    downloaded = await joined_df.open(pd.DataFrame).all()
    print("Downloaded Data:\n", downloaded)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/dataframes.py*

The `open()` call delegates to the DataFrame handler for the stored format and converts to the requested in-memory type.

## Run the example

Finally, we can define a `main` function to run the tasks defined above and a `__main__` block to execute the workflow:

```
@env.task
async def main():
    raw_df = await create_raw_dataframe ()
    flyte_df = await create_flyte_dataframe ()
    joined_df = await join_data (raw_df, flyte_df)
    await download_data (joined_df)

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/dataframes.py*

## Polars DataFrames

The `flyteplugins-polars` package extends Flyte's DataFrame support to `polars.DataFrame` and `polars.LazyFrame`. Install it alongside the core SDK and it registers automatically. No additional configuration required.

```bash
pip install flyteplugins-polars
```

Both types are serialized as Parquet when passed between tasks, just like other DataFrame backends.

### Setup

```
import polars as pl

import flyte

env = flyte.TaskEnvironment(
    name="polars-dataframes",
    image=flyte.Image.from_debian_base(name="polars").with_pip_packages(
        "flyteplugins-polars>=2.0.0", "polars"
    ),
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)

EMPLOYEE_DATA = {
    "employee_id": [1001, 1002, 1003, 1004, 1005, 1006],
    "name": ["Alice", "Bob", "Charlie", "Diana", "Ethan", "Fiona"],
    "department": ["Engineering", "Engineering", "Marketing", "Finance", "Finance", "Engineering"],
    "salary": [75000, 72000, 50000, 68000, 70000, 80000],
    "years_experience": [5, 4, 2, 6, 5, 7],
}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/polars_dataframes.py*

### Eager DataFrames

Use `pl.DataFrame` when you want immediate evaluation. Flyte serializes it to Parquet on output and deserializes it on input:

```
@env.task
async def create_dataframe() -> pl.DataFrame:
    """Create a Polars DataFrame.

    Polars DataFrames are passed between tasks as serialized Parquet files
    stored in the Flyte blob store — no manual upload required.
    """
    return pl.DataFrame(EMPLOYEE_DATA)

@env.task
async def filter_high_earners(df: pl.DataFrame) -> pl.DataFrame:
    """Filter and enrich a Polars DataFrame."""
    return (
        df.filter(pl.col("salary") > 60000)
        .with_columns(
            (pl.col("salary") / pl.col("years_experience")).alias("salary_per_year")
        )
        .sort("salary", descending=True)
    )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/polars_dataframes.py*

### Lazy DataFrames

Use `pl.LazyFrame` when you want to defer computation and let Polars optimize the full query plan before executing. Flyte handles serialization the same way as `pl.DataFrame`:

```
@env.task
async def create_lazyframe() -> pl.LazyFrame:
    """Create a Polars LazyFrame.

    LazyFrames defer computation until collected, allowing Polars to
    optimize the full query plan. They are serialized to Parquet just
    like DataFrames when passed between tasks.
    """
    return pl.LazyFrame(EMPLOYEE_DATA)

@env.task
async def aggregate_by_department(lf: pl.LazyFrame) -> pl.DataFrame:
    """Aggregate salary statistics by department using a LazyFrame.

    The query plan is built lazily and executed only when collect() is called.
    """
    return (
        lf.group_by("department")
        .agg(
            pl.col("salary").mean().alias("avg_salary"),
            pl.col("salary").max().alias("max_salary"),
            pl.len().alias("headcount"),
        )
        .sort("avg_salary", descending=True)
        .collect()
    )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/polars_dataframes.py*

The `collect()` call in `aggregate_by_department` is what triggers execution of the lazy plan. The `LazyFrame` passed between tasks is serialized as Parquet at that point.

### Run the example

```
@env.task
async def main():
    df = await create_dataframe()
    filtered = await filter_high_earners(df=df)
    print("High earners:")
    print(filtered)

    lf = await create_lazyframe()
    summary = await aggregate_by_department(lf=lf)
    print("Department summary:")
    print(summary)

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/dataframes/polars_dataframes.py*

## See also

To display a DataFrame as an HTML table in a task report, define a `flyte.types.Renderable` for it — see [Rendering a custom type](./reports#rendering-a-custom-type) on the Reports page.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/handling-custom-types ===

# Custom types

Flyte has a rich type system that handles most Python types automatically. However, there are cases where you may want to pass custom types into a run or between actions. By default, if Flyte doesn't recognize a type, it uses Python pickle to serialize the data. While this works, pickle has several drawbacks:

- **Inefficiency**: Pickle can be very inefficient for certain data types
- **Language compatibility**: Pickle is Python-specific and doesn't work with other languages
- **Version fragility**: Pickled data can break between Python versions
- **Opacity**: Pickled data appears as bytes or file links in the UI, with no automatic form generation

Consider types like Polars DataFrames or PyTorch Tensors. Using pickle for these is extremely inefficient compared to native serialization formats like Parquet or tensor-specific formats.

Flyte SDK addresses this by allowing you to create and share type extensions.

## Types of extensions

Flyte supports two types of type extensions:

1. **Type transformers**: For scalar types (integers, strings, files, directories, custom objects)
2. **DataFrame extensions**: For tabular data types that benefit from DataFrame-specific handling

DataFrame types are special because they have associated metadata (columns, schemas), can be serialized to efficient formats like Parquet, support parallel uploads from engines like Spark, and can be partitioned.

## Creating a type transformer

Type transformers convert between Python types and Flyte's internal representation. Here's how to create one for a custom `PositiveInt` type.

### Step 1: Define your custom type

```python
# custom_type.py
class PositiveInt:
    """A wrapper type that only accepts positive integers."""

    def __init__(self, value: int):
        if not isinstance(value, int):
            raise TypeError(f"Expected int, got {type(value).__name__}")
        if value <= 0:
            raise ValueError(f"Expected positive integer, got {value}")
        self._value = value

    @property
    def value(self) -> int:
        return self._value

    def __repr__(self) -> str:
        return f"PositiveInt({self._value})"
```

### Step 2: Create the type transformer

```python
# transformer.py
from typing import Type

from flyteidl2.core import literals_pb2, types_pb2

from flyte import logger
from flyte.types import TypeEngine, TypeTransformer, TypeTransformerFailedError
from my_transformer.custom_type import PositiveInt

class PositiveIntTransformer(TypeTransformer[PositiveInt]):
    """
    Type transformer for PositiveInt that validates and transforms positive integers.
    """

    def __init__(self):
        super().__init__(name="PositiveInt", t=PositiveInt)

    def get_literal_type(self, t: Type[PositiveInt]) -> types_pb2.LiteralType:
        """Returns the Flyte literal type for PositiveInt."""
        return types_pb2.LiteralType(
            simple=types_pb2.SimpleType.INTEGER,
            structure=types_pb2.TypeStructure(tag="PositiveInt"),
        )

    async def to_literal(
        self,
        python_val: PositiveInt,
        python_type: Type[PositiveInt],
        expected: types_pb2.LiteralType,
    ) -> literals_pb2.Literal:
        """Converts a PositiveInt instance to a Flyte Literal."""
        if not isinstance(python_val, PositiveInt):
            raise TypeTransformerFailedError(
                f"Expected PositiveInt, got {type(python_val).__name__}"
            )

        return literals_pb2.Literal(
            scalar=literals_pb2.Scalar(
                primitive=literals_pb2.Primitive(integer=python_val.value)
            )
        )

    async def to_python_value(
        self,
        lv: literals_pb2.Literal,
        expected_python_type: Type[PositiveInt]
    ) -> PositiveInt:
        """Converts a Flyte Literal back to a PositiveInt instance."""
        if not lv.scalar or not lv.scalar.primitive:
            raise TypeTransformerFailedError(
                f"Cannot convert literal {lv} to PositiveInt: missing scalar primitive"
            )

        value = lv.scalar.primitive.integer
        try:
            return PositiveInt(value)
        except (TypeError, ValueError) as e:
            raise TypeTransformerFailedError(
                f"Cannot convert value {value} to PositiveInt: {e}"
            )

    def guess_python_type(
        self,
        literal_type: types_pb2.LiteralType
    ) -> Type[PositiveInt]:
        """Guesses the Python type from a Flyte literal type."""
        if (
            literal_type.simple == types_pb2.SimpleType.INTEGER
            and literal_type.structure
            and literal_type.structure.tag == "PositiveInt"
        ):
            return PositiveInt
        raise ValueError(f"Cannot guess PositiveInt from literal type {literal_type}")
```

### Step 3: Register the transformer

Create a registration function that can be called to register your transformer:

```python
def register_positive_int_transformer():
    """Register the PositiveIntTransformer in the TypeEngine."""
    TypeEngine.register(PositiveIntTransformer())
    logger.info("Registered PositiveIntTransformer in TypeEngine")
```

## Distributing type plugins

To share your type transformer as an installable package, configure it as a Flyte plugin using entry points.

### Configure pyproject.toml

Add the entry point to your `pyproject.toml`:

```toml
[project]
name = "my_transformer"
version = "0.1.0"
description = "Custom type transformer"
requires-python = ">=3.10"
dependencies = []

[project.entry-points."flyte.plugins.types"]
my_transformer = "my_transformer.transformer:register_positive_int_transformer"
```

The entry point group `flyte.plugins.types` tells Flyte to automatically load this transformer when the package is installed.

### Automatic loading

When your plugin package is installed, Flyte automatically loads the type transformer at runtime. This happens during `flyte.init()` or `flyte.init_from_config()`.

## Controlling plugin loading

Loading many type plugins can add overhead to initialization. You can disable automatic plugin loading:

```python
import flyte

# Disable automatic loading of type transformer plugins
flyte.init(load_plugin_type_transformers=False)
```

By default, `load_plugin_type_transformers` is `True`.

## Using custom types in tasks

Once registered, use your custom type like any built-in type:

```python
import flyte
from my_transformer.custom_type import PositiveInt

env = flyte.TaskEnvironment(name="custom_types")

@env.task
async def process_positive(value: PositiveInt) -> int:
    """Process a positive integer."""
    return value.value * 2

if __name__ == "__main__":
    flyte.init_from_config()

    # The custom type works seamlessly
    run = flyte.run(process_positive, value=PositiveInt(42))
    run.wait()
    print(run.outputs()[0])  # 84
```

## DataFrame extensions

For tabular data types, Flyte provides a specialized extension mechanism through `flyte.io.DataFrame`. DataFrame extensions support:

- Automatic conversion to/from Parquet format
- Column metadata and schema information
- Parallel uploads from distributed engines
- Partitioning support

DataFrame extensions use encoders and decoders from `flyte.io.extend`. Documentation for creating DataFrame extensions is coming soon.

## Best practices

1. **Use specific types over pickle**: Define type transformers for any custom types used frequently in your workflows
2. **Keep transformers lightweight**: Avoid expensive operations in `to_literal` and `to_python_value`
3. **Add validation**: Validate data in your transformer to catch errors early
4. **Use meaningful tags**: The `TypeStructure.tag` helps identify your type in the Flyte UI
5. **Be judicious with plugins**: Only install the plugins you need to minimize initialization overhead

## See also

To render a custom type as HTML in a task report, define a `flyte.types.Renderable` for it — see [Rendering a custom type](./reports#rendering-a-custom-type) on the Reports page.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/custom-context ===

# Custom context

Custom context provides a mechanism for implicitly passing configuration and metadata through your entire task execution hierarchy without adding parameters to every task. It is ideal for cross-cutting concerns such as tracing, environment metadata, or experiment identifiers.

Think of custom context as **execution-scoped metadata** that automatically flows from parent to child tasks.

> **📝 Note**
>
> In Flyte 1 the runtime context was accessed with `current_context()` (the `flytekit.current_context` API). Flyte 2 uses `flyte.ctx()`, shown below.

## Overview

Custom context is an implicit key-value configuration map that is automatically available to tasks during execution. It is stored in the blob store of your Union/Flyte instance together with the task’s inputs, making it available across tasks without needing to pass it explicitly.

You can access it in a Flyte task via:

```python
flyte.ctx().custom_context
```

Custom context is fundamentally different from standard task inputs. Task inputs are explicit, strongly typed parameters that you declare as part of a task’s signature. They directly influence the task’s computation and therefore participate in Flyte’s caching and reproducibility guarantees.

Custom context, on the other hand, is implicit metadata. It consists only of string key/value pairs, is not part of the task signature, and does not affect task caching. Because it is injected by the Flyte runtime rather than passed as a formal input, it should be used only for environmental or contextual information, not for data that changes the logical output of a task.

## When to use it and when not to

Custom context is perfect when you need metadata, not domain data, to flow through your tasks.

Good use cases:

- Tracing IDs, span IDs
- Experiment or run metadata
- Environment region, cluster ID
- Logging correlation keys
- Feature flags
- Session IDs for 3rd-party APIs (e.g., an LLM session)

Avoid using for:

- Business/domain data
- Inputs that change task outputs
- Anything affecting caching or reproducibility
- Large blobs of data (keep it small)

It is the cleanest mechanism when you need something available everywhere, but not logically an input to the computation.

## Setting custom context

There are two ways to set custom context for a Flyte run:

1. Set it once for the entire run when you launch (`with_runcontext`): this establishes the base context for the execution
2. Set or override it inside task code using `flyte.custom_context(...)` context manager: this changes the active context for that task block and any nested tasks called from it

Both are legitimate and complementary. The important behavioral rules to understand are:

- `with_runcontext(...)` sets the run-level base. Values provided here are available everywhere unless overridden later. Use this for metadata that should apply to most or all tasks in the run (experiment name, top-level trace id, run id, etc.).
- `flyte.custom_context(...)` is used inside task code to set or override values for that scope. It does affect nested tasks invoked while that context is active. In practice this means you can override run-level entries, add new keys for downstream tasks, or both.
- Merging & precedence: contexts are merged; when the same key appears in multiple places the most recent/innermost value wins (i.e., values set by `flyte.custom_context(...)` override the run-level values from `with_runcontext(...)` for the duration of that block).

### Run-level context

Set base metadata once when starting the run:

```
import flyte

env = flyte.TaskEnvironment("custom-context-example")

@env.task
async def leaf_task() -> str:
    # Reads run-level context
    print("leaf sees:", flyte.ctx().custom_context)
    return flyte.ctx().custom_context.get("trace_id")

@env.task
async def root() -> str:
    return await leaf_task()

if __name__ == "__main__":
    flyte.init_from_config()
    # Base context for the entire run
    flyte.with_runcontext(custom_context={"trace_id": "root-abc", "experiment": "v1"}).run(root)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/custom-context/run_context.py*

Output (every task sees the base keys unless overridden):

```bash
leaf sees: {"trace_id": "root-abc", "experiment": "v1"}
```

### Overriding inside a task (local override that affects nested tasks)

Use `flyte.custom_context(...)` inside a task to override or add keys for downstream calls:

```
@env.task
async def downstream() -> str:
    print("downstream sees:", flyte.ctx().custom_context)
    return flyte.ctx().custom_context.get("trace_id")

@env.task
async def parent() -> str:
    print("parent initial:", flyte.ctx().custom_context)

    # Override the trace_id for the nested call(s)
    with flyte.custom_context(trace_id="child-override"):
        val = await downstream()     # downstream sees trace_id="child-override"

    # After the context block, run-level values are back
    print("parent after:", flyte.ctx().custom_context)
    return val
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/custom-context/override_context.py*

If the run was started with `{"trace_id": "root-abc"}`, this prints:

```bash
parent initial: {"trace_id": "root-abc"}
downstream sees: {"trace_id": "child-override"}
parent after: {"trace_id": "root-abc"}
```

Note that the override affected the nested downstream task because it was invoked while the `flyte.custom_context` block was active.

### Adding new keys for nested tasks

You can add keys (not just override):

```python
with flyte.custom_context(experiment="exp-blue", run_group="g-7"):
    await some_task()   # some_task sees both base keys + the new keys
```

## Accessing custom context

Always via the Flyte runtime:

```python
ctx = flyte.ctx().custom_context
value = ctx.get("key")
```

You can access the custom context using either `flyte.ctx().custom_context` or the shorthand `flyte.get_custom_context()`, which returns the same dictionary of key/value pairs.

Values are always strings, so parse as needed:

```python
timeout = int(ctx["timeout_seconds"])
```

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/abort-tasks ===

# Abort and cancel actions

When running complex workflows, you may need to stop actions that are no longer needed.
This can happen when one branch of your workflow makes others redundant, when a task fails and its siblings should not continue, or when you need to manually intervene in a running workflow.

Flyte provides three mechanisms for stopping actions:

- **Automatic cleanup**: When a root action completes, all its in-progress descendant actions are automatically aborted.
- **Programmatic cancellation**: Cancel specific `asyncio` tasks from within your workflow code.
- **External abort**: Stop individual actions via the CLI, the UI, or the API.

For background on runs and actions, see [Runs and actions](https://www.union.ai/docs/latest/flyte/user-guide/get-started/core-concepts/runs-and-actions).

## Action lifetime

The lifetime of all actions in a [run](https://www.union.ai/docs/latest/flyte/user-guide/get-started/core-concepts/runs-and-actions) is tied to the lifetime of the root action (the first task that was invoked).
When the root action exits (whether it succeeds, fails, or returns early) all in-progress descendant actions are automatically aborted and no new actions can be enqueued.

This means you don't need to manually clean up child actions. Flyte handles it for you.

Consider this example where `main` exits after 10 seconds, but it has spawned a `sleep_for` action that is set to run for 30 seconds:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# main = "main"
# params = "seconds = 30"
# ///

import asyncio

import flyte

env = flyte.TaskEnvironment(name="action_lifetime")

@env.task
async def do_something():
    print("Doing something")
    await asyncio.sleep(5)
    print("Finished doing something")

@env.task
async def sleep_for(seconds: int):
    print(f"Sleeping for {seconds} seconds")
    try:
        await asyncio.sleep(seconds)
        await do_something()
    except asyncio.CancelledError:
        print("sleep_for was cancelled")
        return
    print(f"Finished sleeping for {seconds} seconds")

@env.task
async def main(seconds: int):
    print("Starting main")
    asyncio.create_task(sleep_for(seconds))
    await asyncio.sleep(10)
    print("Main finished")

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, seconds=30)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/abort-tasks/action_lifetime.py*

When `main` returns after 10 seconds, the `sleep_for` action (which still has 20 seconds remaining) is automatically aborted.
The `sleep_for` task receives an `asyncio.CancelledError`, giving it a chance to handle the cancellation gracefully.

## Canceling actions programmatically

As a workflow author, you can cancel specific in-progress actions by canceling their corresponding `asyncio` tasks.
This is useful in scenarios like hyperparameter optimization (HPO), where one action converges to the desired result and the remaining actions can be stopped to save compute.

To cancel actions programmatically:

1. Launch actions using `asyncio.create_task()` and retain references to the returned task objects.
2. When the desired condition is met, call `.cancel()` on the tasks you want to stop.

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# main = "main"
# params = "n = 30, f = 10.0"
# ///

import asyncio

import flyte
import flyte.errors

env = flyte.TaskEnvironment("cancel")

@env.task
async def sleepers(f: float, n: int):
    await asyncio.sleep(f)

@env.task
async def failing_task(f: float):
    raise ValueError("I will fail!")

@env.task
async def main(n: int, f: float):
    sleeping_tasks = []
    for i in range(n):
        sleeping_tasks.append(asyncio.create_task(sleepers(f, i)))

    await asyncio.sleep(f)
    try:
        await failing_task(f)
        await asyncio.gather(*sleeping_tasks)
    except flyte.errors.RuntimeUserError as e:
        if e.code == "ValueError":
            print(f"Received ValueError, canceling {len(sleeping_tasks)} sleeping tasks")
            for t in sleeping_tasks:
                t.cancel()
        return

if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(main, 30, 10.0))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/abort-tasks/cancel_tasks.py*

In this code:

* The `main` task launches 30 `sleepers` actions in parallel using `asyncio.create_task()`.
* It then calls `failing_task`, which raises a `ValueError`.
* The error is caught as a `flyte.errors.RuntimeUserError` (since user-raised exceptions are wrapped by Flyte).
* On catching the error, `main` cancels all sleeping tasks by calling `.cancel()` on each one, freeing their compute resources.

This pattern lets you react to runtime conditions and stop unnecessary work. For more on handling errors within workflows, see [Error handling](./error-handling).

## External abort

Sometimes you need to stop an action manually, outside the workflow code itself. You can abort individual actions using the CLI, the UI, or the API.

When an action is externally aborted, the parent action that awaits it receives a `flyte.errors.ActionAbortedError`. You can catch this error to handle the abort gracefully.

### Aborting via the CLI

To abort a specific action:

```bash
flyte abort <run-name> <action-name>
```

Use `--project` and `--domain` to target a specific [project-domain pair](https://www.union.ai/docs/latest/flyte/user-guide/get-started/core-concepts/projects-and-domains).
For all available options, see the [CLI reference](https://www.union.ai/docs/latest/flyte/api-reference/flyte-cli).

### Handling external aborts

When using `asyncio.gather()` with `return_exceptions=True`, externally aborted actions return an `ActionAbortedError` instead of raising it. This lets you inspect results and handle aborts on a per-action basis:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# main = "main"
# params = "n = 10, sleep_for = 30.0"
# ///

import asyncio

import flyte
import flyte.errors

env = flyte.TaskEnvironment("external_abort")

@env.task
async def long_sleeper(sleep_for: float):
    await asyncio.sleep(sleep_for)

@env.task
async def main(n: int, sleep_for: float) -> str:
    coros = [long_sleeper(sleep_for) for _ in range(n)]
    results = await asyncio.gather(*coros, return_exceptions=True)
    for i, r in enumerate(results):
        if isinstance(r, flyte.errors.ActionAbortedError):
            print(f"Action [{i}] was externally aborted")
    return "Hello World!"

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, 10, 30.0)
    print(run.url)
    run.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/abort-tasks/external_abort.py*

In this code:

* The `main` task launches 10 `long_sleeper` actions in parallel.
* If any action is externally aborted (via the CLI, the UI, or the API) while running, `asyncio.gather` captures the `ActionAbortedError` as a result instead of propagating it.
* The `main` task iterates over the results and logs which actions were aborted.
* Because the abort is handled, `main` can continue executing and return its result normally.

Without `return_exceptions=True`, an external abort would raise `ActionAbortedError` directly, which you can handle with a standard `try...except` block.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/container-tasks ===

Container tasks are one of Flyte's superpowers. They allow you to execute tasks using any container image without requiring the Flyte SDK to be installed in that container. This means you can run code written in any language, execute shell scripts, or even use pre-built containers pulled directly from the internet while still maintaining Flyte's data orchestration capabilities.

## What are container tasks?

A container task is a special type of Flyte task that executes arbitrary container images. Unlike standard `@task` decorated functions that require the Flyte SDK, container tasks can run:

- Code written in any programming language (Rust, Go, Java, R, etc.)
- Legacy containers with unsupported Python versions
- Pre-built bioinformatics or scientific computing containers
- Shell scripts and command-line tools
- Dynamically generated code in sandboxed environments

## How data flows in and out

The magic of container tasks lies in Flyte's **copilot sidecar system**. When you execute a container task, Flyte:

1. Launches your specified container alongside a copilot sidecar container
2. Uses shared Kubernetes pod volumes to pass data between containers
3. Reads inputs from `input_data_dir` and writes outputs to `output_data_dir`
4. Automatically handles serialization and deserialization of typed data

This means you can construct workflows where some tasks are container tasks while others are Python functions, and data will flow between them.

## Basic usage

Here's a simple example that runs a shell command in an Alpine container:

```python
import flyte
from flyte.extras import ContainerTask

greeting_task = ContainerTask(
    name="echo_and_return_greeting",
    image=flyte.Image.from_base("alpine:3.18"),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"name": str},
    outputs={"greeting": str},
    command=[
        "/bin/sh",
        "-c",
        "echo 'Hello, my name is {{.inputs.name}}.' | tee -a /var/outputs/greeting"
    ],
)
```

### Template syntax for inputs

Container tasks support template-style references to inputs using the syntax `{{.inputs.<input_name>}}`. This gets replaced with the actual input value at runtime:

```python
command=["/bin/sh", "-c", "echo 'Processing {{.inputs.user_id}}' > /var/outputs/result"]
```

### Using container tasks in workflows

Container tasks integrate with Python tasks:

```python
container_env = flyte.TaskEnvironment.from_task("container_env", greeting_task)
env = flyte.TaskEnvironment(name="hello_world", depends_on=[container_env])

@env.task
async def say_hello(name: str = "flyte") -> str:
    print("Hello container task")
    return await greeting_task(name=name)
```

## Advanced: Passing files and directories

Container tasks can accept `File` and `Dir` inputs. For these types, use path-based syntax (not template syntax) in your commands:

```python
from flyte.io import File
import pathlib

code_runner = ContainerTask(
    name="python_code_runner",
    image="ghcr.io/astral-sh/uv:debian-slim",
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"script.py": File, "a": int, "b": int},
    outputs={"result": int},
    command=[
        "/bin/sh",
        "-c",
        "uv run /var/inputs/script.py {{.inputs.a}} {{.inputs.b}} > /var/outputs/result"
    ],
)

@env.task
async def execute_script() -> int:
    path = pathlib.Path(__file__).parent / "my_script.py"
    script_file = await File.from_local(path)
    return await code_runner(**{"script.py": script_file, "a": 10, "b": 20})
```

Note that when passing files, the input key can include the filename (e.g., `"script.py"`), and you reference it in the command as `/var/inputs/script.py`.

## Use case: Agentic sandbox execution

Container tasks are perfect for running AI-generated code in isolated environments. You can generate a data analysis script dynamically and execute it safely:

```python
import flyte
from flyte.extras import ContainerTask
from flyte.io import File
import pathlib

env = flyte.TaskEnvironment(name="agentic_sandbox")

@env.task
async def run_generated_code(script_content: str, param_a: int, param_b: int) -> int:
    # Define a container task that runs arbitrary Python code
    sandbox = ContainerTask(
        name="code_sandbox",
        image="ghcr.io/astral-sh/uv:debian-slim",
        input_data_dir="/var/inputs",
        output_data_dir="/var/outputs",
        inputs={"script": File, "a": int, "b": int},
        outputs={"result": int},
        command=[
            "/bin/sh",
            "-c",
            "uv run --script /var/inputs/script {{.inputs.a}} {{.inputs.b}} > /var/outputs/result"
        ],
    )

    # Save the generated script to a temporary file
    temp_path = pathlib.Path("/tmp/generated_script.py")
    temp_path.write_text(script_content)

    # Execute it in the sandbox
    script_file = await File.from_local(temp_path)
    return await sandbox(script=script_file, a=param_a, b=param_b)
```

This pattern allows you to:

- Generate code using LLMs or other AI systems
- Execute it in a controlled, isolated environment
- Capture results and integrate them back into your workflow
- Maintain full observability and reproducibility

## Use case: Legacy and specialized containers

Many scientific and bioinformatics tools are distributed as pre-built containers. Container tasks let you integrate them directly:

```python
# Run a bioinformatics tool
blast_task = ContainerTask(
    name="run_blast",
    image="ncbi/blast:latest",
    input_data_dir="/data",
    output_data_dir="/results",
    inputs={"query": File, "database": str},
    outputs={"alignments": File},
    command=[
        "blastn",
        "-query", "/data/query",
        "-db", "{{.inputs.database}}",
        "-out", "/results/alignments",
        "-outfmt", "6"
    ],
)

# Run legacy code with an old Python version
legacy_task = ContainerTask(
    name="legacy_python",
    image="python:2.7",  # Unsupported Python version
    input_data_dir="/app/inputs",
    output_data_dir="/app/outputs",
    inputs={"data_file": File},
    outputs={"processed": File},
    command=[
        "python",
        "/legacy_app/process.py",
        "/app/inputs/data_file",
        "/app/outputs/processed"
    ],
)
```

## Use case: Multi-language workflows

Build workflows that span multiple languages:

```python
# Rust task for high-performance computation
rust_task = ContainerTask(
    name="rust_compute",
    image="rust:1.75",
    inputs={"n": int},
    outputs={"result": int},
    input_data_dir="/inputs",
    output_data_dir="/outputs",
    command=["./compute_binary", "{{.inputs.n}}"],
)

# Python task for orchestration
@env.task
async def multi_lang_workflow(iterations: int) -> dict:
    # Call Rust task for heavy computation
    computed = await rust_task(n=iterations)

    # Process results in Python
    processed = await python_analysis_task(computed)

    return {"rust_result": computed, "analysis": processed}
```

## Configuration options

### ContainerTask parameters

- **name**: Unique identifier for the task
- **image**: Container image to use (string or `Image` object)
- **command**: Command to execute in the container (list of strings)
- **inputs**: Dictionary mapping input names to types
- **outputs**: Dictionary mapping output names to types
- **input_data_dir**: Directory where Flyte writes input data (default: `/var/inputs`)
- **output_data_dir**: Directory where Flyte reads output data (default: `/var/outputs`)
- **arguments**: Additional command arguments (list of strings)
- **metadata_format**: Format for metadata serialization (`"JSON"`, `"YAML"`, or `"PROTO"`)
- **local_logs**: Whether to print container logs during local execution (default: `True`)

### Supported input/output types

Container tasks support all standard Flyte types:

- Primitives: `str`, `int`, `float`, `bool`
- Temporal: `datetime.datetime`, `datetime.timedelta`
- File system: `File`, `Dir`
- Complex types: dataclasses, Pydantic models (serialized as JSON/YAML/PROTO)

## Best practices

1. **Use specific image tags**: Prefer `alpine:3.18` over `alpine:latest` for reproducibility
2. **Keep containers focused**: Each container task should do one thing well
3. **Handle errors gracefully**: Ensure your container commands exit with appropriate status codes
4. **Test locally first**: Container tasks can run locally with Docker, making debugging easier
5. **Consider image size**: Smaller images lead to faster task startup times
6. **Document input/output contracts**: Clearly specify what data flows in and out

## Local execution

Container tasks require Docker to be installed and running on your local machine. When you run them locally, Flyte will:

1. Pull the specified image (if not already available)
2. Mount local directories for inputs and outputs
3. Stream container logs to your console
4. Extract outputs after container completion

This makes it easy to develop and test container tasks before deploying to a remote cluster.

## When to use container tasks

Choose container tasks when you need to:

- Run code in languages other than Python
- Execute pre-built tools or legacy applications
- Isolate potentially unsafe code (AI-generated scripts)
- Use specific runtime environments or dependencies
- Integrate external tools without Python wrappers
- Execute shell scripts or command-line utilities

For Python code that can use the Flyte SDK, standard `@task` decorated functions are usually simpler and more efficient.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/links ===

# Links

Links let you add clickable URLs to tasks that appear in the Flyte UI. Use them to connect tasks to external tools like experiment trackers, monitoring dashboards or custom internal services.

![Links in the Flyte UI](https://www.union.ai/docs/latest/flyte/_static/images/integrations/wandb/single_node_auto_flyte.png)

You can attach links to tasks in two ways:

- **Statically** in the task decorator with `links=`
- **Dynamically** at call time with `task.override(links=...)`

`Link` is a Python [Protocol](https://docs.python.org/3/library/typing.html#typing.Protocol) that you subclass to define how URLs are generated. The Weights & Biases plugin provides a [built-in link implementation](https://www.union.ai/docs/latest/flyte/api-reference/integrations/wandb/wandb) as an example.

## Creating a link

To create a link, subclass `Link` as a dataclass and implement the `get_link()` method. The method returns the URL string to display in the UI:

```python
from dataclasses import dataclass

import flyte
from flyte import Link

@dataclass
class GrafanaLink(Link):
    dashboard_url: str
    name: str = "Grafana"

    def get_link(
        self,
        run_name: str,
        project: str,
        domain: str,
        context: dict,
        parent_action_name: str,
        action_name: str,
        pod_name: str,
        **kwargs,
    ) -> str:
        return f"{self.dashboard_url}?var-pod={pod_name}"

env = flyte.TaskEnvironment(...)

@env.task(links=(GrafanaLink(dashboard_url="https://grafana.example.com/d/abc123"),))
def my_task() -> str:
    return "done"
```

The link appears as a clickable "Grafana" link in the Flyte UI for every execution of `my_task`.

## Using execution metadata

The `get_link()` method receives execution metadata that you can use to construct dynamic URLs. Here's an example modeled on the [built-in Wandb](https://www.union.ai/docs/latest/flyte/user-guide/integrations/wandb/_index) link that uses the `context` dict to resolve a run ID:

```python
from dataclasses import dataclass
from typing import Optional

from flyte import Link

@dataclass
class Wandb(Link):
    project: str
    entity: str
    id: Optional[str] = None
    name: str = "Weights & Biases"

    def get_link(
        self,
        run_name: str,
        project: str,
        domain: str,
        context: dict[str, str],
        parent_action_name: str,
        action_name: str,
        pod_name: str,
        **kwargs,
    ) -> str:
        run_id = self.id or context.get("wandb_id", run_name)
        return f"https://wandb.ai/{self.entity}/{self.project}/runs/{run_id}"
```

The `name` attribute controls the display label in the UI.

See the [`get_link()` API reference](https://www.union.ai/docs/latest/flyte/api-reference/flyte-sdk/flyte/link) for more details. Note that `action_name` and `pod_name` are template variables (`{{.actionName}}` and `{{.podName}}`) that are populated by the backend at runtime.

## Dynamic links with override

Use `task.override(links=...)` to set links at runtime. This is useful when link parameters depend on runtime values like run IDs or configuration:

```python
import os

import flyte
from flyteplugins.wandb import Wandb

env = flyte.TaskEnvironment(...)

WANDB_PROJECT = "my-ml-project"
WANDB_ENTITY = "my-team"

@env.task
def train_model(config: dict) -> dict:
    # Training logic here
    return {"accuracy": 0.95}

@env.task
async def main(wandb_id: str) -> dict:
    result = train_model.override(
        links=(
            Wandb(
                project=WANDB_PROJECT,
                entity=WANDB_ENTITY,
                id=wandb_id,
            ),
        )
    )(config={"lr": 0.001})

    return result

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, wandb_id="my-run-id")
```

The `override` approach lets you attach links with values that are only known at runtime, such as dynamically generated run IDs.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/reports ===

# Reports

The reports feature allows you to display and update custom output in the UI during task execution.

> **📝 Note**
>
> Reports are the Flyte 2 successor to **Decks** in Flyte 1. Where Flyte 1 used `enable_deck=True` and the `flytekit.Deck` API, Flyte 2 uses `report=True` and the `flyte.report` API described below.

First, you set the `report=True` flag in the task decorator. This enables the reporting feature for that task.
Within a task with reporting enabled, a `flyte.report.Report` object is created automatically.

> [!NOTE] Import `flyte.report` explicitly
> `flyte.report` is a submodule that `import flyte` does **not** import automatically.
> You must import it explicitly:
>
> ```python
> import flyte.report
> ```
>
> Without this, calls like `flyte.report.replace()` or `flyte.report.flush()` raise
> `AttributeError: module 'flyte' has no attribute 'report'`, most commonly hit in local or
> notebook runs. This applies to all `flyte.*` submodules: import the specific submodule you use,
> not just the top-level `flyte` package.

A `Report` object contains one or more tabs, each of which contains HTML.
You can write HTML to an existing tab and create new tabs to organize your content.
Initially, the `Report` object has one tab (the default tab, named `main`) with no content.

To write content:

- `flyte.report.log()` appends HTML content directly to the default tab.
- `flyte.report.replace()` replaces the content of the default tab with new HTML.

To get or create a new tab:

- `flyte.report.get_tab()` allows you to specify a unique name for the tab, and it will return the existing tab if it already exists or create a new one if it doesn't.
  It returns a `flyte.report._report.Tab`

You can `log()` or `replace()` HTML on the `Tab` object just as you can directly on the `Report` object.

To access the current `Report` object directly — for example, to enumerate its tabs or assemble the final HTML — call `flyte.report.current_report()`.

Finally, you send the report to the Flyte backend and make it visible in the UI:

- `flyte.report.flush()` dispatches the report (in its current state) to the backend.

You do **not** have to call `flyte.report.flush()` explicitly at the end of a task: when a task with `report=True` finishes, Flyte automatically performs a final flush for you.
Calling `flyte.report.flush()` yourself is only necessary when you want to *stream* updates to the UI while the task is still running (see **Tasks > Build tasks > Reports > Streaming example** below).

## A simple example

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# main = "main"
# params = ""
# ///

import flyte
import flyte.report

env = flyte.TaskEnvironment(name="reports_example")

@env.task(report=True)
async def task1():
    await flyte.report.replace.aio("<p>The quick, brown fox jumps over a lazy dog.</p>")
    tab2 = flyte.report.get_tab("Tab 2")
    tab2.log("<p>The quick, brown dog jumps over a lazy fox.</p>")
    await flyte.report.flush.aio()

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(task1)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/simple.py*

Here we define a task `task1` that uses `flyte.report.replace()` to set the content of the default tab, then creates a new tab named "Tab 2" with `flyte.report.get_tab()` and logs additional HTML content to it.
Finally, `flyte.report.flush()` is called to send the report to the backend.

## A more complex example

Here is another example.
We import the necessary modules, set up the task environment, define the main task with reporting enabled and define the data generation function:

```
import json
import random

import flyte
import flyte.report

env = flyte.TaskEnvironment(
    name="globe_visualization",
)

@env.task(report=True)
async def generate_globe_visualization():
    await flyte.report.replace.aio(get_html_content())
    await flyte.report.flush.aio()

def generate_globe_data():
    """Generate sample data points for the globe"""
    cities = [
        {"city": "New York", "country": "USA", "lat": 40.7128, "lng": -74.0060},
        {"city": "London", "country": "UK", "lat": 51.5074, "lng": -0.1278},
        {"city": "Tokyo", "country": "Japan", "lat": 35.6762, "lng": 139.6503},
        {"city": "Sydney", "country": "Australia", "lat": -33.8688, "lng": 151.2093},
        {"city": "Paris", "country": "France", "lat": 48.8566, "lng": 2.3522},
        {"city": "São Paulo", "country": "Brazil", "lat": -23.5505, "lng": -46.6333},
        {"city": "Mumbai", "country": "India", "lat": 19.0760, "lng": 72.8777},
        {"city": "Cairo", "country": "Egypt", "lat": 30.0444, "lng": 31.2357},
        {"city": "Moscow", "country": "Russia", "lat": 55.7558, "lng": 37.6176},
        {"city": "Beijing", "country": "China", "lat": 39.9042, "lng": 116.4074},
        {"city": "Lagos", "country": "Nigeria", "lat": 6.5244, "lng": 3.3792},
        {"city": "Mexico City", "country": "Mexico", "lat": 19.4326, "lng": -99.1332},
        {"city": "Bangkok", "country": "Thailand", "lat": 13.7563, "lng": 100.5018},
        {"city": "Istanbul", "country": "Turkey", "lat": 41.0082, "lng": 28.9784},
        {"city": "Buenos Aires", "country": "Argentina", "lat": -34.6118, "lng": -58.3960},
        {"city": "Cape Town", "country": "South Africa", "lat": -33.9249, "lng": 18.4241},
        {"city": "Dubai", "country": "UAE", "lat": 25.2048, "lng": 55.2708},
        {"city": "Singapore", "country": "Singapore", "lat": 1.3521, "lng": 103.8198},
        {"city": "Stockholm", "country": "Sweden", "lat": 59.3293, "lng": 18.0686},
        {"city": "Vancouver", "country": "Canada", "lat": 49.2827, "lng": -123.1207},
    ]

    categories = ["high", "medium", "low", "special"]

    data_points = []
    for city in cities:
        data_point = {**city, "value": random.randint(10, 100), "category": random.choice(categories)}
        data_points.append(data_point)

    return data_points
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/globe_visualization.py*

We then define the HTML content for the report:

```python
def get_html_content():
    data_points = generate_globe_data()
    html_content = f"""
    <!DOCTYPE html>
    <html lang="en">
    ...
    </html>
    """
    return html_content
```

(We exclude it here due to length. You can find it in the [source file](https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/globe_visualization.py)).

Finally, we run the workflow:

```
if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(generate_globe_visualization)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/globe_visualization.py*

When the workflow runs, the report will be visible in the UI:

![Globe visualization](https://www.union.ai/docs/latest/flyte/_static/images/user-guide/globe_visualization.png)

## Streaming example

Above we demonstrated reports that are sent to the UI once, at the end of the task execution.
But, you can also stream updates to the report during task execution and see the display update in real-time.

You do this by calling `flyte.report.flush()` periodically during task execution, instead of just at the end.
As a shortcut, you can also pass `do_flush=True` to `flyte.report.log()` or `flyte.report.replace()` to flush immediately after writing the content.

> [!NOTE]
> In the earlier examples we explicitly call `flyte.report.flush()` to send the report to the UI.
> As noted above, that final flush is optional: it happens automatically when the task completes.
> For streaming reports, on the other hand, calling `flyte.report.flush()` periodically (or passing `do_flush=True`
> to `flyte.report.log()` / `flyte.report.replace()`) is what makes the intermediate updates appear.

First we import the necessary modules, and set up the task environment:

```
import asyncio
import json
import math
import random
import time
from datetime import datetime
from typing import List

import flyte
import flyte.report

env = flyte.TaskEnvironment(name="streaming_reports")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/streaming_reports.py*

Next we define the HTML content for the report:

```python
DATA_PROCESSING_DASHBOARD_HTML = """
...
"""
```

(We exclude it here due to length. You can find it in the [source file](
https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/streaming_reports.py)).

Finally, we define the task that renders the report (`data_processing_dashboard`), the driver task of the workflow (`main`), and the run logic:

```
@env.task(report=True)
async def data_processing_dashboard(total_records: int = 50000) -> str:
    """
    Simulates a data processing pipeline with real-time progress visualization.
    Updates every second for approximately 1 minute.
    """
    await flyte.report.log.aio(DATA_PROCESSING_DASHBOARD_HTML, do_flush=True)

    # Simulate data processing
    processed = 0
    errors = 0
    batch_sizes = [800, 850, 900, 950, 1000, 1050, 1100]  # Variable processing rates

    start_time = time.time()

    while processed < total_records:
        # Simulate variable processing speed
        batch_size = random.choice(batch_sizes)

        # Add some processing delays occasionally
        if random.random() < 0.1:  # 10% chance of slower batch
            batch_size = int(batch_size * 0.6)
            await flyte.report.log.aio("""
            <script>addActivity("⚠️ Detected slow processing batch, optimizing...");</script>
            """, do_flush=True)
        elif random.random() < 0.05:  # 5% chance of error
            errors += random.randint(1, 5)
            await flyte.report.log.aio("""
            <script>addActivity("❌ Processing errors detected, retrying failed records...");</script>
            """, do_flush=True)
        else:
            await flyte.report.log.aio(f"""
            <script>addActivity("✅ Successfully processed batch of {batch_size} records");</script>
            """, do_flush=True)

        processed = min(processed + batch_size, total_records)
        current_time = time.time()
        elapsed = current_time - start_time
        rate = int(batch_size) if elapsed < 1 else int(processed / elapsed)
        success_rate = ((processed - errors) / processed) * 100 if processed > 0 else 100

        # Update dashboard
        await flyte.report.log.aio(f"""
        <script>
            updateDashboard({processed}, {total_records}, {rate}, {success_rate});
        </script>
        """, do_flush=True)

        print(f"Processed {processed:,} records, Errors: {errors}, Rate: {rate:,}"
              f" records/sec, Success Rate: {success_rate:.2f}%", flush=True)
        await asyncio.sleep(1)  # Update every second

        if processed >= total_records:
            break

    # Final completion message
    total_time = time.time() - start_time
    avg_rate = int(total_records / total_time)

    await flyte.report.log.aio(f"""
    <script>addActivity("🎉 Processing completed successfully!");</script>
    <div style="background-color: #d4edda; border: 1px solid #c3e6cb; color: #155724; padding: 20px; border-radius: 8px; margin-top: 20px;">
        <h3>🎉 Processing Complete!</h3>
        <ul>
            <li><strong>Total Records:</strong> {total_records:,}</li>
            <li><strong>Processing Time:</strong> {total_time:.1f} seconds</li>
            <li><strong>Average Rate:</strong> {avg_rate:,} records/second</li>
            <li><strong>Success Rate:</strong> {success_rate:.2f}%</li>
            <li><strong>Errors Handled:</strong> {errors}</li>
        </ul>
    </div>
    """, do_flush=True)
    print(f"Data processing completed: {processed:,} records processed with {errors} errors.", flush=True)

    return f"Processed {total_records:,} records successfully"

@env.task
async def main():
    """
    Main task to run both reports.
    """
    await data_processing_dashboard(total_records=50000)

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/reports/streaming_reports.py*

The key to the live update ability is the `while` loop that appends Javascript to the report. The Javascript calls execute on append to the document and update it.

When the workflow runs, you can see the report updating in real-time in the UI:

![Data Processing Dashboard](https://www.union.ai/docs/latest/flyte/_static/images/user-guide/data_processing_dashboard.png)

## Rendering a custom type

The examples above build report HTML by hand. When you have a **custom type** — or a DataFrame, or a `StructuredDataset` — that you render the same way in many places, you can define a reusable **renderer** for the type and attach it to the type, instead of repeating the HTML-building logic at every call site.

A renderer is any class that satisfies the `flyte.types.Renderable` protocol: it implements a single `to_html(self, value) -> str` method that returns an HTML fragment for a value of your type.

```python
from flyte.types import Renderable

class Molecule:
    def __init__(self, name: str, smiles: str):
        self.name = name
        self.smiles = smiles

class MoleculeRenderer(Renderable):
    """A Renderable for the Molecule type."""

    def to_html(self, mol: Molecule) -> str:
        return f"<h2>{mol.name}</h2><pre>{mol.smiles}</pre>"
```

You attach the renderer to the type with `typing.Annotated`, then dispatch a value through its attached renderer with `flyte.types.TypeEngine.to_html()`. Log the resulting HTML to the report just like any other content:

```python
from typing import Annotated

import flyte
import flyte.report
from flyte.types import TypeEngine

env = flyte.TaskEnvironment(name="custom_renderer")

# Attaching the renderer to the type is the "registration".
RenderedMolecule = Annotated[Molecule, MoleculeRenderer()]

@env.task(report=True)
async def show_molecule() -> Molecule:
    mol = Molecule("caffeine", "CN1C=NC2=C1C(=O)N(C(=O)N2C)C")

    # Dispatch the value through the renderer attached to RenderedMolecule.
    html = TypeEngine.to_html(mol, RenderedMolecule)
    await flyte.report.log.aio(html)
    await flyte.report.flush.aio()

    return mol

if __name__ == "__main__":
    flyte.init_from_config()
    print(flyte.run(show_molecule).url)
```

`TypeEngine.to_html()` finds the `Renderable` attached to the type via `Annotated`, calls its `to_html()`, and returns the HTML string — which you then send to the report with `flyte.report.log()` (or `replace()`). The same pattern works for a DataFrame or `StructuredDataset`: annotate the type with a renderer that turns the frame into an HTML table.

> [!NOTE]
> The report contains only what you explicitly `log()` or `replace()`. Returning a value whose type has a renderer attached does **not** by itself add it to the report — render the value and log the HTML, as shown above.

Flyte's SDK also implements a few renderers of this kind internally — for pandas and PyArrow DataFrames and for Markdown strings. These aren't exposed as public API (only the `flyte.types.Renderable` protocol is), so treat them as examples of the same pattern rather than importable helpers.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/notebooks ===

# Notebooks

Flyte is designed to work with Jupyter notebooks, allowing you to write and execute workflows directly within a notebook environment.

## Iterating on and running a workflow

Download the following notebook file and open it in your favorite Jupyter environment: [interactive.ipynb](https://www.union.ai/docs/latest/flyte/_static/public/interactive.ipynb)

<!-- TODO: add back when working
📥 [interactive.ipynb](/_static/public/interactive.ipynb)
-->

In this example we have a simple workflow defined in our notebook.
You can iterate on the code in the notebook while running each cell in turn.

Note that the `flyte.init()` call at the top of the notebook looks like this:

```python
flyte.init(
    endpoint="https://union.example.com",
    org="example_org",
    project="example_project",
    domain="development",
)
```

You will have to adjust it to match your Union server endpoint, organization, project, and domain.

## Accessing runs and downloading logs

Similarly, you can download the following notebook file and open it in your favorite Jupyter environment: [remote.ipynb](https://www.union.ai/docs/latest/flyte/_static/public/remote.ipynb)

<!-- TODO: add back when working
📥 [remote.ipynb](/_static/public/remote.ipynb)
-->

In this example we use the `flyte.remote` package to list existing runs, access them, and download their details and logs.

For a guide on working with runs, actions, inputs, and outputs, see [Interact with runs and actions](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-deployment/interacting-with-runs).

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/remote-tasks ===

# Remote tasks

Remote tasks let you use previously deployed tasks without importing their code or dependencies. This enables teams to share and reuse tasks without managing complex dependency chains or container images.

## Prerequisites

Remote tasks must be deployed before you can use them. See the [task deployment guide](../task-deployment/_index) for details.

## Basic usage

Use `flyte.remote.Task.get()` to reference a deployed task:

```python
import flyte
import flyte.remote

env = flyte.TaskEnvironment(name="my_env")

# Get the latest version of a deployed task
data_processor = flyte.remote.Task.get(
    "data_team.spark_analyzer",
    auto_version="latest"
)

# Use it in your task
@env.task
async def my_task(data_path: str) -> flyte.io.DataFrame:
    # Call the reference task like any other task
    result = await data_processor(input_path=data_path)
    return result
```

You can run this directly without deploying it:

```bash
flyte run my_workflow.py my_task --data_path s3://my-bucket/data.parquet
```

## Understanding lazy loading

Remote tasks use **lazy loading** to keep module imports fast and enable flexible client configuration. When you call `flyte.remote.Task.get()`, it returns a lazy reference that doesn't actually fetch the task from the server until the first invocation.

### When tasks are fetched

The remote task is fetched from the server only when:

- You call `flyte.run()` with the task
- You call `flyte.deploy()` with code that uses the task
- You invoke the task with the `()` operator inside another task
- You explicitly call `.fetch()` on the lazy reference

```python
import flyte.remote

# This does NOT make a network call - returns a lazy reference
data_processor = flyte.remote.Task.get(
    "data_team.spark_analyzer",
    auto_version="latest"
)

# The task is fetched here when you invoke it
run = flyte.run(data_processor, input_path="s3://my-bucket/data.parquet")
```

### Benefits of lazy loading

**Fast module loading**: Since no network calls are made during import, your Python modules load quickly even when referencing many remote tasks.

**Late binding**: You can call `flyte.init()` after importing remote tasks, and the correct client will be bound when the task is actually invoked:

```python
import flyte
import flyte.remote

# Load remote task reference at module level
data_processor = flyte.remote.Task.get(
    "data_team.spark_analyzer",
    auto_version="latest"
)

# Initialize the client later
flyte.init_from_config()

# The task uses the client configured above
run = flyte.run(data_processor, input_path="s3://data.parquet")
```

### Error handling

Because of lazy loading, if a referenced task doesn't exist, you won't get an error when calling `get()`. Instead, the error occurs during invocation, raising a `flyte.errors.RemoteTaskNotFoundError`:

```python
import flyte
import flyte.remote
import flyte.errors

# This succeeds even if the task doesn't exist
data_processor = flyte.remote.Task.get(
    "nonexistent.task",
    auto_version="latest"
)

try:
    # Error occurs here during invocation
    run = flyte.run(data_processor, input_path="s3://data.parquet")
except flyte.errors.RemoteTaskNotFoundError as e:
    print(f"Task not found or invocation failed: {e}")
    # Handle the error - perhaps use a fallback task
    # or notify the user that the task needs to be deployed
```

You can also catch errors when using remote tasks within other tasks:

```python
import flyte.errors

@env.task
async def pipeline_with_fallback(data_path: str) -> dict:
    try:
        # Try to use the remote task
        result = await data_processor(input_path=data_path)
        return {"status": "success", "result": result}
    except flyte.errors.RemoteTaskNotFoundError as e:
        # Fallback to local processing
        print(f"Remote task failed: {e}, using local fallback")
        return {"status": "fallback", "result": local_process(data_path)}
    except flyte.errors.RemoteTaskUsageError as e:
        raise ValueError(f"Bad Usage of remote task, maybe arguments dont match!")
```

### Eager fetching with `fetch()`

While lazy loading is convenient, you can explicitly fetch a task upfront using the `fetch()` method. This is useful for:

- **Catching errors early**: Validate that the task exists before execution starts
- **Caching**: Avoid the network call on first invocation when running multiple times
- **Service initialization**: Pre-load tasks when your service starts

```python
import flyte
import flyte.remote
import flyte.errors

# Get the lazy reference
data_processor = flyte.remote.Task.get(
    "data_team.spark_analyzer",
    auto_version="latest"
)

try:
    # Eagerly fetch the task details
    task_details = data_processor.fetch()

    # Now the task is cached - subsequent calls won't hit the remote service
    # You can pass either the original reference or task_details to flyte.run
    run1 = flyte.run(data_processor, input_path="s3://data1.parquet")
    run2 = flyte.run(task_details, input_path="s3://data2.parquet")

except flyte.errors.RemoteTaskNotFoundError as e:
    print(f"Task not found failed at startup: {e}")
    raise
except flyte.errors.RemoteTaskUsageError as e:
    print(f"Task run validation failed....")
    # Handle the error before any execution attempts
```

For async contexts, use `await fetch.aio()`:

```python
import flyte.remote

async def initialize_service():
    processor_ref = flyte.remote.Task.get(
        "data_team.spark_analyzer",
        auto_version="latest"
    )

    try:
        # Fetch asynchronously
        task_details = await processor_ref.fetch.aio()
        print(f"Task {task_details.name} loaded successfully")
        return processor_ref  # Return the cached reference
    except flyte.errors.RemoteTaskNotFoundError as e:
        print(f"Failed to load task: {e}")
        raise

# Initialize once at service startup
cached_processor = None

async def startup():
    global cached_processor
    cached_processor = await initialize_service()

# Later in your service
async def process_request(data_path: str):
    # The task is already cached from initialization
    # No network call on first invocation
    run = flyte.run(cached_processor, input_path=data_path)
    return run
```

**When to use eager fetching**:

- **Service startup**: Fetch all remote tasks during initialization to validate they exist and cache them
- **Multiple invocations**: If you'll invoke the same task many times, fetch once to cache it
- **Fail-fast validation**: Catch configuration errors before execution begins

**When lazy loading is better**:

- **Single-use tasks**: If you only invoke the task once, lazy loading is simpler
- **Import-time overhead**: Keep imports fast by deferring network calls
- **Conditional usage**: If the task may not be needed, don't fetch it upfront

### Module-level vs dynamic loading

**Module-level loading (recommended)**: Load remote tasks at the module level for cleaner, more maintainable code:

```python
import flyte.remote

# Module-level - clear and maintainable
data_processor = flyte.remote.Task.get(
    "data_team.spark_analyzer",
    auto_version="latest"
)

@env.task
async def my_task(data_path: str):
    return await data_processor(input_path=data_path)
```

**Dynamic loading**: You can also load remote tasks dynamically within a task if needed:

```python
@env.task
async def dynamic_pipeline(task_name: str, data_path: str):
    # Load the task based on runtime parameters
    processor = flyte.remote.Task.get(
        f"data_team.{task_name}",
        auto_version="latest"
    )

    try:
        result = await processor(input_path=data_path)
        return result
    except flyte.errors.RemoteTaskNotFoundError as e:
        raise ValueError(f"Task {task_name} not found: {e}")
```

## Complete example

This example shows how different teams can collaborate using remote tasks.

### Team A: Spark environment

Team A maintains Spark-based data processing tasks:

```python
# spark_env.py
from dataclasses import dataclass
import flyte

env = flyte.TaskEnvironment(name="spark_env")

@dataclass
class AnalysisResult:
    mean_value: float
    std_dev: float

@env.task
async def analyze_data(data_path: str) -> AnalysisResult:
    # Spark code here (not shown)
    return AnalysisResult(mean_value=42.5, std_dev=3.2)

@env.task
async def compute_score(result: AnalysisResult) -> float:
    # More Spark processing
    return result.mean_value / result.std_dev
```

Deploy the Spark environment:

```bash
flyte deploy spark_env/
```

### Team B: ML environment

Team B maintains PyTorch-based ML tasks:

```python
# ml_env.py
from pydantic import BaseModel
import flyte

env = flyte.TaskEnvironment(name="ml_env")

class PredictionRequest(BaseModel):
    feature_x: float
    feature_y: float

class Prediction(BaseModel):
    score: float
    confidence: float
    model_version: str

@env.task
async def run_inference(request: PredictionRequest) -> Prediction:
    # PyTorch model inference (not shown)
    return Prediction(
        score=request.feature_x * 2.5,
        confidence=0.95,
        model_version="v2.1"
    )
```

Deploy the ML environment:

```bash
flyte deploy ml_env/
```

### Team C: Orchestration

Team C builds a workflow using remote tasks from both teams without needing Spark or PyTorch dependencies:

```python
# orchestration_env.py
import flyte.remote

env = flyte.TaskEnvironment(name="orchestration")

# Reference tasks from other teams
analyze_data = flyte.remote.Task.get(
    "spark_env.analyze_data",
    auto_version="latest"
)

compute_score = flyte.remote.Task.get(
    "spark_env.compute_score",
    auto_version="latest"
)

run_inference = flyte.remote.Task.get(
    "ml_env.run_inference",
    auto_version="latest"
)

@env.task
async def orchestrate_pipeline(data_path: str) -> float:
    # Use Spark tasks without Spark dependencies
    analysis = await analyze_data(data_path=data_path)

    # Access attributes from the result
    # (Flyte creates a fake type that allows attribute access)
    print(f"Analysis: mean={analysis.mean_value}, std={analysis.std_dev}")

    data_score = await compute_score(result=analysis)

    # Use ML task without PyTorch dependencies
    # Pass Pydantic models as dictionaries
    prediction = await run_inference(
        request={
            "feature_x": analysis.mean_value,
            "feature_y": data_score
        }
    )

    # Access Pydantic model attributes
    print(f"Prediction: {prediction.score} (confidence: {prediction.confidence})")

    return prediction.score
```

Run the orchestration task directly (no deployment needed):

**Using Python API**:

```python
if __name__ == "__main__":
    flyte.init_from_config()

    run = flyte.run(
        orchestrate_pipeline,
        data_path="s3://my-bucket/data.parquet"
    )

    print(f"Execution URL: {run.url}")
    # You can wait for the execution
    run.wait()

    # You can then retrieve the outputs
    print(f"Pipeline result: {run.outputs()}")
```

**Using CLI**:

```bash
flyte run orchestration_env.py orchestrate_pipeline --data_path s3://my-bucket/data.parquet
```

## Invoke remote tasks in a script.

You can also run any remote task directly using a script in a similar way

```python
import flyte
import flyte.models
import flyte.remote

flyte.init_from_config()

# Fetch the task
remote_task = flyte.remote.Task.get("package-example.calculate_average", auto_version="latest")

# Create a run, note keyword arguments are required currently. In the future this will accept positional args based on the declaration order, but, we still recommend to use keyword args.
run = flyte.run(remote_task, numbers=[1.0, 2.0, 3.0])

print(f"Execution URL: {run.url}")
# you can view the phase

print(f"Current Phase: {run.phase}")
# You can wait for the execution
run.wait()

# Only available after flyte >= 2.0.0b39
print(f"Current phase: {run.phase}")

# Phases can be compared to
if run.phase == flyte.models.ActionPhase.SUCCEEDED:
    print(f"Run completed!")

# You can then retrieve the outputs
print(f"Pipeline result: {run.outputs()}")
```

## Why use remote tasks?

Remote tasks solve common collaboration and dependency management challenges:

**Cross-team collaboration**: Team A has deployed a Spark task that analyzes large datasets. Team B needs this analysis for their ML pipeline but doesn't want to learn Spark internals, install Spark dependencies, or build Spark-enabled container images. With remote tasks, Team B simply references Team A's deployed task.

**Platform reusability**: Platform teams can create common, reusable tasks (data validation, feature engineering, model serving) that other teams can use without duplicating code or managing complex dependencies.

**Microservices for data workflows**: Remote tasks work like microservices for long-running tasks or agents, enabling secure sharing while maintaining isolation.

## When to use remote tasks

Use remote tasks when you need to:

- Use functionality from another team without their dependencies
- Share common tasks across your organization
- Build reusable platform components
- Avoid dependency conflicts between different parts of your workflow
- Create modular, maintainable data pipelines

## How remote tasks work

### Security model

Remote tasks run in the **caller's project and domain** using the caller's compute resources, but execute with the **callee's service accounts, IAM roles, and secrets**. This ensures:

- Tasks are secure from misuse
- Resource usage is properly attributed
- Authentication and authorization are maintained
- Collaboration remains safe and controlled

### Type system

Remote tasks use Flyte's default types as inputs and outputs. Flyte's type system translates data between tasks without requiring the original dependencies:

| Remote Task Type | Flyte Type |
|-------------------|------------|
| DataFrames (`pandas`, `polars`, `spark`, etc.) | `flyte.io.DataFrame` |
| Object store files | `flyte.io.File` |
| Object store directories | `flyte.io.Dir` |
| Pydantic models | Dictionary (Flyte creates a representation) |

Any DataFrame type (pandas, polars, spark) automatically becomes `flyte.io.DataFrame`, allowing data exchange between tasks using different DataFrame libraries. You can also write custom integrations or explore Flyte's plugin system for additional types.

For Pydantic models specifically, you don't need the exact model locally. Pass a dictionary as input, and Flyte will handle the translation.

## Versioning options

Reference tasks support flexible versioning:

**Specific version**:

```python
task = flyte.remote.Task.get(
    "team_a.process_data",
    version="v1.2.3"
)
```

**Latest version** (`auto_version="latest"`):

```python
# Always use the most recently deployed version
task = flyte.remote.Task.get(
    "team_a.process_data",
    auto_version="latest"
)
```

**Current version** (`auto_version="current"`):

```python
# Use the same version as the calling task's deployment
# Useful when all environments deploy with the same version
# Can only be used from within a task context
task = flyte.remote.Task.get(
    "team_a.process_data",
    auto_version="current"
)
```

## Customizing remote tasks

Remote tasks can be customized by overriding various properties without modifying the original deployed task. This allows you to adjust resource requirements, retry strategies, caching behavior, and more based on your specific use case.

### Available overrides

The `override()` method on remote tasks accepts the following parameters:

- **short_name** (`str`): A short name for the task instance
- **resources** (`flyte.Resources`): CPU, memory, GPU, and storage limits
- **retries** (`int | flyte.RetryStrategy`): Number of retries or retry strategy
- **timeout** (`flyte.TimeoutType`): Task execution timeout
- **env_vars** (`Dict[str, str]`): Environment variables to set
- **secrets** (`flyte.SecretRequest`): Secrets to inject
- **max_inline_io_bytes** (`int`): Maximum size for inline IO in bytes
- **cache** (`flyte.Cache`): Cache behavior and settings
- **queue** (`str`): Execution queue to use

### Override examples

**Increase resources for a specific use case**:

```python
import flyte.remote

# Get the base task
data_processor = flyte.remote.Task.get(
    "data_team.spark_analyzer",
    auto_version="latest"
)

# Override with more resources for large dataset processing
large_data_processor = data_processor.override(
    resources=flyte.Resources(
        cpu="16",
        memory="64Gi",
        storage="200Gi"
    )
)

@env.task
async def process_large_dataset(data_path: str):
    # Use the customized version
    return await large_data_processor(input_path=data_path)
```

**Add retries and timeout**:

```python
# Override with retries and timeout for unreliable operations
reliable_processor = data_processor.override(
    retries=3,
    timeout="2h"
)

@env.task
async def robust_pipeline(data_path: str):
    return await reliable_processor(input_path=data_path)
```

**Configure caching**:

```python
# Override cache settings
cached_processor = data_processor.override(
    cache=flyte.Cache(
        behavior="override",
        version_override="v2",
        serialize=True
    )
)
```

**Set environment variables and secrets**:

```python
# Override with custom environment and secrets
custom_processor = data_processor.override(
    env_vars={
        "LOG_LEVEL": "DEBUG",
        "REGION": "us-west-2"
    },
    secrets=flyte.SecretRequest(
        secrets={"api_key": "my-secret-key"}
    )
)
```

**Multiple overrides**:

```python
# Combine multiple overrides
production_processor = data_processor.override(
    short_name="prod_spark_analyzer",
    resources=flyte.Resources(cpu="8", memory="32Gi"),
    retries=5,
    timeout="4h",
    env_vars={"ENV": "production"},
    queue="high-priority"
)

@env.task
async def production_pipeline(data_path: str):
    return await production_processor(input_path=data_path)
```

### Chain overrides

You can chain multiple `override()` calls to incrementally adjust settings:

```python
# Start with base task
processor = flyte.remote.Task.get("data_team.analyzer", auto_version="latest")

# Add resources
processor = processor.override(resources=flyte.Resources(cpu="4", memory="16Gi"))

# Add retries for production
if is_production:
    processor = processor.override(retries=5, timeout="2h")

# Use the customized task
result = await processor(input_path="s3://data.parquet")
```

## Best practices

### 1. Use meaningful task names

Remote tasks are accessed by name, so use clear, descriptive naming:

```python
# Good
customer_segmentation = flyte.remote.Task.get("ml_platform.customer_segmentation")

# Avoid
task1 = flyte.remote.Task.get("team_a.task1")
```

### 2. Document task interfaces

Since remote tasks abstract away implementation details, clear documentation of inputs, outputs, and behavior is essential:

```python
@env.task
async def process_customer_data(
    customer_ids: list[str],
    date_range: tuple[str, str]
) -> flyte.io.DataFrame:
    """
    Process customer data for the specified date range.

    Args:
        customer_ids: List of customer IDs to process
        date_range: Tuple of (start_date, end_date) in YYYY-MM-DD format

    Returns:
        DataFrame with processed customer features
    """
    ...
```

### 3. Prefer module-level loading

Load remote tasks at the module level rather than inside functions for cleaner code:

```python
import flyte.remote

# Good - module level
data_processor = flyte.remote.Task.get("team.processor", auto_version="latest")

@env.task
async def my_task(data: str):
    return await data_processor(input=data)
```

This approach:

- Makes dependencies clear and discoverable
- Reduces code duplication
- Works well with lazy loading (no performance penalty)

Dynamic loading within tasks is also supported when you need runtime flexibility.

### 4. Handle versioning thoughtfully

- Use `auto_version="latest"` during development for rapid iteration
- Use specific versions in production for stability and reproducibility
- Use `auto_version="current"` when coordinating multienvironment deployments

### 5. Deploy remote tasks first

Always deploy the remote tasks before using them. Tasks that reference them can be run directly without deployment:

Deploy the remote task environments first:

```bash
flyte deploy spark_env/
flyte deploy ml_env/
```

Then run the orchestration task directly (no deployment needed):

```bash
flyte run orchestration_env.py orchestrate_pipeline
```

If you want to deploy the orchestration task as well (for scheduled runs or to be referenced by other tasks), deploy it after its dependencies:

```bash
flyte deploy orchestration_env/
```

## Limitations

1. **Lazy error detection**: Because of lazy loading, errors about missing or invalid tasks only occur during invocation, not when calling `get()`. You'll receive a `flyte.errors.RemoteTaskNotFoundError` if the task doesn't exist and `flyte.errors.RemoteTaskUsageError` if it can't be invoked in the way you are passing either arguments or overrides.

2. **Type fidelity**: While Flyte translates types, you work with Flyte's representation of Pydantic models, not the exact original types

3. **Deployment order**: Referenced tasks must be deployed before tasks that reference them can be invoked

4. **Context requirement**: Using `auto_version="current"` requires running within a task context

5. **Dictionary inputs**: Pydantic models must be passed as dictionaries, which loses compile-time type checking

6. **No positional arguments**: Remote tasks currently only support keyword arguments (this may change in future versions)

## Next steps

- Learn about [task deployment](../task-deployment/_index)
- Explore [task environments and configuration](../task-configuration/_index)

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/error-handling ===

# Error handling

One of the key features of Flyte 2 is the ability to recover from user-level errors in a workflow execution.
This includes out-of-memory errors, timeouts, oversized inline I/O, and other exceptions.

In a distributed system with heterogeneous compute, certain types of errors are expected and even, in a sense, acceptable.
Flyte 2 recognizes this and allows you to handle them gracefully as part of your workflow logic.

This ability is a direct result of the fact that workflows are now written in regular Python,
giving you all the power and flexibility of Python error handling.
When a task fails, Flyte surfaces the failure to the calling task as a typed exception that you can catch with a
standard `try...except` block and respond to however you like: retry with more resources, fall back to a different
code path, or clean up and re-raise.

## How Flyte represents failures

When a downstream task fails, the failure propagates to the awaiting parent task as an exception from the
`flyte.errors` module. Every native exception derives from a small hierarchy of base classes:

- `flyte.errors.BaseRuntimeError`: the root of all Flyte runtime errors.
- `flyte.errors.RuntimeUserError`: the failure was caused by your code (a bug, an exception you raised, an
  out-of-memory condition, and so on). An exception you raise inside a task, say a `ValueError`, is wrapped and
  surfaces to the parent as a `flyte.errors.RuntimeUserError`.
- `flyte.errors.RuntimeSystemError`: the failure was caused by the platform rather than your code.
- `flyte.errors.RuntimeUnknownError`: the failure could not be classified as a user or system error.
Every concrete error carries a `code` attribute (a short, stable string identifier, often the exception's class name, e.g. `"TaskTimeoutError"`) that you can
inspect when logging or branching. Because the errors form a hierarchy, you can catch broadly
(`except flyte.errors.RuntimeUserError`) or narrowly (`except flyte.errors.OOMError`), depending on how specific
your recovery logic needs to be.

## Catching and recovering from errors

The most common pattern is to catch a specific exception and re-run the failing task with a different
configuration. The following example intentionally triggers an out-of-memory error, catches the
`flyte.errors.OOMError`, and retries the task with more memory:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
# ]
# main = "main"
# params = ""
# ///

import asyncio

import flyte
import flyte.errors

env = flyte.TaskEnvironment(name="fail", resources=flyte.Resources(cpu=1, memory="250Mi"))

@env.task
async def oomer(x: int):
    large_list = [0] * 100000000
    print(len(large_list))

@env.task
async def always_succeeds() -> int:
    await asyncio.sleep(1)
    return 42

@env.task
async def main() -> int:
    try:
        await oomer(2)
    except flyte.errors.OOMError as e:
        print(f"Failed with oom trying with more resources: {e}, of type {type(e)}, {e.code}")
        try:
            await oomer.override(resources=flyte.Resources(cpu=1, memory="1Gi"))(5)
        except flyte.errors.OOMError as e:
            print(f"Failed with OOM Again giving up: {e}, of type {type(e)}, {e.code}")
            raise e
    finally:
        await always_succeeds()

    return await always_succeeds()

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/error-handling/error_handling.py*

In this code, we do the following:

* Import the necessary modules, including `flyte.errors`.
* Set up the task environment with a modest resource allocation of 1 CPU and 250 MiB of memory.
* Define two tasks: `oomer`, which allocates a large list and is likely to run out of memory, and
  `always_succeeds`, which always returns cleanly.
* Define the `main` task (the top-level workflow task) that contains the failure-recovery logic.

The `try...except` block in `main` runs `oomer`. If it exhausts memory, `main` catches the
`flyte.errors.OOMError` and retries by calling `oomer.override(resources=...)` with a larger memory allocation.
If the retry also runs out of memory, `main` gives up and re-raises the error. The `finally` block runs
`always_succeeds` regardless of the outcome.

This type of dynamic error handling lets you gracefully recover from user-level errors in your workflows using
patterns you already know from ordinary Python. For a complete, self-tuning version of this pattern that caches
the optimal memory setting across runs, see the
[`resource_tuner` example](https://github.com/flyteorg/flyte-sdk/blob/main/examples/advanced/resource_tuner.py).

> [!NOTE] Programmatic recovery vs. automatic retries
> Catching an exception and re-running a task is *programmatic* recovery: you decide what to do differently on
> the next attempt. This is distinct from Flyte's *automatic* retries (`retries=N` on a task), which simply
> re-run the same attempt unchanged. The two compose: automatic retries handle transient failures, while a
> `try...except` handles failures you want to respond to deliberately. See
> [Retries and timeouts](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/retries-and-timeouts).

## Limiting inline I/O

Small task inputs and outputs are passed *inline* (embedded directly in the task's metadata) rather than offloaded
to blob storage. This is fast, but very large inline values are undesirable, so each task has a ceiling on the
size of its inline I/O. You set this ceiling with the `max_inline_io_bytes` parameter on `@env.task`, and Flyte
raises a `flyte.errors.InlineIOMaxBytesBreached` when an input or output exceeds it:

```python
import flyte
import flyte.errors

env = flyte.TaskEnvironment(
    name="large_inline_io",
    resources=flyte.Resources(cpu=1, memory="250Mi"),
)

@env.task(max_inline_io_bytes=100 * 1024)  # Limit inline I/O to 100 KiB
async def printer_task(x: str) -> str:
    print(f"Printer task received: {x}")
    return x

@env.task
async def large_inline_io() -> str:
    small = await printer_task("Hello, world!")
    print(f"Small string result: {small}")

    # A large string that exceeds the 100 KiB inline limit
    large_string = "A" * 10**6  # ~1 MiB
    try:
        return await printer_task(large_string)
    except flyte.errors.InlineIOMaxBytesBreached as e:
        print(f"Inline I/O limit breached: {e}")
        raise
```

The small string passes through, but the ~1 MiB string breaches the 100 KiB limit and raises
`flyte.errors.InlineIOMaxBytesBreached`. When you expect large values, raise `max_inline_io_bytes` or pass the
data as a `flyte.io.File` or `flyte.io.Dir` so it is offloaded to blob storage instead of travelling inline. A
runnable version of this example is available as
[`large_inline_io.py`](https://github.com/flyteorg/flyte-sdk/blob/main/examples/advanced/large_inline_io.py).

## Natively-supported exceptions

Flyte raises typed exceptions for the failure modes it recognizes, so you can catch exactly the condition you
care about. The most commonly caught errors are:

| Exception | Raised when |
|---|---|
| `flyte.errors.OOMError` | A task exceeds its memory allocation. |
| `flyte.errors.TaskTimeoutError` | A task runs longer than its configured timeout. |
| `flyte.errors.InlineIOMaxBytesBreached` | An input or output exceeds the task's `max_inline_io_bytes` limit. |
| `flyte.errors.RetriesExhaustedError` | A task fails after all of its automatic retries are used up. |
| `flyte.errors.TaskInterruptedError` | A task running on interruptible (spot) compute is preempted. |
| `flyte.errors.ActionAbortedError` | An action is aborted externally via the CLI, UI, or API. |
| `flyte.errors.ImagePullBackOffError` | The task's container image cannot be pulled. |
| `flyte.errors.NonRecoverableError` | A failure that should not be retried, regardless of the retry budget. |

These all derive from `flyte.errors.RuntimeUserError`, so a single `except flyte.errors.RuntimeUserError`
catches any of them when you want uniform handling. This is only a selection. For the complete catalog of
catchable exception classes, see the [`flyte.errors` API reference](https://www.union.ai/docs/latest/flyte/user-guide/api-reference/flyte-sdk/flyte.errors/_index).

## Related pages

- [Retries and timeouts](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/retries-and-timeouts): configure automatic retries and execution time limits.
- [Abort and cancel actions](./abort-tasks): stop actions programmatically or externally, and handle `flyte.errors.ActionAbortedError`.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/traces ===

# Traces

The `@flyte.trace` decorator provides fine-grained observability and resumption capabilities for functions called within your Flyte workflows.
Traces are used on **helper functions** that tasks call to perform specific operations like API calls, data processing, or computations.
Traces are particularly useful for [managing the challenges of non-deterministic behavior in workflows](https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/gotchas-and-caveats), allowing you to track execution details and resume from failures.

## What are traced functions for?

At the top level, Flyte workflows are composed of **tasks**. But it is also common practice to break down complex task logic into smaller, reusable functions by defining helper functions that tasks call to perform specific operations.

Any helper functions defined or imported into the same file as a task definition are automatically uploaded to the Flyte environment alongside the task when it is deployed.

At the task level, observability and resumption of failed executions is provided by caching, but what if you want these capabilities at a more granular level, for the individual operations that tasks perform?

This is where **traced functions** come in. By decorating helper functions with `@flyte.trace`, you enable:

- **Detailed observability**: Track execution time, inputs/outputs, and errors for each function call.
- **Fine-grained resumption**: If a workflow fails, resume from the last successful traced function instead of re-running the entire task.
Each traced function is effectively a checkpoint within its task.

Here is an example:

```
import asyncio

import flyte

env = flyte.TaskEnvironment("env")

@flyte.trace
async def call_llm(prompt: str) -> str:
    await asyncio.sleep(0.1)
    return f"LLM response for: {prompt}"

@flyte.trace
async def process_data(data: str) -> dict:
    await asyncio.sleep(0.2)
    return {"processed": data, "status": "completed"}

@env.task
async def research_workflow(topic: str) -> dict:
    llm_result = await call_llm(f"Generate research plan for: {topic}")
    processed_data = await process_data(llm_result)
    return {"topic": topic, "result": processed_data}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/traces/task_vs_trace.py*

## What gets traced

Traces capture detailed execution information:

- **Execution time**: How long each function call takes.
- **Inputs and outputs**: Function parameters and return values.
- **Checkpoints**: State that enables workflow resumption.

### Errors are not recorded

Only successful trace executions are recorded in the checkpoint system. When a traced function fails, the exception propagates up to your task code where you can handle it with standard error handling patterns.

### Supported function types

The trace decorator works with:

- **Asynchronous functions**: Functions defined with `async def`.
- **Generator functions**: Functions that `yield` values.
- **Async generators**: Functions that `async yield` values.

> [!NOTE]
> Currently tracing only works for asynchronous functions. Tracing of synchronous functions is coming soon.

```
@flyte.trace
async def async_api_call(topic: str) -> dict:
    # Asynchronous API call
    await asyncio.sleep(0.1)
    return {"data": ["item1", "item2", "item3"], "status": "success"}

@flyte.trace
async def stream_data(items: list[str]):
    # Async generator function for streaming
    for item in items:
        await asyncio.sleep(0.02)
        yield f"Processing: {item}"

@flyte.trace
async def async_stream_llm(prompt: str):
    # Async generator for streaming LLM responses
    chunks = ["Research shows", " that machine learning", " continues to evolve."]
    for chunk in chunks:
        await asyncio.sleep(0.05)
        yield chunk

@env.task
async def research_workflow(topic: str) -> dict:
    llm_result = await async_api_call(topic)

    # Collect async generator results
    processed_data = []
    async for item in stream_data(llm_result["data"]):
        processed_data.append(item)

    llm_stream = []
    async for chunk in async_stream_llm(f"Summarize research on {topic}"):
        llm_stream.append(chunk)

    return {
        "topic": topic,
        "processed_data": processed_data,
        "llm_summary": "".join(llm_stream)
    }
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/traces/function_types.py*

## Task orchestration pattern

The typical Flyte workflow follows this pattern:

```
@flyte.trace
async def search_web(query: str) -> list[dict]:
    # Search the web and return results
    await asyncio.sleep(0.1)
    return [{"title": f"Article about {query}", "content": f"Content on {query}"}]

@flyte.trace
async def summarize_content(content: str) -> str:
    # Summarize content using LLM
    await asyncio.sleep(0.1)
    return f"Summary of {len(content.split())} words"

@flyte.trace
async def extract_insights(summaries: list[str]) -> dict:
    # Extract insights from summaries
    await asyncio.sleep(0.1)
    return {"insights": ["key theme 1", "key theme 2"], "count": len(summaries)}

@env.task
async def research_pipeline(topic: str) -> dict:
    # Each helper function creates a checkpoint
    search_results = await search_web(f"research on {topic}")

    summaries = []
    for result in search_results:
        summary = await summarize_content(result["content"])
        summaries.append(summary)

    final_insights = await extract_insights(summaries)

    return {
        "topic": topic,
        "insights": final_insights,
        "sources_count": len(search_results)
    }
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/traces/pattern.py*

**Benefits of this pattern:**

- If `search_web` succeeds but `summarize_content` fails, resumption skips the search step
- Each operation is independently observable and debuggable
- Clear separation between workflow coordination (task) and execution (traced functions)

## Relationship to caching and checkpointing

Understanding how traces work with Flyte's other execution features:

| Feature | Scope | Purpose | Default Behavior |
|---------|-------|---------|------------------|
| **Task Caching** | Entire task execution (`@env.task`) | Skip re-running tasks with same inputs | Enabled (`cache="auto"`) |
| **Traces** | Individual helper functions | Observability and fine-grained resumption | Manual (requires `@flyte.trace`) |
| **Checkpointing** | Workflow state | Resume workflows from failure points | Automatic when traces are used |

### How they work together

<!-- TODO
Lets use better typing for all of these examples, we have the opportunity to make this right for our users
-->

```
@flyte.trace
async def traced_data_cleaning(dataset_id: str) -> List[str]:
    # Creates checkpoint after successful execution.
    await asyncio.sleep(0.2)
    return [f"cleaned_record_{i}_{dataset_id}" for i in range(100)]

@flyte.trace
async def traced_feature_extraction(data: List[str]) -> dict:
    # Creates checkpoint after successful execution.
    await asyncio.sleep(0.3)
    return {
        "features": [f"feature_{i}" for i in range(10)],
        "feature_count": len(data),
        "processed_samples": len(data)
    }

@flyte.trace
async def traced_model_training(features: dict) -> dict:
    # Creates checkpoint after successful execution.
    await asyncio.sleep(0.4)
    sample_count = features["processed_samples"]
    # Mock accuracy based on sample count
    accuracy = min(0.95, 0.7 + (sample_count / 1000))
    return {
        "accuracy": accuracy,
        "epochs": 50,
        "model_size": "125MB"
    }

@env.task(cache="auto")  # Task-level caching enabled
async def data_pipeline(dataset_id: str) -> dict:
    # 1. If this exact task with these inputs ran before,
    #    the entire task result is returned from cache

    # 2. If not cached, execution begins and each traced function
    #    creates checkpoints for resumption
    cleaned_data = await traced_data_cleaning(dataset_id)      # Checkpoint 1
    features = await traced_feature_extraction(cleaned_data)   # Checkpoint 2
    model_results = await traced_model_training(features)      # Checkpoint 3

    # 3. If workflow fails at step 3, resumption will:
    #    - Skip traced_data_cleaning (checkpointed)
    #    - Skip traced_feature_extraction (checkpointed)
    #    - Re-run only traced_model_training

    return {"dataset_id": dataset_id, "accuracy": model_results["accuracy"]}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/traces/caching_vs_checkpointing.py*

### Execution flow

1. **Task Submission**: Task is submitted with input parameters
2. **Cache Check**: Flyte checks if identical task execution exists in cache
3. **Cache Hit**: If cached, return cached result immediately (no traces needed)
4. **Cache Miss**: Begin fresh execution
5. **Trace Checkpoints**: Each `@flyte.trace` function creates resumption points
6. **Failure Recovery**: If workflow fails, resume from last successful checkpoint
7. **Task Completion**: Final result is cached for future identical inputs

<!--
Clarify what actually happens on error vs success with traces

## Error handling and observability

Traces capture comprehensive execution information for debugging and monitoring:

```
@flyte.trace
async def risky_api_call(endpoint: str, data: dict) -> dict:
    """API call that might fail - traces capture errors."""
    try:
        response = await api_client.post(endpoint, json=data)
        return response.json()
    except Exception as e:
        # Error is automatically captured in trace
        logger.error(f"API call failed: {e}")
        raise

@env.task
async def error_handling() -> dict:
    try:
        result = await risky_api_call("/process", {"invalid": "data"})
        return {"status": "success", "result": result}
    except Exception as e:
        # The error is recorded in the trace for debugging
        return {"status": "error", "message": str(e)}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/traces/error_handling.py*

**What traces capture:**
- **Execution time**: Duration of each function call
- **Inputs and outputs**: Function parameters and return values
- **Checkpoints**: State that enables workflow resumption from successful executions
- **Action IDs**: Unique identifiers for each execution

**Error handling:**
- Errors from traced functions are not recorded in checkpoints
- Exceptions propagate to your task code for standard error handling
- The error_handling example shows how to catch and handle these exceptions in your task

TODO:
Ketan Umare:
we should show an example where tasks and traces can be used interchangeably

## Examples in practice

### LLM pipeline with traces

```python
import flyte

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

@flyte.trace
async def call_llm(prompt: str, model: str = "gpt-4") -> str:
    """Call LLM with specified model."""
    response = await llm_client.chat(prompt, model=model)
    return response

@flyte.trace
async def extract_entities(text: str) -> list[str]:
    """Extract named entities from text."""
    entities = await nlp_service.extract_entities(text)
    return entities

@env.task
async def process_documents(documents: list[str]) -> dict:
    """Process multiple documents through LLM pipeline."""
    results = []

    for doc in documents:
        # Each call is traced for monitoring and resumption
        summary = await call_llm(f"Summarize: {doc}")
        entities = await extract_entities(summary)

        results.append({
            "document": doc,
            "summary": summary,
            "entities": entities
        })

    return {"processed_documents": results, "total_count": len(results)}
```

This comprehensive tracing system provides visibility into your workflow execution while enabling robust error recovery and resumption capabilities.
-->

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/intra-task-checkpoints ===

# Intra-task checkpoints

Long-running tasks (model training especially) can fail partway through: a spot
instance is reclaimed, a pod is evicted, an out-of-memory error kills the process.
When the task is retried, it normally starts over from the beginning.

Intra-task checkpoints let a task save in-progress state to object storage as it
runs and load that state at the start of the next attempt, so a retry resumes from
where the previous attempt left off instead of repeating completed work.

## The checkpoint object

Inside a running task, `flyte.ctx().checkpoint` returns a `flyte.Checkpoint` (or `None`
when checkpointing isn't configured) bound to the action's checkpoint location in object storage:

- **Save**: `await checkpoint.save(...)` (async tasks) or `checkpoint.save_sync(...)`
  (sync tasks and synchronous framework callbacks). Accepts raw `bytes`, a file path,
  or a directory path; a directory is stored as a single compressed archive.
- **Load**: `await checkpoint.load()` or `checkpoint.load_sync()`. Returns a local
  `pathlib.Path` to the restored file or directory tree, or `None` when there is no
  previous checkpoint (i.e. on the first attempt).
- `flyte.latest_checkpoint(root, glob_pattern="**/last.ckpt")` finds the newest checkpoint file under a restored
  directory tree, useful for frameworks like PyTorch Lightning that write
  `last.ckpt` files into a directory.

Checkpoints only matter when the task can run more than once, so give the task
retries with `@env.task(retries=...)`. Each retry attempt sees the checkpoint saved
by the attempt before it.

> [!NOTE] Checkpoints vs. caching vs. traces
>
> - **Task caching** skips an entire task when it has already run with the same inputs.
> - **[Traces](./traces)** checkpoint at the boundaries of helper functions called by a task.
> - **Intra-task checkpoints** save state *within* a single task body (mid-loop,
>   mid-epoch) across retry attempts of the same action.

## Basic usage

The simplest checkpoint is a raw byte payload. This task counts up to
`n_iterations`, saving its progress on every iteration. A simulated failure kills
it partway through; the retry loads the saved counter and continues rather than
restarting from zero:

### Async

```
import flyte

env = flyte.TaskEnvironment(name="checkpoint_generic")

RETRIES = 3

@env.task(retries=RETRIES)
async def use_checkpoint(n_iterations: int = 10) -> int:
    checkpoint = flyte.ctx().checkpoint

    # Load the previous attempt's checkpoint, if any.
    # On the first attempt there is none, so load() returns None.
    path = await checkpoint.load()
    start = int(path.read_bytes()) if path else 0

    failure_interval = n_iterations // RETRIES
    index = start
    for index in range(start, n_iterations):
        if index > start and index % failure_interval == 0:
            # Simulate a failure so the next attempt resumes from the checkpoint
            raise RuntimeError(f"Simulated failure at iteration {index}")
        # Persist progress to object storage.
        await checkpoint.save(f"{index + 1}".encode())
    return index
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/generic_checkpoint.py*

### Sync

```
import flyte

env = flyte.TaskEnvironment(name="checkpoint_generic_sync")

RETRIES = 3

@env.task(retries=RETRIES)
def use_checkpoint(n_iterations: int = 10) -> int:
    checkpoint = flyte.ctx().checkpoint

    # Load the previous attempt's checkpoint, if any.
    # On the first attempt there is none, so load_sync() returns None.
    path = checkpoint.load_sync()
    start = int(path.read_bytes()) if path else 0

    failure_interval = n_iterations // RETRIES
    index = start
    for index in range(start, n_iterations):
        if index > start and index % failure_interval == 0:
            # Simulate a failure so the next attempt resumes from the checkpoint
            raise RuntimeError(f"Simulated failure at iteration {index}")
        # Persist progress to object storage.
        checkpoint.save_sync(f"{index + 1}".encode())
    return index
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/generic_checkpoint_sync.py*

Running this with `n_iterations=10` produces three failed attempts and one
successful one. Each attempt fails later than the last, because each one starts
from the checkpoint its predecessor saved.

## Checkpointing ML training frameworks

The same pattern applies to real training loops, whatever the framework:

1. **Load** the previous attempt's checkpoint at the start of the task; if one
   exists, restore the model/optimizer state and work out where to resume.
2. **Save** a checkpoint at a regular interval (every epoch or every N steps)
   as training progresses.

Frameworks with their own checkpoint files (PyTorch Lightning, Hugging Face
`Trainer`) already write them to a local directory; there you hook their callback
system and mirror that directory to the Flyte checkpoint, then feed the restored
directory back to the framework's native resume mechanism.

### PyTorch

Save the model state dict, optimizer state, and epoch counter with `torch.save`
after each epoch, and restore all three with `torch.load` on retry:

<br>

```
@env.task(retries=RETRIES)
async def train_linear(epochs: int = 10) -> float:
    checkpoint = flyte.ctx().checkpoint

    model = nn.Linear(4, 1)
    opt = torch.optim.SGD(model.parameters(), lr=0.01)

    # Resume model, optimizer, and epoch from the previous attempt, if any.
    prev = await checkpoint.load()
    if prev:
        blob = torch.load(prev, map_location="cpu", weights_only=False)
        model.load_state_dict(blob["model"])
        opt.load_state_dict(blob["opt"])
        start = int(blob["epoch"]) + 1
    else:
        start = 0

    wpath = pathlib.Path("pytorch_linear") / "training.pt"
    wpath.parent.mkdir(parents=True, exist_ok=True)

    failure_interval = epochs // RETRIES
    for epoch in range(start, epochs):
        x = torch.randn(8, 4)
        y = torch.randn(8, 1)
        loss = torch.nn.functional.mse_loss(model(x), y)
        opt.zero_grad()
        loss.backward()
        opt.step()

        if epoch > start and epoch % failure_interval == 0:
            # Simulate a failure so the next attempt resumes from the checkpoint
            raise RuntimeError(f"Simulated failure at epoch {epoch}")

        # Save model, optimizer, and epoch state to object storage.
        torch.save(
            {"model": model.state_dict(), "opt": opt.state_dict(), "epoch": epoch},
            wpath,
        )
        await checkpoint.save(wpath)

    with torch.no_grad():
        return float(model(torch.ones(1, 4)).squeeze().item())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/pytorch_checkpoint.py*

### PyTorch Lightning

Lightning already writes `last.ckpt` through its `ModelCheckpoint` callback.
Subclass it to mirror the checkpoint directory to Flyte after each epoch
(Lightning callbacks are synchronous, so use `flyte.Checkpoint.save_sync`):

<br>

```
class FlyteLightningCheckpointCallback(ModelCheckpoint):
    """A `ModelCheckpoint` that mirrors `dirpath` to the Flyte checkpoint after each epoch."""

    def __init__(self, flyte_checkpoint: flyte.Checkpoint, *, dirpath: str | pathlib.Path, **kwargs) -> None:
        super().__init__(dirpath=str(dirpath), **kwargs)
        self._flyte_checkpoint = flyte_checkpoint

    @override
    def on_train_epoch_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
        super().on_train_epoch_end(trainer, pl_module)
        if self.dirpath:
            # Lightning callbacks are synchronous, so use save_sync
            self._flyte_checkpoint.save_sync(pathlib.Path(self.dirpath))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/pytorch_lightning_checkpoint.py*

In the task, restore the previous tree, pick the newest `last.ckpt` with
`flyte.latest_checkpoint(restored_root, glob_pattern="**/last.ckpt")`, and hand it to `Trainer.fit(ckpt_path=...)`; Lightning
restores the model, optimizer, and epoch from there:

<br>

```
@env.task(retries=RETRIES)
def train_lightning(max_epochs: int = 10) -> float:
    checkpoint = flyte.ctx().checkpoint

    ckpt_dir = pathlib.Path("pl_checkpoints")
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    # Restore the previous attempt's checkpoint tree and find the newest last.ckpt.
    resume_ckpt = None
    prev = checkpoint.load_sync()
    if prev:
        last = flyte.latest_checkpoint(prev)
        if last:
            resume_ckpt = str(last)

    model = TinyModule()
    mc = FlyteLightningCheckpointCallback(
        checkpoint,
        dirpath=ckpt_dir,
        filename="last",
        save_last=True,
        save_top_k=1,
    )
    trainer = L.Trainer(
        max_epochs=max_epochs,
        enable_checkpointing=True,
        callbacks=[mc],
        enable_progress_bar=True,
        logger=False,
        accelerator="cpu",
        devices=1,
    )
    trainer.fit(model, make_loader(), ckpt_path=resume_ckpt)

    with torch.no_grad():
        return float(model(torch.ones(1, FEATURES)).squeeze().item())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/pytorch_lightning_checkpoint.py*

### Hugging Face Trainer

`transformers.Trainer` writes `checkpoint-<step>` directories under its
`output_dir` (here, every epoch via `save_strategy="epoch"`). A `TrainerCallback`
mirrors that directory to Flyte after each save:

<br>

```
class FlyteTrainerCheckpointCallback(TrainerCallback):
    """Mirror the Trainer's `output_dir` to the Flyte checkpoint after each epoch."""

    def __init__(self, checkpoint: flyte.Checkpoint, output_dir: pathlib.Path) -> None:
        self._checkpoint = checkpoint
        self._output_dir = output_dir

    def on_epoch_end(self, args, state, control, **kwargs) -> None:
        # Trainer callbacks are synchronous, so use save_sync
        self._checkpoint.save_sync(self._output_dir)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/huggingface_trainer_checkpoint.py*

On retry, restore the tree, locate the last Hugging Face checkpoint with
`get_last_checkpoint`, and pass it to `trainer.train(resume_from_checkpoint=...)`:

<br>

```
@env.task(retries=RETRIES)
def train_transformers(max_epochs: int = 10) -> float:
    checkpoint = flyte.ctx().checkpoint

    ckpt_dir = pathlib.Path("hf_trainer")
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    # Restore the previous attempt's checkpoint tree and find the last HF checkpoint.
    hf_resume = None
    prev = checkpoint.load_sync()
    if prev:
        hf_resume = get_last_checkpoint(str(prev))

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID, num_labels=2)

    args = TrainingArguments(
        output_dir=str(ckpt_dir),
        num_train_epochs=max_epochs,
        per_device_train_batch_size=4,
        save_strategy="epoch",
        save_total_limit=2,
        logging_steps=1,
        report_to="none",
        seed=42,
        use_cpu=True,
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=ToyTextDataset(tokenizer),
        data_collator=DataCollatorWithPadding(tokenizer),
        callbacks=[FlyteTrainerCheckpointCallback(checkpoint, ckpt_dir)],
    )
    trainer.train(resume_from_checkpoint=hf_resume)

    model.eval()
    device = next(model.parameters()).device
    with torch.no_grad():
        batch = tokenizer(
            "classification example for inference",
            return_tensors="pt",
            truncation=True,
            max_length=32,
            padding="max_length",
        )
        batch = {k: v.to(device) for k, v in batch.items()}
        logits = model(**batch).logits
        return float(logits[0, 1].item())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/huggingface_trainer_checkpoint.py*

### scikit-learn

For estimators that support incremental training with `partial_fit`, pickle the
estimator together with a progress counter after each training chunk. A retry
unpickles the bundle and continues from the next chunk:

<br>

```
@env.task(retries=RETRIES)
async def incremental_sgd(chunks: int = 10) -> float:
    checkpoint = flyte.ctx().checkpoint

    # Resume the estimator and progress from the previous attempt, if any.
    prev = await checkpoint.load()
    if prev:
        bundle = pickle.loads(prev.read_bytes())
        start = bundle["chunks_done"]
        clf = bundle["clf"]
    else:
        start = 0
        clf = SGDClassifier(max_iter=1, tol=None, random_state=0)

    bundle_path = pathlib.Path("sklearn_partial") / "sgd_bundle.pkl"
    bundle_path.parent.mkdir(parents=True, exist_ok=True)

    rng = np.random.default_rng(0)
    classes = np.array([0, 1])

    failure_interval = chunks // RETRIES
    for i in range(start, chunks):
        x = rng.standard_normal((32, 8))
        y = (x[:, 0] + x[:, 1] > 0).astype(int)
        clf.partial_fit(x, y, classes=classes)

        if i > start and i % failure_interval == 0:
            # Simulate a failure so the next attempt resumes from the checkpoint
            raise RuntimeError(f"Simulated failure at chunk {i}")

        # Pickle the estimator plus progress and save it to object storage.
        bundle_path.write_bytes(pickle.dumps({"clf": clf, "chunks_done": i + 1}))
        await checkpoint.save(bundle_path)

    x_test = rng.standard_normal((64, 8))
    y_test = (x_test[:, 0] + x_test[:, 1] > 0).astype(int)
    return float(clf.score(x_test, y_test))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/sklearn_partial_checkpoint.py*

### Unsloth

LoRA fine-tuning with [Unsloth](https://unsloth.ai/) and `trl.SFTTrainer` uses the
same callback-and-resume pattern as the Hugging Face `Trainer`, since `SFTTrainer`
is built on it. Unsloth requires an NVIDIA, AMD, or Intel GPU, so the task
environment requests one:

<br>

```
@env.task(retries=RETRIES)
def train_unsloth_sft(max_epochs: int = 10) -> float:
    from trl import SFTConfig, SFTTrainer
    from unsloth import FastLanguageModel

    checkpoint = flyte.ctx().checkpoint

    ckpt_dir = pathlib.Path("unsloth_sft")
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    # Restore the previous attempt's checkpoint tree and find the last HF checkpoint.
    hf_resume = None
    prev = checkpoint.load_sync()
    if prev:
        hf_resume = get_last_checkpoint(str(prev))

    max_seq_length = 512
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=MODEL_NAME,
        max_seq_length=max_seq_length,
        dtype=None,
        load_in_4bit=True,
    )
    model = FastLanguageModel.get_peft_model(
        model,
        r=16,
        target_modules=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ],
        lora_alpha=16,
        lora_dropout=0.0,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=42,
    )

    args = SFTConfig(
        output_dir=str(ckpt_dir),
        num_train_epochs=max_epochs,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=1,
        save_strategy="epoch",
        save_total_limit=2,
        logging_steps=1,
        report_to="none",
        seed=42,
        dataset_text_field="text",
        max_length=max_seq_length,
    )

    trainer = SFTTrainer(
        model=model,
        args=args,
        train_dataset=tiny_instruction_dataset(),
        processing_class=tokenizer,
        callbacks=[FlyteTrainerCheckpointCallback(checkpoint, ckpt_dir)],
    )
    trainer.train(resume_from_checkpoint=hf_resume)

    model.eval()
    device = next(model.parameters()).device
    with torch.no_grad():
        batch = tokenizer(
            "classification example for inference",
            return_tensors="pt",
            truncation=True,
            max_length=32,
            padding="max_length",
        )
        batch = {k: v.to(device) for k, v in batch.items()}
        logits = model(**batch).logits
        return float(logits[0, 1].mean().item())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/unsloth_sft_checkpoint.py*

> [!NOTE] Simulated failures in the runnable examples
> The full example files for the basic, PyTorch, and scikit-learn cases inject a
> failure at a regular interval (`failure_interval`) so you can watch the retries
> resume from the checkpoint. In production code you would drop those lines; real
> failures (preemptions, OOMs, crashes) trigger the same resume path.

## How checkpoints are stored

Each action attempt gets a checkpoint prefix in the object store configured for
your cluster. `flyte.Checkpoint.save` uploads a file as-is, stores a directory as
a gzip-compressed tarball, and accepts raw `bytes` as a single blob.
`flyte.Checkpoint.load` downloads the previous attempt's object into a local
temporary workspace and returns the path: a restored directory tree, or the path
to the single restored file.

Saving repeatedly overwrites the same checkpoint object, so the cost of frequent
checkpointing is upload bandwidth, not unbounded storage growth. Checkpoint how
often you can afford to lose work: every epoch is typical for training loops.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/grouping-actions ===

# Grouping actions

Groups are an organizational feature in Flyte that allow you to logically cluster related task invocations (called "actions") for better visualization and management in the UI.
Groups help you organize task executions into manageable, hierarchical structures regardless of whether you're working with large fanouts or smaller, logically-related sets of operations.

## What are groups?

Groups provide a way to organize task invocations into logical units in the Flyte UI.
When you have multiple task executions (whether from large [fanouts](./fanout), sequential operations, or any combination of tasks), groups help organize them into manageable units.

### The problem groups solve

Without groups, complex workflows can become visually overwhelming in the Flyte UI:

- Multiple task executions appear as separate nodes, making it hard to see the high-level structure
- Related operations are scattered throughout the workflow graph
- Debugging and monitoring becomes difficult when dealing with many individual task executions

Groups solve this by:

- **Organizing actions**: Multiple task executions within a group are presented as a hierarchical "folder" structure
- **Improving UI visualization**: Instead of many individual nodes cluttering the view, you see logical groups that can be collapsed or expanded
- **Aggregating status information**: Groups show aggregated run status (success/failure) of their contained actions when you hover over them in the UI
- **Maintaining execution parallelism**: Tasks still run concurrently as normal, but are organized for display

### How groups work

Groups are declared using the `flyte.group` context manager.
Any task invocations that occur within the `with flyte.group()` block are automatically associated with that group:

```python
with flyte.group("my-group-name"):
    # All task invocations here belong to "my-group-name"
    result1 = await task_a(data)
    result2 = await task_b(data)
    result3 = await task_c(data)
```

The key points about groups:

1. **Context-based**: Use the `with flyte.group("name"):` context manager.
2. **Organizational tool**: Task invocations within the context are grouped together in the UI.
3. **UI folders**: Groups appear as collapsible/expandable folders in the Flyte UI run tree.
4. **Status aggregation**: Hover over a group in the UI to see aggregated success/failure information.
5. **Execution unchanged**: Tasks still execute in parallel as normal; groups only affect organization and visualization.

**Important**: Groups do not aggregate outputs. Each task execution still produces its own individual outputs. Groups are purely for organization and UI presentation.

## Common grouping patterns

### Sequential operations

Group related sequential operations that logically belong together:

```
@env.task
async def data_pipeline(raw_data: str) -> str:
    with flyte.group("data-validation"):
        validated_data = await process_data(raw_data, "validate_schema")
        validated_data = await process_data(validated_data, "check_quality")
        validated_data = await process_data(validated_data, "remove_duplicates")

    with flyte.group("feature-engineering"):
        features = await process_data(validated_data, "extract_features")
        features = await process_data(features, "normalize_features")
        features = await process_data(features, "select_features")

    with flyte.group("model-training"):
        model = await process_data(features, "train_model")
        model = await process_data(model, "validate_model")
        final_model = await process_data(model, "save_model")

    return final_model
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/grouping-actions/grouping.py*

### Parallel processing with groups

Groups work well with parallel execution patterns:

```
@env.task
async def parallel_processing_example(n: int) -> str:
    tasks = []

    with flyte.group("parallel-processing"):
        # Collect all task invocations first
        for i in range(n):
            tasks.append(process_item(i, "transform"))

        # Execute all tasks in parallel
        results = await asyncio.gather(*tasks)

    # Convert to string for consistent return type
    return f"parallel_results: {results}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/grouping-actions/grouping.py*

### Multi-phase workflows

Use groups to organize different phases of complex workflows:

```
@env.task
async def multi_phase_workflow(data_size: int) -> str:
    # First phase: data preprocessing
    preprocessed = []
    with flyte.group("preprocessing"):
        for i in range(data_size):
            preprocessed.append(process_item(i, "preprocess"))
        phase1_results = await asyncio.gather(*preprocessed)

    # Second phase: main processing
    processed = []
    with flyte.group("main-processing"):
        for result in phase1_results:
            processed.append(process_item(result, "transform"))
        phase2_results = await asyncio.gather(*processed)

    # Third phase: postprocessing
    postprocessed = []
    with flyte.group("postprocessing"):
        for result in phase2_results:
            postprocessed.append(process_item(result, "postprocess"))
        final_results = await asyncio.gather(*postprocessed)

    # Convert to string for consistent return type
    return f"multi_phase_results: {final_results}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/grouping-actions/grouping.py*

### Conditional grouping

Groups can be used with conditional logic:

```
@env.task
async def conditional_processing(use_advanced_features: bool, input_data: str) -> str:
    base_result = await process_data(input_data, "basic_processing")

    if use_advanced_features:
        with flyte.group("advanced-features"):
            enhanced_result = await process_data(base_result, "advanced_processing")
            optimized_result = await process_data(enhanced_result, "optimize_result")
            return optimized_result
    else:
        with flyte.group("basic-features"):
            simple_result = await process_data(base_result, "simple_processing")
            return simple_result
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/grouping-actions/grouping.py*

## Key insights

Groups are primarily an organizational and UI visualization tool; they don't change how your tasks execute or aggregate their outputs, but they help organize related task invocations (actions) into collapsible folder-like structures for better workflow management and display. The aggregated status information (success/failure rates) is visible when hovering over group folders in the UI.

Groups make your Flyte workflows more maintainable and easier to understand, especially when working with complex workflows that involve multiple logical phases or large numbers of task executions. They serve as organizational "folders" in the UI's call stack tree, allowing you to collapse sections to reduce visual distraction while still seeing aggregated status information on hover.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/fanout ===

# Fanout

Flyte is designed to scale, allowing you to run workflows with large fanouts.
When you need to execute many tasks in parallel (such as processing a large dataset or running hyperparameter sweeps), Flyte provides powerful patterns to implement these operations efficiently.

> **📝 Note**
>
> In Flyte 1, mapping a task over many inputs used `map_task()` (the `flytekit.map_task` API). In Flyte 2, fan out with `asyncio.gather()` or `flyte.map()`.

This page covers the general `asyncio.gather` fanout pattern. For applying the *same* task to every
item of a list (the direct successor to Flyte 1's `map_task`), see [Mapping over inputs](./map).

That page also covers concurrency limits and error handling.

## Understanding fanout

A "fanout" pattern occurs when you spawn multiple tasks concurrently.
Each task runs in its own container and contributes an output that you later collect.
The most common way to implement this is using the [`asyncio.gather`](https://docs.python.org/3/library/asyncio-task.html#asyncio.gather) function.

In Flyte terminology, each individual task execution is called an "action": this represents a specific invocation of a task with particular inputs. When you call a task multiple times in a loop, you create multiple actions.

## Example

We start by importing our required packages, defining our Flyte environment, and creating a simple task that fetches user data from a mock API.

```
import asyncio
from typing import List, Tuple

import flyte

env = flyte.TaskEnvironment("fanout_env")

@env.task
async def fetch_data(user_id: int) -> dict:
    """Simulate fetching user data from an API - good for parallel execution."""
    # Simulate network I/O delay
    await asyncio.sleep(0.1)
    return {
        "user_id": user_id,
        "name": f"User_{user_id}",
        "score": user_id * 10,
        "data": f"fetched_data_{user_id}"
    }
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/fanout/fanout.py*

### Parallel execution

Next we implement the most common fanout pattern, which is to collect task invocations and execute them in parallel using `asyncio.gather()`:

```
@env.task
async def parallel_data_fetching(user_ids: List[int]) -> List[dict]:
    """Fetch data for multiple users in parallel - ideal for I/O bound operations."""
    tasks = []

    # Collect all fetch tasks - these can run in parallel since they're independent
    for user_id in user_ids:
        tasks.append(fetch_data(user_id))

    # Execute all fetch operations in parallel
    results = await asyncio.gather(*tasks)
    return results
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/fanout/fanout.py*

### Running the example

To actually run our example, we create a main guard that initializes Flyte and runs our main driver task:

```
if __name__ == "__main__":
    flyte.init_from_config()
    user_ids = [1, 2, 3, 4, 5]
    r = flyte.run(parallel_data_fetching, user_ids)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/fanout/fanout.py*

## How Flyte handles concurrency and parallelism

In the example we use a standard `asyncio.gather()` pattern.
When this pattern is used in a normal Python environment, the tasks would execute **concurrently** (cooperatively sharing a single thread through the event loop), but not in true **parallel** (multiple CPU cores simultaneously).

However, **Flyte transforms this concurrency model into true parallelism**. When you use `asyncio.gather()` in a Flyte task:

1. **Flyte acts as a distributed event loop**: Instead of scheduling coroutines on a single machine, Flyte schedules each task action to run in its own container across the cluster
2. **Concurrent becomes parallel**: What would be cooperative multitasking in regular Python becomes true parallel execution across multiple machines
3. **Native Python patterns**: You use familiar `asyncio` patterns, but Flyte automatically distributes the work

This means that when you write:

```python
results = await asyncio.gather(fetch_data(1), fetch_data(2), fetch_data(3))
```

Instead of three coroutines sharing one CPU, you get three separate containers running simultaneously, each with their own CPU, memory, and resources. Flyte bridges the gap between Python's concurrency model and distributed parallel computing, allowing for massive scalability while maintaining the familiar async/await programming model.

## Iterative fanout: recursive feature elimination

Fanout isn't limited to a single parallel burst; you can fan out **repeatedly**, using the results of one round to shape the next.
A good real-world example is [recursive feature elimination (RFE)](https://github.com/flyteorg/flyte-sdk/blob/main/examples/ml/rfe.py), a feature-selection technique that repeatedly trains a model with one candidate feature held out, drops the feature whose removal least hurts the score, and repeats until a single feature remains.
Every iteration is itself a fanout: for each remaining feature, a `train` action runs in parallel with that feature dropped, scored by cross-validation.

The `train` task evaluates the model with a single feature held out and returns its cross-validated score:

```python
@worker.task
async def train(features: list[str], drop: str) -> float:
    features.remove(drop)

    X, y = fetch_california_housing(as_frame=True, return_X_y=True)
    fold = KFold(n_splits=5, random_state=42, shuffle=True)
    model = LinearRegression()

    scores = cross_val_score(estimator=model, X=X[features], y=y, cv=fold, scoring="r2")
    return float(scores.mean())
```

The `rfe` driver task runs the elimination loop.
Each round wraps its fanout in a `flyte.group` context (see [Grouping actions](./grouping-actions)) so the iterations appear as collapsible folders in the UI, and uses `asyncio.gather()` to evaluate every candidate feature in parallel:

```python
@worker.task
async def rfe():
    x, _y = fetch_california_housing(as_frame=True, return_X_y=True)
    features = list(x.columns)

    for i in range(len(features) - 1):
        with flyte.group(f"iteration-{i}"):
            runs = {feature: train(list(features), drop=feature) for feature in features}
            values = await asyncio.gather(*(runs[feature] for feature in runs))
            scores = dict(zip(runs.keys(), values))
            best = max(scores, key=scores.get)
            features.remove(best)
```

Because each `train` call becomes its own action, every iteration's candidate evaluations run as separate containers in true parallel, while grouping keeps the nested rounds organized in the run tree.

> **📝 Note**
>
> The full runnable example lives in the [Flyte SDK repository](https://github.com/flyteorg/flyte-sdk/blob/main/examples/ml/rfe.py). From a local checkout of the `flyte-sdk` repository, run it with `uv run --prerelease=allow examples/ml/rfe.py` (the command uses a repo-relative path).

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/map ===

# Mapping over inputs

`flyte.map` applies a single task to every item of one or more input iterables, running the
invocations in parallel across the cluster and yielding their results **in input order**.
It is the structured way to [fan out](./fanout) uniform work: instead of assembling a list of
coroutines by hand and passing them to `asyncio.gather`, you hand `flyte.map` the task and the
inputs and it produces one action per item.

Use `flyte.map` when every item goes through the *same* task. For fanning out across *different*
tasks, or for full control over how invocations are assembled, use `asyncio.gather`. See
[Fanout](./fanout).

## Minimal example

From a **synchronous** task, iterate the results with a plain `for` loop:

```python
from typing import List

import flyte

env = flyte.TaskEnvironment(name="map-example")

@env.task
def process(x: int) -> str:
    return f"result-{x}"

@env.task
def main(n: int) -> List[str]:
    results: List[str] = []
    for r in flyte.map(process, range(n)):
        if isinstance(r, Exception):
            raise r
        results.append(r)
    return results

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, 10)
    print(run.url)
```

Each item in `range(n)` becomes its own action, running in its own container, and the results come
back in the same order as the inputs.

## Mapping from an async task: `flyte.map.aio`

`flyte.map` returns a synchronous iterator. Inside an **async** task, use `flyte.map.aio`, which
returns an async iterator you consume with `async for`:

```python
@env.task
async def main(n: int) -> List[str]:
    results: List[str] = []
    async for r in flyte.map.aio(process, range(n)):
        if isinstance(r, Exception):
            raise r
        results.append(r)
    return results
```

`flyte.map.aio` works over both async and sync tasks, so you can call an existing synchronous task
in parallel from an async context without rewriting it: useful when migrating a Flyte 1.x
`map_task` or integrating legacy sync code.

## Signature and parameters

```python
flyte.map(
    func,               # the task (or functools.partial) to apply to each item
    *args,              # one or more iterables, zipped item-by-item into func's arguments
    group_name=None,    # optional name for the group of mapped actions (UI grouping)
    concurrency=0,      # max actions in flight at once; 0 means unbounded (all at once)
    return_exceptions=True,
)
```

- **`func`**: the task to map. It receives one item per invocation. To hold some arguments
  constant across the map, wrap it with `functools.partial` (see **Tasks > Build tasks > Mapping over inputs > Binding constant arguments with `functools.partial`**).
- **`*args`**: one or more input iterables. With multiple iterables they are **zipped**: the
  *i*-th invocation receives the *i*-th element of each, matching `func`'s positional parameters in
  order.
- **`group_name`**: groups the resulting actions under a single label in the UI (see
  [Grouping actions](./grouping-actions)).
- **`concurrency`**: the maximum number of actions in flight at any moment. `0` (the default)
  submits everything at once. A positive value bounds the fan-out with a worker pool, so memory
  stays proportional to `concurrency` rather than to the total number of items. See
  [Controlling parallel execution](./controlling-parallelism).
- **`return_exceptions`**: when `True` (the default), a failed invocation yields the raised
  exception as its result instead of aborting the whole map; check each result with
  `isinstance(r, Exception)`. When `False`, the first failure stops iteration and raises.

Results are always yielded **in the order of the inputs**, regardless of the order in which the
individual actions finish.

## Limiting concurrency

For rate-limited APIs, GPU quotas, or connection limits, cap how many actions run at once with the
`concurrency` parameter:

```python
async for r in flyte.map.aio(call_llm_api, prompts, concurrency=3):
    ...
```

Only three actions are in flight at a time; as each completes, the next input is submitted. For a
full comparison of `flyte.map(concurrency=N)` against `asyncio.Semaphore`, see
[Controlling parallel execution](./controlling-parallelism).

## Handling errors

By default (`return_exceptions=True`) the map runs to completion even if some invocations fail, and
each failure surfaces as an exception object in the results stream:

```python
@env.task
def maybe_fail(x: int) -> str:
    if x == 2:
        raise ValueError("bad input")
    return f"ok-{x}"

@env.task
def main(n: int) -> None:
    for r in flyte.map(maybe_fail, range(n)):
        if isinstance(r, Exception):
            print(f"error: {r}")
        else:
            print(r)
```

Set `return_exceptions=False` to fail fast instead: iteration raises on the first failed action.

## Binding constant arguments with `functools.partial`

Often you want to map over one argument while holding others constant. Bind the constants with
`functools.partial`, leaving exactly one parameter free. That's the one `flyte.map` varies:

```python
from functools import partial

import flyte

env = flyte.TaskEnvironment(name="map-partial")

@env.task
def score(compound_id: str, model_name: str, batch_id: str) -> str:
    return f"{compound_id}:{model_name}:{batch_id}"

@env.task
def main() -> None:
    compounds = [str(i) for i in range(3)]
    scorer = partial(score, model_name="v2", batch_id="run-42")
    # compound_id is the only parameter left unbound, so it is what map varies.
    results = list(flyte.map(scorer, compounds))
    print("\n".join(results))
```

`flyte.map` inserts each mapped value **positionally, right after the partial's bound positional
arguments**, and requires **exactly one** parameter to be left unbound. Above, `model_name` and
`batch_id` are bound as keywords, so the mapped value fills the first slot: `compound_id`. To vary a
*later* parameter, bind the ones before it positionally and the ones after it by keyword. For
example, `partial(score, "compound-1", batch_id="run-42")` maps `model_name`. `flyte.map` raises a
`TypeError` if more or fewer than one parameter is left unbound, or if the mapped positional slot is
also bound as a keyword.

## When to use `flyte.map`

Reach for `flyte.map` when:

- Every item goes through the **same** task.
- You want built-in, in-order result collection and per-item error capture.
- You want simple, declarative concurrency control via the `concurrency` parameter.

Use [`asyncio.gather`](./fanout) instead when you are fanning out across **different** tasks in one
batch, or when you need full control over how the coroutines are assembled. Use an
[`asyncio.Semaphore`](./controlling-parallelism) when different task types in the same batch need
different concurrency limits.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/consuming-a-message-queue ===

# Consuming a message queue

A common production pattern is a **queue consumer**: a long-running task that pulls messages
from an external message queue (such as [AWS SQS](https://aws.amazon.com/sqs/)) and processes
each message concurrently.
Flyte 2 expresses this naturally by combining three building blocks you have already seen:

- **Async tasks**: the consumer loop is an `async def` task that awaits I/O against the queue.
- [**Fanout**](./fanout): each received message is dispatched to its own `process_message` action with `asyncio.create_task()`, so processing runs in parallel across the cluster.
- [**Reusable containers**](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/reusable-containers): a `ReusePolicy` keeps a warm pool of replicas ready, so messages are processed without per-message container cold-start.

The complete, runnable source for this example, a producer (`generator.py`) and a consumer
(`processor.py`), lives in the Flyte SDK repository under
[`examples/queue-reader`](https://github.com/flyteorg/flyte-sdk/tree/main/examples/queue-reader).

> [!NOTE]
> This example relies on [reusable containers](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/reusable-containers) (`flyte.ReusePolicy`), which are only available when running your Flyte code on a Union backend.
> See [Reusable containers](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/reusable-containers) for details.

> [!NOTE]
> This example reads from AWS SQS and therefore requires an SQS queue and AWS credentials
> available to the running task (here the queue is passed as an ARN via the `QUEUE_ARN`
> environment variable). The Flyte pattern shown below applies to any external queue: swap
> the SQS client calls for your queue's client.

## The consumer

### Define the task environment

The consumer runs in a [reusable](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/reusable-containers) `TaskEnvironment`.
`replicas=3` keeps a warm pool of at least two replicas to avoid starvation while the parent
consumer task occupies one, and `idle_ttl=300` shuts the pool down after five minutes of
inactivity.
The image is built from the script's own inline dependencies with
`flyte.Image.from_uv_script`, plus the `unionai-reuse` runtime library that reusable containers
require:

```python
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte",
#    "aioboto3>=11.3.0",
#    "asyncio",
# ]
# ///

import asyncio
import json
import os
from typing import List

import aioboto3

import flyte

env = flyte.TaskEnvironment(
    name="sqs_processor",
    resources=flyte.Resources(memory="500Mi", cpu=1),
    image=flyte.Image.from_uv_script(
        __file__,
        name="flyte",
    ).with_pip_packages("unionai-reuse>=0.1.3"),
    reusable=flyte.ReusePolicy(
        replicas=3,  # 1 for the consumer loop + 2 workers, so processing never starves
        idle_ttl=300,  # Idle time to keep the task environment alive
    ),
)

# The queue is passed as an ARN via the QUEUE_ARN environment variable.
DEFAULT_QUEUE_ARN = os.getenv("QUEUE_ARN")

def get_queue_url_from_arn(queue_arn: str) -> str:
    """Convert an SQS ARN to a queue URL."""
    parts = queue_arn.split(":")
    region = parts[3]
    account = parts[4]
    queue_name = parts[5]

    return f"https://sqs.{region}.amazonaws.com/{account}/{queue_name}"
```

### Process a single message

Each message is handled by its own task. These tasks run in parallel across the reusable pool,
bounded by the number of worker replicas: with `replicas=3` and the default `concurrency=1`, the
parent consumer loop occupies one replica and the other two each process a single message at a
time, so about two messages are handled concurrently. To let a single replica handle more than one
message at once, raise `concurrency` above 1:

```python
@env.task
async def process_message(message: dict) -> str:
    """Process a single message asynchronously and return the extracted word."""
    body = json.loads(message["Body"])
    word = body.get("word", "unknown")
    print(f"Task Processing message {body.get('message_id')}: {word}")
    return word
```

### The consumer loop

The driver task long-polls the queue, and for each message it receives it **dispatches a
`process_message` action with `asyncio.create_task()`** rather than awaiting it inline. This
is what fans the work out in parallel. It deletes each message once processing has started, then
awaits all dispatched tasks with `asyncio.gather()`:

```python
@env.task
async def main(queue_arn: str = DEFAULT_QUEUE_ARN, max_messages: int = 10) -> List[str]:
    queue_url = get_queue_url_from_arn(queue_arn)
    session = aioboto3.Session(region_name="us-east-2")

    results = []
    tasks = []
    messages_received = 0

    async with session.client("sqs") as sqs:
        while messages_received < max_messages:
            response = await sqs.receive_message(
                QueueUrl=queue_url,
                AttributeNames=["All"],
                MaxNumberOfMessages=1,   # one message at a time
                WaitTimeSeconds=20,      # long-polling timeout (max 20 seconds)
            )

            messages = response.get("Messages", [])
            if not messages:
                continue

            message = messages[0]
            messages_received += 1

            # Fan out: dispatch processing as a parallel action.
            process_task = asyncio.create_task(process_message(message))
            tasks.append(process_task)

            # Delete the message once we've started processing it.
            await sqs.delete_message(QueueUrl=queue_url, ReceiptHandle=message["ReceiptHandle"])

    # Wait for all dispatched processing tasks to complete.
    if tasks:
        completed_tasks = await asyncio.gather(*tasks)
        results.extend(completed_tasks)

    return results
```

### Run it

Initialize Flyte from your config and run the consumer remotely:

```python
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, queue_arn=DEFAULT_QUEUE_ARN, max_messages=10)
    print(run.url)
```

## The producer

To exercise the consumer, the example includes a standalone
[`generator.py`](https://github.com/flyteorg/flyte-sdk/tree/main/examples/queue-reader) that
pushes ten JSON messages onto the same SQS queue with `boto3`. It is an ordinary Python script,
not a Flyte task. Any producer that writes to the queue will do.

## Notes and gotchas

- **Delete after receive, not after processing completes.** The example deletes each message as
  soon as it dispatches the processing task. If a `process_message` action can fail and you need
  at-least-once semantics, delete the message only after the task succeeds instead.
- **`max_messages` bounds the run.** The consumer loop here stops after `max_messages`. For a
  continuously running consumer, drive it on a [trigger](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/triggers) or remove
  the bound and manage the task lifecycle explicitly.
- **Reusable containers require a Union backend.** See
  [Reusable containers](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/reusable-containers) for the `ReusePolicy` parameters
  (`replicas`, `concurrency`, `idle_ttl`, `scaledown_ttl`) and their capacity math.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/controlling-parallelism ===

# Controlling parallel execution

When you [fan out](./fanout) to many tasks, you often need to limit how many run at the same time.
Common reasons include rate-limited APIs, GPU quotas, database connection limits, or simply avoiding overwhelming a downstream service.

Flyte 2 provides two ways to control concurrency:
[`asyncio.Semaphore`](https://docs.python.org/3/library/asyncio-sync.html#asyncio.Semaphore) for fine-grained control,
and `flyte.map` with a built-in `concurrency` parameter for simpler cases.

## The problem: unbounded parallelism

A straightforward `asyncio.gather` launches every task at once.
If you are calling an external API that allows only a few concurrent requests, this can cause throttling or errors:

```
import asyncio

import flyte

env = flyte.TaskEnvironment("controlling_parallelism")

@env.task
async def call_llm_api(prompt: str) -> str:
    """Simulate calling a rate-limited LLM API."""
    # In a real workflow, this would call an external API.
    # The API might allow only a few concurrent requests.
    await asyncio.sleep(0.5)
    return f"Response to: {prompt}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/controlling-parallelism/controlling_parallelism.py*

```
@env.task
async def process_all_at_once(prompts: list[str]) -> list[str]:
    """Send all requests in parallel with no concurrency limit.

    This can overwhelm a rate-limited API, causing errors or throttling.
    """
    results = await asyncio.gather(*[call_llm_api(p) for p in prompts])
    return list(results)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/controlling-parallelism/controlling_parallelism.py*

With eight prompts, this fires eight concurrent API calls.
That works fine when there are no limits, but will fail when the API enforces a concurrency cap.

## Using asyncio.Semaphore

An `asyncio.Semaphore` acts as a gate: only a fixed number of tasks can pass through at a time.
The rest wait until a slot opens up.

```
@env.task
async def process_batch_with_semaphore(
    prompts: list[str],
    max_concurrent: int = 3,
) -> list[str]:
    """Process prompts in parallel, limiting concurrency with a semaphore.

    At most `max_concurrent` calls to the API run at any given time.
    The remaining tasks wait until a slot is available.
    """
    semaphore = asyncio.Semaphore(max_concurrent)

    async def limited_call(prompt: str) -> str:
        async with semaphore:
            return await call_llm_api(prompt)

    results = await asyncio.gather(*[limited_call(p) for p in prompts])
    return list(results)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/controlling-parallelism/controlling_parallelism.py*

The pattern is:

1. Create a semaphore with the desired limit.
2. Wrap each task call in an inner async function that acquires the semaphore before calling and releases it after.
3. Pass all wrapped calls to `asyncio.gather`.

All eight tasks are submitted immediately, but the Flyte orchestrator only allows three to run in parallel.
As each one completes, the next waiting task starts.

> [!NOTE]
> The semaphore controls how many tasks execute concurrently on the Flyte cluster.
> Each task still runs in its own container with its own resources: the semaphore simply limits how many containers are active at a time.

## Using flyte.map with concurrency

For uniform work (applying the same task to a list of inputs), `flyte.map` with the `concurrency` parameter is simpler:

```
@env.task
async def process_batch_with_map(prompts: list[str]) -> list[str]:
    """Process prompts using flyte.map with a built-in concurrency limit.

    This is the simplest approach when every item goes through the same task.
    """
    results = list(flyte.map(call_llm_api, prompts, concurrency=3))
    return results
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/controlling-parallelism/controlling_parallelism.py*

This achieves the same concurrency limit with less boilerplate.

For the full `flyte.map` treatment (signature, return order, error handling, and partials), see
[Mapping over inputs](./map).

## Running the example

```
if __name__ == "__main__":
    flyte.init_from_config()
    prompts = [
        "Summarize this text",
        "Translate to French",
        "Extract key points",
        "Generate a title",
        "Write a conclusion",
        "List the main topics",
        "Identify the tone",
        "Suggest improvements",
    ]
    r = flyte.run(process_batch_with_semaphore, prompts)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/controlling-parallelism/controlling_parallelism.py*

## When to use each approach

Use **`flyte.map(concurrency=N)`** when:

- Every item goes through the same task.
- You want the simplest possible code.

Use **`asyncio.Semaphore`** when:

- You need different concurrency limits for different task types within the same workflow.
- You want to combine concurrency control with error handling (e.g., `asyncio.gather(*tasks, return_exceptions=True)`).
- You are calling multiple different tasks in one parallel batch.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/streaming-map-reduce ===

# Streaming map-reduce

When you [fan out](./fanout) with [`asyncio.gather`](https://docs.python.org/3/library/asyncio-task.html#asyncio.gather), you wait for **every** task to finish before doing anything with the results.
For a map-reduce workload that is wasteful: the reduce step sits idle until the slowest mapper returns.

A better pattern is to process results **as they complete** — accumulating them into batches and kicking off reduce operations incrementally, while the remaining map tasks are still running.
This is a *gradual* (or *streaming*) map-reduce, and it is built on the standard-library [`asyncio.as_completed`](https://docs.python.org/3/library/asyncio-task.html#asyncio.as_completed) function.

## When to use it

Reach for streaming map-reduce when:

- Map tasks have **uneven durations**, so waiting for the slowest one wastes time the faster ones could spend reducing.
- You are processing a **large number of items** and want to reduce in batches rather than holding every intermediate result in memory at once.
- The reduce step is **associative** — batch results can themselves be reduced into a final result (counts, sums, aggregations, embeddings, inference outputs).

If you simply need all results before a single reduce, plain `asyncio.gather` (see [Fanout](./fanout)) is simpler.
If your goal is to *cap* how many map tasks run at once, see [Controlling parallel execution](./controlling-parallelism); the two patterns compose.

## Example

We define an environment and two tasks: one that maps over a single item, and one that reduces a batch of results.

```python
import asyncio
import random

import flyte

env = flyte.TaskEnvironment(
    name="streaming_map_reduce",
    resources=flyte.Resources(cpu="1"),
)

@env.task
async def process_item(item: str) -> str:
    print(f"Processing {item}", flush=True)
    # Simulate varying processing times so results finish out of order.
    await asyncio.sleep(random.uniform(1, 5))
    return f"processed_{item}"

@env.task
async def reduce_batch(items: list[str]) -> str:
    print(f"Reducing batch of {len(items)} items")
    return f"reduced_batch_of_{len(items)}_items"
```

### The driver task

The driver fans out all the map tasks up front, then walks the results in completion order with `asyncio.as_completed`.
Each time a batch fills up, it launches a `reduce_batch` action **without blocking** — the loop keeps consuming newly completed map results while the reduce runs.

```python
@env.task
async def streaming_reduce_processing() -> str:
    input_items = [f"item_{i}" for i in range(100)]

    # Fan out: start every item task immediately.
    tasks = [asyncio.create_task(process_item(item)) for item in input_items]

    batch_size = 10
    accumulated_values: list[str] = []
    reducers: list[asyncio.Task] = []

    print(f"Started {len(tasks)} tasks, will reduce in batches of {batch_size}")

    # Consume results as each task finishes, rather than waiting for all of them.
    for task in asyncio.as_completed(tasks):
        result = await task
        accumulated_values.append(result)

        # Once a batch has accumulated, kick off a reduce without blocking the loop.
        if len(accumulated_values) >= batch_size:
            print(f"Triggering reduce for batch of {len(accumulated_values)}")
            reducer_task = asyncio.create_task(reduce_batch(accumulated_values.copy()))
            reducers.append(reducer_task)
            accumulated_values.clear()

    # Reduce any stragglers that did not fill a full batch.
    if accumulated_values:
        print(f"Handling final batch of {len(accumulated_values)} stragglers")
        reducers.append(asyncio.create_task(reduce_batch(accumulated_values)))

    # Wait for every batch reduce to finish.
    reduced_results = await asyncio.gather(*reducers)

    # Combine the batch results into a single final result.
    final_result = await reduce_batch(reduced_results)

    print(f"Completed {len(reducers)} reduce operations, final result: {final_result}")
    return final_result
```

### Running the example

```python
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(streaming_reduce_processing)
    print(run.url)
```

## How it works

The key building blocks are all standard `asyncio`:

1. **`asyncio.create_task(process_item(item))`** schedules each map action. Because `process_item` is a Flyte task, each of these runs in its own container on the cluster — the fanout is real distributed parallelism, not single-machine concurrency (see [Fanout](./fanout) for how Flyte turns `asyncio` into distributed execution).
2. **`asyncio.as_completed(tasks)`** yields the task handles in the order they *finish*, not the order they were submitted. This is what lets the driver react to the fastest map results first.
3. **`asyncio.create_task(reduce_batch(...))`** launches each reduce as its own Flyte action and appends it to `reducers` without awaiting it, so map consumption and reduction overlap.
4. **`asyncio.gather(*reducers)`** joins all the in-flight batch reduces before the final combine step.

The result is a pipeline where reduce work begins as soon as the first batch of map results is ready, instead of after the last map task returns.

> [!NOTE]
> `as_completed` returns awaitables in completion order but gives you no control over *how many* map tasks run at once — it schedules all of them.
> To bound the map fanout as well, combine this pattern with an `asyncio.Semaphore` or `flyte.map(concurrency=...)` from [Controlling parallel execution](./controlling-parallelism).

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/structured-concurrency-anyio ===

# Structured concurrency with anyio

Flyte builds a task's dependency graph from what you `await` — not from any particular async library.
`asyncio` is the default and the one used throughout the [Fanout](./fanout), [Controlling parallel execution](./controlling-parallelism), and [Task dependencies and ordering](./task-dependencies) guides, but it is not the only option.
Any structured-concurrency runtime that drives coroutines works, and [`anyio`](https://anyio.readthedocs.io/) is a popular one.
Its **task groups** give you a top-level alternative to raw `asyncio.gather` / `asyncio.create_task`, with clearer lifetime and error-propagation semantics.

Use `anyio` when you want structured concurrency — a scope that owns the tasks it spawns, waits for all of them on exit, and cancels the siblings automatically if one fails — instead of tracking `asyncio.create_task` handles by hand.

## The task-group pattern

An `anyio` task group is an `async with` block. You spawn work into it with `start_soon`, and the block does not exit until every spawned task has finished. Because task groups don't return the spawned tasks' values directly, this example uses [`aioresult`](https://aioresult.readthedocs.io/)'s `ResultCapture` to collect each result.

We define a reusable environment and a simple per-item task. `anyio` and `aioresult` are ordinary pip dependencies, so we add them to the image:

```python
from dataclasses import dataclass

import aioresult
import anyio

import flyte

env = flyte.TaskEnvironment(
    name="anyio_batch",
    resources=flyte.Resources(cpu="1"),
    image=flyte.Image.from_debian_base(name="anyio").with_pip_packages("anyio", "aioresult"),
)

@dataclass
class InferenceRequest:
    feature_a: float
    feature_b: float

@env.task
async def predict_one(request: InferenceRequest) -> float:
    # A dummy linear model: 2 * feature_a + 3 * feature_b + bias(=1.0)
    return 2.0 * request.feature_a + 3.0 * request.feature_b + 1.0
```

The driver task fans out over the batch inside a task group:

```python
@env.task
async def predict_batch(requests: list[InferenceRequest]) -> list[float]:
    captured = []
    async with anyio.create_task_group() as tg:
        # Start each prediction; they run at the same time.
        for req in requests:
            captured.append(aioresult.ResultCapture.start_soon(tg, predict_one, req))
    # The `async with` block has exited, so every task has completed.
    return [c.result() for c in captured]
```

What happens here mirrors an `asyncio.gather` fanout, but with structured-concurrency guarantees:

1. **`start_soon` schedules each `predict_one`** into the group. As with `asyncio`, Flyte runs each action in its own container, so the batch executes in true parallel across the cluster — the runtime you use to express concurrency does not change how Flyte distributes the work.
2. **Leaving the `async with` block is the fan-in edge.** The group blocks until all spawned tasks finish, exactly as `await asyncio.gather(...)` would. `predict_batch` cannot return until every prediction is in.
3. **`ResultCapture` collects the return values**, which you read with `.result()` after the group closes.

> [!NOTE]
> Task groups give you cancellation for free: if any task in the group raises, `anyio` cancels the remaining siblings and propagates the error out of the `async with` block. You get the "cancel the rest on failure" behavior that requires manual `.cancel()` bookkeeping with `asyncio` (see [Abort and cancel actions](./abort-tasks#canceling-actions-programmatically)).

## When to use anyio

Reach for `anyio` when:

- You want **structured concurrency** — spawned work is scoped to a block, awaited on exit, and cancelled together on error — rather than manually pairing `asyncio.create_task` handles with `asyncio.gather`.
- Your code (or a library you depend on) already uses `anyio` or `trio`, and you want one consistent concurrency model.

Stay with `asyncio` when:

- You just need to fan out and collect results — `await asyncio.gather(...)` is simpler (see [Fanout](./fanout)).
- You need fine-grained, dependency-driven scheduling where different consumers await different producers (see [Task dependencies and ordering](./task-dependencies)).

Either way, the underlying model is the same: Flyte reads the dependency graph from your `await`s and turns concurrent coroutines into distributed parallel actions.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/task-dependencies ===

# Task dependencies and ordering

Flyte 1 built a workflow's DAG (directed acyclic graph) explicitly: you declared nodes and wired their edges with the `>>` operator or `create_node`.
Flyte 2 has no such API. Instead, **the dependency graph is inferred from the data you `await`**.
When you await one task's result and pass it into another, Flyte records the edge; tasks that share no data run independently.
This page shows how to express the ordering patterns you used to build by hand — sequencing, fan-out, fan-in, and fine-grained dependency-driven scheduling — using ordinary Python `asyncio`.

If you are coming from Flyte 1, read [Parallelism and fan-out](https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/parallelism) first for the migration mapping.

## The dependency graph is implicit

There is nothing special to learn: a data dependency *is* the edge.

```python
import asyncio

import flyte

env = flyte.TaskEnvironment(name="pipeline")

@env.task
async def extract() -> str:
    return "raw"

@env.task
async def transform(data: str) -> str:
    return f"transformed({data})"

@env.task
async def load(data: str) -> str:
    return f"loaded({data})"

@env.task
async def main() -> str:
    raw = await extract()          # runs first
    clean = await transform(raw)   # waits for extract — it consumes `raw`
    return await load(clean)       # waits for transform — it consumes `clean`
```

Each `await` means "wait for this to finish before continuing," so a chain of `await`s that pass results downstream runs sequentially — exactly like a linear Flyte 1 workflow. You never declare the edges; passing `raw` into `transform` and `clean` into `load` *is* the DAG.

## Ordering without a data dependency

Sometimes you need task `B` to run after task `A` even though `B` does not consume `A`'s output — for example, `A` writes to a store that `B` reads out-of-band, or `A` must finish before you send a notification.

Because ordering comes from `await`, you force it simply by awaiting `A` before invoking `B`:

```python
@env.task
async def main() -> str:
    result = await transform(await extract())
    await notify()   # runs after transform completes, though it uses no upstream output
    return result
```

You do not need a special "run after" construct — a preceding `await` is the ordering primitive.

## Fan-out and fan-in

**Fan-out** launches independent tasks concurrently; **fan-in** collects their results into a single downstream task.
Use [`asyncio.gather`](https://docs.python.org/3/library/asyncio-task.html#asyncio.gather) to await several tasks at once — Flyte runs each in its own container in parallel (see [Fanout](./fanout)):

```python
@env.task
async def combine(a: str, b: str, c: str) -> str:
    return f"{a} + {b} + {c}"

@env.task
async def main() -> str:
    # Fan-out: three tasks start together and run in parallel
    a, b, c = await asyncio.gather(extract(), extract(), extract())
    # Fan-in: one task depends on all three results
    return await combine(a, b, c)
```

The `await asyncio.gather(...)` establishes the fan-in edge: `combine` cannot start until all three upstream tasks have produced their results.

## Dependency-driven scheduling

The pattern that most often motivates "replicating DAG behavior" is **fine-grained scheduling**: a diamond or fork where each downstream task should start the moment *its own* upstreams finish, without waiting for unrelated slow tasks.

Consider three producers of different durations and four consumers with different dependencies:

- `needs_short` depends on `short` only
- `needs_medium` depends on `medium` only
- `needs_long` depends on `long` only
- `needs_all` depends on all three

A single `await asyncio.gather(short, medium, long)` before starting any consumer would make every consumer wait for the slowest producer. To let each consumer start as early as possible, start the producers as [`asyncio.create_task`](https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task) handles, then wrap each consumer in a small helper coroutine that awaits only the handles it needs. Launch all the helpers together with `asyncio.gather`:

```python
@env.task
async def long_producer() -> str:
    await asyncio.sleep(10)
    return "long"

@env.task
async def medium_producer() -> str:
    await asyncio.sleep(3)
    return "medium"

@env.task
async def short_producer() -> str:
    await asyncio.sleep(1)
    return "short"

@env.task
async def needs_short(x: str) -> str:
    return f"needs_short({x})"

@env.task
async def needs_medium(x: str) -> str:
    return f"needs_medium({x})"

@env.task
async def needs_long(x: str) -> str:
    return f"needs_long({x})"

@env.task
async def needs_all(long: str, medium: str, short: str) -> str:
    return f"needs_all({long}, {medium}, {short})"

@env.task
async def main() -> str:
    # Start all producers concurrently and keep their handles
    long_task = asyncio.create_task(long_producer())
    medium_task = asyncio.create_task(medium_producer())
    short_task = asyncio.create_task(short_producer())

    # Each helper awaits only the producers it actually depends on,
    # so it starts as soon as those specific producers finish.
    async def run_needs_short() -> str:
        return await needs_short(await short_task)

    async def run_needs_medium() -> str:
        return await needs_medium(await medium_task)

    async def run_needs_long() -> str:
        return await needs_long(await long_task)

    async def run_needs_all() -> str:
        long_r, medium_r, short_r = await asyncio.gather(long_task, medium_task, short_task)
        return await needs_all(long_r, medium_r, short_r)

    # Launch every branch concurrently; each resolves on its own dependencies
    results = await asyncio.gather(
        run_needs_short(),   # starts after ~1s
        run_needs_medium(),  # starts after ~3s
        run_needs_long(),    # starts after ~10s
        run_needs_all(),     # starts after ~10s
    )
    return str(results)
```

`needs_short` starts about a second in, as soon as `short_producer` returns — it does not wait for the 10-second `long_producer`. Awaiting an `asyncio` task handle more than once is safe: the handle caches its result, so `long_task` can feed both `run_needs_long` and `run_needs_all` without re-running the producer.

> [!NOTE]
> Reach for helper coroutines that each `await` their specific handles, rather than a manual completion loop that inspects [`asyncio.as_completed`](https://docs.python.org/3/library/asyncio-task.html#asyncio.as_completed) and dispatches downstream tasks by hand. Hand-rolled dispatch loops are easy to get wrong — a mis-tracked "has this fired yet?" check can launch the same downstream task twice. Let the dependency edges fall out of `await` instead.

## When to reach for `as_completed`

Use [`asyncio.as_completed`](https://docs.python.org/3/library/asyncio-task.html#asyncio.as_completed) when you want to process results **in completion order** — for example, streaming each result into a running reduction as it lands — rather than to encode a fixed dependency graph:

```python
@env.task
async def main() -> list[str]:
    tasks = [asyncio.create_task(short_producer()) for _ in range(10)]
    processed = []
    for task in asyncio.as_completed(tasks):
        result = await task
        processed.append(result)  # handle each result the moment it arrives
    return processed
```

For a worked streaming/reduce example, see [Fanout](./fanout) and [Controlling parallel execution](./controlling-parallelism).

## Summary

- Flyte 2 has no explicit DAG-construction API; dependencies come from the data you `await`.
- Sequence tasks by awaiting them in order — a preceding `await` orders even tasks that share no data.
- Fan out with `asyncio.gather`; fan in by awaiting several results into one downstream task.
- For fine-grained scheduling, keep producer handles from `asyncio.create_task` and have each consumer await only the handles it depends on, so it starts as early as possible.
- Prefer letting `await` express the graph over hand-rolled completion-tracking loops.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/conditions ===

# External conditions

An **external condition** is a first-class action that pauses a running task until an external
signal arrives: a human approval, a callback from an external service, or a value supplied at
runtime. The paused action stays observable, resumable, and governable like any other action, so
you no longer need polling loops or side processes to wait on something the workflow can't produce
itself.

Inside a task, `await flyte.new_condition.aio(...)` registers a condition action and returns a
handle; `await handle.wait.aio()` blocks the task until the condition is signaled and returns the
typed payload. (`new_condition` and `wait` are sync-by-default; in an `async def` task use their
`.aio()` form.)

## Supported types

A condition declares a `data_type`, which determines what a signal must supply and what `wait()`
returns:

| `data_type` | `wait()` returns | A signal value of |
|---|---|---|
| `bool` (default) | `True` / `False` | `true` / `false` |
| `int` | Python `int` | an integer literal |
| `float` | Python `float` | a decimal literal |
| `str` | Python `str` | any string |

## Example: human approval

A typed approval gate with a timeout (the most common use case):

```python
from datetime import timedelta
import flyte

env = flyte.TaskEnvironment("approvals")

@env.task
async def etl_pipeline():
    staged = await transform()

    approval = await flyte.new_condition.aio(
        "prod_write_approval",
        prompt="Approve writing staged data to production?",
        data_type=bool,
        timeout=timedelta(hours=24),
    )
    if not await approval.wait.aio():
        raise RuntimeError("Pipeline rejected by reviewer")

    await write_to_prod(staged)
```

The task pauses at `await approval.wait.aio()` until someone signals the condition (see
**Tasks > Build tasks > External conditions > Signaling a condition**). If the timeout elapses with no signal, `wait()`
raises `flyte.errors.ConditionTimedoutError`.

## Example: string input at runtime

A condition can collect a typed value (not just a yes/no) and feed it back into the workflow.
Here the task waits for a free-form string before continuing:

```python
import flyte

env = flyte.TaskEnvironment("conditions")

@env.task
async def deploy_with_reason():
    reason = await flyte.new_condition.aio(
        "deploy_reason",
        prompt="Enter a deployment reason to continue:",
        data_type=str,
    )
    note: str = await reason.wait.aio()
    # `note` now holds the string a human supplied — use it downstream.
    await record_audit(note)
```

## Parameters

```python
flyte.new_condition(
    name,
    prompt="Approve?",
    prompt_type="text",
    data_type=bool,
    description="",
    timeout=None,
    webhook=None,
)
```

| Parameter | Type | Default | Description |
|---|---|---|---|
| `name` | `str` | required | Identifier for the condition within the parent action. Signal it with this name (`flyte signal condition <run> <name>`) or look it up with `flyte.remote.Condition.get("<name>", ...)`. |
| `prompt` | `str` | `"Approve?"` | Human-readable text shown in the UI signal form. |
| `prompt_type` | `"text"` \| `"markdown"` | `"text"` | How the prompt is rendered. |
| `data_type` | `type` | `bool` | Payload type: one of `bool`, `int`, `float`, `str`. Determines what `wait()` returns and what a signal must supply. |
| `description` | `str` | `""` | Longer explanation rendered alongside the prompt. |
| `timeout` | `timedelta` \| `int` \| `float` \| `None` | `None` | Maximum wait. If it elapses with no signal, `wait()` raises `flyte.errors.ConditionTimedoutError`. |

An optional advanced `webhook` parameter accepts a `flyte.ConditionWebhook` so the backend POSTs a
callback URL when the condition is created; see the API reference for details.

## Signaling a condition

A condition is satisfied by delivering exactly one typed signal of its declared `data_type`.

### From the CLI

```shell
# List the conditions on a run, optionally scoped to one parent action.
flyte get condition <run-name>
flyte get condition <run-name> <parent-action>

# Signal a specific condition by its name. Omit the value for an interactive typed prompt.
flyte signal condition <run-name> <condition-name> true
```

The value is coerced to the condition's declared `data_type` (`true`/`false` for `bool`, integer
literals for `int`, decimal literals for `float`, any string for `str`).

### From Python (remote)

```python
import flyte.remote

condition = flyte.remote.Condition.get(
    "prod_write_approval",
    run_name="deploy-run-123",
    action_name="etl_pipeline",
)
condition.signal(True)
```

`flyte.remote.Condition.listall(run_name=...)` enumerates the conditions on a run.

## Timeout with a fallback

```python
from datetime import timedelta
import flyte
from flyte.errors import ConditionTimedoutError

env = flyte.TaskEnvironment("conditions")

@env.task
async def with_default():
    cond = await flyte.new_condition.aio("threshold", data_type=float, timeout=timedelta(minutes=30))
    try:
        threshold = await cond.wait.aio()
    except ConditionTimedoutError:
        threshold = 0.5   # proceed with a default if no one responds
```

## Errors

| Situation | Raised |
|---|---|
| Timeout elapses before a signal | `flyte.errors.ConditionTimedoutError` |
| Creating a condition whose `name` already exists in the action | `flyte.errors.ConditionAlreadyExistsError` |
| Condition fails during execution | `flyte.errors.ConditionFailedError` |
| Signal value doesn't match `data_type` | `TypeError` (client-side, before any call) |
| `wait()` called outside a task context | `RuntimeError` |

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/unit-testing ===

Unit testing is essential for ensuring your Flyte tasks work correctly. Flyte 2.0 provides flexible testing approaches that allow you to test both your business logic and Flyte-specific features like type transformations and caching.

## Understanding task invocation

When working with functions decorated with `@env.task`, there are two ways to invoke them, each with different behavior:

### Direct function invocation

When you call a task directly like a regular Python function:

```python
result = my_task(x=10, y=20)
```

**Flyte features are NOT invoked**, including:

- Type transformations and serialization
- Caching
- Data validation

This behaves exactly like calling a regular Python function, making it ideal for testing your business logic.

### Using `flyte.run()`

When you invoke a task using `flyte.run()`:

```python
run = flyte.run(my_task, x=10, y=20)
result = run.outputs()
```

**Flyte features ARE invoked**, including:

- Type transformations and serialization
- Data validation
- Type checking (raises `flyte.errors` if types are not supported or restricted)

This allows you to test Flyte-specific behavior like serialization and caching.

## Testing business logic

For most unit tests, you want to verify your business logic works correctly. Use **direct function invocation** for this:

```python
import flyte

env = flyte.TaskEnvironment("my_env")

@env.task
def add(a: int, b: int) -> int:
    return a + b

def test_add():
    result = add(a=3, b=5)
    assert result == 8
```

### Testing async tasks

Async tasks work the same way with direct invocation:

```python
import pytest

@env.task
async def subtract(a: int, b: int) -> int:
    return a - b

@pytest.mark.asyncio
async def test_subtract():
    result = await subtract(a=10, b=4)
    assert result == 6
```

### Testing nested tasks

When tasks call other tasks, direct invocation continues to work without any Flyte overhead:

```python
@env.task
def nested(a: int, b: int) -> int:
    return add(a, b)  # Calls the add task directly

def test_nested():
    result = nested(3, 5)
    assert result == 8
```

## Testing type transformations and serialization

When you need to test how Flyte handles data types, serialization, or caching, use `flyte.run()`:

```python
@pytest.mark.asyncio
async def test_add_with_flyte_run():
    run = flyte.run(add, 3, 5)
    assert run.outputs() == 8
```

### Testing type restrictions

Some types may not be supported or may be restricted. Use `flyte.run()` to test that these restrictions are enforced:

```python
from typing import Tuple
import flyte.errors

@env.task
def not_supported_types(x: Tuple[str, str]) -> str:
    return x[0]

@pytest.mark.asyncio
async def test_not_supported_types():
    # Direct invocation works fine
    result = not_supported_types(x=("a", "b"))
    assert result == "a"

    # flyte.run enforces type restrictions
    with pytest.raises(flyte.errors.RestrictedTypeError):
        flyte.run(not_supported_types, x=("a", "b"))
```

### Testing nested tasks with serialization

You can also test nested task execution with Flyte's full machinery:

```python
@pytest.mark.asyncio
async def test_nested_with_run():
    run = flyte.run(nested, 3, 5)
    assert run.outputs() == 8
```

## Testing traced functions

Functions decorated with `@flyte.trace` can be tested similarly to tasks:

```python
@flyte.trace
async def traced_multiply(a: int, b: int) -> int:
    return a * b

@pytest.mark.asyncio
async def test_traced_multiply():
    result = await traced_multiply(a=6, b=7)
    assert result == 42
```

## Best practices

1. **Test logic with direct invocation**: For most unit tests, call tasks directly to test your business logic without Flyte overhead.

2. **Test serialization with `flyte.run()`**: Use `flyte.run()` when you need to verify:
   - Type transformations work correctly
   - Data serialization/deserialization
   - Caching behavior
   - Type restrictions are enforced

3. **Use standard testing frameworks**: Flyte tasks work with pytest, unittest, and other Python testing frameworks.

4. **Test async tasks properly**: Use `@pytest.mark.asyncio` for async tasks and await their results.

5. **Mock external dependencies**: Use standard Python mocking techniques for external services, databases, etc.

## Quick reference

| Test Scenario | Method | Example |
|--------------|--------|---------|
| Business logic (sync) | Direct call | `result = task(x=10)` |
| Business logic (async) | Direct await | `result = await task(x=10)` |
| Type transformations | `flyte.run()` | `r = flyte.run(task, x=10)` |
| Data serialization | `flyte.run()` | `r = flyte.run(task, x=10)` |
| Caching behavior | `flyte.run()` | `r = flyte.run(task, x=10)` |
| Type restrictions | `flyte.run()` + pytest.raises | `pytest.raises(flyte.errors.RestrictedTypeError)` |

## Example test suite

Here's a complete example showing different testing approaches:

```python
import pytest
import flyte
import flyte.errors

env = flyte.TaskEnvironment("test_env")

@env.task
def add(a: int, b: int) -> int:
    return a + b

@env.task
async def subtract(a: int, b: int) -> int:
    return a - b

# Test business logic directly
def test_add_logic():
    result = add(a=3, b=5)
    assert result == 8

@pytest.mark.asyncio
async def test_subtract_logic():
    result = await subtract(a=10, b=4)
    assert result == 6

# Test with Flyte serialization
@pytest.mark.asyncio
async def test_add_serialization():
    run = flyte.run(add, 3, 5)
    assert run.outputs() == 8

@pytest.mark.asyncio
async def test_subtract_serialization():
    run = flyte.run(subtract, a=10, b=4)
    assert run.outputs() == 6
```

## Future improvements

The Flyte SDK team is actively working on improvements for advanced unit testing scenarios, particularly around initialization and setup for complex test cases. Additional utilities and patterns may be introduced in future releases to further simplify unit testing.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/higher-order-functions ===

# Higher-order functions

A *higher-order function* is a function that takes other functions as arguments or returns them. Because Flyte 2 tasks execute as native Python and can be [passed as arguments](./other-features#passing-tasks-and-functions-as-arguments) like any other callable, you can write higher-order functions that operate on **tasks themselves**: reusable orchestration components that wrap a task with retry, fallback, batching, or fault-tolerance logic, without changing the task's business logic.

This is possible because Flyte 2 workflows run as ordinary Python:

- **Tasks are callables.** You can accept a task as a parameter and `await` it, `.override(...)` its resources, or hand it to `asyncio`.
- **Arbitrary nesting.** A task can invoke other tasks at any depth, so an orchestration wrapper can drive a task from inside another task.
- **Native control flow.** Loops, conditionals, and `try`/`except` work directly on task results (task outputs are plain Python objects, not promises), so a wrapper can inspect a result or catch an exception and react.

> [!NOTE] Higher-order functions are plain functions, not tasks
> The wrappers below are **not** decorated with `@env.task`. They are regular `async` Python functions that orchestrate tasks. You call them from inside a driver task (an `@env.task`), which is where the actual task invocations happen. Keep the reusable orchestration logic in a plain function so it can be applied to any task.

The patterns on this page are drawn from the runnable [`higher_order_patterns`](https://github.com/flyteorg/flyte-sdk/tree/main/examples/higher_order_patterns) examples in the Flyte SDK repository.

## Fallback runner

Run a primary task and, if it fails with a matching exception, automatically fall back to an alternative task. Useful for degrading to a cheaper model, a different region, or a simpler algorithm when the preferred path fails.

```python
from typing import Callable, List, Optional, Type, TypeVar

R = TypeVar("R")

async def run_with_fallback(
    primary_task: Callable[..., R],
    fallback_task: Callable[..., R],
    *args,
    fallback_exceptions: Optional[List[Type[Exception]]] = None,
    **kwargs,
) -> R:
    try:
        return await primary_task(*args, **kwargs)
    except Exception as e:
        # Fall back only on the exceptions we opted into (None means any).
        should_fallback = fallback_exceptions is None or any(
            isinstance(e, exc) for exc in fallback_exceptions
        )
        if not should_fallback:
            raise
        return await fallback_task(*args, **kwargs)
```

Call it from a driver task, passing the two tasks as arguments:

```python
import flyte
import flyte.errors

env = flyte.TaskEnvironment("fallback")

@env.task
async def primary(x: int) -> int:
    # Business logic that may fail, e.g. raise ValueError(...) on bad input.
    ...

@env.task
async def backup(x: int) -> int:
    ...

@env.task
async def main(x: int) -> int:
    return await run_with_fallback(primary, backup, x, fallback_exceptions=[flyte.errors.RuntimeUserError])
```

Note the `fallback_exceptions` list holds `flyte.errors` types, not bare Python exceptions. An exception raised inside a task does not reach the parent as its original Python type: Flyte wraps it as a `flyte.errors` type (a `ValueError` raised in a task surfaces to the caller as a `flyte.errors.RuntimeUserError` whose `code` is `"ValueError"`). So `isinstance`/type-matching in a wrapper must target the `flyte.errors.*` hierarchy; matching on `ValueError` here would never fire and the fallback would never run. See [Error handling](./error-handling) for how failures propagate.

## Retry with increasing memory (OOM retrier)

Retry a task with progressively larger memory allocations when it hits an out-of-memory error, so you don't have to hard-code a worst-case memory request. The wrapper uses `.override()` to raise the task's `flyte.Resources` on each attempt and catches `flyte.errors.OOMError`.

```python
import flyte
import flyte.errors

async def retry_with_memory(
    task_fn,
    *args,
    initial_memory_mi: int = 250,
    increment_mi: int = 200,
    max_memory_mi: int = 4096,
    cpu: int = 1,
    **kwargs,
):
    current = initial_memory_mi
    while current <= max_memory_mi:
        try:
            return await task_fn.override(
                resources=flyte.Resources(cpu=cpu, memory=f"{current}Mi")
            )(*args, **kwargs)
        except flyte.errors.OOMError:
            if current >= max_memory_mi:
                break
            current = min(current + increment_mi, max_memory_mi)
    raise RuntimeError(f"Task still OOMing at {max_memory_mi}Mi")
```

Because the wrapper only takes the task and its arguments, it works with any task:

```python
@env.task
async def process(data: list[int]) -> int:
    # Business logic that may run out of memory on large inputs.
    return sum(data)

@env.task
async def main(data: list[int]) -> int:
    return await retry_with_memory(process, data, initial_memory_mi=500, max_memory_mi=8192)
```

See [Error handling](./error-handling) for more on `flyte.errors.OOMError` and resource-based recovery.

## Circuit breaker

Run a task over many items in parallel, but stop early ("open the circuit") once failures exceed a threshold, so a systemic problem doesn't burn resources on every remaining item. It launches all invocations with `asyncio.create_task`, processes them as they complete, and cancels the rest when the limit is crossed.

```python
import asyncio
from typing import Callable, List, Optional, TypeVar

T = TypeVar("T")
R = TypeVar("R")

class CircuitBreakerError(Exception):
    """Raised when too many failures occur."""

async def circuit_breaker_execute(
    task_fn: Callable[[T], R], items: List[T], max_failures: int = 3
) -> List[Optional[R]]:
    tasks = [asyncio.create_task(task_fn(item)) for item in items]
    results: List[Optional[R]] = [None] * len(items)
    failures = 0
    pending = set(tasks)

    while pending:
        done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
        for task in done:
            idx = tasks.index(task)
            if task.exception():
                failures += 1
                if failures > max_failures:
                    for remaining in pending:
                        remaining.cancel()
                    raise CircuitBreakerError(
                        f"{failures} failures exceed limit of {max_failures}"
                    )
            else:
                results[idx] = task.result()
    return results
```

Failed items come back as `None`; if the failure threshold is crossed, the remaining tasks are cancelled and `CircuitBreakerError` is raised. See [Fanout](./fanout) for the basics of running tasks in parallel and [Controlling parallel execution](./controlling-parallelism) for bounding concurrency.

## Auto batcher

Split a large input into batches, run a map task over each batch in parallel, then combine the results with a reduce step. This bounds how many invocations are in flight at once while still processing everything.

```python
import asyncio
from typing import Any, Callable, List, TypeVar

T = TypeVar("T")
R = TypeVar("R")

def create_batches(data: List[T], batch_size: int) -> List[List[T]]:
    return [data[i : i + batch_size] for i in range(0, len(data), batch_size)]

async def batch_map_reduce(
    map_fn: Callable[[T], R],
    reduce_fn: Callable[[List[R]], Any],
    data: List[T],
    batch_size: int = 10,
) -> Any:
    all_results: List[R] = []
    for batch in create_batches(data, batch_size):
        coros = [asyncio.create_task(map_fn(item)) for item in batch]
        all_results.extend(await asyncio.gather(*coros))
    return reduce_fn(all_results)
```

The map step is a task; the reduce step can be a task or a plain function:

```python
@env.task
async def square(x: int) -> int:
    return x * x

@env.task
async def main(data: list[int]) -> int:
    return await batch_map_reduce(square, sum, data, batch_size=25)
```

For a first-class parallel-map primitive, see `flyte.map` in [Fanout](./fanout).

## Composing the patterns

Because each wrapper is just a function that takes a task, you can layer them (for example, wrap a task in the OOM retrier and then hand *that* to the fallback runner) to build orchestration behavior out of small, reusable pieces without touching the underlying task code.

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/other-features ===

This section covers advanced programming patterns and techniques for working with Flyte tasks.

## Task forwarding

When one task calls another task using the normal invocation syntax (e.g., `await inner_task(x)`), Flyte creates a durable action that's recorded in the UI with data passed through the metadata store. However, if you want to execute a task in the same Python VM without this overhead, use the `.forward()` method.

**When to use**: You want to avoid durability overhead and execute task logic directly in the current VM.

```python
import flyte

env = flyte.TaskEnvironment("my-env")

@env.task
async def inner_task(x: int) -> int:
    return x + 1

@env.task
async def outer_task(x: int) -> int:
    # Executes in same VM, no durable action created
    v = await inner_task.forward(x=10)

    # Creates a durable action, recorded in UI
    return await inner_task(v)
```

The `.forward()` method works with both sync and async tasks:

```python
@env.task
def sync_inner_task(x: int) -> int:
    return x + 1

@env.task
def sync_outer_task(x: int) -> int:
    # Direct execution, no remote call
    v = sync_inner_task.forward(x=10)
    return sync_inner_task(v)
```

## Passing tasks and functions as arguments

You can pass both Flyte tasks and regular Python functions as arguments to other tasks. Flyte handles this through pickling, so the code appears as pickled data in the UI.

```python
import typing
import flyte

env = flyte.TaskEnvironment("udfs")

@env.task
async def add_one_udf(x: int) -> int:
    return x + 1

# Regular async function (not a task)
async def fn_add_two_udf(x: int) -> int:
    return x + 2

@env.task
async def run_udf(x: int, udf: typing.Callable[[int], typing.Awaitable[int]]) -> int:
    return await udf(x)

@env.task
async def main() -> list[int]:
    # Pass a Flyte task as an argument
    result_one = await run_udf(5, add_one_udf)

    # Pass a regular function as an argument
    result_two = await run_udf(5, fn_add_two_udf)

    return [result_one, result_two]
```

**Note**: Both tasks and regular functions are serialized via pickling when passed as arguments.

## Custom action names

By default, actions in the UI use the task's function name. You can provide custom, user-friendly names using the `short_name` parameter.

### Set at task definition

```python
import flyte

env = flyte.TaskEnvironment("friendly_names")

@env.task(short_name="my_task")
async def some_task() -> str:
    return "Hello, Flyte!"
```

### Override at call time

```python
@env.task(short_name="entrypoint")
async def main() -> str:
    # Uses the default short_name "my_task"
    s = await some_task()

    # Overrides to use "my_name" for this specific action
    return s + await some_task.override(short_name="my_name")()
```

This is useful when the same task is called multiple times with different contexts, making the UI more readable.

## Invoking async functions from sync tasks

When migrating from Flyte 1.x to 2.0, you may have legacy sync code that needs to call async functions. Use `nest_asyncio.apply()` to enable `asyncio.run()` within sync tasks.

```python
import asyncio
import nest_asyncio
import flyte

env = flyte.TaskEnvironment(
    "async_in_sync",
    image=flyte.Image.from_debian_base().with_pip_packages("nest_asyncio"),
)

# Apply at module level
nest_asyncio.apply()

async def async_function() -> str:
    await asyncio.sleep(1)
    return "done"

@env.task
def sync_task() -> str:
    # Now you can use asyncio.run() in a sync task
    return asyncio.run(async_function())
```

**Important**:

- Call `nest_asyncio.apply()` at the module level before defining tasks
- Add `nest_asyncio` to your image dependencies
- This is particularly useful during migration when you have mixed sync/async code

## Async and sync task interoperability

When migrating from older sync-based code to async tasks, or when working with mixed codebases, you need to call sync tasks from async parent tasks. Flyte provides the `.aio` method on every task (even sync ones) to enable this.

### Calling sync tasks from async tasks

Every sync task automatically has an `.aio` property that returns an async-compatible version:

```python
import flyte

env = flyte.TaskEnvironment("mixed-tasks")

@env.task
def sync_task(x: int) -> str:
    """Legacy sync task"""
    return f"Processed {x}"

@env.task
async def async_task(x: int) -> str:
    """New async task that calls legacy sync task"""
    # Use .aio to call sync task from async context
    result = await sync_task.aio(x)
    return result
```

### Using with `flyte.map.aio()`

When you need to call sync tasks in parallel from an async context, use `flyte.map.aio()`:

```python
from typing import List
import flyte

env = flyte.TaskEnvironment("map-example")

@env.task
def sync_process(x: int) -> str:
    """Synchronous processing task"""
    return f"Task {x}"

@env.task
async def async_main(n: int) -> List[str]:
    """Async task that maps over sync task"""
    results = []

    # Map over sync task from async context
    async for result in flyte.map.aio(sync_process, range(n)):
        if isinstance(result, Exception):
            raise result
        results.append(result)

    return results
```

**Why this matters**: This pattern is powerful when migrating from Flyte 1.x or integrating legacy sync tasks with new async code. You don't need to rewrite all sync tasks at once; they can be called from async contexts.

## Using AnyIO in async tasks

Flyte async tasks support `anyio` for structured concurrency as an alternative to `asyncio.gather()`.

```python
import anyio
import aioresult
import flyte

env = flyte.TaskEnvironment(
    "anyio_example",
    image=flyte.Image.from_debian_base().with_pip_packages("anyio", "aioresult"),
)

@env.task
async def process_item(x: int) -> int:
    return x * 2

@env.task
async def batch_process(items: list[int]) -> list[int]:
    captured_results = []

    async with anyio.create_task_group() as tg:
        # Start multiple tasks concurrently
        for item in items:
            captured_results.append(
                aioresult.ResultCapture.start_soon(tg, process_item, item)
            )

    # Extract results
    return [r.result() for r in captured_results]
```

**Note**: You can use anyio's task groups, timeouts, and other structured concurrency primitives within Flyte async tasks.