# From Flyte 1 to 2
> This bundle contains all pages in the From Flyte 1 to 2 section.
> Source: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2 ===

# From Flyte 1 to 2

> **📝 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.

Flyte 2 represents a fundamental shift in how Flyte workflows are written and executed.

## Pure Python execution

Write workflows in pure Python, enabling a more natural development experience and removing the constraints of a
domain-specific language (DSL).

### Sync Python

```
import flyte

env = flyte.TaskEnvironment("sync_example_env")

@env.task
def hello_world(name: str) -> str:
    return f"Hello, {name}!"

@env.task
def main(name: str) -> str:
    for i in range(10):
        hello_world(name)
    return "Done"

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

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/sync_example.py*

### Async Python

```
import asyncio
import flyte

env = flyte.TaskEnvironment("async_example_env")

@env.task
async def hello_world(name: str) -> str:
    return f"Hello, {name}!"

@env.task
async def main(name: str) -> str:
    results = []
    for i in range(10):
        results.append(hello_world(name))
    await asyncio.gather(*results)
    return "Done"

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

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/async_example.py*

As you can see in the hello world example, workflows can be constructed at runtime, allowing for more flexible and
adaptive behavior. Flyte 2 supports:

- Python's asynchronous programming model to express parallelism.
- Python's native error handling with `try-except` to overridden configurations, like resource requests.
- Python's native conditional `if`/`elif`/`else` syntax and `for`/`while` loops.

## Simplified API

The new API is more intuitive, with fewer abstractions to learn and a focus on simplicity.

| Use case                      | Flyte 1                     | Flyte 2                                 |
| ----------------------------- | --------------------------- | --------------------------------------- |
| Environment management        | `N/A`                       | `TaskEnvironment`                       |
| Perform basic computation     | `@task`                     | `@env.task`                             |
| Combine tasks into a workflow | `@workflow`                 | `@env.task`                             |
| Create dynamic workflows      | `@dynamic`                  | `@env.task`                             |
| Fanout parallelism            | `flytekit.map`              | Python `for` loop with `asyncio.gather` |
| Conditional execution         | `flytekit.conditional`      | Python `if-elif-else`                   |
| Catching workflow failures    | `@workflow(on_failure=...)` | Python `try-except`                     |

There is no `@workflow` decorator. Instead, "workflows" are authored through a pattern of tasks calling tasks.
Tasks are defined within environments, which encapsulate the context and resources needed for execution.

## Fine-grained reproducibility and recoverability

As in Flyte 1, Flyte 2 supports caching at the task level (via `@env.task(cache=...)`), but it further enables recovery at the finer-grained, sub-task level through a feature called tracing (via `@flyte.trace`).

```
import flyte

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

@flyte.trace
async def call_llm(prompt: str) -> str:
    return "Initial response from LLM"

@env.task
async def finalize_output(output: str) -> str:
    return "Finalized output"

@env.task(cache=flyte.Cache(behavior="auto"))
async def main(prompt: str) -> str:
    output = await call_llm(prompt)
    output = await finalize_output(output)
    return output

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(main, prompt="Prompt to LLM")
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/trace.py*

Here `call_llm` runs in the same container as `main` and acts as an automated checkpoint with full observability in the UI.
If the task fails due to a system error (e.g., node preemption or infrastructure failure), Flyte can recover and replay from the
last successful trace rather than restarting from the beginning.

Note that tracing is distinct from caching: traces are recovered only if there is a system failure
whereas with cached outputs are persisted for reuse across separate runs.

## Improved remote functionality

Flyte 2 provides full management of the workflow lifecycle through a standardized API through the CLI and the Python SDK.

| Use case      | CLI                | Python SDK          |
| ------------- | ------------------ | ------------------- |
| Run a task    | `flyte run ...`    | `flyte.run(...)`    |
| Deploy a task | `flyte deploy ...` | `flyte.deploy(...)` |

You can also fetch and run remote (previously deployed) tasks within the course of a running workflow.

```
import flyte
from flyte import remote

env_1 = flyte.TaskEnvironment(name="env_1")
env_2 = flyte.TaskEnvironment(name="env_2")
env_1.add_dependency(env_2)

@env_2.task
async def remote_task(x: str) -> str:
    return "Remote task processed: " + x

@env_1.task
async def main() -> str:
    remote_task_ref = remote.Task.get("env_2.remote_task", auto_version="latest")
    r = await remote_task_ref(x="Hello")
    return "main called remote and recieved: " + r

if __name__ == "__main__":
    flyte.init_from_config()
    d = flyte.deploy(env_1)
    print(d[0].summary_repr())
    r = flyte.run(main)
    print(r.name)
    print(r.url)
    r.wait()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/remote.py*

## Native notebook support

Author and run workflows and fetch workflow metadata (I/O and logs) directly from Jupyter notebooks.

![Native Notebook](https://www.union.ai/docs/latest/flyte/_static/images/user-guide/notebook.png)

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/overview ===

# Migration overview

This section walks through the Flyte 1 workload patterns you already know — data ETL, model training, hyperparameter sweeps, batch inference — and their Flyte 2 equivalents. Every pattern is a complete, runnable v1↔v2 example pair in the [`unionai-examples`](https://github.com/unionai/unionai-examples/tree/main/v2/user-guide/migration/flyte-2) repository.

Two conceptual shifts motivate almost every change — **pure Python execution** and the **asynchronous model** — after which most migrations come down to a couple of mechanical moves. This page covers both shifts, the terminology mapping, and a quick-reference module.

## Pure Python execution

In Flyte 1, `@workflow` functions were constrained to a DSL subset of Python that compiled to a static DAG. In Flyte 2 there is **no `@workflow` decorator**: everything is a `@env.task`, and a "workflow" is simply a task that calls other tasks. Orchestration runs as real Python at runtime, so loops, conditionals, and `try`/`except` work anywhere.

| Flyte 1 | Flyte 2 |
| --- | --- |
| `@workflow` functions are constrained to a subset of Python defining a static DAG. | **No `@workflow` decorator**: your top-level "workflow" is just a task that calls other tasks. |
| `@task` functions had the full power of Python, but only within a single container execution. | `@env.task`s call other tasks and build dynamic structures with any Python construct, anywhere. |
| Workflows compiled to static DAGs at registration time. | Workflows are tasks calling tasks; compile-time safety is coming via `compiled_task`. |

### Flyte 1

```python
import flytekit

image = flytekit.ImageSpec(
    name="hello-world-image",
    packages=["requests"],
)

@flytekit.task(container_image=image)
def mean(data: list[float]) -> float:
    return sum(list) / len(list)

@flytekit.workflow
def main(data: list[float]) -> float:
    output = mean(data)

    # ❌ performing trivial operations in a workflow is not allowed
    # output = output / 100

    # ❌ if/else is not allowed
    # if output < 0:
    #     raise ValueError("Output cannot be negative")

    return output
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/pure-python/flyte_1.py*

### Flyte 2

```
import flyte

env = flyte.TaskEnvironment(
    "hello_world",
    image=flyte.Image.from_debian_base().with_pip_packages("requests"),
)

@env.task
def mean(data: list[float]) -> float:
    return sum(data) / len(data)

@env.task
def main(data: list[float]) -> float:
    output = mean(data)

    # ✅ performing trivial operations in a workflow is allowed
    output = output / 100

    # ✅ if/else is allowed
    if output < 0:
        raise ValueError("Output cannot be negative")

    return output
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/pure-python/flyte_2.py*

This unlocks workflows that adapt to runtime conditions, native `try`/`except` error handling, and the intuitive composition of ordinary Python functions.

## Asynchronous model

Flyte 2 is built on Python's `asyncio`, with a crucial twist: **the Flyte orchestrator acts as the event loop**, scheduling awaited tasks across distributed infrastructure. This makes `async`/`await` the natural way to express parallelism.

| | Flyte 1 | Flyte 2 |
| --- | --- | --- |
| Parallelism | The DSL auto-parallelized independent tasks; the `map` operator ran a task over many inputs. | Python's `asyncio` expresses parallelism, with the Flyte orchestrator acting as the event loop across distributed infrastructure. |

The core async keywords carry Flyte-specific meaning:

- **`async def`** declares a coroutine.
- **`await`** signals where a task can be scheduled in parallel — not just an I/O yield point.
- **`asyncio.gather`** tells the orchestrator that a set of tasks are independent and can be distributed across separate compute resources.

Consider this pattern for parallel data processing:

```
import asyncio
import flyte

env = flyte.TaskEnvironment("data_pipeline")

@env.task
async def process_chunk(chunk_id: int, data: str) -> str:
    # This could be any computational work - CPU or I/O bound
    await asyncio.sleep(1)  # Simulating work
    return f"Processed chunk {chunk_id}: {data}"

@env.task
async def parallel_pipeline(data_chunks: list[str]) -> list[str]:
    # Create coroutines for all chunks
    tasks = []
    for i, chunk in enumerate(data_chunks):
        tasks.append(process_chunk(i, chunk))

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

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/async/async.py*

In standard Python this mainly benefits I/O-bound work; in Flyte 2 the orchestrator schedules each `process_chunk` task on its own pod.

### True parallelism for all workloads

Async syntax in Flyte 2 is **not just for I/O-bound operations**. When the orchestrator encounters `await asyncio.gather(...)`, it runs those independent tasks simultaneously across compute resources — achieving true parallelism for CPU-bound work (model training, heavy math), I/O-bound work (queries, API calls), and mixed workloads alike.

### Calling sync tasks from async tasks

You don't need to rewrite existing synchronous code. Flyte automatically "asyncifies" sync functions; just call them from an async context with `.aio()`:

```
@env.task
def legacy_computation(x: int) -> int:
    # Existing synchronous function works unchanged
    return x * x + 2 * x + 1

@env.task
async def modern_workflow(numbers: list[int]) -> list[int]:
    # Call sync tasks from async context using .aio()
    tasks = []
    for num in numbers:
        tasks.append(legacy_computation.aio(num))

    results = await asyncio.gather(*tasks)
    return results
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/async/async.py*

### The `flyte.map` function: Familiar patterns

For code that used Flyte 1's `map`, `flyte.map` is a direct replacement that works in both sync and async contexts:

### Sync Map

```
@env.task
def sync_map_example(n: int) -> list[str]:
    # Synchronous version for easier migration
    results = []
    for result in flyte.map(process_item, range(n)):
        if isinstance(result, Exception):
            raise result
        results.append(result)
    return results
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/async/async.py*

### Async Map

```
@env.task
async def async_map_example(n: int) -> list[str]:
    # Async version using flyte.map - exact pattern from SDK examples
    results = []
    async for result in flyte.map.aio(process_item, range(n), return_exceptions=True):
        if isinstance(result, Exception):
            raise result
        results.append(result)
    return results
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/flyte-2/async/async.py*

It provides dual interfaces (`flyte.map.aio()` and `flyte.map()`), `return_exceptions` for graceful failure handling (a more flexible replacement for Flyte 1's `min_success_ratio`), automatic UI grouping, and optional concurrency limits.

## Terminology and concept mapping

Several Flyte 1 concepts were renamed or reshaped in Flyte 2. The table below maps the ones you'll meet most often.

| Flyte 1 | Flyte 2 | Notes |
|---|---|---|
| `flytekit` (package) | `flyte` (package) | The Python SDK was renamed; imports change from `import flytekit` to `import flyte`. |
| `pyflyte` (CLI) | `flyte` (CLI) | The command-line tool was renamed. |
| `@task` / `@workflow` / `@dynamic` | `@env.task` | A single task decorator off a `flyte.TaskEnvironment`. Workflows and dynamic tasks are no longer distinct constructs: everything is a task, and orchestration is plain Python. |
| `map_task()` | `flyte.map()` | Plus `asyncio.gather()` for async fan-out. |
| `conditional()` | native `if` / `elif` / `else` | Branching is now ordinary Python control flow, not a DSL. |
| `ImageSpec` | `flyte.Image` | Container image definition. |
| `current_context()` | `flyte.ctx()` | Runtime context access. |
| `FlyteFile` / `FlyteDirectory` | `flyte.io.File` / `flyte.io.Dir` | Offloaded file and directory references. |
| `StructuredDataset` | `flyte.io.DataFrame` | Offloaded tabular data. |
| `LaunchPlan` | `flyte.Trigger` | Scheduling and parameterized entry points. |
| `CronSchedule` | `flyte.Cron` | Cron-based scheduling, used with a `flyte.Trigger`. |
| Decks (`enable_deck=True`) | Reports (`report=True`) | Custom HTML rendered in the UI during/after a run. See [Reports](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/reports). |

## Package imports

The package is renamed from `flytekit` to `flyte`, and the workflow/dynamic/map_task imports disappear:

### Flyte 1

```python
import flytekit
from flytekit import task, workflow, dynamic, map_task
from flytekit import ImageSpec, Resources, Secret
from flytekit import current_context, LaunchPlan, CronSchedule
```

### Flyte 2

```python
import flyte
from flyte import TaskEnvironment, Resources, Secret
from flyte import Image, Trigger, Cron
```

## The two mechanical changes behind (almost) every migration

Most of a migration comes down to two moves:

### 1. Move task configuration into a `TaskEnvironment`

Instead of configuring the image, resources, and caching on each task decorator, configure them once on a `flyte.TaskEnvironment` and share it across tasks:

```python
env = flyte.TaskEnvironment(
    name="training",
    image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn", "pandas"),
    resources=flyte.Resources(cpu="2", memory="4Gi"),
    cache="auto",
)
```

### 2. Replace `@task` / `@workflow` / `@dynamic` with `@env.task`

Every decorated function becomes an `@env.task`. There is no separate workflow or dynamic construct: a "workflow" is simply a task that calls other tasks, and orchestration is plain Python.

> **📝 Note**
>
> The `env` in `@env.task` is just the variable you assigned your `TaskEnvironment` to. Name it whatever you like.

## In this section

The migration patterns are grouped by theme. Start with **Tasks and workflows**, then jump to whatever your workload needs:

- **[Tasks and workflows](./tasks-and-workflows)** — the structural shift: `@task`/`@workflow` → `@env.task`, sequential ordering, nested "subworkflows", and the `@task` → `TaskEnvironment` parameter mapping.
- **[Task configuration](./configuration)** — moving image/resources/cache to the `TaskEnvironment`, GPUs, secrets, caching, and scheduling with triggers.
- **[CLI and configuration](./cli-and-configuration)** — `pyflyte` → `flyte` command mapping and config-file changes.
- **[Control flow](./control-flow)** — `conditional()` and `@dynamic` become plain Python `if`/loops, and `on_failure` becomes `try`/`except`.
- **[Parallelism and fan-out](./parallelism)** — `map_task` → `flyte.map` / `asyncio.gather`, plus a data-backfill example.
- **[Data types and I/O](./data-io)** — `FlyteFile`/`FlyteDirectory` → `flyte.io.File`/`Dir`, `StructuredDataset` → `flyte.io.DataFrame`, dataclasses, and an ETL example.
- **[ML workloads](./ml-workloads)** — small-model training, hyperparameter optimization, deep learning, batch inference, and an end-to-end pipeline.
- **[New in Flyte 2](./new-in-flyte-2)** — patterns that weren't possible in Flyte 1 at all: real-time model serving, batch inference, apps, and sandboxed code execution.

- **[Hybrid v1 and v2 pipelines](./hybrid-pipelines)** — calling between v1 and v2 in both directions during the transition.
- **[Gotchas and caveats](./gotchas-and-caveats)** — common gotchas plus the deeper caveats of the new execution model, including non-deterministic behavior and keeping orchestration lightweight.

## Quick reference

A minimal Flyte 2 module, end to end:

```python
import asyncio
import flyte

# 1. Define an image
image = (
    flyte.Image.from_debian_base(python_version=(3, 11))
    .with_pip_packages("pandas", "numpy")
)

# 2. Create a TaskEnvironment
env = flyte.TaskEnvironment(
    name="my_env",
    image=image,
    resources=flyte.Resources(cpu="1", memory="2Gi"),
)

# 3. Define tasks
@env.task
async def process(x: int) -> int:
    return x * 2

# 4. Define the entrypoint task
@env.task
async def main(items: list[int]) -> list[int]:
    results = await asyncio.gather(*[process(x) for x in items])
    return list(results)

# 5. Run it
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main, items=[1, 2, 3, 4, 5])
    print(run.url)
    run.wait()
```

```bash
# CLI
flyte run my_module.py main --items '[1,2,3,4,5]'   # remote (default)
flyte run --local my_module.py main --items '[1,2,3,4,5]'
flyte deploy my_module.py my_env
```

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/tasks-and-workflows ===

# Tasks and workflows

The biggest structural change in Flyte 2 is that **everything is a task**. The `@task`, `@workflow`, and `@dynamic` decorators all collapse into a single `@env.task` on a `flyte.TaskEnvironment`, and a "workflow" is just a task that calls other tasks. This page covers the basic structure; see [Task configuration](./configuration) for the environment settings and [Migration](./overview) for the big picture.

## Hello world: tasks and workflows

A `@task` plus `@workflow` becomes two `@env.task`s, where the entrypoint task calls the others. Sequential calls are naturally ordered — no `>>` operator required.

### Flyte 1

```python
import flytekit

@flytekit.task
def say_hello(name: str) -> str:
    return f"Hello, {name}!"

@flytekit.task
def to_upper(greeting: str) -> str:
    return greeting.upper()

@flytekit.workflow
def main(name: str) -> str:
    greeting = say_hello(name=name)
    return to_upper(greeting=greeting)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/hello_world_v1.py*

### Flyte 2

```
import flyte

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

@env.task
def say_hello(name: str) -> str:
    return f"Hello, {name}!"

@env.task
def to_upper(greeting: str) -> str:
    return greeting.upper()

# The "workflow" is now just a task that calls other tasks.
@env.task
def main(name: str) -> str:
    greeting = say_hello(name)
    return to_upper(greeting)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/hello_world_v2.py*

## Chaining and ordering

In Flyte 1 you sometimes used `>>` to force ordering between tasks with no data dependency. In Flyte 2, sequential (synchronous) calls run in the order they're written, and `await`ing async tasks in sequence does the same. The `>>` operator is gone.

### Flyte 1

```python
from flytekit import task, workflow

@task
def clear_staging_table() -> None:
    # Side effect only: truncate the staging table.
    print("cleared staging table")

@task
def load_into_staging() -> None:
    # Side effect only: load fresh rows into staging.
    print("loaded staging table")

@task
def publish_to_prod() -> None:
    # Side effect only: swap staging into the production table.
    print("published to prod")

@workflow
def main() -> None:
    clear = clear_staging_table()
    load = load_into_staging()
    publish = publish_to_prod()

    # These tasks pass no data between them, so use the >> operator to force
    # ordering: clear must finish before load, which must finish before publish.
    clear >> load >> publish
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/chained_tasks_v1.py*

### Flyte 2

```
import flyte

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

@env.task
def clear_staging_table() -> None:
    print("cleared staging table")

@env.task
def load_into_staging() -> None:
    print("loaded staging table")

@env.task
def publish_to_prod() -> None:
    print("published to prod")

# Sequential (synchronous) calls run in the order they're written, even when no
# data flows between them. The Flyte 1 `>>` ordering operator is gone.
@env.task
def main() -> None:
    clear_staging_table()
    load_into_staging()
    publish_to_prod()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/chained_tasks_v2.py*

## Subworkflows

A `@workflow` invoked by another `@workflow` (for example, a reusable preprocessing pipeline) becomes a task that calls other tasks — nest them as deeply as you like.

### Flyte 1

```python
from flytekit import task, workflow

@task
def impute(value: float) -> float:
    # Replace missing/negative sentinel values with 0.
    return value if value >= 0 else 0.0

@task
def scale(value: float) -> float:
    return value / 100.0

@workflow
def preprocess(value: float) -> float:
    imputed = impute(value=value)
    return scale(value=imputed)

@workflow
def main(raw_value: float) -> float:
    return preprocess(value=raw_value)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/subworkflow_v1.py*

### Flyte 2

```
import flyte

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

@env.task
def impute(value: float) -> float:
    # Replace missing/negative sentinel values with 0.
    return value if value >= 0 else 0.0

@env.task
def scale(value: float) -> float:
    return value / 100.0

# A preprocessing "subworkflow" is just a task that calls other tasks.
@env.task
def preprocess(value: float) -> float:
    imputed = impute(value)
    return scale(imputed)

@env.task
def main(raw_value: float) -> float:
    return preprocess(raw_value)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/subworkflow_v2.py*

## TaskEnvironment configuration

The `TaskEnvironment` holds the configuration that Flyte 1 spread across the `@task` decorator. The task decorator can still override a few settings per-task.

```python
import flyte

env = flyte.TaskEnvironment(
    name="my_env",                           # Required: unique name
    image=flyte.Image.from_debian_base(),    # Or a string, or "auto"
    resources=flyte.Resources(
        cpu="2",
        memory="4Gi",
        gpu="A100:1",
        disk="10Gi",
    ),
    env_vars={"LOG_LEVEL": "INFO"},
    secrets=[flyte.Secret(key="api-key", as_env_var="API_KEY")],
    cache="auto",                            # "auto", "override", "disable", or a Cache object
    reusable=flyte.ReusePolicy(replicas=5, idle_ttl=60),
    interruptible=True,
)

# The task decorator can override some settings:
@env.task(
    short_name="my_task",   # Display name
    cache="disable",        # Override cache
    retries=3,              # Retry count
    timeout=3600,           # Seconds or a timedelta
    report=True,            # Generate an HTML report
)
def my_task(x: int) -> int:
    return x
```

## Parameter mapping: `@task` → `TaskEnvironment` + `@env.task`

| Flyte 1 `@task` parameter | Flyte 2 location | Notes |
|---|---|---|
| `container_image` | `TaskEnvironment(image=...)` | Env-level only |
| `requests` | `TaskEnvironment(resources=...)` | Env-level only |
| `limits` | `TaskEnvironment(resources=...)` | Combined with requests (single value) |
| `environment` | `TaskEnvironment(env_vars=...)` | Env-level only |
| `secret_requests` | `TaskEnvironment(secrets=...)` | Env-level only |
| `cache` | Both | Can override at task level |
| `cache_version` | `flyte.Cache(version_override=...)` | In a `Cache` object |
| `retries` | `@env.task(retries=...)` | Task-level only |
| `timeout` | `@env.task(timeout=...)` | Task-level only |
| `interruptible` | Both | Can override at task level |
| `pod_template` | Both | Can override at task level |
| `deprecated` | N/A | Not in Flyte 2 |
| `docs` | `@env.task(docs=...)` | Task-level only |

For image, resource, secret, and caching detail, see [Task configuration](./configuration).

## Next

- [Task configuration](./configuration) — image, resources, caching, secrets, and scheduling
- [Control flow](./control-flow) — conditionals, dynamic behavior, and error handling
- [Parallelism and fan-out](./parallelism) — running tasks in parallel

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/configuration ===

# Task configuration

In Flyte 1, image, resources, caching, secrets, and scheduling were configured per-task on the `@task` decorator or per-workflow on a `LaunchPlan`. In Flyte 2 most of this moves to the `flyte.TaskEnvironment`, so it's declared once and shared. See [Migration](./overview) for the overall approach.

## Image, resources, and caching

Image, resources, and caching move from the `@task` decorator to the `TaskEnvironment`. Per-task settings like `retries` and `timeout` stay on `@env.task`. Note that `mem` is renamed to `memory`, and there are no separate `requests`/`limits` — a single `Resources` value serves as both.

### Flyte 1

```python
from datetime import timedelta

import flytekit
from flytekit import Resources

image = flytekit.ImageSpec(
    name="training-image",
    packages=["scikit-learn", "pandas"],
)

@flytekit.task(
    container_image=image,
    requests=Resources(cpu="2", mem="4Gi"),
    limits=Resources(cpu="4", mem="8Gi"),
    cache=True,
    cache_version="1.0",
    retries=3,
    timeout=timedelta(minutes=30),
)
def train_epoch(step: int) -> float:
    # A stand-in for a training step that returns the current loss.
    return 1.0 / (step + 1)

@flytekit.workflow
def main(step: int) -> float:
    return train_epoch(step=step)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/task_config_v1.py*

### Flyte 2

```
from datetime import timedelta

import flyte

# Image, resources, and caching move to the TaskEnvironment, so they are declared
# once and shared by every task in the environment.
env = flyte.TaskEnvironment(
    name="training",
    image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn", "pandas"),
    resources=flyte.Resources(cpu="2", memory="4Gi"),  # "memory", not "mem"
    cache="auto",
)

# retries and timeout stay on the task decorator.
@env.task(retries=3, timeout=timedelta(minutes=30))
def train_epoch(step: int) -> float:
    # A stand-in for a training step that returns the current loss.
    return 1.0 / (step + 1)

@env.task
def main(step: int) -> float:
    return train_epoch(step)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/task_config_v2.py*

For the full `@task` → `TaskEnvironment` + `@env.task` parameter mapping, see [Tasks and workflows](./tasks-and-workflows).

## Container images

Flyte 1's `ImageSpec` is replaced by Flyte 2's `flyte.Image` with a fluent builder API.

```python
from flyte import Image

image = (
    Image.from_debian_base(name="my-image", registry="ghcr.io/myorg", python_version=(3, 11))
    .with_pip_packages("pandas", "numpy")
    .with_apt_packages("curl", "git")
    .with_env_vars({"MY_VAR": "value"})
)
```

Instead of one constructor with many arguments, you start from a base and chain builder methods:

| Constructor | Use case |
|---|---|
| `Image.from_debian_base()` | Most common; includes the Flyte SDK |
| `Image.from_base(image_uri)` | Start from any existing image |
| `Image.from_dockerfile(path)` | Complex custom builds |
| `Image.from_uv_script(path)` | UV-based projects |

Common chainable builder methods: `.with_pip_packages(...)`, `.with_requirements(path)`, `.with_uv_project(path)`, `.with_apt_packages(...)`, `.with_commands([...])`, `.with_source_file(path, dst=...)`, `.with_source_folder(path, dst=...)`, `.with_env_vars({...})`, and `.with_workdir(...)`.

| Flyte 1 `ImageSpec` | Flyte 2 `Image` | Notes |
|---|---|---|
| `name` | `name` (constructor) | Same |
| `registry` | `registry` (constructor) | Same |
| `python_version` | `python_version` (tuple) | `"3.11"` becomes `(3, 11)` |
| `packages` | `.with_pip_packages()` | Method instead of param |
| `apt_packages` | `.with_apt_packages()` | Method instead of param |
| `requirements` | `.with_requirements()` | Supports txt, poetry.lock, uv.lock |
| `env` | `.with_env_vars()` | Method instead of param |
| `commands` | `.with_commands()` | Method instead of param |
| `copy` / `source_root` | `.with_source_file()` / `.with_source_folder()` | More explicit methods |
| `base_image` | `Image.from_base()` | Different constructor |
| `builder` | Config file or `flyte.init()` | Global setting |
| `platform` | `platform` (constructor) | Tuple: `("linux/amd64", "linux/arm64")` |

For a private registry, create an image-pull secret and reference it:

```shell
flyte create secret --type image_pull my-registry-secret --from-file ~/.docker/config.json
```

```python
image = Image.from_debian_base(
    registry="private.registry.com",
    name="my-image",
    registry_secret="my-registry-secret",
)
```

See [Container images](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/container-images) for more.

## Resources

A single `flyte.Resources` value serves as both request and limit — there are no separate `requests`/`limits`. Several parameters were renamed.

| Flyte 1 | Flyte 2 | Notes |
|---|---|---|
| `cpu="1"` | `cpu="1"` | Same |
| `mem="2Gi"` | `memory="2Gi"` | Renamed |
| `gpu="1"` | `gpu="A100:1"` | `Type:count` format |
| `ephemeral_storage="10Gi"` | `disk="10Gi"` | Renamed |
| N/A | `shm="auto"` | New: shared memory |

GPU type and count are combined into one string, replacing the separate Flyte 1 `accelerator=` argument:

```python
env = flyte.TaskEnvironment(
    name="gpu_env",
    resources=flyte.Resources(
        cpu="4",
        memory="32Gi",
        gpu="A100:2",              # Type:count
        # gpu="A100 80G:1"         # 80GB variant
        # gpu=flyte.GPU("A100", count=1, partition="1g.5gb")   # MIG partition
    ),
)
```

Supported GPU types include A10, A10G, A100, A100 80G, B200, H100, H200, L4, L40s, T4, V100, RTX PRO 6000, and GB10. See [Resources](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/resources) for more.

## Caching

Caching is enabled at the env level with `cache="auto"` (or per-task on `@env.task`). The explicit `cache_version` string moves into a `flyte.Cache` object.

| Behavior | Description |
|---|---|
| `"auto"` | Cache results and reuse if available |
| `"override"` | Always execute and overwrite the cache |
| `"disable"` | No caching (default for a `TaskEnvironment`) |

```python
# Flyte 1: @task(cache=True, cache_version="1.0")
# Flyte 2:
@env.task(cache="auto")
def cached_task(x: int) -> int:
    return x * 2

# Advanced control (replaces cache_version, serialize, ignored_inputs, ...)
@env.task(cache=flyte.Cache(
    behavior="auto",
    version_override="v1.0",
    serialize=True,
    ignored_inputs=("debug",),
))
def advanced(x: int, debug: bool = False) -> int:
    return x * 2
```

See [Caching](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/caching) for more.

## Secrets

Secrets move from `secret_requests` on the task to `secrets` on the `TaskEnvironment`, and you read them from environment variables instead of `current_context().secrets` — for example, an API key for a model registry or hosted LLM.

### Flyte 1

```python
from flytekit import task, workflow, Secret, current_context

@task(secret_requests=[Secret(group="openai", key="api_key")])
def call_api() -> str:
    token = current_context().secrets.get(group="openai", key="api_key")
    return f"token has {len(token)} chars"

@workflow
def main() -> str:
    return call_api()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/secrets_v1.py*

### Flyte 2

```
import os

import flyte

# Secrets are declared on the TaskEnvironment and injected as environment
# variables (instead of read through current_context().secrets).
env = flyte.TaskEnvironment(
    name="secrets",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
)

@env.task
def call_api() -> str:
    token = os.getenv("OPENAI_API_KEY", "")
    return f"token has {len(token)} chars"

@env.task
def main() -> str:
    return call_api()
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/secrets_v2.py*

A `flyte.Secret` can be mounted as an environment variable or as a file, and the access convention changes:

```python
flyte.Secret(key="openai-key", as_env_var="OPENAI_API_KEY")   # mount as env var
flyte.Secret(key="access-key", group="aws")                    # env var: AWS_ACCESS_KEY
flyte.Secret(key="ssl-cert", mount="/etc/flyte/secrets")       # mount as a file
```

| Flyte 1 pattern | Flyte 2 pattern |
|---|---|
| `ctx.secrets.get(key="mykey", group="mygroup")` | `os.environ["MYGROUP_MYKEY"]` (auto-named) |
| `ctx.secrets.get(key="mykey")` | `os.environ["MY_SECRET"]` (with `as_env_var="MY_SECRET"`) |

Create and manage secrets from the CLI:

```bash
flyte create secret MY_SECRET_KEY --value my_secret_value
flyte create secret MY_SECRET_KEY --from-file /path/to/secret
flyte get secret
flyte delete secret MY_SECRET_KEY
```

See [Secrets](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/secrets) for more.

## Scheduling

A `LaunchPlan` with a `CronSchedule` (say, a nightly retraining job) becomes a `flyte.Trigger` attached directly to the task. Use `flyte.TriggerTime` to bind the scheduled fire time to an input, and deploy the trigger with `flyte deploy`.

### Flyte 1

```python
from flytekit import task, workflow, LaunchPlan, CronSchedule

@task
def retrain(kickoff_time: str) -> str:
    return f"retrained model at {kickoff_time}"

@workflow
def main(kickoff_time: str) -> str:
    return retrain(kickoff_time=kickoff_time)

# A LaunchPlan attaches a schedule (and default inputs) to a workflow.
nightly_retrain = LaunchPlan.get_or_create(
    workflow=main,
    name="nightly_retrain",
    schedule=CronSchedule(
        schedule="0 2 * * *",  # 2 AM daily
        kickoff_time_input_arg="kickoff_time",
    ),
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/scheduling_v1.py*

### Flyte 2

```
from datetime import datetime

import flyte

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

# A Trigger replaces LaunchPlan + CronSchedule. It is attached directly to the
# task and deployed with it (flyte deploy). flyte.TriggerTime binds the
# scheduled fire time to a task input.
nightly_retrain = flyte.Trigger(
    name="nightly_retrain",
    automation=flyte.Cron("0 2 * * *"),  # 2 AM daily
    inputs={"trigger_time": flyte.TriggerTime},
    auto_activate=True,
)

@env.task(triggers=nightly_retrain)
def main(trigger_time: datetime = datetime(2024, 1, 1, 2, 0)) -> str:
    return f"retrained model at {trigger_time.isoformat()}"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/scheduling_v2.py*

Triggers support `flyte.Cron("0 9 * * *", timezone="America/New_York")` and `flyte.FixedRate(timedelta(hours=1))` as automations, plus convenience constructors like `flyte.Trigger.hourly()` and `flyte.Trigger.daily()`. See [Triggers](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/triggers) for more.

## Next

- [CLI and configuration](./cli-and-configuration) — `pyflyte` → `flyte` command and config-file mapping
- [Control flow](./control-flow) — conditionals, dynamic behavior, and error handling
- [Data types and I/O](./data-io) — files, DataFrames, and dataclasses

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/cli-and-configuration ===

# CLI and configuration

The command-line tool is renamed from `pyflyte` to `flyte`, and the config file is trimmed down. See [Migration](./overview) for the overall approach.

## CLI command mapping

| Flyte 1 | Flyte 2 | Notes |
|---|---|---|
| `pyflyte run` | `flyte run` | Similar, different flags |
| `pyflyte run --remote` | `flyte run` | Remote is the default in Flyte 2 |
| `pyflyte run` (local) | `flyte run --local` | Local execution is now explicit |
| `pyflyte register` | `flyte deploy` | Different concept |
| `pyflyte package` | N/A | Not needed in Flyte 2 |
| `pyflyte serialize` | N/A | Not needed in Flyte 2 |

### Running tasks

### Flyte 1

```shell
# Local
pyflyte run my_module.py my_workflow --arg1 value1

# Remote
pyflyte --config config.yaml run --remote my_module.py my_workflow --arg1 value1
```

### Flyte 2

```shell
# Remote (default)
flyte run my_module.py my_task --arg1 value1

# Local
flyte run --local my_module.py my_task --arg1 value1

# With an explicit config file
flyte --config config.yaml run my_module.py my_task --arg1 value1
```

### Deploying

In Flyte 1 you registered a module; in Flyte 2 you deploy task environments.

### Flyte 1

```shell
pyflyte register my_module.py -p my-project -d development
```

### Flyte 2

```shell
# Deploy a task environment
flyte deploy my_module.py my_env --project my-project --domain development

# Deploy all environments in a file
flyte deploy --all my_module.py

# Deploy with an explicit version, or recursively
flyte deploy --version v1.0.0 my_module.py my_env
flyte deploy --recursive --all ./src
```

### Key flag differences

| Flyte 1 flag | Flyte 2 flag | Notes |
|---|---|---|
| `--remote` | (default) | Remote is the default |
| `--copy-all` | `--copy-style all` | File copying |
| N/A | `--copy-style loaded_modules` | Default: only imported modules |
| `-p, --project` | `--project` | Same |
| `-d, --domain` | `--domain` | Same |
| `-i, --image` | `--image` | Same format |
| N/A | `--follow, -f` | Follow execution logs |

## Configuration files

The config file lives in the same place (`~/.flyte/config.yaml`), but the environment variable changes from `FLYTECTL_CONFIG` to `FLYTE_CONFIG`, and the format is simpler.

### Flyte 1

```yaml
admin:
  endpoint: dns:///your-cluster.hosted.unionai.cloud
  insecure: false
  authType: Pkce
```

### Flyte 2

```yaml
admin:
  endpoint: dns:///your-cluster.hosted.unionai.cloud

image:
  builder: remote  # or "local"

task:
  domain: development
  org: your-org
  project: your-project
```

| Setting | Flyte 1 | Flyte 2 |
|---|---|---|
| Endpoint | `admin.endpoint` | `admin.endpoint` |
| Auth type | `admin.authType` | Auto-detected (PKCE default) |
| Project | CLI flag `-p` | `task.project` (default) |
| Domain | CLI flag `-d` | `task.domain` (default) |
| Organization | CLI flag `--org` | `task.org` (default) |
| Image builder | N/A | `image.builder` (`local` or `remote`) |

### Configuring in code

```python
import flyte

# From a config file (auto-discovers, or pass a path)
flyte.init_from_config()
flyte.init_from_config("path/to/config.yaml")

# Programmatically
flyte.init(
    endpoint="flyte.example.com",
    project="my-project",
    domain="development",
)
```

For API-key authentication in non-interactive environments, use `flyte.init_from_api_key()` — see [Run on a remote cluster](https://www.union.ai/docs/latest/flyte/user-guide/get-started/run-modes/running-remote).

## Next

- [Control flow](./control-flow) — conditionals, dynamic behavior, and error handling
- [Hybrid v1 and v2 pipelines](./hybrid-pipelines) — calling between v1 and v2 during the transition

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/control-flow ===

# Control flow

Flyte 1 expressed branching, dynamic fan-out, and failure handling through DSL constructs (`conditional()`, `@dynamic`, `@workflow(on_failure=...)`). In Flyte 2 these are all ordinary Python, because orchestration runs as real Python at runtime. See [Migration](./overview) for the overall approach.

## Conditional execution

The `conditional()` DSL becomes ordinary Python `if` / `elif` / `else` — for example, choosing a model based on dataset size.

### Flyte 1

```python
from flytekit import task, workflow, conditional

@task
def train_gradient_boosting(n_rows: int) -> str:
    return f"trained gradient boosting on {n_rows} rows"

@task
def train_logistic_regression(n_rows: int) -> str:
    return f"trained logistic regression on {n_rows} rows"

@workflow
def main(n_rows: int) -> str:
    # Pick the model based on dataset size.
    return (
        conditional("model_choice")
        .if_(n_rows > 10_000)
        .then(train_gradient_boosting(n_rows=n_rows))
        .else_()
        .then(train_logistic_regression(n_rows=n_rows))
    )
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/conditional_v1.py*

### Flyte 2

```
import flyte

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

@env.task
def train_gradient_boosting(n_rows: int) -> str:
    return f"trained gradient boosting on {n_rows} rows"

@env.task
def train_logistic_regression(n_rows: int) -> str:
    return f"trained logistic regression on {n_rows} rows"

# Branching is now ordinary Python control flow -- no conditional() DSL.
@env.task
def main(n_rows: int) -> str:
    if n_rows > 10_000:
        return train_gradient_boosting(n_rows)
    return train_logistic_regression(n_rows)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/conditional_v2.py*

## Dynamic workflows

`@dynamic` existed so a task could generate a variable number of subtask calls at runtime (e.g. one per data partition discovered at runtime). In Flyte 2 every task can do this natively, so `@dynamic` simply disappears — loop over runtime data in an ordinary `@env.task`.

### Flyte 1

```python
from flytekit import task, workflow, dynamic

@task
def list_partitions(n: int) -> list[int]:
    return list(range(n))

@task
def process_partition(partition_id: int) -> int:
    # Aggregate one data partition.
    return partition_id * 2

@dynamic
def process_all(partitions: list[int]) -> list[int]:
    results = []
    for partition_id in partitions:
        results.append(process_partition(partition_id=partition_id))
    return results

@workflow
def main(n: int) -> list[int]:
    partitions = list_partitions(n=n)
    return process_all(partitions=partitions)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/dynamic_v1.py*

### Flyte 2

```
import flyte

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

@env.task
def process_partition(partition_id: int) -> int:
    # Aggregate one data partition.
    return partition_id * 2

# No @dynamic decorator needed: a plain task can loop over runtime data (e.g. a
# variable number of partitions discovered at runtime) and call other tasks.
@env.task
def main(n: int) -> list[int]:
    return [process_partition(partition_id) for partition_id in range(n)]
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/dynamic_v2.py*

## Error handling

Flyte 1's `@workflow(on_failure=...)` handler becomes ordinary Python `try` / `except` — catch a failed training run, run cleanup, and recover or re-raise.

### Flyte 1

```python
from flytekit import task, workflow

@task
def train_fold(max_depth: int) -> float:
    if max_depth <= 0:
        raise ValueError("max_depth must be positive")
    # Return validation accuracy for this hyperparameter.
    return 0.90 + 0.001 * max_depth

@task
def notify_failure() -> None:
    print("training run failed -- sending alert")

# The on_failure handler runs if any node in the workflow fails. There is no
# try/except inside a Flyte 1 workflow.
@workflow(on_failure=notify_failure)
def main(max_depth: int) -> float:
    return train_fold(max_depth=max_depth)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/error_handling_v1.py*

### Flyte 2

```
import flyte

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

@env.task
async def train_fold(max_depth: int) -> float:
    if max_depth <= 0:
        raise ValueError("max_depth must be positive")
    return 0.90 + 0.001 * max_depth

# Failure handling is ordinary Python try/except -- no on_failure handler.
@env.task
async def main(max_depth: int) -> float:
    try:
        return await train_fold(max_depth)
    except ValueError as e:
        print(f"invalid hyperparameter ({e}); falling back to a safe default")
        # Recover with a safe default instead of failing the whole run.
        return await train_fold(max_depth=6)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/error_handling_v2.py*

Flyte 2 also exposes typed errors, so you can catch a specific failure and retry with more resources — a common need for memory-hungry training jobs:

```python
try:
    return await train_fold(sample_size)
except flyte.errors.OOMError:
    # Retry the same task with a larger memory request.
    return await train_fold.override(
        resources=flyte.Resources(memory="16Gi")
    )(sample_size)
```

See [Error handling](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/error-handling) for more.

## Next

- [Parallelism and fan-out](./parallelism) — running many tasks in parallel
- [ML workloads](./ml-workloads) — training, HPO, and inference

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/parallelism ===

# Parallelism and fan-out

Flyte 1's `map_task` becomes `flyte.map`, and the idiomatic Flyte 2 approach to fan-out is Python `async`/`await` with `asyncio.gather`. See the [Asynchronous model](./overview#asynchronous-model) guide for the concepts, and [Migration](./overview) for the overall approach.

## Fan-out: `map_task`

`map_task()` becomes `flyte.map()`, a near drop-in replacement. The one catch: `flyte.map` returns a generator, so wrap it in `list()`. For new code, the idiomatic approach is Python `async`/`await` with `asyncio.gather()`, which gives you finer control over concurrency and error handling.

### Flyte 1

```python
from functools import partial

from flytekit import task, workflow, map_task

@task
def get_shards(n: int) -> list[int]:
    return list(range(n))

@task
def score_shard(shard_id: int, model_version: int) -> int:
    # Score one shard of records with the given model version.
    return shard_id * model_version

@workflow
def main(n: int, model_version: int) -> list[int]:
    shards = get_shards(n=n)
    return map_task(
        partial(score_shard, model_version=model_version),
        concurrency=10,
    )(shard_id=shards)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/map_task_v1.py*

### Flyte 2 (flyte.map)

```
@env.task
def score_shard(shard_id: int, model_version: int) -> int:
    # Score one shard of records with the given model version.
    return shard_id * model_version

@env.task
def main(n: int, model_version: int) -> list[int]:
    bound = partial(score_shard, model_version=model_version)
    # flyte.map is a drop-in for map_task, but it returns a generator, so wrap
    # it in list() to materialize the results.
    return list(flyte.map(bound, range(n), concurrency=10))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/map_task_v2.py*

### Flyte 2 (asyncio.gather)

```
@env.task
async def score_shard_async(shard_id: int, model_version: int) -> int:
    return shard_id * model_version

@env.task
async def main_async(n: int, model_version: int) -> list[int]:
    # asyncio.gather is the idiomatic Flyte 2 way to fan out.
    coros = [score_shard_async(i, model_version) for i in range(n)]
    return list(await asyncio.gather(*coros))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/map_task_v2.py*

See the [Asynchronous model](./overview#asynchronous-model) guide for the concepts behind async execution.

### Choosing `flyte.map` vs `asyncio.gather`

| Feature | `flyte.map` (sync) | `asyncio.gather` (async) |
|---|---|---|
| Syntax | `list(flyte.map(fn, items))` | `await asyncio.gather(*tasks)` |
| Concurrency limit | Built-in `concurrency=N` | Use `asyncio.Semaphore` |
| Streaming / as-completed | No | Yes, via `asyncio.as_completed()` |
| Error handling | `return_exceptions=True` | Check return type |

Use **`flyte.map`** for the smallest change from Flyte 1 `map_task`, or when you're stuck in synchronous code. Use **`asyncio.gather`** for new code where you want streaming results or fine-grained concurrency control.

### Concurrency control and error handling

`map_task`'s `concurrency` and `min_success_ratio` become an `asyncio.Semaphore` and `return_exceptions=True`:

```python
import asyncio

@env.task
async def main(items: list[int], max_concurrent: int = 5) -> list[str]:
    sem = asyncio.Semaphore(max_concurrent)

    async def process_with_limit(item: int) -> str:
        async with sem:
            return await process_item(item)

    tasks = [process_with_limit(i) for i in items]
    results = await asyncio.gather(*tasks, return_exceptions=True)

    return [r for r in results if not isinstance(r, Exception)]
```

## Data backfills

Reprocessing a range of dates is a textbook `@dynamic` use case in Flyte 1, because the number of days is only known at runtime. In Flyte 2 it's a plain task that builds the date range and fans the days out with `asyncio.gather`.

### Flyte 1

```python
from datetime import date, timedelta

from flytekit import task, workflow, dynamic

@task
def process_day(day: str) -> int:
    # Reprocess a single day's partition; return the row count.
    return len(day)

# @dynamic is needed because the number of days is only known at runtime.
@dynamic
def backfill(start: str, days: int) -> list[int]:
    base = date.fromisoformat(start)
    results = []
    for i in range(days):
        day = (base + timedelta(days=i)).isoformat()
        results.append(process_day(day=day))
    return results

@workflow
def main(start: str, days: int) -> list[int]:
    return backfill(start=start, days=days)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/data_backfill_v1.py*

### Flyte 2

```
import asyncio
from datetime import date, timedelta

import flyte

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

@env.task
async def process_day(day: str) -> int:
    # Reprocess a single day's partition; return the row count.
    return len(day)

# A plain task builds the date range at runtime and fans the days out in
# parallel with asyncio.gather -- no @dynamic and no map_task needed.
@env.task
async def main(start: str, days: int) -> list[int]:
    base = date.fromisoformat(start)
    coros = [
        process_day((base + timedelta(days=i)).isoformat())
        for i in range(days)
    ]
    return list(await asyncio.gather(*coros))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/data_backfill_v2.py*

For fine-grained concurrency control (semaphores, `as_completed`, error handling), see [Fanout](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/fanout).

## Next

- [Data types and I/O](./data-io) — files, DataFrames, and dataclasses
- [ML workloads](./ml-workloads) — training, HPO, and batch inference

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/data-io ===

# Data types and I/O

Flyte 2 renames the offloaded-data types and makes their I/O `async`, but the mental model is the same: pass lightweight references to large data between tasks. See [Migration](./overview) for the overall approach.

## Files and directories

`FlyteFile` and `FlyteDirectory` become `flyte.io.File` and `flyte.io.Dir` — the way you pass model artifacts and datasets between tasks. The I/O is now `async`: use `await File.from_local(...)` to upload and `async with file.open(...)` to read. Like their Flyte 1 counterparts, these are lightweight references to offloaded data, not the materialized bytes.

### Flyte 1

```python
import os

from flytekit import task, workflow, current_context
from flytekit.types.file import FlyteFile

@task
def write_file(content: str) -> FlyteFile:
    path = os.path.join(current_context().working_directory, "out.txt")
    with open(path, "w") as f:
        f.write(content)
    return FlyteFile(path=path)

@task
def read_file(f: FlyteFile) -> str:
    with open(f.download()) as fh:
        return fh.read()

@workflow
def main(content: str) -> str:
    f = write_file(content=content)
    return read_file(f=f)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/files_v1.py*

### Flyte 2

```
import flyte
from flyte.io import File

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

@env.task
async def write_file(content: str) -> File:
    with open("out.txt", "w") as f:
        f.write(content)
    # File.from_local uploads the file to blob storage and returns a reference
    # (a lightweight pointer, not the materialized bytes).
    return await File.from_local("out.txt")

@env.task
async def read_file(f: File) -> str:
    async with f.open("rb") as fh:
        return (await fh.read()).decode("utf-8")

@env.task
async def main(content: str) -> str:
    f = await write_file(content)
    return await read_file(f)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/files_v2.py*

See [Files and directories](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/files-and-directories) for more.

## DataFrames

`StructuredDataset` becomes `flyte.io.DataFrame`. Construct one with `flyte.io.DataFrame.from_df(df)` and read it back with `await df.open(pandas.DataFrame).all()`.

### Flyte 1

```python
import pandas as pd
from flytekit import task, workflow
from flytekit.types.structured import StructuredDataset

@task
def make_df() -> StructuredDataset:
    df = pd.DataFrame({"employee_id": [1, 2, 3], "salary": [50000, 60000, 70000]})
    return StructuredDataset(dataframe=df)

@task
def total_payroll(sd: StructuredDataset) -> float:
    df = sd.open(pd.DataFrame).all()
    return float(df["salary"].sum())

@workflow
def main() -> float:
    return total_payroll(sd=make_df())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/dataframe_v1.py*

### Flyte 2

```
import pandas as pd
import flyte
import flyte.io

env = flyte.TaskEnvironment(
    name="dataframe",
    image=flyte.Image.from_debian_base().with_pip_packages("pandas", "pyarrow"),
)

@env.task
async def make_df() -> flyte.io.DataFrame:
    df = pd.DataFrame({"employee_id": [1, 2, 3], "salary": [50000, 60000, 70000]})
    # StructuredDataset becomes flyte.io.DataFrame.
    return flyte.io.DataFrame.from_df(df)

@env.task
async def total_payroll(fdf: flyte.io.DataFrame) -> float:
    df = await fdf.open(pd.DataFrame).all()
    return float(df["salary"].sum())

@env.task
async def main() -> float:
    return await total_payroll(await make_df())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/dataframe_v2.py*

See [DataFrames](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/dataframes) for more.

## Dataclasses and structured types

Flyte 1 required a `@dataclass_json` mixin for dataclass I/O. In Flyte 2, plain dataclasses (and Pydantic `BaseModel`s) work directly as task inputs and outputs — handy for passing around a training config.

### Flyte 1

```python
from dataclasses import dataclass

from dataclasses_json import dataclass_json
from flytekit import task, workflow

@dataclass_json
@dataclass
class TrainingConfig:
    learning_rate: float
    n_estimators: int
    max_depth: int = 6

@task
def make_config(learning_rate: float, n_estimators: int) -> TrainingConfig:
    return TrainingConfig(learning_rate=learning_rate, n_estimators=n_estimators)

@task
def train(config: TrainingConfig) -> str:
    return (
        f"trained with lr={config.learning_rate}, "
        f"n_estimators={config.n_estimators}, max_depth={config.max_depth}"
    )

@workflow
def main(learning_rate: float, n_estimators: int) -> str:
    config = make_config(learning_rate=learning_rate, n_estimators=n_estimators)
    return train(config=config)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/dataclasses_v1.py*

### Flyte 2

```
from dataclasses import dataclass

import flyte

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

# Plain dataclasses work directly as task I/O -- no @dataclass_json mixin needed.
# Pydantic BaseModels work the same way.
@dataclass
class TrainingConfig:
    learning_rate: float
    n_estimators: int
    max_depth: int = 6

@env.task
def make_config(learning_rate: float, n_estimators: int) -> TrainingConfig:
    return TrainingConfig(learning_rate=learning_rate, n_estimators=n_estimators)

@env.task
def train(config: TrainingConfig) -> str:
    return (
        f"trained with lr={config.learning_rate}, "
        f"n_estimators={config.n_estimators}, max_depth={config.max_depth}"
    )

@env.task
def main(learning_rate: float, n_estimators: int) -> str:
    config = make_config(learning_rate, n_estimators)
    return train(config)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/dataclasses_v2.py*

## Data ETL

Putting the data types together: extract, clean, aggregate, and write out a feature table. `StructuredDataset` becomes `flyte.io.DataFrame`, and the tasks become `async`.

### Flyte 1

```python
import pandas as pd
from flytekit import task, workflow
from flytekit.types.structured import StructuredDataset

@task
def extract() -> pd.DataFrame:
    # Read raw transaction records (stand-in for a real source).
    return pd.DataFrame(
        {
            "user_id": [1, 1, 2, 3, 3, 3],
            "amount": [10.0, 5.0, 20.0, 7.5, 2.5, 1.0],
        }
    )

@task
def transform(df: pd.DataFrame) -> StructuredDataset:
    # Clean and aggregate into a per-user feature table.
    df = df[df["amount"] > 0]
    agg = df.groupby("user_id", as_index=False)["amount"].sum()
    return StructuredDataset(dataframe=agg)

@task
def load(sd: StructuredDataset) -> int:
    df = sd.open(pd.DataFrame).all()
    return len(df)

@workflow
def main() -> int:
    raw = extract()
    features = transform(df=raw)
    return load(sd=features)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/data_etl_v1.py*

### Flyte 2

```
import pandas as pd
import flyte
import flyte.io

env = flyte.TaskEnvironment(
    name="data_etl",
    image=flyte.Image.from_debian_base().with_pip_packages("pandas", "pyarrow"),
)

@env.task
async def extract() -> pd.DataFrame:
    # Read raw transaction records (stand-in for a real source).
    return pd.DataFrame(
        {
            "user_id": [1, 1, 2, 3, 3, 3],
            "amount": [10.0, 5.0, 20.0, 7.5, 2.5, 1.0],
        }
    )

@env.task
async def transform(df: pd.DataFrame) -> flyte.io.DataFrame:
    # Clean and aggregate into a per-user feature table.
    df = df[df["amount"] > 0]
    agg = df.groupby("user_id", as_index=False)["amount"].sum()
    # StructuredDataset becomes flyte.io.DataFrame.
    return flyte.io.DataFrame.from_df(agg)

@env.task
async def load(sd: flyte.io.DataFrame) -> int:
    df = await sd.open(pd.DataFrame).all()
    return len(df)

@env.task
async def main() -> int:
    raw = await extract()
    features = await transform(raw)
    return await load(features)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/data_etl_v2.py*

## Next

- [ML workloads](./ml-workloads) — training, HPO, and batch inference
- [Parallelism and fan-out](./parallelism) — processing partitions in parallel

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/ml-workloads ===

# ML workloads

The core patterns compose into the workloads data scientists and ML engineers run every day. Each of these is a complete v1→v2 pair in the [examples repo](https://github.com/unionai/unionai-examples/tree/main/v2/user-guide/migration/flyte-2). See [Migration](./overview) for the overall approach.

## Small model training (scikit-learn / XGBoost)

Train a model, persist it as a `File`, and evaluate it. Image, resources, and caching move to the `TaskEnvironment`; `FlyteFile` becomes `flyte.io.File`.

### Flyte 1

```python
import os

import joblib
from flytekit import task, workflow, ImageSpec, Resources, current_context
from flytekit.types.file import FlyteFile
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier

image = ImageSpec(
    name="xgb-image",
    packages=["xgboost", "scikit-learn", "joblib"],
)

@task(container_image=image, requests=Resources(cpu="2", mem="4Gi"))
def train_model(n_estimators: int, max_depth: int) -> FlyteFile:
    data = load_breast_cancer()
    X_train, _, y_train, _ = train_test_split(data.data, data.target, random_state=42)
    model = XGBClassifier(n_estimators=n_estimators, max_depth=max_depth)
    model.fit(X_train, y_train)

    model_path = os.path.join(current_context().working_directory, "model.json")
    joblib.dump(model, model_path)
    return FlyteFile(path=model_path)

@task(container_image=image)
def evaluate(model_file: FlyteFile) -> float:
    model = joblib.load(model_file.download())
    data = load_breast_cancer()
    _, X_test, _, y_test = train_test_split(data.data, data.target, random_state=42)
    return float(model.score(X_test, y_test))

@workflow
def main(n_estimators: int, max_depth: int) -> float:
    model = train_model(n_estimators=n_estimators, max_depth=max_depth)
    return evaluate(model_file=model)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/train_xgboost_v1.py*

### Flyte 2

```
import os

import joblib
import flyte
from flyte.io import File
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier

env = flyte.TaskEnvironment(
    name="train_xgboost",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "xgboost", "scikit-learn", "joblib"
    ),
    resources=flyte.Resources(cpu="2", memory="4Gi"),
)

@env.task
async def train_model(n_estimators: int, max_depth: int) -> File:
    data = load_breast_cancer()
    X_train, _, y_train, _ = train_test_split(data.data, data.target, random_state=42)
    model = XGBClassifier(n_estimators=n_estimators, max_depth=max_depth)
    model.fit(X_train, y_train)

    model_path = os.path.join(os.getcwd(), "model.json")
    joblib.dump(model, model_path)
    return await File.from_local(model_path)

@env.task
async def evaluate(model_file: File) -> float:
    local_path = await model_file.download()
    model = joblib.load(local_path)
    data = load_breast_cancer()
    _, X_test, _, y_test = train_test_split(data.data, data.target, random_state=42)
    return float(model.score(X_test, y_test))

@env.task
async def main(n_estimators: int, max_depth: int) -> float:
    model = await train_model(n_estimators, max_depth)
    return await evaluate(model)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/train_xgboost_v2.py*

## Hyperparameter optimization

Fan out one training run per hyperparameter, then pick the best. In Flyte 1 the grid search runs through `map_task` and the "pick the best" step must itself be a task. In Flyte 2 you `gather` the runs and select the winner in plain Python.

### Flyte 1

```python
from flytekit import task, workflow, map_task
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score

@task
def get_grid() -> list[int]:
    return [2, 4, 8, 16]

@task
def train_eval(max_depth: int) -> float:
    data = load_iris()
    model = RandomForestClassifier(max_depth=max_depth, random_state=42)
    scores = cross_val_score(model, data.data, data.target, cv=3)
    return float(scores.mean())

@task
def best_score(scores: list[float]) -> float:
    return max(scores)

@workflow
def main() -> float:
    grid = get_grid()
    # Fan out one training run per hyperparameter value.
    scores = map_task(train_eval)(max_depth=grid)
    return best_score(scores=scores)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/hpo_v1.py*

### Flyte 2

```
import asyncio

import flyte
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score

env = flyte.TaskEnvironment(
    name="hpo",
    image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn"),
)

@env.task
async def train_eval(max_depth: int) -> float:
    data = load_iris()
    model = RandomForestClassifier(max_depth=max_depth, random_state=42)
    scores = cross_val_score(model, data.data, data.target, cv=3)
    return float(scores.mean())

@env.task
async def main() -> dict:
    grid = [2, 4, 8, 16]
    # Fan out one training run per hyperparameter value...
    scores = await asyncio.gather(*[train_eval(d) for d in grid])
    # ...then pick the best in plain Python (impossible in a Flyte 1 workflow).
    best_idx = max(range(len(scores)), key=lambda i: scores[i])
    return {"best_max_depth": grid[best_idx], "best_score": scores[best_idx]}
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/hpo_v2.py*

## Large model training (deep learning)

GPU configuration moves to the `TaskEnvironment`: the Flyte 1 `Resources(gpu="1")` plus a separate `accelerator=T4` become a single `gpu="T4:1"` string on `flyte.Resources`.

### Flyte 1

```python
from flytekit import task, workflow, ImageSpec, Resources
from flytekit.extras.accelerators import T4
import torch
import torch.nn as nn

image = ImageSpec(
    name="dl-image",
    packages=["torch"],
)

@task(
    container_image=image,
    requests=Resources(cpu="4", mem="16Gi", gpu="1"),
    accelerator=T4,
)
def train(epochs: int) -> float:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = nn.Linear(10, 1).to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()

    X = torch.randn(128, 10, device=device)
    y = torch.randn(128, 1, device=device)

    loss = torch.tensor(0.0)
    for _ in range(epochs):
        optimizer.zero_grad()
        loss = loss_fn(model(X), y)
        loss.backward()
        optimizer.step()
    return float(loss.item())

@workflow
def main(epochs: int) -> float:
    return train(epochs=epochs)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/train_deep_learning_v1.py*

### Flyte 2

```
import flyte
import torch
import torch.nn as nn

# GPU type and count go in a single "T4:1"-style string. For multi-node
# distributed training, wrap the training task with the torch elastic plugin.
env = flyte.TaskEnvironment(
    name="train_deep_learning",
    image=flyte.Image.from_debian_base().with_pip_packages("torch"),
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="T4:1"),
)

@env.task
async def train(epochs: int) -> float:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = nn.Linear(10, 1).to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()

    X = torch.randn(128, 10, device=device)
    y = torch.randn(128, 1, device=device)

    loss = torch.tensor(0.0)
    for _ in range(epochs):
        optimizer.zero_grad()
        loss = loss_fn(model(X), y)
        loss.backward()
        optimizer.step()
    return float(loss.item())

@env.task
async def main(epochs: int) -> float:
    return await train(epochs)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/train_deep_learning_v2.py*

For multi-node distributed training (PyTorch elastic, etc.), see [Resources](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/resources) and the plugin integrations.

## Batch inference

Load a trained model once and score many batches in parallel. `map_task` with a `partial`-bound model becomes `asyncio.gather` over the batches, reusing the same model reference.

### Flyte 1

```python
import os
from functools import partial

import joblib
from flytekit import task, workflow, map_task, ImageSpec, current_context
from flytekit.types.file import FlyteFile
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

image = ImageSpec(name="inference-image", packages=["scikit-learn", "joblib"])

@task(container_image=image)
def train_model() -> FlyteFile:
    data = load_iris()
    model = RandomForestClassifier().fit(data.data, data.target)
    model_path = os.path.join(current_context().working_directory, "model.joblib")
    joblib.dump(model, model_path)
    return FlyteFile(path=model_path)

@task(container_image=image)
def get_batches() -> list[list[list[float]]]:
    data = load_iris()
    rows = data.data.tolist()
    # Split the rows into batches of 30.
    return [rows[i : i + 30] for i in range(0, len(rows), 30)]

@task(container_image=image)
def score_batch(model_file: FlyteFile, batch: list[list[float]]) -> list[int]:
    model = joblib.load(model_file.download())
    return [int(p) for p in model.predict(batch)]

@workflow
def main() -> list[list[int]]:
    model = train_model()
    batches = get_batches()
    return map_task(partial(score_batch, model_file=model))(batch=batches)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/batch_inference_v1.py*

### Flyte 2

```
import asyncio
import os

import joblib
import flyte
from flyte.io import File
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

env = flyte.TaskEnvironment(
    name="batch_inference",
    image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn", "joblib"),
)

@env.task
async def train_model() -> File:
    data = load_iris()
    model = RandomForestClassifier().fit(data.data, data.target)
    model_path = os.path.join(os.getcwd(), "model.joblib")
    joblib.dump(model, model_path)
    return await File.from_local(model_path)

@env.task
async def score_batch(model_file: File, batch: list[list[float]]) -> list[int]:
    local_path = await model_file.download()
    model = joblib.load(local_path)
    return [int(p) for p in model.predict(batch)]

@env.task
async def main() -> list[list[int]]:
    model = await train_model()
    rows = load_iris().data.tolist()
    batches = [rows[i : i + 30] for i in range(0, len(rows), 30)]
    # Score every batch in parallel, reusing the same model reference.
    coros = [score_batch(model, batch) for batch in batches]
    return list(await asyncio.gather(*coros))
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/batch_inference_v2.py*

## A complete example: end-to-end ML pipeline

Putting it together — a load / train / evaluate pipeline shows the image, resources, caching, file I/O, and orchestration changes in one place.

### Flyte 1

```python
import os

import joblib
import pandas as pd
from flytekit import task, workflow, ImageSpec, Resources, current_context
from flytekit.types.file import FlyteFile
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

image = ImageSpec(
    name="ml-image",
    packages=["pandas", "scikit-learn", "joblib"],
)

@task(
    container_image=image,
    requests=Resources(cpu="2", mem="4Gi"),
    cache=True,
    cache_version="1.0",
)
def load_data() -> pd.DataFrame:
    data = load_iris(as_frame=True)
    df = data.frame
    df["species"] = data.target
    return df

@task(container_image=image)
def train_model(data: pd.DataFrame) -> FlyteFile:
    model = RandomForestClassifier()
    X = data.drop("species", axis=1)
    y = data["species"]
    model.fit(X, y)

    model_path = os.path.join(current_context().working_directory, "model.joblib")
    joblib.dump(model, model_path)
    return FlyteFile(path=model_path)

@task(container_image=image)
def evaluate(model_file: FlyteFile, data: pd.DataFrame) -> float:
    model = joblib.load(model_file.download())
    X = data.drop("species", axis=1)
    y = data["species"]
    return float(model.score(X, y))

@workflow
def main() -> float:
    data = load_data()
    model = train_model(data=data)
    return evaluate(model_file=model, data=data)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/ml_pipeline_v1.py*

### Flyte 2

```
import os

import joblib
import pandas as pd
import flyte
from flyte.io import File
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

# Image, resources, and cache are set once on the TaskEnvironment.
env = flyte.TaskEnvironment(
    name="ml_pipeline",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "pandas", "scikit-learn", "joblib"
    ),
    resources=flyte.Resources(cpu="2", memory="4Gi"),
    cache="auto",
)

@env.task
async def load_data() -> pd.DataFrame:
    data = load_iris(as_frame=True)
    df = data.frame
    df["species"] = data.target
    return df

@env.task
async def train_model(data: pd.DataFrame) -> File:
    model = RandomForestClassifier()
    X = data.drop("species", axis=1)
    y = data["species"]
    model.fit(X, y)

    model_path = os.path.join(os.getcwd(), "model.joblib")
    joblib.dump(model, model_path)
    return await File.from_local(model_path)

@env.task
async def evaluate(model_file: File, data: pd.DataFrame) -> float:
    local_path = await model_file.download()
    model = joblib.load(local_path)
    X = data.drop("species", axis=1)
    y = data["species"]
    return float(model.score(X, y))

# The "workflow" is just an orchestrating task.
@env.task
async def main() -> float:
    data = await load_data()
    model = await train_model(data)
    return await evaluate(model, data)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/migration/flyte-2/ml_pipeline_v2.py*

## Next

- [New in Flyte 2](./new-in-flyte-2) — real-time serving, apps, and sandboxing
- [Gotchas and caveats](./gotchas-and-caveats) — caveats of the new execution model

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/new-in-flyte-2 ===

# New in Flyte 2: patterns that weren't possible in Flyte 1

Flyte 1 was a batch orchestration system: everything ran as a finite DAG that started, did work, and finished. Flyte 2 keeps all of that and adds **long-running services**, **high-throughput batch inference**, and **sandboxed code execution** — so the same project that trains your model can also serve it, host a dashboard, saturate a GPU, or safely run LLM-generated code. There is no v1 counterpart to migrate here; these are net-new capabilities.

## Real-time inference and model serving

Instead of scoring a batch and exiting, you can stand up an always-on REST endpoint from a `FastAPIAppEnvironment` and deploy it with `flyte.deploy`. The app can load a model artifact produced by one of your training tasks.

```
app = FastAPI(title="ML Model API")

# Define request/response models
class PredictionRequest(BaseModel):
    feature1: float
    feature2: float
    feature3: float

class PredictionResponse(BaseModel):
    prediction: float
    probability: float

# Load model (you would typically load this from storage)
model = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global model
    model_path = os.getenv("MODEL_PATH", "/app/models/model.joblib")
    # In production, load from your storage
    if os.path.exists(model_path):
        with open(model_path, "rb") as f:
            model = joblib.load(f)
    yield

@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    # Make prediction
    # prediction = model.predict([[request.feature1, request.feature2, request.feature3]])

    # Dummy prediction for demo
    prediction = 0.85
    probability = 0.92

    return PredictionResponse(
        prediction=prediction,
        probability=probability,
    )

env = FastAPIAppEnvironment(
    name="ml-model-api",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
        "scikit-learn",
        "pydantic",
        "joblib",
    ),
    parameters=[
        flyte.app.Parameter(
            name="model_file",
            value=flyte.io.File.from_existing_remote("s3://bucket/models/model.joblib"),
            mount="/app/models",
            env_var="MODEL_PATH",
        ),
    ],
    resources=flyte.Resources(cpu=2, memory="2Gi"),
    requires_auth=False,
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/fastapi/ml_model_serving.py*

See [Build apps](https://www.union.ai/docs/latest/flyte/user-guide/migration/apps/build-apps/_index) and [Serve and deploy apps](https://www.union.ai/docs/latest/flyte/user-guide/migration/apps/serve-and-deploy-apps/_index).

## LLM serving

For serving large language models, the `flyteplugins-vllm` integration gives you a production-grade vLLM server (with autoscaling to zero) in a few lines.

```
vllm_app = VLLMAppEnvironment(
    name="my-llm-app",
    model_hf_path="Qwen/Qwen3-0.6B",  # HuggingFace model path
    model_id="qwen3-0.6b",  # Model ID exposed by vLLM
    resources=flyte.Resources(
        cpu="4",
        memory="16Gi",
        gpu="L40s:1",  # GPU required for LLM serving
        disk="10Gi",
    ),
    scaling=flyte.app.Scaling(
        replicas=(0, 1),
        scaledown_after=300,  # Scale down after 5 minutes of inactivity
    ),
    requires_auth=False,
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/build-apps/vllm/basic_vllm.py*

See [vLLM](https://www.union.ai/docs/latest/flyte/user-guide/apps/native-app-integrations/vllm-app) and the other [Native app integrations](https://www.union.ai/docs/latest/flyte/user-guide/migration/apps/native-app-integrations/_index) (SGLang, Streamlit, FastAPI).

## App serving (dashboards and APIs)

Any web app — a Streamlit dashboard for exploring results, a Gradio demo, a Flask backend — runs as an `AppEnvironment`. Configure the image, resources, port, autoscaling, and a custom subdomain, then `flyte.serve` it.

```
app_env = flyte.app.AppEnvironment(
    name="hello-world-app",
    image=image,
    args=["streamlit", "hello", "--server.port", "8080"],
    port=8080,
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    requires_auth=False,
    domain=flyte.app.Domain(subdomain="hello"),
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/hello-world-app.py*

See [Introducing apps](https://www.union.ai/docs/latest/flyte/user-guide/get-started/core-concepts/introducing-apps) and [Configure apps](https://www.union.ai/docs/latest/flyte/user-guide/migration/apps/configure-apps/_index).

## Dynamic batching for GPU inference

For in-process batch inference, `DynamicBatcher` from `flyte.extras` keeps an expensive GPU saturated: async producers load and preprocess data concurrently while a single consumer feeds the model in optimally-sized batches, with built-in backpressure. This replaces the Flyte 1 pattern of standing up a separate inference server just to get request batching.

```python
import asyncio
from flyte.extras import DynamicBatcher

async with DynamicBatcher(
    process_fn=run_inference,   # takes a batch, returns results in the same order
    target_batch_cost=1000,     # cost budget per batch
    max_batch_size=64,          # hard cap on records per batch
    batch_timeout_s=0.05,       # max wait before dispatching a partial batch
) as batcher:
    futures = [await batcher.submit(record) for record in records]
    results = await asyncio.gather(*futures)
```

`submit()` is non-blocking and returns a `Future`; when the queue is full it applies backpressure automatically. See [Batch inference](https://www.union.ai/docs/latest/flyte/user-guide/run-scaling/batch-inference), which also covers `TokenBatcher` for token-aware LLM batching.

## Sandboxed code execution

`flyte.sandbox.create()` runs arbitrary Python code or shell commands inside an ephemeral, single-use Docker container — built on demand from declared dependencies, executed once, then discarded. Only declared inputs go in and only declared outputs come back, which makes it the safe way to run **untrusted code, most importantly code generated by an LLM**.

```
# sandbox_environment provides the base runtime for code sandboxes.
# Include it in depends_on so the sandbox runtime is available when tasks execute.
env = flyte.TaskEnvironment(
    name="sandbox-demo",
    image=flyte.Image.from_debian_base(name="sandbox-demo"),
    depends_on=[sandbox_environment],
)

# Auto-IO mode: pure computation. The code string runs in an isolated sandbox;
# only the declared inputs go in and only the declared outputs come back.
sum_sandbox = flyte.sandbox.create(
    name="sum-to-n",
    code="total = sum(range(n + 1)) if conditional else 0",
    inputs={"n": int, "conditional": bool},
    outputs={"total": int},
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/sandboxing/code_sandbox.py*

Call it from a task with `await sum_sandbox.run.aio(n=10, conditional=True)`. See [Code sandboxing](https://www.union.ai/docs/latest/flyte/user-guide/agents/sandboxing/code-sandboxing). This also powers **code mode** (programmatic tool calling), where an agent writes a whole program instead of emitting one tool call at a time — see [Programmatic tool calling for agents](https://www.union.ai/docs/latest/flyte/user-guide/agents/sandboxing/code-mode).

## Next

- [Gotchas and caveats](./gotchas-and-caveats) — caveats of the new execution model
- [Hybrid v1 and v2 pipelines](./hybrid-pipelines) — calling between v1 and v2 during the transition

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/hybrid-pipelines ===

# Hybrid v1 and v2 pipelines

Migrations rarely happen all at once. For a while you'll have Flyte 1 and Flyte 2 workloads running side by side, and you'll want them to call each other: a Flyte 1 workflow that kicks off a newly ported Flyte 2 task, or a Flyte 2 task that triggers a workflow that hasn't been migrated yet.

You can bridge the two in both directions. The idea is the same each way: one task installs **both** SDKs, authenticates to the **other** control plane, fetches the entity it wants to run, and launches it.

> **📝 Note**
>
> The bridging task acts as a driver: it authenticates to a remote control plane and launches work there. Keep it lightweight and focused on orchestration — see [Gotchas and caveats](./gotchas-and-caveats).

## Running a Flyte 2 task from a Flyte 1 workflow

The bridge is a single Flyte 1 task — call it `launch_v2_from_v1` — that runs the Flyte 2 client.

**High-level steps:**

1. Give the `launch_v2_from_v1` task an image with **both** `flytekit` (Flyte 1) and `flyte` (Flyte 2) installed.
2. Give it a Flyte 2 **API key** (see **Migration > From Flyte 1 to 2 > Hybrid v1 and v2 pipelines > Running a Flyte 2 task from a Flyte 1 workflow > 1. Create the API key and store it** below — how you create one differs between Union and open-source Flyte). The key is the `export FLYTE_API_KEY="..."` value it produces.
3. Make the key available to the task, either by storing it as a secret it can read, or by injecting it as the `FLYTE_API_KEY` environment variable.
4. Authenticate inside the task with `flyte.init_from_api_key()`.
5. Fetch the deployed Flyte 2 task with `flyte.remote.Task.get(...)` and run it with `flyte.run(...)`.

### 1. Create the API key and store it

Obtain a Flyte 2 API key from your control plane's authentication setup and store it as a secret the bridging task can read — see [Run on a remote cluster](https://www.union.ai/docs/latest/flyte/user-guide/get-started/run-modes/running-remote) for the authentication options.

You can also make the same value available to your Flyte 1 task as a secret through your existing Flyte 1 secret workflow, or set it directly as the `FLYTE_API_KEY` environment variable on the task.

### 2. Write the bridging task

```python
import flytekit
from flytekit import task, workflow, ImageSpec, Secret, current_context

# The bridge image needs BOTH the v1 (flytekit) and v2 (flyte) SDKs.
bridge_image = ImageSpec(
    name="v1-to-v2-bridge",
    packages=["flytekit", "flyte"],
)

@task(
    container_image=bridge_image,
    secret_requests=[Secret(group="flyte", key="flyte_api_key")],
)
def launch_v2_from_v1(x: int) -> str:
    import flyte
    import flyte.remote

    # Authenticate to the Flyte 2 control plane with the API key.
    # Option A: read the mounted secret and pass it explicitly.
    api_key = current_context().secrets.get(group="flyte", key="flyte_api_key")
    flyte.init_from_api_key(api_key=api_key)

    # Option B: if FLYTE_API_KEY is set as an env var, no argument is needed:
    #     flyte.init_from_api_key()

    # Fetch the deployed Flyte 2 task and run it.
    remote_v2_task = flyte.remote.Task.get(
        "my_v2_env.process",
        auto_version="latest",
    )
    run = flyte.run(remote_v2_task, x=x)
    run.wait()  # optional: block until the v2 run finishes
    return run.url

@workflow
def main(x: int) -> str:
    return launch_v2_from_v1(x=x)
```

The referenced Flyte 2 task (`my_v2_env.process` above) must be **deployed** before the bridge runs — `flyte.remote.Task.get()` looks it up by `env_name.task_name`. See [Remote tasks](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/remote-tasks) for versioning options (`auto_version="latest"`, `version="v1.2.3"`) and `flyte.run` details.

> **📝 Note**
>
> `flyte.init_from_api_key()` is required here — do **not** use `flyte.init_from_config()`, which reads a `config.yaml` that has no API-key field. See [Run on a remote cluster](https://www.union.ai/docs/latest/flyte/user-guide/get-started/run-modes/running-remote) for the authentication methods.

## Running a Flyte 1 workflow from a Flyte 2 task

The reverse works the same way: a Flyte 2 task installs the Flyte 1 client and uses `FlyteRemote` to launch a Flyte 1 workflow.

**High-level steps:**

1. Give the `launch_v1_from_v2` `TaskEnvironment` an image with the Flyte 1 client (`flytekit`) installed.
2. Provide the task with credentials for the Flyte 1 control plane (a config file or API key, supplied as a secret).
3. Instantiate a `FlyteRemote` client inside the task.
4. Fetch the Flyte 1 workflow with `fetch_workflow(...)`.
5. Launch it with `execute(...)`.

```python
import flyte

env = flyte.TaskEnvironment(
    name="v2_to_v1_bridge",
    # The image needs the Flyte 1 client installed.
    image=flyte.Image.from_debian_base().with_pip_packages("flytekit"),
    # Supply credentials for the Flyte 1 control plane (config or API key).
    secrets=[flyte.Secret(key="v1_client_secret", as_env_var="V1_CLIENT_SECRET")],
)

@env.task
async def launch_v1_from_v2(x: int) -> str:
    from flytekit.remote import FlyteRemote
    from flytekit.configuration import Config

    # Point the client at your Flyte 1 cluster.
    remote = FlyteRemote(
        config=Config.for_endpoint(endpoint="my-v1-cluster.example.com"),
        default_project="flytesnacks",
        default_domain="development",
    )

    # Fetch the deployed Flyte 1 workflow and execute it.
    wf = remote.fetch_workflow(name="my_v1_module.main", version="v1.2.3")
    execution = remote.execute(wf, inputs={"x": x}, wait=True)
    return execution.id.name
```

## Considerations

- **Both SDKs in one image.** The bridging task installs `flytekit` and `flyte` together. Pin versions and watch for dependency conflicts; keep the bridge image minimal.
- **Deploy the callee first.** For the v1→v2 direction, the Flyte 2 task must be deployed (`flyte deploy`) before `flyte.remote.Task.get()` can resolve it. For the v2→v1 direction, the Flyte 1 workflow must be registered on its cluster.
- **Wait vs. fire-and-forget.** Both `run.wait()` (v2) and `execute(..., wait=True)` (v1) block until the launched run finishes. Omit them to launch and return immediately, then poll or hand off the execution URL.
- **Credentials cross a boundary.** The bridge authenticates to a *different* control plane than the one it runs on. Store the API key or client credentials as a secret — never hard-code them. See [Secrets](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-configuration/secrets) and [Run on a remote cluster](https://www.union.ai/docs/latest/flyte/user-guide/get-started/run-modes/running-remote).
- **Keep the bridge lightweight.** Like any orchestrating task, it should mostly launch and assemble results rather than do heavy compute — see [Gotchas and caveats](./gotchas-and-caveats).

## See also

- [Remote tasks](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/remote-tasks) — fetching and running deployed Flyte 2 tasks
- [Run on a remote cluster](https://www.union.ai/docs/latest/flyte/user-guide/get-started/run-modes/running-remote) — authentication methods, including `flyte.init_from_api_key()`

- [Migration](./overview) — mapping Flyte 1 workload patterns to Flyte 2

=== PAGE: https://www.union.ai/docs/latest/flyte/user-guide/migration/flyte-2/gotchas-and-caveats ===

# Gotchas and caveats

Flyte 2 is a substantial change from Flyte 1: each Python task can act as its own engine, launching sub-tasks and assembling their outputs. That flexibility is powerful, but it warrants some caveats when authoring your tasks.

This page starts with a quick list of common gotchas you'll hit during a migration, then covers the deeper execution-model considerations in detail.

## Common gotchas

- **`flyte.map` returns a generator.** Wrap it in `list()` to materialize results, unlike `map_task` which returned a list directly.
- **`memory`, not `mem`.** The `Resources` parameter was renamed, and there are no separate `requests`/`limits` — a single value serves as both.
- **GPUs use a `"T4:1"` string.** Type and count are combined; the separate `accelerator=` argument is gone.
- **Image, resources, and cache live on the `TaskEnvironment`.** Set them once at the env level instead of repeating them on every task decorator.
- **`current_context()` is gone.** Read secrets from environment variables and use `flyte.ctx()` for runtime context.
- **The `>>` ordering operator is gone.** Sequential (sync) calls and sequential `await`s are naturally ordered.
- **Retries no longer have a platform cap.** In Flyte 1 the control plane capped attempts at 3; in Flyte 2 total attempts equal `retries + 1`. Audit any large `retries` values before deploying.
- **You can only `await` async tasks.** Call a sync task from an async context with `.aio()`; see the [Asynchronous model](./overview#asynchronous-model).
- **Pick an entrypoint task name.** There's no `@workflow`, so the top-level task is just a task (commonly `main`); run it with `flyte run module.py main`.
- **Type annotations are more lenient.** Flyte 2 will pickle untyped I/O rather than rejecting it at registration.
- **Keep orchestration lightweight.** A task that calls other tasks acts as a driver pod. Avoid heavy CPU work in it — see **Migration > From Flyte 1 to 2 > Gotchas and caveats > Driver pod requirements** below.

## Non-deterministic behavior

When a task launches another task, a new Action ID is determined.
This ID is a hash of the inputs to the task, the task definition itself, along with some other information.
The fact that this ID is consistently hashed is important when it comes to things like recovery and replay.

For example, assume you have the following tasks

```python
@env.task
async def t1():
    val = get_int_input()
    await t2(int=val)

@env.task
async def t2(val: int): ...
```

If you run `t1`, and it launches the downstream `t2` task, and then the pod executing `t1` fails, when Flyte restarts `t1` it will automatically detect that `t2` is still running and will just use that.
If `t2` ends up finishing in the interim, those results would just be used.

However, if you introduce non-determinism into the picture, then that guarantee is no longer there.
To give a contrived example:

```python
@env.task
async def t1():
    val = get_int_input()
    now = datetime.now()

    if now.second % 2 == 0:
        await t2(int=val)
    else:
        await t3(int=val)
```

Here, depending on what time it is, either `t2` or `t3` may end up running.
In the earlier scenario, if `t1` crashes unexpectedly, and Flyte retries the execution, a different downstream task may get kicked off instead.

### Dealing with non-determinism

As a developer, the best way to manage non-deterministic behavior (if it is unavoidable) is to be able to observe it and see exactly what is happening in your code. Flyte 2 provides precisely the tool needed to enable this: Traces.

With this feature you decorate the sub-task functions in your code with `@trace`, enabling checkpointing, reproducibility and recovery at a fine-grained level. See [Traces](https://www.union.ai/docs/latest/flyte/user-guide/tasks/task-programming/traces) for more details.

## Type safety

In Flyte 1, the top-level workflow was defined by a Python-like DSL that was compiled into a static DAG composed of tasks, each of which was, internally, defined in real Python.
The system was able to guarantee type safety across task boundaries because the task definitions were static and the inputs and outputs were defined in a way that Flytekit could validate them.

In Flyte 2, the top-level workflow is defined by Python code that runs at runtime (unless using a compiled task).
This means that the system can no longer guarantee type safety at the workflow level.

Happily, the Python ecosystem has evolved considerably since Flyte 1, and Python type hints are now a standard way to define types.

Consequently, in Flyte 2, developers should use Python type hints and type checkers like `mypy` to ensure type safety at all levels, including the top-most task (i.e., the "workflow" level).

## No global state

A core principle of Flyte 2 (that is also shared with Flyte 1) is that you should not try to maintain global state across your workflow.
It will not be translated across tasks containers,

In a single process Python program, global variables are available across functions.
In the distributed execution model of Flyte, each task runs in its own container, and each container is isolated from the others.

If there is some state that needs to be preserved, it must be reconstructable through repeated deterministic execution.

## Driver pod requirements

Tasks don't have to kick off downstream tasks of course and may themselves represent a leaf level atomic unit of compute.
However, when tasks do run other tasks, and more so if they assemble the outputs of those other tasks, then that parent task becomes a driver
pod of sorts.
In Flyte 1, this assembling of intermediate outputs was done by Flyte Propeller.
In 2, it's done by the parent task.

This means that the pod running your parent task must be appropriately sized, and should ideally not be CPU-bound, otherwise it slow down downstream evaluation and kickoff of tasks.

For example, if you had this also scenario,

```python
@env.task
async def t_main():
    await t1()
    local_cpu_intensive_function()
    await t2()
```

The pod running `t_main` will hang in between tasks `t1` and `t2`. Your parent tasks should ideally focus only on orchestration.

## OOM risk from materialized I/O

Something maybe more nuanced is that if you're not using the soon-to-be-released ref mode, outputs are actually
materialized. That is, if you have the following scenario,

```python
@env.task
async def produce_1gb_list() -> List[float]: ...

@env.task
async def t1():
    list_floats = produce_1gb_list()
    t2(floats=list_floats)
```

The pod running `t1` needs to have memory to handle that 1 GB of floats. Those numbers will be materialized in that pod's memory.
This can lead to out of memory issues.

Note that `flyte.io.File`, `flyte.io.Dir` and `flyte.io.DataFrame` will not suffer from this because while those are materialized, they're only materialized as pointers to offloaded data, so their memory footprint is much lower.