Skip to content

Bind TaskFlow stub-task call arguments in the Go SDK runtime - #70209

Open
jason810496 wants to merge 3 commits into
apache:mainfrom
jason810496:feature/go-sdk/taskflow-arg-binding
Open

Bind TaskFlow stub-task call arguments in the Go SDK runtime#70209
jason810496 wants to merge 3 commits into
apache:mainfrom
jason810496:feature/go-sdk/taskflow-arg-binding

Conversation

@jason810496

@jason810496 jason810496 commented Jul 22, 2026

Copy link
Copy Markdown
Member

Why

#69757 ships the Python-side contract: a @task.stub TaskFlow call is serialized as an ordered arg-binding spec and returned by ti_run as TIRunContext.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-writing GetXCom calls with hard-coded upstream task ids:

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


with DAG(...):
    transform("uk", extract())  # extract() is a normal Python @task
// The runtime 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

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:

@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
    via_flat_map(config={"region": "eu-west-1", "count": 3})  # dict decoded whole into a struct
    via_struct_map(payload={"region": "eu-west-1", "count": 3})  # dict onto a struct's map field
    via_plain_map(labels={"team": "data", "tier": "gold"})  # dict into a plain Go map

On the Go side a task declares either flat positional data parameters or a single struct whose fields bind by name (mixing the two shapes is rejected at registration as too ambiguous):

// Flat parameters: bound in declaration order after the injectables
// (sdk.TIRunContext, *slog.Logger, context.Context, client interfaces).
// Arity or declared-type mismatches fail the task before its body runs.
func ViaFlatArgs(ctx sdk.TIRunContext, log *slog.Logger,
    name string, count int, ratio float64, enabled bool,
    tags []string, config Config, numbers []int, note *string) (any, error)

// Sole struct parameter: exported fields bind per call-argument name.
type ViaStructArgTagInput struct {
    Region    string  `arg:"region_code"` // explicit argument name via tag
    Threshold float64 `arg:"threshold"`   // untagged fields bind their verbatim Go field name
}
func ViaStructArgTag(ctx sdk.TIRunContext, log *slog.Logger, input ViaStructArgTagInput) (any, error)

Binding semantics, mirroring positional vs keyword calls:

  • Flat parameters bind by position after injectables. Literals, XComs, and captured defaults
    fill every data parameter, with None decoded into nil-capable types.
  • Sole structs bind by field name or arg: tag. Unmatched fields keep their zero value,
    from_default entries may go unclaimed, and explicitly passed unclaimed arguments fail.
  • Whole-value structs decode one unclaimed argument into an untagged struct. Tagged structs do
    not use this fallback.
  • Invalid bindings fail before task code runs. Missing or malformed specs, unsupported kinds,
    schema mismatches, and strict struct decode errors are reported instead of silently zero-filling.

How

New pkg/binding package:

  • Analyze classifies parameters once at registration (injectables vs JSON-decodable data parameters, or the single name-bound struct), recursing through nested types so an undecodable one fails when the bundle is built rather than on every execution.
  • Resolve binds the wire spec onto them. The wire union surfaces as a sealed sum type (binding.XComArg/binding.LiteralArg) whose variants and DataType vocabulary are defined in terms of the generated genmodels schema types, so the runtime types cannot drift from the wire model.

Rejected Alternative

  • Ad hoc xcom:/xcom-key: struct tags were considered and dropped. A task that needs an extra XCom still asks the injected client explicitly.

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

A Go task could only reach an upstream task's output by hand-writing a GetXCom
call against a hard-coded task id, duplicating wiring the Dag file already
owns and breaking silently whenever that upstream was renamed.

apache#69757 ships the Python half: a `@task.stub` TaskFlow call is captured at Dag
serialization as an ordered arg-binding spec and returned by ti_run. Consuming
it here lets a Go task function take the Dag's literals and upstream XComs as
ordinary typed parameters.

A function declares either flat positional parameters or a single struct whose
fields bind by name -- kwarg-style, so an unmatched field keeps its zero value
while an argument no field claims fails the task. Signature problems are caught
once at registration; per-execution arity, type and spec errors fail the task
before its body runs, replacing a silent reflect.Zero fill.
Several binding problems stayed quiet until they were expensive or confusing.
A parameter nesting an undecodable value failed on every execution rather than
once when the bundle was built. A struct carrying `arg:` tags whose single
argument no tag matched was decoded whole into the struct, so a typo'd tag
surfaced as a decode error naming the Go type instead of the argument that
matched nothing. A value_schema or from_default of the wrong wire shape was
indistinguishable from an absent one, disabling the declared-type check or
turning a captured stub default into an argument the author supposedly passed.

The XCom whole-value decode and the concurrent multi-pull failure path were
also reachable from a Dag but exercised only in their simplest shape, and the
package docs re-explained the whole binding model at three sites, burying the
rules they were meant to state.
@jason810496
jason810496 force-pushed the feature/go-sdk/taskflow-arg-binding branch 2 times, most recently from 1ddfc95 to 0507c8a Compare August 18, 2026 08:51
The Edge Worker's execution API carries no argument spec at all, so failing a
task there for an argument-count mismatch blamed the Dag author for a limit of
the transport. Keeping data parameters at their Go zero values is how those
tasks behaved before binding existed, and that path is in maintenance rather
than gaining the spec.

Registration rejected struct shapes that decode without complaint -- a struct
carrying a callback alongside its data never needed the callback filled -- and
because registering a task panics, one such signature took its whole bundle
down at startup rather than the single task.

Adding a defaulted parameter to a stub is backwards compatible in Python, and
has to stay so for the Go functions already bound to that stub: the captured
default reaches the wire but needs no Go parameter to receive it.

Untagged fields matched a Go field name verbatim, which no idiomatic snake_case
stub parameter can produce, so tags were mandatory in practice and a mismatch
quietly fell back to decoding the argument whole. Folding case and underscores
makes the untagged form usable, and embedded structs now contribute their
fields the way encoding/json has all along.

A type that decodes itself from JSON also passed registration only to be
rejected at run time by a schema check judging it on its Go kind.
@jason810496
jason810496 force-pushed the feature/go-sdk/taskflow-arg-binding branch from 0507c8a to be8f718 Compare August 19, 2026 06:08

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.

Should we add a hook to ensure this file is up-to-date? (Or at least not incorrectly regenerated.)

@jason810496 jason810496 Aug 19, 2026

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.

IIRC, the last consensus we had for this is to bump the generated file manually instead of having fully automate static check.

Since the one who change the Task SDK schema might not able to change the Lang SDK side generated schema.

@jason810496 jason810496 Aug 19, 2026

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.

(Or at least not incorrectly regenerated.)

Make sense for this part, the best way I can come up with is vendoring the the supervisor schema JSON. But this introduce another question: How can we ensure the vendored supervisor schema JSON is correct?

We already had e2e test to coverage all the features in Go SDK, so I think it's fine not to have a static check to guard the auto generated schema. If it's a malformed schema, it can't even pass the e2e test.

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.

2 participants