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 guide for the concepts, and Migration 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 1Flyte 2 (flyte.map)Flyte 2 (asyncio.gather)
map_task_v1.py
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)
map_task_v2.py
@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))
map_task_v2.py
@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))

See the 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:

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 1Flyte 2
data_backfill_v1.py
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)
data_backfill_v2.py
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))

For fine-grained concurrency control (semaphores, as_completed, error handling), see Fanout.

Next