jason810496 commented on code in PR #69757: URL: https://github.com/apache/airflow/pull/69757#discussion_r3747619192
########## airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py: ########## @@ -0,0 +1,93 @@ +# 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 the ``@task.stub`` TaskFlow call, stored in the serialized +Dag, and delivered to the lang-SDK runtime via ``TIRunContext.arg_bindings``. +""" + +from __future__ import annotations + +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 + +# A named, titled alias (like TaskArgBinding below) kept as free-form JSON rather than a +# typed model, so unknown JSON-schema keywords survive re-serialization along the way. +ArgValueSchema = TypeAliasType( + "ArgValueSchema", Annotated[dict[str, JsonValue], Field(title="ArgValueSchema")] +) +"""JSON-schema fragment constraining the value a stub-task argument binds to; generated +by pydantic from the stub annotation, carried verbatim, unknown keywords ignored.""" + + +class _ArgBindingBase(BaseModel): + """Fields every :class:`TaskArgBinding` variant carries, regardless of ``kind``.""" + + name: str + """The stub function's parameter name this binding fills, in declaration order.""" + + value_schema: ArgValueSchema | None = None + """Schema fragment from the stub function's annotation; omitted when unconstrained.""" + + +class XComArgBinding(_ArgBindingBase): + """One positional stub-task argument pulled from an upstream task's XCom.""" + + # No default: it would drop ``kind`` from ``required``, and the generated task-sdk + # client then types it ``Literal | None``, invalid as a tagged-union discriminator. + kind: Literal["xcom"] Review Comment: I tried the following three approaches and would like to keep the current, as-is form to avoid corrupting the generated discriminated union. `kind` is the discriminator (`Field(discriminator="kind")`) on the `TaskArgBinding`. That union round-trips through `datamodel-code-generator`, which regenerates `task-sdk/src/airflow/sdk/api/datamodels/_generated.py` from this model's OpenAPI schema. All three options below were tested against the real generator (`uv run --active --group codegen --project apache-airflow-task-sdk --directory task-sdk -s dev/generate_task_sdk_models.py`). #### 1. Plain default ```python class XComArgBinding(_ArgBindingBase): kind: Literal["xcom"] = "xcom" ``` `datamodel-code-generator` regenerates the client field with an added `| None`: ```diff - kind: Annotated[Literal["xcom"], Field(title="Kind")] + kind: Annotated[Literal["xcom"] | None, Field(title="Kind")] = "xcom" ``` Because a default drops the field from the OpenAPI schema's `required` list, and the generator's response to "not required" is "wrap it in `Optional`" — regardless of the fact that the underlying type is a single-member `Literal` that **needs to stay bare for the discriminator to work**. `airflow` itself, for example `airflow.sdk.exceptions` reaches `airflow.sdk`, which eagerly touches `_generated.py`now blows up at class-definition time: ``` File ".../task-sdk/src/airflow/sdk/api/datamodels/_generated.py", line 812, in <module> class TaskArgBinding(RootModel[XComArgBinding | LiteralArgBinding]): File ".../pydantic/_internal/_model_construction.py", line 256, in __new__ complete_model_class( ... File ".../pydantic/_internal/_discriminated_union.py", line 473, in _infer_discriminator_values_for_inner_schema raise PydanticUserError( pydantic.errors.PydanticUserError: Model 'XComArgBinding' needs field 'kind' to be of type `Literal` For further information visit https://errors.pydantic.dev/2.13/u/discriminator-needs-literal ``` #### 2. `Field(init=False, default=...)` ```python class XComArgBinding(_ArgBindingBase): kind: Literal["xcom"] = Field(init=False, default="xcom") ``` `init=False` only removes the field from `__init__`'s signature; it does nothing to whether the field is *required* for schema/validation purposes, which is what `datamodel-code-generator` and the discriminator care about. ```diff - kind: Annotated[Literal["xcom"], Field(title="Kind")] + kind: Annotated[Literal["xcom"] | None, Field(title="Kind")] = "xcom" ``` Same `PydanticUserError: Model 'XComArgBinding' needs field 'kind' to be of type Literal` on import. `init=False` and a bare default are schema-equivalent as far as `required` is concerned. #### 3. Bare `Literal`, no default (current, as-is) ```python class XComArgBinding(_ArgBindingBase): # No default: it would drop `kind` from `required`, and the generated task-sdk # client then types it `Literal | None`, invalid as a tagged-union discriminator. kind: Literal["xcom"] ``` `kind` stays in the OpenAPI schema's `required` list, so `datamodel-code-generator` emits it as a bare, non-optional `Literal`: ```python kind: Annotated[Literal["xcom"], Field(title="Kind")] ``` `import airflow.sdk.api.datamodels._generated` succeeds cleanly. Every `TaskArgBinding` entry is built server-side as a plain dict literal in `build_arg_bindings` (`{"name": name, "kind": "xcom", ...}`), never via `XComArgBinding(...)`, so there's no call site that would benefit from a constructor-level default in the first place. ########## 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: Sorry for leaving that hanging and I left it as `None` to keep the `None` semantic as "regular Python Task" and the argument-less `@task.stub` call. The Lang SDK consumer side could even short circuit the further parsing based on the `None` semantic. ########## airflow-core/tests/unit/serialization/test_dag_serialization.py: ########## Review Comment: The last time we bump -- https://github.com/apache/airflow/pull/55836/changes#diff-2f03f28d3201c630dfbae758caf95ada61ab10328de789fc464f6e27c4afb2d0 is for the Airflow 3 major release. It seems we're unclear about the serialization versioning. For example, should we use - incremental versioning (just let it be `4` even though we're still in Airflow 3) - semantic versioning (e.g. `v3_3_0_to_v3_3_4` ). Then should we also deal add `v3_3_1_to_v3_3_4` between each patch release? - or follow the existing dated-base versioning like the execution API between each major version? I tracked the issue in https://github.com/apache/airflow/issues/71364 as follow-up. I need some more design for this and perhaps raise out as a discussion (offline discussion first for sure). ########## airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py: ########## @@ -372,6 +381,110 @@ 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() + + response = client.patch(f"/execution/task-instances/{tis['transform'].id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, Review Comment: Duplicate our offline discussion here: > know "how to unpack the XCom" into native object > > But that needs to be done at runtime after fetching the XCom -- we can't tell exactly what it might be at dag parse time. Yes, I know. Even without the data type field, the SDK runtime can still parse into the native type. However, this means we treat the SDK runtime side annotation (kind of a client side) as source of truth for unpacking. I prefer to keep the data type on an xcom arg with the following reasons: 1. Treat / unpack based on the server side (airflow core serialization) source of truth. 2. Make it possible to raise the early error when the type un-match on the Dag processing time (after AIP-85) for mixed Lang Dag 3. In the native Dag era, we can leverage this field to distinguish whether the upstream and downstream type is match or not in Dag processing time. ########## providers/standard/src/airflow/providers/standard/decorators/stub.py: ########## @@ -18,19 +18,283 @@ from __future__ import annotations import ast -from collections.abc import Callable +import copy +import datetime +import inspect +import json +import types +import typing +from collections.abc import Callable, Collection, Mapping +from functools import cache from typing import TYPE_CHECKING, Any +try: + from pydantic import PydanticUserError, TypeAdapter + from pydantic.json_schema import GenerateJsonSchema +except ImportError: + # Airflow 3 always ships pydantic but Airflow 2.x base installs do not; without it, + # stub args carry no value schemas and runtimes keep their decode-only fallback. + GenerateJsonSchema = object # type: ignore[assignment,misc] + TypeAdapter = None # type: ignore[assignment,misc] + PydanticUserError = None # type: ignore[assignment,misc] + from airflow.providers.common.compat.sdk import ( + KNOWN_CONTEXT_KEYS, + XCOM_RETURN_KEY, DecoratedOperator, + MappedOperator, + PlainXComArg, TaskDecorator, + XComArg, task_decorator_factory, ) if TYPE_CHECKING: from airflow.providers.common.compat.sdk import Context +class _ValueSchemaGenerator(GenerateJsonSchema): + """ + Pydantic's stock JSON-schema generation plus OpenAPI's fixed-width numeric formats. + + A foreign runtime decodes numbers into machine types, which the bare + ``integer``/``number`` type names cannot convey; ``format`` is an annotation per + JSON schema, so runtimes that don't know these names simply skip them. + """ + + def int_schema(self, schema): + return {**super().int_schema(schema), "format": "int64"} + + def float_schema(self, schema): + return {**super().float_schema(schema), "format": "double"} + + +# Most-derived first: datetime subclasses date, so it must be matched before date. +_TEMPORAL_BASES = (datetime.datetime, datetime.date, datetime.time, datetime.timedelta) + + +def _normalize_temporal_annotation(annotation: Any) -> Any: + """ + Map temporal subclasses (e.g. ``pendulum.DateTime``) to their stdlib base. + + Applied recursively through unions and containers, and only as a retry when direct + schema generation fails, so temporal types carrying their own pydantic schema keep it. + """ + # Parametrized generics must be detected before the plain-class branch: on Python + # 3.10, isinstance(list[X], type) is True and issubclass silently consults the + # origin, so the class branch would return list[X] unnormalized. + origin = typing.get_origin(annotation) + args = typing.get_args(annotation) + if origin is not None and args: + normalized = tuple(_normalize_temporal_annotation(arg) for arg in args) + if normalized == args: + return annotation + if origin in (typing.Union, types.UnionType): + return typing.Union[normalized] # noqa: UP007 -- runtime construction from a tuple + return origin[normalized] + if isinstance(annotation, type): + return next((base for base in _TEMPORAL_BASES if issubclass(annotation, base)), annotation) + return annotation + + +def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: + """ + Build the JSON-schema fragment for one stub parameter annotation, via pydantic. + + The pydantic-generated schema ships verbatim, so runtimes must treat it as + open-vocabulary JSON schema. Returns ``None`` when the annotation constrains nothing + (missing, ``Any``, bare ``None``) or pydantic cannot generate a schema for it; the + binding then omits ``value_schema`` and the foreign runtime falls back to a + decode-only check. + """ + if TypeAdapter is None: + return None + if annotation is inspect.Parameter.empty or annotation is None or annotation is Any: + return None + if annotation is type(None): + # get_type_hints normalizes a bare ``None`` annotation to NoneType; a parameter + # that can only ever be None constrains nothing worth shipping. + return None + try: + schema = _generate_value_schema(annotation) + except TypeError: + # Unhashable annotations cannot key the cache; generate directly. Any pydantic + # failure inside the body degrades to None there, so this retry never re-raises. + schema = _generate_value_schema.__wrapped__(annotation) + # Deep-copy so callers embedding the fragment never alias the cached dict. + return copy.deepcopy(schema) if schema else None + + +@cache +def _generate_value_schema(annotation: Any) -> dict[str, Any] | None: + """ + Generate the schema for one annotation, cached for the process lifetime. + + TypeAdapter construction is one of pydantic's most expensive operations and + annotations are static, so re-parses of the same Dag file must not re-pay it. + """ + # Reached only when pydantic is installed (``_infer_value_schema`` guards on + # ``TypeAdapter is None``), so ``PydanticUserError`` is a real exception class here. + # It is the base of PydanticSchemaGenerationError and PydanticInvalidForJsonSchema and + # covers annotations pydantic rejects outright (e.g. bare ClassVar); TypeError catches + # the exotic generics pydantic chokes on with a plain TypeError. Either way, "pydantic + # cannot schema this" degrades to no schema rather than failing Dag parsing. + try: + return TypeAdapter(annotation).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticUserError, TypeError): + normalized = _normalize_temporal_annotation(annotation) + if normalized is annotation: + return None + try: + return TypeAdapter(normalized).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticUserError, TypeError): + return None + + +def _validate_stub_signature(signature: inspect.Signature, task_id: str) -> None: + 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)" + ) + + +def _resolve_param_annotations(python_callable: Callable, signature: inspect.Signature) -> dict[str, Any]: + """Map each parameter to its parse-time-resolvable annotation (``Parameter.empty`` when not).""" + 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 resolve(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 + + return {name: resolve(name, param) for name, param in signature.parameters.items()} + + +def _ensure_json_literal(value: Any, task_id: str, name: str) -> None: + if next(XComArg.iter_xcom_references(value), None) is not None: + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a collection with an " + "upstream task output nested inside it; only a direct XComArg argument can cross " + "the language boundary -- pass the upstream output as its own argument" + ) + 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; pass it in its JSON form instead" + ) + + +def _validate_xcom_value(value: Any, task_id: str, name: str) -> bool: + """Validate an XComArg argument, returning True when it is a bindable direct upstream output.""" + if isinstance(value, PlainXComArg): + if value.key != XCOM_RETURN_KEY: + 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" + ) Review Comment: We only focus on the TaskFlow here and SDK runtime side even hard coded the `return_value` value. Or do you think both side should still handle the custom XCom key case? I lean toward focusing on TaskFlow default `return_value` only. ########## airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py: ########## @@ -0,0 +1,67 @@ +# 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. +"""Business logic backing the task-instance execution routes.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + from airflow.models.dagbag import DBDagBag + +# Task types (``TaskInstance.operator``, the operator class name) whose tasks carry a +# lang-SDK ``arg_bindings`` spec. Used to gate the serialized-Dag lookup so regular tasks +# never pay for it. The gate matches exact class names; a new lang-SDK operator adds its +# name here. +LANG_SDK_OPERATORS = frozenset({"_StubOperator"}) Review Comment: The whole implementation is migrated to core serialization layer, so it's now all "flag-based" instead of relying on the naming. (I plan to write an ADR for this design trade-off.) ```python if (task := dag.task_dict.get(ti.task_id)) is None or not task.is_stub: return None return task.arg_bindings ``` -- 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]
