Skip to content

Support dynamic task mapping on stub tasks - #70570

Open
jason810496 wants to merge 15 commits into
apache:mainfrom
jason810496:feature/lang-sdk/taskflow-stub-dag-mapped
Open

jason810496 wants to merge 15 commits into
apache:mainfrom
jason810496:feature/lang-sdk/taskflow-stub-dag-mapped

Conversation

@jason810496

@jason810496 jason810496 commented Jul 28, 2026 •

Copy link
Copy Markdown
Member

Support dynamic task mapping on stub tasks

A @task.stub can now be mapped with every Python mapping primitive. At ti_run the API server works out which upstream row and element each map index reads and sends it in arg_bindings, as ADR-0016 decides, so no Lang SDK re-implements the mapping arithmetic. The server sends coordinates, never XCom values.

def clean(row: str) -> str:
    return row.strip()


@task.stub(queue="typescript")
def transform(region: str, row: str, scale: int = 2) -> int: ...


@task.stub(queue="typescript")
def report(totals: list[int]) -> None: ...


rows = extract()
totals = transform.partial(region="eu").expand(row=rows.map(clean))
report(totals)

Argument lifecycle: how parse time and the API server split the work for the example above.

 PARSE TIME  (Dag processor, serialization/stub_arg_bindings.py)
   build_mapped_arg_bindings(transform)  one "param" per parameter, in declaration order
   build_arg_bindings(report)            "literal", "xcom", "sequence", or "param" for a group's expand argument
   rejects: map() not outermost, zip() inside zip(), a whole map()/zip()/concat() result,
            a literal dict expand value, a non-JSON zip() fillvalue, a strict expand_kwargs()
            key that partial() also sets, a mapped task's whole output as an expand_kwargs()
            item, and a producer mapped again inside a shared mapped task group
        |
        v
 SERIALIZED DAG  (metadata DB)
   transform
     "_arg_bindings": [
       {"name": "region", "kind": "param", "value_schema": {"type": "string"}},
       {"name": "row", "kind": "param", "value_schema": {"type": "string"}},
       {"name": "scale", "kind": "param", "value_schema": {"format": "int64", "type": "integer"},
        "default": 2}
     ]
     "partial_kwargs": {"queue": "typescript", "op_kwargs": {"__var": {"region": "eu"}, "__type": "dict"},
       "op_args": []}
     "op_kwargs_expand_input": {"type": "dict-of-lists", "value": {"__var": {
       "row": {"__var": {"type": "map",
         "arg": {"task_id": "extract", "key": "return_value"},
         "callables": ["def clean(row: str) -> str:\n    return row.strip()\n"]}, "__type": "xcomref"}},
       "__type": "dict"}}
   report
     "_arg_bindings": [
       {"name": "totals", "kind": "sequence", "task_id": "transform",
        "value_schema": {"items": {"format": "int64", "type": "integer"}, "type": "array"}}
     ]
        |
        |  ti_run reads the entries and the fields Python already writes
        v
 API SERVER  (execution_api/services/stub_arg_bindings.py)
   param             group_arg -> expand input element -> partial() value -> default
                     (expand_keys orders the expand() cross product)
   xcom / sequence   narrowed to one row when the producer shares the stub's mapped task group
   reads TaskMap lengths only, never an XCom value
        |
        v
 TIRunContext.arg_bindings  (Execution API 2027-02-28)
   transform, map index 2
     {"name": "region", "kind": "literal", "value_schema": {"type": "string"}, "value": "eu"}
     {"name": "row", "kind": "xcom", "value_schema": {"type": "string"}, "task_id": "extract",
      "element_index": 2, "apply_map": true}
     {"name": "scale", "kind": "literal", "value_schema": {"format": "int64", "type": "integer"},
      "value": 2, "from_default": true}
   report
     {"name": "totals", "kind": "sequence", "task_id": "transform",
      "value_schema": {"items": {"format": "int64", "type": "integer"}, "type": "array"}}

Write path: serialization/stub_arg_bindings.py builds _arg_bindings from the live operator at parse time.

 build_arg_bindings(op)              unmapped stub, including a member of a mapped task group
   bind the call to the stub signature, then apply defaults
   for each parameter, in declaration order:
     group's expand argument  -> _build_group_arg_binding     -> param {group_arg: {group_id, key}}
     XComArg                  -> _validate_whole_xcom_value   plain output only
                                 _validate_group_relative     no producer mapped again in a shared group
                                 _build_xcom_binding          -> xcom | sequence (mapped producer)
     anything else            -> _ensure_json_literal         -> literal {value, from_default}
   + value_schema from _infer_value_schema

 build_mapped_arg_bindings(op)       mapped stub, .expand() or .expand_kwargs()
   partial() values         -> _validate_whole_value         whole map/zip/concat rejected
   expand() values          -> _validate_expand_value -> _validate_expand_shape   literal dict rejected
   expand_kwargs() value    -> _validate_expand_kwargs       unknown key, strict overlap with partial()
                               _validate_kwarg_fallbacks     strict=False over an upstream output
   for each parameter       -> _build_param_binding          -> param {default, value_schema, expand_keys}
                               (the values stay in partial_kwargs and op_kwargs_expand_input)

 _validate_expand_shape walks the XComArg tree:
   map()     only as the outermost call
   zip()     never under another zip(), fillvalue must be JSON
   concat()  recurses into its arguments

Read path: execution_api/services/stub_arg_bindings.py derives this map index's bindings at ti_run.

 get_arg_bindings(task, ti)          services/task_instances.py
   not a stub, or no _arg_bindings          -> None
   mapped, client older than 2027-02-28     -> None
   mapped, or in a mapped task group        -> derive_mapped_arg_bindings
   unmapped and in no mapped task group     -> the stored entries as they are
   client older than 2027-02-28 and a binding has a new kind or field
                                            -> UndeliverableMappedArgBindingsError
   UndeliverableMappedArgBindingsError      -> 422 invalid_arg_bindings + warning log

 _bind_entry, per stored entry:
   param + group_arg    -> _bind_group_arg     the group's expand argument at this index
   param, mapped stub   -> _bind_mapped_param  1. expand():        _bind_element
                                               2. expand_kwargs(): _bind_mapping_entry
                                                  (skipped for a strict key partial() sets)
                                               3. partial():       _bind_whole
                                               4. default:         _bind_default
   xcom / sequence      -> _bind_upstream      one row when the producer shares the mapped group
   literal              -> as stored

 building blocks:
   _bind_element    literal list or dict  -> literal (a dict item is [key, value])
                    map(f)                -> _locate(inner) + apply_map
                    other XComArg         -> _locate
   _locate          zip()     -> tuple {parts: [XComRef | LiteralValue (fillvalue)]}
                    concat()  -> the argument owning this index
                    plain     -> xcom {map_index of the producer's row at that offset (mapped)
                                       | element_index (unmapped)}
   _bind_mapping_entry / _get_fallback
                    literal dict          -> the key's value now
                    upstream dict         -> xcom + kwarg + fallback (partial() value or default)
   _get_sub_indexes / _get_length
                    expand() cross product in expand_keys order, and lengths from TaskMap and
                    XCom row counts, never an XCom value
  • Every primitive resolves on the server: the expand() cross product, zip() (a new tuple binding, one part per zipped argument), concat(), expand_kwargs(), and a stub inside a mapped task group.
  • New per-index fields: map_index and element_index locate one element of an upstream output, apply_map marks map() callables still to run on it, and kwarg and fallback serve expand_kwargs() over an upstream output, whose keys are unknown until it runs.
  • A mapped task's whole output is a sequence binding, so an unmapped stub or a partial() argument can read a mapped upstream; parse time used to reject this.
  • Parse time records only what the serialized Dag lacks: declaration order, defaults, value schemas, call arguments of a mapped task group member, and expand_keys, because a JSONB store reorders dict keys and the order fixes the cross product.

Follow-ups, per ADR-0016:

  • Make Lang SDKs refuse mapped stub arg bindings they cannot bind #76967: the Go, TypeScript and Java runtimes raise a not-implemented error for the per-index fields until Lang SDK 1.1.0, and Java moves to supervisor schema 2027-02-28.
  • The supervisor reads xcom and tuple bindings through the XCom backend and hands the runtime a literal.
  • ApplyMap runs a Python Dag's map() callables for a stub TaskHandler.
  • needs_mapped_length lets the supervisor measure a Lang-SDK producer's map length.

Until the supervisor change lands, no runtime reads the new fields, so the stack has to land before the 2027-02-28 versions are released.

Compatibility: the change ships in Execution API version 2027-02-28 and in the supervisor schema version of the same date.

  • An older client gets no arg_bindings for a mapped stub, the old ignored-arguments behaviour. For an unmapped stub whose bindings need the new kinds or fields, such as a sequence or an argument of a mapped task group, ti_run answers 422.
  • The supervisor fails a TI whose bindings need the new kinds or fields when its runtime is older than 2027-02-28, and asks to upgrade the SDK.
  • The Python stub TaskHandler raises NotImplementedError for a binding kind or per-index field that mapped tasks need, instead of reading the wrong value. The Go, TypeScript and Java runtimes do the same in Make Lang SDKs refuse mapped stub arg bindings they cannot bind #76967, pointing at Lang SDK 1.1.0.
  • A mapped stub in a Dag serialized before this change keeps the old ignored-arguments behaviour.

Was generative AI tooling used to co-author this PR?

@github-actions

Copy link
Copy Markdown
Contributor

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.

@github-actions github-actions Bot added the stale Stale PRs per the .github/workflows/stale.yml policy file label Oct 28, 2026
@jason810496 jason810496 removed the area:API Airflow's REST/HTTP API label Nov 2, 2026
@github-actions github-actions Bot removed the stale Stale PRs per the .github/workflows/stale.yml policy file label Nov 3, 2026
@jason810496
jason810496 force-pushed the feature/lang-sdk/taskflow-stub-dag-mapped branch from b786830 to 7518b09 Compare November 3, 2026 07:04
jason810496 and others added 6 commits November 12, 2026 03:12
Restore the mapped (.expand()) stub arg-binding support split out of the
unmapped PR: the Dag serializer captures per-parameter metadata
(declaration order, defaults, value schemas) from the stub signature via
the get_mapped_serialized_fields hook, ti_run derives per-map-index
bindings from it with the map-index decomposition on
SchedulerDictOfListsExpandInput, and XComArgBinding regains the
map_index/element_index delivery fields across the task-sdk and ts-sdk
generated models.
ADR-0016 keeps the mapping arithmetic in core, so a Lang SDK never
re-implements it: the server picks each map index's element and sends it
as an arg binding. The change ships in Execution API 2027-02-28 because
2026-10-30 is already released.
ti_run loads the task instance as a column row, which has no task attribute,
so every mapped stub that needed producer contexts failed with a 500.
An expand_kwargs() mapping entry must override a partial() value, as it does in Python.
Exposing the mapped metadata as a property keeps callers off a private attribute,
and the mapped prefix tells the per-map-index path apart from the parse-time one.

Co-Authored-By: Claude <[email protected]>
Two fields describing the same per-parameter bindings made every reader branch on which one
a stub had, and the dedicated field needed {__type, __var} encoding that nothing in it requires.
A param kind keeps one plain-JSON list in declaration order, matching ADR-0016.

A typing.TypedDict annotation on Python < 3.12 used to ship with no value schema silently,
so it now fails parsing with the fix in the message.

Co-Authored-By: Claude <[email protected]>
The 2026-10-30 supervisor schema still has the arg_bindings field, so a plain model_dump puts it back as its None default after the downgrade removes it.

Co-Authored-By: Claude <[email protected]>
The old name described where the argument gets resolved, not what it is, and read as if it covered only map(). The kind also covers zip() nesting that no single binding can hold.

Co-Authored-By: Claude <[email protected]>
jason810496 and others added 7 commits November 12, 2026 03:14
The supervisor now narrows each binding to its value and the runtime applies map() callables, so the server no longer needs an escape hatch for arguments it cannot finish.
Shapes no single binding can carry fail at parse time instead, where the Dag author sees them, and an expand_kwargs() key now wins over its partial() value as it does in Python.

Co-Authored-By: Claude <[email protected]>
A review of the mapped stub-task bindings found cases where the server and the Python task runner disagree, where a Dag fails only at run time, or where an older client silently runs a task with no arguments.
Stored Dags may re-order expand() keys, strict expand_kwargs() and zip().map() were not honoured, and deterministic binding errors were retried as server errors.

Co-Authored-By: Claude <[email protected]>
A second review found that a mapped upstream with missing rows could bind the wrong element, that old runtimes could still run a task with no arguments, and that some server faults were reported as Dag errors the client must not retry.
The Python task handler binder also accepted binding shapes it cannot honour.

Co-Authored-By: Claude <[email protected]>
Reviewers had to trace where match-pattern captures such as producer came
from, and Any on the validator values hid which expand shapes they accept.

Co-Authored-By: Claude <[email protected]>
The round-trip test only checked _arg_bindings, so losing partial_kwargs,
the expand input or the strict flag in serialization would go unnoticed.

Co-Authored-By: Claude <[email protected]>
A stale group id in a stored binding surfaced as a bare 500 KeyError
instead of the 422 that other undeliverable bindings return.

Co-Authored-By: Claude <[email protected]>
Upstream added a field to the same 2027-02-28 schema version, so the
generated Go and TypeScript models have to include both changes.

Co-Authored-By: Claude <[email protected]>
@jason810496
jason810496 force-pushed the feature/lang-sdk/taskflow-stub-dag-mapped branch from 3f8a2a5 to 89fa2d7 Compare November 12, 2026 03:21
jason810496 and others added 2 commits November 12, 2026 03:55
Keep this change focused on the API server; the Go, TypeScript and Java
runtimes learn to refuse per-index bindings in a stacked PR.

Co-Authored-By: Claude <[email protected]>
Mapped task support for stub tasks is planned, so a binding the handler
cannot read is a missing feature rather than an invalid value.

Co-Authored-By: Claude <[email protected]>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant