Skip to content

Support TaskFlow call syntax on stub tasks for the Lang SDK - #69757

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

Support TaskFlow call syntax on stub tasks for the Lang SDK#69757
jason810496 wants to merge 6 commits into
apache:mainfrom
jason810496:feature/lang-sdk/taskflow-stub-dag

Conversation

@jason810496

@jason810496 jason810496 commented Jul 11, 2026

Copy link
Copy Markdown
Member

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:

Why

@task.stub tasks can only be declared argless today, so a Go task that needs an upstream's output has to hand-write GetXCom calls (with the upstream task_id hard-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:

@task.stub(queue="golang")
def transform(country: str, extracted: dict): ...


with DAG(...):
    transform("uk", extract())  # extract() is a normal Python @task
// The runtime (stacked #70209) binds "uk" onto country and pulls
// extract's XCom into extracted.
func transform(ctx sdk.TIRunContext, log *slog.Logger, country string, extracted map[string]any) error

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's taskflow_binding_dag example):

@task.stub(queue="golang")
def via_flat_args(
    name: str,
    count: int,
    ratio: float,
    enabled: bool,
    tags: list,
    config: dict,
    numbers: list,
    note: str | None = None,
): ...


@dag(dag_id="taskflow_binding_dag")
def taskflow_binding_dag():
    via_flat_args(
        "summary",
        3,
        2.5,
        True,  # positional scalar literals (str/int/float/bool)
        ["metrics", "hourly"],  # array literal
        config=make_config(),  # keyword arg: XCom from another @task.stub
        numbers=make_numbers(),  # XCom binding onto a typed array parameter
    )  # `note` unpassed: its None default is captured as from_default
    region = make_region()
    via_struct_no_tags(RegionCode=region, Threshold=0.75)  # one XCom fanned into several calls
    via_struct_arg_tag(region_code=region, threshold=0.75)  # literal + XCom mixed as kwargs
    via_struct_unmatched_arg(region_code=region)  # defaulted param left unpassed
  • Literals of every JSON shape, positional or keyword; the wire value_schema is the JSON-schema fragment pydantic generates from the parameter annotation (TypeAdapter(annotation).json_schema() with a GenerateJsonSchema subclass layering OpenAPI's int64/double numeric 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 standard date-time/date/time/duration formats, unions→standard anyOf (str | None{"anyOf": [{"type": "string"}, {"type": "null"}]}); annotations pydantic cannot schema (arbitrary classes, unresolvable names) and untyped/Any parameters omit value_schema entirely (decode-only binding).
  • XComArg return values from any upstream — a normal @task or another @task.stub — which also wires the dependency edge (transform("uk", extract()) implies extract >> transform).
  • Keyword arguments normalize to declaration order through signature binding, so the serialized spec is always positional.
  • Defaults left unpassed are captured with from_default: true, letting keyword-style consumers (the Go sdk.TaskInput struct mode) leave them unclaimed.
  • Argless calls behave exactly as before: no spec is serialized, and pre-existing signatures (**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/concat XComArgs, indexing an upstream by custom XCom key, *args/**kwargs or context-key parameter names (only when the call actually passes arguments), and non-JSON-serializable literals (including NaN/Infinity).

How

  • Parse time (providers/standard): _StubOperator binds 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 the common.compat sdk seam.
  • Wire model: each spec entry is one variant of a kind-discriminated union — XComArgBinding (pull of an upstream's return-value XCom by task_id) or LiteralArgBinding (inline JSON value, from_default-flagged when captured from a signature default) — carrying the stub parameter's name and an annotation-derived value_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.
  • Server: ti_run returns the spec as a new optional TIRunContext.arg_bindings field, resolved through the shared DBDagBag only for _StubOperator tasks and stripped for older clients by a new execution-API version 2026-10-30 (Airflow 3.4 target); the supervisor schema gets a mirror version 2026-10-30 with a downgrade path, and the generated ts-sdk models plus the Go SDK's SupervisorSchemaVersion pin follow the schema bump in this PR. The serialized-Dag schema.json documents the optional per-task _arg_bindings property (no SERIALIZER_VERSION bump: optional field, no serialization-logic change).
  • Anything outside the v1 contract fails loudly at parse time (the rejected list above).

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

Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment thread go-sdk/dags/go_examples.py Outdated
Comment thread go-sdk/pkg/binding/binding.go Outdated
Comment thread go-sdk/pkg/binding/binding.go Outdated
Comment thread go-sdk/pkg/binding/binding.go Outdated
Comment thread ts-sdk/src/generated/supervisor.ts Outdated

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will wait until next review then address my own comments to avoid CI-rerun.

Comment thread airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py Outdated
@jason810496
jason810496 marked this pull request as ready for review July 22, 2026 09:17
@jason810496
jason810496 requested a review from uranusjr July 22, 2026 09:17

@ashb ashb left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment thread airflow-core/tests/unit/serialization/test_dag_serialization.py
Comment thread airflow-core/tests/unit/serialization/test_dag_serialization.py Outdated
Comment thread providers/standard/src/airflow/providers/standard/decorators/stub.py Outdated

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks you Ash for the review. I will address the comments shortly.

Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py Outdated
Comment thread airflow-core/tests/unit/serialization/test_dag_serialization.py
Comment thread providers/standard/src/airflow/providers/standard/decorators/stub.py Outdated
Comment thread providers/standard/tests/unit/standard/decorators/test_stub.py
Comment thread providers/standard/pyproject.toml Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
@jason810496
jason810496 requested a review from Copilot July 24, 2026 08:22

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@jason810496
jason810496 force-pushed the feature/lang-sdk/taskflow-stub-dag branch from 7dd3bf9 to 2bafed2 Compare July 31, 2026 15:02
@uranusjr

uranusjr commented Aug 3, 2026

Copy link
Copy Markdown
Member

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.

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.
@jason810496
jason810496 force-pushed the feature/lang-sdk/taskflow-stub-dag branch from a2d48e6 to aad0a01 Compare August 6, 2026 03:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Status: In progress

Development

Successfully merging this pull request may close these issues.

5 participants