ashb commented on code in PR #69757: URL: https://github.com/apache/airflow/pull/69757#discussion_r3637137529
########## airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py: ########## @@ -0,0 +1,102 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Positional-argument binding spec for stub (foreign-runtime) tasks. + +Captured at parse time from a stub task's TaskFlow call (``@task.stub``), stored in the +serialized Dag, and delivered to the lang-SDK runtime through ``TIRunContext.arg_bindings`` +so it can bind the values onto the native task function's parameters. +""" + +from __future__ import annotations + +from enum import Enum +from functools import cache +from typing import Annotated, Literal + +from pydantic import Field, JsonValue, TypeAdapter +from typing_extensions import TypeAliasType + +from airflow.api_fastapi.core_api.base import BaseModel + + +class ArgBindingDataType(str, Enum): + """Language-neutral value type a stub-task argument binds to in the foreign runtime.""" + + STRING = "string" + INTEGER = "integer" + NUMBER = "number" + BOOLEAN = "boolean" + OBJECT = "object" + ARRAY = "array" + ANY = "any" + + +class XComArgBinding(BaseModel): + """One positional stub-task argument pulled from an upstream task's XCom.""" + + kind: Literal["xcom"] Review Comment: Nit? ```suggestion kind: Literal["xcom"] = "xcom" ``` ########## airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py: ########## @@ -0,0 +1,102 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Positional-argument binding spec for stub (foreign-runtime) tasks. + +Captured at parse time from a stub task's TaskFlow call (``@task.stub``), stored in the +serialized Dag, and delivered to the lang-SDK runtime through ``TIRunContext.arg_bindings`` +so it can bind the values onto the native task function's parameters. +""" + +from __future__ import annotations + +from enum import Enum +from functools import cache +from typing import Annotated, Literal + +from pydantic import Field, JsonValue, TypeAdapter +from typing_extensions import TypeAliasType + +from airflow.api_fastapi.core_api.base import BaseModel + + +class ArgBindingDataType(str, Enum): + """Language-neutral value type a stub-task argument binds to in the foreign runtime.""" + + STRING = "string" + INTEGER = "integer" + NUMBER = "number" Review Comment: JSON has no distinction between integer and number. Do we need both types, or is Number alone enough? (In go for example we already have to possibly/conditional convert Integer to different int types (int32, int, in64 etc etc) ########## airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py: ########## @@ -310,6 +332,30 @@ def ti_run( should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries), ) + # Only set for stub (foreign-runtime) tasks with a captured TaskFlow arg + # spec; the route excludes unset fields, keeping regular responses lean. + if ti.operator == _STUB_TASK_TYPE and ( + arg_bindings := _get_arg_bindings(dag_bag, ti.dag_version_id, ti.task_id, session=session) + ): + try: + context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) + except ValidationError: + log.exception( + "Serialized arg_bindings spec failed validation", + dag_id=ti.dag_id, + task_id=ti.task_id, Review Comment: ```suggestion task_id=ti.task_id, dag_version_id=ti.dag_version_id, ``` ########## airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py: ########## @@ -0,0 +1,102 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Positional-argument binding spec for stub (foreign-runtime) tasks. + +Captured at parse time from a stub task's TaskFlow call (``@task.stub``), stored in the +serialized Dag, and delivered to the lang-SDK runtime through ``TIRunContext.arg_bindings`` +so it can bind the values onto the native task function's parameters. +""" + +from __future__ import annotations + +from enum import Enum +from functools import cache +from typing import Annotated, Literal + +from pydantic import Field, JsonValue, TypeAdapter +from typing_extensions import TypeAliasType + +from airflow.api_fastapi.core_api.base import BaseModel + + +class ArgBindingDataType(str, Enum): Review Comment: I wonder if we can reuse something from JSON Schema here instead of our own types... ########## airflow-core/tests/unit/serialization/test_dag_serialization.py: ########## @@ -3405,6 +3405,44 @@ def inner(): assert serialized3["python_callable_name"] == "empty_function" +def test_stub_task_args_round_trip(): + """The stub task's TaskFlow arg spec (``_arg_bindings``) survives Dag serialization.""" + from airflow.sdk import task + + with DAG(dag_id="arg_bindings_dag", schedule=None) as dag: + + @task.stub + def extract(): ... + + @task.stub + def transform(country: str, extracted: dict): ... + + transform("uk", extract()) + + ser_dag = DagSerialization.to_dict(dag) + encoded_tasks = {t[Encoding.VAR]["task_id"]: t[Encoding.VAR] for t in ser_dag["dag"]["tasks"]} + assert "_arg_bindings" not in encoded_tasks["extract"], "argless stubs must not serialize a spec" + assert encoded_tasks["transform"]["_arg_bindings"] == [ + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, + }, + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, + }, + ] + + round_tripped = DagSerialization.from_dict(ser_dag) + assert round_tripped.task_dict["transform"]._arg_bindings == [ + {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, + {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, + ] + assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") or ( + round_tripped.task_dict["extract"]._arg_bindings is None + ) Review Comment: not hasattr or is None feels a bit odd. Pick one to assert. ########## providers/standard/src/airflow/providers/standard/decorators/stub.py: ########## @@ -18,22 +18,170 @@ from __future__ import annotations import ast -from collections.abc import Callable -from typing import TYPE_CHECKING, Any +import inspect +import json +import types +import typing +from collections.abc import Callable, Collection, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Union from airflow.providers.common.compat.sdk import ( + KNOWN_CONTEXT_KEYS, DecoratedOperator, + PlainXComArg, TaskDecorator, + XComArg, task_decorator_factory, ) if TYPE_CHECKING: from airflow.providers.common.compat.sdk import Context +def _infer_data_type(annotation: Any) -> str: + """ + Map a stub function parameter annotation to the language-neutral arg-type vocabulary. + + The returned name is one of the execution API's ``ArgBindingDataType`` values + (``string``/``integer``/``number``/``boolean``/``object``/``array``/``any``); the foreign + runtime type-checks the bound value against it. Anything we cannot classify confidently + maps to ``any`` so binding falls back to a decode-only check. + """ + if annotation is inspect.Parameter.empty or annotation is None or annotation is Any: + return "any" + origin = typing.get_origin(annotation) + if origin is not None: + if origin is Union or origin is types.UnionType: + members = [a for a in typing.get_args(annotation) if a is not type(None)] + if len(members) == 1: + return _infer_data_type(members[0]) + return "any" + annotation = origin + if not isinstance(annotation, type): + return "any" + # bool subclasses int, and str/bytes are Sequences -- order matters. + if issubclass(annotation, bool): + return "boolean" + if issubclass(annotation, int): + return "integer" + if issubclass(annotation, float): + return "number" + if issubclass(annotation, str): + return "string" + if issubclass(annotation, bytes): + return "any" + if issubclass(annotation, (dict, Mapping)): + return "object" + if issubclass(annotation, (list, tuple, set, frozenset, Sequence)): + return "array" + return "any" + + +def _build_arg_bindings( + python_callable: Callable, + op_args: Collection[Any], + op_kwargs: Mapping[str, Any], + task_id: str, +) -> list[dict[str, Any]] | None: + """ + Bind the TaskFlow call arguments to the stub signature and build the ordered arg spec. + + Each spec entry is a plain dict matching one variant of the execution API's + ``TaskArgBinding`` union: an ``XComArgBinding`` (``kind="xcom"``) for upstream TaskFlow + outputs, or a ``LiteralArgBinding`` (``kind="literal"``) for everything else. ``name`` is + always the stub function's parameter name, so a foreign runtime can bind by name (e.g. the + Go SDK's ``sdk.TaskInput`` struct fields) in addition to the existing positional order. + Returns ``None`` for argless calls: the binding contract (including the signature checks + below) applies only once a TaskFlow call actually passes arguments, so pre-TaskFlow stub + Dags whose call arguments were always ignored keep parsing. + """ + if not op_args and not op_kwargs: + return None + + signature = inspect.signature(python_callable) + + for param in signature.parameters.values(): + if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): + raise ValueError( + f"@task.stub task {task_id!r} must declare a fixed number of parameters for the " + f"foreign runtime to bind against; *{param.name} is not supported" + ) + if param.name in KNOWN_CONTEXT_KEYS: + raise ValueError( + f"@task.stub task {task_id!r} parameter {param.name!r} is an Airflow context key; " + "stub signatures declare only data parameters -- the lang-SDK runtime injects its " + "own task context natively (e.g. the Go SDK's sdk.TIRunContext parameter)" + ) + + bound = signature.bind(*op_args, **op_kwargs) + explicitly_bound = set(bound.arguments) + bound.apply_defaults() + + try: + hints = typing.get_type_hints(python_callable) + except (NameError, TypeError): + # Annotations that cannot be resolved at parse time (e.g. names behind + # TYPE_CHECKING with ``from __future__ import annotations``) degrade to "any". + hints = {} + + def get_annotation_for(name: str, param: inspect.Parameter) -> Any: + if name in hints: + return hints[name] + if isinstance(param.annotation, str): + return inspect.Parameter.empty + return param.annotation + + spec: list[dict[str, Any]] = [] + for name, param in signature.parameters.items(): + value = bound.arguments[name] + data_type = _infer_data_type(get_annotation_for(name, param)) + if isinstance(value, PlainXComArg): + if value.key != "return_value": + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} references the XCom key " + f"{value.key!r}; only an upstream task's return value can cross the language " + "boundary -- indexing an output by a custom key is not supported" + ) + spec.append( + { + "name": name, + "kind": "xcom", + "data_type": data_type, + "task_id": value.operator.task_id, + } + ) + continue + if isinstance(value, XComArg): + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a " + f"{type(value).__name__}; only direct upstream task outputs can cross the " + "language boundary -- .map()/.zip()/.concat() results are not supported" + ) + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError): + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a literal of type " + f"{type(value).__name__} that is not JSON-serializable, so it cannot be passed " + "to the foreign runtime" + ) + entry: dict[str, Any] = {"name": name, "kind": "literal", "data_type": data_type, "value": value} + if name not in explicitly_bound: + entry["from_default"] = True + spec.append(entry) + return spec + + class _StubOperator(DecoratedOperator): custom_operator_name: str = "@task.stub" + # Mapped stubs would need per-map-index arg specs, which the foreign runtime cannot Review Comment: "Mapped stubs would need per-map-index arg specs" I don't think this is true -- it's the same function for each mapped index (by design) so each index would receive the same type of arguments. ########## airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py: ########## @@ -0,0 +1,102 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Positional-argument binding spec for stub (foreign-runtime) tasks. + +Captured at parse time from a stub task's TaskFlow call (``@task.stub``), stored in the +serialized Dag, and delivered to the lang-SDK runtime through ``TIRunContext.arg_bindings`` +so it can bind the values onto the native task function's parameters. +""" + +from __future__ import annotations + +from enum import Enum +from functools import cache +from typing import Annotated, Literal + +from pydantic import Field, JsonValue, TypeAdapter +from typing_extensions import TypeAliasType + +from airflow.api_fastapi.core_api.base import BaseModel + + +class ArgBindingDataType(str, Enum): Review Comment: What about Dates/DateTimes? We handle those specially in xcom serialization already don't we?, and it's likely quite common to want to pass those along, so I think we should handle the "convert to target language native date type" rather than just passing it as a string and making task code parse it. ########## airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py: ########## @@ -435,6 +436,13 @@ class TIRunContext(BaseModel): always reflects when the task *first* started, not when it was rescheduled/resumed. """ + arg_bindings: list[TaskArgBinding] | None = None + """ + Ordered positional-argument binding spec for stub (foreign-runtime) tasks. + + ``None`` for regular tasks and for stub tasks that declare no parameters. + """ Review Comment: Hmmmmm, I wonder if this should not allow none, and make it an empty list in that case. I don't think it functionally makes a difference but... 🤔 ########## airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py: ########## @@ -51,9 +51,11 @@ AddTeamNameField, AddVariableKeysEndpoint, ) +from airflow.api_fastapi.execution_api.versions.v2026_07_30 import AddArgBindingsToTIRunContext Review Comment: What Airflow version are we targeting this for release? If its 3.4 then this should be further in the future, say 2026_10_30. Or do you want to make a case that adding support for TaskFlow in go is a bug fix? ########## providers/standard/src/airflow/providers/standard/decorators/stub.py: ########## Review Comment: Given how tightly integrated into the exec API this change is, I'm not sure putting it in standard provider is right -- my first thought is that this should live in/with the dag parsing code, not with the stub operator code? ########## airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py: ########## @@ -370,6 +371,87 @@ async def workload_token(request: Request) -> TIToken: assert extras["scope"] == "execution" assert extras["sub"] == str(ti.id) + def test_ti_run_returns_arg_bindings_for_stub_task(self, client, dag_maker): + """A stub task's TaskFlow arg spec is extracted from the serialized Dag and returned.""" + with dag_maker("test_arg_bindings_dag", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(country: str, extracted: dict, limit: int = 10): ... + + transform("uk", extract()) + + dr = dag_maker.create_dagrun() + tis = {ti.task_id: ti for ti in dr.get_task_instances()} + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + payload = { + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": "2024-09-30T12:00:00Z", + } + + response = client.patch(f"/execution/task-instances/{tis['transform'].id}/run", json=payload) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, + {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, Review Comment: data_type on an xcom arg feels slightly odd here -- we don't/cant know what the type is from the dag defn really, as it depends on what the task actually pushes doesn't it? ```python @task def upstream(): if CONDITION: return {"data": "one"} else return False ``` Extreme case that Golang tasks may never be able to cope with, but regardless, we don't know the type of an xcom. ########## airflow-core/tests/unit/serialization/test_dag_serialization.py: ########## Review Comment: I wonder if introducing arg_bindings should result in a bump in the serialization version? @amoghrajesh WDYT? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
