Flyte 2 Devbox is available today to run a full Flyte backend and UI locally. Preview Flyte 2 for production, hosted on Union.ai

Mapping over inputs

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

Use flyte.map when every item goes through the same task. For fanning out across different tasks, or for full control over how invocations are assembled, use asyncio.gather — see Fanout.

Minimal example

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

from typing import List

import flyte

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


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


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


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

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

Mapping from an async task: flyte.map.aio

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

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

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

Signature and parameters

flyte.map(
    func,               # the task (or functools.partial) to apply to each item
    *args,              # one or more iterables, zipped item-by-item into func's arguments
    group_name=None,    # optional name for the group of mapped actions (UI grouping)
    concurrency=0,      # max actions in flight at once; 0 means unbounded (all at once)
    return_exceptions=True,
)
  • func — the task to map. It receives one item per invocation. To hold some arguments constant across the map, wrap it with functools.partial (see Binding constant arguments).
  • *args — one or more input iterables. With multiple iterables they are zipped: the i-th invocation receives the i-th element of each, matching func’s positional parameters in order.
  • group_name — groups the resulting actions under a single label in the UI (see Grouping actions).
  • concurrency — the maximum number of actions in flight at any moment. 0 (the default) submits everything at once. A positive value bounds the fan-out with a worker pool, so memory stays proportional to concurrency rather than to the total number of items — see Controlling parallel execution.
  • return_exceptions — when True (the default), a failed invocation yields the raised exception as its result instead of aborting the whole map; check each result with isinstance(r, Exception). When False, the first failure stops iteration and raises.

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

Limiting concurrency

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

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

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

Handling errors

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

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


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

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

Binding constant arguments with functools.partial

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

from functools import partial

import flyte

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


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


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

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

When to use flyte.map

Reach for flyte.map when:

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

Use asyncio.gather instead when you are fanning out across different tasks in one batch, or when you need full control over how the coroutines are assembled. Use an asyncio.Semaphore when different task types in the same batch need different concurrency limits.