Support mapped task expansion

[kriben]

Loading N surfaces is a map operation. TaskMaestro could expand one task declaration over a configured mapping:

- task: resinsight.load_regular_surface
  name: load_surfaces
  map:
    over: surfaces
    key_as: surface_name
    value_as: path
  depends_on:
    resinsight: resinsight.connect

The output could be a typed dictionary keyed by the input aliases. Desirable semantics include:

• bounded parallelism;
• fail-fast or collect-all-errors modes;
• per-item lifecycle events and timings;
• stable result ordering;
• correlation IDs derived from the parent task; and
• visualization of mapped instances without an unreadable expanded DAG.

[kriben]

Design and implementation plan for issue #8: mapped task expansion

Summary

Add mapped task execution: one logical workflow task is invoked once for every entry in a configured mapping, and its outputs are aggregated into an ordered dict[str, O] for downstream tasks.

The first version will be deliberately sequential. Bounded concurrency will be handled in a follow-up issue so that the mapping API and execution semantics can be established without introducing thread safety, cancellation, and timeout complications.

Initial scope

The first version will support:

• Mapping over a dict[str, T] stored in JobConfiguration
• Injecting each mapping key and value into named task input fields
• Ordinary upstream dependencies shared by every mapped invocation
• Additional static config_fields shared by every mapped invocation
• Sequential execution in mapping declaration order
• fail_fast and collect_all error modes
• A raw dict[str, O] aggregate output
• Per-item lifecycle events, results, timings, and derived correlation IDs
• Empty mappings
• Python and YAML workflow definitions
• Compact Mermaid visualization

Deferred:

• Concurrent execution and max_concurrency
• Mapping over an upstream task's dynamically produced collection
• Mapping over lists or other iterable types
• Unpacking one mapped value into multiple input fields
• Non-string mapping keys
• Nested workflow composition when the inner workflow's result task is itself mapped

1. User-facing model

Given a task that loads one surface:

class LoadSurfaceInput(BaseModel):
    resinsight: RipsInstance
    surface_name: str
    path: str
    unit: str


class LoadRegularSurface(Task[LoadSurfaceInput, RegularSurface]):
    def run(
        self,
        input: LoadSurfaceInput,
        ctx: ExecutionContext,
    ) -> RegularSurface:
        ...

one logical workflow declaration can execute it for several configured surfaces.

Python API

Add a public immutable TaskMap model and a mapped_over argument to WorkflowBuilder.add_task():

from taskmaestro import TaskMap

workflow = (
    Workflow.builder("create_grid")
    .add_task(ConnectToResInsight)
    .add_task(
        LoadRegularSurface,
        name="load_surfaces",
        depends_on={
            "resinsight": ConnectToResInsight,
        },
        config_fields=["unit"],
        mapped_over=TaskMap(
            over="surfaces",
            key_as="surface_name",
            value_as="path",
            error_mode="fail_fast",
        ),
    )
    .add_task(
        CreateGrid,
        depends_on={
            "surfaces": "load_surfaces",
        },
    )
    .build()
)

Configuration:

job_configuration = JobConfiguration(
    {
        "load_surfaces": {
            "unit": "meters",
            "surfaces": {
                "top": "/data/top.irap",
                "middle": "/data/middle.irap",
                "base": "/data/base.irap",
            },
        }
    }
)

The downstream task receives:

class CreateGridInput(BaseModel):
    surfaces: dict[str, RegularSurface]

The map source surfaces is consumed by TaskMap; it is not itself a field on LoadSurfaceInput and must not be listed in config_fields.

YAML API

Workflow:

workflow:
  name: create_grid
  tasks:
    - task: resinsight.connect

    - task: resinsight.load_regular_surface
      name: load_surfaces
      map:
        over: surfaces
        key_as: surface_name
        value_as: path
        error_mode: fail_fast
      depends_on:
        resinsight: resinsight.connect
      config_fields:
        - unit

    - task: resinsight.create_grid
      depends_on:
        surfaces: load_surfaces

Input:

load_surfaces:
  unit: meters
  surfaces:
    top: /data/top.irap
    middle: /data/middle.irap
    base: /data/base.irap

When YAML infers config_fields from per-task input, it must exclude the field named by map.over.

2. Mapping declaration

Add an immutable public model:

@dataclass(frozen=True)
class TaskMap:
    over: str
    key_as: str
    value_as: str
    error_mode: Literal["fail_fast", "collect_all"] = "fail_fast"

Validation on construction or workflow build:

• over, key_as, and value_as must be non-empty strings.
• key_as and value_as must be different fields.
• error_mode must be fail_fast or collect_all.
• No concurrency option is exposed in this version.

Store map declarations by registered task instance name:

Workflow._task_maps: dict[str, TaskMap]

Expose:

workflow.get_task_map(task_name) -> TaskMap | None
workflow.is_mapped_task(task_name) -> bool
workflow.get_output_annotation(task_name) -> Any

get_output_annotation() returns the task's ordinary declared output type for a normal node and dict[str, O] for a mapped node.

3. Build-time validation

For a mapped task, the workflow builder must validate:

1. key_as exists in the task input model.
2. value_as exists in the task input model.
3. The key field accepts str.
4. The map-injected key and value fields are not also supplied through depends_on.
5. The map-injected fields are not listed in config_fields.
6. All other required task input fields are covered by dependencies, static configuration, or defaults.
7. Downstream consumers validate against the mapped node's effective output type, dict[str, O].
8. A mapped task can participate normally in cycle detection and sink detection because it remains one static DAG node.

The map value's concrete type is not available until a JobConfiguration is provided. Each value is therefore validated against the value_as input field when the job is created or the mapped task begins execution.

Similarly, every mapping key is validated against the key_as field, though version one requires source keys to be strings.

4. Job configuration validation

For every mapped task, validate that:

• The task has an entry in JobConfiguration.
• The over field exists in that entry.
• The over value is a mapping.
• Every key is a string.
• Every key can populate key_as.
• Every value can populate value_as.
• Remaining configured fields correspond to declared config_fields.

The mapping's insertion order is its semantic order.

An empty mapping is valid and produces an empty output dictionary without invoking the task.

Validation should happen as early as practical, preferably in Job construction, so malformed configuration does not leave a running job in a partially executed state.

5. Effective output type

An ordinary task invocation still returns its declared O, and every item output is checked with:

isinstance(output, get_output_type(task_cls))

The logical mapped workflow node produces:

dict[str, O]

For example:

LoadRegularSurface -> RegularSurface
load_surfaces      -> dict[str, RegularSurface]

Workflow validation must use the effective output annotation of the registered task instance rather than always calling get_output_type(task_cls).

This applies to:

• Whole-output dependencies
• Fan-in fields
• Collection dependencies from issue Add possibility that a output of multiple tasks become input for a task #4
• Output-field routing validation
• Mermaid labels
• Workflow result typing

Because mapped aggregate outputs are raw dictionaries rather than BaseModel instances, internal output annotations must be broadened where needed:

Runner outputs: dict[str, object]
TaskResult.output: object | None
Job.result: object | None
Hook aggregate output parameters: object

Individual mapped item outputs remain BaseModel instances under the current Task[I, O] contract.

Initial nested-workflow limitation

A workflow whose final result task is mapped cannot initially be wrapped with Workflow.as_task(), because Task currently requires a BaseModel output while the mapped result is a raw dictionary. Such use should raise a clear WorkflowDefinitionError rather than failing during generic type introspection.

A later change can broaden the task output contract or introduce a typed workflow-boundary wrapper if this use case is needed.

6. Sequential execution

Add focused runner methods rather than expanding the existing main conditional:

_build_task_input(...)
_run_single_task(...)
_run_mapped_task(...)
_run_mapped_item(...)

For a mapped task, the runner performs these steps:

1. Resolve shared upstream dependency values once.
2. Read common static configuration fields once.
3. Read the mapping named by TaskMap.over.
4. Emit the logical mapped-task start event.
5. Iterate through mapping entries in declaration order.
6. Create a fresh task instance for each item.
7. Build input from shared dependencies, common config, and mapped key/value fields.
8. Validate the complete input through the task's Pydantic input model.
9. Create a child execution context.
10. Emit item-start.
11. Apply the task's timeout to that item.
12. Execute and validate the item output.
13. Emit item-complete or item-fail.
14. Store successful output under the original key.
15. After success, emit logical task-complete with the ordered aggregate dictionary.

Conceptually:

outputs: dict[str, RegularSurface] = {}

for surface_name, path in configured_surfaces.items():
    item_input = LoadSurfaceInput.model_validate(
        {
            "resinsight": connection,
            "unit": "meters",
            "surface_name": surface_name,
            "path": path,
        }
    )
    output = LoadRegularSurface().run(item_input, child_context)
    outputs[surface_name] = output

The mapped node is still one node in the static DAG. Downstream tasks run only after mapped execution has completed successfully.

7. Failure semantics

fail_fast (default)

Stop immediately after the first failed item:

top: success
middle: failed
base: not run

Then:

• Emit item-fail for middle.
• Emit logical task-fail for load_surfaces.
• Mark the job failed.
• Do not execute downstream tasks.

collect_all

Continue executing remaining items after failures:

top: success
middle: failed
base: success

After all items have been attempted:

• Emit item events for every attempted item.
• Raise or record one MappedTaskExecutionError containing errors by mapping key.
• Emit logical task-fail.
• Mark the job failed.
• Do not execute downstream tasks.

Add:

class MappedTaskExecutionError(TaskExecutionError):
    errors: dict[str, Exception]

Its string representation should be stable and list failures in configuration order.

Partial aggregate outputs must not be exposed to downstream tasks after either failure mode.

8. Timeouts

Because execution is sequential, the existing signal-based task timeout can be applied independently to each item.

If:

class LoadRegularSurface(...):
    timeout_seconds = 30

then every mapped invocation gets up to 30 seconds.

The job-level timeout continues to cover the complete mapped operation according to the runner's existing timeout semantics.

A timeout is treated like any other item failure:

• fail_fast: stop immediately.
• collect_all: record the timeout and continue.

9. Lifecycle hooks and results

Preserve one logical result per statically declared workflow task in:

job.task_results

The parent result for a successful mapped task contains the aggregate dictionary.

Add per-item results separately:

job.mapped_item_results: dict[str, list[TaskResult]]

Entries use synthetic names for readability:

load_surfaces[top]
load_surfaces[middle]
load_surfaces[base]

Add hook events:

class Event(StrEnum):
    ...
    MAP_ITEM_START = "map_item_start"
    MAP_ITEM_COMPLETE = "map_item_complete"
    MAP_ITEM_FAIL = "map_item_fail"

And no-op BaseHook methods:

def on_map_item_start(self, job, task, key): ...
def on_map_item_complete(self, job, task, key, output): ...
def on_map_item_fail(self, job, task, key, error): ...

Existing hooks remain source-compatible. Built-in logging and timing hooks should be extended to expose item-level information.

The logical mapped task continues to emit the ordinary task start/complete/fail events once, preserving the existing declared-task lifecycle.

10. Child execution contexts

Add a child-context factory:

child = parent_context.child(
    task_name="load_surfaces",
    item_key="top",
)

Each child context should:

• Have a human-readable derived correlation ID, such as parent:load_surfaces:top.
• Have a key-specific scratch directory.
• Reuse the parent's logger.
• Share the parent's service registry.
• Preserve access to the parent correlation ID.

Map keys must be escaped or sanitized when used in correlation IDs and filesystem paths. Different raw keys must not collapse to the same scratch directory; append a stable short hash if sanitization can cause collisions.

Even though execution is initially sequential, each invocation receives a distinct task and context instance. This makes later concurrent execution safer.

11. Mermaid visualization

Do not expand configured items into separate static DAG nodes. Item count belongs to the job configuration and may be large.

Render one compact mapped node:

graph TD
    connect -->|resinsight: RipsInstance| load_surfaces
    job_config -.->|surfaces, unit| load_surfaces
    load_surfaces -->|dict[str, RegularSurface]| create_grid

Loading

Mark the mapped node, for example:

load_surfaces["load_surfaces<br/>map over: surfaces"]

When to_mermaid(job_configuration=...) is called, optionally include the item count:

load_surfaces × 3

Do not include aliases or paths in the graph, since they may be numerous or sensitive.

12. Concurrency follow-up

Concurrency is intentionally postponed. A later issue may add:

map:
  over: surfaces
  key_as: surface_name
  value_as: path
  max_concurrency: 4

The sequential implementation should isolate item execution behind _run_mapped_item() and _run_mapped_task() so it can later be replaced by a bounded executor without changing:

• Workflow syntax
• TaskMap
• Input construction
• Aggregate output shape
• Failure modes
• Per-item events
• Downstream dependencies

The follow-up design must address thread safety of shared services, cancellation, task timeout behavior outside the main thread, and deterministic event ordering.

13. Implementation plan

Phase 1: Mapping model and workflow metadata

Likely files:

• taskmaestro/mapping.py (new)
• taskmaestro/workflow.py
• taskmaestro/__init__.py

Work:

1. Add and export TaskMap.
2. Add mapped_over to WorkflowBuilder.add_task().
3. Store mapping declarations by task instance name.
4. Add mapping accessors and effective output annotation lookup.
5. Validate map field declarations and input completeness.
6. Update all edge type validation to use instance-aware output annotations.

Phase 2: Configuration validation

Likely files:

• taskmaestro/job.py
• taskmaestro/yaml_config.py

Work:

1. Validate map source presence and mapping shape.
2. Validate string keys and item key/value types.
3. Exclude map.over from inferred config_fields.
4. Ensure mapped tasks retrieve configuration even when they have no ordinary config_fields.
5. Support empty mappings.

Phase 3: Sequential execution

Likely files:

• taskmaestro/runner.py
• taskmaestro/exceptions.py

Work:

1. Refactor ordinary input assembly and task invocation into reusable methods.
2. Add sequential mapped execution.
3. Instantiate one task per item.
4. Apply timeout per item.
5. Aggregate outputs in declaration order.
6. Implement fail_fast and collect_all.
7. Add MappedTaskExecutionError.
8. Broaden aggregate output storage types where required.

Phase 4: Item observability

Likely files:

• taskmaestro/context.py
• taskmaestro/job.py
• taskmaestro/hooks/base.py
• taskmaestro/hooks/logging.py
• taskmaestro/hooks/timing.py

Work:

1. Add child contexts.
2. Add item-level hook events.
3. Add job.mapped_item_results.
4. Extend logging and timing hooks.
5. Preserve one parent task result per declared workflow node.

Phase 5: YAML and visualization

Likely files:

• taskmaestro/yaml_config.py
• taskmaestro/visualization.py

Work:

1. Add the YAML map schema.
2. Build the same TaskMap used by the Python API.
3. Render compact mapped nodes.
4. Show the effective aggregate output type.
5. Optionally show item count when job configuration is supplied.

Phase 6: Documentation and examples

Likely files:

• README.md
• examples/resinsight/

Work:

1. Replace repeated surface-loader declarations with one mapped task example.
2. Document Python and YAML syntax.
3. Document ordering, empty mappings, errors, lifecycle events, and timeout behavior.
4. Clearly state that mapped execution is sequential in this version.

14. Test plan

Workflow tests

• Valid mapped task with dependencies and config fields.
• Missing key_as or value_as field.
• Duplicate key_as and value_as.
• Map fields overlapping dependencies.
• Map fields overlapping config_fields.
• Missing required non-map input field.
• Effective mapped output type accepted by dict[str, O] consumer.
• Incompatible downstream dictionary type rejected.
• Mapped task participates correctly in cycle and sink detection.
• Same task class mapped under multiple instance names.

Job configuration tests

• Missing mapped task configuration.
• Missing over field.
• over value is not a mapping.
• Non-string key.
• Invalid key value for key_as.
• Invalid mapped value for value_as.
• Empty mapping.
• Common configuration fields alongside map source.
• YAML-inferred config_fields excludes map.over.

Runner tests

• One invocation per map entry.
• Fresh task instance per item.
• Declaration-order execution and output ordering.
• Shared dependency supplied to every item.
• Common config supplied to every item.
• Correct key and value injection.
• Output type validation per item.
• Empty mapping returns {}.
• fail_fast skips later items.
• collect_all runs all items and aggregates errors.
• Per-item timeout.
• Downstream task receives aggregate dictionary.
• Mapped task can be the job result.

Hook and context tests

• One logical task start and completion event.
• Per-item start and completion events.
• Per-item failure events.
• Stable event order under sequential execution.
• Item timing records.
• Child correlation IDs are deterministic and unique.
• Child scratch directories are unique.
• Services are available in child contexts.

YAML tests

• End-to-end mapped workflow.
• Named mapped task instance.
• All validation failures above.
• fail_fast and collect_all parsing.
• Empty mapping.
• Duplicate mapping aliases rejected.

Visualization tests

• Mapped node is marked without item expansion.
• Effective dict[str, O] output label.
• Configuration edge includes map source.
• Job-aware rendering shows item count.
• Large mappings do not enlarge the graph.

Regression and quality gates

pytest -v --cov=taskmaestro --cov-report=term-missing
ruff format --check taskmaestro tests examples
ruff check taskmaestro tests examples
mypy taskmaestro

Target coverage remains 100%.

Acceptance criteria

Issue #8 is complete when one configured mapping can drive sequential invocations of one declared task:

load_surfaces:
  surfaces:
    top: /data/top.irap
    middle: /data/middle.irap
    base: /data/base.irap

The task must run exactly three times in that order, receiving the corresponding alias and path each time. The logical node must produce:

{
    "top": top_surface,
    "middle": middle_surface,
    "base": base_surface,
}

The aggregate must be type-checked against downstream inputs, item failures must follow the selected error mode, item lifecycle information must be observable, YAML and Python APIs must behave consistently, and the Mermaid graph must remain compact.