Skip to content

Support dynamic task mapping on stub tasks - #70570

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

Support dynamic task mapping on stub tasks#70570
jason810496 wants to merge 41 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

Why

#69757 ships stub TaskFlow arg-binding but rejects .expand() on a stub at parse time. This PR restores that scope: a @task.stub can be dynamically task-mapped, and every map index gets its own arg-binding spec so the foreign runtime receives its element instead of the aggregated output.

Supported dynamic-mapping forms

A mapped @task.stub produces one arg-binding per map index, always in the stub's signature declaration order (not call-site order). Each parameter resolves as one of:

Expanded arguments.expand(param=…), value differs per index:

Form Example Wire binding
Expand over a literal collection transform.expand(country=["uk", "fr", "de"]) literal, element resolved server-side (list → element; dict → [key, value] per item)
Expand over an unmapped upstream's list output transform.expand(extracted=extract()) xcom + element_index=i (pull the single row, take element i)
Expand over a mapped upstream transform.expand(extracted=seed.expand(n=[1, 2])) xcom + map_index=i (pull upstream row i directly)

Multiple expanded arguments — cross product:

Form Example Behavior
.expand(a=…, b=…) combine.expand(a=["x", "y"], b=[1, 2, 3]) → 6 instances the map index is decomposed into one sub-index per kwarg (last varies fastest); each kwarg resolves independently, so the three expand forms above can be mixed in one call

Partial arguments.partial(param=…), constant across every index:

Form Example Wire binding
Partial literal transform.partial(country="uk").expand(…) literal (same value every index)
Partial unmapped-upstream XCom transform.partial(extracted=extract()).expand(…) xcom, whole return value (no sub-index)

Defaulted arguments:

Form Example Wire binding
Unpassed param with a signature default retries: int = 3 left unpassed literal + from_default: true

One DAG exercising every form at once:

@task.stub(queue="golang")
def make_items(): ...                     # unmapped stub → returns a list

@task.stub(queue="golang")
def seed(n: int): ...                     # mapped below → a mapped upstream

@task.stub(queue="golang")
def transform(
    country: str,                         # expand over a literal list
    extracted: dict,                      # expand over an unmapped upstream's list → element_index
    seeded: dict,                         # expand over a mapped upstream           → map_index
    region: str,                          # partial literal (constant)
    config: dict,                         # partial xcom over an unmapped upstream (constant)
    retries: int = 3,                     # unpassed → default captured (from_default)
): ...

@dag(dag_id="mapped_binding_dag")
def mapped_binding_dag():
    transform.partial(
        region="uk",
        config=load_config(),
    ).expand(                             # cross product; expand sources mixed across kwargs
        country=["uk", "fr", "de"],
        extracted=make_items(),
        seeded=seed.expand(n=[1, 2, 3]),
    )

Rejected loudly (parse-time in the provider; re-checked server-side for Dags from other provider versions): .expand_kwargs(); a partial() kwarg over a mapped upstream's aggregated output (would bind the nonexistent map_index=-1 row); .map() / .zip() / concat or custom-key XCom; non-JSON literals; a mapped stub TI still at map_index=-1. A mapped stub in an older-provider Dag (no captured metadata) delivers no bindings and keeps the legacy ignored-args behavior.

How

  • Parse time (providers/standard): a mapped stub never instantiates at parse time, so _StubOperator captures per-parameter metadata (declaration order, defaults, value schemas) via a new optional get_mapped_serialized_fields operator hook. The core serializer calls it at the single point where operator_class/python_callable are still the real objects — everything the server cannot recover from the serialized Dag alone.
  • Wire model: XComArgBinding regains map_index (which upstream row to pull — expand over a mapped upstream) and element_index (take element N of the unmapped list — expand over an unmapped upstream's output).
  • Server (ti_run): for a mapped stub, bindings are derived per map index by decomposing the TI's map_index into one sub-index per expanded kwarg via new SchedulerDictOfListsExpandInput.resolve_expansion_sub_indexes — the server-side twin of the SDK's _expand_mapped_field cross-product (last kwarg varies fastest). Expanded kwargs get map_index/element_index, partial() kwargs and unpassed defaults bind as above.
  • Backward compat: Dags serialized by an older provider carry no metadata → resolve to None → keep the legacy ignored-args behavior (their args were never deliverable). The provider's parse-time rejections are re-checked server-side for Dags produced by other provider versions.
  • Schema: new arg_binding_param definition + optional _mapped_arg_binding_params array on the operator; the inner object stays open so newer metadata keeps validating on older cores. No SERIALIZER_VERSION bump (optional field).

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

Stub Dags could only declare argless tasks, so cross-language dataflow
required hand-written GetXCom calls inside each Go task. Capturing the
TaskFlow call's argument spec at parse time and delivering it through
the Execution API and StartupDetails lets a Go task receive upstream
outputs and Dag-file literals as plain typed parameters, with loud
arity/type errors instead of silently zero-filled values.
"stub_args" leaked the _StubOperator implementation detail into the
wire contract that foreign-language SDKs code-generate against;
"arg bindings" names what the data actually is -- the ordered spec a
runtime binds onto the task function. Renaming now, before the field
ships in a released execution API or supervisor schema version, keeps
the contract clean without any compatibility shims.
The single try/except made Airflow 3.0 (whose SDK predates
KNOWN_CONTEXT_KEYS) fall back to the Airflow 2 import paths and fail;
the arg-capture tests imported airflow.sdk directly, which does not
exist on 2.11; and the .expand() rejection relies on the
supports_expand opt-out that only ships with Airflow 3.4. Also reword
the context-key rejection to stop implying foreign runtimes have no
task context -- the lang SDKs inject their own natively; stub
signatures just must not declare Airflow context parameters.
Only the Multi-Lang stub-task path needs the serialized-dag machinery
and the arg-binding models, so regular task-run requests should not pay
for them: the TaskArgBinding datamodels move to a dedicated module and
the serialized-dag imports become local to the stub lookup. The OpenAPI
schema is unchanged (component names stay the same), which is why no
execution API version bump accompanies this commit.
simple_dag only exercises the minimal binding: one literal and one XCom
argument. The new taskflow_binding_dag locks in the rest of the surface
end to end -- scalar and array literals, keyword arguments, a defaulted
None, and XCom fan-in from two upstream Go tasks bound onto a strict
struct and a typed slice -- with the Go task verifying every bound value
so binding regressions fail the example run loudly.
Naming every stub argument as a separate flat Go parameter gets unwieldy
as the argument count grows, and there was no way for a Go task to pull
an XCom that the Python TaskFlow call itself never passed. A struct that
embeds sdk.TaskInput lets a task bind many arguments by name (or an
explicit ad hoc XCom pull) onto one parameter instead, while the
existing flat/positional binding keeps working unchanged for functions
that don't opt in.

This required adding a name to the wire-level TaskArgBinding spec so a
struct field can look itself up by the Dag's TaskFlow argument name
regardless of declaration order on either side, since Go cannot recover
a plain function parameter's name via reflection the way it can for a
struct's fields.
These tests built stub tasks with the raw stub(fn)(...) call instead of
the @task.stub decorator every real Dag (Go/TS/Java examples) already
uses, so a reader comparing the tests to real usage saw a syntax the
feature doesn't actually ship.
As a plain Literal type alias, the field's generated model came out
under a generic, field-derived name (DataType) in both the task-sdk
client model and the Go SDK's generated types, rather than the
ArgBindingDataType name declared in the source. A real Enum class
carries its own name through codegen, so providers/standard can import
it directly instead of re-deriving the same string vocabulary by hand,
with a hand-written fallback for Airflow 2 where the execution-API
generated models aren't importable.
The combined TaskInput example mixed all three field-binding modes
(arg: tag, no tag, xcom: tag) into one struct, so no single task
demonstrated any one mode in isolation. Split it into via_struct_no_tags,
via_struct_arg_tag, and via_struct_xcom_tag, and renamed combine to
via_flat_args to make the positional/keyword-style split between flat
and struct binding legible at the call-site naming level.

A TaskInput struct field whose name has no matching TaskFlow call
argument now stays at its Go zero value instead of failing the task --
keyword-argument semantics (an unpassed name falls back to its default)
rather than the strict arity check flat, positional parameters get.
via_struct_unmatched_arg exercises this directly.
The flat kind-discriminated shape kept foreign-language codegen simple
but left each variant's contract implicit: task_id was nullable even
though every xcom binding has one, and value/key were dead weight on the
opposite kind. Modelling arg_bindings as a kind-discriminated union
makes the contracts explicit on every wire (OpenAPI, supervisor schema,
Go, TypeScript) - xcom bindings now require task_id - and lets the Go
runtime mirror the split as a sealed sum type instead of branching on a
string field, so malformed specs fail the task before its body runs.
An ad hoc `xcom:"<task-id>"` pull baked the upstream task id into the
compiled Go binary, hiding a data dependency from the Dag file that owns
task wiring on the Python side (the example even needed a manual >> to
order the pull's upstream). Fields now bind exclusively by argument
name -- an `arg:"<name>"` tag, or the snake_cased field name when the
tag is omitted -- so every value a task consumes stays visible in its
TaskFlow call, and a task that needs an extra XCom can still ask for it
explicitly through the injected client.
The snake_cased fallback silently rewrote Go field names into wire
argument names, hiding the cross-language mapping from the reader; and
because an unmatched TaskInput field kwarg-style falls back to its zero
value, a wrong guess about the conversion never failed loudly. Matching
the field name verbatim removes that magic: every snake_case Python
parameter a field binds is now spelled out as an explicit `arg:` tag in
the Go source.

The e2e module also still referenced the via_struct_xcom_tag task
removed with the xcom struct tag, which would have failed the suite
against the current Dag.
Most workloads are not stub operators, so constructing the
discriminated-union adapter at module import made every execution API
process pay for it up front. Moving it next to the TaskArgBinding
models behind a cached getter defers the cost to the first stub-task
run and leaves the _STUB_TASK_TYPE gate as the only stub-specific
module-level state in the route.
The XCom key was always return_value for a TaskFlow call, so the key
field carried no information; it is removed end to end (datamodel,
serialized spec, supervisor schema, generated task-sdk and Go models)
and indexing a stub argument by a custom key now fails at parse time
instead of being silently representable.

Mixing flat positional data parameters with a TaskInput struct in one
Go task signature was too ambiguous to reason about, so Analyze now
rejects it: a function declares one binding shape or the other.

The Go binding sum type and its DataType vocabulary are now defined in
terms of the generated supervisor-schema models rather than hand-written
mirrors, so they cannot drift from the wire contract, and every
via_struct_* example task now binds an XCom-sourced argument
(make_region) alongside a literal so struct-field binding is exercised
with both sources end to end.

The TypeScript supervisor model bump is left out of this PR on purpose.
The 2026-06-30 execution API version already shipped in Airflow 3.3.0, so
appending the arg_bindings migration to it would mutate a released version,
which the execution API versioning policy forbids; the change now opens
version 2026-07-30, matching the supervisor-schema date.

The rest addresses a local multi-reviewer audit of the branch:

- The airflow-go-pack integration test's expected manifest was missing the
  make_region task added to the example bundle, failing go test.
- The per-field cadwyn didnt_exist instructions on XComArgBinding and
  LiteralArgBinding name fields could never apply (arg_bindings is stripped
  wholesale on downgrade) and are dropped on both the execution API and the
  supervisor schema side; the supervisor-schema change class is renamed so
  the two same-named migrations cannot be confused.
- The stub decorator's hand-rolled version-split imports now go through the
  common.compat sdk seam (new PlainXComArg and KNOWN_CONTEXT_KEYS exports).
- The Go runtime validates required wire-spec fields (name, xcom task_id)
  instead of silently binding empty strings, populates the carried Kind
  discriminant, and reports a binding bookkeeping bug as a task error
  instead of panicking the worker.
- The supports_expand opt-out is now covered by task-sdk-level tests, the
  ti_run serialized-dag scan moved onto LazyDeserializedDAG next to its
  sibling accessors, and assorted review nits (exception types, stale
  wording, enum comparisons) are fixed.
A multi-angle review of the branch surfaced gaps at the edges of the
new binding contract:

- An unrecognized serialized spec escaped ti_run as an opaque 500 on
  provider/core version skew; it now returns a structured
  invalid_arg_bindings error, per the route-boundary convention.
- The new parse-time signature checks broke previously importable
  argless stub Dags (e.g. a **kwargs or ti parameter); they now fire
  only when a TaskFlow call actually passes arguments.
- Stubs called with arguments inside a mapped task group serialized a
  spec with no map-index dimension and failed (or mis-bound) at
  runtime; they are now rejected at parse time.
- A Go TaskInput struct had to mirror every stub parameter -- including
  defaulted ones the author never passed -- or fail every run, while an
  empty spec silently zero-filled the whole struct, contradicting the
  documented fail-loud behavior. Literal entries captured from
  signature defaults now carry from_default on the wire (inside the
  still-in-progress 2026-07-30 schema) and may go unclaimed; a spec
  that never arrives fails loudly when the struct declares bindable
  fields.
- NaN/Infinity literals passed the parse-time JSON check only to fail
  far away (or silently bind 0.0); json.dumps now rejects them.
- A malformed spec bypassed ShouldRetry while an equally permanent
  arity error retried; both now share retry semantics.
- ti_run no longer issues two queries and re-parses the serialized-Dag
  blob on every stub-task start: single joinedload query plus a
  per-(dag_version, task) cache of the immutable extracted spec, and
  XCom pulls in the Go binding path now run concurrently.
Reviewing the Python-side arg-binding contract and the Go runtime that
consumes it in one PR ties the core/task-sdk review to Go SDK internals.
Scoping this PR to the contract lets it merge on its own; the Go SDK
consumption (pkg/binding, task-runner dispatch, example bundle, e2e
test) lands stacked on top from feature/go-sdk/taskflow-arg-binding.
The API server already holds a DBDagBag with configurable LRU+TTL
caching of deserialized Dags; a route-local raw-blob accessor plus a
hand-rolled module-level cache duplicated that machinery with its own
eviction story. The serialized _arg_bindings field survives full
deserialization onto the task object, so ti_run can read it off the
dag_bag-resolved Dag version directly, and the
LazyDeserializedDAG.get_task_arg_bindings accessor goes away.
The from_default flag records provenance, not value equality: the Go
TaskInput struct mode fails unclaimed explicit arguments but tolerates
unclaimed defaults, so an author-passed value that happens to equal the
signature default must not be flagged. Pin that boundary at the capture
site.
The supervisor schema bump ships in this PR, and the version-pin guards
(TestSupervisorSchemaVersionMatchesSnapshot, the ts-sdk supervisor
schema hook) rightly insist the pinned constant and the generated
TypeScript models follow the schema in the same change. The Go runtime
consumption of the new arg_bindings field stays in the stacked
feature/go-sdk/taskflow-arg-binding branch.
Review feedback on the arg-binding contract asked to reuse JSON Schema
rather than invent an ArgBindingDataType vocabulary, so foreign runtimes
can validate bound values with plain JSON-schema semantics. Each binding
now ships an optional value_schema fragment carrying the standard type
and format keywords: unions and Optionals map to type lists instead of
degrading to "any", and int/float/datetime/date/time/timedelta
annotations gain int64/double/date-time/date/time/duration formats the
bare type name cannot convey. An unconstrained argument omits the field
entirely, and unknown keywords from newer providers are tolerated.

The kind discriminator stays required: giving it a server-side default
drops it from the OpenAPI required list and datamodel-codegen then emits
Literal | None, which pydantic rejects as a tagged-union discriminator.

Also retargets the execution API and supervisor schema version to
2026-10-30 to match the Airflow 3.4 release train, documents the
optional _arg_bindings task property in the serialized-Dag schema (no
serializer version bump: optional field, unchanged logic), adds the Dag
version id to the ti_run spec-validation failure log, and rewords the
stub .expand() rejection rationale: arg types are uniform across map
indexes, the blocker is per-index value resolution at runtime.
…d mapper

The hand-written annotation-to-fragment mapper duplicated what pydantic
already exposes publicly: TypeAdapter(annotation).json_schema(), with a
GenerateJsonSchema subclass layering the int64/double numeric formats a
foreign runtime needs. Delegating to pydantic removes the bespoke union
logic and buys richer, standard fragments for free -- anyOf for unions
(per-member formats survive mixed unions now), items/additionalProperties
for parameterized containers, and enum for Literal annotations -- while
anything pydantic cannot schema (arbitrary classes anywhere in the
annotation) still degrades to a decode-only binding.

value_schema becomes a free-form JSON object on the wire instead of a
typed model: a typed model silently strips every keyword it does not
know when the spec is re-serialized along the server-to-supervisor
delivery path, which would corrupt exactly the open-vocabulary fragments
this contract promises to carry verbatim.

The provider declares its now-direct pydantic dependency (edge3/http
precedent). Trade-offs accepted: fragment shapes follow the pydantic
version active at parse time, and pendulum.DateTime annotations degrade
to decode-only since pydantic has no schema for arbitrary datetime
subclasses. No new execution-API version: the field's shape changes
inside the still-unreleased 2026-10-30 version this PR introduces.
pendulum.DateTime is the most common temporal annotation in Airflow Dag
code, but pydantic cannot generate schemas for datetime subclasses, so
stub arg specs silently degraded to decode-only checks. Temporal
subclasses now normalize to their stdlib bases (recursively through
unions and containers) before schema generation.

Also drop the standard provider's explicit pydantic dependency
(apache-airflow already provides it), trim the verbose comments
introduced with the arg-binding spec, and shorten the
invalid-arg-bindings error message.
…annotation

Airflow 2.x base installs do not ship pydantic (it is an optional extra
there), so after dropping the standard provider's explicit pydantic
dependency, importing the stub decorator would crash Dag parsing on
such installs. Value schemas now degrade to the decode-only fallback
the wire contract already supports when pydantic is absent.

Annotations pydantic can validate but not schema-ify (e.g. Callable)
raise PydanticInvalidForJsonSchema, which escaped the existing handler
and crashed Dag parsing; it is now caught alongside
PydanticSchemaGenerationError.

Temporal normalization now runs only as a retry after direct schema
generation fails, so temporal subclasses that carry their own pydantic
schema keep it.
The cross-arch pack integration test asserted a hardcoded schema date,
so bumping the Go SDK's supervisor schema version (2026-06-16 ->
2026-10-30 on this branch) broke it in CI while the unit tests, which
inject the version explicitly, kept passing. The packed binary and the
test compile from the same module tree, so referencing
execution.SupervisorSchemaVersion asserts the same value the binary
embeds and future bumps cannot drift.
Banning .expand() on @task.stub blocked a core dynamic-mapping pattern
for foreign-runtime Dags with no workaround. A mapped stub never
instantiates at parse time, so instead of a parse-time capture, ti_run
now derives the per-map-index arg spec from the serialized expand
input, mirroring the task-sdk's DictOfListsExpandInput index
decomposition: literal expands resolve to their element server-side,
expands over a mapped upstream bind that upstream's XCom row via the
new map_index field, and expands over an unmapped upstream's output
carry the new element_index field so the runtime picks the right
element of the pulled list. Value schemas stay parse-time-only and are
omitted for mapped stubs, falling back to the decode-only contract.

The derivation grew into enough business logic that it lives in a new
execution_api services package (mirroring core_api's services layout)
rather than the routes module, and the generic map-index decomposition
sits on SchedulerDictOfListsExpandInput beside its map-length helpers,
mirroring where the task-sdk twin keeps the same arithmetic.

The supports_expand opt-out this branch added to the task-sdk decorator
machinery existed only for the stub ban, so it is reverted. Stubs with
arguments inside a mapped task group stay rejected: those instances
have no expand input of their own to derive bindings from.
On Python 3.10, isinstance(list[X], type) is True and issubclass on the
alias silently consults the origin, so the plain-class branch swallowed
parametrized generics before the origin/args reconstruction could
rewrite their arguments; list[pendulum.DateTime] then degraded to no
value schema at all. Python 3.11+ returns False there, which is why the
regression only surfaced on the 3.10 CI jobs. Detecting parametrized
generics first restores the normalization on every supported version.
The map_index and element_index fields added to XComArgBinding in the
supervisor schema must land together with the generated TypeScript
output, which the check-ts-sdk-supervisor-schema static check enforces
by regenerating and diffing the file.
Main bumped datamodel-code-generator 0.33.0 -> 0.41.0 (apache#69854), which
changes how the generated task-sdk client renders this branch's
ArgValueSchema (RootModel drops the hoisted null), and the merged lock's
pydantic reorders the supervisor schema snapshot $defs. Regenerated the
task-sdk datamodels, supervisor schema snapshot, and ts-sdk supervisor
types as part of the merge so the committed output matches what CI
regenerates on the merged tree.
The mapped-stub decorator test asserted MappedOperator.is_mapped, which
does not exist on Airflow 2.x, so the provider compat suite failed on
2.11 with an AttributeError. Asserting on op_kwargs_expand_input and
partial_kwargs keeps the test meaningful on every supported Airflow
version and additionally pins down what "no parse-time bindings" means.
The server-side derivation for mapped stub tasks typed the task as Any
and left its non-obvious decisions undocumented: the expand_kwargs()
gate, when NotFullyPopulated can actually fire, the partial()/expand()
kwarg partition, and how a runtime consumes element_index. Narrowing to
SerializedMappedOperator via the is_mapped() guard lets mypy check the
mapped-only attribute access, and the comments capture the reasoning
where it applies. Delivering value_schema on mapped bindings stays
deferred, now tracked at
apache#70523.

The serialization round-trip test asserted an either-or for the argless
stub; deserialization never sets _arg_bindings for it, so assert exactly
that.
A stub argument fed a mapped task's combined output silently bound the
unmapped XCom row (map_index=-1), which never exists for a mapped
upstream, so the foreign runtime received nothing where Python TaskFlow
delivers the aggregated list. The wire contract cannot express "pull all
rows", so fail loudly like the other inexpressible constructs: at parse
time for the unmapped call path, and server-side for already-serialized
Dags.

XComArgs inside partial() op_kwargs also deserialize to _XComRef and
were never dereferenced, falling past the XComArg branches entirely; a
partial() kwarg over an unmapped upstream failed spec validation
instead of binding that XCom row. Dereferencing them fixes that and
lets the mapped-upstream rejection see the real reference.
Clearing only the upstream of a queued mapped stub and re-running it to
an empty list records a TaskMap length of 0 while the expanded TI still
exists; decomposing its map index then divided by zero and surfaced as
an opaque catch-all 500. The task-sdk twin of this arithmetic guards
mapped lengths below 1, so mirror it and route the failure through the
structured invalid_arg_bindings error like every other undeliverable
binding.
ti_run derived arg bindings for every stub task regardless of the
client's negotiated API version, so a stub Dag using a construct the
derivation rejects (e.g. expand_kwargs) went from running with its args
ignored to hard-failing with a 500 after a server upgrade -- even for
clients whose responses have arg_bindings stripped anyway. The cadwyn
migration only pops the field from successful responses; it cannot gate
the computation, so consult the negotiated version before deriving.
A mapped stub never instantiates at parse time, so ti_run derived its
arg bindings blind to the stub signature: the spec came out in call-site
dict order while the wire contract promises declaration order (a
positional binder like the Go SDK's flat mode then receives swapped
values), and parameters filled from signature defaults were silently
dropped where the unmapped path ships from_default entries. Declaration
order, defaults, and value schemas can only come from the real function,
which exists nowhere but the Dag processor, so expose a classmethod the
core serializer can call while serializing the mapped operator (wired up
in a follow-up commit). Building the metadata also validates the mapping
at parse time, so expand_kwargs() on a parameterful stub, partial()
kwargs over a mapped upstream, and mappings that do not bind to the
signature fail as Dag import errors instead of per-TI 500s at run time.
The server-side derivation for mapped stubs was blind to the stub
signature: it emitted the spec in call-site dict order while the wire
contract promises declaration order (a positional binder like the Go
SDK's flat mode then receives silently swapped values), dropped
parameters filled from signature defaults where the unmapped path ships
from_default entries, and could not attach value schemas.

The Dag serializer now consults an optional operator-class hook while
serializing a mapped operator -- the one point where operator_class and
python_callable are still the real objects -- and stores the stub's
per-parameter metadata under _mapped_arg_binding_params. ti_run walks
that metadata in declaration order, fills expanded, partial, and
defaulted parameters alike, and carries each parameter's value schema,
closing the mapped/unmapped contract gap
(apache#70523). Dags serialized
without the metadata (an older provider) keep the legacy ignored-args
behavior instead of receiving order-uncertain bindings, and the old
derivation's rejections stay as backstops for such Dags.
TypeAdapter construction is one of pydantic's most expensive operations
and ran fresh for every annotated stub parameter on every Dag file
re-parse, which the Dag processor repeats continuously. Annotations are
static, so cache the generated fragment per annotation for the process
lifetime, deep-copying on the way out so embedded specs never alias the
cache, and falling back to uncached generation for unhashable
annotations.
Both binding validators duplicated the canonical XCom return-value key
as a string literal, evading constant-based refactors and diverging from
the neighboring code (serialization's xcom_arg already compares against
XCOM_RETURN_KEY). The constant is importable in both contexts:
common.compat.sdk re-exports it for the provider, airflow.models.xcom
for core.
Reviewing unmapped TaskFlow delivery and per-map-index derivation
together made the PR hard to land, so this PR narrows to the unmapped
contract: mapped (.expand()) stubs keep the released ignored-args
behavior (they capture no parse-time spec, so ti_run naturally delivers
no bindings), documented on the stub decorator. The mapped derivation --
the serializer capture hook, per-parameter metadata, map-index
decomposition, and their tests -- moves wholesale to the stacked
follow-up branch feature/lang-sdk/taskflow-stub-dag-mapped.
XComArgBinding carried map_index and element_index for the per-map-index
delivery that now lands in the stacked follow-up branch; the unmapped
path never sets either, so this PR ships the contract without them. The
task-sdk client models, supervisor schema snapshot, and ts-sdk types are
regenerated accordingly; the follow-up re-adds the fields with its
derivation.
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.
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