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.
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)
@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))
@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.
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)
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
- Data types and I/O — files, DataFrames, and dataclasses
- ML workloads — training, HPO, and batch inference