Bind TaskFlow stub-task call arguments in the Go SDK runtime - #70209
Draft
jason810496 wants to merge 52 commits into
Draft
Bind TaskFlow stub-task call arguments in the Go SDK runtime#70209jason810496 wants to merge 52 commits into
jason810496 wants to merge 52 commits into
Conversation
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.
boring-cyborg
Bot
added
area:API
1 task
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.
jason810496
force-pushed
the
feature/go-sdk/taskflow-arg-binding
branch
from
July 22, 2026 03:55
81c8af8 to
68363f9
Compare
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.
jason810496
force-pushed
the
feature/go-sdk/taskflow-arg-binding
branch
from
July 24, 2026 11:16
983d0f5 to
83a2210
Compare
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.
Split out of the arg-binding contract PR (feature/lang-sdk/ taskflow-stub-dag) so the Python-side wire model can merge on its own review track; this stacked branch carries the Go SDK consumption of TIRunContext.arg_bindings and the end-to-end coverage.
Task authors previously had to embed the zero-size sdk.TaskInput marker to opt a struct into per-field, name-based TaskFlow argument binding. The marker was redundant: a function whose sole data parameter is a struct is unambiguous, and the per-execution argument spec (argument names, arity, and from_default) is enough to choose between binding fields by name and decoding one argument whole. Removing it drops boilerplate every keyword-style task had to carry and that had no analogue in the Python or Java SDKs.
Exercise the same one-struct task signature resolving field-by-field and whole-value across two executions, pinning the per-execution choice that replaced the sdk.TaskInput marker.
… SDK example The taskflow_binding_dag example covered scalar/array/object arguments but not the two ways a single dict binds once the sdk.TaskInput marker was dropped. Add via_flat_map (one dict decoded whole into a struct) and via_struct_map (one dict bound by name onto a struct's map field) so both paths are demonstrated and verified end to end.
The stub arg-binding contract this branch builds on replaced the coarse data_type enum with a per-argument JSON-schema fragment (value_schema). Rebasing onto it makes the Go runtime consume that fragment: the generated models drop ArgBindingDataType for ArgValueSchema, and the type check reads the schema's "type" keyword, skipping the check (decode-only) when the fragment is absent or carries no plain-string type -- matching the Python side, which omits value_schema for unconstrained annotations.
A .expand() stub delivers per-map-index XCom bindings that select a specific upstream row (map_index) or an element of the pulled list (element_index), but the Go runtime dropped both fields and always pulled the whole unmapped row -- so an expanded task received the entire upstream output instead of its element, or failed to find a mapped upstream's row. Read both fields off the wire and honor them when pulling: forward map_index to GetXCom and take element_index out of the pulled sequence, so dynamic task mapping over TaskFlow arguments works.
…end to end The example bundle covered scalar/struct/map argument shapes but not a mapped (.expand()) stub, which is where the runtime's map_index/element_index handling actually matters. Add make_items (a list producer) and a mapped via_expand stub fanned out over it, so each mapped instance binds its own element; the e2e asserts every instance received the element at its map_index.
A Go stub task whose list output a downstream .expand() maps over never recorded a TaskMap, because a foreign runtime cannot inspect the Dag to know its return value feeds a mapping and so never reported a mapped_length. Without it the scheduler could not determine the expansion size and left the mapped task upstream_failed. The execution API now flags such a stub task via a new TIRunContext.has_mapped_dependants field (derived from the serialized Dag, mirroring the Python task runner's iter_mapped_dependants check and its not-is_mapped guard). When set, the supervisor records the mapped_length of the runtime's mappable return value on its behalf -- the foreign-runtime analogue of the task runner's _push_xcom_if_needed. Regenerates the task-sdk datamodels, supervisor schema snapshot, Go models, and the ts-sdk supervisor schema.
jason810496
force-pushed
the
feature/go-sdk/taskflow-arg-binding
branch
from
July 26, 2026 07:52
3f276fe to
25e2d12
Compare
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.
Brings in the base branch's merge with main and its CI fixes. Main's datamodel-code-generator 0.41.0 bump and the merged lock's pydantic reorder the generated task-sdk datamodels, supervisor schema snapshot, ts-sdk supervisor types, and Go models, so all of them are regenerated here as part of the merge.
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.
The unmapped PR narrows to delivering plain TaskFlow call arguments; consuming per-map-index bindings (map_index row selection, element_index extraction), the via_expand example and e2e coverage, and the has_mapped_dependants/mapped_length supervisor recording move wholesale to the stacked follow-up branch feature/go-sdk/taskflow-arg-binding-mapped.
The supervisor schema no longer carries map_index/element_index on XComArgBinding (they move to the mapped follow-up), so the generated models and the binding package's scope note follow suit.
1 task
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
feature/lang-sdk/taskflow-stub-dag; only the tip commit is new here).Why
#69757 ships the Python-side contract: a
@task.stubTaskFlow call is serialized as an ordered arg-binding spec and returned byti_runasTIRunContext.arg_bindings. This PR makes the Go SDK actually consume it — Go task functions receive the Dag file's literals and upstream XComs as typed parameters instead of hand-writingGetXComcalls with hard-coded upstream task ids:Supported TaskFlow syntax
The example bundle's
taskflow_binding_dag(go-sdk/dags/go_examples.py+go-sdk/example/bundle/taskflowbinding/) exercises the full surface end to end:On the Go side a task declares either flat positional data parameters or one keyword-style
sdk.TaskInputstruct (mixing the two shapes is rejected at registration as too ambiguous):Binding semantics, mirroring positional vs keyword calls:
Noneonto a pointer arrives asnil.sdk.TaskInputfields are keyword-like: a field whose name/tag matches no call argument stays at its Go zero value; an argument the Dag author explicitly passed that no field claims fails the task;from_defaultentries (captured stub defaults) may go unclaimed.reflect.Zerofill.How
pkg/bindingpackage:Analyzeclassifies parameters once at registration (injectables vs JSON-decodable data parameters, or the singleTaskInputstruct);Resolvebinds the wire spec onto them. The wire union surfaces as a sealed sum type (binding.XComArg/binding.LiteralArg) whose variants andDataTypevocabulary are defined in terms of the generatedgenmodelsschema types, so the runtime types cannot drift from the wire model.TaskWithArgsdispatch in the task runner with strict wire-spec validation (missing/emptynameor xcomtask_idfails the task before its body runs); arity and malformed-spec errors share retry semantics.xcom:/xcom-key:struct tags were considered and dropped — they would hard-code upstream task ids in the Go binary, hiding wiring the Dag file owns; a task that needs an extra XCom still asks the injected client explicitly.jsonstruct tags so typed XComs round-trip.genmodelsfrom the supervisor schema2026-07-30+SupervisorSchemaVersionbump.What
go-sdk:pkg/binding(flat positional binding +sdk.TaskInputstruct-field injection); task-runner dispatch; regenerated genmodels; msgpackjson-tag support; example bundle (taskflow_binding_dagcovering the full binding surface, with everyvia_struct_*task binding an XCom-sourced field frommake_regionalongside a literal); README + ADR 0003 updated.airflow-e2e-tests:test_go_sdk_taskflow_binding.pydrives the example Dag end to end and asserts every task bound its arguments.Was generative AI tooling used to co-author this PR?