Support TaskFlow call syntax on stub tasks for the Lang SDK - #69757
Support TaskFlow call syntax on stub tasks for the Lang SDK#69757jason810496 wants to merge 6 commits into
Conversation
7b528b3 to
bec0e8d
Compare
fca5aa1 to
0b464a1
Compare
6c609d9 to
c21137f
Compare
82a0aa5 to
7e2d0f2
Compare
jason810496
left a comment
There was a problem hiding this comment.
I will wait until next review then address my own comments to avoid CI-rerun.
ashb
left a comment
There was a problem hiding this comment.
Overall I like the direction. Almost all of my comments I can be challenged on, don't just make the changes if you think the current way is better/more correct
jason810496
left a comment
There was a problem hiding this comment.
Thanks you Ash for the review. I will address the comments shortly.
7dd3bf9 to
2bafed2
Compare
|
I can’t comment on the implementation, but the proposed interface looks very reasonable to me. It seems to me the described programming interface has not been fully implemented in this PR, and I don’t have enough knowledge to say whether this is on the right track or not tbh. Maybe it would be a good idea to edit the PR description to ground what exactly should be expected here. |
2bafed2 to
1b7e1d6
Compare
1b7e1d6 to
73bd2f8
Compare
ed70fb2 to
a2d48e6
Compare
The @task.stub TaskFlow support in providers-standard imports KNOWN_CONTEXT_KEYS, PlainXComArg, MappedOperator and the decorator base classes through the compat layer so the provider keeps working down to Airflow 2.11. Those symbols first ship in common-compat 1.19.0 (1.18.0 was released from main in the meantime without them), so the version is cut here for the standard provider's pin to resolve.
Stub tasks silently ignored TaskFlow call arguments, so a Dag author could not hand literals or upstream XCom results to a lang-SDK runtime. The decorator now binds the call to the stub's signature at parse time and captures an ordered arg spec (literal values and direct upstream XCom references, with pydantic-derived JSON value schemas) that serializes with the Dag, while rejecting what cannot cross the language boundary: custom XCom keys, aggregated mapped outputs, non-JSON literals, and stubs with arguments inside mapped task groups. Mapped (.expand()) stubs capture no spec and keep the legacy behavior until a follow-up delivers per-map-index bindings.
TIRunContext gains an arg_bindings field so a lang-SDK runtime receives the stub task's TaskFlow arg spec at startup. ti_run derives it from the serialized Dag only for stub operators, so regular tasks never pay for the lookup, and only for clients on the new API version -- gated on the Cadwyn VersionChangeWithSideEffects.is_applied check rather than a date comparison -- so stub Dags that predate arg bindings keep running against older clients, for which the version migration strips the field.
StartupDetails in the supervisor wire schema carries the new arg_bindings so foreign runtimes receive the spec at task startup, with a version migration that strips it for runtimes pinned to the previous schema. The Go and TS SDKs regenerate against the new schema version; the Go arg-binding runtime itself lands in a stacked follow-up PR.
An XComArg buried in a list or dict literal fell through to the JSON check, whose "pass it in its JSON form instead" advice is impossible to follow for a task output. Detect nested references up front and point the author at the working alternative: pass the upstream output as its own argument.
When a PR cuts a new provider version while the previous version is still being voted on, only the rcN tags exist on the apache remote - the final tag is pushed after the vote passes. The changes-table walk in _get_all_changes_for_package assumed every past version has a final tag and crashed with git exit 128 in that window, breaking CI for any PR that bumps a provider version during a release wave.
a2d48e6 to
aad0a01
Compare
Scope
This PR ships the Python side only of the contract: parse-time capture of the TaskFlow call into a serialized arg-binding spec, the wire model, and its delivery to SDK runtimes through the Execution API and supervisor schema.
Nothing in this PR binds arguments inside a task runtime — the Go snippet below illustrates the consumer and lives in the stacked follow-up.
Stacked on top of this PR:
arg_bindings).expand()/.partial()rejection below)Why
@task.stubtasks can only be declared argless today, so a Go task that needs an upstream's output has to hand-writeGetXComcalls (with the upstreamtask_idhard-coded in Go, duplicating the wiring the Dag file already expresses). This PR ships the Python side of making the natural TaskFlow call work across the language boundary:Supported TaskFlow syntax
Every form below parses, serializes, and round-trips through the execution API (provider capture matrix in
test_stub.py; the Dag shown is #70209'staskflow_binding_dagexample):value_schemais the JSON-schema fragment pydantic generates from the parameter annotation (TypeAdapter(annotation).json_schema()with aGenerateJsonSchemasubclass layering OpenAPI'sint64/doublenumeric formats):str→{"type": "string"},int→{"type": "integer", "format": "int64"},dict[str, int]→{"type": "object", "additionalProperties": {...}},list[int]→{"type": "array", "items": {...}},Literal["a", "b"]→{"type": "string", "enum": [...]},datetime/date/time/timedelta→string with the standarddate-time/date/time/durationformats, unions→standardanyOf(str | None→{"anyOf": [{"type": "string"}, {"type": "null"}]}); annotations pydantic cannot schema (arbitrary classes, unresolvable names) and untyped/Anyparameters omitvalue_schemaentirely (decode-only binding).@taskor another@task.stub— which also wires the dependency edge (transform("uk", extract())impliesextract >> transform).from_default: true, letting keyword-style consumers (the Gosdk.TaskInputstruct mode) leave them unclaimed.**kwargs, context-key parameter names) keep parsing.Rejected loudly at parse time (v1 scope):
.expand()/.partial()on a stub, stubs called with arguments inside a mapped task group,map/zip/concatXComArgs, indexing an upstream by custom XCom key,*args/**kwargsor context-key parameter names (only when the call actually passes arguments), and non-JSON-serializable literals (including NaN/Infinity).How
_StubOperatorbinds the TaskFlow call to the stub's signature and serializes an ordered positional-arg spec with the Dag (_arg_bindings); the cross-version imports it needs are routed through thecommon.compatsdk seam.kind-discriminated union —XComArgBinding(pull of an upstream's return-value XCom bytask_id) orLiteralArgBinding(inline JSON value,from_default-flagged when captured from a signature default) — carrying the stub parameter'snameand an annotation-derivedvalue_schema. The fragment is deliberately free-form (dict[str, JsonValue], not a typed model) so every keyword pydantic generated survives the server → supervisor → runtime trip verbatim; consumers validate the keywords they understand and ignore the rest, per JSON-schema semantics. The exact fragment shape follows the pydantic version active at Dag-parse time.ti_runreturns the spec as a new optionalTIRunContext.arg_bindingsfield, resolved through the sharedDBDagBagonly for_StubOperatortasks and stripped for older clients by a new execution-API version2026-10-30(Airflow 3.4 target); the supervisor schema gets a mirror version2026-10-30with a downgrade path, and the generated ts-sdk models plus the Go SDK'sSupervisorSchemaVersionpin follow the schema bump in this PR. The serialized-Dagschema.jsondocuments the optional per-task_arg_bindingsproperty (noSERIALIZER_VERSIONbump: optional field, no serialization-logic change).Was generative AI tooling used to co-author this PR?