Scale AI Agents & Workflows Across Containers Using asyncio

AI engineering tip of the week: Scale AI Agents & Workflows Across Containers Using asyncio

Flyte 2 lets you use `asyncio.gather()`, the same tool you'd use in any async Python app to fan out tasks in parallel. Each awaited task call becomes a distributed task on your cluster, running in its own container with its own resources.

Basic fan-out with asyncio.gather

import asyncio
import flyte

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

@env.task
async def process_item(item: int) -> str:
    await asyncio.sleep(1)  # simulate work
    return f"processed_{item}"

@env.task
async def main(items: list[int]) -> list[str]:
    tasks = [process_item(item) for item in items]
    results = await asyncio.gather(*tasks)
    return list(results)

That’s real parallelism. Each `process_item` call gets its own container, and they all run at the same time. If you have 100 items, you get 100 parallel tasks. No special syntax, no DSL, just Python.

Limit concurrency with a semaphore

Don’t want to flood your cluster with 10,000 tasks at once? Use a standard `asyncio.Semaphore`:

@env.task
async def main_throttled(items: list[int], max_concurrent: int = 50) -> list[str]:
    semaphore = asyncio.Semaphore(max_concurrent)

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

    tasks = [process_with_limit(item) for item in items]
    results = await asyncio.gather(*tasks)
    return list(results)

At most 50 tasks run simultaneously. The rest wait their turn. Same pattern you’d use in a web scraper or API client.

Handle errors without losing everything

By default, `asyncio.gather` cancels everything if one task fails. Use return_exceptions=True to keep the rest going:

@env.task
async def main_resilient(items: list[int]) -> list[str]:
    tasks = [process_item(item) for item in items]
    results = await asyncio.gather(*tasks, return_exceptions=True)

    processed = []
    for i, result in enumerate(results):
        if isinstance(result, Exception):
            print(f"Item {items[i]} failed: {result}")
            processed.append(f"FAILED:{items[i]}")
        else:
            processed.append(result)
    return processed

Kick off runs immediately with create_task

Calling `process_item(5)` doesn’t start anything. It returns a coroutine that sits until something schedules it and `gather()` is what does the scheduling. So nothing launches until you reach the `gather` line, and then you’re stuck there until everything finishes.

`asyncio.create_task()` hands the coroutine to the event loop right away. The container starts spinning up, control comes back to you, and you await the result whenever you actually need it:

@env.task
async def agent_with_warm_start(question: str) -> str:
    # Both containers start launching now -- no await, no gather.
    index = asyncio.create_task(load_vector_index("docs-v3"))
    tools = asyncio.create_task(warm_tool_sandbox())

    # Meanwhile, do local work. The cluster is already busy.
    plan = build_plan(question)

    # Collect the results only when you need them.
    return await run_agent(plan, await index, await tools)

This is the pattern to reach for when a slow dependency (a model load, an index build, a sandbox warm-up) is needed later in the task but could have started at the top. With `gather`, that work waits. With `create_task`, it overlaps.

It also unlocks patterns `gather` can’t express, like racing coroutines and cancelling the losers.

Stream results as they complete

Don’t want to wait for all 100 tasks to finish before you start processing? Use `asyncio.as_completed()`:

@env.task
async def main_streaming(items: list[int]) -> list[str]:
    tasks = [asyncio.create_task(process_item(item)) for item in items]

    results = []
    for completed_task in asyncio.as_completed(tasks):
        result = await completed_task
        results.append(result)
        print(f"Got result: {result} ({len(results)}/{len(items)})")
    return results

Results stream in as they finish. This is great for reducing downstream work while upstream tasks are still running.

Compose different tasks together

Unlike `flyte.map()` (which runs the same task on different inputs), asyncio.gather lets you run completely different tasks in parallel:

@env.task
async def fetch_data() -> str:
    return "raw data"

@env.task
async def load_model() -> str:
    return "model_v2"

@env.task
async def pipeline() -> dict:
    data, model = await asyncio.gather(
        fetch_data(),
        load_model(),
    )
    return {"data": data, "model": model}

Two different tasks, running at the same time, results collected when both finish. This is the bread and butter of efficient pipeline design.

Flyte 2 also still offers `flyte.map()` as fanout method too.

When to use gather vs flyte.map

`asyncio.gather()` is best when:

• You need different tasks running in parallel
• You want to stream results with `as_completed()`
• You need flexible composition and error handling

`flyte.map()` is best when:

• You have the same task running on many inputs
• You want a built-in concurrency limit `(concurrency=N)`
• You need to work with sync (non-async) tasks

Full async docs: https://www.union.ai/docs/v2/union/user-guide/task-programming/fanout/

See what’s happening in the Flyte Community

Latest from the blog

• Fine-tune an LLM with LoRA & QLoRA in a Flyte Pipeline - Read on union.ai
• From Billions to Bytes: The Science of Shrinking Neural Networks - Read on Union.ai
• Container-enabled asyncio is all you need (to build Pythonic AI workflows at scale) - Read on Union.ai
• Flyte MCP: give your local coding agent control-plane superpowers - Read on Union.ai

Recent talks & recordings

• Building Code Mode Agents - Watch on YouTube
• LLM fine-tuning with GRPO - Watch on YouTube
• LLM fine-tuning with LoRA & QLoRA - Watch on YouTube

Upcoming events

• July 23rd: Seattle Agent Loop Hacknight: Build Agents with Flyte - RSVP on Luma
• July 28th: Seattle TwelveLabs + Qdrant: AI Systems for Video Embeddings and Search - RSVP on Luma

Releases & updates

• Flyte 2 OSS: Backend Devbox and Reimagined UI - Read on Union
• June’s release brought first-class agents with memory and tool approval, SDK-authored MCP servers, backoff retries and per-attempt timeouts, multi-pod log streaming, and beta queues and events APIs. - Read the Release notes

<div class="button-group is-center"><a class="button" target="_blank" rel="noopener noreferrer" href="https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox/">Download Devbox</a></div>

From the community

• AI Book Club: RAG with Python Cookbook - RSVP on Luma
• Open-Source Text-to-Speech: “Natural” Voices - RSVP on Luma

That’s all for this week! —Sage Elliott