diff --git a/.agents/skills/airflow-java-sdk/SKILL.md b/.agents/skills/airflow-java-sdk/SKILL.md index 8edb742679535..9921832eafee5 100644 --- a/.agents/skills/airflow-java-sdk/SKILL.md +++ b/.agents/skills/airflow-java-sdk/SKILL.md @@ -50,8 +50,9 @@ subclasses must only import from `org.apache.airflow.sdk`; any import of ## Bundle composition and coordinator discovery -A **bundle** is a directory of JAR files (typically `build/bundle/`) placed on the coordinator's -`jars_root`. The coordinator scans the directory at task-dispatch time to find: +A **bundle** is a directory of JAR files (typically `build/bundle/`) placed in the Dag bundle named +by the coordinator's `task_handler_bundle_name` (the task's own Dag bundle when unset). The +coordinator scans that Dag bundle at task-dispatch time to find: 1. **`Main-Class`** (standard JAR manifest attribute) — the fully-qualified class name of the entry point that the coordinator invokes with `java -classpath … --comm … --logs …`. @@ -65,7 +66,7 @@ A **bundle** is a directory of JAR files (typically `build/bundle/`) placed on t `runtimeClasspath` and copies it into the shadow JAR manifest. In thin-JAR mode (`fatJar = false`), the value stays in the `airflow-sdk` JAR deployed alongside the bundle JAR. -The Python coordinator (`JavaCoordinator`) scans every JAR under `jars_root` with +The Python coordinator (`JavaCoordinator`) scans every JAR in that Dag bundle with `_JarInfo.find()`, reads `META-INF/MANIFEST.MF` out of each ZIP, and collects `Main-Class` and `Airflow-Supervisor-Schema-Version` from whichever JARs carry them. The resolved schema version is then passed as the `schema_version` return value from `_build_execute_task_command`, which @@ -74,7 +75,7 @@ the base `SubprocessCoordinator` uses to negotiate the supervisor wire protocol. If `main_class` is set explicitly on the `JavaCoordinator` instance (via `[sdk] coordinators` kwargs), the scan uses it as a filter; otherwise the first JAR with a `Main-Class` attribute wins. Either way, `Airflow-Supervisor-Schema-Version` must be present in at least one JAR in -`jars_root` or startup fails. +the Dag bundle or startup fails. Every JAR in the Dag bundle goes on one classpath. --- @@ -117,7 +118,7 @@ E2E_TEST_MODE=java_sdk uv run --project airflow-e2e-tests pytest \ `coordinator.py` extends `SubprocessCoordinator`. The only method subclasses must implement is `_build_execute_task_command`, which returns `(argv, schema_version)`. Look at the existing -implementation for how `jars_root`, `java_executable`, `jvm_args`, and `main_class` are +implementation for how the scanned Dag bundle, `java_executable`, `jvm_args`, and `main_class` are assembled into the command. Do not reach into the JVM process from Python beyond what this method provides. diff --git a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md index c2acaadbedf2a..5dc9518068130 100644 --- a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md +++ b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md @@ -90,7 +90,7 @@ DagFileProcessorProcess(etl.py) ← manager spawn ├── _serialize_dags(bag) → is_stub tasks carry arg_bindings (ADR-0007) │ │ ┌─────────────────────────────────────────────────────────────────────┐ - │ │ Step 1: Collect the dag_ids to ask about — every Dag in this │ + │ │ Step 1: Collect the Dags to validate — every Dag in this │ │ │ file with at least one is_stub task → ["etl"] │ │ └─────────────────────────────────────────────────────────────────────┘ │ @@ -123,27 +123,26 @@ DagFileProcessorProcess(etl.py) ← manager spawn │ │ └── BundleScanner.scanBundles(roots) │ │ │ → "etl" → ResolvedBundle(analytics.jar, mainClass, ...) │ │ │ │ - │ │ Group the dag_ids by (coordinator, artifact) — one group, one │ + │ │ Group the stub tasks by (coordinator, artifact) — one group, one │ │ │ process, one request │ │ └─────────────────────────────────────────────────────────────────────┘ │ │ ┌─────────────────────────────────────────────────────────────────────┐ │ │ Step 4: Query each group — one request, one response │ │ │ │ - │ │ SDKTaskHandlerProcessorProcess.start( │ - │ │ target=_parse_task_handler_entrypoint, │ + │ │ LangSDKTaskHandlerProcessorProcess.start( │ + │ │ target=_start_task_handler_runtime_entrypoint, │ │ │ coordinator=JavaCoordinator("jdk-11"), │ │ │ path=analytics.jar) │ │ │ │ │ │ │ ├── in the child: _build_parse_task_handler_command() │ │ │ │ coordinator.parse_task_handler() — spawn JVM │ │ │ │ │ - │ │ │ ──TaskHandlerParseRequest(file=analytics.jar, │ - │ │ │ dag_ids=["etl"])─────▶ JVM │ + │ │ │ ──TaskHandlerParseRequest(file=analytics.jar)─────▶ JVM │ │ │ │ (ToSDKTaskHandlerProcessor) │ │ │ │ │ - │ │ │ JVM answers from its own TaskHandler registrations │ - │ │ │ whose dagId is one of the requested ids │ + │ │ │ JVM answers with every TaskHandler registration │ + │ │ │ in the artifact, or {} when there is none │ │ │ │ │ │ │ │ ◀─TaskHandlerParsingResult(task_handlers={ │ │ │ │ "etl": [extract, transform, load]})─────── JVM │ @@ -162,7 +161,9 @@ DagFileProcessorProcess(etl.py) ← manager spawn │ │ Python Dag "etl" (stub tasks) TaskHandlerDeclaration │ │ │ ────────────────────────────── ────────────────────────────── │ │ │ task_id ↔ task_id (sets must match) │ - │ │ arg_bindings[*].name ↔ params[*].name (in order) │ + │ │ arg_bindings[*] ↔ params[*] (per binding) │ + │ │ by position, or by folded or exact name │ + │ │ unmatched by name → a warning, not an error │ │ │ arg_bindings[*].value_schema ↔ params[*].value_schema │ │ │ compared only where neither side is null │ │ │ │ @@ -175,8 +176,8 @@ DagFileProcessorProcess(etl.py) ← manager spawn ``` The parse owns validation, not an importer. `PythonDagImporter` returns `airflow.sdk.DAG` objects and knows nothing about coordinators or queues, so `@task.stub` keeps working for -any importer that can produce a Dag carrying stub tasks. `_parse_file` is also the only place where the whole file's Dags are visible at once, which is what lets one request cover -every `dag_id` that resolved to the same artifact ([ADR-0012](0012-lang-sdk-parse-protocol.md)). +any importer that can produce a Dag carrying stub tasks. `_parse_file` is also the only place where the whole file's Dags are visible at once, which is what lets one request per +artifact serve every Dag in the file whose stubs resolved to it ([ADR-0012](0012-lang-sdk-parse-protocol.md)). Resolution goes through the coordinator registry, not the filesystem, so the Python Dag and the Lang-SDK artifact **do not need to be in the same DagBundle**. Nothing here needs an `airflow.sdk.DAG` round-trip either — validation compares against the Dag the Python parser already built. Appendix B states exactly what is compared. @@ -186,7 +187,7 @@ Resolution goes through the coordinator registry, not the filesystem, so the Pyt | Caller | Coordinator call | What comes back | Action | |--------------------------------------------|------------------------------------------------------|------------------------------------------|--------------------------------------------------| | `_parse_file` → `PythonDagImporter` | — (the Python file is parsed in process) | its own parsed Dags | PERSIST | -| `_parse_file`, per (coordinator, artifact) | `parse_task_handler`, scoped to that group's dag_ids | `TaskHandlerParsingResult` | VALIDATE only — not a Dag, so nothing to persist | +| `_parse_file`, per (coordinator, artifact) | `parse_task_handler`, for every handler it registers | `TaskHandlerParsingResult` | VALIDATE only — not a Dag, so nothing to persist | | `_parse_file` → `JavaDagImporter` | `parse_dag` | `DagFileParsingResult`, native Dags only | PERSIST | There is no fourth row. A `TaskHandlerRef` has no Dag, so no `DagImporter` — and nothing reading a `DagImporter`'s results — ever sees one. @@ -197,8 +198,9 @@ There is no fourth row. A `TaskHandlerRef` has no Dag, so no `DagImporter` — a - No importer knows about coordinators. `PythonDagImporter` is unchanged by this ADR; the stub-to-handler comparison sits in `_parse_file`, above every importer. - A mixed-language `dag_id` never appears in Dag processing results. No `Dag` registration exists for a `dag_id` a Python file already owns, so everything downstream sees exactly one record per `dag_id`, with no flag to interpret. -- Stub/implementation mismatches — missing handler, extra handler, parameter name or order, incompatible schema — surface as import errors against the Python file at parse time, - alongside the errors the parse already reports. An unannotated stub argument is checked by name and position only. +- Stub/implementation mismatches — missing handler, extra handler, a `positional` argument count that does not match, incompatible schema — surface as import errors against + the Python file at parse time, alongside the errors the parse already reports. A `named` argument or parameter that matches nothing is only logged as a warning, because the + runtime runs the task anyway. An unannotated stub argument is checked only for how it binds, by position or by name. - The Python Dag and Lang-SDK artifact can live in different DagBundles. - A single Dag can have stubs targeting different queues, some Java, some Go. Each resolves to its own coordinator instance, and validation unions their declarations per `dag_id` before comparing task ids. @@ -234,12 +236,16 @@ with one `bundle.serve()`. One artifact can hold both, so neither the file nor t ### Appendix B — What is compared -Handler declarations from every process spawned in Step 4 are unioned per `dag_id` before comparison, since one Dag's stubs can target several queues. +Handler declarations from every process spawned in Step 4 are unioned per `dag_id` before comparison, since one Dag's stubs can target several queues. Each reply covers +every Dag its artifact registers handlers for; only the parsed file's Dags are compared. - `task_id` sets must match exactly. A missing or extra handler is an error. -- `arg_bindings[*].name` against `params[*].name`, in order — both sides bind positionally. +- `arg_bindings` against `params` as each declaration's `binding` says: by position for `positional`, by name for `named`, compared case-insensitively with underscores ignored + unless `exact_name` is set. A `positional` count mismatch is an error. Under `named`, an argument or parameter that matches nothing is only logged as a warning, and a lone + unmatched argument, which may be the whole value, is not ([ADR-0012](0012-lang-sdk-parse-protocol.md) Appendix B). When `params` is `None`, only the handler's presence is + checked. - `arg_bindings[*].value_schema` against `params[*].value_schema`, compared only where neither side is null. An unannotated `@task.stub` parameter produces `null` today, so a strict comparison would make every untyped stub argument a parse error. -Any mismatch is reported against the Python file, which is the definition the author can act on, and travels back on `DagFileParsingResult.import_errors` with everything else the +Every error is reported against the Python file, which is the definition the author can act on, and travels back on `DagFileParsingResult.import_errors` with everything else the parse found. diff --git a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md index e144b87445b58..34c18508e4567 100644 --- a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md +++ b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md @@ -26,8 +26,8 @@ Proposed ## Context The Dag processor asks a Lang-SDK runtime two different questions. "Which Dags does this artifact define?" is answered over the messages [ADR-0004](0004-dag-parsing.md) already -defines. "Which task handlers does this artifact register for a `dag_id` Python already owns?" has no answer in those messages, because a `TaskHandler` registration carries no Dag -([ADR-0011](0011-mixed-language-dag-processing.md)). +defines. "Which task handlers does this artifact register, each for a `dag_id` Python already owns?" has no answer in those messages, because a `TaskHandler` registration carries +no Dag ([ADR-0011](0011-mixed-language-dag-processing.md)). This ADR defines the request that carries the second question, the subprocess classes that carry both, and the two parse-side entry points on the coordinator. @@ -82,7 +82,6 @@ relay is only type-safe because both hops speak the same pair, which is the reas ``` class TaskHandlerParseRequest: file: str # the artifact resolved for this coordinator - dag_ids: list[str] # every Dag in the parsed file with stub tasks that resolved here bundle_path: Path bundle_name: str type: Literal["TaskHandlerParseRequest"] @@ -96,20 +95,22 @@ class TaskHandlerParsingResult: class TaskHandlerDeclaration: task_id: str - params: list[TaskHandlerParam] # ordered — arg_bindings are positional + binding: Literal["positional", "named"] # how stub-task arguments bind to params + params: list[TaskHandlerParam] | None # ordered; the order matters only for "positional". None: the runtime cannot list them class TaskHandlerParam: - name: str + name: str | None # None: the runtime has no name for this positional parameter value_schema: ArgValueSchema | None = None - required: bool # the handler declares no default + exact_name: bool = False # match as spelled, not case-insensitively with underscores ignored ``` -One request carries every `dag_id` that resolved to the same artifact under the same coordinator, so a file whose stubs all target one runtime costs one process. A `dag_id` the -artifact registers nothing for is **omitted** from `task_handlers` rather than returned empty: the key set is not required to match `dag_ids`, because it is the union across -coordinators that has to cover the stubs ([ADR-0011](0011-mixed-language-dag-processing.md)). +The request names no Dags. The runtime answers with every task handler the artifact registers, keyed by `dag_id`, and with `{}` when it registers none. The answer must depend only +on the artifact, never on the request, so the Dag processor can cache it per artifact ([ADR-0013](0013-persisted-task-handler-bindings.md)). One request per (coordinator, artifact) +pair serves every Dag in the parsed file whose stubs resolved there, so a file whose stubs all target one runtime costs one process. The key set is not required to match the file's +Dags: it can hold other files' Dags, and it is the union across coordinators that has to cover the stubs ([ADR-0011](0011-mixed-language-dag-processing.md)). `value_schema` reuses the `ArgValueSchema` definition `arg_bindings` already carries ([ADR-0007](0007-taskflow-across-language-boundary.md)), so both sides of a comparison are the -same type. Two properties matter to validation: the field is nullable on both sides, and `params` is ordered. Appendix B says what that forces. +same type. Two properties matter to validation: the field is nullable on both sides, and each declaration names its binding mode. Appendix B says what that forces. `task_handlers` is the counterpart to `DagFileParsingResult.serialized_dags`, but fully typed. `serialized_dags` is `list[LazyDeserializedDAG]`, which is an opaque object in the schema snapshot. A handler declaration carries no Dag, so it code-generates and schema-validates in every SDK, and nothing on this path needs a DagSerialization implementation. @@ -129,8 +130,8 @@ WatchedSubprocess │ └── coordinator.parse_dag() — spawn runtime, forward fd 0 ⇄ comm socket │ same request and result types as its base class │ - └── SDKTaskHandlerProcessorProcess (new — ADR-0011) - target = _parse_task_handler_entrypoint + └── LangSDKTaskHandlerProcessorProcess (new — ADR-0011) + target = _start_task_handler_runtime_entrypoint └── coordinator.parse_task_handler() — same forwarding TaskHandlerParseRequest → TaskHandlerParsingResult ``` @@ -208,13 +209,23 @@ bodies that never differ. Reusing `ToManager` leaves one reply union with one ne A boolean on `DagFileParseRequest` was the other alternative. It cannot work: a `DagRef` and a `TaskHandlerRef` are different payloads, not two subsets of one, so the flag would select between shapes the result type cannot both hold. -### Appendix B — What the nullable, ordered parameter list forces +### Appendix B — What the nullable schema and the binding mode force `value_schema` is nullable on both sides. An unannotated `@task.stub` parameter produces `value_schema: null` today, so validation compares schemas only where neither side is null, and falls back to name-and-arity otherwise. A strict comparison would turn every untyped stub argument into a parse error. -`params` is ordered because `LiteralArgBinding` and `XComArgBinding` are each documented as "one positional stub-task argument". Position is part of the contract, not incidental, -and both sides bind positionally. +The stub side always has names and positions: `LiteralArgBinding` and `XComArgBinding` are each documented as "one positional stub-task argument". The handler side binds the way +its runtime does, so each declaration names its `binding` and the check follows it: + +- `positional`: by position. Names are informative only, and absent where the runtime has none (Go flat params). Java's `TaskArgs` binds this way although it has names. An + argument count that matches `params` neither with every argument nor after dropping the defaulted ones, or a value type a parameter does not accept, is an import error. +- `named`: by name in any order, case-insensitively with underscores ignored unless `exact_name` is set (Go `arg:` tags, explicit Java names). A Go struct, tagged or untagged, + and Java's `TaskInput` bind this way. An argument no parameter takes, or a parameter no argument fills, is logged as a warning, and the task still runs: the runtimes allow + both, and an unfilled field keeps its default. When no parameter matches and exactly one argument was passed, it may be the whole value and is not warned about, unless no + parameter is declared, a parameter sets `exact_name` (a Go `arg:` tag), or the argument cannot be an object. A value type a parameter does not accept is an import error. + +`params` is `None` when the runtime cannot list a handler's parameters, and then only the handler's presence is checked. TypeScript declares `named` with `params: None`: its +types are erased, so a handler cannot list what it takes. A declaration carries no class, method, or source location. [ADR-0006](0006-no-lang-sdk-source-display.md) rules out Lang-SDK source display, and putting it on the wire would invite a consumer to render it. It carries no `dag_id` either — the `task_handlers` key supplies it, so a declaration cannot disagree with the bucket it arrived in. diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index cbcace7df5231..e38c4c3c3cb58 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -21,7 +21,7 @@ ## Status -Proposed +Accepted ## Context @@ -157,31 +157,34 @@ backs many handlers and its fingerprint must have exactly one value. ```sql CREATE TABLE lang_sdk_task_handler_artifact ( - id UUID NOT NULL, - bundle_name VARCHAR(250) NOT NULL, -- the task_handler_bundle_name it was found in - relative_fileloc VARCHAR(2000) NOT NULL, -- path within that bundle - size_bytes BIGINT NOT NULL, -- cheap fingerprint tier - cache_digest VARCHAR(64) NOT NULL, -- content fingerprint tier; see "The fast path" - last_probed_at TIMESTAMP NOT NULL, - PRIMARY KEY (id), - CONSTRAINT lstha_bundle_fileloc_uq UNIQUE (bundle_name, relative_fileloc) + id UUID NOT NULL, + bundle_name VARCHAR(250) NOT NULL, -- the task_handler_bundle_name it was found in + relative_fileloc VARCHAR(2000) NOT NULL, -- path within that bundle + relative_fileloc_hash VARCHAR(32) NOT NULL, -- md5 of relative_fileloc; the path is too long to index + size_bytes BIGINT NOT NULL, -- cheap fingerprint tier + cache_digest VARCHAR(128) NULL, -- content fingerprint tier, NULL if none is stored; see "The fast path" + task_handlers JSON NOT NULL, -- {dag_id: [TaskHandlerDeclaration, ...]}, the probe answer + last_probed_at TIMESTAMP NOT NULL, + CONSTRAINT lang_sdk_task_handler_artifact_pkey PRIMARY KEY (id), + CONSTRAINT lang_sdk_task_handler_artifact_bundle_fileloc_uq UNIQUE (bundle_name, relative_fileloc_hash) ); CREATE TABLE lang_sdk_task_handler ( - dag_id VARCHAR(250) NOT NULL, - task_id VARCHAR(250) NOT NULL, - artifact_id UUID NOT NULL, - dag_bundle_name VARCHAR(250) NOT NULL, -- the *Python* file that owns this row - dag_relative_fileloc VARCHAR(2000) NOT NULL, -- ditto - handler_params JSON NOT NULL, -- list[TaskHandlerParam], ordered - PRIMARY KEY (dag_id, task_id), - CONSTRAINT lsth_dag_fkey FOREIGN KEY (dag_id) + dag_id VARCHAR(250) NOT NULL, + task_id VARCHAR(250) NOT NULL, + artifact_id UUID NOT NULL, + dag_bundle_name VARCHAR(250) NOT NULL, -- the *Python* file that owns this row + dag_relative_fileloc VARCHAR(2000) NOT NULL, -- ditto + dag_relative_fileloc_hash VARCHAR(32) NOT NULL, -- md5 of dag_relative_fileloc + CONSTRAINT lang_sdk_task_handler_pkey PRIMARY KEY (dag_id, task_id), + CONSTRAINT lang_sdk_task_handler_dag_id_fkey FOREIGN KEY (dag_id) REFERENCES dag (dag_id) ON DELETE CASCADE, - CONSTRAINT lsth_artifact_fkey FOREIGN KEY (artifact_id) + CONSTRAINT lang_sdk_task_handler_artifact_id_fkey FOREIGN KEY (artifact_id) REFERENCES lang_sdk_task_handler_artifact (id) ); -CREATE INDEX idx_lsth_dag_file ON lang_sdk_task_handler (dag_bundle_name, dag_relative_fileloc); -CREATE INDEX idx_lsth_artifact_id ON lang_sdk_task_handler (artifact_id); +CREATE INDEX idx_lang_sdk_task_handler_dag_file + ON lang_sdk_task_handler (dag_bundle_name, dag_relative_fileloc_hash); +CREATE INDEX idx_lang_sdk_task_handler_artifact_id ON lang_sdk_task_handler (artifact_id); ``` `PRIMARY KEY (dag_id, task_id)` is the conflict guard: two artifacts claiming the same task cannot @@ -194,11 +197,10 @@ joining through `DagModel`: `dag.relative_fileloc` is not indexed, and the codeb that querying it means "a sequential scan of dag" ([`reassign_dags_with_unconfigured_bundles`](https://github.com/apache/airflow/blob/79991cd4db0c9346a28b23c453377f6df0c6b4ed/airflow-core/src/airflow/dag_processing/bundles/manager.py#L478)). -`handler_params` stores what the runtime declared, so a changed Python file can be re-validated -against a cached declaration with no subprocess. It is deliberately **not** called `arg_bindings`: -that name already denotes the Python side of the comparison (`XComArgBinding` / `LiteralArgBinding`, -carrying wiring and values, [`build_arg_bindings`](https://github.com/apache/airflow/blob/79991cd4db0c9346a28b23c453377f6df0c6b4ed/airflow-core/src/airflow/serialization/stub_arg_bindings.py#L221-L286)), -and reusing it would make the validation read as comparing a thing to itself. +`task_handlers` caches the artifact's whole probe answer, every handler it registers keyed by +`dag_id`, so a changed Python file can be re-validated against it with no subprocess. It is written +only from a probe, together with `size_bytes` and `cache_digest`, so the answer always belongs to the +fingerprint beside it. A binding row is then only a mapping from a stub task to its artifact. `cache_digest` is **opaque and coordinator-defined**, not "SHA-256 of the file". @@ -208,15 +210,16 @@ and reusing it would make the validation read as comparing a thing to itself. already knows about. The parse child processor subprocesses run in the client context without a database connection ([`_parse_file_entrypoint`](https://github.com/apache/airflow/blob/79991cd4db0c9346a28b23c453377f6df0c6b4ed/airflow-core/src/airflow/dag_processing/processor.py#L208-L232)), so the prior cache state must be pushed down from the manager. The alternative is a dedicated -Execution API for the child processor process to retrieve `KnownSDKTaskHandlerArtifact` itself, which +Execution API for the child processor process to retrieve `TaskHandlerArtifact` itself, which was rejected on blast radius. ```python -class KnownSDKTaskHandlerArtifact(BaseModel): +class TaskHandlerArtifact(BaseModel): bundle_name: str relative_fileloc: str size_bytes: int - cache_digest: str + cache_digest: str | None # None: the artifact stores no fingerprint, so it is always probed + task_handlers: dict[str, list[TaskHandlerDeclaration]] # the probe answer: dag_id -> declarations class DagFileParseRequest(BaseModel): @@ -224,14 +227,11 @@ class DagFileParseRequest(BaseModel): bundle_path: Path bundle_name: str callback_requests: list[CallbackRequest] - known_artifacts: list[KnownSDKTaskHandlerArtifact] = [] # new + known_artifacts: list[TaskHandlerArtifact] = [] # new type: Literal["DagFileParseRequest"] ``` -`known_artifacts` is scoped by *artifact* bundle, not by Dag file, so the manager reads it **once per -parsing loop** for every configured `task_handler_bundle_name` and pushes the same list to every -child. Ten Dag files resolving against one twenty-jar bundle therefore probe that bundle once in -total, not once each. +`known_artifacts` is scoped by *artifact* bundle, not by Dag file, so the manager reads it **once per parsing loop** for every bundle that can hold task handlers: the `task_handler_bundle_name` of each coordinator a queue routes to, plus the Dag bundles it parses when one of them sets none and so reads the task's own Dag bundle. A coordinator no queue routes to runs no stub task, so its bundle is not read. Each child gets the rows of the named bundles whose team matches its Dag bundle's team, and of its own Dag bundle when a coordinator falls back to it, never those of another Dag bundle, which no coordinator reads for it. The team rule exists because an answer one file's parse records is trusted by every file that reads it. Teams match only when equal: without `[core] multi_team` every named bundle is in scope, and with it a team-less Dag bundle sees only team-less named bundles. One recorded answer serves every file, so ten Dag files resolving against one bundle probe a changed artifact once in total, apart from the children already started when it changed. **Dag-parsing child → coordinator subprocess.** Introduced here. The child spawns the runtime and forwards bytes in both directions, decoding nothing; the process that spawned the parse decodes the @@ -241,7 +241,6 @@ one member, and `ToManager` gains `SDKTaskHandlerParsingResult`. ```python class SDKTaskHandlerParseRequest(BaseModel): # parent -> runtime, on ToSDKTaskHandlerProcessor file: str # the candidate artifact being probed - dag_ids: list[str] # every Dag in this file with stub tasks routed here bundle_path: Path bundle_name: str type: Literal["SDKTaskHandlerParseRequest"] @@ -257,18 +256,19 @@ class SDKTaskHandlerParsingResult(BaseModel): # runtime -> parent, on ToManager class TaskHandlerDeclaration(BaseModel): task_id: str - params: list[TaskHandlerParam] # ordered; arg bindings are positional + binding: Literal["positional", "named"] # how arguments bind to params + params: list[TaskHandlerParam] | None # ordered; None: the runtime cannot list them class TaskHandlerParam(BaseModel): - name: str + name: str | None # None: the runtime has no name for this positional parameter value_schema: JSONSchema | None = None - required: bool # the handler declares no default + exact_name: bool = False # match as spelled, not case-insensitively with underscores ignored ``` -A `dag_id` the artifact registers nothing for is **omitted** from `task_handlers` rather than returned -empty, so a probe that matches nothing is distinguishable from a probe that matched a Dag with zero -tasks. +`task_handlers` maps every `dag_id` the artifact registers a handler for, and is `{}` when it registers +none. The answer must not depend on the request, because the manager records it once and serves it to +every Dag file that resolves against the artifact. **Dag-parsing child → manager.** [`DagFileParsingResult`](https://github.com/apache/airflow/blob/79991cd4db0c9346a28b23c453377f6df0c6b4ed/airflow-core/src/airflow/dag_processing/processor.py#L133-L145) gains the resolved bindings. @@ -347,12 +347,17 @@ stub Dag is validated against bindings that have not been written yet. ``` DagProcessorManager [reads DB] │ - │ once per parsing loop, per configured task_handler_bundle_name: - │ SELECT bundle_name, relative_fileloc, size_bytes, cache_digest + │ once per parsing loop, for every bundle that can hold task handlers + │ (each routed coordinator's task_handler_bundle_name, plus the Dag + │ bundles when one sets none and reads the task's own Dag bundle): + │ SELECT bundle_name, relative_fileloc, size_bytes, cache_digest, + │ task_handlers │ FROM lang_sdk_task_handler_artifact - │ WHERE bundle_name IN (:configured bundles) ──▶ known_artifacts + │ WHERE bundle_name IN (:those bundles) ──▶ known_artifacts │ - │ per file about to be dispatched: + │ per file about to be dispatched: the rows of the bundles in its scope + │ (named bundles of its team, plus its own Dag bundle when a routed + │ coordinator sets none) │ (the child re-reads nothing; it has no DB) │ ├── DagFileParseRequest(file=etl.py, bundle_*, known_artifacts=[...]) @@ -387,17 +392,18 @@ DagFileProcessorProcess(etl.py) [no DB — client con │ anything differs ──▶ probe │ ├─6─ PROBE, per differing candidate — one subprocess - │ SDKTaskHandlerProcessorProcess.start( - │ target=_parse_task_handler_entrypoint, + │ LangSDKTaskHandlerProcessorProcess.start( + │ target=_start_task_handler_runtime_entrypoint, │ coordinator=JavaCoordinator("jdk-17"), │ path=) - │ ──SDKTaskHandlerParseRequest(file=…, dag_ids=["etl"])──▶ runtime + │ ──SDKTaskHandlerParseRequest(file=…)──▶ runtime │ ◀─SDKTaskHandlerParsingResult(task_handlers={"etl": […]})── runtime │ Get* from the runtime is relayed up ToManager unchanged │ ├─7─ VALIDATE per dag_id, unioned across coordinators │ task_id sets must match exactly - │ arg_bindings[*].name ↔ handler_params[*].name, in order + │ arg_bindings[*] ↔ handler_params[*], per the declaration's binding: + │ by position, or by folded or exact name │ arg_bindings[*].schema ↔ handler_params[*].value_schema, │ compared only where neither is null │ two candidates claiming one (dag_id, task_id) → import error diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst index bd2a3e4ac761a..f35012832d6e2 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst @@ -48,7 +48,8 @@ Prerequisites * Go 1.24 or later to build and pack bundles. This is a build-time requirement only; the worker that runs a packed bundle needs no Go toolchain, because the bundle is a self-contained native executable. -* The packed bundle must be accessible from the Airflow worker, under a directory the coordinator scans. +* The packed bundle must be accessible from the Airflow worker and the Dag processor, in the Dag bundle the + coordinator scans, and built for the operating system and CPU architecture of both. * The ``apache-airflow-task-sdk`` package (installed with Airflow) provides the coordinator; no additional Python packages are needed. @@ -178,33 +179,46 @@ to pass on with ``bundle.Register(reports.Handlers()...)``. Coordinator configuration ~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Register the coordinator and route the queue to it under ``[sdk]`` in ``airflow.cfg`` (or the equivalent -``AIRFLOW__SDK__*`` environment variables): +Register a Dag bundle for the packed bundles, register the coordinator, and route the queue to it in +``airflow.cfg`` (or the equivalent ``AIRFLOW__*`` environment variables): .. code-block:: ini + [dag_processor] + dag_bundle_config_list = [ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + { + "name": "go-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/opt/airflow/go-task-handlers"} + } + ] + [sdk] coordinators = { "go": { "classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", - "kwargs": {"executables_root": ["~/airflow/executable-bundles"]} + "kwargs": {"task_handler_bundle_name": "go-task-handlers"} } } queue_to_coordinator = {"golang": "go"} -``executables_root`` is one or more directories the coordinator scans for bundles; ``queue_to_coordinator`` -routes stub tasks with ``queue="golang"`` to this Go coordinator. See :ref:`go-sdk/coordinator-config` for -the full list of accepted ``kwargs``. +``task_handler_bundle_name`` names the Dag bundle the coordinator scans for packed bundles; +``queue_to_coordinator`` routes stub tasks with ``queue="golang"`` to this Go coordinator. See +:ref:`go-sdk/coordinator-config` for the full list of accepted ``kwargs`` and how bundles are located. There is no separate Go worker to run: the Airflow worker forks the bundle binary once per task instance. .. note:: - The coordinator is part of the Airflow worker, so the ``[sdk]`` config (and the bundle files in - ``executables_root``) only need to be present wherever tasks actually execute. With ``CeleryExecutor``, - setting it on the Celery workers is sufficient. With ``LocalExecutor``, tasks run inside the scheduler - process, so it must be set where the scheduler can read it. The API server and Dag processor do not need - it. + The ``[sdk]`` config and the packed bundle files must be present wherever tasks execute and on the Dag + processor. With ``CeleryExecutor``, tasks execute on the Celery workers; with ``LocalExecutor``, they run + inside the scheduler process. The Dag processor checks the stub tasks of each Python Dag against the task + handlers the packed bundles register, so it runs them too and needs bundles built for its operating + system and CPU architecture. The API server does not need any of it. Register the Dag bundle in + ``[dag_processor] dag_bundle_config_list`` on every component, like your other Dag bundles: the worker + and the Dag processor resolve ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config + is read it is rejected if the name is missing there. Writing tasks ------------- @@ -444,24 +458,24 @@ Build and pack in one step; any flags after ``--`` are forwarded verbatim to ``g go tool airflow-go-pack ./example/bundle -- -trimpath -tags=prod -Use ``--output `` to write the packed bundle straight into a directory the coordinator scans -(``executables_root``): +Use ``--output `` to write the packed bundle straight into the directory of the Dag bundle the +coordinator scans (see `Deploying`_): .. code-block:: bash - go tool airflow-go-pack --output ~/airflow/executable-bundles/sample-dag-bundle ./example/bundle + go tool airflow-go-pack --output /opt/airflow/go-task-handlers/sample-dag-bundle ./example/bundle Cross-platform builds ~~~~~~~~~~~~~~~~~~~~~~~ -The worker that runs a bundle often uses a different operating system or CPU architecture than your build -machine (for example, deploying to a Linux host from an Apple-silicon ``darwin/arm64`` laptop). Pass -``--goos`` / ``--goarch`` and the packer cross-builds for you: +The worker and the Dag processor that run a bundle often use a different operating system or CPU +architecture than your build machine (for example, deploying to a Linux host from an Apple-silicon +``darwin/arm64`` laptop). Pass ``--goos`` / ``--goarch`` and the packer cross-builds for you: .. code-block:: bash go tool airflow-go-pack --goos linux --goarch amd64 \ - --output ~/airflow/executable-bundles/sample-dag-bundle \ + --output /opt/airflow/go-task-handlers/sample-dag-bundle \ ./example/bundle Alternatively, pack a pre-built binary with ``--executable`` / ``--source``. The packer normally execs the @@ -485,12 +499,16 @@ with ``--airflow-metadata``: Deploying ~~~~~~~~~ -Copy or mount the packed bundle into a directory listed in the coordinator's ``executables_root``. The -:class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` scans those directories recursively, +Copy or mount the packed bundle into the Dag bundle named by the coordinator's ``task_handler_bundle_name``. +The :class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` scans that Dag bundle recursively, matches the incoming ``dag_id`` against each bundle's manifest, verifies the bundle's integrity hash, and launches the matching bundle. Bundles are identified by the trailer magic, not by filename (no extension on Linux/macOS, ``.exe`` on Windows), so the file name on the worker is irrelevant. +Only files with the executable bit set are considered, so the packed bundles need a Dag bundle that keeps +it. A ``LocalDagBundle`` over a directory you manage does; object-store Dag bundles such as ``S3DagBundle`` +do not. + .. _go-sdk/coordinator-config: :class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` configuration @@ -506,16 +524,13 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the * - Parameter - Default - Description - * - ``executables_root`` - - *(optional)* - - One or more directories scanned recursively for executable bundles. Accepts a string, - a path, or a list of strings/paths. When omitted, bundles are located through a Dag - bundle instead (see the note below). Explicitly setting this option to ``null`` or - an empty list is invalid. - * - ``dag_bundle_name`` - - *(auto: task's own bundle)* - - Name of a configured Dag bundle to load executable bundles from. Mutually exclusive - with ``executables_root``. + * - ``task_handler_bundle_name`` + - *(task's own Dag bundle)* + - Name of the Dag bundle scanned recursively for executable bundles. It is used only by + mixed-language Dags, to locate the task handlers for the ``@task.stub`` tasks of a Python Dag; + Dags defined natively in a language SDK do not use it. It must be registered in + ``[dag_processor] dag_bundle_config_list``. It is checked when the ``[sdk]`` configuration is + loaded, so a typo fails there rather than on the first task. * - ``task_startup_timeout`` - ``10.0`` - Seconds to wait for the bundle subprocess to connect after launch. Increase this if your @@ -523,15 +538,15 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the .. note:: - **Locating bundles.** ``executables_root`` and ``dag_bundle_name`` are mutually exclusive, - and both are optional: + **Locating bundles.** The packed bundles for the ``@task.stub`` tasks of a Python Dag are read from a Dag + bundle, so they are delivered, refreshed and versioned by the same machinery as your Dags. - * Set ``executables_root`` to scan explicit filesystem directories you manage yourself. - * Set ``dag_bundle_name`` to load bundles from a configured Dag bundle, so they are delivered - and versioned through the same bundle machinery as your Dags. The task uses the version that - bundle is on when it starts, pinned for the whole task. - * Leave both unset (the default) to load bundles from the **task's own** Dag bundle, pinned - to the version the run was created with. + * The expected layout is a separate Dag bundle for the packed bundles, named by + ``task_handler_bundle_name``, rather than the Dag bundle that holds your ``.py`` files. The task + uses the version that Dag bundle is on when it starts, pinned for the whole task. + * If ``task_handler_bundle_name`` is unset, packed bundles are read from the **task's own** Dag + bundle, pinned to the version the run was created with. + * The Dag bundle must keep the executable bit (see `Deploying`_). .. _go-sdk/limitations: diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/index.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/index.rst index 79a0cb3df1f4d..159911d84c9f3 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/index.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/index.rst @@ -155,9 +155,9 @@ Coordinators are registered in ``airflow.cfg`` (or via environment variables) un } } - The ``classpath`` value must be importable by the worker. The ``kwargs`` are passed directly - to the coordinator's constructor. See the language-specific guide for the accepted kwargs - of each coordinator (e.g. :ref:`java-sdk/coordinator-config` for + The ``classpath`` value must be importable by the worker and the Dag processor. The ``kwargs`` + are passed directly to the coordinator's constructor. See the language-specific guide for the + accepted kwargs of each coordinator (e.g. :ref:`java-sdk/coordinator-config` for :class:`~airflow.sdk.coordinators.java.JavaCoordinator`). ``extra`` is an optional object for any additional information you want to associate with a diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index 3f69772e4d92d..eafd542df6955 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -41,8 +41,8 @@ Prerequisites * JDK 11 or later is required on the machine that builds the Java project. A local Gradle installation is only needed to generate the Gradle Wrapper for a new project. -* JRE 11 or later must be available on the Airflow worker nodes. -* The compiled task JAR(s) and JVM dependencies must be accessible from the worker. +* JRE 11 or later must be available on the Airflow worker nodes and the Dag processor. +* The compiled task JAR(s) and JVM dependencies must be accessible from the worker and the Dag processor. * The ``apache-airflow-task-sdk`` package (installed with Airflow) provides the coordinator; no additional Python packages are needed. @@ -56,7 +56,8 @@ deployment process are the same. See :ref:`java-sdk/interface-api` for the inter The Python Dag source and the Java Gradle project are independent. They do not need to be in the same repository or have any particular relative filesystem layout. The Dag follows the deployment's normal Dag -delivery process; only the compiled Java bundle is deployed from the Gradle project to ``jars_root``. +delivery process; only the compiled Java bundle is deployed from the Gradle project, into a separate Dag +bundle that holds the JARs. Define the Python Dag ~~~~~~~~~~~~~~~~~~~~~ @@ -245,29 +246,45 @@ Deploy ``sales_pipeline.py`` separately through the deployment's normal Dag deli that process might sync it to ``${AIRFLOW_HOME}/dags/`` or package it in a Dag bundle; neither location is inside or relative to ``sales-pipeline-java/``. -Configure Airflow so the coordinator scans the parent JAR directory recursively and routes the ``java`` queue -to it. Add the following ``[sdk]`` section to the file selected by ``AIRFLOW_CONFIG`` (by default, -``${AIRFLOW_HOME}/airflow.cfg``), or set the equivalent ``AIRFLOW__SDK__*`` environment variables: +Configure Airflow to register the JAR directory as a Dag bundle, point the coordinator at it, and route the +``java`` queue to the coordinator. Add the following sections to the file selected by ``AIRFLOW_CONFIG`` (by +default, ``${AIRFLOW_HOME}/airflow.cfg``), or set the equivalent ``AIRFLOW__*`` environment variables: .. code-block:: ini + [dag_processor] + dag_bundle_config_list = [ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + { + "name": "java-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/opt/airflow/jars/sales-pipeline"} + } + ] + [sdk] coordinators = { "java": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", - "kwargs": {"jars_root": ["/opt/airflow/jars"]} + "kwargs": {"task_handler_bundle_name": "java-task-handlers"} } } queue_to_coordinator = {"java": "java"} ``java`` is a user-chosen coordinator name, not a reserved value. The value assigned to the queue in -``queue_to_coordinator`` must match a key in ``coordinators``. - -Restart the affected Airflow components after changing this configuration. The coordinator config and JARs -must be available wherever tasks execute. With ``CeleryExecutor``, that means the Celery workers; with -``LocalExecutor``, tasks run in subprocesses on the scheduler's host. The API server and Dag processor do not -need the JARs, while the Dag processor must receive ``sales_pipeline.py`` through the separate Dag delivery -process. +``queue_to_coordinator`` must match a key in ``coordinators``, and ``task_handler_bundle_name`` must match a +Dag bundle name in ``dag_bundle_config_list``. See :ref:`java-sdk/coordinator-config` for how JARs are +located. + +Restart the affected Airflow components after changing this configuration. The coordinator config, the JARs +and a JRE must be available wherever tasks execute and on the Dag processor. With ``CeleryExecutor``, tasks +execute on the Celery workers; with ``LocalExecutor``, they run in subprocesses on the scheduler's host. The +Dag processor checks the stub tasks of ``sales_pipeline.py`` against the task handlers the JARs register, so +it runs them too. The API server does not need any of it. Register the Dag bundle in +``[dag_processor] dag_bundle_config_list`` on every component, like your other Dag bundles: the worker and +the Dag processor resolve ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config is read +it is rejected if the name is missing there. The Dag processor still receives ``sales_pipeline.py`` through +the separate Dag delivery process. After Airflow has parsed the Dag, trigger it from the UI or command line: @@ -726,7 +743,7 @@ configuration): "java-jdk17": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": { - "jars_root": ["/opt/airflow/jars"], + "task_handler_bundle_name": "java-task-handlers", "jvm_args": ["-Djava.util.logging.config.file=/opt/airflow/logging.properties"] } } @@ -851,9 +868,10 @@ Then run: ./gradlew bundle -The ``build/bundle/`` directory contains all required JAR(s). Copy or mount it into the directory pointed to -by ``jars_root`` in the coordinator configuration. :class:`~airflow.sdk.coordinators.java.JavaCoordinator` -scans ``jars_root`` recursively and builds the classpath automatically. +The ``build/bundle/`` directory contains all required JAR(s). Copy or mount it into the Dag bundle named by +``task_handler_bundle_name`` in the coordinator configuration. +:class:`~airflow.sdk.coordinators.java.JavaCoordinator` scans that Dag bundle recursively and builds the +classpath automatically. .. note:: @@ -865,8 +883,7 @@ scans ``jars_root`` recursively and builds the classpath automatically. The plugin generates a fat JAR with the `Shadow `__ plugin by default. This is generally a good idea since you only deploy one JAR file to avoid dependency issues between projects. If this does not suit you, set ``fatJar = false`` in ``airflowBundle`` to produce thin JARs instead. The rest of the - process stays the same, but you will need to put all dependency JARs somewhere Airflow can find with - ``jars_root``. + process stays the same, but you will need to put all dependency JARs in the same Dag bundle. .. _java-sdk/build/maven: @@ -956,8 +973,8 @@ Then run: mvn package -The fat JAR is written to ``target/-.jar``. Copy it to the directory configured as -``jars_root`` in your coordinator. +The fat JAR is written to ``target/-.jar``. Copy it into the Dag bundle named by +``task_handler_bundle_name`` in your coordinator. **Option 2: thin JAR with separate dependencies** @@ -1017,8 +1034,8 @@ Then run: mvn package -``target/bundle/`` will contain the thin JAR and all runtime dependency JARs. Point ``jars_root`` at -this directory. +``target/bundle/`` will contain the thin JAR and all runtime dependency JARs. Copy or mount this +directory into the Dag bundle named by ``task_handler_bundle_name``. .. note:: @@ -1044,16 +1061,13 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the * - Parameter - Default - Description - * - ``jars_root`` - - *(optional)* - - One or more directories scanned recursively for ``.jar`` files. Accepts a string, - a path, or a list of strings/paths. When omitted, JARs are located through a Dag - bundle instead (see the note below). Explicitly setting this option to ``null`` or - an empty list is invalid. - * - ``dag_bundle_name`` - - *(auto: task's own bundle)* - - Name of a configured Dag bundle to load JARs from. Mutually exclusive with - ``jars_root``. + * - ``task_handler_bundle_name`` + - *(task's own Dag bundle)* + - Name of the Dag bundle scanned recursively for ``.jar`` files. It is used only by + mixed-language Dags, to locate the task handlers for the ``@task.stub`` tasks of a Python Dag; + Dags defined natively in a language SDK do not use it. It must be registered in + ``[dag_processor] dag_bundle_config_list``. It is checked when the ``[sdk]`` configuration is + loaded, so a typo fails there rather than on the first task. * - ``java_executable`` - ``"java"`` - Path to the ``java`` binary. Defaults to ``java`` on ``$PATH``. @@ -1062,10 +1076,9 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Extra JVM arguments such as ``["-Xmx1g", "-Dsome.property=value"]``. * - ``main_class`` - *(auto-detect)* - - Explicit entry-point class. If omitted, the coordinator scans for a JAR whose - manifest sets ``Main-Class`` — in ``jars_root`` when set, otherwise across the - resolved Dag bundle. If multiple executable JARs match the result is - non-deterministic; set ``main_class`` explicitly in that case. + - Explicit entry-point class. If omitted, the coordinator scans the Dag bundle for a JAR + whose manifest sets ``Main-Class``. If more than one JAR in that Dag bundle sets it, which + one runs is non-deterministic, so set ``main_class`` explicitly in that case. * - ``task_startup_timeout`` - ``10.0`` - Seconds to wait for the JVM subprocess to connect after launch. Increase this if your @@ -1073,22 +1086,25 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the .. note:: - **Locating JARs.** ``jars_root`` and ``dag_bundle_name`` are mutually exclusive, and both - are optional: + **Locating JARs.** The JARs for the ``@task.stub`` tasks of a Python Dag are read from a Dag bundle, so + they are delivered, refreshed and versioned by the same machinery as your Dags. - * Set ``jars_root`` to scan explicit filesystem directories you manage yourself. - * Set ``dag_bundle_name`` to load JARs from a configured Dag bundle, so they are delivered and - versioned through the same bundle machinery as your Dags. The task uses the version that bundle + * The expected layout is a separate Dag bundle for the JARs, named by ``task_handler_bundle_name``, + rather than the Dag bundle that holds your ``.py`` files. The task uses the version that Dag bundle is on when it starts, pinned for the whole task. - * Leave both unset (the default) to load JARs from the **task's own** Dag bundle, pinned to - the version the run was created with. + * If ``task_handler_bundle_name`` is unset, JARs are read from the **task's own** Dag bundle, pinned + to the version the run was created with. + * Every JAR in the Dag bundle goes on one classpath, so all handlers in it share one set of + dependencies. To isolate conflicting dependency versions, put the handlers in a second Dag bundle + served by a second coordinator on its own queue. .. note:: The ``[sdk]`` configuration is read at startup, so changes to ``coordinators`` or ``queue_to_coordinator`` (for example adding ``jvm_args``) only take effect after you restart the - scheduler (or ``airflow standalone``). A rebuilt bundle JAR, by contrast, is picked up on the next - task launch without a restart, because a fresh JVM is spawned per task instance. + components that read it: the workers (the scheduler with ``LocalExecutor``), the Dag processor, or + ``airflow standalone``. A rebuilt bundle JAR, by contrast, is picked up on the next task launch without a + restart, because a fresh JVM is spawned per task instance. .. _java-sdk/java-executable: @@ -1110,7 +1126,7 @@ point ``java_executable`` at it explicitly: "java-jdk17": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": { - "jars_root": ["/opt/airflow/jars"], + "task_handler_bundle_name": "java-task-handlers", "java_executable": "/opt/homebrew/opt/openjdk@17/bin/java" } } diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst index 47070c581da15..9364305ee2ddd 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -48,9 +48,9 @@ The SDK is the ``apache-airflow-ts-sdk`` package (ESM-only). It is currently in Prerequisites ------------- -* Node.js 22 or later must be available on the Airflow worker nodes. +* Node.js 22 or later must be available on the Airflow worker nodes and the Dag processor. * The packed bundle (a single ``bundle.min.mjs`` file, see :ref:`typescript-sdk/build`) must be accessible - from the worker, under a directory the coordinator scans. + from the worker and the Dag processor, in the Dag bundle the coordinator scans. * The ``apache-airflow-task-sdk`` package (installed with Airflow) provides the coordinator; no additional Python packages are needed. * In the TypeScript project, install the ``apache-airflow-ts-sdk`` npm package to author task handlers: @@ -209,34 +209,48 @@ check. Coordinator configuration ~~~~~~~~~~~~~~~~~~~~~~~~~ -Register the coordinator and route the queue to it under ``[sdk]`` in ``airflow.cfg`` (or the equivalent -``AIRFLOW__SDK__*`` environment variables): +Register a Dag bundle for the packed bundles, register the coordinator, and route the queue to it in +``airflow.cfg`` (or the equivalent ``AIRFLOW__*`` environment variables): .. code-block:: ini + [dag_processor] + dag_bundle_config_list = [ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + { + "name": "ts-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/opt/airflow/ts-bundles"} + } + ] + [sdk] coordinators = { "ts": { "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", - "kwargs": {"bundles_root": ["/opt/airflow/ts-bundles"]} + "kwargs": {"task_handler_bundle_name": "ts-task-handlers"} } } queue_to_coordinator = {"typescript": "ts"} -``bundles_root`` is one or more directories the coordinator scans for bundles; ``queue_to_coordinator`` -routes stub tasks with ``queue="typescript"`` to this coordinator. See -:ref:`typescript-sdk/coordinator-config` for the full list of accepted ``kwargs``. +``task_handler_bundle_name`` names the Dag bundle the coordinator scans for packed bundles; +``queue_to_coordinator`` routes stub tasks with ``queue="typescript"`` to this coordinator. See +:ref:`typescript-sdk/coordinator-config` for the full list of accepted ``kwargs`` and how the bundle is +located. There is no separate Node.js worker to run: the Airflow worker launches the bundle with ``node`` once per task instance. .. note:: - The coordinator runs inside the Airflow worker, so the ``[sdk]`` config (and the packed ``*.min.mjs`` - bundles in ``bundles_root``) only need to be present wherever tasks actually execute. With - ``CeleryExecutor``, setting them on the Celery workers is sufficient. With ``LocalExecutor``, tasks run - inside the scheduler process, so they must be present where the scheduler can read them. The API server - and Dag processor do not need them. + The ``[sdk]`` config, the packed ``*.min.mjs`` bundles and Node.js must be present wherever tasks execute + and on the Dag processor. With ``CeleryExecutor``, tasks execute on the Celery workers; with + ``LocalExecutor``, they run inside the scheduler process. The Dag processor checks the stub tasks of each + Python Dag against the task handlers the packed bundles register, so it runs them too. The API server + does not need any of it. Register the Dag bundle in ``[dag_processor] dag_bundle_config_list`` on every + component, like your other Dag bundles: the worker and the Dag processor resolve + ``task_handler_bundle_name`` through it, and wherever the ``[sdk]`` config is read it is rejected if the + name is missing there. .. _typescript-sdk/native-dag: @@ -499,7 +513,7 @@ that only supply utilities or types are not embedded. npx airflow-ts-pack src/main.ts --outdir dist Use ``--outdir `` to choose the output directory (default ``dist``), ``--outfile `` to name the -artifact exactly, which helps when one ``bundles_root`` holds several bundles, and ``--source `` to set +artifact exactly, which helps when one Dag bundle holds several bundles, and ``--source `` to set the source name displayed in the Airflow UI (default: the entry file's basename). ``--outdir`` and ``--outfile`` are mutually exclusive, and an ``--outfile`` name must end in ``.min.mjs`` so the coordinator can find it. @@ -507,12 +521,11 @@ can find it. Deploying ~~~~~~~~~ -Copy or mount the bundle into a directory listed in the coordinator's ``bundles_root``. -:class:`~airflow.sdk.coordinators.node.NodeCoordinator` searches the configured directories in order, -recursively, and launches the first integrity-verified ``*.min.mjs`` bundle whose metadata declares the task -instance's Dag. The artifact's name does not matter beyond that suffix, so one root can hold several bundles -and a Dag is routed to whichever declares it. If multiple bundles declare the same Dag, the first configured -root wins, and within a root the first in sorted path order. +Copy or mount the bundle into the Dag bundle named by the coordinator's ``task_handler_bundle_name``. +:class:`~airflow.sdk.coordinators.node.NodeCoordinator` searches that Dag bundle recursively and launches the +first integrity-verified ``*.min.mjs`` bundle whose metadata declares the task instance's Dag. The artifact's +name does not matter beyond that suffix, so one Dag bundle can hold several bundles and a Dag is routed to +whichever declares it. If multiple bundles declare the same Dag, the first in sorted path order wins. .. _typescript-sdk/coordinator-config: @@ -529,16 +542,13 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the * - Parameter - Default - Description - * - ``bundles_root`` - - *(optional)* - - One or more directories searched recursively, in order, for an integrity-verified ``*.min.mjs`` - bundle that declares the requested Dag. Accepts a string, a path, or a list of strings/paths. When - omitted, the bundle is located through a Dag bundle instead (see the note below). Explicitly setting - this option to ``null`` or an empty list is invalid. - * - ``dag_bundle_name`` - - *(auto: task's own bundle)* - - Name of a configured Dag bundle to load the ``*.min.mjs`` bundle from. Mutually exclusive with - ``bundles_root``. + * - ``task_handler_bundle_name`` + - *(task's own Dag bundle)* + - Name of the Dag bundle searched recursively for an integrity-verified ``*.min.mjs`` bundle that + declares the requested Dag. It is used only by mixed-language Dags, to locate the task handlers for + the ``@task.stub`` tasks of a Python Dag; Dags defined natively in a language SDK do not use it. It + must be registered in ``[dag_processor] dag_bundle_config_list``. It is checked when the ``[sdk]`` + configuration is loaded, so a typo fails there rather than on the first task. * - ``node_executable`` - ``"node"`` - Path to the ``node`` binary. Defaults to ``node`` on ``$PATH``. @@ -549,15 +559,14 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the .. note:: - **Locating the bundle.** ``bundles_root`` and ``dag_bundle_name`` are mutually exclusive, and both - are optional: + **Locating the bundle.** The packed bundles for the ``@task.stub`` tasks of a Python Dag are read from a + Dag bundle, so they are delivered, refreshed and versioned by the same machinery as your Dags. - * Set ``bundles_root`` to scan explicit filesystem directories you manage yourself. - * Set ``dag_bundle_name`` to load the bundle from a configured Dag bundle, so it is delivered - and versioned through the same bundle machinery as your Dags. The task uses the version that - bundle is on when it starts, pinned for the whole task. - * Leave both unset (the default) to load the bundle from the **task's own** Dag bundle, pinned to the - version the run was created with. + * The expected layout is a separate Dag bundle for the packed bundles, named by + ``task_handler_bundle_name``, rather than the Dag bundle that holds your ``.py`` files. The task uses + the version that Dag bundle is on when it starts, pinned for the whole task. + * If ``task_handler_bundle_name`` is unset, the bundle is read from the **task's own** Dag bundle, pinned + to the version the run was created with. Limitations ----------- diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index afd3ed01e8c32..70847aee7d4d3 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,10 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``90e4d18ccadf`` (head) | ``e5a91c7f42b3`` | ``3.4.0`` | Add timetable_asset_gated to DagModel. | +| ``f7ed13533d23`` (head) | ``90e4d18ccadf`` | ``3.4.0`` | Add lang_sdk_task_handler_artifact and lang_sdk_task_handler | +| | | | tables. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``90e4d18ccadf`` | ``e5a91c7f42b3`` | ``3.4.0`` | Add timetable_asset_gated to DagModel. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``e5a91c7f42b3`` | ``ca8499dc1004`` | ``3.4.0`` | Add language column to dag_code. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/newsfragments/73970.significant.rst b/airflow-core/newsfragments/73970.significant.rst new file mode 100644 index 0000000000000..886d1d6ca0a8f --- /dev/null +++ b/airflow-core/newsfragments/73970.significant.rst @@ -0,0 +1,5 @@ +Language SDK coordinators load task handlers from a Dag bundle + +The experimental ``jars_root`` (Java) and ``executables_root`` (Go) coordinator kwargs are removed. +Put the artifacts in a Dag bundle registered in ``[dag_processor] dag_bundle_config_list`` and set ``task_handler_bundle_name`` in the coordinator's ``kwargs``, or leave it unset to use the task's own Dag bundle. +A config that still sets a removed kwarg fails with ``Cannot instantiate coordinator ''``. diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 34849d5eb5e13..eafaecdda739b 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2191,6 +2191,19 @@ sdk: example, two ``JavaCoordinator`` instances pinned to different JDK versions). + The language SDK coordinators read the task handlers for the + ``@task.stub`` tasks of a Python Dag from the Dag bundle named by the + ``task_handler_bundle_name`` kwarg. Dags defined natively in a language + SDK do not use it. It must name a bundle in + ``[dag_processor] dag_bundle_config_list`` and is checked when this + option is loaded. When it is unset, a task reads its handlers from its + own Dag bundle. + + The Dag processor needs this option too, with the language runtime the + coordinator starts and the files of that Dag bundle, because it checks + the stub tasks of each Python Dag against the task handlers those files + register. The API server does not need it. + An entry may also carry an optional ``extra`` object for additional information associated with the coordinator that the coordinator itself does not receive; other components read it as needed. For example, @@ -2205,7 +2218,7 @@ sdk: "jdk-17": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": { - "jars_root": ["/opt/airflow/java-bundles"], + "task_handler_bundle_name": "java-task-handlers", "java_executable": "/usr/lib/jvm/java-17-openjdk/bin/java", "jvm_args": ["-Xmx1024m"] }, @@ -2218,7 +2231,7 @@ sdk: "go-sdk": { "classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": { - "executables_root": ["/opt/airflow/executable-bundles"] + "task_handler_bundle_name": "go-task-handlers" } } } diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 8c4fd0588eb87..fa8705493215c 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -39,6 +39,7 @@ import attrs import structlog +from pydantic import ValidationError from sqlalchemy import select, update from sqlalchemy.exc import OperationalError from sqlalchemy.orm import load_only @@ -55,7 +56,11 @@ ) from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.dag_processing.collection import update_dag_parsing_results_in_db -from airflow.dag_processing.processor import DagFileParsingResult, DagFileProcessorProcess +from airflow.dag_processing.processor import ( + DagFileParsingResult, + DagFileProcessorProcess, + TaskHandlerArtifact, +) from airflow.models.asset import remove_references_to_deleted_dags from airflow.models.dag import DagModel from airflow.models.dagbag import DagPriorityParsingRequest @@ -63,8 +68,10 @@ from airflow.models.dagwarning import DagWarning from airflow.models.db_callback_request import DbCallbackRequest from airflow.models.errors import ParseImportError +from airflow.models.lang_sdk_task_handler import LangSDKTaskHandlerArtifact from airflow.observability.metrics import stats_utils from airflow.sdk import SecretCache +from airflow.sdk.execution_time.coordinator import get_coordinator_manager # noqa: SDK001 from airflow.sdk.log import init_log_file, logging_processors from airflow.typing_compat import assert_never from airflow.utils.file import list_py_file_paths, might_contain_dag @@ -151,6 +158,16 @@ def normalized_file_path_for_stats(self) -> str: return normalize_name_for_stats(str(self.rel_path), log_warning=False) +class _TaskHandlerBundles(NamedTuple): + """Where the coordinators that queues route to read task-handler artifacts from.""" + + named: frozenset[str] + """Bundles named by a coordinator's ``task_handler_bundle_name``.""" + + include_dag_bundles: bool + """Whether a coordinator without ``task_handler_bundle_name`` reads the task's own Dag bundle.""" + + def _config_int_factory(section: str, key: str): return functools.partial(conf.getint, section, key) @@ -1444,7 +1461,97 @@ def client(self) -> Client: client.base_url = "http://in-process.invalid./" return client - def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess: + @functools.cached_property + def _task_handler_bundles(self) -> _TaskHandlerBundles: + try: + bundle_names = get_coordinator_manager().get_task_handler_bundle_names().values() + except Exception: + # A Lang-SDK configuration error must not stop Python-only parsing. + self.log.exception( + "Cannot read [sdk] coordinators; Dag files are parsed without known task-handler artifacts" + ) + return _TaskHandlerBundles(named=frozenset(), include_dag_bundles=False) + return _TaskHandlerBundles( + named=frozenset(name for name in bundle_names if name is not None), + include_dag_bundles=None in bundle_names, + ) + + def get_known_task_handler_artifacts( + self, bundle_names: Collection[str] + ) -> dict[str, list[TaskHandlerArtifact]]: + """ + Return the recorded task-handler artifacts in *bundle_names*, keyed by bundle name. + + Default implementation reads from the metadata DB; override to source them from an API. + """ + return self._get_known_task_handler_artifacts_from_db(bundle_names) + + @provide_session + def _get_known_task_handler_artifacts_from_db( + self, bundle_names: Collection[str], *, session: Session = NEW_SESSION + ) -> dict[str, list[TaskHandlerArtifact]]: + rows = session.execute( + select( + LangSDKTaskHandlerArtifact.bundle_name, + LangSDKTaskHandlerArtifact.relative_fileloc, + LangSDKTaskHandlerArtifact.size_bytes, + LangSDKTaskHandlerArtifact.cache_digest, + LangSDKTaskHandlerArtifact.task_handlers, + ) + .where(LangSDKTaskHandlerArtifact.bundle_name.in_(bundle_names)) + .order_by(LangSDKTaskHandlerArtifact.bundle_name, LangSDKTaskHandlerArtifact.relative_fileloc) + ) + known: dict[str, list[TaskHandlerArtifact]] = defaultdict(list) + for row in rows: + try: + artifact = TaskHandlerArtifact.model_validate(row._asdict()) + except ValidationError as exc: + # Left out, the artifact is probed again and its new answer replaces the row. + self.log.warning( + "Ignoring a recorded task-handler artifact that fails validation", + bundle_name=row.bundle_name, + relative_fileloc=row.relative_fileloc, + error=str(exc), + ) + continue + known[row.bundle_name].append(artifact) + return dict(known) + + def _query_known_task_handler_artifacts(self) -> dict[str, list[TaskHandlerArtifact]]: + """Return the recorded artifacts of every bundle that can hold task handlers for this manager's Dag files.""" + bundle_names = set(self._task_handler_bundles.named) + if self._task_handler_bundles.include_dag_bundles: + bundle_names.update(bundle.name for bundle in self._dag_bundles) + if not bundle_names: + return {} + try: + return self.get_known_task_handler_artifacts(bundle_names) + except Exception: + # Without known artifacts the children probe their candidates, so parsing keeps running. + self.log.exception( + "Cannot read the recorded task-handler artifacts; Dag files are parsed without them" + ) + return {} + + def _get_task_handler_artifact_bundle_names(self, dag_bundle_name: str) -> frozenset[str]: + """Return the bundles whose task-handler artifacts a Dag file in *dag_bundle_name* may read and record.""" + bundles = self._task_handler_bundles + teams = self._get_team_names(bundles.named | {dag_bundle_name}) + scope = {name for name in bundles.named if teams.get(name) == teams.get(dag_bundle_name)} + if bundles.include_dag_bundles: + scope.add(dag_bundle_name) + return frozenset(scope) + + def _select_known_task_handler_artifacts( + self, known_artifacts: dict[str, list[TaskHandlerArtifact]], dag_bundle_name: str + ) -> list[TaskHandlerArtifact]: + """Return the artifacts a Dag file in *dag_bundle_name* may resolve its stub tasks against.""" + bundle_names = self._get_task_handler_artifact_bundle_names(dag_bundle_name) + return [artifact for name in sorted(bundle_names) for artifact in known_artifacts.get(name, ())] + + def _create_process( + self, dag_file: DagFileInfo, *, known_artifacts: Sequence[TaskHandlerArtifact] = () + ) -> DagFileProcessorProcess: id = uuid7() callback_to_execute_for_file = self._callback_to_execute.pop(dag_file, []) @@ -1457,6 +1564,7 @@ def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess: bundle_name=dag_file.bundle_name, dag_file_rel_path=str(dag_file.rel_path), callbacks=callback_to_execute_for_file, + known_artifacts=known_artifacts, selector=self.selector, logger=logger, logger_filehandle=logger_filehandle, @@ -1467,6 +1575,9 @@ def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess: def _start_new_processes(self): """Start more processors if we have enough slots and files to process.""" bundle_to_team = self._get_team_names({file.bundle_name for file in self._file_queue}) + # Read once per loop, and only when a child starts: rows persisted by earlier children are + # then visible to later ones, which a cache held across a whole pass would hide. + known_artifacts: dict[str, list[TaskHandlerArtifact]] | None = None while self._parallelism > len(self._processors) and self._file_queue: file, _ = self._file_queue.popitem(last=False) @@ -1474,7 +1585,12 @@ def _start_new_processes(self): if file in self._processors: continue - processor = self._create_process(file) + if known_artifacts is None: + known_artifacts = self._query_known_task_handler_artifacts() + processor = self._create_process( + file, + known_artifacts=self._select_known_task_handler_artifacts(known_artifacts, file.bundle_name), + ) stats.incr( "dag_processing.processes", tags=prune_dict( diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index 80e6d7161dec2..408d108869437 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -23,12 +23,13 @@ import traceback from collections.abc import Callable, Sequence from pathlib import Path -from typing import TYPE_CHECKING, Annotated, BinaryIO, ClassVar, Literal +from typing import TYPE_CHECKING, Annotated, Any, BinaryIO, ClassVar, Literal import attrs from pydantic import BaseModel, Field, TypeAdapter from airflow._shared.observability.metrics import stats +from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import ArgValueSchema # noqa: TC001 from airflow.callbacks.callback_requests import ( CallbackRequest, DagCallbackRequest, @@ -92,10 +93,90 @@ from airflow.sdk.definitions.context import Context from airflow.sdk.definitions.dag import DAG from airflow.sdk.definitions.mappedoperator import MappedOperator - from airflow.sdk.execution_time.supervisor import RequestHandler, RequestResult + from airflow.sdk.execution_time.supervisor import RequestHandler, RequestResult, ResponseSent from airflow.typing_compat import Self +class DagFileParsingResult(BaseModel): + """ + Result of DAG File Parsing. + + This is the result of a successful DAG parse, in this class, we gather all serialized DAGs, + import errors and warnings to send back to the scheduler to store in the DB. + """ + + fileloc: str + serialized_dags: list[LazyDeserializedDAG] + warnings: list | None = None + import_errors: dict[str, str] | None = None + parsed_definitions: list[str] = Field(default_factory=list) + """Bundle-relative locations of the Dag definitions imported from ``fileloc``.""" + dag_source_codes: dict[str, DagSourceCode] = Field(default_factory=dict) + """Source code of the parsed Dags, keyed by Dag fileloc.""" + type: Literal["DagFileParsingResult"] = "DagFileParsingResult" + + +TaskHandlerBindingMode = Literal["positional", "named"] + + +class TaskHandlerParam(BaseModel): + """One parameter of a task handler.""" + + name: str | None + """``None`` when the runtime has no name for this positional parameter.""" + + value_schema: ArgValueSchema | None = None + """JSON Schema of the values the parameter accepts; ``None`` when the handler does not constrain it.""" + + exact_name: bool = False + """Whether ``name`` matches only as spelled, not case-insensitively with underscores ignored.""" + + +class TaskHandlerDeclaration(BaseModel): + """A task handler that a Lang-SDK artifact registers for one task.""" + + task_id: str + + binding: TaskHandlerBindingMode + """ + How stub-task arguments bind to ``params``. + + - ``positional``: by position; names are informative only. An argument count that matches ``params`` + neither with every argument nor after dropping the defaulted ones, or a value type the param does not + accept, makes the Dag fail to import. + - ``named``: by name in any order, case-insensitively with underscores ignored unless ``exact_name`` is set. + An argument no param takes, or a param no argument fills, is logged as a warning, and the task still + runs. When no param matches and exactly one argument was passed, it may be the whole value and is not + warned about, unless ``params`` is empty, a param sets ``exact_name``, or the argument cannot be an + object. A value type a param does not accept makes the Dag fail to import. + """ + + # The title keeps Go's generated type for this list apart from the HITL ``Params`` map. + params: Annotated[list[TaskHandlerParam] | None, Field(title="Task Handler Params")] + """ + In declaration order; the order is significant only for ``positional`` binding. + + ``None`` when the runtime cannot list the handler's parameters, so only the handler's presence is checked. + """ + + +class TaskHandlerArtifact(BaseModel): + """A Lang-SDK artifact and every task handler it registers.""" + + bundle_name: str + + relative_fileloc: str = Field(max_length=2000) + """Path of the artifact within its bundle.""" + + size_bytes: int + + cache_digest: str | None = Field(max_length=128) + """Opaque content fingerprint defined by the coordinator; ``None`` when the artifact stores none, so it is always probed.""" + + task_handlers: dict[str, list[TaskHandlerDeclaration]] + """Every Dag id the artifact registers a task handler for, with those handlers.""" + + class DagFileParseRequest(BaseModel): """ Request for DAG File Parsing. @@ -113,30 +194,52 @@ class DagFileParseRequest(BaseModel): """Bundle name for team-specific executor validation.""" callback_requests: list[CallbackRequest] = Field(default_factory=list) + + known_artifacts: list[TaskHandlerArtifact] = Field(default_factory=list) + """The recorded task-handler artifacts, with their answers, that this file's stub tasks may resolve against.""" + type: Literal["DagFileParseRequest"] = "DagFileParseRequest" -class DagFileParsingResult(BaseModel): +class TaskHandlerParseRequest(BaseModel): """ - Result of DAG File Parsing. + Request for Task Handler Parsing. - This is the result of a successful DAG parse, in this class, we gather all serialized DAGs, - import errors and warnings to send back to the scheduler to store in the DB. + Asks a Lang-SDK runtime for every task handler an artifact registers. + """ + + file: str + """The artifact to ask.""" + + bundle_path: Path + + bundle_name: str + + type: Literal["TaskHandlerParseRequest"] = "TaskHandlerParseRequest" + + +class TaskHandlerParsingResult(BaseModel): + """ + Result of Task Handler Parsing. + + Every task handler a Lang-SDK artifact registers, keyed by Dag id. + + The answer depends only on the artifact, never on the request. """ fileloc: str - serialized_dags: list[LazyDeserializedDAG] - warnings: list | None = None + + task_handlers: dict[str, list[TaskHandlerDeclaration]] + """Every Dag id the artifact registers a task handler for; ``{}`` when it registers none.""" + import_errors: dict[str, str] | None = None - parsed_definitions: list[str] = Field(default_factory=list) - """Bundle-relative locations of the Dag definitions imported from ``fileloc``.""" - dag_source_codes: dict[str, DagSourceCode] = Field(default_factory=dict) - """Source code of the parsed Dags, keyed by Dag fileloc.""" - type: Literal["DagFileParsingResult"] = "DagFileParsingResult" + warnings: list | None = None + type: Literal["TaskHandlerParsingResult"] = "TaskHandlerParsingResult" ToManager = Annotated[ DagFileParsingResult + | TaskHandlerParsingResult | GetConnection | GetVariable | GetVariableKeys @@ -155,9 +258,9 @@ class DagFileParsingResult(BaseModel): Field(discriminator="type"), ] -ToDagProcessor = Annotated[ - DagFileParseRequest - | ConnectionResult +# Answers to the child's requests, whichever parse it was started for. +_ParseSideResponses = ( + ConnectionResult | VariableResult | VariableKeysResult | TaskStatesResult @@ -169,8 +272,13 @@ class DagFileParsingResult(BaseModel): | XComCountResponse | XComResult | XComSequenceIndexResult - | XComSequenceSliceResult, - Field(discriminator="type"), + | XComSequenceSliceResult +) + +ToDagProcessor = Annotated[DagFileParseRequest | _ParseSideResponses, Field(discriminator="type")] + +ToSDKTaskHandlerProcessor = Annotated[ + TaskHandlerParseRequest | _ParseSideResponses, Field(discriminator="type") ] @@ -573,18 +681,17 @@ def in_process_api_server() -> InProcessExecutionAPI: @attrs.define(kw_only=True) -class DagFileProcessorProcess(WatchedSubprocess, LoggingMixin): +class BaseDagFileProcessorProcess(WatchedSubprocess, LoggingMixin): """ - Parses dags with Task SDK API. + Parse one Dag file in a child process for the Dag processor manager. - This class provides a wrapper and management around a subprocess to parse a specific DAG file. - - Since DAGs are written with the Task SDK, we need to parse them in a task SDK process such that - we can use the Task SDK definitions when serializing. This prevents potential conflicts with classes - in core Airflow. + The child's output goes to the file's parse log, and its requests are answered with + :attr:`client`. The parse is done once the child has exited and all its sockets are closed; + :attr:`parsing_result` then holds what it sent. Subclasses start the child and send it the + parse request. """ - logger_filehandle: BinaryIO + logger_filehandle: BinaryIO | None = None parsing_result: DagFileParsingResult | None = None decoder: ClassVar[TypeAdapter[ToManager]] = TypeAdapter[ToManager](ToManager) had_callbacks: bool = False # Track if this process was started with callbacks to prevent stale DAG detection false positives @@ -595,60 +702,6 @@ class DagFileProcessorProcess(WatchedSubprocess, LoggingMixin): bundle_name: str dag_file_rel_path: str - @classmethod - def start( # type: ignore[override] - cls, - *, - path: str | os.PathLike[str], - bundle_path: Path, - bundle_name: str, - dag_file_rel_path: str, - callbacks: list[CallbackRequest], - target: Callable[[], None] = _parse_file_entrypoint, - client: Client, - **kwargs, - ) -> Self: - logger = kwargs["logger"] - - # Parsing DAG files runs user code that can trigger macOS-unsafe ObjC - # initialization (secret backends, connection/variable lookups, HTTP - # clients). Fork+exec a clean interpreter there. Tests override `target` - # with a stub to exercise the base infrastructure; keep bare fork for those. - use_exec = target is _parse_file_entrypoint and supervisor._should_use_exec() - - # Pre-importing only helps the bare-fork child (it inherits the imports via - # copy-on-write). An exec'd child re-imports from scratch, so skip it there - # to avoid leaking user modules into the long-lived processor manager. - if not use_exec: - _pre_import_airflow_modules(os.fspath(path), logger) - - proc: Self = super().start( - target=target, - client=client, - bundle_name=bundle_name, - dag_file_rel_path=dag_file_rel_path, - use_exec=use_exec, - **kwargs, - ) - proc.had_callbacks = bool(callbacks) # Track if this process had callbacks - proc._on_child_started(callbacks, path, bundle_path, bundle_name) - return proc - - def _on_child_started( - self, - callbacks: list[CallbackRequest], - path: str | os.PathLike[str], - bundle_path: Path, - bundle_name: str, - ) -> None: - msg = DagFileParseRequest( - file=os.fspath(path), - bundle_path=bundle_path, - bundle_name=bundle_name, - callback_requests=callbacks, - ) - self.send_msg(msg, request_id=0) - def _get_target_loggers(self) -> tuple[FilteringBoundLogger, ...]: base = super()._get_target_loggers() if not self.subprocess_logs_to_stdout: @@ -674,11 +727,11 @@ def _create_log_forwarder( def _handle_parsing_result( self, msg: DagFileParsingResult, log: FilteringBoundLogger, req_id: int - ) -> RequestResult: + ) -> RequestResult | ResponseSent: self.parsing_result = msg return None, {} - _request_handlers: ClassVar[dict[type[BaseModel], RequestHandler[DagFileProcessorProcess]]] = { + _request_handlers: ClassVar[dict[type[BaseModel], RequestHandler[Any]]] = { **WatchedSubprocess._get_shared_request_handlers( DeleteVariable, GetConnection, @@ -720,6 +773,8 @@ def wait(self) -> int: def close(self): self.cleanup_sockets_after_kill() + if self.logger_filehandle is None: + return try: self.logger_filehandle.close() except OSError: @@ -728,3 +783,76 @@ def close(self): self.dag_file_rel_path, exc_info=True, ) + + +@attrs.define(kw_only=True) +class DagFileProcessorProcess(BaseDagFileProcessorProcess): + """ + Parses dags with Task SDK API. + + This class provides a wrapper and management around a subprocess to parse a specific DAG file. + + Since DAGs are written with the Task SDK, we need to parse them in a task SDK process such that + we can use the Task SDK definitions when serializing. This prevents potential conflicts with classes + in core Airflow. + """ + + logger_filehandle: BinaryIO + + @classmethod + def start( # type: ignore[override] + cls, + *, + path: str | os.PathLike[str], + bundle_path: Path, + bundle_name: str, + dag_file_rel_path: str, + callbacks: list[CallbackRequest], + known_artifacts: Sequence[TaskHandlerArtifact] = (), + target: Callable[[], None] = _parse_file_entrypoint, + client: Client, + **kwargs, + ) -> Self: + logger = kwargs["logger"] + + # Parsing DAG files runs user code that can trigger macOS-unsafe ObjC + # initialization (secret backends, connection/variable lookups, HTTP + # clients). Fork+exec a clean interpreter there. Tests override `target` + # with a stub to exercise the base infrastructure; keep bare fork for those. + use_exec = target is _parse_file_entrypoint and supervisor._should_use_exec() + + # Pre-importing only helps the bare-fork child (it inherits the imports via + # copy-on-write). An exec'd child re-imports from scratch, so skip it there + # to avoid leaking user modules into the long-lived processor manager. + if not use_exec: + _pre_import_airflow_modules(os.fspath(path), logger) + + proc: Self = super().start( + target=target, + client=client, + bundle_name=bundle_name, + dag_file_rel_path=dag_file_rel_path, + use_exec=use_exec, + **kwargs, + ) + proc.had_callbacks = bool(callbacks) # Track if this process had callbacks + proc._on_child_started(callbacks, path, bundle_path, bundle_name, known_artifacts=known_artifacts) + return proc + + def _on_child_started( + self, + callbacks: list[CallbackRequest], + path: str | os.PathLike[str], + bundle_path: Path, + bundle_name: str, + *, + known_artifacts: Sequence[TaskHandlerArtifact] = (), + ) -> None: + msg = DagFileParseRequest( + file=os.fspath(path), + bundle_path=bundle_path, + bundle_name=bundle_name, + callback_requests=callbacks, + known_artifacts=list(known_artifacts), + ) + self.send_msg(msg, request_id=0) diff --git a/airflow-core/src/airflow/dag_processing/task_handler_processor.py b/airflow-core/src/airflow/dag_processing/task_handler_processor.py new file mode 100644 index 0000000000000..20a60d0a200cd --- /dev/null +++ b/airflow-core/src/airflow/dag_processing/task_handler_processor.py @@ -0,0 +1,605 @@ +# 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. +"""Ask the Lang-SDK runtime of a coordinator which task handlers an artifact registers.""" + +from __future__ import annotations + +import contextlib +import functools +import os +import selectors +import signal +import time +from pathlib import Path +from socket import MSG_DONTWAIT, socket +from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal, cast, get_args + +import attrs +import msgspec +import psutil +from pydantic import BaseModel, Field, TypeAdapter +from uuid6 import uuid7 + +from airflow import settings +from airflow.configuration import conf +from airflow.dag_processing.processor import ( + BaseDagFileProcessorProcess, + DagFileParsingResult, + TaskHandlerParseRequest, + TaskHandlerParsingResult, + ToManager, +) +from airflow.sdk.coordinators._subprocess import _is_connection_from_pid, _start_server +from airflow.sdk.exceptions import AirflowRuntimeError +from airflow.sdk.execution_time import supervisor, task_runner +from airflow.sdk.execution_time.comms import CommsDecoder, ErrorResponse, MaskSecret, _RequestFrame +from airflow.sdk.execution_time.coordinator import get_coordinator_manager +from airflow.sdk.execution_time.supervisor import ( + ResponseSent, + length_prefixed_frame_reader, + make_buffered_socket_reader, + process_log_messages_from_subprocess, + register_request_method, +) + +if TYPE_CHECKING: + from collections.abc import Generator + + from structlog.typing import FilteringBoundLogger + + from airflow.sdk.api.client import Client + from airflow.sdk.execution_time.supervisor import RequestHandler, RequestResult + from airflow.typing_compat import Self + +# How long a runtime may keep running after its parse result, as Node does while a handle stays open. +_EXIT_GRACE_PERIOD = 5.0 + + +# StartTaskHandlerRuntime, LangSDKRuntimeSchemaVersion and LangSDKRuntimeStartFailed pass only between +# the probe and its forked child before the exec, so they are not part of the supervisor schema the +# runtimes speak. + + +class StartTaskHandlerRuntime(BaseModel): + """Ask the parse child to exec the runtime of the coordinator configured under *coordinator*.""" + + file: str + bundle_path: Path + coordinator: str + """The coordinator's key in ``[sdk] coordinators``.""" + comm_address: tuple[str, int] + logs_address: tuple[str, int] + type: Literal["StartTaskHandlerRuntime"] = "StartTaskHandlerRuntime" + + +class LangSDKRuntimeSchemaVersion(BaseModel): + """The schema version and the import timeout of the runtime the parse child is about to exec.""" + + schema_version: str | None + import_timeout: float | None = None + """Seconds from the start of the parse; ``None`` means no timeout.""" + type: Literal["LangSDKRuntimeSchemaVersion"] = "LangSDKRuntimeSchemaVersion" + + +class LangSDKRuntimeStartFailed(BaseModel): + """Why the parse child could not start the runtime.""" + + error: str + type: Literal["LangSDKRuntimeStartFailed"] = "LangSDKRuntimeStartFailed" + + +def _get_import_timeout(path: str) -> float | None: + """Return the ``get_dagbag_import_timeout`` policy's timeout for *path*; ``None`` means none.""" + timeout = settings.get_dagbag_import_timeout(path) + if not isinstance(timeout, (int, float)): + raise TypeError(f"Value ({timeout}) from get_dagbag_import_timeout must be int or float") + return timeout if timeout > 0 else None + + +def _start_task_handler_runtime_entrypoint() -> None: + """Exec the runtime that probes the artifact named by the start request, or report why it cannot start.""" + os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "client" + # fd 0 becomes the runtime's stdin, so the request channel moves to a close-on-exec copy. + comms = CommsDecoder[StartTaskHandlerRuntime, LangSDKRuntimeSchemaVersion | LangSDKRuntimeStartFailed]( + socket=socket(fileno=os.dup(0)), + body_decoder=TypeAdapter(StartTaskHandlerRuntime), + ) + devnull = os.open(os.devnull, os.O_RDONLY) + os.dup2(devnull, 0) + os.close(devnull) + + msg = comms._get_response() + if not isinstance(msg, StartTaskHandlerRuntime): + raise RuntimeError(f"Required first message to be a StartTaskHandlerRuntime, it was {msg}") + + def report_schema_version(schema_version: str | None) -> None: + comms.send(LangSDKRuntimeSchemaVersion(schema_version=schema_version, import_timeout=import_timeout)) + + try: + # The policy is user code: it runs in this child, where a failure is only this file's import error. + import_timeout = _get_import_timeout(msg.file) + get_coordinator_manager().get_coordinator(msg.coordinator).parse_task_handler( + path=Path(msg.file), + bundle_path=msg.bundle_path, + comm_address=msg.comm_address, + logs_address=msg.logs_address, + report_schema_version=report_schema_version, + ) + except Exception as e: + comms.send(LangSDKRuntimeStartFailed(error=f"{type(e).__name__}: {e}")) + + +_Channel = Literal["comm", "logs"] + + +class _ReadsWithoutWaiting: + """ + The runtime's comm socket, with reads that return at once instead of waiting for data. + + A runtime that stops in the middle of a frame then cannot block the caller's loop, which keeps + checking the import timeout. Replies are sent on the socket itself, which stays blocking. + """ + + def __init__(self, sock: socket) -> None: + self._sock = sock + + def recv(self, bufsize: int) -> bytes: + return self._sock.recv(bufsize, MSG_DONTWAIT) + + def recv_into(self, buffer: memoryview) -> int: + return self._sock.recv_into(buffer, 0, MSG_DONTWAIT) + + +# The requests relayed to the Dag processor without a client: those it answers, but not the Dag parse's +# result. MaskSecret is not relayed: its handler masks the secret here, and mask_secret sends it on to a +# parent. +_PARENT_REQUESTS = frozenset(BaseDagFileProcessorProcess._request_handlers) - { + DagFileParsingResult, + MaskSecret, +} + + +@attrs.define(kw_only=True) +class LangSDKTaskHandlerProcessorProcess(BaseDagFileProcessorProcess): + """ + Ask a coordinator's Lang-SDK runtime for every task handler an artifact registers. + + The forked parse child finds the coordinator, reports the runtime's schema version and execs the + runtime. The runtime connects back to two listeners this process owns and answers the + ``TaskHandlerParseRequest`` itself, so the request is sent once it has connected. A failed start, + a missing result, an invalid frame or message, or a timeout is an import error on the result, + keyed by the artifact's path in its Dag bundle. Processes the runtime leaves in its process group + are killed when it exits. + + The runtime's requests are answered with :attr:`client`. Without one, as in a Dag-parsing child, + they are relayed up the supervisor channel of the process this runs in, or get an error when there + is none. On Linux the runtime is killed when the thread that started this process exits, so start + it from a thread that outlives the probe. + """ + + client: Client | None = None # type: ignore[assignment] + """Answers the runtime's requests; without one, they are relayed to the parent process if there is one.""" + + parsing_result: TaskHandlerParsingResult | None = None # type: ignore[assignment] + + decoder = TypeAdapter( + Annotated[ + LangSDKRuntimeSchemaVersion | LangSDKRuntimeStartFailed | get_args(ToManager)[0], + Field(discriminator="type"), + ] + ) + + coordinator: str + """The coordinator's key in ``[sdk] coordinators``.""" + + _listeners: dict[_Channel, socket] + _parse_request: TaskHandlerParseRequest + _runtime_schema_version: str | None = attrs.field(default=None, init=False) + _import_timeout: float | None = attrs.field(default=None, init=False) + _schema_version_reported: bool = attrs.field(default=False, init=False) + _parsing_result_monotonic: float | None = attrs.field(default=None, init=False) + _unverified_connections: list[tuple[socket, _Channel]] = attrs.field(factory=list, init=False) + + @classmethod + def start( # type: ignore[override] + cls, + *, + coordinator: str, + path: str | os.PathLike[str], + bundle_path: Path, + bundle_name: str, + artifact_rel_path: str, + **kwargs, + ) -> Self: + """ + Start probing the artifact at *path* for every task handler it registers. + + *bundle_path* and *bundle_name* are those of the Dag bundle holding the artifact, and + *artifact_rel_path* is the artifact's path in it. + """ + listeners: dict[_Channel, socket] = {"comm": _start_server(), "logs": _start_server()} + try: + for listener in listeners.values(): + listener.setblocking(False) + parse_request = TaskHandlerParseRequest( + file=os.fspath(path), bundle_path=bundle_path, bundle_name=bundle_name + ) + proc = super().start( + target=_start_task_handler_runtime_entrypoint, + use_exec=supervisor._should_use_exec(), + new_process_group=True, + coordinator=coordinator, + bundle_name=bundle_name, + dag_file_rel_path=artifact_rel_path, + listeners=listeners, + parse_request=parse_request, + **kwargs, + ) + except BaseException: + for listener in listeners.values(): + listener.close() + raise + try: + for channel, listener in listeners.items(): + proc._open_sockets[listener] = f"{channel}-listener" + proc.selector.register( + listener, + selectors.EVENT_READ, + (functools.partial(proc._accept_connection, channel=channel), proc._on_socket_closed), + ) + proc.send_msg( + StartTaskHandlerRuntime( + file=parse_request.file, + bundle_path=bundle_path, + coordinator=coordinator, + comm_address=listeners["comm"].getsockname()[:2], + logs_address=listeners["logs"].getsockname()[:2], + ), + request_id=0, + ) + except BaseException: + proc._kill_runtime() + proc.close() + raise + return proc + + @classmethod + def run( + cls, + *, + coordinator: str, + path: str | os.PathLike[str], + bundle_path: Path, + bundle_name: str, + artifact_rel_path: str, + logger: FilteringBoundLogger, + ) -> TaskHandlerParsingResult: + """ + Probe the artifact at *path* as :meth:`start` does, and wait for the result. + + The artifact's ``get_dagbag_import_timeout`` bounds the probe, and + ``[dag_processor] dag_file_processor_timeout`` until the parse child resolves it. + """ + return cls._run_to_completion( + coordinator=coordinator, + path=path, + bundle_path=bundle_path, + bundle_name=bundle_name, + artifact_rel_path=artifact_rel_path, + logger=logger, + ) + + @classmethod + def _run_to_completion( + cls, *, logger: FilteringBoundLogger, **start_kwargs: Any + ) -> TaskHandlerParsingResult: + """ + Probe outside a caller's selector loop and wait for the result. + + The import timeout the parse child reports bounds the probe, and + ``[dag_processor] dag_file_processor_timeout`` until it is reported. + """ + processor_timeout = conf.getfloat("dag_processor", "dag_file_processor_timeout") + with selectors.DefaultSelector() as selector: + proc = cls.start(id=uuid7(), selector=selector, logger=logger, **start_kwargs) + try: + while not proc.is_ready: + # is_ready applies the import timeout once the parse child has reported it. + if ( + not proc._schema_version_reported + and time.monotonic() - proc.start_time > processor_timeout + ): + proc._time_out(processor_timeout) + break + proc._service_subprocess(max_wait_time=0.1) + except BaseException: + proc._kill_runtime() + raise + finally: + proc.close() + return cast("TaskHandlerParsingResult", proc.parsing_result) + + def _accept_connection(self, listener: socket, *, channel: _Channel) -> bool: + try: + conn, _ = listener.accept() + except (BlockingIOError, InterruptedError): + return True + conn.setblocking(True) + self._unverified_connections.append((conn, channel)) + self._verify_connections() + return True + + def _verify_connections(self) -> None: + """ + Use each accepted connection once it is confirmed to come from the runtime. + + A connection that is not visible yet stays pending and is checked again on the next + ``is_ready`` poll, so the caller's loop never waits here. + """ + pending = [] + for conn, channel in self._unverified_connections: + if channel not in self._listeners: + # The runtime already connected this channel. + conn.close() + continue + try: + owned = _is_connection_from_pid(conn, self.pid) + except OSError: + conn.close() + continue + if not owned: + pending.append((conn, channel)) + continue + self._close_listener(channel) + if channel == "comm": + self._register_comm(conn) + else: + self._register_logs(conn) + self._unverified_connections = pending + + def _close_listener(self, channel: _Channel) -> None: + if (listener := self._listeners.pop(channel, None)) is not None: + self._on_socket_closed(listener) + listener.close() + + def _close_listeners(self) -> None: + """Close the listeners of a runtime that did not connect, and connections never verified.""" + for channel in list(self._listeners): + self._close_listener(channel) + for conn, _ in self._unverified_connections: + conn.close() + self._unverified_connections = [] + + def _register_comm(self, conn: socket) -> None: + self.stdin = conn + self._open_sockets[conn] = "requests" + read_frame, on_close = length_prefixed_frame_reader( + self._handle_valid_requests(), on_close=self._on_socket_closed + ) + + def read_valid_frame(sock: socket) -> bool: + try: + return read_frame(cast("socket", _ReadsWithoutWaiting(sock))) + except BlockingIOError: + # The rest of the frame has not arrived; the reader keeps what it has read so far. + return True + except msgspec.DecodeError as e: + # A frame that does not decode would otherwise escape the caller's selector loop. + self._fail_on_invalid_message(f"The Lang-SDK runtime sent an invalid frame: {e}") + return False + + self.selector.register(conn, selectors.EVENT_READ, (read_valid_frame, on_close)) + # The parse child reports the version and waits for the reply before it execs the runtime, + # so the version is known here. It is set only now, so the child's messages are not migrated. + self._subprocess_schema_version = self._runtime_schema_version + self.send_msg(self._parse_request, request_id=0) + + def _handle_valid_requests(self) -> Generator[None, _RequestFrame, None]: + """ + Pass each request on to ``handle_requests``, or kill the runtime at one that does not validate. + + ``handle_requests`` would only log such a request, and the runtime would wait for a reply. The + runtime speaks ``ToManager`` only; the start messages come from the parse child. + """ + requests = self.handle_requests(self.process_log) + next(requests) + while True: + frame = yield + try: + BaseDagFileProcessorProcess.decoder.validate_python(self._deserialize_request(frame.body)) + except ValueError as e: + self._fail_on_invalid_message( + f"The Lang-SDK runtime sent a message that does not validate: {e}" + ) + return + requests.send(frame) + + def _fail_on_invalid_message(self, message: str) -> None: + """Kill the runtime; *message* is the import error unless a parse result was already received.""" + if self.parsing_result is None: + self._set_import_error(message) + else: + self.process_log.warning( + "Ignoring an invalid message from the Lang-SDK runtime after its parse result", error=message + ) + self._kill_runtime() + + def _register_logs(self, conn: socket) -> None: + self._open_sockets[conn] = "logs" + self.selector.register( + conn, + selectors.EVENT_READ, + make_buffered_socket_reader( + process_log_messages_from_subprocess(self._get_target_loggers()), + on_close=self._on_socket_closed, + ), + ) + + def _set_import_error(self, message: str) -> None: + self.parsing_result = TaskHandlerParsingResult( + fileloc=self._parse_request.file, + task_handlers={}, + import_errors={self.dag_file_rel_path: message}, + ) + + def _handle_runtime_schema_version( + self, msg: LangSDKRuntimeSchemaVersion, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + if self._schema_version_reported: + self._reject_request(msg, log, req_id) + return ResponseSent.ALREADY_SENT + self._runtime_schema_version = msg.schema_version + self._import_timeout = msg.import_timeout + self._schema_version_reported = True + return None, {} + + def _handle_start_failed( + self, msg: LangSDKRuntimeStartFailed, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + self._set_import_error(f"Cannot start the Lang-SDK runtime: {msg.error}") + return None, {} + + def _handle_parsing_result( # type: ignore[override] + self, msg: TaskHandlerParsingResult, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + if self.parsing_result is not None: + log.warning("Ignoring another parse result from the Lang-SDK runtime", fileloc=msg.fileloc) + self.send_msg( + None, + request_id=req_id, + error=ErrorResponse(detail={"message": "A parse result was already received"}), + ) + return ResponseSent.ALREADY_SENT + self.parsing_result = msg + self._parsing_result_monotonic = time.monotonic() + return None, {} + + _request_handlers: ClassVar[dict[type[BaseModel], RequestHandler[Any]]] = { + **{ + message_type: handler + for message_type, handler in BaseDagFileProcessorProcess._request_handlers.items() + if message_type is not DagFileParsingResult + }, + **dict( + [ + register_request_method(LangSDKRuntimeSchemaVersion, _handle_runtime_schema_version), + register_request_method(LangSDKRuntimeStartFailed, _handle_start_failed), + register_request_method(TaskHandlerParsingResult, _handle_parsing_result), + ] + ), + } + + def _handle_request(self, msg, log: FilteringBoundLogger, req_id: int) -> None: + if self.client is None and type(msg) in _PARENT_REQUESTS: + self._relay_request(msg, req_id) + return + super()._handle_request(msg, log, req_id) + + def _relay_request(self, msg: BaseModel, req_id: int) -> None: + """Answer the runtime's request through the supervisor channel of this process, if it has one.""" + comms = getattr(task_runner, "SUPERVISOR_COMMS", None) + if comms is None: + self.send_msg( + None, + request_id=req_id, + error=ErrorResponse( + detail={"message": f"{type(msg).__name__} is answered only in the Dag processor"} + ), + ) + return + try: + response = comms.send(msg) + except AirflowRuntimeError as e: + self.send_msg(None, request_id=req_id, error=e.error) + return + # Only the fields the parent sent, under their wire names, so the runtime gets the same body. + self.send_msg(response, request_id=req_id, exclude_unset=True, by_alias=True) + + @property + def is_ready(self) -> bool: + self._verify_connections() + if ( + self._parsing_result_monotonic is not None + and self._exit_code is None + and time.monotonic() - self._parsing_result_monotonic > _EXIT_GRACE_PERIOD + ): + self.process_log.warning("The Lang-SDK runtime did not exit after its parse result; killing it") + self._kill_runtime() + if ( + self._import_timeout is not None + and self.parsing_result is None + and self._exit_code is None + and time.monotonic() - self.start_time > self._import_timeout + ): + self._time_out(self._import_timeout) + if self._check_subprocess_exit() is None: + return False + self._close_listeners() + if ( + self._open_sockets + and self._import_timeout is not None + and time.monotonic() - self.start_time > self._import_timeout + ): + # A process the runtime left outside its process group holds these open. + self._time_out(self._import_timeout) + self.cleanup_sockets_after_kill() + if not super().is_ready: + return False + if self.parsing_result is None: + self._set_import_error( + f"The Lang-SDK runtime exited with code {self._exit_code} without a parse result" + ) + return True + + def _time_out(self, timeout: float) -> None: + if self.parsing_result is None: + self._set_import_error( + f"The Lang-SDK runtime did not parse {self._parse_request.file} within {timeout}s" + ) + self._kill_runtime() + + def _check_subprocess_exit( + self, raise_on_timeout: bool = False, expect_signal: None | int = None + ) -> int | None: + if self._exit_code is None and self._is_runtime_exited(): + # Until the exited runtime is reaped below, its pid, and so its process group id, cannot be + # reused, so this reaches only the processes it left behind. + with contextlib.suppress(ProcessLookupError, PermissionError): + os.killpg(self.pid, signal.SIGKILL) + return super()._check_subprocess_exit(raise_on_timeout=raise_on_timeout, expect_signal=expect_signal) + + def _is_runtime_exited(self) -> bool: + """Return whether the runtime has exited and is not reaped yet.""" + try: + return psutil.Process(self.pid).status() == psutil.STATUS_ZOMBIE + except psutil.Error: + return False + + def _kill_runtime(self) -> None: + """Kill the runtime and wait for it, without servicing its sockets, whose handler may have failed.""" + if self._exit_code is not None: + return + try: + self._signal_subprocess(signal.SIGKILL) + self._exit_code = self._process.wait(timeout=None) + except (self._process.ProcessNotFound, ProcessLookupError): + self._exit_code = -1 + + def close(self) -> None: + # A listener has nothing to drain, and cleanup would call its accept handler forever. + self._close_listeners() + super().close() diff --git a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py new file mode 100644 index 0000000000000..c9679ede25bec --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_add_lang_sdk_task_handler_tables.py @@ -0,0 +1,91 @@ +# +# 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. + +""" +Add lang_sdk_task_handler_artifact and lang_sdk_task_handler tables. + +Revision ID: f7ed13533d23 +Revises: 90e4d18ccadf +Create Date: 2026-09-30 14:19:35.109799 + +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +from airflow.migrations.db_types import StringID +from airflow.utils.sqlalchemy import UtcDateTime + +# revision identifiers, used by Alembic. +revision = "f7ed13533d23" +down_revision = "90e4d18ccadf" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + + +def upgrade(): + """Add lang_sdk_task_handler_artifact and lang_sdk_task_handler tables.""" + op.create_table( + "lang_sdk_task_handler_artifact", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column("bundle_name", StringID(), nullable=False), + sa.Column("relative_fileloc", sa.String(length=2000), nullable=False), + sa.Column("relative_fileloc_hash", sa.String(length=32), nullable=False), + sa.Column("size_bytes", sa.BigInteger(), nullable=False), + sa.Column("cache_digest", sa.String(length=128), nullable=True), + sa.Column("task_handlers", sa.JSON(), nullable=False), + sa.Column("last_probed_at", UtcDateTime(), nullable=False), + sa.PrimaryKeyConstraint("id", name=op.f("lang_sdk_task_handler_artifact_pkey")), + sa.UniqueConstraint( + "bundle_name", + "relative_fileloc_hash", + name=op.f("lang_sdk_task_handler_artifact_bundle_fileloc_uq"), + ), + ) + op.create_table( + "lang_sdk_task_handler", + sa.Column("dag_id", StringID(), nullable=False), + sa.Column("task_id", StringID(), nullable=False), + sa.Column("artifact_id", sa.Uuid(), nullable=False), + sa.Column("dag_bundle_name", StringID(), nullable=False), + sa.Column("dag_relative_fileloc", sa.String(length=2000), nullable=False), + sa.Column("dag_relative_fileloc_hash", sa.String(length=32), nullable=False), + sa.PrimaryKeyConstraint("dag_id", "task_id", name=op.f("lang_sdk_task_handler_pkey")), + sa.ForeignKeyConstraint( + columns=("dag_id",), + refcolumns=["dag.dag_id"], + name="lang_sdk_task_handler_dag_id_fkey", + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + columns=("artifact_id",), + refcolumns=["lang_sdk_task_handler_artifact.id"], + name="lang_sdk_task_handler_artifact_id_fkey", + ), + sa.Index("idx_lang_sdk_task_handler_dag_file", "dag_bundle_name", "dag_relative_fileloc_hash"), + sa.Index("idx_lang_sdk_task_handler_artifact_id", "artifact_id"), + ) + + +def downgrade(): + """Drop lang_sdk_task_handler and lang_sdk_task_handler_artifact tables.""" + op.drop_table("lang_sdk_task_handler") + op.drop_table("lang_sdk_task_handler_artifact") diff --git a/airflow-core/src/airflow/models/__init__.py b/airflow-core/src/airflow/models/__init__.py index f5bffb4e57173..f806537c684df 100644 --- a/airflow-core/src/airflow/models/__init__.py +++ b/airflow-core/src/airflow/models/__init__.py @@ -71,6 +71,7 @@ def import_all_models(): import airflow.models.dagwarning import airflow.models.deadline_alert import airflow.models.errors + import airflow.models.lang_sdk_task_handler import airflow.models.revoked_token import airflow.models.serialized_dag import airflow.models.task_state_store diff --git a/airflow-core/src/airflow/models/lang_sdk_task_handler.py b/airflow-core/src/airflow/models/lang_sdk_task_handler.py new file mode 100644 index 0000000000000..86cb020f3f1d7 --- /dev/null +++ b/airflow-core/src/airflow/models/lang_sdk_task_handler.py @@ -0,0 +1,128 @@ +# +# 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. +from __future__ import annotations + +from collections.abc import Callable +from datetime import datetime +from typing import TYPE_CHECKING, Any +from uuid import UUID + +import sqlalchemy as sa +import uuid6 +from sqlalchemy import BigInteger, ForeignKeyConstraint, Index, String, UniqueConstraint, Uuid +from sqlalchemy.orm import Mapped, mapped_column, validates + +from airflow._shared.timezones import timezone +from airflow.models.base import Base, StringID +from airflow.utils.hashlib_wrapper import md5 +from airflow.utils.sqlalchemy import UtcDateTime + +if TYPE_CHECKING: + from sqlalchemy.engine.default import DefaultExecutionContext + + +def compute_fileloc_hash(relative_fileloc: str) -> str: + """ + Return the md5 hex digest that the task handler tables index in place of a relative file location. + + A 2000-character path exceeds MySQL's 3072-byte key limit and Postgres's 2704-byte btree entry limit. + """ + return md5(relative_fileloc.encode()).hexdigest() + + +# The column default fills a hash on INSERT, ORM or Core. ORM updates resync it in each model's +# @validates hook instead of an onupdate default, which would fire on every UPDATE, even one that +# leaves the path alone. A Core UPDATE or upsert that changes a path must set its hash itself. +def _build_fileloc_hash_default(path_key: str) -> Callable[[DefaultExecutionContext], str]: + def compute(context: DefaultExecutionContext) -> str: + return compute_fileloc_hash(context.get_current_parameters()[path_key]) + + return compute + + +class LangSDKTaskHandlerArtifact(Base): + """ + A Language SDK artifact in a task handler Dag bundle, cached with its fingerprint and probe answer. + + ``task_handlers`` is the runtime's whole answer, ``{dag_id: [TaskHandlerDeclaration as JSON, ...]}``. + It is written only from a probe and together with the fingerprint, so the two always match. + """ + + __tablename__ = "lang_sdk_task_handler_artifact" + + id: Mapped[UUID] = mapped_column(Uuid(), primary_key=True, default=uuid6.uuid7) + bundle_name: Mapped[str] = mapped_column(StringID(), nullable=False) + relative_fileloc: Mapped[str] = mapped_column(String(2000), nullable=False) + relative_fileloc_hash: Mapped[str] = mapped_column( + String(32), nullable=False, default=_build_fileloc_hash_default("relative_fileloc") + ) + size_bytes: Mapped[int] = mapped_column(BigInteger, nullable=False) + # NULL means the artifact stores no fingerprint and is probed again on every parse. + cache_digest: Mapped[str | None] = mapped_column(String(128), nullable=True) + task_handlers: Mapped[dict[str, list[dict[str, Any]]]] = mapped_column(sa.JSON(), nullable=False) + last_probed_at: Mapped[datetime] = mapped_column(UtcDateTime, nullable=False, default=timezone.utcnow) + + __table_args__ = ( + UniqueConstraint( + bundle_name, + relative_fileloc_hash, + name="lang_sdk_task_handler_artifact_bundle_fileloc_uq", + ), + ) + + @validates("relative_fileloc") + def _sync_relative_fileloc_hash(self, key: str, relative_fileloc: str) -> str: + self.relative_fileloc_hash = compute_fileloc_hash(relative_fileloc) + return relative_fileloc + + +class LangSDKTaskHandler(Base): + """The binding of one stub task, by ``(dag_id, task_id)``, to the artifact that registers its handler.""" + + __tablename__ = "lang_sdk_task_handler" + + dag_id: Mapped[str] = mapped_column(StringID(), primary_key=True) + task_id: Mapped[str] = mapped_column(StringID(), primary_key=True) + artifact_id: Mapped[UUID] = mapped_column(Uuid(), nullable=False) + dag_bundle_name: Mapped[str] = mapped_column(StringID(), nullable=False) + dag_relative_fileloc: Mapped[str] = mapped_column(String(2000), nullable=False) + dag_relative_fileloc_hash: Mapped[str] = mapped_column( + String(32), nullable=False, default=_build_fileloc_hash_default("dag_relative_fileloc") + ) + + __table_args__ = ( + ForeignKeyConstraint( + (dag_id,), + ["dag.dag_id"], + name="lang_sdk_task_handler_dag_id_fkey", + ondelete="CASCADE", + ), + # No ON DELETE: deleting an artifact that a handler still references must fail. + ForeignKeyConstraint( + (artifact_id,), + ["lang_sdk_task_handler_artifact.id"], + name="lang_sdk_task_handler_artifact_id_fkey", + ), + Index("idx_lang_sdk_task_handler_dag_file", dag_bundle_name, dag_relative_fileloc_hash), + Index("idx_lang_sdk_task_handler_artifact_id", artifact_id), + ) + + @validates("dag_relative_fileloc") + def _sync_dag_relative_fileloc_hash(self, key: str, dag_relative_fileloc: str) -> str: + self.dag_relative_fileloc_hash = compute_fileloc_hash(dag_relative_fileloc) + return dag_relative_fileloc diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index f160cb0597d84..689d57ded759f 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "90e4d18ccadf", + "3.4.0": "f7ed13533d23", } # Prefix used to identify tables holding data moved during migration. diff --git a/airflow-core/tests/unit/dag_processing/fake_task_handler_runtime.py b/airflow-core/tests/unit/dag_processing/fake_task_handler_runtime.py new file mode 100644 index 0000000000000..151d30d4a8dfb --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/fake_task_handler_runtime.py @@ -0,0 +1,120 @@ +# +# 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. +"""A coordinator whose artifacts name the command that parses them, and a runtime a test can play.""" + +from __future__ import annotations + +import contextlib +import json +import socket +from pathlib import Path +from typing import TYPE_CHECKING, Any +from unittest import mock + +import attrs +from pydantic import TypeAdapter + +from airflow.dag_processing.processor import ( + TaskHandlerParseRequest, + TaskHandlerParsingResult, + ToManager, + ToSDKTaskHandlerProcessor, +) +from airflow.sdk.coordinators._subprocess import SubprocessCoordinator +from airflow.sdk.execution_time import supervisor +from airflow.sdk.execution_time.comms import CommsDecoder +from airflow.sdk.execution_time.coordinator import reset_coordinator_manager + +from tests_common.test_utils.config import conf_vars + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator, Sequence + + +@attrs.define(kw_only=True) +class FakeCoordinator(SubprocessCoordinator): + """ + The artifact's JSON names the command that parses it. + + ``argv`` is the command, ``schema_version`` its schema version, and ``command_error`` an error to + raise instead. + """ + + def _build_parse_task_handler_command(self, *, path: Path) -> tuple[list[str], str | None]: + spec = json.loads(path.read_text()) + if error := spec.get("command_error"): + raise FileNotFoundError(error) + return spec.get("argv", ["/bin/false"]), spec.get("schema_version") + + +@contextlib.contextmanager +def fake_coordinator(**kwargs: Any) -> Iterator[None]: + """ + Configure a ``FakeCoordinator`` as ``fake``, with fresh coordinators inside and after the block. + + The parse child is a bare fork even on macOS, so it sees the test's patches and config. + """ + spec = {"fake": {"classpath": f"{__name__}.FakeCoordinator", "kwargs": kwargs}} + reset_coordinator_manager() + try: + with ( + conf_vars({("sdk", "coordinators"): json.dumps(spec)}), + mock.patch.object(supervisor, "_should_use_exec", return_value=False), + ): + yield + finally: + reset_coordinator_manager() + + +def write_artifact(path: Path, **spec: Any) -> Path: + path.write_text(json.dumps(spec)) + return path + + +def play_runtime( + reply: Callable[[TaskHandlerParseRequest, CommsDecoder], TaskHandlerParsingResult | None], + *, + schema_version: str | None = None, + log_lines: Sequence[dict[str, Any]] = (), +) -> Callable[..., None]: + """ + Return a ``parse_task_handler`` that plays the runtime in the parse child instead of exec'ing one. + + It connects back as a runtime does, writes *log_lines* to its logs channel and answers the parse + request with what *reply* returns; ``None`` sends no result. + """ + + def parse_task_handler( + coordinator, *, comm_address, logs_address, report_schema_version, **kwargs + ) -> None: + report_schema_version(schema_version) + with ( + socket.create_connection(comm_address) as comm, + socket.create_connection(logs_address) as logs, + ): + for line in log_lines: + logs.sendall(json.dumps(line).encode() + b"\n") + comms = CommsDecoder[ToSDKTaskHandlerProcessor, ToManager]( + socket=comm, body_decoder=TypeAdapter(ToSDKTaskHandlerProcessor) + ) + request = comms._get_response() + assert isinstance(request, TaskHandlerParseRequest) + if (result := reply(request, comms)) is not None: + comms.send(result) + + return parse_task_handler diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index a3624e1d4a143..1f0a66cbf303d 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -61,6 +61,9 @@ DagFileParseRequest, DagFileParsingResult, DagFileProcessorProcess, + TaskHandlerArtifact, + TaskHandlerDeclaration, + TaskHandlerParam, _parse_file, ) from airflow.models import DagModel, DbCallbackRequest @@ -68,10 +71,12 @@ from airflow.models.dag_version import DagVersion from airflow.models.dagbundle import DagBundleModel from airflow.models.dagcode import DagCode +from airflow.models.lang_sdk_task_handler import LangSDKTaskHandlerArtifact from airflow.models.serialized_dag import SerializedDagModel from airflow.models.team import Team from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import DAG as SdkDAG +from airflow.sdk.execution_time.coordinator import reset_coordinator_manager from airflow.sdk.importers import DagSourceCode from airflow.serialization.serialized_objects import LazyDeserializedDAG from airflow.utils.net import get_hostname @@ -1859,9 +1864,10 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): @pytest.mark.usefixtures("testing_dag_bundle") @pytest.mark.parametrize( - ("callbacks", "path", "expected_body"), + ("callbacks", "known_artifacts", "path", "expected_body"), [ pytest.param( + [], [], "/opt/airflow/dags/test_dag.py", { @@ -1869,8 +1875,61 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): "bundle_path": "/opt/airflow/dags", "bundle_name": "testing", "callback_requests": [], + "known_artifacts": [], + "type": "DagFileParseRequest", + }, + ), + pytest.param( + [], + [ + TaskHandlerArtifact( + bundle_name="java-task-handlers", + relative_fileloc="etl.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={ + "etl": [ + TaskHandlerDeclaration( + task_id="extract", + binding="positional", + params=[TaskHandlerParam(name=None)], + ) + ] + }, + ) + ], + "/opt/airflow/dags/etl.py", + { + "file": "/opt/airflow/dags/etl.py", + "bundle_path": "/opt/airflow/dags", + "bundle_name": "testing", + "callback_requests": [], + "known_artifacts": [ + { + "bundle_name": "java-task-handlers", + "relative_fileloc": "etl.jar", + "size_bytes": 1024, + "cache_digest": "ab12", + "task_handlers": { + "etl": [ + { + "task_id": "extract", + "binding": "positional", + "params": [ + { + "name": None, + "value_schema": None, + "exact_name": False, + } + ], + } + ] + }, + } + ], "type": "DagFileParseRequest", }, + id="known-artifacts", ), pytest.param( [ @@ -1884,6 +1943,7 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): is_failure_callback=False, ) ], + [], "/opt/airflow/dags/dag_callback_dag.py", { "file": "/opt/airflow/dags/dag_callback_dag.py", @@ -1903,17 +1963,22 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): "type": "DagCallbackRequest", } ], + "known_artifacts": [], "type": "DagFileParseRequest", }, ), ], ) - def test_serialize_callback_requests(self, callbacks, path, expected_body): + def test_serialize_parse_request(self, callbacks, known_artifacts, path, expected_body): from airflow.sdk.execution_time.comms import _ResponseFrame processor, read_socket = self.mock_processor() processor._on_child_started( - callbacks, path, bundle_path=Path("/opt/airflow/dags"), bundle_name="testing" + callbacks, + path, + bundle_path=Path("/opt/airflow/dags"), + bundle_name="testing", + known_artifacts=known_artifacts, ) read_socket.settimeout(0.1) @@ -2623,6 +2688,7 @@ def test_callback_queue(self, mock_get_logger, configure_testing_dag_bundle): bundle_name="testing", dag_file_rel_path=str(dag2_path.rel_path), callbacks=[dag2_req1], + known_artifacts=[], selector=mock.ANY, logger=mock_logger, logger_filehandle=mock_filehandle, @@ -2636,6 +2702,7 @@ def test_callback_queue(self, mock_get_logger, configure_testing_dag_bundle): bundle_name="testing", dag_file_rel_path=str(dag1_path.rel_path), callbacks=[dag1_req1, dag1_req2], + known_artifacts=[], selector=mock.ANY, logger=mock_logger, logger_filehandle=mock_filehandle, @@ -4102,6 +4169,333 @@ def test_a_sweep_pays_the_fixed_cost_once_per_call(self, session, testing_dag_bu ) +JAVA_TASK_HANDLERS = "java-task-handlers" + + +def _coordinator(task_handler_bundle_name: str | None = None) -> dict: + kwargs = ( + {} if task_handler_bundle_name is None else {"task_handler_bundle_name": task_handler_bundle_name} + ) + return {"classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": kwargs} + + +def _routed_coordinators(coordinators: dict[str, dict]) -> dict[tuple[str, str], str]: + """Return the ``[sdk]`` options that configure *coordinators* and route a queue to each.""" + return { + ("sdk", "coordinators"): json.dumps(coordinators), + ("sdk", "queue_to_coordinator"): json.dumps({f"queue-{key}": key for key in coordinators}), + } + + +def _known_artifact(bundle_name: str) -> TaskHandlerArtifact: + return TaskHandlerArtifact( + bundle_name=bundle_name, + relative_fileloc=f"{bundle_name}.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={ + "etl": [ + TaskHandlerDeclaration( + task_id="extract", + binding="positional", + params=[TaskHandlerParam(name=None)], + ) + ] + }, + ) + + +GO_TASK_HANDLERS = "go-task-handlers" +TS_TASK_HANDLERS = "ts-task-handlers" +TEAM_COORDINATORS = { + "java": _coordinator(JAVA_TASK_HANDLERS), + "go": _coordinator(GO_TASK_HANDLERS), + "ts": _coordinator(TS_TASK_HANDLERS), +} + + +class TestKnownTaskHandlerArtifacts: + @pytest.fixture(autouse=True) + def _reset_coordinator_manager(self): + reset_coordinator_manager() + yield + reset_coordinator_manager() + + @pytest.fixture + def manager(self, configure_dag_bundles, tmp_path): + with ( + conf_vars({("core", "load_examples"): "False"}), + configure_dag_bundles({"dags-a": tmp_path, "dags-b": tmp_path, JAVA_TASK_HANDLERS: tmp_path}), + ): + manager = DagFileProcessorManager(max_runs=1) + manager._dag_bundles = list(DagBundlesManager().get_all_dag_bundles()) + yield manager + + @conf_vars(_routed_coordinators({"java": _coordinator(JAVA_TASK_HANDLERS), "go": _coordinator()})) + @mock.patch.object(DagFileProcessorProcess, "start") + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_start_new_processes_reads_known_artifacts_once_per_loop(self, get_known, start, manager): + known = {name: [_known_artifact(name)] for name in ("dags-a", "dags-b", JAVA_TASK_HANDLERS)} + get_known.return_value = known + files = [ + DagFileInfo(bundle_name=bundle_name, rel_path=Path(rel_path), bundle_path=TEST_DAGS_FOLDER) + for bundle_name, rel_path in (("dags-a", "one.py"), ("dags-a", "two.py"), ("dags-b", "three.py")) + ] + manager._parallelism = 3 + manager._file_queue = OrderedDict.fromkeys(files) + + manager._start_new_processes() + + get_known.assert_called_once() + assert { + call.kwargs["dag_file_rel_path"]: call.kwargs["known_artifacts"] for call in start.call_args_list + } == { + "one.py": [*known["dags-a"], *known[JAVA_TASK_HANDLERS]], + "two.py": [*known["dags-a"], *known[JAVA_TASK_HANDLERS]], + "three.py": [*known["dags-b"], *known[JAVA_TASK_HANDLERS]], + } + + manager._processors.clear() + manager._file_queue = OrderedDict.fromkeys(files[:1]) + manager._start_new_processes() + + assert get_known.call_count == 2 + + @pytest.mark.parametrize( + ("coordinators", "teams", "expected_bundle_names"), + [ + pytest.param( + TEAM_COORDINATORS, + {}, + {JAVA_TASK_HANDLERS, GO_TASK_HANDLERS, TS_TASK_HANDLERS}, + id="multi-team-off", + ), + pytest.param( + TEAM_COORDINATORS, + {"dags-a": "team-a", JAVA_TASK_HANDLERS: "team-a", GO_TASK_HANDLERS: "team-b"}, + {JAVA_TASK_HANDLERS}, + id="same-team", + ), + pytest.param( + TEAM_COORDINATORS, + {JAVA_TASK_HANDLERS: "team-a", GO_TASK_HANDLERS: "team-b"}, + {TS_TASK_HANDLERS}, + id="no-team", + ), + pytest.param( + {**TEAM_COORDINATORS, "fallback": _coordinator()}, + {"dags-a": "team-a", JAVA_TASK_HANDLERS: "team-a", GO_TASK_HANDLERS: "team-b"}, + {JAVA_TASK_HANDLERS, "dags-a"}, + id="own-bundle-fallback", + ), + ], + ) + def test_known_artifacts_are_scoped_by_team( + self, configure_dag_bundles, tmp_path, coordinators, teams, expected_bundle_names + ): + bundle_names = ("dags-a", JAVA_TASK_HANDLERS, GO_TASK_HANDLERS, TS_TASK_HANDLERS) + with ( + conf_vars({("core", "load_examples"): "False", **_routed_coordinators(coordinators)}), + configure_dag_bundles(dict.fromkeys(bundle_names, tmp_path)), + mock.patch.object(DagFileProcessorManager, "_get_team_names", autospec=True, return_value=teams), + ): + selected = DagFileProcessorManager(max_runs=1)._select_known_task_handler_artifacts( + {name: [_known_artifact(name)] for name in bundle_names}, "dags-a" + ) + + assert {artifact.bundle_name for artifact in selected} == expected_bundle_names + + @conf_vars(_routed_coordinators({"java": _coordinator(JAVA_TASK_HANDLERS)})) + @mock.patch.object(DagFileProcessorProcess, "start") + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_start_new_processes_keeps_own_bundle_out_without_a_fallback(self, get_known, start, manager): + # An override may return bundles it was not asked for. + get_known.return_value = {name: [_known_artifact(name)] for name in ("dags-a", JAVA_TASK_HANDLERS)} + manager._file_queue = OrderedDict.fromkeys( + [DagFileInfo(bundle_name="dags-a", rel_path=Path("one.py"), bundle_path=TEST_DAGS_FOLDER)] + ) + + manager._start_new_processes() + + assert start.call_args.kwargs["known_artifacts"] == [_known_artifact(JAVA_TASK_HANDLERS)] + + @conf_vars(_routed_coordinators({"java": _coordinator(JAVA_TASK_HANDLERS)})) + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_start_new_processes_reads_nothing_without_a_child_to_start(self, get_known, manager): + manager._start_new_processes() + + get_known.assert_not_called() + + @pytest.mark.parametrize( + ("coordinators", "expected_bundle_names"), + [ + pytest.param({"java": _coordinator(JAVA_TASK_HANDLERS)}, {JAVA_TASK_HANDLERS}, id="named-only"), + pytest.param( + {"java": _coordinator(JAVA_TASK_HANDLERS), "go": _coordinator()}, + {JAVA_TASK_HANDLERS, "dags-a", "dags-b"}, + id="own-bundle-fallback", + ), + ], + ) + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_query_asks_only_for_bundles_that_can_hold_artifacts( + self, get_known, manager, coordinators, expected_bundle_names + ): + with conf_vars(_routed_coordinators(coordinators)): + manager._query_known_task_handler_artifacts() + + get_known.assert_called_once_with(manager, expected_bundle_names) + + @pytest.mark.parametrize( + "unrouted", + [ + _coordinator("ghost"), + { + "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", + "kwargs": {"task_handler_bundle_name": ["ghost"]}, + }, + _coordinator(), + ], + ids=["unconfigured", "not-a-string", "own-bundle"], + ) + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_query_leaves_out_the_bundle_of_an_unrouted_coordinator(self, get_known, manager, unrouted): + with conf_vars( + { + ("sdk", "coordinators"): json.dumps( + {"java": _coordinator(JAVA_TASK_HANDLERS), "unrouted": unrouted} + ), + ("sdk", "queue_to_coordinator"): json.dumps({"queue-java": "java"}), + } + ): + manager._query_known_task_handler_artifacts() + + get_known.assert_called_once_with(manager, {JAVA_TASK_HANDLERS}) + + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_query_skips_the_read_without_coordinators(self, get_known, manager): + assert manager._query_known_task_handler_artifacts() == {} + + get_known.assert_not_called() + + @conf_vars(_routed_coordinators({"java": _coordinator("ghost")})) + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_invalid_coordinator_config_is_logged_once_and_sends_nothing(self, get_known, manager, caplog): + assert manager._query_known_task_handler_artifacts() == {} + assert manager._query_known_task_handler_artifacts() == {} + + get_known.assert_not_called() + event = "Cannot read [sdk] coordinators; Dag files are parsed without known task-handler artifacts" + assert [entry["log_level"] for entry in caplog.entries if entry["event"] == event] == ["error"] + + @conf_vars(_routed_coordinators({"java": _coordinator(JAVA_TASK_HANDLERS)})) + @mock.patch.object(DagFileProcessorProcess, "start") + @mock.patch.object(DagFileProcessorManager, "get_known_task_handler_artifacts", autospec=True) + def test_start_new_processes_sends_no_known_artifacts_when_the_read_fails( + self, get_known, start, manager, caplog + ): + get_known.side_effect = OperationalError("SELECT", {}, Exception("connection lost")) + manager._parallelism = 2 + manager._file_queue = OrderedDict.fromkeys( + DagFileInfo(bundle_name="dags-a", rel_path=Path(rel_path), bundle_path=TEST_DAGS_FOLDER) + for rel_path in ("one.py", "two.py") + ) + + manager._start_new_processes() + + get_known.assert_called_once() + assert [call.kwargs["known_artifacts"] for call in start.call_args_list] == [[], []] + assert { + "event": "Cannot read the recorded task-handler artifacts; Dag files are parsed without them", + "log_level": "error", + } in caplog + + def test_get_known_task_handler_artifacts_from_db_skips_a_row_that_fails_validation( + self, manager, session, caplog + ): + valid = _known_artifact(JAVA_TASK_HANDLERS) + session.add_all( + [ + LangSDKTaskHandlerArtifact(**valid.model_dump(mode="json")), + LangSDKTaskHandlerArtifact( + bundle_name=JAVA_TASK_HANDLERS, + relative_fileloc="broken.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={"etl": [{"task_id": "extract", "binding": "unknown"}]}, + ), + ] + ) + session.flush() + + known = manager._get_known_task_handler_artifacts_from_db({JAVA_TASK_HANDLERS}, session=session) + + assert known == {JAVA_TASK_HANDLERS: [valid]} + assert { + "event": "Ignoring a recorded task-handler artifact that fails validation", + "bundle_name": JAVA_TASK_HANDLERS, + "relative_fileloc": "broken.jar", + "log_level": "warning", + } in caplog + + def test_get_known_task_handler_artifacts_from_db_orders_by_bundle_and_path(self, manager, session): + session.add_all( + LangSDKTaskHandlerArtifact( + bundle_name=bundle_name, + relative_fileloc=rel_path, + size_bytes=1024, + cache_digest=None, + task_handlers={}, + ) + for bundle_name, rel_path in ( + (JAVA_TASK_HANDLERS, "c.jar"), + ("dags-a", "b.jar"), + (JAVA_TASK_HANDLERS, "a.jar"), + (JAVA_TASK_HANDLERS, "b.jar"), + ) + ) + session.flush() + + known = manager._get_known_task_handler_artifacts_from_db( + {JAVA_TASK_HANDLERS, "dags-a"}, session=session + ) + + assert [ + (bundle_name, [artifact.relative_fileloc for artifact in artifacts]) + for bundle_name, artifacts in known.items() + ] == [("dags-a", ["b.jar"]), (JAVA_TASK_HANDLERS, ["a.jar", "b.jar", "c.jar"])] + + def test_get_known_task_handler_artifacts_from_db_reads_only_the_given_bundles(self, manager, session): + session.add_all( + LangSDKTaskHandlerArtifact( + bundle_name=name, + relative_fileloc=f"{name}.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={ + "etl": [ + { + "task_id": "extract", + "binding": "positional", + "params": [{"name": None}], + } + ] + }, + ) + for name in (JAVA_TASK_HANDLERS, "dags-a", "unconfigured") + ) + session.flush() + + known = manager._get_known_task_handler_artifacts_from_db( + {JAVA_TASK_HANDLERS, "dags-a"}, session=session + ) + + assert known == { + JAVA_TASK_HANDLERS: [_known_artifact(JAVA_TASK_HANDLERS)], + "dags-a": [_known_artifact("dags-a")], + } + + class TestMultiTeamMetrics: """Tests for team_name tag on dag processing metrics in multi-team mode.""" diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 0cb61aa02866a..2de2b89aa313b 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -27,13 +27,13 @@ import uuid import zipfile from collections.abc import Callable, Iterable -from socket import socketpair +from socket import socket, socketpair from typing import TYPE_CHECKING, Any, BinaryIO from unittest.mock import MagicMock, PropertyMock, patch import pytest import structlog -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from sqlalchemy import select from structlog.typing import FilteringBoundLogger @@ -54,11 +54,18 @@ from airflow.dag_processing.dagbag import DagBag from airflow.dag_processing.manager import DagFileProcessorManager, process_parse_results from airflow.dag_processing.processor import ( + BaseDagFileProcessorProcess, DagFileParseRequest, DagFileParsingResult, DagFileProcessorProcess, + TaskHandlerArtifact, + TaskHandlerDeclaration, + TaskHandlerParam, + TaskHandlerParseRequest, + TaskHandlerParsingResult, ToDagProcessor, ToManager, + ToSDKTaskHandlerProcessor, _execute_callbacks, _execute_dag_callbacks, _execute_email_callbacks, @@ -615,6 +622,82 @@ def test_parse_file_entrypoint_parses_dag_callbacks(mocker): ] +def _make_known_artifact_body(**overrides) -> dict: + return { + "bundle_name": "java-task-handlers", + "relative_fileloc": "etl.jar", + "size_bytes": 1024, + "cache_digest": "ab12", + "task_handlers": { + "etl": [{"task_id": "extract", "binding": "positional", "params": [{"name": None}]}] + }, + **overrides, + } + + +def test_parse_file_entrypoint_decodes_known_artifacts(): + r, w = socketpair() + frame = comms._ResponseFrame( + id=0, + body={ + "file": "/files/dags/etl.py", + "bundle_path": "/files/dags", + "bundle_name": "testing", + "known_artifacts": [ + _make_known_artifact_body(), + _make_known_artifact_body(relative_fileloc="empty.jar", cache_digest=None, task_handlers={}), + ], + "type": "DagFileParseRequest", + }, + ) + w.sendall(frame.as_bytes()) + + # The same decoder _parse_file_entrypoint builds. + decoder = comms.CommsDecoder[ToDagProcessor, ToManager]( + socket=r, + body_decoder=TypeAdapter[ToDagProcessor](ToDagProcessor), + ) + + msg = decoder._get_response() + assert isinstance(msg, DagFileParseRequest) + assert msg.known_artifacts == [ + TaskHandlerArtifact( + bundle_name="java-task-handlers", + relative_fileloc="etl.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={ + "etl": [ + TaskHandlerDeclaration( + task_id="extract", + binding="positional", + params=[TaskHandlerParam(name=None)], + ) + ] + }, + ), + TaskHandlerArtifact( + bundle_name="java-task-handlers", + relative_fileloc="empty.jar", + size_bytes=1024, + cache_digest=None, + task_handlers={}, + ), + ] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + pytest.param("cache_digest", "a" * 129, id="cache-digest"), + pytest.param("relative_fileloc", "a" * 2001, id="rel-path"), + ], +) +def test_known_artifact_rejects_values_wider_than_their_column(field, value): + with pytest.raises(ValidationError, match=field): + TaskHandlerArtifact.model_validate(_make_known_artifact_body(**{field: value})) + + def test_parse_file_with_dag_callbacks(spy_agency): from airflow import DAG @@ -2354,15 +2437,70 @@ def get_type_names(union_type): + "\n\nEither handle these types in ToDagProcessor or update in_task_runner_but_not_in_dag_processing_process list." ) + def test_parse_request_unions_differ_only_in_request(self): + dag_processor_types = set(typing.get_args(typing.get_args(ToDagProcessor)[0])) + task_handler_types = set(typing.get_args(typing.get_args(ToSDKTaskHandlerProcessor)[0])) + + assert dag_processor_types - task_handler_types == {DagFileParseRequest} + assert task_handler_types - dag_processor_types == {TaskHandlerParseRequest} + + +class TestTaskHandlerDeclaration: + @pytest.mark.parametrize( + ("binding", "param", "expected"), + [ + pytest.param("positional", {"name": None}, TaskHandlerParam(name=None), id="positional"), + pytest.param( + "named", + {"name": "day", "exact_name": True}, + TaskHandlerParam(name="day", exact_name=True), + id="named", + ), + ], + ) + def test_decodes_binding(self, binding, param, expected): + declaration = TaskHandlerDeclaration.model_validate( + {"task_id": "extract", "binding": binding, "params": [param]} + ) + + assert declaration.binding == binding + assert declaration.params == [expected] + + def test_decodes_unlisted_params(self): + declaration = TaskHandlerDeclaration.model_validate( + {"task_id": "extract", "binding": "named", "params": None} + ) + + assert declaration.params is None + + @pytest.mark.parametrize( + "declaration", + [ + pytest.param({"task_id": "extract", "params": []}, id="missing"), + pytest.param({"task_id": "extract", "binding": "keyword", "params": []}, id="unknown"), + pytest.param( + {"task_id": "extract", "binding": "named_or_whole", "params": []}, id="named_or_whole" + ), + pytest.param({"task_id": "extract", "binding": "named_open", "params": []}, id="named_open"), + ], + ) + def test_rejects_invalid_binding(self, declaration): + with pytest.raises(ValidationError, match="binding"): + TaskHandlerDeclaration.model_validate(declaration) + class TestDagFileProcessorProcess: def test_registered_message_types(self): - expected = set(typing.get_args(typing.get_args(ToManager)[0])) + expected = set(typing.get_args(typing.get_args(ToManager)[0])) - {TaskHandlerParsingResult} assert set(DagFileProcessorProcess._request_handlers) == expected @pytest.mark.parametrize( "message_type", - sorted(set(typing.get_args(typing.get_args(ToManager)[0])) - {DagFileParsingResult}, key=str), + sorted( + set(typing.get_args(typing.get_args(ToManager)[0])) + - {DagFileParsingResult, TaskHandlerParsingResult}, + key=str, + ), ids=lambda message_type: message_type.__name__, ) def test_reuses_shared_request_handlers(self, message_type): @@ -2399,6 +2537,38 @@ def test_dispatch_parsing_result(self, send_msg, proc): assert proc.parsing_result is result send_msg.assert_called_once_with(proc, None, request_id=42, error=None) + @patch.object(DagFileProcessorProcess, "send_msg", autospec=True) + def test_rejects_task_handler_parsing_result(self, send_msg, proc): + msg = proc.decoder.validate_python( + { + "type": "TaskHandlerParsingResult", + "fileloc": "handlers.jar", + "task_handlers": { + "etl": [ + { + "task_id": "extract", + "binding": "positional", + "params": [ + {"name": "day", "value_schema": {"type": "string"}}, + {"name": "limit"}, + ], + } + ] + }, + } + ) + assert isinstance(msg, TaskHandlerParsingResult) + + proc._handle_request(msg, structlog.get_logger(), req_id=42) + + send_msg.assert_called_once_with( + proc, + None, + request_id=42, + error=comms.ErrorResponse(detail={"status_code": 400, "message": "Unhandled request"}), + ) + assert proc.parsing_result is None + @patch.object(DagFileProcessorProcess, "send_msg", autospec=True) def test_previous_successful_run_uses_process_id(self, send_msg, proc): proc.client.task_instances.get_previous_successful_dagrun.return_value = ( @@ -2546,6 +2716,25 @@ def test_handle_request_get_variable_masks_value_with_key(self, proc): } +class TestBaseDagFileProcessorProcess: + @patch.object(BaseDagFileProcessorProcess, "cleanup_sockets_after_kill", autospec=True) + def test_close_without_log_file_handle_cleans_up_sockets(self, cleanup_sockets_after_kill): + proc = BaseDagFileProcessorProcess( + process_log=structlog.get_logger(), + id=uuid.uuid4(), + pid=1234, + process=MagicMock(spec=supervisor.ProcessTracker), + stdin=MagicMock(spec=socket), + client=MagicMock(spec=Client), + bundle_name="mybundle", + dag_file_rel_path="dags/my_dag.py", + ) + + proc.close() + + cleanup_sockets_after_kill.assert_called_once_with(proc) + + class TestMultiTeamCallbackMetrics: """Tests for team_name tag on dag.callback_exceptions in multi-team mode.""" diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py new file mode 100644 index 0000000000000..6b39cfd89bcc5 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor.py @@ -0,0 +1,855 @@ +# +# 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. +from __future__ import annotations + +import contextlib +import io +import os +import selectors +import signal +import socket +import sys +import threading +import time +import uuid +from pathlib import Path +from unittest.mock import ANY, MagicMock, patch + +import psutil +import pytest +import structlog +from pydantic import TypeAdapter +from structlog.typing import FilteringBoundLogger + +from airflow.configuration import conf +from airflow.dag_processing.processor import ( + BaseDagFileProcessorProcess, + DagFileParseRequest, + DagFileParsingResult, + DagFileProcessorProcess, + TaskHandlerDeclaration, + TaskHandlerParam, + TaskHandlerParseRequest, + TaskHandlerParsingResult, + ToDagProcessor, + ToManager, +) +from airflow.dag_processing.task_handler_processor import ( + LangSDKRuntimeSchemaVersion, + LangSDKTaskHandlerProcessorProcess, + _get_import_timeout, +) +from airflow.sdk.api.client import Client, VariableOperations +from airflow.sdk.api.datamodels._generated import VariableResponse +from airflow.sdk.coordinators._subprocess import SubprocessCoordinator +from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType +from airflow.sdk.execution_time import supervisor, task_runner +from airflow.sdk.execution_time.comms import ( + CommsDecoder, + ErrorResponse, + GetVariable, + MaskSecret, + VariableResult, + _RequestFrame, +) + +from tests_common.test_utils.config import conf_vars +from unit.dag_processing.fake_task_handler_runtime import ( + FakeCoordinator, + fake_coordinator, + play_runtime, + write_artifact, +) + +# The oldest supervisor schema version, so the parse request is downgraded. +OLDEST_SCHEMA_VERSION = "2026-06-16" + + +def _reply_with(*task_ids: str, **result): + """Reply with a handler for each of *task_ids* under the Dag ``etl``.""" + + def reply(request: TaskHandlerParseRequest, comms) -> TaskHandlerParsingResult: + declarations = [ + TaskHandlerDeclaration(task_id=task_id, binding="positional", params=[]) for task_id in task_ids + ] + return TaskHandlerParsingResult(fileloc=request.file, task_handlers={"etl": declarations}, **result) + + return reply + + +def _get_task_ids(result: TaskHandlerParsingResult) -> list[str]: + return [declaration.task_id for declaration in result.task_handlers["etl"]] + + +def _get_open_fds() -> set[int]: + # Without /proc, as on macOS, this is empty, so the fd leak checks pass trivially. + return {int(fd) for fd in os.listdir("/proc/self/fd")} if os.path.isdir("/proc/self/fd") else set() + + +# A runtime that exits once a child has left its process group; the child keeps the runtime's +# stdout and stderr open and writes its pid to argv[1]. +_LEAVE_A_CHILD_OUTSIDE_THE_GROUP = """ +import os, sys, time +read_fd, write_fd = os.pipe() +if os.fork(): + os.read(read_fd, 1) + os._exit(0) +os.setsid() +with open(sys.argv[1], "w") as f: + f.write(str(os.getpid())) +os.write(write_fd, b"x") +time.sleep(30) +""" + + +def _is_running(pid: int) -> bool: + try: + return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +@pytest.fixture(autouse=True) +def _coordinator(): + with fake_coordinator(): + yield + + +@pytest.fixture +def supervisor_comms(monkeypatch): + """Give this process a supervisor channel, as a Dag-parsing child has.""" + comms = MagicMock(spec=CommsDecoder) + monkeypatch.setattr(task_runner, "SUPERVISOR_COMMS", comms, raising=False) + return comms + + +def _start(tmp_path, selector, *, client: Client | None = None, **spec) -> LangSDKTaskHandlerProcessorProcess: + return LangSDKTaskHandlerProcessorProcess.start( + id=uuid.uuid4(), + coordinator="fake", + path=write_artifact(tmp_path / "etl.artifact", **spec), + bundle_path=tmp_path, + bundle_name="task-handlers", + artifact_rel_path="etl.artifact", + selector=selector, + logger=structlog.get_logger(), + client=client, + ) + + +@pytest.fixture +def parse(tmp_path): + """Probe ``etl.artifact`` under a caller's selector loop, and check that nothing is left open.""" + + def _parse(**kwargs) -> LangSDKTaskHandlerProcessorProcess: + fds_before = _get_open_fds() + with selectors.DefaultSelector() as selector: + proc = _start(tmp_path, selector, **kwargs) + deadline = time.monotonic() + 30 + while not proc.is_ready: + assert time.monotonic() < deadline, "the Lang-SDK parse did not finish" + proc._service_subprocess(max_wait_time=0.1) + assert selector.get_map() == {} + proc.close() + assert _get_open_fds() <= fds_before + return proc + + return _parse + + +def _block_until_killed(comms) -> None: + """Hang as a stuck runtime does, until it is killed or the parent closes the comm socket.""" + comms.socket.recv(1) + + +def _send_an_invalid_frame(request, comms) -> None: + comms.socket.sendall(bytes.fromhex("00000003c1c1c1")) + _block_until_killed(comms) + + +def _send_a_frame(comms, body: dict) -> None: + comms.socket.sendall(_RequestFrame(id=1, body=body).as_bytes()) + _block_until_killed(comms) + + +def _send_an_invalid_result(request, comms) -> None: + _send_a_frame( + comms, {"type": "TaskHandlerParsingResult", "fileloc": request.file, "task_handlers": "none"} + ) + + +def _send_a_start_message(request, comms) -> None: + _send_a_frame(comms, LangSDKRuntimeSchemaVersion(schema_version=None).model_dump()) + + +def _run(tmp_path, *, coordinator: str = "fake", **spec) -> TaskHandlerParsingResult: + return LangSDKTaskHandlerProcessorProcess.run( + coordinator=coordinator, + path=write_artifact(tmp_path / "etl.artifact", **spec), + bundle_path=tmp_path, + bundle_name="task-handlers", + artifact_rel_path="etl.artifact", + logger=structlog.get_logger(), + ) + + +class TestLangSDKTaskHandlerProcessorProcess: + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_collects_the_result_the_runtime_returns( + self, mock_parse_task_handler, parse, tmp_path, cap_structlog + ): + mock_parse_task_handler.side_effect = play_runtime( + _reply_with("extract"), + schema_version=OLDEST_SCHEMA_VERSION, + log_lines=[{"event": "Registering handlers", "level": "info"}], + ) + + proc = parse() + + assert proc.parsing_result.fileloc == os.fspath(tmp_path / "etl.artifact") + assert proc.parsing_result.import_errors is None + assert _get_task_ids(proc.parsing_result) == ["extract"] + assert proc._subprocess_schema_version == OLDEST_SCHEMA_VERSION + assert "Registering handlers" in cap_structlog + + @patch("airflow.dag_processing.task_handler_processor._is_connection_from_pid", autospec=True) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_connection_is_used_once_it_is_verified(self, mock_parse_task_handler, mock_owned, tmp_path): + mock_parse_task_handler.side_effect = play_runtime(_reply_with("extract")) + mock_owned.return_value = False + with selectors.DefaultSelector() as selector: + proc = _start(tmp_path, selector) + child_stdin = proc.stdin + deadline = time.monotonic() + 30 + while len(proc._unverified_connections) < 2: + assert proc.stdin is child_stdin, "an unverified connection was used" + assert time.monotonic() < deadline, "the runtime did not connect" + proc._service_subprocess(max_wait_time=0.1) + assert proc.stdin is child_stdin + + mock_owned.return_value = True + while not proc.is_ready: + assert time.monotonic() < deadline, "the Lang-SDK parse did not finish" + proc._service_subprocess(max_wait_time=0.1) + proc.close() + + assert _get_task_ids(proc.parsing_result) == ["extract"] + + @patch.object( + LangSDKTaskHandlerProcessorProcess, + "send_msg", + autospec=True, + side_effect=OSError("cannot send the start request"), + ) + def test_a_start_that_fails_after_the_fork_leaves_nothing_behind(self, mock_send_msg, tmp_path): + fds_before = _get_open_fds() + children_before = {child.pid for child in psutil.Process().children()} + + with selectors.DefaultSelector() as selector: + with pytest.raises(OSError, match="cannot send the start request"): + _start(tmp_path, selector) + assert selector.get_map() == {} + + assert {child.pid for child in psutil.Process().children()} == children_before + assert _get_open_fds() <= fds_before + + @patch.object(BaseDagFileProcessorProcess, "start", autospec=True, side_effect=OSError("fork failed")) + def test_a_start_that_fails_before_the_fork_closes_its_listeners(self, mock_start, tmp_path): + fds_before = _get_open_fds() + + with selectors.DefaultSelector() as selector: + with pytest.raises(OSError, match="fork failed"): + _start(tmp_path, selector) + + listeners = mock_start.call_args.kwargs["listeners"] + assert [listener.fileno() for listener in listeners.values()] == [-1, -1] + assert _get_open_fds() <= fds_before + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_requests_are_answered_by_the_client(self, mock_parse_task_handler, parse): + def reply(request, comms): + variable = comms.send(GetVariable(key="probe_var")) + return _reply_with(variable.value)(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + client = MagicMock(spec=Client) + client.variables = MagicMock(spec=VariableOperations) + client.variables.get.return_value = VariableResponse(key="probe_var", value="from-the-client") + + proc = parse(client=client) + + assert _get_task_ids(proc.parsing_result) == ["from-the-client"] + + @pytest.mark.parametrize( + ("spec", "reply", "error"), + [ + pytest.param( + {"command_error": "no runtime"}, + None, + "Cannot start the Lang-SDK runtime: FileNotFoundError: no runtime", + id="command-not-resolved", + ), + pytest.param( + {"argv": ["/no/such/runtime"]}, + None, + "Cannot start the Lang-SDK runtime: FileNotFoundError: " + "[Errno 2] No such file or directory: '/no/such/runtime'", + id="exec-failed", + ), + pytest.param( + {"argv": ["/bin/sh", "-c", "exit 3"]}, + None, + "The Lang-SDK runtime exited with code 3 without a parse result", + id="exits-before-connecting", + ), + pytest.param( + {}, + lambda request, comms: None, + "The Lang-SDK runtime exited with code 0 without a parse result", + id="exits-without-a-result", + ), + pytest.param( + {}, + _send_an_invalid_frame, + "The Lang-SDK runtime sent an invalid frame: MessagePack data is malformed: " + "invalid opcode '\\xc1' (byte 0)", + id="invalid-frame", + ), + ], + ) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_failed_parse_is_an_import_error( + self, mock_parse_task_handler, parse, tmp_path, spec, reply, error + ): + mock_parse_task_handler.side_effect = ( + play_runtime(reply) if reply else SubprocessCoordinator.parse_task_handler + ) + + proc = parse(**spec) + + assert proc.parsing_result == TaskHandlerParsingResult( + fileloc=os.fspath(tmp_path / "etl.artifact"), + task_handlers={}, + import_errors={"etl.artifact": error}, + ) + + def test_a_coordinator_that_is_not_configured_is_an_import_error(self, tmp_path): + result = _run(tmp_path, coordinator="missing") + + assert result == TaskHandlerParsingResult( + fileloc=os.fspath(tmp_path / "etl.artifact"), + task_handlers={}, + import_errors={ + "etl.artifact": "Cannot start the Lang-SDK runtime: " + "InvalidCoordinatorError: No coordinator 'missing' in [sdk] coordinators" + }, + ) + + @patch.object( + FakeCoordinator, + "_build_parse_task_handler_command", + SubprocessCoordinator._build_parse_task_handler_command, + ) + def test_a_coordinator_that_does_not_parse_task_handlers_is_an_import_error(self, parse): + proc = parse() + + assert proc.parsing_result.import_errors == { + "etl.artifact": "Cannot start the Lang-SDK runtime: " + "NotImplementedError: FakeCoordinator does not parse task handlers" + } + + @pytest.mark.parametrize( + ("reply", "error"), + [ + pytest.param( + _send_an_invalid_result, + "TaskHandlerParsingResult.task_handlers\n Input should be a valid dictionary", + id="result", + ), + pytest.param( + _send_a_start_message, "does not match any of the expected tags", id="start-message" + ), + ], + ) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_message_that_does_not_validate_is_an_import_error( + self, mock_parse_task_handler, parse, reply, error + ): + mock_parse_task_handler.side_effect = play_runtime(reply) + + proc = parse() + + [message] = proc.parsing_result.import_errors.values() + assert message.startswith("The Lang-SDK runtime sent a message that does not validate: ") + assert error in message + assert proc._exit_code == -signal.SIGKILL + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_dag_file_parsing_result_is_not_a_task_handler_result(self, mock_parse_task_handler, tmp_path): + def reply(request, comms): + with pytest.raises(AirflowRuntimeError, match="Unhandled request"): + comms.send(DagFileParsingResult(fileloc=request.file, serialized_dags=[])) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path) + + assert result.import_errors == { + "etl.artifact": "The Lang-SDK runtime exited with code 0 without a parse result" + } + + @patch.object( + FakeCoordinator, "parse_task_handler", autospec=True, side_effect=play_runtime(_send_an_invalid_frame) + ) + def test_killing_the_runtime_is_not_reported_as_out_of_memory( + self, mock_parse_task_handler, parse, cap_structlog + ): + proc = parse() + + assert proc._exit_code == -signal.SIGKILL + assert not any("Likely out of memory" in str(entry.get("event")) for entry in cap_structlog.entries) + + @patch("airflow.dag_processing.task_handler_processor._EXIT_GRACE_PERIOD", 0.5) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_runtime_that_runs_on_after_its_result_is_killed( + self, mock_parse_task_handler, parse, cap_structlog + ): + def reply(request, comms): + comms.send(_reply_with("extract")(request, comms)) + _block_until_killed(comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == ["extract"] + assert proc._exit_code == -signal.SIGKILL + assert "The Lang-SDK runtime did not exit after its parse result; killing it" in cap_structlog + + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=5) + def test_what_the_runtime_leaves_in_its_process_group_is_killed_when_it_exits( + self, mock_timeout, parse, tmp_path + ): + pid_file = tmp_path / "leftover.pid" + + proc = parse(argv=["/bin/sh", "-c", f"sleep 30 & echo $! > {pid_file}; exit 0"]) + + leftover = int(pid_file.read_text()) + try: + assert proc.parsing_result.import_errors == { + "etl.artifact": "The Lang-SDK runtime exited with code 0 without a parse result" + } + # A poll: the leftover is not a child to wait for, and time_machine cannot reach another process. + deadline = time.monotonic() + 10 + while _is_running(leftover): + assert time.monotonic() < deadline, "the runtime's leftover process is still running" + time.sleep(0.05) + finally: + with contextlib.suppress(ProcessLookupError): + os.kill(leftover, signal.SIGKILL) + + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + def test_the_import_timeout_holds_after_the_runtime_exits(self, mock_timeout, parse, tmp_path): + pid_file = tmp_path / "leftover.pid" + + try: + proc = parse(argv=[sys.executable, "-c", _LEAVE_A_CHILD_OUTSIDE_THE_GROUP, os.fspath(pid_file)]) + finally: + if pid_file.exists(): + os.kill(int(pid_file.read_text()), signal.SIGKILL) + + assert proc.parsing_result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s" + } + assert proc._exit_code == 0 + + @pytest.mark.parametrize( + ("policy", "error"), + [ + pytest.param(RuntimeError("policy bug"), "RuntimeError: policy bug", id="raises"), + pytest.param( + lambda path: "30", + "TypeError: Value (30) from get_dagbag_import_timeout must be int or float", + id="not-a-number", + ), + ], + ) + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True) + def test_a_failing_import_timeout_policy_is_an_import_error(self, mock_timeout, parse, policy, error): + mock_timeout.side_effect = policy + + proc = parse() + + assert proc.parsing_result.import_errors == { + "etl.artifact": f"Cannot start the Lang-SDK runtime: {error}" + } + + @pytest.mark.skipif(not Path("/proc/self/fd").is_dir(), reason="reads /proc") + @pytest.mark.parametrize("use_exec", [False, True], ids=["fork", "spawn"]) + def test_the_runtime_inherits_only_its_standard_streams(self, monkeypatch, tmp_path, use_exec): + if use_exec: + # The spawned interpreter finds the coordinator again from its environment. + monkeypatch.setattr(supervisor, "_should_use_exec", lambda: True) + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(sys.path)) + monkeypatch.setenv("AIRFLOW__SDK__COORDINATORS", conf.get("sdk", "coordinators")) + with selectors.DefaultSelector() as selector: + proc = _start(tmp_path, selector, argv=["/bin/sh", "-c", "exec sleep 30"]) + deadline = time.monotonic() + 30 + while psutil.Process(proc.pid).name() != "sleep": + assert time.monotonic() < deadline, "the runtime did not start" + proc._service_subprocess(max_wait_time=0.1) + fd_dir = Path(f"/proc/{proc.pid}/fd") + fds = {fd.name: os.readlink(fd) for fd in fd_dir.iterdir()} + proc.kill(signal.SIGKILL) + proc.close() + + assert sorted(fds) == ["0", "1", "2"] + assert fds["0"] == "/dev/null" + + +def _probe_from_a_dag_parsing_child() -> None: + """Stand in for ``_parse_file_entrypoint``: probe the file it is asked to parse, and return the result.""" + comms_decoder = CommsDecoder[ToDagProcessor, ToManager](body_decoder=TypeAdapter(ToDagProcessor)) + request = comms_decoder._get_response() + assert isinstance(request, DagFileParseRequest) + task_runner.SUPERVISOR_COMMS = comms_decoder # type: ignore[assignment] + + result = LangSDKTaskHandlerProcessorProcess.run( + coordinator="fake", + path=request.file, + bundle_path=request.bundle_path, + bundle_name=request.bundle_name, + artifact_rel_path="etl.artifact", + logger=structlog.get_logger(logger_name="task"), + ) + comms_decoder.send( + DagFileParsingResult( + fileloc=request.file, serialized_dags=[], warnings=[result.model_dump(mode="json")] + ) + ) + + +class TestRequestsWithoutAClient: + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_request_is_relayed_to_the_parent(self, mock_parse_task_handler, parse, supervisor_comms): + def reply(request, comms): + variable = comms.send(GetVariable(key="probe_var")) + return _reply_with(variable.value)(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply, schema_version=OLDEST_SCHEMA_VERSION) + # As decoded from the parent's frame, which always carries its type. + supervisor_comms.send.return_value = VariableResult.model_validate( + {"key": "probe_var", "value": "from-the-parent", "type": "VariableResult"} + ) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == ["from-the-parent"] + supervisor_comms.send.assert_called_once_with(GetVariable(key="probe_var")) + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_the_parent_s_error_reaches_the_runtime(self, mock_parse_task_handler, parse, supervisor_comms): + def reply(request, comms): + with pytest.raises(AirflowRuntimeError) as ctx: + comms.send(GetVariable(key="probe_var")) + error = ctx.value.error + return _reply_with(f"{error.error.value}:{error.detail['key']}")(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + supervisor_comms.send.side_effect = AirflowRuntimeError( + ErrorResponse(error=ErrorType.VARIABLE_NOT_FOUND, detail={"key": "probe_var"}) + ) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == [f"{ErrorType.VARIABLE_NOT_FOUND.value}:probe_var"] + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_request_gets_an_error_without_a_parent(self, mock_parse_task_handler, monkeypatch, tmp_path): + monkeypatch.delattr(task_runner, "SUPERVISOR_COMMS", raising=False) + + def reply(request, comms): + with pytest.raises(AirflowRuntimeError) as ctx: + comms.send(GetVariable(key="probe_var")) + return _reply_with(ctx.value.error.detail["message"])(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path) + + assert _get_task_ids(result) == ["GetVariable is answered only in the Dag processor"] + + @patch("airflow.sdk.log._secrets_masker", autospec=True) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_secret_is_masked_here_and_by_the_parent( + self, mock_parse_task_handler, mock_secrets_masker, parse, supervisor_comms + ): + def reply(request, comms): + comms.send(MaskSecret(value="probe-secret", name="probe_conn")) + return _reply_with("extract")(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + proc = parse() + + assert _get_task_ids(proc.parsing_result) == ["extract"] + mock_secrets_masker.return_value.add_mask.assert_called_once_with("probe-secret", "probe_conn") + supervisor_comms.send.assert_called_once_with(MaskSecret(value="probe-secret", name="probe_conn")) + + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_runtime_request_is_answered_by_the_dag_processor(self, mock_parse_task_handler, tmp_path): + def reply(request, comms): + variable = comms.send(GetVariable(key="probe_var")) + return _reply_with(variable.value)(request, comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + client = MagicMock(spec=Client) + client.variables = MagicMock(spec=VariableOperations) + client.variables.get.return_value = VariableResponse(key="probe_var", value="from-the-dag-processor") + artifact = write_artifact(tmp_path / "etl.artifact") + + proc = DagFileProcessorProcess.start( + id=1, + path=artifact, + bundle_path=tmp_path, + bundle_name="task-handlers", + dag_file_rel_path="etl.artifact", + callbacks=[], + target=_probe_from_a_dag_parsing_child, + logger=structlog.get_logger(), + logger_filehandle=io.BytesIO(), + client=client, + ) + deadline = time.monotonic() + 30 + while not proc.is_ready: + assert time.monotonic() < deadline, "the Dag-parsing child did not finish" + proc._service_subprocess(max_wait_time=0.1) + proc.close() + + client.variables.get.assert_called_once_with("probe_var") + [probe_result] = proc.parsing_result.warnings + assert TaskHandlerParsingResult.model_validate(probe_result) == TaskHandlerParsingResult( + fileloc=os.fspath(artifact), + task_handlers={ + "etl": [ + TaskHandlerDeclaration(task_id="from-the-dag-processor", binding="positional", params=[]) + ] + }, + ) + + +class TestRun: + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_returns_every_handler_the_runtime_registers(self, mock_parse_task_handler, tmp_path): + def reply(request, comms): + # Echo what the request carries, so the test can check it. + param = TaskHandlerParam(name=request.bundle_name, value_schema={"type": "string"}) + return TaskHandlerParsingResult( + fileloc=request.file, + task_handlers={ + dag_id: [ + TaskHandlerDeclaration( + task_id=os.fspath(request.bundle_path), binding="positional", params=[param] + ) + ] + for dag_id in ("etl", "report") + }, + ) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path) + + param = TaskHandlerParam(name="task-handlers", value_schema={"type": "string"}) + declaration = TaskHandlerDeclaration( + task_id=os.fspath(tmp_path), binding="positional", params=[param] + ) + assert result == TaskHandlerParsingResult( + fileloc=os.fspath(tmp_path / "etl.artifact"), + task_handlers={"etl": [declaration], "report": [declaration]}, + ) + + @pytest.mark.execution_timeout(30) + @pytest.mark.parametrize("connected", [False, True], ids=["before-connecting", "after-connecting"]) + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + @patch.object( + LangSDKTaskHandlerProcessorProcess, + "close", + autospec=True, + side_effect=LangSDKTaskHandlerProcessorProcess.close, + ) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_parse_past_the_import_timeout_is_killed( + self, mock_parse_task_handler, mock_close, mock_timeout, tmp_path, connected + ): + mock_parse_task_handler.side_effect = ( + play_runtime(lambda request, comms: _block_until_killed(comms)) + if connected + else SubprocessCoordinator.parse_task_handler + ) + # A runtime that never connects leaves both listeners open when it is killed. + fds_before = _get_open_fds() + + result = _run(tmp_path, argv=["/bin/sh", "-c", "exec sleep 60"]) + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s" + } + [proc] = [c.args[0] for c in mock_close.call_args_list] + assert proc._exit_code == -9 + assert not proc._open_sockets + assert _get_open_fds() <= fds_before + + @pytest.mark.execution_timeout(30) + @patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1) + @patch.object(FakeCoordinator, "parse_task_handler", autospec=True) + def test_a_runtime_that_stops_mid_frame_is_timed_out( + self, mock_parse_task_handler, mock_timeout, tmp_path + ): + def reply(request, comms): + comms.socket.sendall((100).to_bytes(4, byteorder="big")) + _block_until_killed(comms) + + mock_parse_task_handler.side_effect = play_runtime(reply) + + result = _run(tmp_path) + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s" + } + + @pytest.mark.execution_timeout(30) + @conf_vars({("dag_processor", "dag_file_processor_timeout"): "1"}) + @patch.object( + FakeCoordinator, + "_build_parse_task_handler_command", + autospec=True, + side_effect=lambda self, *, path: threading.Event().wait(), + ) + def test_the_dag_file_processor_timeout_applies_until_the_import_timeout_is_reported( + self, mock_build_command, tmp_path + ): + result = _run(tmp_path) + + assert result.import_errors == { + "etl.artifact": f"The Lang-SDK runtime did not parse {tmp_path / 'etl.artifact'} within 1.0s" + } + + +@pytest.mark.parametrize(("configured", "expected"), [(30, 30), (0.5, 0.5), (0, None), (-1, None)]) +@patch("airflow.settings.get_dagbag_import_timeout", autospec=True) +def test_only_a_positive_import_timeout_applies(mock_timeout, configured, expected): + mock_timeout.return_value = configured + + assert _get_import_timeout("/b/etl.artifact") == expected + mock_timeout.assert_called_once_with("/b/etl.artifact") + + +def _make_process(**kwargs) -> LangSDKTaskHandlerProcessorProcess: + return LangSDKTaskHandlerProcessorProcess( + id=uuid.uuid4(), + pid=1, + stdin=MagicMock(spec=socket.socket), + process=MagicMock(spec=supervisor.ProcessTracker), + process_log=MagicMock(spec=FilteringBoundLogger), + selector=MagicMock(spec=selectors.BaseSelector), + bundle_name="task-handlers", + dag_file_rel_path="etl.artifact", + coordinator="fake", + listeners={}, + parse_request=TaskHandlerParseRequest( + file="/b/etl.artifact", bundle_path=Path("/b"), bundle_name="task-handlers" + ), + **kwargs, + ) + + +def _build_result(*task_ids: str) -> TaskHandlerParsingResult: + return TaskHandlerParsingResult( + fileloc="/b/etl.artifact", + task_handlers={ + "etl": [ + TaskHandlerDeclaration(task_id=task_id, binding="positional", params=[]) + for task_id in task_ids + ] + }, + ) + + +@patch.object(LangSDKTaskHandlerProcessorProcess, "send_msg", autospec=True) +def test_the_first_parse_result_wins(mock_send_msg): + proc = _make_process() + + proc._handle_request(_build_result("first"), structlog.get_logger(), 1) + proc._handle_request(_build_result("second"), structlog.get_logger(), 2) + + assert _get_task_ids(proc.parsing_result) == ["first"] + assert mock_send_msg.call_args.kwargs["error"].detail == { + "message": "A parse result was already received" + } + + +@pytest.mark.parametrize( + "invalid_frame", + [ + pytest.param(bytes.fromhex("00000003c1c1c1"), id="does-not-decode"), + pytest.param( + _RequestFrame( + id=2, + body={ + "type": "TaskHandlerParsingResult", + "fileloc": "/b/etl.artifact", + "task_handlers": "none", + }, + ).as_bytes(), + id="does-not-validate", + ), + ], +) +@patch.object(LangSDKTaskHandlerProcessorProcess, "_kill_runtime", autospec=True) +@patch.object(LangSDKTaskHandlerProcessorProcess, "send_msg", autospec=True) +def test_an_invalid_message_after_the_parse_result_keeps_it(mock_send_msg, mock_kill_runtime, invalid_frame): + proc = _make_process() + runtime, conn = socket.socketpair() + with runtime, conn: + proc._register_comm(conn) + read_frame, _ = proc.selector.register.call_args.args[2] + runtime.sendall(_RequestFrame(id=1, body=_build_result("extract").model_dump(mode="json")).as_bytes()) + assert read_frame(conn) + runtime.sendall(invalid_frame) + assert not read_frame(conn) + + assert _get_task_ids(proc.parsing_result) == ["extract"] + assert proc.parsing_result.import_errors is None + proc.process_log.warning.assert_called_once_with( + "Ignoring an invalid message from the Lang-SDK runtime after its parse result", error=ANY + ) + mock_kill_runtime.assert_called_once_with(proc) + + +@patch.object(LangSDKTaskHandlerProcessorProcess, "send_msg", autospec=True) +def test_the_schema_version_is_reported_once(mock_send_msg): + proc = _make_process() + + proc._handle_request( + LangSDKRuntimeSchemaVersion(schema_version=OLDEST_SCHEMA_VERSION), structlog.get_logger(), 1 + ) + proc._handle_request(LangSDKRuntimeSchemaVersion(schema_version=None), structlog.get_logger(), 2) + + assert proc._runtime_schema_version == OLDEST_SCHEMA_VERSION + assert mock_send_msg.call_args.kwargs["error"].detail["message"] == "Unhandled request" diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py new file mode 100644 index 0000000000000..6dd1694d17ac9 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_go.py @@ -0,0 +1,168 @@ +# 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. +"""Pack the Go SDK example bundle, probe it for its task handlers, and read its cache digest.""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +from typing import TYPE_CHECKING +from unittest import mock + +import pytest +import structlog + +from airflow.dag_processing.processor import TaskHandlerDeclaration, TaskHandlerParam +from airflow.dag_processing.task_handler_processor import LangSDKTaskHandlerProcessorProcess +from airflow.sdk.coordinators.executable.coordinator import read_cache_digest +from airflow.sdk.execution_time import supervisor +from airflow.sdk.execution_time.coordinator import reset_coordinator_manager + +from tests_common.test_utils.config import conf_vars +from tests_common.test_utils.paths import AIRFLOW_ROOT_PATH + +if TYPE_CHECKING: + from pathlib import Path + +pytestmark = pytest.mark.skipif( + os.environ.get("AIRFLOW_LANG_SDK_REAL_PROBE_TESTS") != "1", + reason="set AIRFLOW_LANG_SDK_REAL_PROBE_TESTS=1 to build and probe a real Go bundle", +) + +STRING = {"type": "string"} +INT64 = {"type": "integer", "format": "int64"} +DOUBLE = {"type": "number", "format": "double"} + + +def _make_nullable(schema: dict) -> dict: + return {"anyOf": [schema, {"type": "null"}]} + + +GO_SDK_PATH = AIRFLOW_ROOT_PATH / "go-sdk" + + +def _pack(bundle: Path, *flags: str) -> Path: + completed = subprocess.run( + ["go", "tool", "airflow-go-pack", "--output", os.fspath(bundle), *flags, "./example/bundle"], + cwd=GO_SDK_PATH, + env={**os.environ, "CGO_ENABLED": "0"}, + capture_output=True, + text=True, + check=False, + ) + assert completed.returncode == 0, completed.stderr + return bundle + + +@pytest.fixture(scope="module") +def go_bundle(tmp_path_factory) -> Path: + if shutil.which("go") is None: + pytest.skip("needs a Go toolchain on PATH") + return _pack(tmp_path_factory.mktemp("go-task-handlers") / "example_dags") + + +@pytest.fixture(autouse=True) +def _go_coordinator(monkeypatch, tmp_path): + # Without cgo the Go SDK reads the current user from USER and HOME, and fails without them. + for name, value in (("USER", "airflow"), ("HOME", os.fspath(tmp_path))): + if not os.environ.get(name): + monkeypatch.setenv(name, value) + spec = {"go": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator"}} + reset_coordinator_manager() + try: + # A bare fork, so the parse child sees this config. + with ( + conf_vars({("sdk", "coordinators"): json.dumps(spec)}), + mock.patch.object(supervisor, "_should_use_exec", autospec=True, return_value=False), + ): + yield + finally: + reset_coordinator_manager() + + +def test_probes_the_task_handlers_of_a_packed_go_bundle(go_bundle): + result = LangSDKTaskHandlerProcessorProcess.run( + coordinator="go", + path=go_bundle, + bundle_path=go_bundle.parent, + bundle_name="go-task-handlers", + artifact_rel_path=go_bundle.name, + logger=structlog.get_logger(), + ) + + assert result.import_errors is None + assert set(result.task_handlers) == { + "simple_dag", + "concurrent_xcom_dag", + "taskflow_binding_dag", + "variable_write_dag", + } + assert [d.task_id for d in result.task_handlers["simple_dag"]] == ["extract", "transform", "load"] + + declarations = {d.task_id: d for d in result.task_handlers["taskflow_binding_dag"]} + assert declarations["via_flat_args"] == TaskHandlerDeclaration( + task_id="via_flat_args", + binding="positional", + params=[ + TaskHandlerParam(name=None, value_schema=schema) + for schema in ( + STRING, + INT64, + DOUBLE, + {"type": "boolean"}, + _make_nullable({"type": "array", "items": STRING}), + {"type": "object"}, + _make_nullable({"type": "array", "items": INT64}), + _make_nullable(STRING), + ) + ], + ) + assert declarations["via_struct_arg_tag"] == TaskHandlerDeclaration( + task_id="via_struct_arg_tag", + binding="named", + params=[ + TaskHandlerParam(name="region_code", exact_name=True, value_schema=STRING), + TaskHandlerParam(name="threshold", exact_name=True, value_schema=DOUBLE), + ], + ) + # The Dag passes this handler an argument it does not take, which only warns under named binding. + assert declarations["via_struct_more_args"] == TaskHandlerDeclaration( + task_id="via_struct_more_args", + binding="named", + params=[TaskHandlerParam(name="region_code", exact_name=True, value_schema=STRING)], + ) + assert declarations["via_flat_map"] == TaskHandlerDeclaration( + task_id="via_flat_map", + binding="named", + params=[ + TaskHandlerParam(name="Region", value_schema=STRING), + TaskHandlerParam(name="Count", value_schema=INT64), + ], + ) + + +def test_a_repack_keeps_the_cache_digest_until_a_source_byte_changes(go_bundle, tmp_path): + digest = read_cache_digest(go_bundle) + assert digest is not None + + assert read_cache_digest(_pack(tmp_path / "unchanged")) == digest + + source = tmp_path / "main.go" + source.write_bytes((GO_SDK_PATH / "example" / "bundle" / "main.go").read_bytes() + b"\n") + assert read_cache_digest(_pack(tmp_path / "changed", "--source", os.fspath(source))) != digest diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py new file mode 100644 index 0000000000000..a76ae128c17d8 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py @@ -0,0 +1,209 @@ +# +# 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. +"""Probe a bundle built with the Java SDK's Gradle plugin through a real JVM.""" + +from __future__ import annotations + +import json +import os +import re +import shutil +import subprocess +import zipfile +from typing import TYPE_CHECKING +from unittest import mock + +import pytest +import structlog + +from airflow.dag_processing.processor import ( + TaskHandlerDeclaration, + TaskHandlerParam, + TaskHandlerParsingResult, +) +from airflow.dag_processing.task_handler_processor import LangSDKTaskHandlerProcessorProcess +from airflow.sdk.coordinators.java.coordinator import _parse_manifest +from airflow.sdk.execution_time import supervisor +from airflow.sdk.execution_time.coordinator import reset_coordinator_manager + +from tests_common.pytest_plugin import AIRFLOW_ROOT_PATH +from tests_common.test_utils.config import conf_vars + +if TYPE_CHECKING: + from collections.abc import Iterator + from pathlib import Path + +JAVA_SDK_PATH = AIRFLOW_ROOT_PATH / "java-sdk" + +# The example's sources, built as a thin bundle against the SDK as an included build, so nothing is +# published. The plugin's fat mode looks the SDK up as a published airflow-sdk artifact. +_SETTINGS = """ +pluginManagement { + includeBuild("../java-sdk") +} +includeBuild("../java-sdk") +rootProject.name = "probe-example" +""" + +_BUILD = """ +plugins { + id("org.apache.airflow.sdk") +} +repositories { + mavenCentral() +} +dependencies { + annotationProcessor("org.apache.airflow:processor") + implementation("org.apache.airflow:sdk") + implementation("org.apache.airflow:jpl") +} +sourceSets { + main { + java.srcDir("../java-sdk/example/src/java") + } +} +airflowBundle { + mainClass = "org.apache.airflow.example.ExampleBundleBuilder" + fatJar = false +} +""" + +pytestmark = pytest.mark.skipif( + os.environ.get("AIRFLOW_LANG_SDK_REAL_PROBE_TESTS") != "1", + reason="set AIRFLOW_LANG_SDK_REAL_PROBE_TESTS=1 to probe real Lang-SDK artifacts", +) + + +@pytest.fixture(scope="module") +def example_bundle(tmp_path_factory) -> Path: + """Build ``java-sdk/example`` from a copy of the SDK sources, so the checkout gets no build output.""" + if not JAVA_SDK_PATH.is_dir(): + pytest.skip("the Java SDK sources are absent") + if shutil.which("java") is None or shutil.which("javac") is None: + pytest.skip("needs a JDK") + root = tmp_path_factory.mktemp("java-sdk") + sdk = root / "java-sdk" + shutil.copytree(JAVA_SDK_PATH, sdk, ignore=shutil.ignore_patterns("build", ".gradle", ".kotlin")) + project = root / "probe-example" + project.mkdir() + (project / "settings.gradle").write_text(_SETTINGS) + (project / "build.gradle").write_text(_BUILD) + subprocess.run([sdk / "gradlew", "--no-daemon", "--quiet", "bundle"], cwd=project, check=True) + return project / "build" / "bundle" + + +@pytest.fixture +def java_coordinator() -> Iterator[None]: + spec = {"java": {"classpath": "airflow.sdk.coordinators.java.JavaCoordinator"}} + reset_coordinator_manager() + try: + # The parse child is a bare fork even on macOS, so it sees this config. + with ( + conf_vars({("sdk", "coordinators"): json.dumps(spec)}), + mock.patch.object(supervisor, "_should_use_exec", autospec=True, return_value=False), + ): + yield + finally: + reset_coordinator_manager() + + +def _make_nullable(schema: dict) -> dict: + return {"anyOf": [schema, {"type": "null"}]} + + +INT32 = {"type": "integer", "format": "int32"} +INT64 = {"type": "integer", "format": "int64"} +FLOAT = {"type": "number", "format": "float"} +DOUBLE = {"type": "number", "format": "double"} +STRING = {"type": "string"} + + +def _declare_positional(task_id: str, *params: tuple[str, dict]) -> TaskHandlerDeclaration: + return TaskHandlerDeclaration( + task_id=task_id, + binding="positional", + params=[TaskHandlerParam(name=name, value_schema=schema) for name, schema in params], + ) + + +def _declare_named(task_id: str, *params: tuple[str, dict, bool]) -> TaskHandlerDeclaration: + return TaskHandlerDeclaration( + task_id=task_id, + binding="named", + params=[ + TaskHandlerParam(name=name, value_schema=schema, exact_name=exact) + for name, schema, exact in params + ], + ) + + +@pytest.mark.usefixtures("java_coordinator") +def test_a_built_bundle_declares_its_task_handlers(example_bundle): + jar = example_bundle / "probe-example.jar" + + result = LangSDKTaskHandlerProcessorProcess.run( + coordinator="java", + path=jar, + bundle_path=example_bundle, + bundle_name="java-task-handlers", + artifact_rel_path=jar.name, + logger=structlog.get_logger(), + ) + + assert result == TaskHandlerParsingResult( + fileloc=os.fspath(jar), + task_handlers={ + "java_xcom_casting_example": [ + _declare_named("produce_number"), + _declare_positional("widen_to_long", ("value", INT64)), + _declare_positional("widen_to_double", ("value", DOUBLE)), + _declare_named("produce_nothing"), + _declare_positional("consume_nullable", ("value", _make_nullable(INT32))), + _declare_named("produce_fraction"), + _declare_positional("consume_float", ("value", FLOAT)), + _declare_positional( + "consume_double_list", + ("values", _make_nullable({"type": "array", "items": _make_nullable(DOUBLE)})), + ), + ], + "java_interface_example": [ + _declare_named("extract"), + _declare_named("transform", ("extracted", INT64, False)), + _declare_named( + "summarize", ("region_code", _make_nullable(STRING), True), ("transformed", INT64, False) + ), + ], + "java_annotation_example": [ + _declare_named("extract"), + _declare_positional("transform", ("extracted", INT64)), + _declare_positional("load", ("transformed", INT64)), + _declare_named( + "report", ("runLabel", _make_nullable(STRING), False), ("transformed", INT64, False) + ), + _declare_named("concurrent"), + ], + }, + ) + assert list(result.task_handlers) == [ + "java_interface_example", + "java_annotation_example", + "java_xcom_casting_example", + ] + with zipfile.ZipFile(jar) as zf: + manifest = _parse_manifest(zf.read("META-INF/MANIFEST.MF")) + assert re.fullmatch(r"[0-9a-f]{64}", manifest["airflow-cache-digest"]) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py new file mode 100644 index 0000000000000..05cd0a2e4a50a --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py @@ -0,0 +1,138 @@ +# +# 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. +"""Probe a bundle packed by the TypeScript SDK through the real Node runtime.""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +from typing import TYPE_CHECKING +from unittest import mock + +import pytest +import structlog + +from airflow.dag_processing.processor import TaskHandlerDeclaration, TaskHandlerParsingResult +from airflow.dag_processing.task_handler_processor import LangSDKTaskHandlerProcessorProcess +from airflow.sdk.execution_time import supervisor +from airflow.sdk.execution_time.coordinator import reset_coordinator_manager + +from tests_common.pytest_plugin import AIRFLOW_ROOT_PATH +from tests_common.test_utils.config import conf_vars + +if TYPE_CHECKING: + from collections.abc import Iterator + from pathlib import Path + +TS_SDK_PATH = AIRFLOW_ROOT_PATH / "ts-sdk" +# The SDK's `engines` requirement. +MIN_NODE_MAJOR = 22 +COORDINATORS = { + "ts": { + "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", + "kwargs": {"node_executable": shutil.which("node")}, + } +} + + +def _get_toolchain_problem() -> str | None: + """Return why this machine cannot pack the example bundle, or ``None`` when it can.""" + if not TS_SDK_PATH.is_dir(): + return "the TypeScript SDK sources are absent" + node = shutil.which("node") + if node is None or shutil.which("pnpm") is None: + return "needs node and pnpm" + try: + version = subprocess.run( + [node, "--version"], capture_output=True, text=True, check=True + ).stdout.strip() + major = int(version.lstrip("v").split(".")[0]) + except (OSError, subprocess.CalledProcessError, ValueError) as e: + return f"cannot read the Node.js version: {e}" + if major < MIN_NODE_MAJOR: + return f"needs Node.js {MIN_NODE_MAJOR} or later, found {version}" + return None + + +pytestmark = pytest.mark.skipif( + os.environ.get("AIRFLOW_LANG_SDK_REAL_PROBE_TESTS") != "1", + reason="set AIRFLOW_LANG_SDK_REAL_PROBE_TESTS=1 to probe real Lang-SDK artifacts", +) + + +@pytest.fixture(scope="module") +def example_bundle(tmp_path_factory) -> Path: + """Pack ``ts-sdk/example`` from a copy of the SDK sources, so the checkout gets no build output.""" + if (problem := _get_toolchain_problem()) is not None: + pytest.skip(problem) + sdk = tmp_path_factory.mktemp("ts-sdk") / "ts-sdk" + shutil.copytree(TS_SDK_PATH, sdk, ignore=shutil.ignore_patterns("node_modules", "dist", ".pnpm-store")) + env = {**os.environ, "CI": "true", "COREPACK_ENABLE_DOWNLOAD_PROMPT": "0"} + # The example links the SDK, so it is installed again once the SDK is built. + for cwd, args in [ + (sdk, ["install", "--frozen-lockfile"]), + (sdk, ["run", "build"]), + (sdk / "example", ["install"]), + (sdk / "example", ["run", "build"]), + ]: + subprocess.run(["pnpm", *args], cwd=cwd, env=env, check=True) + return sdk / "example" / "dist" / "bundle.min.mjs" + + +@pytest.fixture +def fresh_coordinator_manager() -> Iterator[None]: + reset_coordinator_manager() + yield + reset_coordinator_manager() + + +def _declare(task_id: str) -> TaskHandlerDeclaration: + return TaskHandlerDeclaration(task_id=task_id, binding="named", params=None) + + +@pytest.mark.usefixtures("fresh_coordinator_manager") +@conf_vars({("sdk", "coordinators"): json.dumps(COORDINATORS)}) +@mock.patch.object(supervisor, "_should_use_exec", autospec=True, return_value=False) +def test_a_packed_bundle_declares_its_task_handlers(mock_should_use_exec, example_bundle): + result = LangSDKTaskHandlerProcessorProcess.run( + coordinator="ts", + path=example_bundle, + bundle_path=example_bundle.parent, + bundle_name="ts-task-handlers", + artifact_rel_path=example_bundle.name, + logger=structlog.get_logger(), + ) + + assert result == TaskHandlerParsingResult( + fileloc=os.fspath(example_bundle), + task_handlers={ + "typescript_example": [ + _declare("build_message"), + _declare("read_connection"), + _declare("write_and_delete_variable"), + ], + "typescript_taskflow_example": [ + _declare("summarize"), + _declare("report"), + _declare("build_message"), + ], + }, + ) + assert list(result.task_handlers) == ["typescript_example", "typescript_taskflow_example"] diff --git a/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py new file mode 100644 index 0000000000000..510e2117dd299 --- /dev/null +++ b/airflow-core/tests/unit/models/test_lang_sdk_task_handler.py @@ -0,0 +1,233 @@ +# +# 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. +from __future__ import annotations + +from typing import Any + +import pytest +import uuid6 +from sqlalchemy import delete, insert, select +from sqlalchemy.exc import IntegrityError + +from airflow.models.dag import DagModel +from airflow.models.lang_sdk_task_handler import ( + LangSDKTaskHandler, + LangSDKTaskHandlerArtifact, + compute_fileloc_hash, +) + +pytestmark = pytest.mark.db_test + +ARTIFACT_BUNDLE = "java-task-handlers" +# 6000 bytes in UTF-8, over both the MySQL key limit and the Postgres btree entry limit. +LONG_NON_ASCII_PATH = "任" * 2000 +TASK_HANDLERS = { + "etl": [ + { + "task_id": "extract", + "binding": "positional", + "params": [ + {"name": "path", "value_schema": {"type": "string"}, "required": True, "exact_name": True}, + {"name": None, "value_schema": None, "required": False, "exact_name": False}, + ], + } + ], +} + + +def _make_artifact( + *, + bundle_name: str = ARTIFACT_BUNDLE, + relative_fileloc: str = "etl.jar", + cache_digest: str | None = "0" * 64, + task_handlers: dict[str, list[dict[str, Any]]] | None = None, +) -> LangSDKTaskHandlerArtifact: + return LangSDKTaskHandlerArtifact( + bundle_name=bundle_name, + relative_fileloc=relative_fileloc, + size_bytes=1024, + cache_digest=cache_digest, + task_handlers=TASK_HANDLERS if task_handlers is None else task_handlers, + ) + + +def _make_handler(*, dag_id: str, artifact: LangSDKTaskHandlerArtifact) -> LangSDKTaskHandler: + return LangSDKTaskHandler( + dag_id=dag_id, + task_id="extract", + artifact_id=artifact.id, + dag_bundle_name="testing", + dag_relative_fileloc=f"{dag_id}.py", + ) + + +def _add_dags_and_artifact(*dag_ids: str, session) -> LangSDKTaskHandlerArtifact: + session.add_all(DagModel(dag_id=dag_id, bundle_name="testing") for dag_id in dag_ids) + artifact = _make_artifact() + session.add(artifact) + session.flush() + return artifact + + +def test_deleting_dag_deletes_its_handlers(testing_dag_bundle, session): + artifact = _add_dags_and_artifact("dag_a", "dag_b", session=session) + session.add_all(_make_handler(dag_id=dag_id, artifact=artifact) for dag_id in ("dag_a", "dag_b")) + session.flush() + + session.execute(delete(DagModel).where(DagModel.dag_id == "dag_a")) + + assert session.scalars(select(LangSDKTaskHandler.dag_id)).all() == ["dag_b"] + assert session.scalars(select(LangSDKTaskHandlerArtifact.id)).all() == [artifact.id] + + +@pytest.mark.parametrize( + "task_handlers", + [ + pytest.param(TASK_HANDLERS, id="handlers"), + pytest.param({}, id="none"), + ], +) +def test_artifact_stores_its_task_handlers(session, task_handlers): + session.add(_make_artifact(task_handlers=task_handlers)) + session.flush() + session.expire_all() + + assert session.scalar(select(LangSDKTaskHandlerArtifact.task_handlers)) == task_handlers + + +def test_artifact_can_store_no_cache_digest(session): + session.add(_make_artifact(cache_digest=None)) + session.flush() + session.expire_all() + + assert session.scalar(select(LangSDKTaskHandlerArtifact.cache_digest)) is None + + +def test_deleting_referenced_artifact_fails(testing_dag_bundle, session): + artifact = _add_dags_and_artifact("dag_a", session=session) + session.add(_make_handler(dag_id="dag_a", artifact=artifact)) + session.flush() + + with pytest.raises(IntegrityError): + session.execute( + delete(LangSDKTaskHandlerArtifact).where(LangSDKTaskHandlerArtifact.id == artifact.id) + ) + + +@pytest.mark.parametrize( + "relative_fileloc", + [ + pytest.param("etl.jar", id="short"), + pytest.param(LONG_NON_ASCII_PATH, id="long-non-ascii"), + ], +) +def test_artifact_path_is_unique_per_bundle(session, relative_fileloc): + session.add_all( + [ + _make_artifact(relative_fileloc=relative_fileloc), + _make_artifact(bundle_name="go-task-handlers", relative_fileloc=relative_fileloc), + ] + ) + session.flush() + + session.add(_make_artifact(relative_fileloc=relative_fileloc)) + with pytest.raises(IntegrityError): + session.flush() + + +def test_core_insert_fills_fileloc_hashes(testing_dag_bundle, session): + session.add(DagModel(dag_id="dag_a", bundle_name="testing")) + session.flush() + artifact_id = uuid6.uuid7() + + session.execute( + insert(LangSDKTaskHandlerArtifact).values( + id=artifact_id, + bundle_name=ARTIFACT_BUNDLE, + relative_fileloc="etl.jar", + size_bytes=1024, + cache_digest="0" * 64, + task_handlers={}, + ) + ) + session.execute( + insert(LangSDKTaskHandler).values( + dag_id="dag_a", + task_id="extract", + artifact_id=artifact_id, + dag_bundle_name="testing", + dag_relative_fileloc="dags/etl.py", + ) + ) + + assert session.scalar(select(LangSDKTaskHandlerArtifact.relative_fileloc_hash)) == compute_fileloc_hash( + "etl.jar" + ) + assert session.scalar(select(LangSDKTaskHandler.dag_relative_fileloc_hash)) == compute_fileloc_hash( + "dags/etl.py" + ) + + +def _add_artifact(*, session) -> LangSDKTaskHandlerArtifact: + artifact = _make_artifact() + session.add(artifact) + session.flush() + return artifact + + +def _add_handler(*, session) -> LangSDKTaskHandler: + handler = _make_handler(dag_id="dag_a", artifact=_add_dags_and_artifact("dag_a", session=session)) + session.add(handler) + session.flush() + return handler + + +@pytest.mark.parametrize( + ("add_row", "path_attr", "hash_column", "other_attr", "other_value"), + [ + pytest.param( + _add_artifact, + "relative_fileloc", + LangSDKTaskHandlerArtifact.relative_fileloc_hash, + "size_bytes", + 2048, + id="artifact", + ), + pytest.param( + _add_handler, + "dag_relative_fileloc", + LangSDKTaskHandler.dag_relative_fileloc_hash, + "dag_bundle_name", + "other-bundle", + id="handler", + ), + ], +) +def test_fileloc_hash_follows_orm_updates( + testing_dag_bundle, session, add_row, path_attr, hash_column, other_attr, other_value +): + row = add_row(session=session) + original_path = getattr(row, path_attr) + + setattr(row, other_attr, other_value) + session.flush() + assert session.scalar(select(hash_column)) == compute_fileloc_hash(original_path) + + setattr(row, path_attr, "moved/file") + session.flush() + assert session.scalar(select(hash_column)) == compute_fileloc_hash("moved/file") diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py b/airflow-core/tests/unit/utils/test_db_cleanup.py index 375461b6848ac..41ec3ff7de9a7 100644 --- a/airflow-core/tests/unit/utils/test_db_cleanup.py +++ b/airflow-core/tests/unit/utils/test_db_cleanup.py @@ -1087,6 +1087,7 @@ def test_no_models_missing(self): # leave alone - per-asset key/value state, upserted in place (PK is asset_id+key), # so it is bounded and current, not accumulating history; removed with its asset "asset_state_store", + "lang_sdk_task_handler_artifact", # parse-time cache of task handler artifacts, not run data # Purged indirectly: each of these hangs off a cleaned table by an # ON DELETE CASCADE foreign key, so the rows go when the parent does. # cascade from dag_run once the partition run has fired; while it is still @@ -1098,6 +1099,7 @@ def test_no_models_missing(self): "hitl_detail", # cascade from task_instance "hitl_detail_history", # cascade from task_instance_history "job_team", # cascade from job + "lang_sdk_task_handler", # cascade from dag "task_inlet_asset_reference", # cascade from dag } diff --git a/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java b/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java index ad33a15708a46..1c15b425932df 100644 --- a/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java +++ b/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java @@ -19,7 +19,6 @@ package org.apache.airflow.e2e; -import java.util.List; import org.apache.airflow.sdk.*; import org.jetbrains.annotations.NotNull; @@ -27,7 +26,7 @@ * Bundle for the runner-behaviour E2E tests: deliberately broken task classes that exercise * instantiation failures, and a task that round-trips Airflow Variables through the supervisor. */ -public class TestBundleBuilder implements BundleBuilder { +public class TestBundleBuilder { public static class MissingNoArgConstructor implements Task { public MissingNoArgConstructor(String unused) {} @@ -60,19 +59,15 @@ public void execute(@NotNull Context context, Client client) { } } - @NotNull - @Override - public Iterable getDags() { - var uninstantiable = new DagDef("java_uninstantiable"); - uninstantiable.addTask("missing_no_arg_constructor", MissingNoArgConstructor.class); - uninstantiable.addTask("non_static_inner", NonStaticInner.class); - var variableWrite = new DagDef("java_variable_write"); - variableWrite.addTask("write_and_delete", WriteAndDeleteVariable.class); - return List.of(uninstantiable, variableWrite); + public static Bundle build() { + return new Bundle() + .register( + "java_uninstantiable", "missing_no_arg_constructor", MissingNoArgConstructor.class) + .register("java_uninstantiable", "non_static_inner", NonStaticInner.class) + .register("java_variable_write", "write_and_delete", WriteAndDeleteVariable.class); } public static void main(String[] args) { - var bundle = new TestBundleBuilder().build(); - Server.create(args).serve(bundle); + Server.create(args).serve(build()); } } diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py index 39ff3fd0a64b8..1282326cc9a1c 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py @@ -293,6 +293,22 @@ def _setup_xcom_object_storage_integration(dot_env_file, tmp_dir): ] +def _build_dag_bundle_config(artifact_bundles: dict[str, str]) -> str: + """Return a ``dag_bundle_config_list`` of the Dags folder plus a ``LocalDagBundle`` per artifact path. + + Registration is what makes a coordinator's ``task_handler_bundle_name`` resolvable on the + worker. The artifact directories are mounted only there; the Dag processor finds no files + in them. + """ + local_bundle = "airflow.dag_processing.bundles.local.LocalDagBundle" + bundles = [{"name": "dags-folder", "classpath": local_bundle, "kwargs": {}}] + bundles.extend( + {"name": name, "classpath": local_bundle, "kwargs": {"path": path}} + for name, path in artifact_bundles.items() + ) + return json.dumps(bundles) + + def _run_java_sdk_gradle(workdir, *gradle_argv, capture_output=False, native=False): """Run the Java SDK Gradle wrapper natively or inside the pinned JDK container. @@ -416,7 +432,7 @@ def _setup_java_sdk_integration(dot_env_file, tmp_dir): copyfile(JAVA_DOCKERFILE_PATH, tmp_dir / "Dockerfile.java") # Copy each bundle's JARs into its own directory; the compose bind-mounts - # expose them to the worker, and each JavaCoordinator globs its own dir. + # expose them to the worker, where each is registered as its own Dag bundle. copytree(JAVA_SDK_EXAMPLE_LIBS_PATH, tmp_dir / "java-jars") copytree(SCALA_SPARK_EXAMPLE_LIBS_PATH, tmp_dir / "scala-jars") copytree(JAVA_TEST_BUNDLE_LIBS_PATH, tmp_dir / "java-test-jars") @@ -456,20 +472,27 @@ def _setup_java_sdk_integration(dot_env_file, tmp_dir): ) # One JavaCoordinator per queue on the same worker image, each serving its - # own bundle. The scala-jdk entry pins main_class (Spark's large classpath - # makes Main-Class discovery ambiguous) and carries Spark's Java 17 module - # openings, a small driver heap, and a longer startup timeout for its large - # dependency classpath. + # own artifact bundle (one bundle is one classpath). The scala-jdk entry pins + # main_class (Spark's large classpath makes Main-Class discovery ambiguous) + # and carries Spark's Java 17 module openings, a small driver heap, and a + # longer startup timeout for its large dependency classpath. + dag_bundle_config = _build_dag_bundle_config( + { + "java-task-handlers": "/opt/airflow/java-jars", + "scala-task-handlers": "/opt/airflow/scala-jars", + "java-test-task-handlers": "/opt/airflow/java-test-jars", + } + ) coordinator_config = json.dumps( { "java-jdk": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", - "kwargs": {"jars_root": ["/opt/airflow/java-jars"]}, + "kwargs": {"task_handler_bundle_name": "java-task-handlers"}, }, "scala-jdk": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": { - "jars_root": ["/opt/airflow/scala-jars"], + "task_handler_bundle_name": "scala-task-handlers", "main_class": "org.apache.airflow.example.ScalaSparkBundleBuilder", "jvm_args": ["-Xmx512m", *_SPARK_JAVA_MODULE_OPTIONS], "task_startup_timeout": 60.0, @@ -477,7 +500,7 @@ def _setup_java_sdk_integration(dot_env_file, tmp_dir): }, "java-test-jdk": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", - "kwargs": {"jars_root": ["/opt/airflow/java-test-jars"]}, + "kwargs": {"task_handler_bundle_name": "java-test-task-handlers"}, }, } ) @@ -502,6 +525,7 @@ def _setup_java_sdk_integration(dot_env_file, tmp_dir): dot_env_file.write_text( f"AIRFLOW_UID={os.getuid()}\n" # Single-quote the JSON values so Docker Compose reads them literally. + f"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='{dag_bundle_config}'\n" f"AIRFLOW__SDK__COORDINATORS='{coordinator_config}'\n" f"AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{queue_to_coordinator}'\n" f"AIRFLOW_CONN_TEST_HTTP='{test_http_conn}'\n" @@ -598,9 +622,9 @@ def _setup_go_sdk_integration(dot_env_file, tmp_dir): """Set up the go_sdk E2E test mode. Compiles the Go SDK example bundle into a self-contained executable bundle - via the ``airflow-go-pack`` tooling, drops it into the directory the - ``ExecutableCoordinator`` scans, copies the Python stub Dag, and writes the - coordinator configuration. + via the ``airflow-go-pack`` tooling, drops it into the directory registered + as the ``go-task-handlers`` Dag bundle, copies the Python stub Dag, and + writes the coordinator configuration. The packed bundle is a statically linked native executable (built with ``CGO_ENABLED=0``), so the stock Airflow worker image can exec it directly @@ -624,13 +648,14 @@ def _setup_go_sdk_integration(dot_env_file, tmp_dir): copyfile(GO_SDK_DAGS_PATH / "go_examples.py", tmp_dir / "dags" / "go_examples.py") # Coordinator registry: maps the logical name "go-sdk" to ExecutableCoordinator, - # which scans executables_root for the packed bundle by dag_id. + # which scans the go-task-handlers Dag bundle for the packed bundle by dag_id. # Queue mapping: routes tasks on the "golang" queue to "go-sdk". + dag_bundle_config = _build_dag_bundle_config({"go-task-handlers": "/opt/airflow/go-bundles"}) coordinator_config = json.dumps( { "go-sdk": { "classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", - "kwargs": {"executables_root": ["/opt/airflow/go-bundles"]}, + "kwargs": {"task_handler_bundle_name": "go-task-handlers"}, } } ) @@ -639,6 +664,7 @@ def _setup_go_sdk_integration(dot_env_file, tmp_dir): dot_env_file.write_text( f"AIRFLOW_UID={os.getuid()}\n" # Single-quote the JSON values so Docker Compose reads them literally. + f"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='{dag_bundle_config}'\n" f"AIRFLOW__SDK__COORDINATORS='{coordinator_config}'\n" f"AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{queue_to_coordinator}'\n" # Connection and variable read by the Go example bundle tasks. @@ -730,12 +756,13 @@ def _setup_ts_sdk_integration(dot_env_file, tmp_dir): for dag_file in ("typescript_example.py", "typescript_taskflow_example.py"): copyfile(TS_SDK_EXAMPLE_PATH / "dags" / dag_file, tmp_dir / "dags" / dag_file) + dag_bundle_config = _build_dag_bundle_config({"ts-task-handlers": "/opt/airflow/ts-bundles"}) coordinator_config = json.dumps( { "ts": { "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", "kwargs": { - "bundles_root": ["/opt/airflow/ts-bundles"], + "task_handler_bundle_name": "ts-task-handlers", "node_executable": "/opt/nodejs/node", }, } @@ -747,6 +774,7 @@ def _setup_ts_sdk_integration(dot_env_file, tmp_dir): f"AIRFLOW_UID={os.getuid()}\n" f"NODE_IMAGE={NODE_IMAGE}\n" # single-quoted so Docker Compose reads the JSON literally + f"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='{dag_bundle_config}'\n" f"AIRFLOW__SDK__COORDINATORS='{coordinator_config}'\n" f"AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{queue_to_coordinator}'\n" "AIRFLOW_CONN_TYPESCRIPT_EXAMPLE_HTTP=http://user:pass@example.com/\n" diff --git a/contributing-docs/30_new_language_sdk.rst b/contributing-docs/30_new_language_sdk.rst index 34c4cdbb4a27a..46ded9b149369 100644 --- a/contributing-docs/30_new_language_sdk.rst +++ b/contributing-docs/30_new_language_sdk.rst @@ -121,10 +121,32 @@ The method returns a ``(command, subprocess_schema_version)`` pair: across SDK versions. See `Supervisor Schema`_ below. Call ``self._get_scan_roots()`` to retrieve the artifact directories the base -class has already resolved from the coordinator's configured source — an -explicit filesystem root, a named Dag bundle (``dag_bundle_name``), or the -task's own bundle. Subclasses should scan those roots rather than reading the -configured root directly. +class has already resolved: the Dag bundle named by ``task_handler_bundle_name``, +or the task's own Dag bundle when it is unset, pinned for the whole task. +Subclasses should scan those directories rather than locating artifacts +themselves. + +SubprocessCoordinator: implementing ``_build_parse_task_handler_command`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +When a Python Dag file has stub tasks, the Dag processor asks the artifact that +implements them which task handlers it registers, and checks each stub task +against its handler. A coordinator opts in by building the command that starts +its runtime for one artifact: + +.. code-block:: python + + def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: ... + +*path* is the artifact the Dag processor picked, so the method does not search +for one. The returned pair follows the rules of ``_build_execute_task_command``: +no ``--comm`` or ``--logs`` flags, and the schema version the runtime +understands. The runtime then answers as described in +`Answering TaskHandlerParseRequest`_. + +The default raises ``NotImplementedError``, so a coordinator that does not +implement it cannot be probed. ``ExecutableCoordinator`` implements it for +executable bundles. Supervisor Schema ~~~~~~~~~~~~~~~~~ @@ -195,8 +217,10 @@ as soon as possible. The supervisor verifies that the connecting peer belongs to the launched process tree, so the SDK MUST connect from the same process or one of its descendants. -Once both connections are accepted, the supervisor sends a ``StartupDetails`` -message on the comm socket to initiate execution. +Once both connections are accepted, the supervisor sends the first message on +the comm socket. ``StartupDetails`` starts a task. ``TaskHandlerParseRequest`` +comes from the Dag processor instead, and asks which task handlers the runtime +registers (see `Answering TaskHandlerParseRequest`_). Wire protocol ~~~~~~~~~~~~~ @@ -273,6 +297,40 @@ calls themselves) can already produce log records before the ``--logs`` socket in stage 3 exists to carry them. See `Logging`_ below for how to handle that gap. +Answering ``TaskHandlerParseRequest`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The SDK sends back one ``TaskHandlerParsingResult``, waits for the supervisor's +response to it, and exits. No user code runs; the answer comes from the SDK's +own task handler registrations. It MUST declare every registered handler and +depend only on the artifact, never on the request. + +* ``task_handlers`` maps each Dag id the artifact registers handlers for to + their declarations, in registration order. +* Each ``TaskHandlerDeclaration`` states in ``binding`` how stub-task arguments + bind to its ``params``: + + * ``positional``: by position. An argument count or value type the handler + cannot take makes the Dag fail to import. + * ``named``: by name in any order, ignoring case and underscores unless a + param sets ``exact_name``. An argument or param that matches nothing is + logged as a warning and the task still runs, so the runtime must accept + both. When no param matches and exactly one argument was passed, it may be + the whole value and is not warned about, unless ``params`` is empty, a + param sets ``exact_name``, or the argument cannot be an object. A value + type a param does not accept makes the Dag fail to import. + +* ``params`` lists the handler's parameters in order, ``[]`` when it has none, + or is ``null`` when the SDK cannot list them. Then only the handler's + presence is checked. +* Each ``TaskHandlerParam`` has a ``name`` (``null`` when the SDK has no name + for a positional parameter) and a ``value_schema``: the JSON Schema of the + values it accepts, in the vocabulary ``@task.stub`` uses for Python + annotations, or ``null`` when the SDK cannot state one. +* An empty ``task_handlers`` mapping is sent empty, never as ``null``. + +``go-sdk/pkg/execution`` is a reference implementation. + Logging ~~~~~~~ diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt index b906ba258554b..dd4fa33cb2381 100644 --- a/generated/known_sdk_imports_in_core.txt +++ b/generated/known_sdk_imports_in_core.txt @@ -6,6 +6,7 @@ airflow-core/src/airflow/configuration.py::1 airflow-core/src/airflow/dag_processing/dagbag.py::3 airflow-core/src/airflow/dag_processing/manager.py::4 airflow-core/src/airflow/dag_processing/processor.py::14 +airflow-core/src/airflow/dag_processing/task_handler_processor.py::8 airflow-core/src/airflow/exceptions.py::1 airflow-core/src/airflow/executors/base_executor.py::3 airflow-core/src/airflow/jobs/triggerer_job_runner.py::18 diff --git a/go-sdk/README.md b/go-sdk/README.md index a6dedde08049b..adb0feed4c11a 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -250,15 +250,15 @@ the full range of task states, and alternate XCom backends without implementing go tool airflow-go-pack ./example/bundle -- -trimpath -tags=prod ``` - Use `--output ` to write the packed bundle straight into a directory the coordinator scans - (`executables_root`), and pass extra `go build` flags after `--`. + Use `--output ` to write the packed bundle straight into the directory of the Dag bundle the + coordinator scans (`task_handler_bundle_name`), and pass extra `go build` flags after `--`. For cross-compiling (e.g. deploy to a Linux host from an Apple-silicon (darwin/arm64) machine), pass `--goos`/`--goarch` and the packer cross-builds for you: ```bash go tool airflow-go-pack --goos linux --goarch amd64 \ - --output ~/airflow/executable-bundles/sample-dag-bundle \ + --output /opt/airflow/go-task-handlers/sample-dag-bundle \ ./example/bundle ``` @@ -280,24 +280,37 @@ the full range of task states, and alternate XCom backends without implementing > `tool github.com/apache/airflow/go-sdk/cmd/airflow-go-pack` to your bundle module's `go.mod` and run > it with `go tool airflow-go-pack`. This pins the packer version per project. -- Register the coordinator and route the queue to it, under `[sdk]` in `airflow.cfg` (or the equivalent - `AIRFLOW__SDK__*` env vars): +- Register a Dag bundle for the packed bundles, register the coordinator, and route the queue to it, in + `airflow.cfg` (or the equivalent `AIRFLOW__*` env vars): ```ini + [dag_processor] + dag_bundle_config_list = [ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + {"name": "go-task-handlers", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {"path": "/opt/airflow/go-task-handlers"}} + ] + [sdk] - coordinators = {"go": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": {"executables_root": ["~/airflow/executable-bundles"]}}} + coordinators = {"go": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": {"task_handler_bundle_name": "go-task-handlers"}}} queue_to_coordinator = {"golang": "go"} ``` - `executables_root` is one or more directories the coordinator scans for bundles; `queue_to_coordinator` - routes stub tasks with `queue="golang"` to this Go coordinator. + `task_handler_bundle_name` names the Dag bundle the coordinator scans for packed bundles (the task's own + Dag bundle when unset). It is used only by mixed-language Dags, to locate the task handlers for the + `@task.stub` tasks of a Python Dag; Dags defined natively in a language SDK do not use it. + `queue_to_coordinator` routes stub tasks with `queue="golang"` to this Go coordinator. Only files with the + executable bit are considered, so use a Dag bundle that keeps it: a `LocalDagBundle` does, object-store + Dag bundles such as `S3DagBundle` do not. > [!IMPORTANT] - > The coordinator is part of the Airflow worker, so the `[sdk]` config (and the bundle files in - > `executables_root`) only need to be present wherever tasks actually execute. With `CeleryExecutor`, - > setting it on the Celery workers is sufficient. With `LocalExecutor`, tasks run inside the scheduler - > process, so it must be set where the scheduler can read it. The API server and Dag processor do not - > need it. + > The `[sdk]` config and the packed bundle files must be present wherever tasks execute and on the Dag + > processor. With `CeleryExecutor`, tasks execute on the Celery workers; with `LocalExecutor`, they run + > inside the scheduler process. The Dag processor checks the stub tasks of each Python Dag against the + > task handlers the packed bundles register, so it runs them too and needs bundles built for its + > operating system and CPU architecture. The API server does not need any of it. Register the Dag bundle + > in `[dag_processor] dag_bundle_config_list` on every component, like your other Dag bundles: the worker + > and the Dag processor resolve `task_handler_bundle_name` through it, and wherever the `[sdk]` config is + > read it is rejected if the name is missing there. - Deploy the matching Python stub Dag (above) into Airflow. There is no separate Go worker to run: the Airflow worker forks the bundle binary once per task instance. @@ -376,6 +389,9 @@ Python supervisor / task runner protocol on the comm socket, with structured JSON-line logs on the logs socket. - The Python runtime is the worker. It proxies every `GetConnection` / `GetVariable` / `GetXCom` / `SetXCom` call through to the Execution API. The Go binary just runs the task function. +- The Dag processor launches the binary the same way to check a Python Dag's stub tasks. It sends a + `TaskHandlerParseRequest`, and the binary replies with every task handler it registers, keyed by Dag + id in registration order, and how each handler's parameters bind. Then it exits. The Go side of the protocol is implemented in `pkg/execution/`. On the Python side it is the `ExecutableCoordinator` in `task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py`. diff --git a/go-sdk/airflow/bundle.go b/go-sdk/airflow/bundle.go index 8f82b07121138..58f04354ea358 100644 --- a/go-sdk/airflow/bundle.go +++ b/go-sdk/airflow/bundle.go @@ -125,18 +125,15 @@ func (b *BundleRef) Register(items ...Registerable) { } // taskHandlerMap holds the registered task handlers by dag_id and task_id. -// It also keeps registration order. The --airflow-metadata manifest lists the tasks of each Dag -// in that order. +// It also keeps registration order. The --airflow-metadata manifest, and the reply to the Dag +// processor's task handler parse, list the tasks of each Dag in that order. type taskHandlerMap struct { mu sync.RWMutex handlers map[string]map[string]bundle.Task order []bundle.TaskHandlerInfo } -var ( - _ bundle.Bundle = (*taskHandlerMap)(nil) - _ bundle.EnumerableBundle = (*taskHandlerMap)(nil) -) +var _ bundle.Registry = (*taskHandlerMap)(nil) func (m *taskHandlerMap) add(dagID, taskID string, task bundle.Task) { m.mu.Lock() diff --git a/go-sdk/airflow/serve.go b/go-sdk/airflow/serve.go index bcfb33ab934d1..910cd5a1398f2 100644 --- a/go-sdk/airflow/serve.go +++ b/go-sdk/airflow/serve.go @@ -57,8 +57,9 @@ const ( // The command-line flags of the executable decide what Serve does. // With --airflow-metadata it prints the bundle's manifest and returns, which is how // airflow-go-pack reads the Dag and task ids of the registered task handlers. -// With --comm and --logs, which the Airflow supervisor passes, it runs one task over the -// coordinator protocol. +// With --comm and --logs, which Airflow passes, it speaks the coordinator protocol: it either +// runs one task, or tells the Dag processor which task handlers it registers for the Dags the +// Dag processor asks about. // // main must exit with a non-zero status when Serve returns an error, because the exit status // is how the supervisor learns that the task failed: diff --git a/go-sdk/cmd/airflow-go-pack/inspect_test.go b/go-sdk/cmd/airflow-go-pack/inspect_test.go index 8208136e80176..5841f3d360e7a 100644 --- a/go-sdk/cmd/airflow-go-pack/inspect_test.go +++ b/go-sdk/cmd/airflow-go-pack/inspect_test.go @@ -42,7 +42,7 @@ func TestInspectCmd(t *testing.T) { " - \"t1\"\n", ) bundle := filepath.Join(dir, "bundle") - require.NoError(t, writeBundle(exe, bundle, source, manifest)) + require.NoError(t, writeBundle(exe, bundle, source, fixedManifest(manifest))) for _, tc := range []struct { name string diff --git a/go-sdk/cmd/airflow-go-pack/pack.go b/go-sdk/cmd/airflow-go-pack/pack.go index f4be6fa8c4578..23d542832d5bb 100644 --- a/go-sdk/cmd/airflow-go-pack/pack.go +++ b/go-sdk/cmd/airflow-go-pack/pack.go @@ -19,6 +19,8 @@ package main import ( "bytes" + "crypto/sha256" + "encoding/hex" "errors" "fmt" "go/ast" @@ -173,20 +175,34 @@ func runPack(stdout, stderr io.Writer, opts *packOptions) error { } warnOnSuspiciousIDs(stderr, meta) - manifest, err := renderManifest(meta, filepath.Base(sourcePath)) - if err != nil { - return fmt.Errorf("rendering manifest: %w", err) - } sourceBytes, err := os.ReadFile(sourcePath) if err != nil { return fmt.Errorf("reading source file: %w", err) } + // The digests are taken over the staged copy, the file the trailer's + // binary_sha256 is computed from, so the two agree even if the executable + // changes while packing. + renderWithDigests := func(stagedExec string) ([]byte, error) { + baseManifest, err := renderManifest(meta, filepath.Base(sourcePath), nil) + if err != nil { + return nil, fmt.Errorf("rendering manifest: %w", err) + } + digests, err := computeDigests(stagedExec, sourceBytes, baseManifest) + if err != nil { + return nil, err + } + manifest, err := renderManifest(meta, filepath.Base(sourcePath), &digests) + if err != nil { + return nil, fmt.Errorf("rendering manifest: %w", err) + } + return manifest, nil + } // Assemble the bundle through a temp file and atomically move it into // place: we never mutate the build artefact or the user-supplied // --executable, and a failed pack never leaves a truncated or half-written // file at output. - if err := writeBundle(execPath, output, sourceBytes, manifest); err != nil { + if err := writeBundle(execPath, output, sourceBytes, renderWithDigests); err != nil { return err } @@ -462,12 +478,53 @@ func runIntrospect(execPath string, flag string) ([]byte, error) { return stdout.Bytes(), nil } +// cacheDigestDomain starts the cache digest's input, so the digest cannot equal +// a SHA-256 taken of the same bytes for another purpose, and a later definition +// can change the version. +const cacheDigestDomain = "airflow-go-pack cache digest v1\x00" + +// bundleDigests is the manifest's digests mapping, each a lower-case hex SHA-256. +type bundleDigests struct { + // Integrity covers the binary region; it equals the trailer's binary_sha256. + Integrity string + // Cache covers every region of the bundle except its own value. It is + // SHA-256(cacheDigestDomain || SHA-256(binary) || SHA-256(source) || + // SHA-256(manifest rendered without digests)), with each inner digest as + // its 32 raw bytes. Packing the same inputs again gives the same value, and + // a change to any of them gives a different one. + Cache string +} + +func computeDigests(execPath string, source, baseManifest []byte) (bundleDigests, error) { + binaryHash, err := bundlefooter.HashFile(execPath) + if err != nil { + return bundleDigests{}, fmt.Errorf("hashing executable %s: %w", execPath, err) + } + sourceHash := sha256.Sum256(source) + manifestHash := sha256.Sum256(baseManifest) + + h := sha256.New() + h.Write([]byte(cacheDigestDomain)) + h.Write(binaryHash[:]) + h.Write(sourceHash[:]) + h.Write(manifestHash[:]) + return bundleDigests{ + Integrity: hex.EncodeToString(binaryHash[:]), + Cache: hex.EncodeToString(h.Sum(nil)), + }, nil +} + // renderManifest serialises the airflow-metadata manifest as deterministic, // sorted-key YAML matching airflow-metadata.schema.json. It injects the schema's -// source field (the filename the manifest is built from), which the producer's -// Manifest omits because only the packer knows it; every other field is copied -// from the introspected manifest verbatim. -func renderManifest(meta airflowmetadata.Manifest, sourceName string) ([]byte, error) { +// source field (the filename the manifest is built from) and, when digests is +// not nil, the digests mapping, both of which the producer's Manifest omits +// because only the packer knows them; every other field is copied from the +// introspected manifest verbatim. +func renderManifest( + meta airflowmetadata.Manifest, + sourceName string, + digests *bundleDigests, +) ([]byte, error) { version := meta.AirflowBundleMetadataVersion if version == "" { version = airflowmetadata.FormatVersion @@ -514,9 +571,21 @@ func renderManifest(meta airflowmetadata.Manifest, sourceName string) ([]byte, e }, }, scalar("source"), quotedScalar(sourceName), - scalar("dags"), dagsNode, }, } + if digests != nil { + manifest.Content = append(manifest.Content, + scalar("digests"), + &yaml.Node{ + Kind: yaml.MappingNode, + Content: []*yaml.Node{ + scalar("integrity"), quotedScalar(digests.Integrity), + scalar("cache"), quotedScalar(digests.Cache), + }, + }, + ) + } + manifest.Content = append(manifest.Content, scalar("dags"), dagsNode) root.Content = []*yaml.Node{manifest} var buf bytes.Buffer @@ -613,11 +682,16 @@ func sameFile(a, b string) (bool, error) { // writeBundle assembles the bundle at output by copying the executable to a // temporary file in output's directory, appending the source+manifest footer -// to that copy, then atomically renaming it into place. Writing through a +// to that copy, then atomically renaming it into place. renderMetadata builds +// the manifest from the copy before the footer is appended. Writing through a // temp file keeps a failed pack from leaving a truncated or half-written // artefact at output, and guarantees the file being copied is never the same // open file as the destination. -func writeBundle(execPath, output string, source, metadata []byte) error { +func writeBundle( + execPath, output string, + source []byte, + renderMetadata func(stagedExec string) ([]byte, error), +) error { outDir := filepath.Dir(output) // The temp file and the atomic rename both live in output's directory, so it // must exist. Create it for the user (e.g. --output ./bin/bundle with no @@ -641,6 +715,10 @@ func writeBundle(execPath, output string, source, metadata []byte) error { if err := copyFile(execPath, tmpPath, 0o755); err != nil { return fmt.Errorf("writing %s: %w", output, err) } + metadata, err := renderMetadata(tmpPath) + if err != nil { + return err + } if err := bundlefooter.Append(tmpPath, source, metadata); err != nil { return err } diff --git a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go index e1cc29fe68e67..1d75ab6c41936 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go @@ -19,6 +19,8 @@ package main import ( "bytes" + "crypto/sha256" + "encoding/hex" "encoding/json" "io" "os" @@ -138,6 +140,12 @@ func TestPack_CrossArchExecutableWithMetadataFile(t *testing.T) { sdkVersion := m[1] assert.Regexp(t, `^(\(devel\)|v[0-9].*)$`, sdkVersion, `sdk.version must be "(devel)" or a v-prefixed module version`) + // The cache digest covers the environment-dependent binary and sdk.version + // too, so it is folded in the same way. + cacheLine := regexp.MustCompile(`(?m)^ cache: "([0-9a-f]{64})"$`) + c := cacheLine.FindStringSubmatch(string(metadata)) + require.NotNil(t, c, "manifest must contain a digests.cache line:\n%s", metadata) + crossHash := sha256.Sum256(crossBytes) expectedManifest := `airflow_bundle_metadata_version: "1.0" sdk: @@ -145,6 +153,9 @@ sdk: version: "` + sdkVersion + `" supervisor_schema_version: "` + execution.SupervisorSchemaVersion + `" source: "main.go" +digests: + integrity: "` + hex.EncodeToString(crossHash[:]) + `" + cache: "` + c[1] + `" dags: concurrent_xcom_dag: tasks: diff --git a/go-sdk/cmd/airflow-go-pack/pack_test.go b/go-sdk/cmd/airflow-go-pack/pack_test.go index 999d5c0f01771..ce5f90894dc4c 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_test.go @@ -18,6 +18,9 @@ package main import ( + "bytes" + "crypto/sha256" + "encoding/hex" "errors" "io" "os" @@ -28,6 +31,7 @@ import ( "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" "github.com/apache/airflow/go-sdk/internal/airflowmetadata" "github.com/apache/airflow/go-sdk/internal/bundlefooter" @@ -47,9 +51,9 @@ func TestRenderManifest_DeterministicDagOrdering(t *testing.T) { }, } - got1, err := renderManifest(meta, "main.go") + got1, err := renderManifest(meta, "main.go", nil) require.NoError(t, err) - got2, err := renderManifest(meta, "main.go") + got2, err := renderManifest(meta, "main.go", nil) require.NoError(t, err) assert.Equal(t, got1, got2, "manifest should be byte-identical for identical input") @@ -87,7 +91,7 @@ func TestRenderManifest_QuotesValuesNotKeys(t *testing.T) { }, } - got, err := renderManifest(meta, "main.go") + got, err := renderManifest(meta, "main.go", nil) require.NoError(t, err) // Task values that look like scalars are quoted. @@ -98,6 +102,36 @@ func TestRenderManifest_QuotesValuesNotKeys(t *testing.T) { assert.NotContains(t, string(got), `"my_dag"`) } +func TestRenderManifest_Digests(t *testing.T) { + meta := airflowmetadata.Manifest{ + AirflowBundleMetadataVersion: "1.0", + SDK: airflowmetadata.SDK{ + Language: "go", + Version: "0.1.0", + SupervisorSchemaVersion: "2026-06-16", + }, + Dags: map[string]airflowmetadata.Dag{"my_dag": {Tasks: []string{"t1"}}}, + } + + got, err := renderManifest(meta, "main.go", &bundleDigests{Integrity: "aa11", Cache: "bb22"}) + require.NoError(t, err) + + assert.Equal(t, `airflow_bundle_metadata_version: "1.0" +sdk: + language: "go" + version: "0.1.0" + supervisor_schema_version: "2026-06-16" +source: "main.go" +digests: + integrity: "aa11" + cache: "bb22" +dags: + my_dag: + tasks: + - "t1" +`, string(got)) +} + func TestRenderManifest_EmptyDags(t *testing.T) { meta := airflowmetadata.Manifest{ AirflowBundleMetadataVersion: "1.0", @@ -108,7 +142,7 @@ func TestRenderManifest_EmptyDags(t *testing.T) { }, Dags: map[string]airflowmetadata.Dag{}, } - got, err := renderManifest(meta, "main.go") + got, err := renderManifest(meta, "main.go", nil) require.NoError(t, err) assert.Contains(t, string(got), "dags: {}") } @@ -417,6 +451,96 @@ func TestRunPack_AcceptsYAMLMetadataFile(t *testing.T) { ) } +// fixedManifest is a writeBundle renderMetadata that ignores the staged executable. +func fixedManifest(manifest []byte) func(string) ([]byte, error) { + return func(string) ([]byte, error) { return manifest, nil } +} + +// packDigests packs the given executable, source and JSON manifest in dir, and +// returns the digests the packed bundle's manifest records. +func packDigests(t *testing.T, dir string, exe, source, meta []byte) (integrity, cache string) { + t.Helper() + exePath := filepath.Join(dir, "exe") + require.NoError(t, os.WriteFile(exePath, exe, 0o755)) + sourcePath := filepath.Join(dir, "main.go") + require.NoError(t, os.WriteFile(sourcePath, source, 0o644)) + metaPath := filepath.Join(dir, "airflow-metadata.json") + require.NoError(t, os.WriteFile(metaPath, meta, 0o644)) + out := filepath.Join(dir, "bundle") + + require.NoError(t, runPack(io.Discard, io.Discard, &packOptions{ + executable: exePath, + source: sourcePath, + airflowMetadata: metaPath, + output: out, + })) + + _, gotMeta, err := bundlefooter.Read(out) + require.NoError(t, err) + var parsed struct { + Digests struct { + Integrity string `yaml:"integrity"` + Cache string `yaml:"cache"` + } `yaml:"digests"` + } + require.NoError(t, yaml.Unmarshal(gotMeta, &parsed)) + return parsed.Digests.Integrity, parsed.Digests.Cache +} + +func TestRunPack_RecordsDigests(t *testing.T) { + exe := []byte("binary-bytes") + source := []byte("package main\nfunc main() {}\n") + meta := []byte(`{"airflow_bundle_metadata_version":"1.0",` + + `"sdk":{"language":"go","version":"0.1.0","supervisor_schema_version":"2026-06-16"},` + + `"dags":{"my_dag":{"tasks":["t1"]}}}`) + + integrity, cache := packDigests(t, t.TempDir(), exe, source, meta) + + binaryHash := sha256.Sum256(exe) + assert.Equal( + t, + hex.EncodeToString(binaryHash[:]), + integrity, + "integrity is the trailer's binary_sha256", + ) + assert.Len(t, cache, sha256.Size*2) + assert.NotEqual(t, integrity, cache) + + t.Run("a repack of the same inputs", func(t *testing.T) { + gotIntegrity, gotCache := packDigests(t, t.TempDir(), exe, source, meta) + assert.Equal(t, integrity, gotIntegrity) + assert.Equal(t, cache, gotCache) + }) + + oneByte := func(b []byte, i int) []byte { + changed := bytes.Clone(b) + changed[i] ^= 1 + return changed + } + for name, tc := range map[string]struct { + exe, source, meta []byte + integrityChanges bool + }{ + "one source byte": {exe: exe, source: oneByte(source, len(source)-2), meta: meta}, + "one binary byte": {exe: oneByte(exe, 0), source: source, meta: meta, integrityChanges: true}, + "the manifest": { + exe: exe, + source: source, + meta: bytes.Replace(meta, []byte(`"0.1.0"`), []byte(`"0.1.1"`), 1), + }, + } { + t.Run(name, func(t *testing.T) { + gotIntegrity, gotCache := packDigests(t, t.TempDir(), tc.exe, tc.source, tc.meta) + assert.NotEqual(t, cache, gotCache) + if tc.integrityChanges { + assert.NotEqual(t, integrity, gotIntegrity) + } else { + assert.Equal(t, integrity, gotIntegrity) + } + }) + } +} + // With no explicit --output, packing "./bundle" from a package dir named // "bundle" derives a default output that collides with the pre-built binary. func TestRunPack_RejectsDefaultOutputAliasingExecutable(t *testing.T) { @@ -490,6 +614,26 @@ func TestTargetPlatform(t *testing.T) { }) } +func TestWriteBundle_RendersMetadataFromTheStagedCopy(t *testing.T) { + dir := t.TempDir() + exe := filepath.Join(dir, "input-bin") + require.NoError(t, os.WriteFile(exe, []byte("binary-bytes"), 0o755)) + + var staged string + var stagedBytes []byte + require.NoError(t, writeBundle(exe, filepath.Join(dir, "bundle"), []byte("source"), + func(path string) ([]byte, error) { + staged = path + var err error + stagedBytes, err = os.ReadFile(path) + return []byte("manifest"), err + }, + )) + + assert.NotEqual(t, exe, staged) + assert.Equal(t, []byte("binary-bytes"), stagedBytes) +} + // When the --output parent directory does not exist, the packer must create it // instead of failing with an opaque temp-file error. func TestWriteBundle_CreatesMissingOutputDir(t *testing.T) { @@ -498,7 +642,10 @@ func TestWriteBundle_CreatesMissingOutputDir(t *testing.T) { require.NoError(t, os.WriteFile(exe, []byte("binary-bytes"), 0o755)) output := filepath.Join(dir, "bin", "nested", "bundle") - require.NoError(t, writeBundle(exe, output, []byte("source"), []byte("manifest"))) + require.NoError( + t, + writeBundle(exe, output, []byte("source"), fixedManifest([]byte("manifest"))), + ) info, err := os.Stat(output) require.NoError(t, err, "bundle should be written into the created directory") diff --git a/go-sdk/example/bundle/Justfile b/go-sdk/example/bundle/Justfile index 5a12c4407aed1..07df82a2280a6 100644 --- a/go-sdk/example/bundle/Justfile +++ b/go-sdk/example/bundle/Justfile @@ -29,8 +29,9 @@ build: pack # DAG/task identity via --airflow-metadata, and appends the source plus # airflow-metadata.yaml plus AFBNDL01 trailer. The output is a single # self-contained executable bundle, named after the bundle's package -# directory and written to the current directory. Drop it into -# ExecutableCoordinator's executables_root to deploy. +# directory and written to the current directory. Drop it into the +# Dag bundle named by ExecutableCoordinator's task_handler_bundle_name +# to deploy. pack: @echo "Packing example DAG bundle..." go tool airflow-go-pack --output ../../bin/example_dags diff --git a/go-sdk/internal/airflowmetadata/airflowmetadata.go b/go-sdk/internal/airflowmetadata/airflowmetadata.go index d11412d722c44..2045daebb17f5 100644 --- a/go-sdk/internal/airflowmetadata/airflowmetadata.go +++ b/go-sdk/internal/airflowmetadata/airflowmetadata.go @@ -28,8 +28,8 @@ const FormatVersion = "1.0" // Manifest is the shape printed by a bundle binary's --airflow-metadata flag // (YAML by default, JSON under --format json). It mirrors -// airflow-metadata.schema.json minus the source field, which only the packer -// can resolve from the build inputs. The yaml tags let the packer's +// airflow-metadata.schema.json minus the source and digests fields, which only +// the packer can resolve from the build inputs. The yaml tags let the packer's // --airflow-metadata flag decode a captured manifest in either JSON or YAML // (the airflow-metadata.yaml in a bundle). type Manifest struct { diff --git a/go-sdk/internal/bundle/task.go b/go-sdk/internal/bundle/task.go index d2353369be8e3..245f2d79308ee 100644 --- a/go-sdk/internal/bundle/task.go +++ b/go-sdk/internal/bundle/task.go @@ -26,6 +26,7 @@ import ( "runtime" "github.com/apache/airflow/go-sdk/pkg/binding" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" "github.com/apache/airflow/go-sdk/pkg/sdkcontext" "github.com/apache/airflow/go-sdk/sdk" ) @@ -35,6 +36,9 @@ import ( // airflow.DagRef.Task wrap a plain Go function into a Task. type Task interface { Execute(ctx context.Context, logger *slog.Logger, args []binding.Arg) error + // Declare describes the parameters the task binds, for the Dag processor to check the + // Python stub task against. + Declare(taskID string) genmodels.TaskHandlerDeclaration } // Bundle looks up a registered task by dag_id and task_id. The coordinator @@ -57,6 +61,13 @@ type EnumerableBundle interface { ListTaskHandlers() []TaskHandlerInfo } +// Registry is what a bundle binary serves: a task run looks its handler up, and a task +// handler parse lists the handlers. +type Registry interface { + Bundle + EnumerableBundle +} + type taskFunction struct { fn reflect.Value fullName string @@ -107,6 +118,10 @@ func (f *taskFunction) Execute( return f.call(ctx, sdkClient, reflectArgs, logger) } +func (f *taskFunction) Declare(taskID string) genmodels.TaskHandlerDeclaration { + return f.plan.Declare(taskID) +} + func clientFrom(ctx context.Context) (sdk.Client, error) { client, ok := ctx.Value(sdkcontext.SdkClientContextKey).(sdk.Client) if !ok { diff --git a/go-sdk/internal/bundlefooter/footer.go b/go-sdk/internal/bundlefooter/footer.go index 6641ab3af08bc..d87d1b8b57389 100644 --- a/go-sdk/internal/bundlefooter/footer.go +++ b/go-sdk/internal/bundlefooter/footer.go @@ -98,7 +98,7 @@ func Append(execPath string, source, metadata []byte) error { // Compute binary_sha256 over the file as it stands now: at this point the // whole file is the binary region, so the digest matches what a reader // recomputes over [0, source_start) after the append. - binaryHash, err := hashFile(execPath) + binaryHash, err := HashFile(execPath) if err != nil { return fmt.Errorf("bundlefooter: hashing binary region of %s: %w", execPath, err) } @@ -247,8 +247,10 @@ func decodeTrailer(b [TrailerSize]byte) (Trailer, error) { return t, nil } -// hashFile computes SHA-256 over the entire contents of the file at path. -func hashFile(path string) ([sha256.Size]byte, error) { +// HashFile computes SHA-256 over the entire contents of the file at path. For an +// executable that has not had its trailer appended yet, that is the binary_sha256 +// Append records. +func HashFile(path string) ([sha256.Size]byte, error) { var sum [sha256.Size]byte f, err := os.Open(path) if err != nil { diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go index 066f479f1bbee..40be16698c83f 100644 --- a/go-sdk/pkg/binding/binding.go +++ b/go-sdk/pkg/binding/binding.go @@ -23,7 +23,8 @@ // Captured defaults may go unclaimed, and a sole untagged struct can decode one // unclaimed argument as a whole value. // -// Analyze validates a function once. Resolve binds each execution. +// Analyze validates a function once. Resolve binds each execution. Declare describes the +// parameters for the Dag processor's check of a stub task against its handler. // AnalyzePositional builds a plan in which a sole struct binds positionally too, as one whole // argument. package binding diff --git a/go-sdk/pkg/binding/declaration.go b/go-sdk/pkg/binding/declaration.go new file mode 100644 index 0000000000000..66e73cb814727 --- /dev/null +++ b/go-sdk/pkg/binding/declaration.go @@ -0,0 +1,207 @@ +// 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. + +package binding + +import ( + "encoding/json" + "math" + "reflect" + "time" + + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +// Declare describes the parameters that the task function's TaskFlow arguments fill, so the Dag +// processor can check a stub task against its Go handler before the task runs. +// +// Flat parameters bind by position and carry no name. A lone struct, tagged or not, binds by +// field name, exactly for an `arg:`-tagged field and ignoring case and underscores otherwise. +func (p *Plan) Declare(taskID string) genmodels.TaskHandlerDeclaration { + decl := genmodels.TaskHandlerDeclaration{ + TaskID: taskID, + Binding: genmodels.TaskHandlerDeclarationBindingPositional, + } + // Not nil: a null list would tell the Dag processor that the params cannot be listed. + params := genmodels.TaskHandlerParams{} + for _, plan := range p.params { + switch plan.kind { + case paramData: + params = append(params, genmodels.TaskHandlerParam{ValueSchema: valueSchema(plan.typ)}) + case paramLoneStruct: + decl.Binding = genmodels.TaskHandlerDeclarationBindingNamed + for _, sf := range plan.fields { + params = append(params, genmodels.TaskHandlerParam{ + Name: sf.argName, + ExactName: sf.tagged, + ValueSchema: valueSchema(sf.fieldType), + }) + } + } + } + decl.Params = ¶ms + return decl +} + +// valueSchema returns the JSON Schema of the values a parameter of type t accepts, or nil when +// no schema can state them. +// +// It states only what decoding into t enforces, so a value the schema rejects is one the task +// would fail on. The exception is a null element or map value, which decoding turns into the +// zero value; like the runtime's own argument check, the schema states null only at the top level. +// The vocabulary is the one build_arg_bindings emits for a stub parameter's Python annotation, +// so the common pairs, such as int and int, come out equal. +func valueSchema(t reflect.Type) *genmodels.ArgValueSchema { + fragment := schemaFragment(t, map[reflect.Type]bool{}) + if fragment == nil { + return nil + } + schema := make(genmodels.ArgValueSchema, len(fragment)) + for k, v := range fragment { + schema[k] = v + } + return &schema +} + +var ( + timeType = reflect.TypeFor[time.Time]() + jsonNumberType = reflect.TypeFor[json.Number]() +) + +func schemaFragment(t reflect.Type, visiting map[reflect.Type]bool) map[string]any { + // A type that contains itself is left unconstrained where it recurs. + if visiting[t] { + return nil + } + visiting[t] = true + defer delete(visiting, t) + + switch t.Kind() { + case reflect.Pointer: + return nullable(schemaFragment(t.Elem(), visiting)) + case reflect.Slice, reflect.Map: + // Decoding null into a slice or map leaves it nil. + return nullable(nonNullFragment(t, visiting)) + default: + return nonNullFragment(t, visiting) + } +} + +func nonNullFragment(t reflect.Type, visiting map[reflect.Type]bool) map[string]any { + switch { + case t == timeType: + return map[string]any{"type": "string", "format": "date-time"} + // json.Number takes a number or a numeric string. + case t == jsonNumberType: + return nil + case implementsJSONUnmarshaler(t): + return nil + case implementsTextUnmarshaler(t): + return map[string]any{"type": "string"} + } + + switch t.Kind() { + case reflect.String: + return map[string]any{"type": "string"} + case reflect.Bool: + return map[string]any{"type": "boolean"} + case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64: + return signedIntFragment(t.Bits()) + case reflect.Uint, + reflect.Uint8, + reflect.Uint16, + reflect.Uint32, + reflect.Uint64, + reflect.Uintptr: + return unsignedIntFragment(t.Bits()) + case reflect.Float32: + return map[string]any{"type": "number", "format": "float"} + case reflect.Float64: + return map[string]any{"type": "number", "format": "double"} + case reflect.Slice: + // A byte slice takes a base64 string as well as an array. + if t.Elem().Kind() == reflect.Uint8 { + return nil + } + return arrayFragment(t.Elem(), visiting) + case reflect.Array: + return arrayFragment(t.Elem(), visiting) + case reflect.Map: + var values any = true + if fragment := schemaFragment(t.Elem(), visiting); fragment != nil { + values = fragment + } + return map[string]any{"type": "object", "additionalProperties": values} + case reflect.Struct: + return map[string]any{"type": "object"} + default: + return nil + } +} + +func signedIntFragment(bits int) map[string]any { + switch bits { + case 64: + return map[string]any{"type": "integer", "format": "int64"} + case 32: + return map[string]any{"type": "integer", "format": "int32"} + default: + return map[string]any{ + "type": "integer", + "minimum": -(int64(1) << (bits - 1)), + "maximum": int64(1)<<(bits-1) - 1, + } + } +} + +func unsignedIntFragment(bits int) map[string]any { + return map[string]any{ + "type": "integer", + "minimum": uint64(0), + "maximum": uint64(math.MaxUint64) >> (64 - bits), + } +} + +func arrayFragment(elem reflect.Type, visiting map[reflect.Type]bool) map[string]any { + items := schemaFragment(elem, visiting) + if items == nil { + items = map[string]any{} + } + return map[string]any{"type": "array", "items": items} +} + +func nullable(fragment map[string]any) map[string]any { + if fragment == nil { + return nil + } + if branches, isUnion := fragment["anyOf"].([]any); isUnion { + for _, branch := range branches { + if b, ok := branch.(map[string]any); ok && b["type"] == "null" { + return fragment + } + } + } + return map[string]any{"anyOf": []any{fragment, map[string]any{"type": "null"}}} +} + +func implementsJSONUnmarshaler(t reflect.Type) bool { + return t.Implements(jsonUnmarshalerType) || reflect.PointerTo(t).Implements(jsonUnmarshalerType) +} + +func implementsTextUnmarshaler(t reflect.Type) bool { + return t.Implements(textUnmarshalerType) || reflect.PointerTo(t).Implements(textUnmarshalerType) +} diff --git a/go-sdk/pkg/binding/declaration_test.go b/go-sdk/pkg/binding/declaration_test.go new file mode 100644 index 0000000000000..3e4ef2f693383 --- /dev/null +++ b/go-sdk/pkg/binding/declaration_test.go @@ -0,0 +1,246 @@ +// 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. + +package binding + +import ( + "encoding/json" + "math" + "reflect" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/apache/airflow/go-sdk/internal/contexttest" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +type declRegion string + +type declLevel int + +func (l *declLevel) UnmarshalText([]byte) error { return nil } + +type declCustom struct{ V int } + +func (c *declCustom) UnmarshalJSON([]byte) error { return nil } + +type declTree []declTree + +type declEmbedded struct{ Count int } + +type declTaggedInput struct { + Region string `arg:"region_code"` + Threshold float64 + declEmbedded +} + +type declUntaggedInput struct { + RegionCode string + Labels map[string]string +} + +var ( + stringFragment = map[string]any{"type": "string"} + int64Fragment = map[string]any{"type": "integer", "format": "int64"} + doubleFragment = map[string]any{"type": "number", "format": "double"} +) + +func nullOr(fragment map[string]any) map[string]any { + return map[string]any{"anyOf": []any{fragment, map[string]any{"type": "null"}}} +} + +func schemaOf(fragment map[string]any) *genmodels.ArgValueSchema { + schema := make(genmodels.ArgValueSchema, len(fragment)) + for k, v := range fragment { + schema[k] = v + } + return &schema +} + +func TestValueSchema(t *testing.T) { + for name, tc := range map[string]struct { + typ reflect.Type + want map[string]any + }{ + "string": {reflect.TypeFor[string](), stringFragment}, + "named string": {reflect.TypeFor[declRegion](), stringFragment}, + "bool": {reflect.TypeFor[bool](), map[string]any{"type": "boolean"}}, + "int": {reflect.TypeFor[int](), int64Fragment}, + "int64": {reflect.TypeFor[int64](), int64Fragment}, + "int32": { + reflect.TypeFor[int32](), + map[string]any{"type": "integer", "format": "int32"}, + }, + "int8": { + reflect.TypeFor[int8](), + map[string]any{"type": "integer", "minimum": int64(-128), "maximum": int64(127)}, + }, + "int16": { + reflect.TypeFor[int16](), + map[string]any{"type": "integer", "minimum": int64(-32768), "maximum": int64(32767)}, + }, + "uint8": { + reflect.TypeFor[uint8](), + map[string]any{"type": "integer", "minimum": uint64(0), "maximum": uint64(255)}, + }, + "uint32": { + reflect.TypeFor[uint32](), + map[string]any{"type": "integer", "minimum": uint64(0), "maximum": uint64(math.MaxUint32)}, + }, + "uint64": { + reflect.TypeFor[uint64](), + map[string]any{"type": "integer", "minimum": uint64(0), "maximum": uint64(math.MaxUint64)}, + }, + "float32": { + reflect.TypeFor[float32](), + map[string]any{"type": "number", "format": "float"}, + }, + "float64": {reflect.TypeFor[float64](), doubleFragment}, + "time.Time": { + reflect.TypeFor[time.Time](), + map[string]any{"type": "string", "format": "date-time"}, + }, + "pointer to time.Time": { + reflect.TypeFor[*time.Time](), + nullOr(map[string]any{"type": "string", "format": "date-time"}), + }, + "text unmarshaler on the pointer": {reflect.TypeFor[declLevel](), stringFragment}, + "json unmarshaler": {reflect.TypeFor[declCustom](), nil}, + "json.Number": {reflect.TypeFor[json.Number](), nil}, + "byte slice": {reflect.TypeFor[[]byte](), nil}, + "any": {reflect.TypeFor[any](), nil}, + "pointer to any": {reflect.TypeFor[*any](), nil}, + "complex": {reflect.TypeFor[complex128](), nil}, + "slice": { + reflect.TypeFor[[]string](), + nullOr(map[string]any{"type": "array", "items": stringFragment}), + }, + "slice of any": { + reflect.TypeFor[[]any](), + nullOr(map[string]any{"type": "array", "items": map[string]any{}}), + }, + "array": { + reflect.TypeFor[[3]int](), + map[string]any{"type": "array", "items": int64Fragment}, + }, + "map": { + reflect.TypeFor[map[string]int](), + nullOr(map[string]any{"type": "object", "additionalProperties": int64Fragment}), + }, + "map of any": { + reflect.TypeFor[map[string]any](), + nullOr(map[string]any{"type": "object", "additionalProperties": true}), + }, + "struct": {reflect.TypeFor[declUntaggedInput](), map[string]any{"type": "object"}}, + "pointer": {reflect.TypeFor[*int](), nullOr(int64Fragment)}, + "pointer to pointer": {reflect.TypeFor[**int](), nullOr(int64Fragment)}, + "pointer to slice": { + reflect.TypeFor[*[]int](), + nullOr(map[string]any{"type": "array", "items": int64Fragment}), + }, + "recursive": { + reflect.TypeFor[declTree](), + nullOr(map[string]any{"type": "array", "items": map[string]any{}}), + }, + } { + t.Run(name, func(t *testing.T) { + got := valueSchema(tc.typ) + if tc.want == nil { + assert.Nil(t, got) + return + } + assert.Equal(t, schemaOf(tc.want), got) + }) + } +} + +func declare(t *testing.T, fn any) genmodels.TaskHandlerDeclaration { + t.Helper() + plan, err := Analyze(reflect.TypeOf(fn), "testFn") + require.NoError(t, err) + return plan.Declare("task") +} + +func TestDeclareFlatParams(t *testing.T) { + decl := declare(t, func( + actx contexttest.Context, region string, count int, note *string, config declUntaggedInput, + ) error { + return nil + }) + + assert.Equal(t, genmodels.TaskHandlerDeclaration{ + TaskID: "task", + Binding: genmodels.TaskHandlerDeclarationBindingPositional, + Params: &genmodels.TaskHandlerParams{ + {ValueSchema: schemaOf(stringFragment)}, + {ValueSchema: schemaOf(int64Fragment)}, + {ValueSchema: schemaOf(nullOr(stringFragment))}, + {ValueSchema: schemaOf(map[string]any{"type": "object"})}, + }, + }, decl) +} + +func TestDeclareTaggedStruct(t *testing.T) { + decl := declare(t, func(actx contexttest.Context, in declTaggedInput) error { return nil }) + + assert.Equal(t, genmodels.TaskHandlerDeclaration{ + TaskID: "task", + Binding: genmodels.TaskHandlerDeclarationBindingNamed, + Params: &genmodels.TaskHandlerParams{ + {Name: "region_code", ExactName: true, ValueSchema: schemaOf(stringFragment)}, + {Name: "Threshold", ValueSchema: schemaOf(doubleFragment)}, + {Name: "Count", ValueSchema: schemaOf(int64Fragment)}, + }, + }, decl) +} + +func TestDeclareUntaggedStruct(t *testing.T) { + want := genmodels.TaskHandlerDeclaration{ + TaskID: "task", + Binding: genmodels.TaskHandlerDeclarationBindingNamed, + Params: &genmodels.TaskHandlerParams{ + {Name: "RegionCode", ValueSchema: schemaOf(stringFragment)}, + { + Name: "Labels", + ValueSchema: schemaOf(nullOr(map[string]any{ + "type": "object", "additionalProperties": stringFragment, + })), + }, + }, + } + for name, fn := range map[string]any{ + "value": func(actx contexttest.Context, in declUntaggedInput) error { return nil }, + "pointer": func(actx contexttest.Context, in *declUntaggedInput) error { return nil }, + } { + t.Run(name, func(t *testing.T) { + assert.Equal(t, want, declare(t, fn)) + }) + } +} + +func TestDeclareContextOnly(t *testing.T) { + decl := declare(t, func(actx contexttest.Context) error { return nil }) + + assert.Equal(t, genmodels.TaskHandlerDeclaration{ + TaskID: "task", + Binding: genmodels.TaskHandlerDeclarationBindingPositional, + Params: &genmodels.TaskHandlerParams{}, + }, decl) +} diff --git a/go-sdk/pkg/execution/genmodels/discriminators.gen.go b/go-sdk/pkg/execution/genmodels/discriminators.gen.go index 09da3118f55da..0fe7a014ca505 100644 --- a/go-sdk/pkg/execution/genmodels/discriminators.gen.go +++ b/go-sdk/pkg/execution/genmodels/discriminators.gen.go @@ -100,6 +100,8 @@ const ( TypeTICount = "TICount" TypeTaskBreadcrumbsResult = "TaskBreadcrumbsResult" TypeTaskCallbackRequest = "TaskCallbackRequest" + TypeTaskHandlerParseRequest = "TaskHandlerParseRequest" + TypeTaskHandlerParsingResult = "TaskHandlerParsingResult" TypeTaskRescheduleStartDate = "TaskRescheduleStartDate" TypeTaskState = "TaskState" TypeTaskStateStoreResult = "TaskStateStoreResult" @@ -350,6 +352,12 @@ func EnsureType(m any) any { case TaskCallbackRequest: b.Type = TypeTaskCallbackRequest return b + case TaskHandlerParseRequest: + b.Type = TypeTaskHandlerParseRequest + return b + case TaskHandlerParsingResult: + b.Type = TypeTaskHandlerParsingResult + return b case TaskRescheduleStartDate: b.Type = TypeTaskRescheduleStartDate return b diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index dee470082f124..a6b4542757690 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -397,6 +397,9 @@ type DagFileParseRequest struct { // File corresponds to the JSON schema field "file". File string `msgpack:"file"` + // KnownArtifacts corresponds to the JSON schema field "known_artifacts". + KnownArtifacts []TaskHandlerArtifact `msgpack:"known_artifacts,omitempty"` + // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` } @@ -1714,8 +1717,90 @@ type TaskCallbackRequest struct { VersionData *VersionData `msgpack:"version_data,omitempty"` } +type Warnings []interface{} + +type TriggerKwargs map[string]interface{} + +// A Lang-SDK artifact and every task handler it registers. +type TaskHandlerArtifact struct { + // BundleName corresponds to the JSON schema field "bundle_name". + BundleName string `msgpack:"bundle_name"` + + // CacheDigest corresponds to the JSON schema field "cache_digest". + CacheDigest interface{} `msgpack:"cache_digest"` + + // RelativeFileloc corresponds to the JSON schema field "relative_fileloc". + RelativeFileloc string `msgpack:"relative_fileloc"` + + // SizeBytes corresponds to the JSON schema field "size_bytes". + SizeBytes int `msgpack:"size_bytes"` + + // TaskHandlers corresponds to the JSON schema field "task_handlers". + TaskHandlers TaskHandlers `msgpack:"task_handlers"` +} + +type TaskHandlers map[string][]TaskHandlerDeclaration + +// A task handler that a Lang-SDK artifact registers for one task. +type TaskHandlerDeclaration struct { + // Binding corresponds to the JSON schema field "binding". + Binding TaskHandlerDeclarationBinding `msgpack:"binding"` + + // Params corresponds to the JSON schema field "params". + Params *TaskHandlerParams `msgpack:"params"` + + // TaskID corresponds to the JSON schema field "task_id". + TaskID string `msgpack:"task_id"` +} + +type TaskHandlerParams []TaskHandlerParam + +// One parameter of a task handler. +type TaskHandlerParam struct { + // ExactName corresponds to the JSON schema field "exact_name". + ExactName bool `msgpack:"exact_name,omitempty"` + + // Name corresponds to the JSON schema field "name". + Name interface{} `msgpack:"name"` + + // ValueSchema corresponds to the JSON schema field "value_schema". + ValueSchema *ArgValueSchema `msgpack:"value_schema,omitempty"` +} + type TaskIds []string +const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" +const TaskInstanceStateDeferred TaskInstanceState = "deferred" +const TaskInstanceStateSkipped TaskInstanceState = "skipped" +const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" +const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" +const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" +const TaskInstanceStateFailed TaskInstanceState = "failed" +const TaskInstanceStateRestarting TaskInstanceState = "restarting" +const TaskInstanceStateSuccess TaskInstanceState = "success" +const TaskInstanceStateRunning TaskInstanceState = "running" +const TaskInstanceStateQueued TaskInstanceState = "queued" +const TaskInstanceStateScheduled TaskInstanceState = "scheduled" +const TaskInstanceStateRemoved TaskInstanceState = "removed" + +type TaskInstanceState string + +// Variable schema for responses with fields that are needed for Runtime. +type VariableResponse struct { + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key"` + + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` +} + +const TaskHandlerDeclarationBindingNamed TaskHandlerDeclarationBinding = "named" +const TaskHandlerDeclarationBindingPositional TaskHandlerDeclarationBinding = "positional" + +type TaskHandlerDeclarationBinding string + +type TaskOutlets []AssetProfile + // Schema for TaskInstance model with minimal required fields needed for Runtime. type TaskInstance struct { // ContextCarrier corresponds to the JSON schema field "context_carrier". @@ -1749,23 +1834,44 @@ type TaskInstance struct { TryNumber int `msgpack:"try_number"` } -type TaskInstanceState string +// Request for Task Handler Parsing. +// +// Asks a Lang-SDK runtime for every task handler an artifact registers. +type TaskHandlerParseRequest struct { + // BundleName corresponds to the JSON schema field "bundle_name". + BundleName string `msgpack:"bundle_name"` -const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" -const TaskInstanceStateDeferred TaskInstanceState = "deferred" -const TaskInstanceStateFailed TaskInstanceState = "failed" -const TaskInstanceStateQueued TaskInstanceState = "queued" -const TaskInstanceStateRemoved TaskInstanceState = "removed" -const TaskInstanceStateRestarting TaskInstanceState = "restarting" -const TaskInstanceStateRunning TaskInstanceState = "running" -const TaskInstanceStateScheduled TaskInstanceState = "scheduled" -const TaskInstanceStateSkipped TaskInstanceState = "skipped" -const TaskInstanceStateSuccess TaskInstanceState = "success" -const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" -const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" -const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" + // BundlePath corresponds to the JSON schema field "bundle_path". + BundlePath string `msgpack:"bundle_path"` -type TaskOutlets []AssetProfile + // File corresponds to the JSON schema field "file". + File string `msgpack:"file"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` +} + +// Result of Task Handler Parsing. +// +// Every task handler a Lang-SDK artifact registers, keyed by Dag id. +// +// The answer depends only on the artifact, never on the request. +type TaskHandlerParsingResult struct { + // Fileloc corresponds to the JSON schema field "fileloc". + Fileloc string `msgpack:"fileloc"` + + // ImportErrors corresponds to the JSON schema field "import_errors". + ImportErrors *ImportErrors `msgpack:"import_errors,omitempty"` + + // TaskHandlers corresponds to the JSON schema field "task_handlers". + TaskHandlers TaskHandlers `msgpack:"task_handlers"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` + + // Warnings corresponds to the JSON schema field "warnings". + Warnings *Warnings `msgpack:"warnings,omitempty"` +} // Response containing the first reschedule date for a task instance. type TaskRescheduleStartDate struct { @@ -1776,6 +1882,12 @@ type TaskRescheduleStartDate struct { Type string `msgpack:"type,omitempty"` } +type TaskStateState string + +const TaskStateStateFailed TaskStateState = "failed" +const TaskStateStateSkipped TaskStateState = "skipped" +const TaskStateStateRemoved TaskStateState = "removed" + // Update a task's state. // // If a process exits without sending one of these the state will be derived from @@ -1799,12 +1911,6 @@ type TaskState struct { Type string `msgpack:"type,omitempty"` } -type TaskStateState string - -const TaskStateStateFailed TaskStateState = "failed" -const TaskStateStateRemoved TaskStateState = "removed" -const TaskStateStateSkipped TaskStateState = "skipped" - // Response to GetTaskStateStore; wraps the generated API response for supervisor // to worker comms. type TaskStateStoreResult struct { @@ -1854,8 +1960,6 @@ type TriggerDagRun struct { Type string `msgpack:"type,omitempty"` } -type TriggerKwargs map[string]interface{} - type UpdateDagRunNote struct { // Note corresponds to the JSON schema field "note". Note interface{} `msgpack:"note"` @@ -1867,6 +1971,8 @@ type UpdateDagRunNote struct { Type string `msgpack:"type,omitempty"` } +type VersionData map[string]interface{} + // Update the response content part of an existing Human-in-the-loop response. type UpdateHITLDetail struct { // ChosenOptions corresponds to the JSON schema field "chosen_options". @@ -1882,19 +1988,6 @@ type UpdateHITLDetail struct { Type string `msgpack:"type,omitempty"` } -// Variable schema for responses with fields that are needed for Runtime. -type VariableResponse struct { - // Key corresponds to the JSON schema field "key". - Key string `msgpack:"key"` - - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` -} - -type Warnings []interface{} - -type VersionData map[string]interface{} - type ValidateInletsAndOutlets struct { // TIID corresponds to the JSON schema field "ti_id". TIID string `msgpack:"ti_id"` diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index c53bed89695b2..dd37b0347eefd 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -101,7 +101,29 @@ func (b testBundle) LookupTask(dagID, taskID string) (bundle.Task, bool) { return task, ok } -func buildBundle(t *testing.T, register func(testBundle)) bundle.Bundle { +// ListTaskHandlers lists nothing: a map keeps no registration order, and only a task +// handler parse lists handlers. The parse tests use handlerRegistry instead. +func (b testBundle) ListTaskHandlers() []bundle.TaskHandlerInfo { return nil } + +// handlerRegistry lists its task handlers in the order they were added, as airflow.Bundle does. +type handlerRegistry struct { + testBundle + order []bundle.TaskHandlerInfo +} + +func newHandlerRegistry() *handlerRegistry { return &handlerRegistry{testBundle: testBundle{}} } + +func (r *handlerRegistry) add(dagID, taskID string, fn any) { + if _, ok := r.testBundle[dagID]; !ok { + r.AddDag(dagID) + } + r.testBundle[dagID].AddTaskWithName(taskID, fn) + r.order = append(r.order, bundle.TaskHandlerInfo{DagID: dagID, TaskID: taskID}) +} + +func (r *handlerRegistry) ListTaskHandlers() []bundle.TaskHandlerInfo { return r.order } + +func buildBundle(t *testing.T, register func(testBundle)) bundle.Registry { t.Helper() b := testBundle{} register(b) @@ -893,3 +915,114 @@ func TestServeFailureAfterConnectClosesComm(t *testing.T) { _, err = readFrame(commConn) require.Error(t, err) } + +type loadInput struct { + Table string `arg:"table_name"` +} + +// startTaskHandlerParse runs Serve for a registry with two handlers, sends it a +// TaskHandlerParseRequest, and returns the comm connection and the reply frame. +func startTaskHandlerParse(t *testing.T) (net.Conn, IncomingFrame, chan error) { + t.Helper() + commAddr, logsAddr, commCh, logsCh, cleanup := startSupervisor(t) + t.Cleanup(cleanup) + + registry := newHandlerRegistry() + registry.add("etl", "extract", func(actx contexttest.Context, region string, day *int) error { + return nil + }) + registry.add("etl", "load", func(actx contexttest.Context, in loadInput) error { return nil }) + registry.add("etl", "notify", simpleTask) + + done := make(chan error, 1) + go func() { done <- Serve(registry, commAddr, logsAddr) }() + + commConn := <-commCh + t.Cleanup(func() { commConn.Close() }) + logsConn := <-logsCh + t.Cleanup(func() { logsConn.Close() }) + require.NoError(t, commConn.SetDeadline(time.Now().Add(10*time.Second))) + + payload, err := encodeRequest(0, map[string]any{ + "type": "TaskHandlerParseRequest", + "file": "/bundles/go/etl", + "bundle_path": "/bundles/go", + "bundle_name": "go", + }) + require.NoError(t, err) + require.NoError(t, writeFrame(commConn, payload)) + + frame, err := readFrame(commConn) + require.NoError(t, err) + require.True(t, isNilRaw(frame.Err)) + return commConn, frame, done +} + +func TestServeTaskHandlerParseRequestEndToEnd(t *testing.T) { + commConn, frame, done := startTaskHandlerParse(t) + + // Decoded as a map, so the wire's keys, nulls and empty lists are what is compared. + assert.Equal(t, map[string]any{ + "type": "TaskHandlerParsingResult", + "fileloc": "/bundles/go/etl", + "task_handlers": map[string]any{ + "etl": []any{ + map[string]any{ + "task_id": "extract", + "binding": "positional", + "params": []any{ + map[string]any{ + "name": nil, + "value_schema": map[string]any{"type": "string"}, + }, + map[string]any{ + "name": nil, + "value_schema": map[string]any{"anyOf": []any{ + map[string]any{"type": "integer", "format": "int64"}, + map[string]any{"type": "null"}, + }}, + }, + }, + }, + map[string]any{ + "task_id": "load", + "binding": "named", + "params": []any{map[string]any{ + "name": "table_name", + "exact_name": true, + "value_schema": map[string]any{"type": "string"}, + }}, + }, + map[string]any{"task_id": "notify", "binding": "positional", "params": []any{}}, + }, + }, + }, rawToMap(t, frame.Body)) + + // Serve returns only once the parent acknowledges the result. + select { + case err := <-done: + t.Fatalf("Serve returned before the result was acknowledged: %v", err) + case <-time.After(100 * time.Millisecond): + } + require.NoError(t, writeFrame(commConn, encodeResponseFrame(t, frame.ID, nil, nil))) + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return after the result was acknowledged") + } +} + +func TestServeTaskHandlerParseFailsWithoutAcknowledgement(t *testing.T) { + commConn, _, done := startTaskHandlerParse(t) + + require.NoError(t, commConn.Close()) + + select { + case err := <-done: + require.ErrorContains(t, err, "sending task handler parse result") + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return after the parent closed the socket") + } +} diff --git a/go-sdk/pkg/execution/messages.go b/go-sdk/pkg/execution/messages.go index 1dd596201258a..77df9add847a6 100644 --- a/go-sdk/pkg/execution/messages.go +++ b/go-sdk/pkg/execution/messages.go @@ -75,6 +75,12 @@ func decodeIncomingBody(raw msgpack.RawMessage) (any, error) { return nil, fmt.Errorf("decoding StartupDetails: %w", err) } return &msg, nil + case genmodels.TypeTaskHandlerParseRequest: + var msg genmodels.TaskHandlerParseRequest + if err := msgpack.Unmarshal(raw, &msg); err != nil { + return nil, fmt.Errorf("decoding TaskHandlerParseRequest: %w", err) + } + return &msg, nil case genmodels.TypeDagFileParseRequest: var msg genmodels.DagFileParseRequest if err := msgpack.Unmarshal(raw, &msg); err != nil { diff --git a/go-sdk/pkg/execution/messages_test.go b/go-sdk/pkg/execution/messages_test.go index fa81df4ea66d2..3cf8e1a742e85 100644 --- a/go-sdk/pkg/execution/messages_test.go +++ b/go-sdk/pkg/execution/messages_test.go @@ -462,6 +462,29 @@ func TestDecodeIncomingBodyDispatch(t *testing.T) { assert.True(t, ok) }) + t.Run("TaskHandlerParseRequest", func(t *testing.T) { + raw := marshalBody(t, map[string]any{ + "type": "TaskHandlerParseRequest", + "file": "/bundles/go/etl", + "bundle_path": "/bundles/go", + "bundle_name": "go", + }) + result, err := decodeIncomingBody(raw) + require.NoError(t, err) + assert.Equal(t, &genmodels.TaskHandlerParseRequest{ + Type: "TaskHandlerParseRequest", + File: "/bundles/go/etl", + BundlePath: "/bundles/go", + BundleName: "go", + }, result) + }) + + t.Run("malformed TaskHandlerParseRequest", func(t *testing.T) { + raw := marshalBody(t, map[string]any{"type": "TaskHandlerParseRequest", "file": 42}) + _, err := decodeIncomingBody(raw) + assert.ErrorContains(t, err, "decoding TaskHandlerParseRequest") + }) + t.Run("ErrorResponse", func(t *testing.T) { raw := marshalBody(t, map[string]any{"type": "ErrorResponse", "error": "GENERIC_ERROR"}) result, err := decodeIncomingBody(raw) diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index 0f56e005315ca..2596a53ceb823 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -20,8 +20,10 @@ // the Airflow supervisor (Python ExecutableCoordinator), the Serve method of // airflow.BundleRef dispatches here. // -// The first inbound frame on the comm socket is a StartupDetails message -// that drives multi-round task execution. +// The first inbound frame on the comm socket selects what the runtime does. A +// StartupDetails message drives multi-round task execution. A +// TaskHandlerParseRequest from the Dag processor gets one TaskHandlerParsingResult +// that declares every registered task handler. // // See go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md. package execution @@ -48,21 +50,23 @@ import ( const dialTimeout = 30 * time.Second // terminalSendTimeout bounds the write of the final TaskState/SucceedTask -// frame. The supervisor normally drains the comm socket promptly, but a -// half-open connection (the supervisor gone without a clean close) could -// otherwise wedge the runtime on a blocked write; the deadline turns that -// into a fast failure -- and thus a non-zero exit -- instead of a hang. +// frame, and the send and acknowledgement of a task handler parse result. The +// supervisor normally drains the comm socket promptly, but a half-open +// connection (the supervisor gone without a clean close) could otherwise wedge +// the runtime on a blocked write or wait; the deadline turns that into a fast +// failure -- and thus a non-zero exit -- instead of a hang. const terminalSendTimeout = 30 * time.Second // Serve runs the bundle binary in coordinator mode. It dials the supervisor's // comm and logs sockets, installs an slog handler that writes JSON-line // records to the logs connection, and dispatches on the first frame. // -// Serve returns nil on a clean shutdown: the task ran and its terminal -// TaskState/SucceedTask frame was delivered, and the caller should exit 0. A -// non-nil error indicates a protocol-level failure (connection loss, -// malformed frames, unknown first message type) that happens before or -// instead of delivering a terminal frame. +// Serve returns nil on a clean shutdown, and the caller should exit 0: either +// the task ran and its terminal TaskState/SucceedTask frame was delivered, or +// the task handler parse result was delivered and acknowledged. A non-nil error +// indicates a protocol-level failure (connection loss, malformed frames, +// unknown first message type) that happens before or instead of delivering a +// terminal frame. // // Failure-signaling contract: the caller (main) must turn a non-nil error // into a non-zero process exit. The supervisor derives the task's final state @@ -73,7 +77,7 @@ const terminalSendTimeout = 30 * time.Second // fails closed without needing to send a frame; the post-connect paths below // log the reason at Error first so it still reaches the supervisor's log // stream over the already-connected logs socket. -func Serve(b bundle.Bundle, commAddr, logsAddr string) error { +func Serve(b bundle.Registry, commAddr, logsAddr string) error { if commAddr == "" { return fmt.Errorf("missing --comm=host:port argument") } @@ -167,6 +171,18 @@ func Serve(b bundle.Bundle, commAddr, logsAddr string) error { } logger.Debug("Task execution complete") + case *genmodels.TaskHandlerParseRequest: + logger.Debug("Task handler parse mode", "file", msg.File) + result := declareTaskHandlers(b, msg) + sendCtx, cancel := context.WithTimeout(ctx, terminalSendTimeout) + defer cancel() + // Waiting for the acknowledgement means the parent holds the result before this + // process exits and closes the socket under it. + if _, err := comm.Communicate(sendCtx, result); err != nil { + return fmt.Errorf("sending task handler parse result: %w", err) + } + logger.Debug("Task handler parse complete") + default: logger.Error("Unexpected initial message type", "type", fmt.Sprintf("%T", body)) return fmt.Errorf("unexpected initial message type: %T", body) diff --git a/go-sdk/pkg/execution/task_handler_parse.go b/go-sdk/pkg/execution/task_handler_parse.go new file mode 100644 index 0000000000000..1e2d5aca2f5f0 --- /dev/null +++ b/go-sdk/pkg/execution/task_handler_parse.go @@ -0,0 +1,42 @@ +// 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. + +package execution + +import ( + "github.com/apache/airflow/go-sdk/internal/bundle" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +// declareTaskHandlers answers a TaskHandlerParseRequest with the declarations of every +// registered task handler, keyed by Dag id, each Dag's in registration order. +func declareTaskHandlers( + b bundle.Registry, + req *genmodels.TaskHandlerParseRequest, +) genmodels.TaskHandlerParsingResult { + // A nil map would go out as null, which the Dag processor rejects. + handlers := genmodels.TaskHandlers{} + for _, info := range b.ListTaskHandlers() { + // Declare only what a task run could find. + task, ok := b.LookupTask(info.DagID, info.TaskID) + if !ok { + continue + } + handlers[info.DagID] = append(handlers[info.DagID], task.Declare(info.TaskID)) + } + return genmodels.TaskHandlerParsingResult{Fileloc: req.File, TaskHandlers: handlers} +} diff --git a/go-sdk/pkg/execution/task_handler_parse_test.go b/go-sdk/pkg/execution/task_handler_parse_test.go new file mode 100644 index 0000000000000..47237b5a8e719 --- /dev/null +++ b/go-sdk/pkg/execution/task_handler_parse_test.go @@ -0,0 +1,77 @@ +// 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. + +package execution + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/apache/airflow/go-sdk/internal/bundle" + "github.com/apache/airflow/go-sdk/internal/contexttest" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +func TestDeclareTaskHandlers(t *testing.T) { + registry := newHandlerRegistry() + registry.add("etl", "transform", func(actx contexttest.Context, region string) error { + return nil + }) + registry.add("report", "render", simpleTask) + registry.add("etl", "extract", simpleTask) + // Listed, but a task run could not find it. + registry.order = append(registry.order, bundle.TaskHandlerInfo{DagID: "etl", TaskID: "ghost"}) + + result := declareTaskHandlers( + registry, + &genmodels.TaskHandlerParseRequest{File: "/bundles/go/etl"}, + ) + + assert.Equal(t, genmodels.TaskHandlerParsingResult{ + Fileloc: "/bundles/go/etl", + TaskHandlers: genmodels.TaskHandlers{ + "etl": { + { + TaskID: "transform", + Binding: genmodels.TaskHandlerDeclarationBindingPositional, + Params: &genmodels.TaskHandlerParams{ + {ValueSchema: &genmodels.ArgValueSchema{"type": "string"}}, + }, + }, + { + TaskID: "extract", + Binding: genmodels.TaskHandlerDeclarationBindingPositional, + Params: &genmodels.TaskHandlerParams{}, + }, + }, + "report": { + { + TaskID: "render", + Binding: genmodels.TaskHandlerDeclarationBindingPositional, + Params: &genmodels.TaskHandlerParams{}, + }, + }, + }, + }, result) +} + +func TestDeclareTaskHandlersWithoutHandlers(t *testing.T) { + result := declareTaskHandlers(newHandlerRegistry(), &genmodels.TaskHandlerParseRequest{}) + + assert.Equal(t, genmodels.TaskHandlers{}, result.TaskHandlers) +} diff --git a/java-sdk/README.md b/java-sdk/README.md index b8600ab5d44fa..f757301b7e1ce 100644 --- a/java-sdk/README.md +++ b/java-sdk/README.md @@ -93,21 +93,37 @@ Now `cd example` into the example project, and * Put the [DAG with stub tasks](./example/src/resources/dags) to somewhere Airflow can find. -* Ensure the `java` command is available in the same environment the Airflow - task worker is in. +* Ensure the `java` command is available in the same environments the Airflow + task worker and the Dag processor are in. -* Configure Airflow to route tasks in the *java* queue to be run with Java: +* Register the packaged example as a Dag bundle, and configure Airflow to route tasks in the *java* + queue to be run with Java from it: ```bash + export AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='[ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + { + "name": "java-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/opt/airflow/java-sdk/example/build/bundle"} + } + ]' export AIRFLOW__SDK__COORDINATORS='{ "java": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", - "kwargs": {"jars_root": ["/opt/airflow/java-sdk/example/build/bundle"]} + "kwargs": {"task_handler_bundle_name": "java-task-handlers"} } }' export AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{"java": "java"}' ``` + Set these, and make the bundle directory available, for the Dag processor + as well as the worker: the Dag processor checks the stub tasks against the + task handlers the JARs register. The API server does not need them. + `task_handler_bundle_name` is used only by mixed-language Dags, to locate + the task handlers for the `@task.stub` tasks of a Python Dag; Dags defined + natively in a language SDK do not use it. + * Ensure the Connection and Variable needed by the example DAG are available: ```bash @@ -644,9 +660,9 @@ where Airflow can find it. When the Airflow supervisor identifies that a task should run with Java, it launches the JVM application as a subprocess. The flow is: -1. `JavaCoordinator.execute_task()` (Python) scans `jars_root`, builds the - classpath, and spawns `java -cp --comm=: - --logs=:`. +1. `JavaCoordinator.execute_task()` (Python) scans the Dag bundle named by + `task_handler_bundle_name`, builds the classpath, and spawns + `java -cp --comm=: --logs=:`. 2. `Server.kt` connects to both sockets immediately on startup. 3. The supervisor sends a `StartupDetails` MessagePack message; the JVM reads it, looks up the matching task by `dag_id` + `task_id`, and calls the @@ -658,6 +674,14 @@ launches the JVM application as a subprocess. The flow is: 5. On completion (or exception) the JVM sends a `TaskState` message and closes the socket. The JVM process then exits. +The Dag processor also launches each JAR to ask which task handlers it +registers. The first frame is then a `TaskHandlerParseRequest`, and +`execution/TaskHandlerParse.kt` replies with one `TaskHandlerParsingResult`: +every task-handler registration, each with its binding mode and parameters. Flat data parameters bind `positional` and a `TaskInput` +binds `named`. The annotation processor records a handler's parameter names on +its generated task class (`internal/TaskParams.kt`), because Java does not keep +them at run time. The JVM waits for the reply's acknowledgement, then exits. + Log messages produced by the SDK (not by user code) are forwarded over the `--logs` socket so the supervisor can append them to Airflow's log store. diff --git a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java index d75f5762ba12d..7ce2fd4de9200 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java @@ -23,8 +23,7 @@ public class ExampleBundleBuilder { public static Bundle build() { - return new Bundle() - .register(InterfaceExampleBuilder.build()) + return InterfaceExampleBuilder.registerInto(new Bundle()) .register(AnnotationExample.class) .register(XComCastingExample.class) .register(org.apache.airflow.example.nativedag.InterfaceExample.build()); diff --git a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java index 4d0a494d6dcaa..0a07a20ccdc22 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java @@ -89,10 +89,12 @@ public void execute(@NotNull Context context, Client client, SummarizeInput inpu } } - public static DagDef build() { - return new DagDef("java_interface_example") - .addTask("extract", Extract.class) - .addTask("transform", Transform.class) - .addTask("summarize", Summarize.class); + // The Python Dag file owns java_interface_example, so its tasks register as + // task handlers rather than as a Dag declared in Java. + public static Bundle registerInto(Bundle bundle) { + return bundle + .register("java_interface_example", "extract", Extract.class) + .register("java_interface_example", "transform", Transform.class) + .register("java_interface_example", "summarize", Summarize.class); } } diff --git a/java-sdk/plugin/build.gradle.kts b/java-sdk/plugin/build.gradle.kts index edf930bb6fc84..7efcdc36e612c 100644 --- a/java-sdk/plugin/build.gradle.kts +++ b/java-sdk/plugin/build.gradle.kts @@ -25,6 +25,12 @@ plugins { dependencies { implementation("com.gradleup.shadow:shadow-gradle-plugin:9.1.0") // Last supporting Java 11. + + testImplementation(kotlin("test")) +} + +tasks.withType { + useJUnitPlatform() } gradlePlugin { diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt index 95836c12fc5fe..817f7c0a32176 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt @@ -23,6 +23,7 @@ import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar import org.gradle.api.Plugin import org.gradle.api.Project import org.gradle.api.provider.Property +import org.gradle.api.tasks.ClasspathNormalizer import org.gradle.api.tasks.Copy import org.gradle.api.tasks.Input import org.gradle.api.tasks.SourceSetContainer @@ -83,8 +84,8 @@ abstract class AirflowBundleExtension { * * The plugin automatically sets the `Main-Class` metadata, and provides a new * task `bundle` to create a Dag bundle in one command. This builds deploy-ready - * artifacts to `build/bundle/` that can be copied directly into an Airflow Java - * coordinator's `jars_root`. + * artifacts to `build/bundle/` that can be copied directly into the Dag bundle + * named by an Airflow Java coordinator's `task_handler_bundle_name`. * * By default, plugin `com.github.johnrengelman.shadow` is applied automatically * to enable fat JAR build. In this mode, one single JAR with user code and all @@ -98,6 +99,10 @@ abstract class AirflowBundleExtension { * directory instead. In this mode, `Airflow-Supervisor-Schema-Version` lives in * the `airflow-sdk` JAR instead. The bundle JAR still contains `Main-Class`. * + * In both modes the bundle JAR also carries `Airflow-Cache-Digest`, a SHA-256 + * over the bundle's classes, resources, and runtime dependencies. It stays the + * same when a rebuild changes none of them, so Airflow can tell that a JAR it + * has already inspected is unchanged. */ class AirflowSdkPlugin : Plugin { override fun apply(project: Project) { @@ -115,15 +120,25 @@ class AirflowSdkPlugin : Plugin { } } - val classFiles = - project.objects.fileCollection().from( - project.extensions - .getByType(SourceSetContainer::class.java) - .getByName("main") - .output - .classesDirs, - project.configurations.getByName("runtimeClasspath"), - ) + val mainOutput = + project.extensions + .getByType(SourceSetContainer::class.java) + .getByName("main") + .output + val runtimeClasspath = project.configurations.getByName("runtimeClasspath") + val classFiles = project.objects.fileCollection().from(mainOutput.classesDirs, runtimeClasspath) + + fun stampCacheDigest(task: Jar) { + // Plain file collections and values only, which the configuration cache can store. + val outputFiles = project.objects.fileCollection().from(mainOutput) + val classpathFiles = project.objects.fileCollection().from(runtimeClasspath) + val mainClass = ext.mainClass + val fatJar = ext.fatJar.get() + task.doFirst { + val digest = BundleDigest.compute(mainClass.orNull, fatJar, outputFiles, classpathFiles) + task.manifest.attributes(mapOf(BundleDigest.ATTRIBUTE to digest)) + } + } val verifyTask = project.tasks.register("verifyBundleMainClass") { task -> @@ -195,6 +210,7 @@ class AirflowSdkPlugin : Plugin { mapOf("Airflow-Supervisor-Schema-Version" to schemaVersionProvider.get()), ) } + stampCacheDigest(task) } project.tasks.register("bundle", Copy::class.java) { task -> @@ -205,6 +221,15 @@ class AirflowSdkPlugin : Plugin { task.into(project.layout.buildDirectory.dir("bundle")) } } else { + project.tasks.named("jar", Jar::class.java).configure { task -> + // The digest covers the dependencies copied next to the thin JAR, so + // a dependency change has to rebuild it. + task.inputs + .files(runtimeClasspath) + .withPropertyName("airflowCacheDigestClasspath") + .withNormalizer(ClasspathNormalizer::class.java) + stampCacheDigest(task) + } // bundle copies the thin JAR and all runtime dependency JARs into // build/bundle/, mirroring what installDist puts in lib/. project.tasks.register("bundle", Copy::class.java) { task -> diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/BundleDigest.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/BundleDigest.kt new file mode 100644 index 0000000000000..e7c09decb90da --- /dev/null +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/BundleDigest.kt @@ -0,0 +1,113 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk.plugin + +import java.io.File +import java.nio.ByteBuffer +import java.security.MessageDigest +import java.util.zip.ZipException +import java.util.zip.ZipFile + +/** + * The `Airflow-Cache-Digest` a bundle JAR carries: a SHA-256 over what the + * bundle runs, which Airflow compares to tell whether a JAR changed since it + * last asked it for its task handlers. + * + * It hashes content, never archive metadata, so a rebuild that changes + * nothing gives the same value although the JAR's bytes differ (Gradle keeps + * file timestamps in archives by default). It covers the main class, the + * packaging mode, every file of the main source set's output, and every + * runtime dependency. A dependency JAR contributes its entries' names, sizes + * and CRC-32 values from the central directory, so no entry is decompressed. + */ +internal object BundleDigest { + const val ATTRIBUTE = "Airflow-Cache-Digest" + + fun compute( + mainClass: String?, + fatJar: Boolean, + outputs: Iterable, + classpath: Iterable, + ): String { + val writer = Writer() + writer.write("airflow-cache-digest", 1, mainClass.orEmpty(), if (fatJar) "fat" else "thin") + writer.writeTree(outputs) + for (entry in classpath) { + when { + entry.isDirectory -> writer.writeTree(listOf(entry)) + entry.isFile -> writer.writeArchive(entry) + } + } + return writer.hex() + } + + private class Writer { + private val digest = MessageDigest.getInstance("SHA-256") + + // Each token is length-prefixed, so no two token sequences hash alike. + fun write(vararg tokens: Any) { + for (token in tokens) { + val bytes = token.toString().toByteArray(Charsets.UTF_8) + digest.update(ByteBuffer.allocate(Int.SIZE_BYTES).putInt(bytes.size).array()) + digest.update(bytes) + } + } + + fun writeTree(roots: Iterable) { + roots + .filter { it.isDirectory } + .flatMap { root -> + root.walkTopDown().filter { it.isFile }.map { root.toPath().relativize(it.toPath()).joinToString("/") to it } + }.sortedBy { it.first } + .forEach { (path, file) -> write("file", path, sha256(file)) } + } + + fun writeArchive(file: File) { + try { + ZipFile(file).use { zip -> + write("archive", file.name) + zip + .entries() + .asSequence() + .filterNot { it.isDirectory } + .sortedBy { it.name } + .forEach { write("entry", it.name, it.size, it.crc) } + } + } catch (_: ZipException) { + write("file", file.name, sha256(file)) + } + } + + fun hex(): String = digest.digest().joinToString("") { "%02x".format(it) } + + private fun sha256(file: File): String { + val fileDigest = MessageDigest.getInstance("SHA-256") + file.inputStream().use { input -> + val buffer = ByteArray(DEFAULT_BUFFER_SIZE) + while (true) { + val read = input.read(buffer) + if (read < 0) break + fileDigest.update(buffer, 0, read) + } + } + return fileDigest.digest().joinToString("") { "%02x".format(it) } + } + } +} diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/BundleDigestTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/BundleDigestTest.kt new file mode 100644 index 0000000000000..b87dcdb0c6316 --- /dev/null +++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/BundleDigestTest.kt @@ -0,0 +1,121 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk.plugin + +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNotEquals +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.io.TempDir +import java.io.File +import java.util.zip.ZipEntry +import java.util.zip.ZipOutputStream + +class BundleDigestTest { + @TempDir + lateinit var tmp: File + + private fun classesDir( + name: String, + files: Map, + ): File = + tmp.resolve(name).also { dir -> + files.forEach { (path, content) -> dir.resolve(path).apply { parentFile.mkdirs() }.writeText(content) } + } + + private fun jar( + name: String, + entries: List>, + time: Long = 0L, + ): File = + tmp.resolve(name).also { file -> + file.parentFile.mkdirs() + ZipOutputStream(file.outputStream()).use { zip -> + entries.forEach { (path, content) -> + zip.putNextEntry(ZipEntry(path).apply { this.time = time }) + zip.write(content.toByteArray()) + zip.closeEntry() + } + } + } + + private val classes = mapOf("com/example/Main.class" to "main", "com/example/Handler.class" to "handler") + private val depEntries = listOf("dep/A.class" to "a", "dep/B.class" to "b") + + private fun digest( + mainClass: String = "com.example.Main", + fatJar: Boolean = true, + outputs: List = listOf(classesDir("classes", classes)), + classpath: List = listOf(jar("lib/dep.jar", depEntries)), + ) = BundleDigest.compute(mainClass, fatJar, outputs, classpath) + + @Test + @DisplayName("Should give the same SHA-256 for the same content") + fun sameContentSameDigest() { + val first = digest() + assertTrue(Regex("[0-9a-f]{64}").matches(first)) + assertEquals(first, digest()) + } + + @Test + @DisplayName("Should ignore dependency archive timestamps and entry order") + fun ignoresArchiveMetadata() { + val first = digest() + val rezipped = jar("lib/dep.jar", depEntries.reversed(), time = 1_700_000_000_000L) + + assertEquals(first, digest(classpath = listOf(rezipped))) + } + + @Test + @DisplayName("Should change when anything the bundle runs changes") + fun changesWithContent() { + val first = digest() + + assertNotEquals(first, digest(outputs = listOf(classesDir("changed", classes + ("com/example/Handler.class" to "HANDLER"))))) + assertNotEquals(first, digest(classpath = listOf(jar("lib/dep.jar", depEntries), jar("lib/extra.jar", depEntries)))) + assertNotEquals(first, digest(classpath = listOf(jar("lib/dep.jar", listOf("dep/A.class" to "a", "dep/B.class" to "bb"))))) + assertNotEquals(first, digest(mainClass = "com.example.Other")) + assertNotEquals(first, digest(fatJar = false)) + } + + @Test + @DisplayName("Should change when a classpath directory's content changes") + fun changesWithClasspathDirectoryContent() { + val included = classesDir("included", classes) + val first = digest(classpath = listOf(included)) + + included.resolve("com/example/Handler.class").writeText("HANDLER") + + assertNotEquals(first, digest(classpath = listOf(included))) + } + + @Test + @DisplayName("Should change when a non-zip classpath file's content changes") + fun changesWithNonZipClasspathFileContent() { + val native = tmp.resolve("lib/native.so").apply { parentFile.mkdirs() } + native.writeText("native") + val first = digest(classpath = listOf(native)) + + native.writeText("NATIVE") + + assertNotEquals(first, digest(classpath = listOf(native))) + } +} diff --git a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt index 049aba0d0601a..ee6407380fc55 100644 --- a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt +++ b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt @@ -23,6 +23,7 @@ package org.apache.airflow.sdk import com.squareup.javapoet.ClassName import com.squareup.javapoet.CodeBlock +import com.squareup.javapoet.FieldSpec import com.squareup.javapoet.JavaFile import com.squareup.javapoet.MethodSpec import com.squareup.javapoet.ParameterizedTypeName @@ -33,6 +34,7 @@ import org.apache.airflow.sdk.internal.Field import org.apache.airflow.sdk.internal.FieldType import org.apache.airflow.sdk.internal.SchemaFields import org.apache.airflow.sdk.internal.TaskArgs +import org.apache.airflow.sdk.internal.TaskParams import org.apache.airflow.sdk.internal.TypeRef import org.apache.airflow.sdk.internal.foldArgName import org.apache.airflow.sdk.internal.registrarName @@ -129,7 +131,8 @@ class BuilderProcessor : AbstractProcessor() { /** * Generates the registrar for a class of [Builder.TaskHandler] methods: one * [Task] implementation per handler, and a `registerInto` that binds each to - * the Dag and task the annotation names. + * the Dag and task the annotation names. Each implementation carries its + * handler's [TaskParams], so the runtime can report how it binds. * * There is no Dag to build here — the Python Dag file owns it — so this is a * registrar rather than a builder. @@ -162,7 +165,12 @@ class BuilderProcessor : AbstractProcessor() { "@Builder.TaskHandler on '${inner.simpleName}' must name the Dag the Python file declares" } val innerName = inner.simpleName.toString().replaceFirstChar(Char::uppercase) - registrar.addType(buildTask(innerName, inner, el)) + registrar.addType( + buildTask(innerName, inner, el) + .toBuilder() + .addField(taskParamsField(collectDataParams(inner))) + .build(), + ) registerInto.addStatement( $$"bundle.register($S, $S, $L.class)", handler.dag, @@ -415,6 +423,26 @@ class BuilderProcessor : AbstractProcessor() { return params } + /** + * Emits the [TaskParams.FIELD] field recording a handler's data parameters: + * its [TaskInput] type, or each flat parameter's name and declared type. + */ + private fun taskParamsField(dataParams: List): FieldSpec { + val input = dataParams.singleOrNull { it.isTaskInput } + val value = + if (input != null) { + CodeBlock.of($$"$T.input($T.class)", TASK_PARAMS_TYPE, TypeName.get(input.type)) + } else { + val params = + dataParams.map { CodeBlock.of($$"$T.param($S, $L)", TASK_PARAMS_TYPE, it.name, typeToken(TypeName.get(it.type))) } + CodeBlock.of($$"$T.of($L)", TASK_PARAMS_TYPE, CodeBlock.join(params, ", ")) + } + return FieldSpec + .builder(TASK_PARAMS_TYPE, TaskParams.FIELD, Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL) + .initializer(value) + .build() + } + private fun ProcessingEnvironment.isTaskInput(type: TypeMirror): Boolean { val marker = elementUtils.getTypeElement(TASK_INPUT_TYPE.canonicalName()) ?: return false return !type.kind.isPrimitive && typeUtils.isAssignable(type, marker.asType()) @@ -499,6 +527,7 @@ private val CLIENT_TYPE = ClassName.get(Client::class.java) private val CONTEXT_TYPE = ClassName.get(Context::class.java) private val TASK_INPUT_TYPE = ClassName.get(TaskInput::class.java) private val TASK_ARGS_TYPE = ClassName.get(TaskArgs::class.java) +private val TASK_PARAMS_TYPE = ClassName.get(TaskParams::class.java) private val TYPE_REF_TYPE = ClassName.get(TypeRef::class.java) private val ARG_VALUES_TYPE = ClassName.get(ArgValues::class.java) @@ -530,11 +559,16 @@ private fun positionalAccess( ): CodeBlock { val type = TypeName.get(param.type) val reader = if (type.isPrimitive) "require" else "get" - val target = - if (type is ParameterizedTypeName) { - CodeBlock.of($$"new $T<$T>() {}", TYPE_REF_TYPE, type) - } else { - CodeBlock.of($$"$T.class", type.box()) - } - return CodeBlock.of($$"$L.$L($L, $L)", argsLocal, reader, param.position, target) + return CodeBlock.of($$"$L.$L($L, $L)", argsLocal, reader, param.position, typeToken(type.box())) } + +/** + * Emits a runtime token for [type]: a [TypeRef] when it has type arguments, + * which a `Class` literal cannot carry, otherwise its `Class` literal. + */ +private fun typeToken(type: TypeName): CodeBlock = + if (type is ParameterizedTypeName) { + CodeBlock.of($$"new $T<$T>() {}", TYPE_REF_TYPE, type) + } else { + CodeBlock.of($$"$T.class", type) + } diff --git a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt index 6fd52e01b75a4..09d0b09eadad5 100644 --- a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt +++ b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt @@ -858,6 +858,7 @@ class BuilderTest { import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.internal.TaskArgs; + import org.apache.airflow.sdk.internal.TaskParams; /** * Registers {@link TestExample}'s task handlers against the Dags the Python file owns. @@ -869,6 +870,8 @@ class BuilderTest { } public static final class Score implements Task { + public static final TaskParams AIRFLOW_TASK_PARAMS = TaskParams.of(TaskParams.param("rows", long.class)); + @Override public void execute(Context context, Client client) throws Exception { TaskArgs args = TaskArgs.of(context, client, 1); @@ -878,6 +881,8 @@ class BuilderTest { } public static final class Audit implements Task { + public static final TaskParams AIRFLOW_TASK_PARAMS = TaskParams.of(); + @Override public void execute(Context context, Client client) throws Exception { new TestExample().audit(); @@ -888,6 +893,107 @@ class BuilderTest { ) } + @Test + @DisplayName("record each handler's data parameters for the runtime to declare") + fun generateHandlerTaskParams() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import java.util.List; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.Context; + import org.apache.airflow.sdk.TaskInput; + public class TestExample { + public static class ScoreInput implements TaskInput { + public String region; + } + + @Builder.TaskHandler(dag = "etl") + public void flat(Client client, long rows, Integer limit, List regions) {} + + @Builder.TaskHandler(dag = "etl") + public void named(ScoreInput input) {} + + @Builder.TaskHandler(dag = "etl") + public void none(Context context) {} + } + """, + ) + + assertThat(compilation).succeeded() + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleHandlers") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleHandlers", + """ + package org.apache.airflow.example; + + import java.lang.Exception; + import java.lang.Integer; + import java.lang.Long; + import java.lang.Override; + import java.lang.String; + import java.util.List; + import org.apache.airflow.sdk.Bundle; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.Context; + import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.internal.ArgValues; + import org.apache.airflow.sdk.internal.TaskArgs; + import org.apache.airflow.sdk.internal.TaskParams; + import org.apache.airflow.sdk.internal.TypeRef; + + /** + * Registers {@link TestExample}'s task handlers against the Dags the Python file owns. + */ + public final class TestExampleHandlers { + public static void registerInto(Bundle bundle) { + bundle.register("etl", "flat", Flat.class); + bundle.register("etl", "named", Named.class); + bundle.register("etl", "none", None.class); + } + + public static final class Flat implements Task { + public static final TaskParams AIRFLOW_TASK_PARAMS = TaskParams.of( + TaskParams.param("rows", long.class), + TaskParams.param("limit", Integer.class), + TaskParams.param("regions", new TypeRef>() {})); + + @Override + public void execute(Context context, Client client) throws Exception { + TaskArgs args = TaskArgs.of(context, client, 3); + long rows = args.require(0, Long.class); + Integer limit = args.get(1, Integer.class); + List regions = args.get(2, new TypeRef>() {}); + new TestExample().flat(client, rows, limit, regions); + } + } + + public static final class Named implements Task { + public static final TaskParams AIRFLOW_TASK_PARAMS = TaskParams.input(TestExample.ScoreInput.class); + + @Override + public void execute(Context context, Client client) throws Exception { + TestExample.ScoreInput input = ArgValues.bindInput(client, TestExample.ScoreInput.class); + new TestExample().named(input); + } + } + + public static final class None implements Task { + public static final TaskParams AIRFLOW_TASK_PARAMS = TaskParams.of(); + + @Override + public void execute(Context context, Client client) throws Exception { + new TestExample().none(context); + } + } + } + """, + ) + } + @Test @DisplayName("reject a handler that names no Dag") fun rejectHandlerWithoutDag() { @@ -941,6 +1047,7 @@ class BuilderTest { import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.internal.TaskArgs; + import org.apache.airflow.sdk.internal.TaskParams; /** * Registers {@link TestExample.Inner}'s task handlers against the Dags the Python file owns. @@ -951,6 +1058,8 @@ class BuilderTest { } public static final class Score implements Task { + public static final TaskParams AIRFLOW_TASK_PARAMS = TaskParams.of(TaskParams.param("rows", long.class)); + @Override public void execute(Context context, Client client) throws Exception { TaskArgs args = TaskArgs.of(context, client, 1); diff --git a/java-sdk/scala_spark_example/README.md b/java-sdk/scala_spark_example/README.md index f59c01e0e52c6..34f248a017cf1 100644 --- a/java-sdk/scala_spark_example/README.md +++ b/java-sdk/scala_spark_example/README.md @@ -41,7 +41,9 @@ cd scala_spark_example ``` `fatJar` is disabled, so `build/bundle/` holds the bundle JAR plus every runtime -JAR (Spark included) — copy it into a Java coordinator's `jars_root`. +JAR (Spark included). Copy it into the Dag bundle named by a Java coordinator's +`task_handler_bundle_name`. Every JAR in that Dag bundle goes on one classpath, +so give Spark its own Dag bundle, coordinator and queue. ## Running Spark under the Java SDK @@ -57,7 +59,7 @@ and bypasses that launcher, so the coordinator has to pass them itself via "scala-jdk": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": { - "jars_root": ["/path/to/scala-jars"], + "task_handler_bundle_name": "scala-task-handlers", "main_class": "org.apache.airflow.example.ScalaSparkBundleBuilder", "jvm_args": [ "-Xmx512m", diff --git a/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala b/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala index f7f32bad8e296..8c6b3cf28b277 100644 --- a/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala +++ b/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala @@ -19,7 +19,7 @@ package org.apache.airflow.example -import org.apache.airflow.sdk.{Bundle, BundleBuilder, Client, Context, DagDef, Server, Task} +import org.apache.airflow.sdk.{Bundle, Client, Context, Server, Task} import org.apache.logging.log4j.{LogManager, Logger} import org.apache.spark.sql.SparkSession import org.apache.spark.sql.functions.sum @@ -132,18 +132,15 @@ class SparkLoad extends Task { } } -object ScalaSparkExample { - def build(): DagDef = - new DagDef(SparkEtl.DagId) - .addTask(SparkEtl.ExtractTaskId, classOf[SparkExtract]) - .addTask(SparkEtl.TransformTaskId, classOf[SparkTransform]) - .addTask(SparkEtl.LoadTaskId, classOf[SparkLoad]) -} - /** Bundle entry point served to Airflow's Java coordinator. */ -object ScalaSparkBundleBuilder extends BundleBuilder { - override def getDags(): java.lang.Iterable[DagDef] = java.util.List.of(ScalaSparkExample.build()) +object ScalaSparkBundleBuilder { + // The Python Dag file owns the Dag, so its tasks register as task handlers. + def build(): Bundle = + new Bundle() + .register(SparkEtl.DagId, SparkEtl.ExtractTaskId, classOf[SparkExtract]) + .register(SparkEtl.DagId, SparkEtl.TransformTaskId, classOf[SparkTransform]) + .register(SparkEtl.DagId, SparkEtl.LoadTaskId, classOf[SparkLoad]) def main(args: Array[String]): Unit = - Server.create(args).serve(new Bundle(getDags())) + Server.create(args).serve(build()) } diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts index 5e9a74a704638..614398800d4bb 100644 --- a/java-sdk/sdk/build.gradle.kts +++ b/java-sdk/sdk/build.gradle.kts @@ -161,6 +161,10 @@ abstract class GenerateDiscriminatorTask : DefaultTask() { } } +// A version is published only when it is released, so the in-progress one +// (which gains messages during a release cycle) is taken from the monorepo +// snapshot whenever that snapshot declares the configured version. Standalone +// (source-release) builds have no snapshot and download the published file. abstract class SyncSupervisorSchemaTask : DefaultTask() { @get:Input abstract val schemaVersion: Property @@ -171,6 +175,9 @@ abstract class SyncSupervisorSchemaTask : DefaultTask() { @get:Internal abstract val schemaFile: RegularFileProperty + @get:Internal + abstract val monorepoFile: RegularFileProperty + private fun apiVersionOf(file: File): String = if (file.exists()) { com.fasterxml.jackson.databind @@ -186,6 +193,16 @@ abstract class SyncSupervisorSchemaTask : DefaultTask() { fun sync() { val file = schemaFile.get().asFile val version = schemaVersion.get() + val monorepo = monorepoFile.get().asFile + if (apiVersionOf(monorepo) == version) { + if (file.exists() && file.readBytes().contentEquals(monorepo.readBytes())) { + logger.lifecycle("Supervisor Schema matches the monorepo snapshot (api_version=$version).") + } else { + logger.lifecycle("Refreshing Supervisor Schema from ${monorepo.path}") + monorepo.copyTo(file, overwrite = true) + } + return + } if (apiVersionOf(file) == version) { logger.lifecycle("Supervisor Schema is up-to-date (api_version=$version).") return @@ -640,6 +657,7 @@ val syncSupervisorSchema by tasks.registering(SyncSupervisorSchemaTask::class) { schemaVersion = airflowSupervisorSchemaVersion baseUrl = schemaBaseUrl schemaFile = schemaInput + monorepoFile = layout.projectDirectory.file("../../task-sdk/src/airflow/sdk/execution_time/schema/schema.json") } tasks.register("generateDiscriminator") { diff --git a/java-sdk/sdk/schema/schema.json b/java-sdk/sdk/schema/schema.json index b671959c50a00..e7a512de7b574 100644 --- a/java-sdk/sdk/schema/schema.json +++ b/java-sdk/sdk/schema/schema.json @@ -34,7 +34,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -187,34 +194,6 @@ "title": "AssetProfile", "type": "object" }, - "AssetReferenceAssetEventDagRun": { - "additionalProperties": false, - "description": "Schema for AssetModel used in AssetEventDagRunReference.", - "properties": { - "name": { - "title": "Name", - "type": "string" - }, - "uri": { - "title": "Uri", - "type": "string" - }, - "extra": { - "additionalProperties": { - "$ref": "#/$defs/JsonValue" - }, - "title": "Extra", - "type": "object" - } - }, - "required": [ - "name", - "uri", - "extra" - ], - "title": "AssetReferenceAssetEventDagRun", - "type": "object" - }, "AssetResponse": { "description": "Asset schema for responses with fields that are needed for Runtime.", "properties": { @@ -234,7 +213,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -273,7 +259,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -304,7 +297,14 @@ "description": "Response to GetAssetStateStore; wraps the generated API response for supervisor to worker comms.", "properties": { "value": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": { "const": "AssetStateStoreResult", @@ -347,17 +347,10 @@ "description": "Park a task instance awaiting human input (Human-in-the-loop), without a trigger.", "properties": { "state": { - "anyOf": [ - { - "const": "awaiting_input", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "awaiting_input", "default": "awaiting_input", - "title": "State" + "title": "State", + "type": "string" }, "timeout": { "anyOf": [ @@ -380,7 +373,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -513,6 +513,97 @@ "title": "ClearTaskStateStore", "type": "object" }, + "ConnectionResponse": { + "description": "Connection schema for responses with fields that are needed for Runtime.", + "properties": { + "conn_id": { + "title": "Conn Id", + "type": "string" + }, + "conn_type": { + "title": "Conn Type", + "type": "string" + }, + "host": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Host" + }, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Schema" + }, + "login": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Login" + }, + "password": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Password" + }, + "port": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Port" + }, + "extra": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Extra" + } + }, + "required": [ + "conn_id", + "conn_type", + "host", + "schema", + "login", + "password", + "port", + "extra" + ], + "title": "ConnectionResponse", + "type": "object" + }, "ConnectionResult": { "properties": { "conn_id": { @@ -532,7 +623,6 @@ "type": "null" } ], - "default": null, "title": "Host" }, "schema": { @@ -544,7 +634,6 @@ "type": "null" } ], - "default": null, "title": "Schema" }, "login": { @@ -556,7 +645,6 @@ "type": "null" } ], - "default": null, "title": "Login" }, "password": { @@ -568,7 +656,6 @@ "type": "null" } ], - "default": null, "title": "Password" }, "port": { @@ -580,7 +667,6 @@ "type": "null" } ], - "default": null, "title": "Port" }, "extra": { @@ -592,7 +678,6 @@ "type": "null" } ], - "default": null, "title": "Extra" }, "type": { @@ -604,7 +689,13 @@ }, "required": [ "conn_id", - "conn_type" + "conn_type", + "host", + "schema", + "login", + "password", + "port", + "extra" ], "title": "ConnectionResult", "type": "object" @@ -867,6 +958,13 @@ "title": "Callback Requests", "type": "array" }, + "known_artifacts": { + "items": { + "$ref": "#/$defs/TaskHandlerArtifact" + }, + "title": "Known Artifacts", + "type": "array" + }, "type": { "const": "DagFileParseRequest", "default": "DagFileParseRequest", @@ -924,6 +1022,20 @@ "default": null, "title": "Import Errors" }, + "parsed_definitions": { + "items": { + "type": "string" + }, + "title": "Parsed Definitions", + "type": "array" + }, + "dag_source_codes": { + "additionalProperties": { + "$ref": "#/$defs/DagSourceCode" + }, + "title": "Dag Source Codes", + "type": "object" + }, "type": { "const": "DagFileParsingResult", "default": "DagFileParsingResult", @@ -957,7 +1069,6 @@ "type": "null" } ], - "default": null, "title": "Bundle Name" }, "bundle_version": { @@ -969,7 +1080,6 @@ "type": "null" } ], - "default": null, "title": "Bundle Version" }, "relative_fileloc": { @@ -981,7 +1091,6 @@ "type": "null" } ], - "default": null, "title": "Relative Fileloc" }, "owners": { @@ -993,7 +1102,6 @@ "type": "null" } ], - "default": null, "title": "Owners" }, "tags": { @@ -1013,7 +1121,6 @@ "type": "null" } ], - "default": null, "title": "Next Dagrun" }, "type": { @@ -1026,7 +1133,12 @@ "required": [ "dag_id", "is_paused", - "tags" + "bundle_name", + "bundle_version", + "relative_fileloc", + "owners", + "tags", + "next_dagrun" ], "title": "DagResult", "type": "object" @@ -1053,13 +1165,19 @@ "type": "null" } ], - "default": null, "title": "Logical Date" }, "start_date": { - "format": "date-time", - "title": "Start Date", - "type": "string" + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Start Date" }, "end_date": { "anyOf": [ @@ -1071,7 +1189,6 @@ "type": "null" } ], - "default": null, "title": "End Date" }, "state": { @@ -1088,7 +1205,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval Start" }, "data_interval_end": { @@ -1101,7 +1217,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval End" }, "partition_key": { @@ -1113,15 +1228,19 @@ "type": "null" } ], - "default": null, "title": "Partition Key" } }, "required": [ "run_id", "dag_id", + "logical_date", "start_date", - "state" + "end_date", + "state", + "data_interval_start", + "data_interval_end", + "partition_key" ], "title": "DagRunAssetReference", "type": "object" @@ -1176,7 +1295,6 @@ "type": "null" } ], - "default": null, "title": "Logical Date" }, "data_interval_start": { @@ -1189,7 +1307,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval Start" }, "data_interval_end": { @@ -1202,7 +1319,6 @@ "type": "null" } ], - "default": null, "title": "Data Interval End" }, "run_after": { @@ -1220,7 +1336,6 @@ "type": "null" } ], - "default": null, "title": "Start Date" }, "end_date": { @@ -1233,7 +1348,6 @@ "type": "null" } ], - "default": null, "title": "End Date" }, "clear_number": { @@ -1295,7 +1409,6 @@ "type": "null" } ], - "default": null, "title": "Partition Key" }, "partition_date": { @@ -1345,10 +1458,16 @@ "required": [ "dag_id", "run_id", + "logical_date", + "data_interval_start", + "data_interval_end", "run_after", + "start_date", + "end_date", "run_type", "state", - "consumed_asset_events" + "consumed_asset_events", + "partition_key" ], "title": "DagRunResult", "type": "object" @@ -1395,22 +1514,34 @@ "title": "DagRunType", "type": "string" }, + "DagSourceCode": { + "description": "Raw source code and its language identifier for a DAG definition.", + "properties": { + "source_code": { + "title": "Source Code", + "type": "string" + }, + "language": { + "title": "Language", + "type": "string" + } + }, + "required": [ + "source_code", + "language" + ], + "title": "DagSourceCode", + "type": "object" + }, "DeferTask": { "additionalProperties": false, "description": "Update a task instance state to deferred.", "properties": { "state": { - "anyOf": [ - { - "const": "deferred", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "deferred", "default": "deferred", - "title": "State" + "title": "State", + "type": "string" }, "classpath": { "title": "Classpath", @@ -1420,7 +1551,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -1467,7 +1605,14 @@ "anyOf": [ { "additionalProperties": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": "object" }, @@ -3045,7 +3190,17 @@ "title": "InactiveAssetsResult", "type": "object" }, - "JsonValue": {}, + "JsonValue": { + "title": "JsonValue", + "type": [ + "object", + "array", + "string", + "number", + "boolean", + "null" + ] + }, "LazyDeserializedDAG": { "description": "Lazily build information from the serialized DAG structure.\n\nAn object that will present \"enough\" of the DAG like interface to update DAG db models etc, without having\nto deserialize the full DAG and Task hierarchy.", "properties": { @@ -3388,17 +3543,10 @@ "description": "Update a task instance state to reschedule/up_for_reschedule.", "properties": { "state": { - "anyOf": [ - { - "const": "up_for_reschedule", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "up_for_reschedule", "default": "up_for_reschedule", - "title": "State" + "title": "State", + "type": "string" }, "reschedule_date": { "format": "date-time", @@ -3441,17 +3589,10 @@ "description": "Update a task instance state to up_for_retry.", "properties": { "state": { - "anyOf": [ - { - "const": "up_for_retry", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "up_for_retry", "default": "up_for_retry", - "title": "State" + "title": "State", + "type": "string" }, "end_date": { "format": "date-time", @@ -3825,17 +3966,10 @@ "description": "Update a task's state to success. Includes task_outlets and outlet_events for registering asset events.", "properties": { "state": { - "anyOf": [ - { - "const": "success", - "type": "string" - }, - { - "type": "null" - } - ], + "const": "success", "default": "success", - "title": "State" + "title": "State", + "type": "string" }, "end_date": { "format": "date-time", @@ -4029,6 +4163,220 @@ "title": "TaskCallbackRequest", "type": "object" }, + "TaskHandlerArtifact": { + "description": "A Lang-SDK artifact and every task handler it registers.", + "properties": { + "bundle_name": { + "title": "Bundle Name", + "type": "string" + }, + "relative_fileloc": { + "maxLength": 2000, + "title": "Relative Fileloc", + "type": "string" + }, + "size_bytes": { + "title": "Size Bytes", + "type": "integer" + }, + "cache_digest": { + "anyOf": [ + { + "maxLength": 128, + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Cache Digest" + }, + "task_handlers": { + "additionalProperties": { + "items": { + "$ref": "#/$defs/TaskHandlerDeclaration" + }, + "type": "array" + }, + "title": "Task Handlers", + "type": "object" + } + }, + "required": [ + "bundle_name", + "relative_fileloc", + "size_bytes", + "cache_digest", + "task_handlers" + ], + "title": "TaskHandlerArtifact", + "type": "object" + }, + "TaskHandlerDeclaration": { + "description": "A task handler that a Lang-SDK artifact registers for one task.", + "properties": { + "task_id": { + "title": "Task Id", + "type": "string" + }, + "binding": { + "enum": [ + "positional", + "named" + ], + "title": "Binding", + "type": "string" + }, + "params": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/TaskHandlerParam" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Task Handler Params" + } + }, + "required": [ + "task_id", + "binding", + "params" + ], + "title": "TaskHandlerDeclaration", + "type": "object" + }, + "TaskHandlerParam": { + "description": "One parameter of a task handler.", + "properties": { + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Name" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "exact_name": { + "default": false, + "title": "Exact Name", + "type": "boolean" + } + }, + "required": [ + "name" + ], + "title": "TaskHandlerParam", + "type": "object" + }, + "TaskHandlerParseRequest": { + "description": "Request for Task Handler Parsing.\n\nAsks a Lang-SDK runtime for every task handler an artifact registers.", + "properties": { + "file": { + "title": "File", + "type": "string" + }, + "bundle_path": { + "format": "path", + "title": "Bundle Path", + "type": "string" + }, + "bundle_name": { + "title": "Bundle Name", + "type": "string" + }, + "type": { + "const": "TaskHandlerParseRequest", + "default": "TaskHandlerParseRequest", + "title": "Type", + "type": "string" + } + }, + "required": [ + "file", + "bundle_path", + "bundle_name" + ], + "title": "TaskHandlerParseRequest", + "type": "object" + }, + "TaskHandlerParsingResult": { + "description": "Result of Task Handler Parsing.\n\nEvery task handler a Lang-SDK artifact registers, keyed by Dag id.\n\nThe answer depends only on the artifact, never on the request.", + "properties": { + "fileloc": { + "title": "Fileloc", + "type": "string" + }, + "task_handlers": { + "additionalProperties": { + "items": { + "$ref": "#/$defs/TaskHandlerDeclaration" + }, + "type": "array" + }, + "title": "Task Handlers", + "type": "object" + }, + "import_errors": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Import Errors" + }, + "warnings": { + "anyOf": [ + { + "items": {}, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Warnings" + }, + "type": { + "const": "TaskHandlerParsingResult", + "default": "TaskHandlerParsingResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "fileloc", + "task_handlers" + ], + "title": "TaskHandlerParsingResult", + "type": "object" + }, "TaskInstanceState": { "description": "All possible states that a Task Instance can be in.\n\nNote that None is also allowed, so always use this in a type hint with Optional.", "enum": [ @@ -4119,6 +4467,18 @@ ], "default": null, "title": "Rendered Map Index" + }, + "retry_reason": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Retry Reason" } }, "required": [ @@ -4132,7 +4492,14 @@ "description": "Response to GetTaskStateStore; wraps the generated API response for supervisor to worker comms.", "properties": { "value": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": { "const": "TaskStateStoreResult", @@ -4267,6 +4634,38 @@ "title": "TriggerDagRun", "type": "object" }, + "UpdateDagRunNote": { + "properties": { + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "note": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Note" + }, + "type": { + "const": "UpdateDagRunNote", + "default": "UpdateDagRunNote", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti_id", + "note" + ], + "title": "UpdateDagRunNote", + "type": "object" + }, "UpdateHITLDetail": { "description": "Update the response content part of an existing Human-in-the-loop response.", "properties": { @@ -4357,6 +4756,33 @@ "title": "VariableKeysResult", "type": "object" }, + "VariableResponse": { + "additionalProperties": false, + "description": "Variable schema for responses with fields that are needed for Runtime.", + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "value": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Value" + } + }, + "required": [ + "key", + "value" + ], + "title": "VariableResponse", + "type": "object" + }, "VariableResult": { "additionalProperties": false, "properties": { @@ -4373,7 +4799,6 @@ "type": "null" } ], - "default": null, "title": "Value" }, "type": { @@ -4384,7 +4809,8 @@ } }, "required": [ - "key" + "key", + "value" ], "title": "VariableResult", "type": "object" @@ -4416,7 +4842,14 @@ "type": "string" }, "value": { - "$ref": "#/$defs/JsonValue" + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ] }, "type": { "const": "XComResult", @@ -4472,97 +4905,6 @@ "title": "XComSequenceSliceResult", "type": "object" }, - "ConnectionResponse": { - "description": "Connection schema for responses with fields that are needed for Runtime.", - "properties": { - "conn_id": { - "title": "Conn Id", - "type": "string" - }, - "conn_type": { - "title": "Conn Type", - "type": "string" - }, - "host": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Host" - }, - "schema": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Schema" - }, - "login": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Login" - }, - "password": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Password" - }, - "port": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "title": "Port" - }, - "extra": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Extra" - } - }, - "required": [ - "conn_id", - "conn_type", - "host", - "schema", - "login", - "password", - "port", - "extra" - ], - "title": "ConnectionResponse", - "type": "object" - }, "ArgValueSchema": { "additionalProperties": { "$ref": "#/$defs/JsonValue" @@ -4573,11 +4915,6 @@ "LiteralArgBinding": { "description": "One positional stub-task argument carrying an inline literal from the Dag file.", "properties": { - "kind": { - "const": "literal", - "title": "Kind", - "type": "string" - }, "name": { "title": "Name", "type": "string" @@ -4593,6 +4930,11 @@ ], "default": null }, + "kind": { + "const": "literal", + "title": "Kind", + "type": "string" + }, "value": { "anyOf": [ { @@ -4611,8 +4953,8 @@ } }, "required": [ - "kind", - "name" + "name", + "kind" ], "title": "LiteralArgBinding", "type": "object" @@ -4638,11 +4980,6 @@ "XComArgBinding": { "description": "One positional stub-task argument pulled from an upstream task's XCom.", "properties": { - "kind": { - "const": "xcom", - "title": "Kind", - "type": "string" - }, "name": { "title": "Name", "type": "string" @@ -4658,31 +4995,19 @@ ], "default": null }, + "kind": { + "const": "xcom", + "title": "Kind", + "type": "string" + }, "task_id": { "title": "Task Id", "type": "string" - }, - "map_index": { - "default": -1, - "title": "Map Index", - "type": "integer" - }, - "element_index": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Element Index" } }, "required": [ - "kind", "name", + "kind", "task_id" ], "title": "XComArgBinding", @@ -4784,6 +5109,34 @@ "title": "AssetEventDagRunReference", "type": "object" }, + "AssetReferenceAssetEventDagRun": { + "additionalProperties": false, + "description": "Schema for AssetModel used in AssetEventDagRunReference.", + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + }, + "extra": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Extra", + "type": "object" + } + }, + "required": [ + "name", + "uri", + "extra" + ], + "title": "AssetReferenceAssetEventDagRun", + "type": "object" + }, "DagRun": { "additionalProperties": false, "description": "Schema for DagRun model with minimal required fields needed for Runtime.", @@ -5066,6 +5419,11 @@ ], "default": null, "title": "Arg Bindings" + }, + "multi_team": { + "default": false, + "title": "Multi Team", + "type": "boolean" } }, "required": [ @@ -5150,33 +5508,6 @@ ], "title": "TaskInstance", "type": "object" - }, - "VariableResponse": { - "additionalProperties": false, - "description": "Variable schema for responses with fields that are needed for Runtime.", - "properties": { - "key": { - "title": "Key", - "type": "string" - }, - "value": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Value" - } - }, - "required": [ - "key", - "value" - ], - "title": "VariableResponse", - "type": "object" } } } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt index 4b42eb68d4c39..e039f4e8671c6 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt @@ -35,6 +35,8 @@ import org.apache.airflow.sdk.execution.LogSender import org.apache.airflow.sdk.execution.Logger import org.apache.airflow.sdk.execution.comm.ErrorResponse import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.execution.comm.TaskHandlerParseRequest +import org.apache.airflow.sdk.execution.parseTaskHandlers import org.apache.airflow.sdk.execution.runTask import kotlin.text.substringAfterLast import kotlin.text.substringBeforeLast @@ -70,6 +72,10 @@ class ApiError( * Connects this JVM process to the Airflow coordinator and dispatches task * execution requests to the registered [Bundle]. * + * The Dag processor starts the same entry point to ask which task handlers + * the bundle registers. That process replies with every registered handler + * and how each binds its arguments, runs no task, and exits. + * * The typical entry point is: * * ```java @@ -181,6 +187,7 @@ class Server( val frame = coordinator.readMessage() when (val body = frame.body) { is StartupDetails -> runTaskAndReport(bundle, body, coordinator) + is TaskHandlerParseRequest -> parseTaskHandlersAndReport(bundle, body, coordinator) is ErrorResponse -> throw ApiError("[${body.error}] ${body.detail}") else -> throw ApiError("Unexpected initial frame (id=${frame.id})") } @@ -194,4 +201,12 @@ class Server( val result = runTask(bundle, startup, coordinator) coordinator.communicate(result) } + + private suspend fun parseTaskHandlersAndReport( + bundle: Bundle, + request: TaskHandlerParseRequest, + coordinator: CoordinatorComm, + ) { + coordinator.communicate(parseTaskHandlers(bundle, request)) + } } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt new file mode 100644 index 0000000000000..12ac86159d071 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParse.kt @@ -0,0 +1,104 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk.execution + +import org.apache.airflow.sdk.Bundle +import org.apache.airflow.sdk.InputTask +import org.apache.airflow.sdk.Task +import org.apache.airflow.sdk.TaskDef +import org.apache.airflow.sdk.TaskInput +import org.apache.airflow.sdk.execution.comm.TaskHandlerDeclaration.Binding +import org.apache.airflow.sdk.execution.comm.TaskHandlerParseRequest +import org.apache.airflow.sdk.internal.TaskParams +import org.apache.airflow.sdk.internal.argNameOf +import org.apache.airflow.sdk.internal.bindableFields +import org.apache.airflow.sdk.internal.buildValueSchema +import org.apache.airflow.sdk.internal.isPinned +import org.apache.airflow.sdk.internal.resolveInputType +import java.lang.reflect.Modifier +import java.lang.reflect.Type + +/** + * Answers a [TaskHandlerParseRequest] with every task handler [bundle] + * registers, keyed by Dag id in registration order, as a + * TaskHandlerParsingResult body. + * + * Only task-handler registrations count: a Dag declared in Java is its own, + * not a Python Dag's. The generated models cannot hold `task_handlers`, so + * the body is a map. + */ +internal fun parseTaskHandlers( + bundle: Bundle, + request: TaskHandlerParseRequest, +): Map = + linkedMapOf( + "type" to "TaskHandlerParsingResult", + "fileloc" to request.file, + "task_handlers" to bundle.taskHandlers.mapValues { (_, dag) -> dag.tasks.values.map(::declare) }, + ) + +private fun declare(task: TaskDef): Map { + val (binding, params) = bindingOf(task.definition) + return linkedMapOf( + "task_id" to task.id, + "binding" to binding.value(), + "params" to params, + ) +} + +/** + * How [definition] binds the stub call's arguments. Flat parameters bind by + * position; a [TaskInput] binds by name. A task that reads no argument is + * declared as binding by name with no parameters: it ignores an argument + * passed to it, so the Dag processor only warns about one. + */ +private fun bindingOf(definition: Class): Pair>> { + if (InputTask::class.java.isAssignableFrom(definition)) { + return Binding.NAMED to inputParams(resolveInputType(definition)) + } + val declared = taskParamsOf(definition) ?: return Binding.NAMED to emptyList() + declared.input?.let { return Binding.NAMED to inputParams(it) } + if (declared.flat.isEmpty()) return Binding.NAMED to emptyList() + return Binding.POSITIONAL to declared.flat.map { param(it.name, it.type, exactName = false) } +} + +/** The [TaskParams] the annotation processor recorded on a generated task class, if [definition] is one. */ +private fun taskParamsOf(definition: Class<*>): TaskParams? { + val field = + try { + definition.getField(TaskParams.FIELD) + } catch (_: NoSuchFieldException) { + return null + } + return if (Modifier.isStatic(field.modifiers)) field.get(null) as? TaskParams else null +} + +private fun inputParams(inputType: Class): List> = + bindableFields(inputType).map { param(argNameOf(it), it.genericType, exactName = isPinned(it)) } + +private fun param( + name: String, + type: Type, + exactName: Boolean, +): Map = + linkedMapOf( + "name" to name, + "exact_name" to exactName, + ).apply { buildValueSchema(type)?.let { put("value_schema", it) } } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/TaskParams.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/TaskParams.kt new file mode 100644 index 0000000000000..7181b13a87e38 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/TaskParams.kt @@ -0,0 +1,72 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk.internal + +import org.apache.airflow.sdk.TaskInput +import java.lang.reflect.Type + +/** + * @suppress + * + * The data parameters of a task class generated for a `@Builder.TaskHandler` + * method, as the runtime reports them when asked which task handlers a bundle + * registers. The generated class holds it in a static field named [FIELD]. + * + * Java parameter names do not survive compilation, so the annotation + * processor records them here. Public so that generated code can build it; + * not user-facing API. + */ +class TaskParams private constructor( + internal val flat: List, + internal val input: Class?, +) { + /** One flat data parameter: its name in the handler method, and its declared type. */ + class Param internal constructor( + internal val name: String, + internal val type: Type, + ) + + companion object { + /** Name of the static field holding a generated task class's [TaskParams]. */ + const val FIELD = "AIRFLOW_TASK_PARAMS" + + /** Flat data parameters, in declaration order; none for a method that takes no data. */ + @JvmStatic + fun of(vararg params: Param): TaskParams = TaskParams(params.toList(), null) + + /** A single [TaskInput] parameter of [type]. */ + @JvmStatic + fun input(type: Class): TaskParams = TaskParams(emptyList(), type) + + /** A flat parameter whose declared type a `Class` literal expresses; primitives stay primitive. */ + @JvmStatic + fun param( + name: String, + type: Class<*>, + ): Param = Param(name, type) + + /** A flat parameter whose declared type has type arguments. */ + @JvmStatic + fun param( + name: String, + type: TypeRef<*>, + ): Param = Param(name, type.type) + } +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ValueSchema.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ValueSchema.kt new file mode 100644 index 0000000000000..1ffd087998271 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ValueSchema.kt @@ -0,0 +1,160 @@ +/* + * 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. + */ + +@file:Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN") + +package org.apache.airflow.sdk.internal + +import java.lang.reflect.GenericArrayType +import java.lang.reflect.ParameterizedType +import java.lang.reflect.Type +import java.lang.reflect.WildcardType +import java.math.BigDecimal +import java.math.BigInteger +import java.time.Duration +import java.time.Instant +import java.time.LocalDate +import java.time.LocalDateTime +import java.time.LocalTime +import java.time.OffsetDateTime +import java.time.ZonedDateTime +import java.util.UUID + +private typealias Schema = Map + +private val UNCONSTRAINED: Schema = emptyMap() + +/** + * Builds the JSON Schema of the values a parameter or field of [type] + * accepts, in the vocabulary pydantic emits for a Python annotation, or null + * when [type] does not say. + * + * The schema describes the declared type, not every value the decoder + * coerces into it: [ArgValues] also turns `"5"` into a `long`, as pydantic's + * lax mode does for `int`. Only a primitive excludes null, which the decoder + * does enforce. + */ +internal fun buildValueSchema(type: Type): Schema? { + val schema = schemaOf(type) ?: return null + return if (type is Class<*> && type.isPrimitive) schema else nullable(schema) +} + +private fun nullable(schema: Schema): Schema = mapOf("anyOf" to listOf(schema, mapOf("type" to "null"))) + +private fun schemaOf(type: Type): Schema? = + when (type) { + is Class<*> -> classSchema(type) + is ParameterizedType -> parameterizedSchema(type) + is GenericArrayType -> arraySchema(type.genericComponentType) + is WildcardType -> type.upperBounds.firstOrNull()?.let(::schemaOf) + else -> null // A type variable says nothing about its values. + } + +private fun integer(format: String): Schema = mapOf("type" to "integer", "format" to format) + +private fun integerRange( + minimum: Long, + maximum: Long, +): Schema = mapOf("type" to "integer", "minimum" to minimum, "maximum" to maximum) + +private fun number(format: String): Schema = mapOf("type" to "number", "format" to format) + +private fun string(format: String): Schema = mapOf("type" to "string", "format" to format) + +private val SCALARS: Map, Schema> = + buildMap { + fun put( + primitive: Class<*>?, + boxed: Class<*>, + schema: Schema, + ) { + primitive?.let { put(it, schema) } + put(boxed, schema) + } + put(java.lang.Boolean.TYPE, java.lang.Boolean::class.java, mapOf("type" to "boolean")) + put(java.lang.Byte.TYPE, java.lang.Byte::class.java, integerRange(Byte.MIN_VALUE.toLong(), Byte.MAX_VALUE.toLong())) + put(java.lang.Short.TYPE, java.lang.Short::class.java, integerRange(Short.MIN_VALUE.toLong(), Short.MAX_VALUE.toLong())) + put(java.lang.Integer.TYPE, java.lang.Integer::class.java, integer("int32")) + put(java.lang.Long.TYPE, java.lang.Long::class.java, integer("int64")) + put(java.lang.Float.TYPE, java.lang.Float::class.java, number("float")) + put(java.lang.Double.TYPE, java.lang.Double::class.java, number("double")) + put( + java.lang.Character.TYPE, + java.lang.Character::class.java, + mapOf("type" to "string", "minLength" to 1, "maxLength" to 1), + ) + put(null, java.lang.String::class.java, mapOf("type" to "string")) + put(null, BigInteger::class.java, mapOf("type" to "integer")) + put(null, BigDecimal::class.java, mapOf("type" to "number")) + put(null, java.lang.Number::class.java, mapOf("type" to "number")) + put(null, UUID::class.java, string("uuid")) + put(null, OffsetDateTime::class.java, string("date-time")) + put(null, ZonedDateTime::class.java, string("date-time")) + put(null, LocalDateTime::class.java, string("date-time")) + put(null, Instant::class.java, string("date-time")) + put(null, LocalDate::class.java, string("date")) + put(null, LocalTime::class.java, string("time")) + put(null, Duration::class.java, string("duration")) + } + +private fun classSchema(type: Class<*>): Schema? { + SCALARS[type]?.let { return it } + return when { + // A byte array decodes from a base64 string as well as from a list of numbers. + type == ByteArray::class.java -> null + type.isArray -> arraySchema(type.componentType) + type.isEnum -> enumSchema(type) + Collection::class.java.isAssignableFrom(type) -> mapOf("type" to "array", "items" to UNCONSTRAINED) + Map::class.java.isAssignableFrom(type) -> mapOf("type" to "object", "additionalProperties" to true) + // Other platform types (Object, Optional, Date, URI, ...) decode from more + // than one JSON shape, or from none, so they constrain nothing here. + isPlatformType(type) -> null + // A POJO's fields are not described: the schema says only that it takes an object. + else -> mapOf("type" to "object") + } +} + +private fun isPlatformType(type: Class<*>): Boolean = + type.name.startsWith("java.") || type.name.startsWith("javax.") || type.name.startsWith("kotlin.") + +// Read from the constant fields rather than enumConstants, which would run the +// enum's static initializer while the bundle is only being described. getFields +// has no defined order, so the names are sorted. +private fun enumSchema(type: Class<*>): Schema = + mapOf( + "type" to "string", + "enum" to + type.fields + .filter { it.isEnumConstant } + .map { it.name } + .sorted(), + ) + +private fun arraySchema(component: Type): Schema = mapOf("type" to "array", "items" to (buildValueSchema(component) ?: UNCONSTRAINED)) + +private fun parameterizedSchema(type: ParameterizedType): Schema? { + val raw = type.rawType as? Class<*> ?: return null + val arguments = type.actualTypeArguments + return when { + Collection::class.java.isAssignableFrom(raw) && arguments.size == 1 -> arraySchema(arguments[0]) + Map::class.java.isAssignableFrom(raw) && arguments.size == 2 -> + mapOf("type" to "object", "additionalProperties" to (buildValueSchema(arguments[1]) ?: true)) + else -> classSchema(raw) + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt index e436dd2fc6160..65e062b2dedcd 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt @@ -27,16 +27,25 @@ import kotlinx.coroutines.runBlocking import org.apache.airflow.sdk.execution.CoordinatorComm import org.apache.airflow.sdk.execution.Frame import org.apache.airflow.sdk.execution.IncomingFrame +import org.apache.airflow.sdk.execution.RawFrame import org.apache.airflow.sdk.execution.comm.TaskState import org.junit.jupiter.api.Assertions import org.junit.jupiter.api.DisplayName import org.junit.jupiter.api.Test import org.junit.jupiter.api.Timeout import org.msgpack.core.MessagePack +import org.msgpack.core.buffer.ArrayBufferInput import java.io.ByteArrayOutputStream import java.util.concurrent.ArrayBlockingQueue import java.util.concurrent.TimeUnit +private class ServerHandlerTask : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + class ServerTest { private fun hexToBytes(hex: String): ByteArray = hex @@ -112,6 +121,86 @@ class ServerTest { comm.close() } + @Test + @DisplayName("Should answer a task handler parse request and wait for its acknowledgement") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun answersTaskHandlerParseRequest() { + val toServer = ByteChannel(autoFlush = true) + val fromServer = ByteChannel(autoFlush = true) + val comm = CoordinatorComm(toServer, fromServer) + val server = Server(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0)) + val bundle = Bundle().register("etl", "audit", ServerHandlerTask::class.java) + + val reported = ArrayBlockingQueue(1) + val supervisor = + Thread { + runBlocking { + toServer.writeFrame(taskHandlerParseRequestFrame(3, "/bundle/etl.jar")) + val prefix = fromServer.readByteArray(4) + val raw = Frame.decodeRaw(ArrayBufferInput(fromServer.readByteArray(Frame.parseLengthPrefix(prefix).toInt()))) + reported.put(raw) + toServer.writeFrame(ackFrame(raw.id)) + } + } + supervisor.start() + + runBlocking { server.dispatchTask(bundle, comm) } + supervisor.join() + + val body = reported.take().rawBody as Map<*, *> + Assertions.assertEquals("TaskHandlerParsingResult", body["type"]) + Assertions.assertEquals("/bundle/etl.jar", body["fileloc"]) + Assertions.assertEquals( + mapOf("etl" to listOf(mapOf("task_id" to "audit", "binding" to "named", "params" to emptyList()))), + body["task_handlers"], + ) + comm.close() + } + + @Test + @DisplayName("Should fail a task handler parse whose result is never acknowledged") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun failsTaskHandlerParseWithoutAcknowledgement() { + val toServer = ByteChannel(autoFlush = true) + val fromServer = ByteChannel(autoFlush = true) + val comm = CoordinatorComm(toServer, fromServer) + val server = Server(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0)) + + val supervisor = + Thread { + runBlocking { + toServer.writeFrame(taskHandlerParseRequestFrame(3, "/bundle/etl.jar")) + val prefix = fromServer.readByteArray(4) + fromServer.readByteArray(Frame.parseLengthPrefix(prefix).toInt()) + toServer.close() + } + } + supervisor.start() + + Assertions.assertThrows(ApiError::class.java) { + runBlocking { server.dispatchTask(Bundle(), comm) } + } + supervisor.join() + comm.close() + } + + private fun taskHandlerParseRequestFrame( + id: Int, + file: String, + ): ByteArray { + val out = ByteArrayOutputStream() + MessagePack.newDefaultPacker(out).use { packer -> + packer.packArrayHeader(2) + packer.packInt(id) + packer.packMapHeader(4) + packer.packString("type").packString("TaskHandlerParseRequest") + packer.packString("file").packString(file) + packer.packString("bundle_path").packString("/bundle") + packer.packString("bundle_name").packString("java-task-handlers") + } + return out.toByteArray() + } + private companion object { // [2, msg, null] with msg coming from // https://github.com/astronomer/airflow/blob/f39c8da8/task-sdk/tests/task_sdk/execution_time/test_comms.py#L73-L108 diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt new file mode 100644 index 0000000000000..02ae55c96a518 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskHandlerParseTest.kt @@ -0,0 +1,232 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk.execution + +import com.fasterxml.jackson.databind.JsonNode +import com.fasterxml.jackson.databind.ObjectMapper +import org.apache.airflow.sdk.ArgName +import org.apache.airflow.sdk.Bundle +import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.Context +import org.apache.airflow.sdk.DagDef +import org.apache.airflow.sdk.InputTask +import org.apache.airflow.sdk.Task +import org.apache.airflow.sdk.TaskInput +import org.apache.airflow.sdk.execution.comm.TaskHandlerParseRequest +import org.apache.airflow.sdk.internal.TaskParams +import org.apache.airflow.sdk.internal.TypeRef +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import java.io.File + +private val NULLABLE_STRING = mapOf("anyOf" to listOf(mapOf("type" to "string"), mapOf("type" to "null"))) + +/** Stands in for a class the annotation processor generates for a handler with flat data parameters. */ +private class GeneratedFlat : Task { + companion object { + @JvmField + val AIRFLOW_TASK_PARAMS: TaskParams = + TaskParams.of( + TaskParams.param("rows", java.lang.Long.TYPE), + TaskParams.param("regions", object : TypeRef>() {}), + ) + } + + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +/** Stands in for a generated class whose handler takes only the injected Client and Context. */ +private class GeneratedNoData : Task { + companion object { + @JvmField + val AIRFLOW_TASK_PARAMS: TaskParams = TaskParams.of() + } + + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +private open class BaseScoreInput : TaskInput { + @JvmField + var threshold: Double = 0.0 +} + +private class ScoreInput : BaseScoreInput() { + @JvmField + @ArgName("region_code") + var region: String? = null +} + +/** Stands in for a generated class whose handler takes a [TaskInput]. */ +private class GeneratedInput : Task { + companion object { + @JvmField + val AIRFLOW_TASK_PARAMS: TaskParams = TaskParams.input(ScoreInput::class.java) + } + + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +private class Score : InputTask { + override fun execute( + context: Context, + client: Client, + input: ScoreInput, + ) = Unit +} + +private class Audit : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +private fun request() = + TaskHandlerParseRequest().apply { + file = "/bundles/java/etl.jar" + bundlePath = "/bundles/java" + bundleName = "java-task-handlers" + } + +private fun param( + name: String, + schema: Map, + exactName: Boolean = false, +) = mapOf("name" to name, "exact_name" to exactName, "value_schema" to schema) + +private val SCORE_INPUT_PARAMS = + listOf( + param("region_code", NULLABLE_STRING, exactName = true), + param("threshold", mapOf("type" to "number", "format" to "double")), + ) + +// Generated from the Python models, so it is what the Dag processor decodes the reply with. +private val SUPERVISOR_SCHEMA: JsonNode = ObjectMapper().readTree(File("schema/schema.json")).path("${'$'}defs") + +private fun assertConforms( + definition: String, + body: Map<*, *>, +) { + val schema = SUPERVISOR_SCHEMA.path(definition) + val properties = + schema + .path("properties") + .fieldNames() + .asSequence() + .toSet() + val required = schema.path("required").map { it.asText() }.toSet() + assertTrue(properties.containsAll(body.keys), "$definition has no ${body.keys - properties}") + assertTrue(body.keys.containsAll(required), "$definition misses ${required - body.keys}") +} + +class TaskHandlerParseTest { + @Test + @DisplayName("Should declare every task handler in registration order") + fun declaresEveryTaskHandlerInRegistrationOrder() { + val bundle = + Bundle() + .register("etl", "extract", GeneratedFlat::class.java) + .register("etl", "score", GeneratedInput::class.java) + .register("etl", "summarize", Score::class.java) + .register("etl", "audit", Audit::class.java) + .register("etl", "notify", GeneratedNoData::class.java) + .register("report", "audit", Audit::class.java) + .register(DagDef("native").addTask("audit", Audit::class.java)) + + val result = parseTaskHandlers(bundle, request()) + + val regions = + mapOf("anyOf" to listOf(mapOf("type" to "array", "items" to NULLABLE_STRING), mapOf("type" to "null"))) + val expected = + mapOf( + "type" to "TaskHandlerParsingResult", + "fileloc" to "/bundles/java/etl.jar", + "task_handlers" to + mapOf( + "etl" to + listOf( + mapOf( + "task_id" to "extract", + "binding" to "positional", + "params" to + listOf( + param("rows", mapOf("type" to "integer", "format" to "int64")), + param("regions", regions), + ), + ), + mapOf("task_id" to "score", "binding" to "named", "params" to SCORE_INPUT_PARAMS), + mapOf("task_id" to "summarize", "binding" to "named", "params" to SCORE_INPUT_PARAMS), + mapOf("task_id" to "audit", "binding" to "named", "params" to emptyList()), + mapOf("task_id" to "notify", "binding" to "named", "params" to emptyList()), + ), + "report" to listOf(mapOf("task_id" to "audit", "binding" to "named", "params" to emptyList())), + ), + ) + assertEquals(expected, result) + assertEquals(listOf("etl", "report"), (result["task_handlers"] as Map<*, *>).keys.toList()) + } + + @Test + @DisplayName("Should send only the fields the supervisor schema declares, and every required one") + fun conformsToSupervisorSchema() { + val bundle = + Bundle() + .register("etl", "extract", GeneratedFlat::class.java) + .register("etl", "score", GeneratedInput::class.java) + .register("etl", "audit", Audit::class.java) + + val result = parseTaskHandlers(bundle, request()) + + assertConforms("TaskHandlerParsingResult", result) + val bindings = + SUPERVISOR_SCHEMA + .path("TaskHandlerDeclaration") + .path("properties") + .path("binding") + .path("enum") + for (declaration in (result["task_handlers"] as Map<*, *>).values.flatMap { it as List<*> }) { + declaration as Map<*, *> + assertConforms("TaskHandlerDeclaration", declaration) + assertTrue(bindings.any { it.asText() == declaration["binding"] }, "unknown binding ${declaration["binding"]}") + (declaration["params"] as List<*>).forEach { assertConforms("TaskHandlerParam", it as Map<*, *>) } + } + } + + @Test + @DisplayName("Should send an empty map when the bundle registers no task handler") + fun sendsEmptyTaskHandlersWithoutHandlers() { + val bundle = Bundle().register(DagDef("native").addTask("audit", Audit::class.java)) + + val result = parseTaskHandlers(bundle, request()) + + assertEquals(emptyMap(), result["task_handlers"]) + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ValueSchemaTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ValueSchemaTest.kt new file mode 100644 index 0000000000000..7f82c7cddb101 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ValueSchemaTest.kt @@ -0,0 +1,219 @@ +/* + * 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. + */ + +@file:Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN") + +package org.apache.airflow.sdk.internal + +import org.junit.jupiter.api.Assertions.assertAll +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import java.math.BigDecimal +import java.math.BigInteger +import java.time.Duration +import java.time.Instant +import java.time.LocalDate +import java.time.LocalTime +import java.time.OffsetDateTime +import java.util.Date +import java.util.Optional +import java.util.UUID + +private fun nullable(schema: Map) = mapOf("anyOf" to listOf(schema, mapOf("type" to "null"))) + +private val INT64 = mapOf("type" to "integer", "format" to "int64") +private val STRING = mapOf("type" to "string") + +private class Payload + +private enum class Color { RED, GREEN } + +private object EnumInitialization { + var ran = false +} + +private enum class Lazy { + ONLY, + ; + + init { + EnumInitialization.ran = true + } +} + +/** Declares one field per case, so each keeps the generic type a Kotlin declaration compiles to. */ +@Suppress("unused") +private class Declared { + @JvmField var flag: Boolean = false + + @JvmField var boxedFlag: Boolean? = null + + @JvmField var tiny: Byte = 0 + + @JvmField var small: Short = 0 + + @JvmField var count: Int = 0 + + @JvmField var total: Long = 0 + + @JvmField var boxedTotal: Long? = null + + @JvmField var ratio: Float = 0f + + @JvmField var score: Double = 0.0 + + @JvmField var letter: Char = 'a' + + @JvmField var text: String? = null + + @JvmField var big: BigInteger? = null + + @JvmField var exact: BigDecimal? = null + + @JvmField var id: UUID? = null + + @JvmField var color: Color? = null + + @JvmField var at: OffsetDateTime? = null + + @JvmField var instant: Instant? = null + + @JvmField var day: LocalDate? = null + + @JvmField var time: LocalTime? = null + + @JvmField var wait: Duration? = null + + @JvmField var names: List? = null + + @JvmField var anything: List? = null + + @JvmField var words: Array? = null + + @JvmField var totals: LongArray? = null + + @JvmField var bytes: ByteArray? = null + + @JvmField var counts: Map? = null + + @JvmField var payload: Payload? = null + + @JvmField var any: Any? = null + + @JvmField var date: Date? = null + + @JvmField var maybe: Optional? = null + + // Kotlin compiles a non-final type argument of a parameter to a wildcard, `List`. + fun acceptNumbers(numbers: List) = numbers +} + +private fun schemaOf(field: String) = buildValueSchema(Declared::class.java.getField(field).genericType) + +class ValueSchemaTest { + @Test + @DisplayName("Should describe scalars in pydantic's vocabulary, excluding null only for primitives") + fun describesScalars() { + assertAll( + { assertEquals(mapOf("type" to "boolean"), schemaOf("flag")) }, + { assertEquals(nullable(mapOf("type" to "boolean")), schemaOf("boxedFlag")) }, + { assertEquals(mapOf("type" to "integer", "minimum" to -128L, "maximum" to 127L), schemaOf("tiny")) }, + { assertEquals(mapOf("type" to "integer", "minimum" to -32768L, "maximum" to 32767L), schemaOf("small")) }, + { assertEquals(mapOf("type" to "integer", "format" to "int32"), schemaOf("count")) }, + { assertEquals(INT64, schemaOf("total")) }, + { assertEquals(nullable(INT64), schemaOf("boxedTotal")) }, + { assertEquals(mapOf("type" to "number", "format" to "float"), schemaOf("ratio")) }, + { assertEquals(mapOf("type" to "number", "format" to "double"), schemaOf("score")) }, + { assertEquals(mapOf("type" to "string", "minLength" to 1, "maxLength" to 1), schemaOf("letter")) }, + { assertEquals(nullable(STRING), schemaOf("text")) }, + { assertEquals(nullable(mapOf("type" to "integer")), schemaOf("big")) }, + { assertEquals(nullable(mapOf("type" to "number")), schemaOf("exact")) }, + { assertEquals(nullable(mapOf("type" to "string", "format" to "uuid")), schemaOf("id")) }, + { assertEquals(nullable(mapOf("type" to "string", "enum" to listOf("GREEN", "RED"))), schemaOf("color")) }, + ) + } + + @Test + @DisplayName("Should describe temporal types by their ISO-8601 string format") + fun describesTemporalTypes() { + assertAll( + { assertEquals(nullable(mapOf("type" to "string", "format" to "date-time")), schemaOf("at")) }, + { assertEquals(nullable(mapOf("type" to "string", "format" to "date-time")), schemaOf("instant")) }, + { assertEquals(nullable(mapOf("type" to "string", "format" to "date")), schemaOf("day")) }, + { assertEquals(nullable(mapOf("type" to "string", "format" to "time")), schemaOf("time")) }, + { assertEquals(nullable(mapOf("type" to "string", "format" to "duration")), schemaOf("wait")) }, + ) + } + + @Test + @DisplayName("Should describe containers with the schema of what they hold") + fun describesContainers() { + assertAll( + { assertEquals(nullable(mapOf("type" to "array", "items" to nullable(STRING))), schemaOf("names")) }, + { assertEquals(nullable(mapOf("type" to "array", "items" to emptyMap())), schemaOf("anything")) }, + { + assertEquals( + nullable(mapOf("type" to "array", "items" to nullable(mapOf("type" to "number")))), + buildValueSchema(Declared::class.java.getMethod("acceptNumbers", List::class.java).genericParameterTypes[0]), + ) + }, + { assertEquals(nullable(mapOf("type" to "array", "items" to nullable(STRING))), schemaOf("words")) }, + { assertEquals(nullable(mapOf("type" to "array", "items" to INT64)), schemaOf("totals")) }, + { + assertEquals( + nullable(mapOf("type" to "object", "additionalProperties" to nullable(INT64))), + schemaOf("counts"), + ) + }, + { + assertEquals( + nullable(mapOf("type" to "array", "items" to emptyMap())), + buildValueSchema(java.util.List::class.java), + ) + }, + { + assertEquals( + nullable(mapOf("type" to "object", "additionalProperties" to true)), + buildValueSchema(java.util.Map::class.java), + ) + }, + { assertEquals(nullable(mapOf("type" to "object")), schemaOf("payload")) }, + ) + } + + @Test + @DisplayName("Should say nothing for a type that accepts more than one shape, or none") + fun describesNothingForOpenTypes() { + assertAll( + { assertEquals(null, schemaOf("bytes")) }, + { assertEquals(null, schemaOf("any")) }, + { assertEquals(null, schemaOf("date")) }, + { assertEquals(null, schemaOf("maybe")) }, + ) + } + + @Test + @DisplayName("Should list an enum's constants without initializing it") + fun listsEnumConstantsWithoutInitializing() { + assertEquals(nullable(mapOf("type" to "string", "enum" to listOf("ONLY"))), buildValueSchema(Lazy::class.java)) + assertFalse(EnumInitialization.ran) + } +} diff --git a/kubernetes-tests/lang_sdk/README.md b/kubernetes-tests/lang_sdk/README.md index 93583ecbd8b44..f19ae247529eb 100644 --- a/kubernetes-tests/lang_sdk/README.md +++ b/kubernetes-tests/lang_sdk/README.md @@ -36,16 +36,16 @@ End-to-end test that one Dag mixing **Python + Go + Java** tasks runs to success worker pod (from that pod template): initContainer stage_artifacts.py ── S3DagBundle.initialize() ──► pulls go-artifacts / java-artifacts bucket into the shared - emptyDir = executables_root / jars_root + emptyDir = go-task-handlers / java-task-handlers LocalDagBundle base container supervisor → coordinator forks the Go binary / Java jar ``` -Key point: in coordinator mode the worker pod does **not** run the Python task-runner, -so the artifact is not downloaded by the normal Dag-bundle path. The init container runs -`stage_artifacts.py`, which reuses the **DagBundle interface** +Key point: each coordinator's `task_handler_bundle_name` names a `LocalDagBundle` over the +shared `emptyDir`, registered in `dagProcessor.dagBundleConfigList`. The init container fills it +by running `stage_artifacts.py`, which reuses the **DagBundle interface** (`DagBundlesManager().get_bundle(name).initialize()` — the download half of -`task_runner.parse`) to pull the artifact from its S3 bucket into the path the -coordinator scans. +`task_runner.parse`) to pull the artifact from its S3 bucket, then restores the Go binary's +execute bit, which the S3 download drops and the coordinator requires. ## Components @@ -58,7 +58,7 @@ coordinator scans. | `pod_templates/lang_sdk_golang.yaml` | `golang` queue worker pod: prod image + go-artifacts init container. | | `pod_templates/lang_sdk_java.yaml` | `java` queue worker pod: JVM image + java-artifacts init container. | | `manifests/localstack.yaml` | In-cluster S3 (localstack). | -| `config/values.yaml` | Helm overrides: KubernetesExecutor, coordinators (+extra.pod_template_file), queue routing, stub-Dag S3 bundle, AWS conn, scheduler pod-template mount. | +| `config/values.yaml` | Helm overrides: KubernetesExecutor, coordinators (+extra.pod_template_file), queue routing, stub-Dag S3 bundle, artifact Dag bundles, AWS conn, scheduler pod-template mount. | The Go binary, Java jar, and stub Dag share one object store (localstack) but live in **separate buckets** (`go-artifacts`, `java-artifacts`, `dags`). diff --git a/kubernetes-tests/lang_sdk/config/values.yaml b/kubernetes-tests/lang_sdk/config/values.yaml index a4b5b4004b403..6f4d9671ad6a5 100644 --- a/kubernetes-tests/lang_sdk/config/values.yaml +++ b/kubernetes-tests/lang_sdk/config/values.yaml @@ -50,6 +50,10 @@ scheduler: # Configure it here (dagProcessor.dagBundleConfigList) rather than overriding # config.dag_processor.dag_bundle_config_list directly -- the chart renders the config from this and # it is the recommended, drift-safe path. +# +# go-task-handlers / java-task-handlers are the coordinators' task_handler_bundle_name bundles: the +# worker pod's init container stages each artifact bucket into these paths (stage_artifacts.py). +# They exist only in the worker pods, so the dag-processor finds no files in them. dagProcessor: dagBundleConfigList: - name: dags-folder @@ -60,6 +64,14 @@ dagProcessor: kwargs: bucket_name: "dags" aws_conn_id: "aws_localstack" + - name: go-task-handlers + classpath: "airflow.dag_processing.bundles.local.LocalDagBundle" + kwargs: + path: "/opt/airflow/artifacts/go-artifacts" + - name: java-task-handlers + classpath: "airflow.dag_processing.bundles.local.LocalDagBundle" + kwargs: + path: "/opt/airflow/artifacts/java-artifacts" config: # Route the golang/java queues to their coordinators, each carrying an @@ -70,5 +82,5 @@ config: # Single-line JSON: airflow.cfg is an ini file, so a value with embedded # newlines would break the rendered configmap. # yamllint disable-line rule:line-length - coordinators: '{"go-sdk": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": {"executables_root": ["/opt/airflow/artifacts/go-artifacts"]}, "extra": {"pod_template_file": "/opt/airflow/lang_sdk/pod_templates/lang_sdk_golang.yaml"}}, "java-sdk": {"classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": {"jars_root": ["/opt/airflow/artifacts/java-artifacts"]}, "extra": {"pod_template_file": "/opt/airflow/lang_sdk/pod_templates/lang_sdk_java.yaml", "worker_container_repository": "lang-sdk-java-worker", "worker_container_tag": "latest"}}}' + coordinators: '{"go-sdk": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": {"task_handler_bundle_name": "go-task-handlers"}, "extra": {"pod_template_file": "/opt/airflow/lang_sdk/pod_templates/lang_sdk_golang.yaml"}}, "java-sdk": {"classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": {"task_handler_bundle_name": "java-task-handlers"}, "extra": {"pod_template_file": "/opt/airflow/lang_sdk/pod_templates/lang_sdk_java.yaml", "worker_container_repository": "lang-sdk-java-worker", "worker_container_tag": "latest"}}}' queue_to_coordinator: '{"golang": "go-sdk", "java": "java-sdk"}' diff --git a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/CombinedExample.java b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/CombinedExample.java index 63aaa5b6f051d..9e44edca50cfa 100644 --- a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/CombinedExample.java +++ b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/CombinedExample.java @@ -25,21 +25,19 @@ import org.slf4j.LoggerFactory; // Java half of the KubernetesExecutor lang-SDK system test bundle. Registers the -// Java tasks of the shared "lang_sdk_combined" Dag; the Go tasks of the same -// dag_id live in ../go_example and the Python stub Dag in ../dags. The -// coordinator locates this jar by dag_id, so only the Java tasks are registered -// here. -@Builder.Dag(id = "lang_sdk_combined") +// Java task handlers of the shared "lang_sdk_combined" Dag; the Go tasks of the +// same dag_id live in ../go_example and the Python stub Dag in ../dags. The +// Python file owns the Dag, so only the Java tasks are registered here. public class CombinedExample { private static final Logger logger = LoggerFactory.getLogger(CombinedExample.class); - @Builder.Task(id = "java_extract") + @Builder.TaskHandler(dag = "lang_sdk_combined", task = "java_extract") public long extract(Client client) { logger.info("java_extract running"); return new Date().getTime(); } - @Builder.Task(id = "java_transform") + @Builder.TaskHandler(dag = "lang_sdk_combined", task = "java_transform") public void transform(Client client) { var variable = client.getVariable("my_variable"); logger.info("java_transform obtained variable {}", variable); diff --git a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java index c3b511b898ced..064dfebd3566f 100644 --- a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java +++ b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java @@ -19,21 +19,10 @@ package org.apache.airflow.k8sexample; -import java.util.List; import org.apache.airflow.sdk.*; -import org.jetbrains.annotations.NotNull; - -// CombinedExampleBuilder is generated by the annotation processor from the -// @Builder.Dag class CombinedExample. -public class K8sBundleBuilder implements BundleBuilder { - @NotNull - @Override - public Iterable getDags() { - return List.of(CombinedExampleBuilder.build()); - } +public class K8sBundleBuilder { public static void main(String[] args) { - var bundle = new K8sBundleBuilder().build(); - Server.create(args).serve(bundle); + Server.create(args).serve(new Bundle().register(CombinedExample.class)); } } diff --git a/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_golang.yaml b/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_golang.yaml index e3e700d700012..00d6b990690ae 100644 --- a/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_golang.yaml +++ b/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_golang.yaml @@ -19,7 +19,8 @@ # # An init container stages the packed Go bundle from the localstack ``go-artifacts`` # bucket into a shared emptyDir at /opt/airflow/artifacts via the DagBundle -# interface (stage_artifacts.py). That path is the coordinator's executables_root. +# interface (stage_artifacts.py). That path is the go-task-handlers Dag bundle the +# coordinator's task_handler_bundle_name names. # The base image is the plain prod image: a Go bundle is a self-contained static # binary, so no extra runtime is needed. Image placeholders are rendered by the # ``breeze k8s setup-lang-sdk-test`` command. diff --git a/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_java.yaml b/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_java.yaml index eaccd63f4d2ad..ad4456610fb73 100644 --- a/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_java.yaml +++ b/kubernetes-tests/lang_sdk/pod_templates/lang_sdk_java.yaml @@ -19,7 +19,8 @@ # # An init container stages the Java jar from the localstack ``java-artifacts`` # bucket into a shared emptyDir at /opt/airflow/artifacts via the DagBundle -# interface (stage_artifacts.py). That path is the coordinator's jars_root. The +# interface (stage_artifacts.py). That path is the java-task-handlers Dag bundle the +# coordinator's task_handler_bundle_name names. The # base image is the prod image plus a headless JRE (built from Dockerfile.java) # so the JavaCoordinator can exec ``java`` -- distinct from the Go queue's plain # prod image. Image placeholders are rendered by ``breeze k8s setup-lang-sdk-test``. diff --git a/kubernetes-tests/lang_sdk/stage_artifacts.py b/kubernetes-tests/lang_sdk/stage_artifacts.py index 3da9ac6fb5a72..45a9d69882b03 100644 --- a/kubernetes-tests/lang_sdk/stage_artifacts.py +++ b/kubernetes-tests/lang_sdk/stage_artifacts.py @@ -23,10 +23,10 @@ parse), so the Go binary / Java jar stored in an S3 (localstack) bucket is pulled into ``bundle.path`` = ``{dag_bundle_storage_path}/{name}``. -That path is the ``emptyDir`` shared with the worker container and is exactly the -``executables_root`` / ``jars_root`` the coordinator scans. Coordinator mode never -runs the Python task-runner in the worker pod, so without this step the artifact -would never reach the pod. +That path is the ``emptyDir`` shared with the worker container, where it is registered +as the ``LocalDagBundle`` the coordinator's ``task_handler_bundle_name`` names. The +coordinator does not name the S3 bundle directly because its download drops the Go +binary's execute bit, which the coordinator requires; Java is staged the same way. Configuration is read from the environment (set by the pod template): diff --git a/scripts/ci/prek/generate_supervisor_schemas_snapshot.py b/scripts/ci/prek/generate_supervisor_schemas_snapshot.py index a626f32718e06..106f12a836663 100755 --- a/scripts/ci/prek/generate_supervisor_schemas_snapshot.py +++ b/scripts/ci/prek/generate_supervisor_schemas_snapshot.py @@ -27,7 +27,8 @@ The snapshot is the head-version JSON Schema for every Pydantic class on the supervisor schema wire (the union members of ``ToTask``, -``ToSupervisor``, ``ToManager``, ``ToDagProcessor``). +``ToSupervisor``, ``ToManager``, ``ToDagProcessor``, +``ToSDKTaskHandlerProcessor``). The actual dump is delegated to ``dump_supervisor_schemas.py`` (the sibling stdout-only script). If the committed snapshot differs from diff --git a/scripts/ci/prek/sync_java_sdk_supervisor_schema.py b/scripts/ci/prek/sync_java_sdk_supervisor_schema.py index d3f4d271a5f5a..c270920c09ac2 100755 --- a/scripts/ci/prek/sync_java_sdk_supervisor_schema.py +++ b/scripts/ci/prek/sync_java_sdk_supervisor_schema.py @@ -18,11 +18,12 @@ """ Keep the Java SDK's bundled Supervisor Schema in sync with ``airflowSupervisorSchemaVersion``. -The Gradle task ``:sdk:syncSupervisorSchema`` downloads a fresh ``schema.json`` when the -``api_version`` inside it differs from the version declared in ``java-sdk/gradle.properties``. +The Gradle task ``:sdk:syncSupervisorSchema`` copies the Task SDK's snapshot when it declares the +configured ``api_version``, and otherwise downloads the published schema when the ``api_version`` +inside ``schema.json`` differs from the version declared in ``java-sdk/gradle.properties``. Starting Gradle for that comparison costs over a minute on a cold CI runner (wrapper download, JVM start, plugin resolution), so this hook does the same comparison in Python first and only -hands over to Gradle when the two versions actually differ. +hands over to Gradle when a download is needed. """ from __future__ import annotations @@ -37,6 +38,9 @@ JAVA_SDK_DIR = REPO_ROOT / "java-sdk" GRADLE_PROPERTIES = JAVA_SDK_DIR / "gradle.properties" SCHEMA_FILE = JAVA_SDK_DIR / "sdk" / "schema" / "schema.json" +MONOREPO_SCHEMA_FILE = ( + REPO_ROOT / "task-sdk" / "src" / "airflow" / "sdk" / "execution_time" / "schema" / "schema.json" +) VERSION_PATTERN = re.compile(r"^airflowSupervisorSchemaVersion\s*=\s*(\S+)\s*$", re.MULTILINE) @@ -49,20 +53,28 @@ def configured_version() -> str: return match.group(1) -def bundled_version() -> str | None: - if not SCHEMA_FILE.exists(): +def read_api_version(path: Path) -> str | None: + if not path.exists(): return None try: - with SCHEMA_FILE.open() as schema: + with path.open() as schema: return json.load(schema).get("api_version") except json.JSONDecodeError: - # A truncated or conflict-marked file counts as drift; Gradle overwrites it. + # A truncated or conflict-marked file counts as drift; it is overwritten. return None def main() -> int: expected = configured_version() - actual = bundled_version() + if read_api_version(MONOREPO_SCHEMA_FILE) == expected: + snapshot = MONOREPO_SCHEMA_FILE.read_bytes() + if SCHEMA_FILE.exists() and SCHEMA_FILE.read_bytes() == snapshot: + print(f"Supervisor Schema matches the Task SDK snapshot (api_version={expected}).") + else: + print(f"Refreshing Supervisor Schema from {MONOREPO_SCHEMA_FILE.relative_to(REPO_ROOT)}.") + SCHEMA_FILE.write_bytes(snapshot) + return 0 + actual = read_api_version(SCHEMA_FILE) if actual == expected: print(f"Supervisor Schema is up-to-date (api_version={expected}).") return 0 diff --git a/scripts/tests/ci/prek/test_sync_java_sdk_supervisor_schema.py b/scripts/tests/ci/prek/test_sync_java_sdk_supervisor_schema.py new file mode 100644 index 0000000000000..ec60dd043734f --- /dev/null +++ b/scripts/tests/ci/prek/test_sync_java_sdk_supervisor_schema.py @@ -0,0 +1,90 @@ +# 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. +from __future__ import annotations + +import json +import os +import subprocess +from typing import TYPE_CHECKING +from unittest import mock + +import pytest +from ci.prek import sync_java_sdk_supervisor_schema as sync + +if TYPE_CHECKING: + from pathlib import Path + +CONFIGURED_VERSION = "2026-06-16" + + +def _write_schema(path: Path, api_version: str, **body) -> None: + path.write_text(json.dumps({"api_version": api_version, **body})) + + +@pytest.fixture +def schemas(tmp_path, monkeypatch): + """Point the hook at a temp gradle.properties, vendored schema and Task SDK snapshot.""" + gradle_properties = tmp_path / "gradle.properties" + gradle_properties.write_text(f"airflowSupervisorSchemaVersion={CONFIGURED_VERSION}\n") + vendored = tmp_path / "java-sdk" / "schema.json" + snapshot = tmp_path / "task-sdk" / "schema.json" + vendored.parent.mkdir() + snapshot.parent.mkdir() + monkeypatch.setattr(sync, "REPO_ROOT", tmp_path) + monkeypatch.setattr(sync, "GRADLE_PROPERTIES", gradle_properties) + monkeypatch.setattr(sync, "SCHEMA_FILE", vendored) + monkeypatch.setattr(sync, "MONOREPO_SCHEMA_FILE", snapshot) + return vendored, snapshot + + +@mock.patch.object(sync.subprocess, "run", autospec=True) +class TestMain: + def test_leaves_a_vendored_schema_that_matches_the_snapshot(self, mock_run, schemas, capsys): + vendored, snapshot = schemas + _write_schema(snapshot, CONFIGURED_VERSION) + _write_schema(vendored, CONFIGURED_VERSION) + os.utime(vendored, ns=(0, 0)) + + assert sync.main() == 0 + + assert "matches the Task SDK snapshot" in capsys.readouterr().out + assert vendored.stat().st_mtime_ns == 0 + mock_run.assert_not_called() + + def test_copies_the_snapshot_over_a_differing_vendored_schema(self, mock_run, schemas): + vendored, snapshot = schemas + _write_schema(snapshot, CONFIGURED_VERSION, definitions={"New": {}}) + _write_schema(vendored, CONFIGURED_VERSION) + + assert sync.main() == 0 + + assert vendored.read_bytes() == snapshot.read_bytes() + mock_run.assert_not_called() + + def test_runs_gradle_when_the_snapshot_declares_another_version(self, mock_run, schemas): + vendored, snapshot = schemas + _write_schema(snapshot, "2026-01-01") + _write_schema(vendored, "2025-01-01") + mock_run.return_value = subprocess.CompletedProcess(args=[], returncode=0) + + assert sync.main() == 0 + + mock_run.assert_called_once_with( + [str(sync.JAVA_SDK_DIR / "gradlew"), "-p", str(sync.JAVA_SDK_DIR), ":sdk:syncSupervisorSchema"], + check=False, + ) + assert json.loads(vendored.read_text()) == {"api_version": "2025-01-01"} diff --git a/task-sdk/docs/airflow-metadata.schema.json b/task-sdk/docs/airflow-metadata.schema.json index d4a5cd240e712..d2a8f35ba000a 100644 --- a/task-sdk/docs/airflow-metadata.schema.json +++ b/task-sdk/docs/airflow-metadata.schema.json @@ -40,6 +40,23 @@ "description": "Original filename of the primary Dag source file (e.g. 'example.go'). Bundle formats with an embedded source region use this as its display name; other formats treat it as the logical authoring name.", "minLength": 1 }, + "digests": { + "type": "object", + "description": "Digests the packer records, so a consumer can compare a stored value instead of computing one. Consumers ignore keys they do not know.", + "additionalProperties": true, + "properties": { + "integrity": { + "type": "string", + "description": "SHA-256 of the binary region: the same value as the trailer's binary_sha256. Task execution still recomputes the trailer's digest.", + "pattern": "^[0-9a-f]{64}$" + }, + "cache": { + "type": "string", + "description": "Stays the same across packs of the same inputs and changes when any region of the bundle changes. It is read, never recomputed, so it does not prove the bundle is intact. The producer defines how it is computed; consumers only compare it for equality.", + "pattern": "^[0-9a-f]{64}$" + } + } + }, "dags": { "type": "object", "description": "Mapping of dag_id to DAG entry. Every dag_id the bundle exposes must appear here.", diff --git a/task-sdk/docs/executable-bundle-spec.rst b/task-sdk/docs/executable-bundle-spec.rst index 856aeca94e922..83d91fbc5ca77 100644 --- a/task-sdk/docs/executable-bundle-spec.rst +++ b/task-sdk/docs/executable-bundle-spec.rst @@ -221,6 +221,24 @@ Top-level keys: source-view panel and pick a syntax-highlighting mode from the extension. +``digests`` (mapping, optional) + Digests the packer records, so a consumer can compare a stored value + instead of computing one. Both are lower-case hex SHA-256 digests, and + consumers ignore keys they do not know. + + - ``integrity`` (string): SHA-256 of the binary region, the same value as + the trailer's ``binary_sha256``. Task execution still recomputes the + trailer's digest; this copy does not replace that check. + - ``cache`` (string): stays the same across packs of the same inputs and + changes when any region of the bundle changes. The Dag processor reads + it to tell whether a bundle changed since it last probed it. It is read, + never recomputed, so it does not prove the bundle is intact. The + producer defines how it is computed, and consumers only compare it for + equality. ``airflow-go-pack`` computes the SHA-256 of the bytes + ``airflow-go-pack cache digest v1`` and a NUL, followed by the raw + SHA-256 digests of the binary region, the source region, and the + manifest as rendered without the ``digests`` key, in that order. + ``dags`` (mapping, required) Mapping of ``dag_id`` to a *DAG entry*. Every ``dag_id`` the bundle exposes MUST appear here. The scanner uses these keys to match a DAG @@ -283,10 +301,10 @@ Code Signing The bundle format itself does not require OS-level code signing. ``binary_sha256`` provides integrity against truncation, in-flight -corruption, and naive tampering, and Airflow's threat model treats -``executables_root`` as Deployment-Manager-controlled — *authenticity* -(signed by a trusted identity) is a deployment-time concern rather than a -bundle-format one. +corruption, and naive tampering, and Airflow's threat model treats the Dag +bundle holding the bundles as Deployment-Manager-controlled, so +*authenticity* (signed by a trusted identity) is a deployment-time concern +rather than a bundle-format one. **Compressors** such as UPX are NOT supported. They rewrite the file end-to-end, destroying both the trailer and the hash invariant. @@ -299,22 +317,23 @@ yields a byte-identical bundle file (and therefore a stable Deployment Layout ----------------- -Bundle files are placed **as-is** in any of the directories configured as the -``executables_root`` kwarg on the +Bundle files are placed **as-is** in the Dag bundle named by the +``task_handler_bundle_name`` kwarg on the :class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` entry -under ``[sdk] coordinators``. The scanner walks each directory **recursively** -and considers only regular files whose **executable bit is set** for the -invoking user; files without the executable bit are skipped without reading -their trailer. For each candidate it reads the last 64 bytes and treats files -whose magic matches ``"AFBNDL01"`` as bundles. Matched files are then -SHA-256-verified per the reader algorithm; a mismatch demotes the file back -to "ignored, with an error log." Files without the magic are silently -ignored, so non-bundle files (READMEs, dotfiles) MAY share the directory +under ``[sdk] coordinators`` (or, when it is unset, in the task's own Dag +bundle). The scanner walks the Dag bundle **recursively** and considers only +regular files whose **executable bit is set** for the invoking user; files +without the executable bit are skipped without reading their trailer, so the +Dag bundle must preserve it. For each candidate it reads the last 64 bytes and +treats files whose magic matches ``"AFBNDL01"`` as bundles. Matched files are +then SHA-256-verified per the reader algorithm; a mismatch demotes the file +back to "ignored, with an error log." Files without the magic are silently +ignored, so non-bundle files (READMEs, dotfiles) MAY share the Dag bundle without interfering with the scan. :: - /opt/airflow/executable-bundles/ + /opt/airflow/go-task-handlers/ ├── example ├── team-a/ │ └── pipeline diff --git a/task-sdk/docs/lang-sdk-spec.rst b/task-sdk/docs/lang-sdk-spec.rst index fed3c2d6f153f..fb99e68919706 100644 --- a/task-sdk/docs/lang-sdk-spec.rst +++ b/task-sdk/docs/lang-sdk-spec.rst @@ -165,3 +165,23 @@ wrote. How ``fn`` reaches ``Context`` and ``Client`` is each SDK's own choice getters that read the scope opened at step 6 — and all this spec asks is that ``fn`` receives the same pair the SDK built at step 3. The terminal state at step 7 is one of ``SucceedTask``, ``RetryTask``, or ``TaskState``, and it is reported exactly once. + +Task handler parse lifecycle +---------------------------- + +The Dag processor starts the same process to check a Python Dag's ``@task.stub`` tasks against the +``TaskHandler`` registrations that implement them: + +.. code-block:: text + + 1 process start + 2 receive TaskHandlerParseRequest + 3 declare every registered TaskHandler, keyed by dagId, in registration order + 4 send one TaskHandlerParsingResult, and wait for the supervisor's reply to it + 5 exit + +No ``fn`` runs: the answer comes from the registrations alone. It must depend only on the artifact, +never on the request. A declaration states whether the stub task's arguments bind to the handler's +parameters by position or by name, and which values each parameter accepts. An SDK that cannot list +a handler's parameters sends ``null`` for them, and then only the handler's presence is checked. An +artifact that registers no ``TaskHandler`` answers with an empty mapping. diff --git a/task-sdk/docs/ts-bundle-spec.rst b/task-sdk/docs/ts-bundle-spec.rst index 705227efa55b2..03ee63ceea0a6 100644 --- a/task-sdk/docs/ts-bundle-spec.rst +++ b/task-sdk/docs/ts-bundle-spec.rst @@ -24,8 +24,8 @@ This document specifies the bundle format produced by ``airflow-ts-pack`` and co Artifact Name ------------- -A bundle's name must end in ``.min.mjs``. Nothing else about it is significant: the coordinator searches each -configured root recursively and routes on embedded metadata, so one root may hold several differently named bundles. +A bundle's name must end in ``.min.mjs``. Nothing else about it is significant: the coordinator searches its Dag +bundle recursively and routes on embedded metadata, so one Dag bundle may hold several differently named bundles. ``airflow-ts-pack`` writes ``bundle.min.mjs`` by default and accepts ``--outfile`` for any other name ending in that suffix. @@ -153,11 +153,12 @@ documents, and neither is used to execute the bundle. Reader and Selection Algorithm ------------------------------ -For each candidate in ``bundles_root``, the coordinator: +For each candidate in the Dag bundle named by ``task_handler_bundle_name`` (or, when it is unset, the task's own Dag +bundle), the coordinator: -1. Opens it once. Candidates are the files whose name ends in ``.min.mjs``, found by walking each root recursively, - roots in configured order and each directory's entries in sorted order, so selection does not depend on the order - a filesystem returns entries in. Directories are deduplicated by ``(st_dev, st_ino)``, so a symlink loop +1. Opens it once. Candidates are the files whose name ends in ``.min.mjs``, found by walking the Dag bundle + recursively, each directory's entries in sorted order, so selection does not depend on the order a filesystem + returns entries in. Directories are deduplicated by ``(st_dev, st_ino)``, so a symlink loop terminates the walk instead of exhausting the interpreter stack. 2. Reads a bounded first line and decodes the named metadata and code ranges. 3. Reads the bounded metadata line and checks that the declared metadata range matches its physical location. @@ -171,15 +172,14 @@ For each candidate in ``bundles_root``, the coordinator: 9. Resolves the supervisor schema version and selects the first usable match. A missing, unrelated, unreadable, malformed, corrupt, or incompatible earlier candidate does not prevent selection -of a later usable match. When more than one usable bundle declares the same Dag, the first configured match wins. If -none matches, the error identifies the requested Dag, searched roots, and rejected candidates. +of a later usable match. When more than one usable bundle declares the same Dag, the first match in walk order wins. +If none matches, the error identifies the requested Dag, the searched directory, and rejected candidates. -Every ``.min.mjs`` file under a root is therefore opened, and one that is not a usable bundle is named among those -rejected candidates. ``bundles_root`` names directories of deployed Airflow bundles, so unrelated minified modules do -not belong there. +Every ``.min.mjs`` file in the Dag bundle is therefore opened, and one that is not a usable bundle is named among +those rejected candidates. That Dag bundle holds deployed Airflow bundles, so unrelated minified modules do not belong +there. -The coordinator does not cache Dag-to-path routing. It checks root ordering and the current deployed files for each -task selection. It may reuse section digests from a bounded process-local cache when the open file identity, +The coordinator does not cache Dag-to-path routing. It checks the current deployed files for each task selection. It may reuse section digests from a bounded process-local cache when the open file identity, timestamps, size, layout ranges, and declared digests have not changed. Integrity, Authenticity, and Provenance diff --git a/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py index ec6e6419cbab2..b97ed80e8d82e 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py +++ b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py @@ -19,10 +19,9 @@ from __future__ import annotations -import os import pathlib import stat -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any import attrs import structlog @@ -38,24 +37,6 @@ log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators") -class _UnsetArtifactRoots: - """Sentinel distinguishing an omitted artifact-roots option from an empty value.""" - - -ARTIFACT_ROOTS_NOT_CONFIGURED: Final = _UnsetArtifactRoots() - - -def convert_roots( - value: None | os.PathLike[str] | pathlib.Path | list[os.PathLike[str] | pathlib.Path], -) -> list[pathlib.Path]: - """Normalize a coordinator's root-directories kwarg into a list of expanded paths.""" - if value is None: - return [] - if isinstance(value, (str, os.PathLike, pathlib.Path)): - return [pathlib.Path(value).expanduser()] - return [pathlib.Path(v).expanduser() for v in value] - - def walk_files( roots: Iterable[pathlib.Path], *, match: Callable[[pathlib.Path], bool] ) -> Iterator[pathlib.Path]: @@ -103,30 +84,6 @@ def _sorted_children(directory: pathlib.Path) -> list[pathlib.Path]: return [] -def convert_configured_roots( - value: _UnsetArtifactRoots - | None - | os.PathLike[str] - | pathlib.Path - | list[os.PathLike[str] | pathlib.Path], -) -> list[pathlib.Path]: - """Normalize configured roots while rejecting explicitly empty values.""" - if isinstance(value, _UnsetArtifactRoots): - return [] - if isinstance(value, (str, os.PathLike)): - values = [value] - elif value: - values = value - else: - values = [] - if not values or any(not os.fspath(v).strip() for v in values): - raise ValueError( - "Artifact roots must contain at least one path when provided, and each path must be non-empty; " - "omit the option to use the task's Dag bundle." - ) - return convert_roots(values) - - def validate_schema_version(instance, _, value) -> str: """Attrs validator resolving a bundle's supervisor schema version to a known one.""" return get_schema_version_migrator().resolve_version(str(value)) diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index 9e69255ba3d07..6e722a33a86d7 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -28,7 +28,7 @@ from __future__ import annotations import contextlib -import enum +import ctypes import ipaddress import itertools import os @@ -36,24 +36,25 @@ import signal import socket import subprocess +import sys import time -from typing import TYPE_CHECKING, TypeVar, cast +from typing import TYPE_CHECKING, NoReturn, TypeVar, cast import attrs import psutil import structlog from airflow.dag_processing.bundles.base import BundleVersionLock, unpack_bundle_version # noqa: SDK002 -from airflow.dag_processing.bundles.manager import DagBundlesManager # noqa: SDK002 from airflow.sdk.api.datamodels._generated import BundleInfo from airflow.sdk.configuration import conf from airflow.sdk.execution_time.bundles import initialize_ti_bundle from airflow.sdk.execution_time.coordinator import BaseCoordinator +from airflow.sdk.execution_time.schema import get_schema_version_migrator from airflow.sdk.execution_time.supervisor import ActivitySubprocess, NeverRaised, ProcessTracker if TYPE_CHECKING: import pathlib - from collections.abc import Sequence + from collections.abc import Callable, Sequence from structlog.typing import FilteringBoundLogger from typing_extensions import Self @@ -103,21 +104,19 @@ def _socket_address(value: tuple | str) -> tuple[str, int] | None: return host, int(port) -def _connection_owned_by_process_tree( - peer: tuple[str, int], local: tuple[str, int], proc: subprocess.Popen -) -> bool: +def _connection_owned_by_process_tree(peer: tuple[str, int], local: tuple[str, int], pid: int) -> bool: """ Return whether ``peer`` <-> ``local`` is an established connection in the child's process tree. The launched child may itself spawn the process that connects back to the supervisor — a JVM launcher, a shell wrapper, or any runtime that forks a worker — so the connecting peer can legitimately belong to a *descendant* of - ``proc.pid`` rather than ``proc.pid`` itself. Every process in the subtree - rooted at ``proc.pid`` is part of the task and is trusted; a process outside - that subtree (e.g. an unrelated local process racing for the port) is not. + ``pid`` rather than ``pid`` itself. Every process in the subtree rooted at + ``pid`` is part of the task and is trusted; a process outside that subtree + (e.g. an unrelated local process racing for the port) is not. """ try: - root = psutil.Process(proc.pid) + root = psutil.Process(pid) processes = [root, *root.children(recursive=True)] except (psutil.AccessDenied, psutil.NoSuchProcess, psutil.ZombieProcess, OSError): return False @@ -155,13 +154,68 @@ def _is_connection_from_process( return False deadline = time.monotonic() + verify_timeout while True: - if _connection_owned_by_process_tree(peer, local, proc): + if _connection_owned_by_process_tree(peer, local, proc.pid): return True if time.monotonic() >= deadline: return False time.sleep(poll_interval) +def _is_connection_from_pid(conn: socket.socket, pid: int) -> bool: + """ + Return whether the accepted TCP connection comes from ``pid``'s process tree, checking once. + + A connection that is not visible yet is reported as not owned, so the caller retries later + instead of waiting here. + """ + peer = _socket_address(conn.getpeername()) + local = _socket_address(conn.getsockname()) + return peer is not None and local is not None and _connection_owned_by_process_tree(peer, local, pid) + + +def _set_close_on_exec_above_stderr() -> None: + """Mark every file descriptor above 2 close-on-exec, as ``subprocess.Popen(close_fds=True)`` does.""" + fd_dir = "/proc/self/fd" if os.path.isdir("/proc/self/fd") else "/dev/fd" + for name in os.listdir(fd_dir): + if (fd := int(name)) > 2: + # The descriptor listing the directory is already closed. + with contextlib.suppress(OSError): + os.set_inheritable(fd, False) + + +def _build_runtime_env() -> dict[str, str]: + """Return the environment for a language SDK runtime, with the Airflow config it reads resolved.""" + # A language SDK runtime cannot read Airflow's config, so the options it needs are resolved here. + # An option without a config default gets the fallback Python uses. StartupDetails arrives too + # late: logs may already be produced. + return { + **os.environ, + "AIRFLOW__LOGGING__LOGGING_LEVEL": conf.get("logging", "logging_level", fallback="INFO"), + "AIRFLOW__LOGGING__NAMESPACE_LEVELS": conf.get("logging", "namespace_levels", fallback=""), + } + + +# From +_PR_SET_PDEATHSIG = 1 + + +def _set_parent_death_signal() -> None: + """ + On Linux, have the kernel SIGKILL this process when the thread that forked it exits. + + The setting survives the exec, so the runtime dies with the process that reads its sockets, + even when that process is killed without the chance to kill the runtime first. + """ + if sys.platform != "linux": + return + try: + libc = ctypes.CDLL(None, use_errno=True) + if libc.prctl(_PR_SET_PDEATHSIG, signal.SIGKILL, 0, 0, 0) != 0: + log.warning("Failed to set PR_SET_PDEATHSIG", errno=ctypes.get_errno()) + except Exception: + log.warning("Unable to set PR_SET_PDEATHSIG", exc_info=True) + + def _accept_connections( servers: dict[str, socket.socket], drains: dict[str, socket.socket], @@ -312,15 +366,6 @@ def start( # type: ignore[override] stdout_r, stdout_w = tracker.track(*socket.socketpair()) stderr_r, stderr_w = tracker.track(*socket.socketpair()) - # A language SDK runtime cannot read Airflow's config, so propagate the - # resolved log levels via the environment at launch. StartupDetails - # arrives too late, the logs might already be produced by then. - env = { - **os.environ, - "AIRFLOW__LOGGING__LOGGING_LEVEL": conf.get("logging", "logging_level", fallback="INFO"), - "AIRFLOW__LOGGING__NAMESPACE_LEVELS": conf.get("logging", "namespace_levels", fallback=""), - } - proc = subprocess.Popen( [ *command, @@ -329,7 +374,7 @@ def start( # type: ignore[override] ], stdout=stdout_w.fileno(), stderr=stderr_w.fileno(), - env=env, + env=_build_runtime_env(), ) tracker.track(proc) for soc in tracker.untrack(stdout_w, stderr_w): @@ -402,22 +447,6 @@ def _initialize_pinned_bundle(target: BundleInfo, logger: FilteringBoundLogger) return initialize_ti_bundle(BundleInfo(name=target.name, version=version, version_data=version_data)) -class _ArtifactSource(enum.Enum): - """How a subprocess coordinator locates the compiled task artifacts.""" - - EXPLICIT_ROOT = enum.auto() - """An explicit filesystem root (``jars_root`` / ``executables_root`` / ``bundles_root``).""" - NAMED_BUNDLE = enum.auto() - """``dag_bundle_name`` names a configured Dag bundle; its version current at task start is used.""" - TASK_BUNDLE = enum.auto() - """ - Neither is set: artifacts are *co-located* with the Python stub Dag. - - The task's own bundle is scanned, so the compiled artifacts ship in the same - bundle, at the same version, as the Dag that delegates to them. - """ - - @attrs.define(kw_only=True) class SubprocessCoordinator(BaseCoordinator): """ @@ -429,105 +458,52 @@ class SubprocessCoordinator(BaseCoordinator): connections, draining startup output, and tearing everything down on failure — is handled here. + A subclass that also reports its artifacts' task handlers implements + :meth:`_build_parse_task_handler_command`. + :param task_startup_timeout: Maximum time the coordinator waits for the subprocess to connect to both servers, in seconds. The default is 10 seconds. - :param dag_bundle_name: Locate artifacts through a configured Dag bundle rather - than an explicit root. Mutually exclusive with the subclass's explicit root; - if neither is set, the task's own bundle is used. A named bundle resolves to - the version current when the task starts; the task's own bundle uses the - run's version. Either way the resolved version is pinned for the whole task. + :param task_handler_bundle_name: Name of the Dag bundle that holds the compiled + task handlers. It must be registered in ``[dag_processor] dag_bundle_config_list``. + If unset, the task's own Dag bundle is used. A named bundle resolves to the + version current when the task starts; the task's own bundle uses the run's + version. Either way the resolved version is pinned for the whole task. """ task_startup_timeout: float = 10.0 - dag_bundle_name: str | None = None + task_handler_bundle_name: str | None = None - _artifact_source: _ArtifactSource = attrs.field(init=False) - # The subclass's explicit root, recorded at construction so the base can - # resolve roots without knowing the subclass field name. - _configured_roots: list[pathlib.Path] = attrs.field(init=False, factory=list) _active_scan_roots: tuple[pathlib.Path, ...] | None = attrs.field(init=False, default=None) - @property - def _explicit_artifact_roots(self) -> tuple[str, Sequence[pathlib.Path]]: - """ - The subclass's explicit-root kwarg name and its configured value. - - The name is only used in error messages. An empty value — the default, for a - subclass that does not override this — selects task-bundle mode rather than - failing at execute time. - """ - return "root", () - - def __attrs_post_init__(self) -> None: - self._classify_artifact_source() - - def _classify_artifact_source(self) -> None: - """ - Classify and validate how this coordinator locates artifacts (construction time). - - Rejects setting both an explicit root and ``dag_bundle_name``, fails fast - when ``dag_bundle_name`` names a bundle that is not configured, and records - the resulting :class:`_ArtifactSource` and explicit root. - """ - root_kwarg, configured = self._explicit_artifact_roots - if configured and self.dag_bundle_name is not None: - raise ValueError( - f"Set at most one of {root_kwarg!r} or 'dag_bundle_name': {root_kwarg!r} for an " - f"explicit path, 'dag_bundle_name' for a configured Dag bundle, or leave both " - f"unset to scan the task's own bundle." - ) - if configured: - source = _ArtifactSource.EXPLICIT_ROOT - self._configured_roots = list(configured) - elif self.dag_bundle_name is not None: - source = _ArtifactSource.NAMED_BUNDLE - if not DagBundlesManager.is_bundle_configured(self.dag_bundle_name): - raise ValueError( - f"Coordinator 'dag_bundle_name' references unconfigured Dag bundle " - f"{self.dag_bundle_name!r}." - ) - else: - source = _ArtifactSource.TASK_BUNDLE - - self._artifact_source = source - log.debug( - "Coordinator artifact source selected", - mode=source.name, - dag_bundle_name=self.dag_bundle_name, - configured_roots=[str(root) for root in self._configured_roots], - ) - - def _init_root_source( + def _resolve_artifact_bundle( self, bundle_info: BundleInfo, logger: FilteringBoundLogger - ) -> tuple[list[pathlib.Path], BaseDagBundle | None]: + ) -> BaseDagBundle: """ - Resolve the directories to scan for artifacts, dispatched on the classified mode. + Materialize the Dag bundle holding the artifacts for a task of *bundle_info*. - Returns ``(roots, bundle)``: an explicit root yields no bundle (``None``); - a Dag-bundle mode returns the materialized path and the resolved bundle so - :meth:`execute_task` can hold a version lock over it. *logger* is the task - logger, so materialization failures surface in the task log. + That is the bundle named by ``task_handler_bundle_name``, or the task's own + bundle when it is unset. *logger* is the task logger, so materialization + failures surface in the task log. """ - if self._artifact_source is _ArtifactSource.EXPLICIT_ROOT: - return self._configured_roots, None - - if self._artifact_source is _ArtifactSource.NAMED_BUNDLE: - # NAMED_BUNDLE implies dag_bundle_name is set. - target = BundleInfo(name=cast("str", self.dag_bundle_name)) - else: + if self.task_handler_bundle_name is None: target = bundle_info + else: + target = BundleInfo(name=self.task_handler_bundle_name) bundle = _initialize_pinned_bundle(target, logger) path = bundle.path if not path.exists(): raise FileNotFoundError(f"Dag bundle {target.name!r} resolved to {path}, which does not exist.") - return [path], bundle + return bundle def _get_scan_roots(self) -> tuple[pathlib.Path, ...]: - """Return the artifact roots resolved for the active task.""" + """Return the artifact roots resolved for the active task or parse.""" if self._active_scan_roots is None: - raise RuntimeError("_get_scan_roots requires an active task; call it during execute_task.") + raise RuntimeError( + "_get_scan_roots requires an active task or parse; " + "call it during execute_task or parse_task_handler." + ) return self._active_scan_roots def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: @@ -544,11 +520,70 @@ def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], """ raise NotImplementedError + def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: + """ + Build the command that reports the task handlers of the artifact at *path*, with its schema version. + + Subclasses can retrieve the directories to scan for artifacts with + :meth:`_get_scan_roots`; for a parse they are the root of the Dag bundle holding + *path*. The contract is that of :meth:`_build_execute_task_command`: *command* + MUST NOT include the ``--comm`` / ``--logs`` flags. + """ + raise NotImplementedError(f"{type(self).__name__} does not parse task handlers") + + def parse_task_handler( + self, + *, + path: pathlib.Path, + bundle_path: pathlib.Path, + comm_address: tuple[str, int], + logs_address: tuple[str, int], + report_schema_version: Callable[[str | None], None], + ) -> NoReturn: + """ + Replace the current process with the runtime that reports the task handlers of the artifact at *path*. + + Call this in a child process whose standard streams are already set up; the runtime + inherits them and no other file descriptor. The command and its supervisor wire-schema + version are resolved against *bundle_path*, and the version is passed to + *report_schema_version* just before the exec. The runtime connects back to + *comm_address* and *logs_address*. On Linux the runtime dies with the thread that forked + the calling process, and it is not started once that parent has exited. + + :raises Exception: when the command cannot be resolved or started, or the parent has + exited; the process is then unchanged. + """ + with self._set_scan_roots([bundle_path]): + command, schema_version = self._build_parse_task_handler_command(path=path) + if schema_version is not None: + get_schema_version_migrator().resolve_version(schema_version) + argv = [ + *command, + f"--comm={comm_address[0]}:{comm_address[1]}", + f"--logs={logs_address[0]}:{logs_address[1]}", + ] + # Taken before the report, which fails if the parent has already exited. A different parent once the + # death signal is set means the parent exited in between, and the kernel does not send it for that. + parent_pid = os.getppid() + report_schema_version(schema_version) + # Python ignores these at startup and exec keeps ignored signals; subprocess.Popen resets + # them the same way for the task runtime. + for name in ("SIGPIPE", "SIGXFZ", "SIGXFSZ"): + if (sig := getattr(signal, name, None)) is not None: + signal.signal(sig, signal.SIG_DFL) + _set_parent_death_signal() + if os.getppid() != parent_pid: + raise RuntimeError("The process that started the runtime has exited") + _set_close_on_exec_above_stderr() + os.execvpe(argv[0], argv, _build_runtime_env()) + @contextlib.contextmanager def _set_scan_roots(self, roots: Sequence[pathlib.Path]): - """Expose *roots* to the command builder for the duration of the task.""" + """Expose *roots* to the command builder for the duration of the task or parse.""" if self._active_scan_roots is not None: - raise RuntimeError("SubprocessCoordinator.execute_task is not re-entrant.") + raise RuntimeError( + "SubprocessCoordinator.execute_task and parse_task_handler are not re-entrant." + ) self._active_scan_roots = tuple(roots) try: yield @@ -567,20 +602,14 @@ def execute_task( subprocess_logs_to_stdout: bool, **kwargs, ) -> BaseCoordinator.ExecutionResult: - task_logger = logger or log - with contextlib.ExitStack() as stack: - roots, resolved_bundle = self._init_root_source(bundle_info, task_logger) - if resolved_bundle is not None: - # Hold the version lock across start()/wait() so bundle cleanup - # cannot rmtree a version this task is still reading from, - # mirroring task_runner.main() for the Python task path. - stack.enter_context( - BundleVersionLock( - bundle_name=resolved_bundle.name, - bundle_version=resolved_bundle.version, - ) - ) - stack.enter_context(self._set_scan_roots(roots)) + bundle = self._resolve_artifact_bundle(bundle_info, logger or log) + # Hold the version lock across start()/wait() so bundle cleanup cannot + # rmtree a version this task is still reading from, mirroring + # task_runner.main() for the Python task path. + with ( + BundleVersionLock(bundle_name=bundle.name, bundle_version=bundle.version), + self._set_scan_roots([bundle.path]), + ): command, subprocess_schema_version = self._build_execute_task_command(what=what) process = _PopenActivitySubprocess.start( what=what, diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py index 9d6b80b249f81..9ca3b7c5d40db 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -31,9 +31,7 @@ import structlog from airflow.sdk.coordinators._bundle_metadata import ( - ARTIFACT_ROOTS_NOT_CONFIGURED, ResolvedBundle, - convert_configured_roots, extract_supervisor_schema_version, parse_metadata_mapping, ) @@ -247,6 +245,35 @@ def _read_bundle_metadata(path: pathlib.Path) -> dict[str, Any] | None: return None +def read_cache_digest(path: pathlib.Path) -> str | None: + """ + Return the cache digest stored in the metadata of the executable bundle at *path*. + + The digest changes whenever any region of the bundle changes, so the Dag processor can + compare it with the one it saw when it last probed the bundle. It is read, not computed: + the binary region is not hashed, so a match does not prove the bundle is intact. + + Returns ``None`` when the file is not a readable bundle or its metadata has no cache digest, + and the caller should treat the bundle as changed. + """ + try: + with open(path, "rb") as f: + footer = _Footer.read(f, path, os.fstat(f.fileno()).st_size) + if footer is None: + return None + f.seek(footer.metadata_start) + metadata_bytes = f.read(footer.metadata_len) + metadata = parse_metadata_mapping(metadata_bytes, source="bundle metadata") + except (OSError, ValueError) as exc: + log.debug("Cannot read the bundle's cache digest", path=str(path), error=str(exc)) + return None + digests = metadata.get("digests") + if not isinstance(digests, dict): + return None + cache = digests.get("cache") + return cache if isinstance(cache, str) and cache else None + + def _dag_ids(metadata: dict[str, Any]) -> set[str]: dags = metadata.get("dags") if not isinstance(dags, dict): @@ -294,10 +321,10 @@ def _walk_executables( @attrs.define class _Bundle(ResolvedBundle): @classmethod - def find(cls, executables_root: Sequence[pathlib.Path], dag_id: str) -> Self: - log.debug("Finding executable bundles recursively", roots=executables_root) + def find(cls, roots: Sequence[pathlib.Path], dag_id: str) -> Self: + log.debug("Finding executable bundles recursively", roots=roots) rejected: list[tuple[pathlib.Path, str]] = [] - for p in _find_executables(executables_root): + for p in _find_executables(roots): if (metadata := _read_bundle_metadata(p)) is None: continue if dag_id not in _dag_ids(metadata): @@ -310,7 +337,7 @@ def find(cls, executables_root: Sequence[pathlib.Path], dag_id: str) -> Self: rejected.append((p.resolve(), str(exc))) continue - resolved_paths = os.pathsep.join(str(r.resolve()) for r in executables_root) + resolved_paths = os.pathsep.join(str(r.resolve()) for r in roots) if rejected: details = "; ".join(f"{path}: {reason}" for path, reason in rejected) tp = ( @@ -328,34 +355,39 @@ class ExecutableCoordinator(SubprocessCoordinator): """ Coordinator that launches a native executable subprocess for task execution. + The Dag processor also launches an executable bundle to ask which task handlers + it registers for the Dags of a Python file. + Configuration is taken from the ``[sdk] coordinators`` entry that constructs this instance:: "go": { "classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": { - "executables_root": ["~/airflow/executable-bundles"] + "task_handler_bundle_name": "go-task-handlers" } } - :param executables_root: A list of directories scanned for executable - bundles when a Python stub DAG delegates task execution to a native - runtime. See :class:`SubprocessCoordinator` for its interaction with - ``dag_bundle_name``. + :param task_handler_bundle_name: Name of the Dag bundle holding the + executable bundles a Python stub Dag delegates task execution to. It must + be registered in ``[dag_processor] dag_bundle_config_list``. If unset, + the task's own Dag bundle is used. Only files with the executable bit set + are considered. :param task_startup_timeout: Maximum time the coordinator waits for a task process to start, in seconds. The default is 10 seconds. """ - executables_root: list[pathlib.Path] = attrs.field( - default=ARTIFACT_ROOTS_NOT_CONFIGURED, - converter=convert_configured_roots, - ) - - @property - def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: - return "executables_root", self.executables_root - def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: roots = self._get_scan_roots() bundle = _Bundle.find(roots, what.dag_id) return [str(bundle.path)], bundle.schema_version + + def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: + # The same trailer and digest check a task gets before its bundle runs. + if (metadata := _read_bundle_metadata(path)) is None: + raise ValueError( + f"{path} is not a valid executable bundle: it cannot be read, has no AFBNDL01 " + "trailer, or its binary digest or metadata is invalid" + ) + # Absolute, as for a task, so exec never searches PATH for it. + return [os.fspath(path.resolve())], extract_supervisor_schema_version(metadata) diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index 1f0f6e6efd5f9..05cbad27b828d 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -19,9 +19,9 @@ from __future__ import annotations -import email import os import pathlib +import re import stat import zipfile from typing import TYPE_CHECKING @@ -29,11 +29,7 @@ import attrs import structlog -from airflow.sdk.coordinators._bundle_metadata import ( - ARTIFACT_ROOTS_NOT_CONFIGURED, - convert_configured_roots, - validate_schema_version, -) +from airflow.sdk.coordinators._bundle_metadata import validate_schema_version from airflow.sdk.coordinators._subprocess import SubprocessCoordinator if TYPE_CHECKING: @@ -86,15 +82,39 @@ def _iter_dir(directory: pathlib.Path) -> Iterator[pathlib.Path]: return -def _calculate_classpath(jars_root: Sequence[pathlib.Path]) -> str: - jars = (p.as_posix() for p in _find_jars(jars_root)) +def _calculate_classpath(roots: Sequence[pathlib.Path]) -> str: + jars = (p.as_posix() for p in _find_jars(roots)) return os.pathsep.join(sorted(jars)) # Keep output deterministic. +def _parse_manifest(data: bytes) -> dict[str, str]: + """ + Return the main-section attributes of a JAR manifest, keyed by lowercased name. + + A line holds at most 72 bytes, so a longer value continues on lines starting with one space, + which are joined without it. The fold can split a multi-byte character, so values are decoded + only once joined. + """ + attributes: dict[str, bytes] = {} + name: str | None = None + for line in re.split(rb"\r\n|\r|\n", data): + if not line: + break # The main section ends at the first blank line. + if line.startswith(b" "): + if name is not None: + attributes[name] += line[1:] + continue + key, sep, value = line.partition(b":") + name = key.decode("utf-8", errors="replace").lower() if sep else None + if name is not None: + attributes[name] = value.removeprefix(b" ") + return {key: value.decode("utf-8", errors="replace") for key, value in attributes.items()} + + @attrs.define class _JarMetadata: - main_class: str - schema_version: str + main_class: str | None + schema_version: str | None @classmethod def from_jar(cls, path: pathlib.Path) -> Self | None: @@ -105,9 +125,8 @@ def from_jar(cls, path: pathlib.Path) -> Self | None: except KeyError: log.debug("JAR does not contain META-INF/MANIFEST.MF; ignored", path=path) return None - with zf.open(manifest_info) as f: - manifest = email.message_from_binary_file(f) - return cls(manifest["Main-Class"], manifest["Airflow-Supervisor-Schema-Version"]) + manifest = _parse_manifest(zf.read(manifest_info)) + return cls(manifest.get("main-class"), manifest.get("airflow-supervisor-schema-version")) except zipfile.BadZipFile: log.exception("Cannot read JAR; ignored", path=path) return None @@ -167,7 +186,7 @@ class JavaCoordinator(SubprocessCoordinator): "jdk-17": { "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": { - "jars_root": ["~/airflow/jars"], + "task_handler_bundle_name": "java-task-handlers", "java_executable": "/usr/lib/jvm/java-17-openjdk/bin/java", "jvm_args": ["-Xmx1024m"] } @@ -176,16 +195,26 @@ class JavaCoordinator(SubprocessCoordinator): :param java_executable: Path to the ``java`` command (defaults to ``"java"``, which relies on ``$PATH``). :param jvm_args: Extra arguments passed to the JVM (e.g. ``["-Xmx512m"]``). - :param jars_root: A list of directories scanned for JAR bundles. See - :class:`SubprocessCoordinator` for its interaction with ``dag_bundle_name``. + :param task_handler_bundle_name: Name of the Dag bundle holding the JARs. It + must be registered in ``[dag_processor] dag_bundle_config_list``. If + unset, the task's own Dag bundle is used. :param main_class: Explicit entry point to execute with *java_executable*. :param task_startup_timeout: Maximum time the coordinator waits for a task process to start, in seconds. The default is 10 seconds. - If *main_class* is not explicitly set, JavaCoordinator scans *jars_root* to + Every JAR in the bundle goes on the classpath, so one bundle is one + classpath. Handlers that need conflicting dependency versions belong in + separate bundles, each served by its own coordinator and queue. + + If *main_class* is not explicitly set, JavaCoordinator scans the bundle to find an executable JAR (one with Main-Class set in its metadata). If more than one executable JAR is found, it may be nondeterministic which one ends - up being executed. + up being executed, so set *main_class* when more than one JAR in the bundle + declares Main-Class. + + To report the task handlers a JAR registers, the coordinator runs that JAR's + own Main-Class with the whole bundle on the classpath. *main_class* applies + to task execution only. A JAR containing metadata *Airflow-Supervisor-Schema-Version* should also be available to specify the wire schema version. The JAR containing the Java @@ -203,26 +232,25 @@ class JavaCoordinator(SubprocessCoordinator): java_executable: str = "java" jvm_args: list[str] = attrs.field(factory=list) - jars_root: list[pathlib.Path] = attrs.field( - default=ARTIFACT_ROOTS_NOT_CONFIGURED, - converter=convert_configured_roots, - ) main_class: str = "" - @property - def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: - return "jars_root", self.jars_root - def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: # Without main_class, the first executable JAR in walk order wins; tracked at # https://github.com/apache/airflow/issues/71134 roots = self._get_scan_roots() jar = _JarInfo.find(roots, self.main_class) - command = [ - self.java_executable, - "-classpath", - _calculate_classpath(roots), - *self.jvm_args, - jar.main_class, - ] - return command, jar.schema_version + return self._build_java_command(roots, jar.main_class), jar.schema_version + + def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: + metadata = _JarMetadata.from_jar(path) + if metadata is None: + raise ValueError(f"{path} is not an executable JAR: it has no readable manifest") + if not metadata.main_class: + raise ValueError(f"{path} is not an executable JAR: its manifest sets no Main-Class") + roots = self._get_scan_roots() + # A thin JAR leaves the version to the airflow-sdk JAR beside it, where execution finds it too. + schema_version = metadata.schema_version or _JarInfo.find(roots, metadata.main_class).schema_version + return self._build_java_command(roots, metadata.main_class), schema_version + + def _build_java_command(self, roots: Sequence[pathlib.Path], main_class: str) -> list[str]: + return [self.java_executable, "-classpath", _calculate_classpath(roots), *self.jvm_args, main_class] diff --git a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py index dd2a3f6f0416f..1644cf425082a 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py @@ -26,12 +26,7 @@ import attrs import structlog -from airflow.sdk.coordinators._bundle_metadata import ( - ARTIFACT_ROOTS_NOT_CONFIGURED, - ResolvedBundle, - convert_configured_roots, - walk_files, -) +from airflow.sdk.coordinators._bundle_metadata import ResolvedBundle, walk_files from airflow.sdk.coordinators._subprocess import SubprocessCoordinator from airflow.sdk.coordinators.node._bundle_reader import read_bundle @@ -55,11 +50,11 @@ def _is_bundle(path: pathlib.Path) -> bool: @attrs.define class _Bundle(ResolvedBundle): @classmethod - def find(cls, bundles_root: Sequence[pathlib.Path], dag_id: str) -> Self: + def find(cls, roots: Sequence[pathlib.Path], dag_id: str) -> Self: """Return the first verified configured bundle that declares *dag_id*.""" - log.debug("Finding TypeScript bundles recursively", roots=bundles_root, dag_id=dag_id) + log.debug("Finding TypeScript bundles recursively", roots=roots, dag_id=dag_id) rejected: list[tuple[pathlib.Path, str]] = [] - for candidate in walk_files(bundles_root, match=_is_bundle): + for candidate in walk_files(roots, match=_is_bundle): try: metadata = read_bundle(candidate) if dag_id not in metadata.dag_ids: @@ -85,7 +80,7 @@ def find(cls, bundles_root: Sequence[pathlib.Path], dag_id: str) -> Self: log.debug("Selected TypeScript bundle", path=candidate, dag_id=dag_id) return bundle - searched = os.pathsep.join(os.fspath(root) for root in bundles_root) + searched = os.pathsep.join(os.fspath(root) for root in roots) if rejected: details = "; ".join(f"{path}: {reason}" for path, reason in rejected) raise FileNotFoundError( @@ -100,6 +95,8 @@ class NodeCoordinator(SubprocessCoordinator): """ Coordinator that launches a Node.js subprocess for task execution. + It also launches a verified ``*.min.mjs`` bundle to report the task handlers it registers. + Configuration is taken from the ``[sdk] coordinators`` entry that constructs this instance:: @@ -108,31 +105,27 @@ class NodeCoordinator(SubprocessCoordinator): "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", "kwargs": { "node_executable": "node", - "bundles_root": ["/opt/airflow/ts-bundles"], + "task_handler_bundle_name": "ts-task-handlers", }, } } :param node_executable: Path to the ``node`` binary (defaults to ``"node"``, which relies on ``$PATH``). - :param bundles_root: Directories searched recursively, in order, for the first verified - ``*.min.mjs`` bundle declaring the task instance's Dag. See :class:`SubprocessCoordinator` - for its interaction with ``dag_bundle_name``. + :param task_handler_bundle_name: Name of the Dag bundle searched recursively for the first + verified ``*.min.mjs`` bundle declaring the task instance's Dag. It must be registered in + ``[dag_processor] dag_bundle_config_list``. If unset, the task's own Dag bundle is used. :param task_startup_timeout: Maximum time the coordinator waits for a task process to start, in seconds. The default is 10 seconds. """ node_executable: str = "node" - bundles_root: list[pathlib.Path] = attrs.field( - default=ARTIFACT_ROOTS_NOT_CONFIGURED, - converter=convert_configured_roots, - ) - - @property - def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: - return "bundles_root", self.bundles_root def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: roots = self._get_scan_roots() bundle = _Bundle.find(roots, what.dag_id) return [self.node_executable, os.fspath(bundle.path)], bundle.schema_version + + def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: + metadata = read_bundle(path) + return [self.node_executable, os.fspath(path)], metadata.supervisor_schema_version diff --git a/task-sdk/src/airflow/sdk/execution_time/coordinator.py b/task-sdk/src/airflow/sdk/execution_time/coordinator.py index 1469b27888c51..f3d37858503b9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/coordinator.py +++ b/task-sdk/src/airflow/sdk/execution_time/coordinator.py @@ -42,17 +42,19 @@ import functools import os import signal -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, NoReturn import attrs import pydantic import structlog +from airflow.dag_processing.bundles.manager import DagBundlesManager # noqa: SDK002 from airflow.sdk._shared.module_loading import import_string from airflow.sdk.configuration import conf if TYPE_CHECKING: - from collections.abc import Generator, Mapping + import pathlib + from collections.abc import Callable, Generator, Mapping from os import PathLike from structlog.typing import FilteringBoundLogger @@ -107,6 +109,30 @@ def execute_task( """ raise NotImplementedError + def parse_task_handler( + self, + *, + path: pathlib.Path, + bundle_path: pathlib.Path, + comm_address: tuple[str, int], + logs_address: tuple[str, int], + report_schema_version: Callable[[str | None], None], + ) -> NoReturn: + """ + Replace the current process with the runtime that reports the task handlers of the artifact at *path*. + + Call this in a child process whose standard streams are already set up; the runtime inherits + them and no other file descriptor. *bundle_path* is the root of the Dag bundle holding *path*. + The runtime connects back to *comm_address* and *logs_address* and answers a + ``TaskHandlerParseRequest``. Its supervisor wire-schema version is passed to + *report_schema_version* just before the exec. + + A coordinator opts in by overriding this; the default raises :class:`NotImplementedError`. + + :raises Exception: when the runtime cannot be resolved or started. + """ + raise NotImplementedError(f"{type(self).__name__} does not parse task handlers") + class _CoordinatorSpec(pydantic.BaseModel): classpath: str @@ -253,24 +279,31 @@ class CoordinatorManager: @classmethod def from_config(cls) -> Self: - """Load coordinator specs from configuration without initialization.""" + """ + Load coordinator specs from configuration without initialization. + + Every ``queue_to_coordinator`` key and the ``task_handler_bundle_name`` of + every routed coordinator are validated here, so a typo fails at config load + rather than on the first task routed to the coordinator. + """ coordinator_specs = { k: _CoordinatorSpec.model_validate(v) for k, v in conf.getjson("sdk", "coordinators", fallback={}).items() } queue_to_coordinator = conf.getjson("sdk", "queue_to_coordinator", fallback={}) - for key in queue_to_coordinator.values(): - if key not in coordinator_specs: + for key in set(queue_to_coordinator.values()): + if (spec := coordinator_specs.get(key)) is None: raise ValueError(f"[sdk] queue_to_coordinator references invalid coordinator key: {key!r}") + bundle_name = spec.kwargs.get("task_handler_bundle_name") + if bundle_name is not None and ( + not isinstance(bundle_name, str) or not DagBundlesManager.is_bundle_configured(bundle_name) + ): + raise InvalidCoordinatorError( + f"[sdk] coordinators {key!r} sets task_handler_bundle_name={bundle_name!r}, " + f"which is not a bundle in [dag_processor] dag_bundle_config_list" + ) return cls(coordinator_specs=coordinator_specs, queue_to_coordinator=queue_to_coordinator) - def _find_queue(self, key: str) -> BaseCoordinator: - with contextlib.suppress(KeyError): - return self._created_coordinators[key] - spec = self._coordinator_specs[key] - coordinator = self._created_coordinators[key] = import_string(spec.classpath)(**spec.kwargs) - return coordinator - def for_queue(self, queue: str) -> BaseCoordinator: """ Find the coordinator for *queue*. @@ -282,15 +315,32 @@ def for_queue(self, queue: str) -> BaseCoordinator: except KeyError: log.debug("Queue not configured to a coordinator; defaulting to Python", queue=queue) return _build_python_coordinator() + if key not in self._coordinator_specs: + raise InvalidCoordinatorError(f"Queue {queue!r} configured to nonexistent coordinator") + coordinator = self.get_coordinator(key) + log.debug("Coordinator found for queue", coordinator=coordinator, queue=queue) + return coordinator + + def get_coordinator(self, key: str) -> BaseCoordinator: + """ + Return the coordinator configured under *key* in ``[sdk] coordinators``, building it on first use. + + :raises InvalidCoordinatorError: when *key* is not configured, or its coordinator cannot be + imported or built. + """ + with contextlib.suppress(KeyError): + return self._created_coordinators[key] try: - coordinator = self._find_queue(key) + spec = self._coordinator_specs[key] except KeyError: - raise InvalidCoordinatorError(f"Queue {queue!r} configured to nonexistent coordinator") + raise InvalidCoordinatorError(f"No coordinator {key!r} in [sdk] coordinators") + try: + coordinator = import_string(spec.classpath)(**spec.kwargs) except ImportError: raise InvalidCoordinatorError(f"Cannot import coordinator {key!r}") except TypeError: raise InvalidCoordinatorError(f"Cannot instantiate coordinator {key!r}") - log.debug("Coordinator found for queue", coordinator=coordinator, queue=queue) + self._created_coordinators[key] = coordinator return coordinator def extra_for_queue(self, queue: str) -> dict[str, Any] | None: @@ -307,6 +357,21 @@ def extra_for_queue(self, queue: str) -> dict[str, Any] | None: return None return spec.extra + def get_task_handler_bundle_names(self) -> dict[str, str | None]: + """ + Return the ``task_handler_bundle_name`` of each coordinator a queue routes to, keyed by coordinator. + + ``None`` means the coordinator reads task handlers from the task's own Dag bundle. + A coordinator no queue routes to runs no stub task, so it is left out and its bundle is not read. + Only the declarative specs are read; no coordinator is instantiated. + """ + routed_keys = set(self._queue_to_coordinator.values()) + return { + key: spec.kwargs.get("task_handler_bundle_name") + for key, spec in self._coordinator_specs.items() + if key in routed_keys + } + @functools.cache def get_coordinator_manager() -> CoordinatorManager: diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md b/task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md index 84c821bad2bbc..362e628eda531 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md +++ b/task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md @@ -15,10 +15,10 @@ models stay in their semantic homes: `ToTask` and `ToSupervisor` discriminated unions in `airflow.sdk.execution_time.comms`. - Dag-processing channel (manager ↔ parser-supervisor): the - `ToManager` and `ToDagProcessor` discriminated unions in - `airflow.dag_processing.processor`. + `ToManager`, `ToDagProcessor` and `ToSDKTaskHandlerProcessor` + discriminated unions in `airflow.dag_processing.processor`. -`registered_models_by_name()` introspects those four unions on first +`registered_models_by_name()` introspects those five unions on first call, so the snapshot the prek hook commits to `schema.json` always matches the exact set of classes `CommsDecoder` actually decodes against — there is no hand-maintained list to keep in sync. The @@ -58,12 +58,19 @@ version bump, and vice versa. Append the class to the relevant discriminated union in its semantic home — `ToTask` / `ToSupervisor` in `comms.py`, or `ToManager` / -`ToDagProcessor` in `processor.py`. That is the *only* registration -step; `registered_models_by_name()` picks it up automatically the next -time the snapshot hook runs. - -No `VersionChange` entry is required on the first introduction — the -head shape *is* the schema for the new body. +`ToDagProcessor` / `ToSDKTaskHandlerProcessor` in `processor.py`. That +is the *only* registration step; `registered_models_by_name()` picks it +up automatically the next time the snapshot hook runs. + +No `VersionChange` entry is required on the first introduction: a +lang-SDK pinned to an older version never sends the new body, so there +is nothing to strip on the way down. Only changes to bodies that +already existed at an earlier dated version need one. + +`check-supervisor-schemas-versions` still fails locally on a commit +that changes the snapshot without touching anything under `versions/`. +CI's `--all-files` run passes, because it hands the hook the +`versions/` files too. ### Adding a field to a registered body diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/__init__.py index 585bc84a33e2d..d2861f1665bc9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/__init__.py @@ -32,11 +32,12 @@ :func:`registered_models_by_name` resolves a wire-shape ``type`` discriminator to the head Pydantic class. It is computed dynamically -from the four discriminated unions ``ToTask``, ``ToSupervisor`` -(task-execution channel) and ``ToManager``, ``ToDagProcessor`` -(dag-processing channel) so the registry is always in sync with the -actual unions ``CommsDecoder`` decodes against -- no hand-maintained -list to drift. Triggerer unions are intentionally excluded (the +from the five discriminated unions ``ToTask``, ``ToSupervisor`` +(task-execution channel) and ``ToManager``, ``ToDagProcessor``, +``ToSDKTaskHandlerProcessor`` (dag-processing channel) so the registry +is always in sync with the actual unions ``CommsDecoder`` decodes +against -- no hand-maintained list to drift. Triggerer unions are +intentionally excluded (the Triggerer channel is not handled by lang-SDK coordinators today). """ @@ -75,7 +76,7 @@ def registered_models_by_name() -> dict[str, type[BaseModel]]: Map every supervisor schema body's class name to the head Pydantic class. Single source of truth for the registry. Built once by walking the - four discriminated unions the supervisor decodes against; cached + five discriminated unions the supervisor decodes against; cached per-process because the registry only changes when a union member is added in ``comms.py`` or ``processor.py`` (which needs a restart anyway). :func:`resolve_body_class` looks up the wire-shape @@ -91,11 +92,11 @@ def registered_models_by_name() -> dict[str, type[BaseModel]]: single head class, so a name clash is a programmer error that must surface immediately rather than silently picking a winner. """ - from airflow.dag_processing.processor import ToDagProcessor, ToManager + from airflow.dag_processing.processor import ToDagProcessor, ToManager, ToSDKTaskHandlerProcessor from airflow.sdk.execution_time.comms import ToSupervisor, ToTask by_name: dict[str, type[BaseModel]] = {} - for source in (ToTask, ToSupervisor, ToManager, ToDagProcessor): + for source in (ToTask, ToSupervisor, ToManager, ToDagProcessor, ToSDKTaskHandlerProcessor): for model in _iter_model_types(source): existing = by_name.get(model.__name__) if existing is None: diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index d6ee95927efee..e7a512de7b574 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -958,6 +958,13 @@ "title": "Callback Requests", "type": "array" }, + "known_artifacts": { + "items": { + "$ref": "#/$defs/TaskHandlerArtifact" + }, + "title": "Known Artifacts", + "type": "array" + }, "type": { "const": "DagFileParseRequest", "default": "DagFileParseRequest", @@ -4156,6 +4163,220 @@ "title": "TaskCallbackRequest", "type": "object" }, + "TaskHandlerArtifact": { + "description": "A Lang-SDK artifact and every task handler it registers.", + "properties": { + "bundle_name": { + "title": "Bundle Name", + "type": "string" + }, + "relative_fileloc": { + "maxLength": 2000, + "title": "Relative Fileloc", + "type": "string" + }, + "size_bytes": { + "title": "Size Bytes", + "type": "integer" + }, + "cache_digest": { + "anyOf": [ + { + "maxLength": 128, + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Cache Digest" + }, + "task_handlers": { + "additionalProperties": { + "items": { + "$ref": "#/$defs/TaskHandlerDeclaration" + }, + "type": "array" + }, + "title": "Task Handlers", + "type": "object" + } + }, + "required": [ + "bundle_name", + "relative_fileloc", + "size_bytes", + "cache_digest", + "task_handlers" + ], + "title": "TaskHandlerArtifact", + "type": "object" + }, + "TaskHandlerDeclaration": { + "description": "A task handler that a Lang-SDK artifact registers for one task.", + "properties": { + "task_id": { + "title": "Task Id", + "type": "string" + }, + "binding": { + "enum": [ + "positional", + "named" + ], + "title": "Binding", + "type": "string" + }, + "params": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/TaskHandlerParam" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Task Handler Params" + } + }, + "required": [ + "task_id", + "binding", + "params" + ], + "title": "TaskHandlerDeclaration", + "type": "object" + }, + "TaskHandlerParam": { + "description": "One parameter of a task handler.", + "properties": { + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Name" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "exact_name": { + "default": false, + "title": "Exact Name", + "type": "boolean" + } + }, + "required": [ + "name" + ], + "title": "TaskHandlerParam", + "type": "object" + }, + "TaskHandlerParseRequest": { + "description": "Request for Task Handler Parsing.\n\nAsks a Lang-SDK runtime for every task handler an artifact registers.", + "properties": { + "file": { + "title": "File", + "type": "string" + }, + "bundle_path": { + "format": "path", + "title": "Bundle Path", + "type": "string" + }, + "bundle_name": { + "title": "Bundle Name", + "type": "string" + }, + "type": { + "const": "TaskHandlerParseRequest", + "default": "TaskHandlerParseRequest", + "title": "Type", + "type": "string" + } + }, + "required": [ + "file", + "bundle_path", + "bundle_name" + ], + "title": "TaskHandlerParseRequest", + "type": "object" + }, + "TaskHandlerParsingResult": { + "description": "Result of Task Handler Parsing.\n\nEvery task handler a Lang-SDK artifact registers, keyed by Dag id.\n\nThe answer depends only on the artifact, never on the request.", + "properties": { + "fileloc": { + "title": "Fileloc", + "type": "string" + }, + "task_handlers": { + "additionalProperties": { + "items": { + "$ref": "#/$defs/TaskHandlerDeclaration" + }, + "type": "array" + }, + "title": "Task Handlers", + "type": "object" + }, + "import_errors": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Import Errors" + }, + "warnings": { + "anyOf": [ + { + "items": {}, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Warnings" + }, + "type": { + "const": "TaskHandlerParsingResult", + "default": "TaskHandlerParsingResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "fileloc", + "task_handlers" + ], + "title": "TaskHandlerParsingResult", + "type": "object" + }, "TaskInstanceState": { "description": "All possible states that a Task Instance can be in.\n\nNote that None is also allowed, so always use this in a type hint with Optional.", "enum": [ diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index c974f797bb0ad..aec4e0874b728 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -40,6 +40,7 @@ def get_bundle() -> VersionBundle: from airflow.sdk.execution_time.schema.versions.v2026_10_30 import ( AddArgBindingsToSupervisorTIRunContext, AddDagDefinitionsToDagFileParsingResult, + AddKnownArtifactsToDagFileParseRequest, AddRetryReasonToTaskState, ) @@ -50,6 +51,7 @@ def get_bundle() -> VersionBundle: AddArgBindingsToSupervisorTIRunContext, AddRetryReasonToTaskState, AddDagDefinitionsToDagFileParsingResult, + AddKnownArtifactsToDagFileParseRequest, ), Version("2026-06-16"), ) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py index 4aca18cf50cc0..3b810cbaaa254 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py @@ -28,7 +28,7 @@ from cadwyn import VersionChange, schema -from airflow.dag_processing.processor import DagFileParsingResult # noqa: SDK002 +from airflow.dag_processing.processor import DagFileParseRequest, DagFileParsingResult # noqa: SDK002 from airflow.sdk.api.datamodels._generated import TIRunContext from airflow.sdk.execution_time.comms import TaskState @@ -64,3 +64,13 @@ class AddDagDefinitionsToDagFileParsingResult(VersionChange): schema(DagFileParsingResult).field("parsed_definitions").didnt_exist, schema(DagFileParsingResult).field("dag_source_codes").didnt_exist, ) + + +class AddKnownArtifactsToDagFileParseRequest(VersionChange): + """Add ``known_artifacts``, the recorded task-handler artifacts and their answers, to ``DagFileParseRequest``.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + schema(DagFileParseRequest).field("known_artifacts").didnt_exist, + ) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index f2373b3bb835e..ecb3799bfb34a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -390,8 +390,9 @@ def _reopen_std_io_handles(child_stdin, child_stdout, child_stderr): del sock # We open the socket/fd as binary, and then pass it to a TextIOWrapper so that it looks more like a - # normal sys.stdout etc. - binary = os.fdopen(fd, mode + "b") + # normal sys.stdout etc. The handle does not own the fd: a child this process supervises replaces + # it after dup2 onto the same fd, and closing it there would close the child's new stream. + binary = os.fdopen(fd, mode + "b", closefd=False) handle = io.TextIOWrapper(binary, line_buffering=True) setattr(sys, handle_name, handle) diff --git a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py index c20913e26eae3..ffe4e432dd646 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -18,7 +18,7 @@ from __future__ import annotations import hashlib -import pathlib +import json import socket import stat import struct @@ -39,10 +39,12 @@ _BinaryDigestCache, _Bundle, _digest_cache, + read_cache_digest, ) from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess +from tests_common.test_utils.config import conf_vars from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS if not AIRFLOW_V_3_3_PLUS: @@ -361,46 +363,6 @@ def test_logs_when_metadata_is_not_a_mapping(self, tmp_path): class TestExecutableCoordinatorAttributes: - def test_executables_root_accepts_single_path(self, tmp_path): - coordinator = ExecutableCoordinator(executables_root=str(tmp_path)) - assert coordinator.executables_root == [tmp_path] - - def test_executables_root_accepts_list(self, tmp_path): - other = tmp_path / "other" - coordinator = ExecutableCoordinator(executables_root=[str(tmp_path), other]) - assert coordinator.executables_root == [tmp_path, other] - - def test_executables_root_optional_defaults_to_empty(self): - coordinator = ExecutableCoordinator() - assert coordinator.executables_root == [] - assert coordinator.dag_bundle_name is None - - @pytest.mark.parametrize( - "executables_root", - [None, [], "", " ", [""]], - ids=["none", "empty-list", "empty-str", "blank-str", "list-of-empty-str"], - ) - def test_explicit_empty_executables_root_raises(self, executables_root): - with pytest.raises(ValueError, match="and each path must be non-empty"): - ExecutableCoordinator(executables_root=executables_root) - - def test_root_and_dag_bundle_name_are_mutually_exclusive(self, tmp_path): - with pytest.raises(ValueError, match="at most one of 'executables_root' or 'dag_bundle_name'"): - ExecutableCoordinator(executables_root=[tmp_path], dag_bundle_name="artifacts") - - @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") - def test_unconfigured_dag_bundle_name_raises(self, mock_manager): - mock_manager.is_bundle_configured.return_value = False - with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): - ExecutableCoordinator(dag_bundle_name="ghost") - - @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") - def test_configured_dag_bundle_name_accepted(self, mock_manager): - mock_manager.is_bundle_configured.return_value = True - coordinator = ExecutableCoordinator(dag_bundle_name="artifacts") - assert coordinator.dag_bundle_name == "artifacts" - mock_manager.is_bundle_configured.assert_called_once_with("artifacts") - def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): binary = _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"]) coordinator = ExecutableCoordinator() @@ -417,7 +379,7 @@ def test_returns_resolved_executable_and_schema_version(self, tmp_path): binary = _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"]) ti = _make_ti(dag_id="tutorial_dag") - coordinator = ExecutableCoordinator(executables_root=[tmp_path]) + coordinator = ExecutableCoordinator() with coordinator._set_scan_roots([tmp_path]): command, schema_version = coordinator._build_execute_task_command(what=ti) assert command == [str(binary.resolve())] @@ -429,7 +391,7 @@ def test_raises_when_bundle_omits_schema_version(self, tmp_path): _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"], metadata=metadata) ti = _make_ti(dag_id="tutorial_dag") - coordinator = ExecutableCoordinator(executables_root=[tmp_path]) + coordinator = ExecutableCoordinator() with ( coordinator._set_scan_roots([tmp_path]), pytest.raises(FileNotFoundError, match="matching bundles were rejected"), @@ -440,7 +402,7 @@ def test_raises_when_dag_id_not_found(self, tmp_path): _build_bundle(tmp_path / "my_bundle", dag_ids=["other_dag"]) ti = _make_ti(dag_id="tutorial_dag") - coordinator = ExecutableCoordinator(executables_root=[tmp_path]) + coordinator = ExecutableCoordinator() with ( coordinator._set_scan_roots([tmp_path]), pytest.raises(FileNotFoundError, match="cannot find executable bundle"), @@ -448,12 +410,113 @@ def test_raises_when_dag_id_not_found(self, tmp_path): coordinator._build_execute_task_command(what=ti) +_CACHE_DIGEST = "c" * 64 + + +def _make_metadata_with_digests(**digests: str) -> dict: + return {**_make_metadata(["etl"]), "digests": digests} + + +class TestReadCacheDigest: + @pytest.mark.parametrize( + "bundle_kwargs", + [ + pytest.param( + {"metadata": _make_metadata_with_digests(integrity="a" * 64, cache=_CACHE_DIGEST)}, + id="valid-bundle", + ), + pytest.param( + {"metadata": _make_metadata_with_digests(cache=_CACHE_DIGEST), "binary_sha256": b"\x00" * 32}, + id="binary-digest-mismatch-not-checked", + ), + ], + ) + def test_returns_the_stored_digest(self, tmp_path, bundle_kwargs): + bundle = _build_bundle(tmp_path / "etl", **bundle_kwargs) + + assert read_cache_digest(bundle) == _CACHE_DIGEST + + @pytest.mark.parametrize( + "build", + [ + pytest.param(lambda path: _build_bundle(path), id="no-digests"), + pytest.param( + lambda path: _build_bundle(path, metadata=_make_metadata_with_digests(integrity="a" * 64)), + id="no-cache-digest", + ), + pytest.param( + lambda path: _build_bundle(path, metadata=_make_metadata_with_digests(cache="")), + id="empty-cache-digest", + ), + pytest.param(lambda path: _build_bundle(path, metadata=b"\xff\xfe"), id="undecodable-metadata"), + pytest.param(_make_executable, id="not-a-bundle"), + pytest.param(lambda path: path, id="missing-file"), + ], + ) + def test_returns_none_without_a_readable_digest(self, tmp_path, build): + assert read_cache_digest(build(tmp_path / "etl")) is None + + +class TestBuildParseTaskHandlerCommand: + def test_returns_the_bundle_and_its_schema_version(self, tmp_path): + # The Dag ids in the metadata play no part: the Dag processor names the artifact. + bundle = _build_bundle(tmp_path / "etl", dag_ids=["other_dag"]) + + command, schema_version = ExecutableCoordinator()._build_parse_task_handler_command(path=bundle) + + assert command == [str(bundle.resolve())] + assert schema_version == "2026-06-16" + + def test_returns_an_absolute_path(self, tmp_path, monkeypatch): + _build_bundle(tmp_path / "etl") + monkeypatch.chdir(tmp_path) + + command, _ = ExecutableCoordinator()._build_parse_task_handler_command(path=Path("etl")) + + assert command == [str((tmp_path / "etl").resolve())] + + @pytest.mark.parametrize( + "build", + [ + pytest.param(_make_executable, id="no-trailer"), + pytest.param( + lambda path: _build_bundle(path, binary_sha256=b"\x00" * 32), id="binary-digest-mismatch" + ), + ], + ) + def test_rejects_a_file_that_is_not_a_valid_bundle(self, tmp_path, build): + bundle = build(tmp_path / "etl") + + with pytest.raises(ValueError, match="is not a valid executable bundle"): + ExecutableCoordinator()._build_parse_task_handler_command(path=bundle) + + def test_rejects_a_bundle_without_a_schema_version(self, tmp_path): + metadata = _make_metadata(["etl"]) + del metadata["sdk"]["supervisor_schema_version"] + bundle = _build_bundle(tmp_path / "etl", metadata=metadata) + + with pytest.raises(ValueError, match="supervisor_schema_version"): + ExecutableCoordinator()._build_parse_task_handler_command(path=bundle) + + @pytest.fixture def bundles_dir(tmp_path): _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"]) return tmp_path +@pytest.fixture +def go_task_handlers(bundles_dir): + """Register *bundles_dir* as the ``go-task-handlers`` Dag bundle and return that name.""" + bundle = { + "name": "go-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": str(bundles_dir)}, + } + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps([bundle])}): + yield bundle["name"] + + @pytest.fixture def mock_client(make_ti_context): client = MagicMock() @@ -462,10 +525,10 @@ def mock_client(make_ti_context): class TestExecutableCoordinatorExecuteTask: - def _captured_popen_cmd(self, bundles_dir: pathlib.Path, mock_client) -> list[str]: + def _captured_popen_cmd(self, bundle_name: str, mock_client) -> list[str]: """Run execute_task with mocked subprocess and return the command list.""" ti = _make_ti(dag_id="tutorial_dag") - coordinator = ExecutableCoordinator(executables_root=[bundles_dir]) + coordinator = ExecutableCoordinator(task_handler_bundle_name=bundle_name) mock_proc = MagicMock(spec=subprocess.Popen) mock_proc.pid = 12345 @@ -505,14 +568,14 @@ def capture_popen(cmd, **kwargs): assert popen_calls, "subprocess.Popen was not called" return popen_calls[0] - def test_executable_path_is_first_arg(self, bundles_dir, mock_client): - cmd = self._captured_popen_cmd(bundles_dir, mock_client) + def test_executable_path_is_first_arg(self, bundles_dir, go_task_handlers, mock_client): + cmd = self._captured_popen_cmd(go_task_handlers, mock_client) expected = str((bundles_dir / "my_bundle").resolve()) assert cmd[0] == expected - def test_returns_execution_result(self, bundles_dir, mock_client): + def test_returns_execution_result(self, go_task_handlers, mock_client): ti = _make_ti(dag_id="tutorial_dag") - coordinator = ExecutableCoordinator(executables_root=[bundles_dir]) + coordinator = ExecutableCoordinator(task_handler_bundle_name=go_task_handlers) mock_proc = MagicMock(spec=subprocess.Popen) mock_proc.pid = 99999 diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py index c55a2cc891253..bae667d635b8e 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -18,6 +18,7 @@ from __future__ import annotations +import json import os import pathlib import re @@ -34,11 +35,13 @@ JavaCoordinator, _calculate_classpath, _JarInfo, + _parse_manifest, _walk_jars, ) from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess +from tests_common.test_utils.config import conf_vars from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS if not AIRFLOW_V_3_3_PLUS: @@ -222,50 +225,22 @@ def test_skip_logged_when_directory_revisited(self, tmp_path): class TestJavaCoordinatorAttributes: def test_default_kwargs(self): - coordinator = JavaCoordinator(jars_root="/airflow/java-bundles") + coordinator = JavaCoordinator() assert coordinator.java_executable == "java" assert coordinator.jvm_args == [] - assert coordinator.jars_root == [pathlib.Path("/airflow/java-bundles")] + assert coordinator.task_handler_bundle_name is None + # main_class stays optional: the entrypoint is auto-detected from the bundle scan. + assert coordinator.main_class == "" def test_custom_kwargs(self): coordinator = JavaCoordinator( java_executable="/opt/java/bin/java", jvm_args=["-Xmx512m", "-Xms256m"], - jars_root=["/airflow/java-bundles"], + task_handler_bundle_name="java-task-handlers", ) assert coordinator.java_executable == "/opt/java/bin/java" assert coordinator.jvm_args == ["-Xmx512m", "-Xms256m"] - assert coordinator.jars_root == [pathlib.Path("/airflow/java-bundles")] - - def test_jars_root_optional_defaults_to_empty(self): - # main_class stays optional: the entrypoint is auto-detected from the bundle scan. - coordinator = JavaCoordinator() - assert coordinator.jars_root == [] - assert coordinator.dag_bundle_name is None - assert coordinator.main_class == "" - - @pytest.mark.parametrize( - "jars_root", - [None, [], "", " ", [""]], - ids=["none", "empty-list", "empty-str", "blank-str", "list-of-empty-str"], - ) - def test_explicit_empty_jars_root_raises(self, jars_root): - with pytest.raises(ValueError, match="and each path must be non-empty"): - JavaCoordinator(jars_root=jars_root) - - def test_explicit_root_does_not_require_main_class(self, tmp_path): - coordinator = JavaCoordinator(jars_root=tmp_path) - assert coordinator.main_class == "" - - def test_root_and_dag_bundle_name_are_mutually_exclusive(self): - with pytest.raises(ValueError, match="at most one of 'jars_root' or 'dag_bundle_name'"): - JavaCoordinator(jars_root="/airflow/java-bundles", dag_bundle_name="artifacts") - - @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") - def test_unconfigured_dag_bundle_name_raises(self, mock_manager): - mock_manager.is_bundle_configured.return_value = False - with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): - JavaCoordinator(dag_bundle_name="ghost") + assert coordinator.task_handler_bundle_name == "java-task-handlers" def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): _make_jar(tmp_path / "app.jar", main_class="com.example.TaskRunner", schema_version="2026-06-16") @@ -277,12 +252,120 @@ def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): assert schema_version == "2026-06-16" +def test_parse_manifest_joins_continued_values_and_stops_at_the_main_section(): + manifest = ( + b"Manifest-Version: 1.0\r\n" + b"Airflow-Cache-Digest: 6a7cf9d4ddd454e97952b07038a8ebc6ff37c3a52ae176e921\r\n" + b" bc03065d6404cd\r\n" + b"Main-Class: com.example.Main\r\n" + b"\r\n" + b"Name: com/example/\r\n" + b"Sealed: true\r\n" + ) + + assert _parse_manifest(manifest) == { + "manifest-version": "1.0", + "airflow-cache-digest": "6a7cf9d4ddd454e97952b07038a8ebc6ff37c3a52ae176e921bc03065d6404cd", + "main-class": "com.example.Main", + } + + +def test_parse_manifest_decodes_a_character_split_across_a_fold(): + manifest = b"Manifest-Version: 1.0\r\nImplementation-Vendor: Z\xc3\r\n \xbcrich\r\n\r\n" + + assert _parse_manifest(manifest)["implementation-vendor"] == "Z\u00fcrich" + + +class TestJavaCoordinatorParseTaskHandlerCommand: + @pytest.fixture + def bundle(self, tmp_path): + _make_jar(tmp_path / "app.jar", main_class="com.example.App", schema_version="2026-06-16") + _make_jar(tmp_path / "other.jar", main_class="com.example.Other", schema_version="2026-06-16") + (tmp_path / "lib").mkdir() + _make_jar(tmp_path / "lib" / "dep.jar", main_class=None) + return tmp_path + + @pytest.mark.parametrize( + ("jar", "main_class"), [("app.jar", "com.example.App"), ("other.jar", "com.example.Other")] + ) + def test_runs_the_probed_jars_own_main_class_with_the_bundle_on_the_classpath( + self, bundle, jar, main_class + ): + coordinator = JavaCoordinator(jvm_args=["-Xmx1g"], main_class="com.example.Configured") + with coordinator._set_scan_roots([bundle]): + command, schema_version = coordinator._build_parse_task_handler_command(path=bundle / jar) + + classpath = os.pathsep.join( + (bundle / name).as_posix() for name in ("app.jar", "lib/dep.jar", "other.jar") + ) + assert command == ["java", "-classpath", classpath, "-Xmx1g", main_class] + assert schema_version == "2026-06-16" + + def test_a_thin_jar_takes_the_schema_version_from_the_sdk_jar(self, tmp_path): + _make_jar(tmp_path / "app.jar", main_class="com.example.App") + _make_jar(tmp_path / "airflow-sdk.jar", main_class=None, schema_version="2026-10-30") + coordinator = JavaCoordinator() + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_parse_task_handler_command(path=tmp_path / "app.jar") + + assert command[-1] == "com.example.App" + assert schema_version == "2026-10-30" + + def test_reads_a_main_class_continued_on_the_next_manifest_line(self, tmp_path): + main_class = "org.example.a.very.long.package.name.that.needs.folding.BundleMain" + manifest = f"Manifest-Version: 1.0\r\nMain-Class: {main_class[:50]}\r\n {main_class[50:]}\r\n" + manifest += "Airflow-Supervisor-Schema-Version: 2026-06-16\r\n\r\n" + with zipfile.ZipFile(tmp_path / "app.jar", "w") as zf: + zf.writestr("META-INF/MANIFEST.MF", manifest) + coordinator = JavaCoordinator() + with coordinator._set_scan_roots([tmp_path]): + command, _ = coordinator._build_parse_task_handler_command(path=tmp_path / "app.jar") + + assert command[-1] == main_class + + @pytest.mark.parametrize( + ("content", "reason"), + [ + ("no-main-class", "its manifest sets no Main-Class"), + ("no-manifest", "it has no readable manifest"), + ("not-a-zip", "it has no readable manifest"), + ], + ) + def test_a_jar_that_is_not_executable_is_rejected(self, tmp_path, content, reason): + path = tmp_path / "app.jar" + if content == "no-main-class": + _make_jar(path, main_class=None, schema_version="2026-06-16") + elif content == "no-manifest": + with zipfile.ZipFile(path, "w") as zf: + zf.writestr("com/example/App.class", b"") + else: + path.write_text("not a zip") + coordinator = JavaCoordinator() + with ( + coordinator._set_scan_roots([tmp_path]), + pytest.raises(ValueError, match=f"is not an executable JAR: {reason}$"), + ): + coordinator._build_parse_task_handler_command(path=path) + + @pytest.fixture -def jars_root(tmp_path): +def jars_dir(tmp_path): _make_jar(tmp_path.joinpath("app.jar"), main_class="com.example.TaskRunner", schema_version="2026-06-16") return tmp_path +@pytest.fixture +def java_task_handlers(jars_dir): + """Register *jars_dir* as the ``java-task-handlers`` Dag bundle and return that name.""" + bundle = { + "name": "java-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": str(jars_dir)}, + } + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps([bundle])}): + yield bundle["name"] + + @pytest.fixture def mock_client(make_ti_context): client = MagicMock() @@ -293,7 +376,7 @@ def mock_client(make_ti_context): class TestJavaCoordinatorExecuteTask: def _captured_popen_cmd( self, - jars_root: pathlib.Path, + bundle_name: str, mock_client, *, java_executable: str = "java", @@ -304,7 +387,7 @@ def _captured_popen_cmd( coordinator = JavaCoordinator( java_executable=java_executable, jvm_args=jvm_args or [], - jars_root=jars_root, + task_handler_bundle_name=bundle_name, ) mock_proc = MagicMock(spec=subprocess.Popen) @@ -345,56 +428,58 @@ def capture_popen(cmd, **kwargs): assert popen_calls, "subprocess.Popen was not called" return popen_calls[0] - def test_java_executable_is_first_arg(self, jars_root, mock_client): + def test_java_executable_is_first_arg(self, java_task_handlers, mock_client): cmd = self._captured_popen_cmd( - jars_root, mock_client, java_executable="/usr/lib/jvm/java-17/bin/java" + java_task_handlers, mock_client, java_executable="/usr/lib/jvm/java-17/bin/java" ) assert cmd[0] == "/usr/lib/jvm/java-17/bin/java" - def test_classpath_flag_and_value_present(self, jars_root, mock_client): - cmd = self._captured_popen_cmd(jars_root, mock_client) + def test_classpath_flag_and_value_present(self, jars_dir, java_task_handlers, mock_client): + cmd = self._captured_popen_cmd(java_task_handlers, mock_client) assert "-classpath" in cmd cp_idx = cmd.index("-classpath") classpath = cmd[cp_idx + 1] - assert jars_root.joinpath("app.jar").as_posix() in classpath + assert jars_dir.joinpath("app.jar").as_posix() in classpath - def test_main_class_present(self, jars_root, mock_client): - cmd = self._captured_popen_cmd(jars_root, mock_client) + def test_main_class_present(self, java_task_handlers, mock_client): + cmd = self._captured_popen_cmd(java_task_handlers, mock_client) assert "com.example.TaskRunner" in cmd - def test_comm_and_logs_args_present(self, jars_root, mock_client): - cmd = self._captured_popen_cmd(jars_root, mock_client) + def test_comm_and_logs_args_present(self, java_task_handlers, mock_client): + cmd = self._captured_popen_cmd(java_task_handlers, mock_client) comm_args = [a for a in cmd if a.startswith("--comm=")] logs_args = [a for a in cmd if a.startswith("--logs=")] assert len(comm_args) == 1 assert len(logs_args) == 1 - def test_comm_and_logs_contain_port(self, jars_root, mock_client): - cmd = self._captured_popen_cmd(jars_root, mock_client) + def test_comm_and_logs_contain_port(self, java_task_handlers, mock_client): + cmd = self._captured_popen_cmd(java_task_handlers, mock_client) comm_arg = next(a for a in cmd if a.startswith("--comm=")) logs_arg = next(a for a in cmd if a.startswith("--logs=")) # format is host:port assert ":" in comm_arg.split("=", 1)[1] assert ":" in logs_arg.split("=", 1)[1] - def test_jvm_args_inserted_before_main_class(self, jars_root, mock_client): - cmd = self._captured_popen_cmd(jars_root, mock_client, jvm_args=["-Xmx512m", "-Dsome.prop=value"]) + def test_jvm_args_inserted_before_main_class(self, java_task_handlers, mock_client): + cmd = self._captured_popen_cmd( + java_task_handlers, mock_client, jvm_args=["-Xmx512m", "-Dsome.prop=value"] + ) main_idx = cmd.index("com.example.TaskRunner") for jvm_arg in ["-Xmx512m", "-Dsome.prop=value"]: assert jvm_arg in cmd assert cmd.index(jvm_arg) < main_idx - def test_comm_and_logs_after_main_class(self, jars_root, mock_client): - cmd = self._captured_popen_cmd(jars_root, mock_client) + def test_comm_and_logs_after_main_class(self, java_task_handlers, mock_client): + cmd = self._captured_popen_cmd(java_task_handlers, mock_client) main_idx = cmd.index("com.example.TaskRunner") comm_idx = next(i for i, a in enumerate(cmd) if a.startswith("--comm=")) logs_idx = next(i for i, a in enumerate(cmd) if a.startswith("--logs=")) assert comm_idx > main_idx assert logs_idx > main_idx - def test_returns_execution_result(self, jars_root, mock_client): + def test_returns_execution_result(self, java_task_handlers, mock_client): ti = _make_ti() - coordinator = JavaCoordinator(jars_root=jars_root) + coordinator = JavaCoordinator(task_handler_bundle_name=java_task_handlers) mock_proc = MagicMock(spec=subprocess.Popen) mock_proc.pid = 99999 diff --git a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py index 146cb91c1d94c..5c5af582be9ba 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -21,7 +21,6 @@ import json import pathlib from unittest import mock -from unittest.mock import patch import pytest from task_sdk.coordinators.node._bundle_test_utils import ( @@ -60,50 +59,22 @@ def _make_ti(dag_id: str = "test_dag", queue: str = "ts") -> TaskInstance: class TestNodeCoordinatorAttributes: def test_default_kwargs(self): - coordinator = NodeCoordinator(bundles_root="/airflow/ts-bundles") + coordinator = NodeCoordinator() assert coordinator.node_executable == "node" - assert coordinator.bundles_root == [pathlib.Path("/airflow/ts-bundles")] assert coordinator.task_startup_timeout == 10.0 def test_custom_kwargs(self): coordinator = NodeCoordinator( node_executable="/opt/node/bin/node", - bundles_root=["/airflow/ts-bundles", "~/extra-bundles"], + task_handler_bundle_name="ts-task-handlers", task_startup_timeout=30.0, ) assert coordinator.node_executable == "/opt/node/bin/node" - assert coordinator.bundles_root == [ - pathlib.Path("/airflow/ts-bundles"), - pathlib.Path("~/extra-bundles").expanduser(), - ] + assert coordinator.task_handler_bundle_name == "ts-task-handlers" assert coordinator.task_startup_timeout == 30.0 - def test_bundles_root_optional_defaults_to_empty(self): - coordinator = NodeCoordinator() - assert coordinator.bundles_root == [] - assert coordinator.dag_bundle_name is None - - @pytest.mark.parametrize( - "bundles_root", - [None, [], "", " ", [""]], - ids=["none", "empty-list", "empty-str", "blank-str", "list-of-empty-str"], - ) - def test_explicit_empty_bundles_root_raises(self, bundles_root): - with pytest.raises(ValueError, match="and each path must be non-empty"): - NodeCoordinator(bundles_root=bundles_root) - - def test_root_and_dag_bundle_name_are_mutually_exclusive(self): - with pytest.raises(ValueError, match="at most one of 'bundles_root' or 'dag_bundle_name'"): - NodeCoordinator(bundles_root="/airflow/ts-bundles", dag_bundle_name="artifacts") - - @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") - def test_unconfigured_dag_bundle_name_raises(self, mock_manager): - mock_manager.is_bundle_configured.return_value = False - with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): - NodeCoordinator(dag_bundle_name="ghost") - def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): bundle = write_bundle(tmp_path, "test_dag") coordinator = NodeCoordinator() @@ -116,10 +87,7 @@ def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): class TestNodeCoordinatorExecuteTaskCommand: def test_selects_bundle_by_dag_id(self, tmp_path): selected = write_bundle(tmp_path, "sales") - coordinator = NodeCoordinator( - node_executable="/opt/node/bin/node", - bundles_root=tmp_path, - ) + coordinator = NodeCoordinator(node_executable="/opt/node/bin/node") with coordinator._set_scan_roots([tmp_path]): command, schema_version = coordinator._build_execute_task_command(what=_make_ti(dag_id="sales")) @@ -129,10 +97,7 @@ def test_selects_bundle_by_dag_id(self, tmp_path): def test_build_execute_task_command_returns_node_bundle_and_schema_version(self, tmp_path): bundle = write_bundle(tmp_path, "test_dag") - coordinator = NodeCoordinator( - node_executable="/opt/node/bin/node", - bundles_root=tmp_path, - ) + coordinator = NodeCoordinator(node_executable="/opt/node/bin/node") with coordinator._set_scan_roots([tmp_path]): command, schema_version = coordinator._build_execute_task_command(what=_make_ti()) @@ -141,6 +106,44 @@ def test_build_execute_task_command_returns_node_bundle_and_schema_version(self, assert schema_version == SCHEMA_VERSION +def _write_plain_file(root: pathlib.Path) -> pathlib.Path: + path = root / BUNDLE_NAME + path.write_bytes(b"export {};\n") + return path + + +def _write_corrupted_bundle(root: pathlib.Path) -> pathlib.Path: + bundle = write_bundle(root, "test_dag") + # The last byte before the trailing newline is in the code region. + mutate_byte(bundle, len(bundle.read_bytes()) - 2) + return bundle + + +class TestNodeCoordinatorParseTaskHandlerCommand: + def test_returns_node_and_the_bundle_schema_version(self, tmp_path): + # A bundle registering no Dag: the command does not depend on the Dags it declares. + bundle = write_bundle(tmp_path) + coordinator = NodeCoordinator(node_executable="/opt/node/bin/node") + + command, schema_version = coordinator._build_parse_task_handler_command(path=bundle) + + assert command == ["/opt/node/bin/node", str(bundle)] + assert schema_version == SCHEMA_VERSION + + @pytest.mark.parametrize( + ("write", "error"), + [ + pytest.param(_write_plain_file, "has no airflow bundle layout", id="not-a-bundle"), + pytest.param(_write_corrupted_bundle, "code SHA-256 mismatch", id="corrupted-code"), + ], + ) + def test_rejects_a_bundle_that_fails_verification(self, tmp_path, write, error): + path = write(tmp_path) + + with pytest.raises(ValueError, match=error): + NodeCoordinator()._build_parse_task_handler_command(path=path) + + class TestBundleFind: @pytest.mark.parametrize( "name", diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 5297f56dead4f..8db8cbb186fd7 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -20,6 +20,7 @@ import contextlib import os import pathlib +import signal import socket import subprocess import sys @@ -32,14 +33,14 @@ import pytest from uuid6 import uuid7 -from airflow.dag_processing.bundles.base import BundleVersion +from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion from airflow.sdk.api.client import Client, TaskInstanceOperations from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance from airflow.sdk.coordinators._subprocess import ( SubprocessCoordinator, _accept_connections, - _ArtifactSource, _connection_owned_by_process_tree, + _is_connection_from_pid, _is_connection_from_process, _PopenActivitySubprocess, _ResourceTracker, @@ -495,13 +496,22 @@ def test_rejects_when_ownership_never_confirmed(self): ) def test_owned_by_tree_returns_false_when_process_gone(self): - mock_proc = MagicMock(spec=subprocess.Popen) - mock_proc.pid = 999999 with patch( "airflow.sdk.coordinators._subprocess.psutil.Process", side_effect=psutil.NoSuchProcess(999999), ): - assert _connection_owned_by_process_tree(("127.0.0.1", 1), ("127.0.0.1", 2), mock_proc) is False + assert _connection_owned_by_process_tree(("127.0.0.1", 1), ("127.0.0.1", 2), 999999) is False + + @pytest.mark.parametrize("owned", [True, False]) + @patch("airflow.sdk.coordinators._subprocess._connection_owned_by_process_tree", autospec=True) + def test_is_connection_from_pid_checks_once(self, mock_owned, owned): + conn = MagicMock(spec=socket.socket) + conn.getpeername.return_value = ("127.0.0.1", 5000) + conn.getsockname.return_value = ("127.0.0.1", 6000) + mock_owned.return_value = owned + + assert _is_connection_from_pid(conn, 999) is owned + mock_owned.assert_called_once_with(("127.0.0.1", 5000), ("127.0.0.1", 6000), 999) class TestResourceTracker: @@ -573,26 +583,37 @@ def test_untrack_unknown_object_does_not_raise(self): class _StubSubprocessCoordinator(SubprocessCoordinator): """Minimal SubprocessCoordinator subclass used to exercise the base machinery. - ``explicit_roots`` defaults to a real path so the coordinator classifies as - EXPLICIT_ROOT and execute_task resolves without touching a Dag bundle; pass - ``explicit_roots=[]`` to exercise TASK_BUNDLE mode. Roots handed to the - command builder are recorded in ``recorded_roots`` so wiring can be asserted. + Roots handed to the command builder are recorded in ``recorded_roots`` so + wiring can be asserted. """ command: list[str] schema_version: str | None = None - explicit_roots: list[pathlib.Path] = attrs.field(factory=lambda: [pathlib.Path(".")]) recorded_roots: list[list[pathlib.Path]] = attrs.field(init=False, factory=list) - @property - def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: - return "explicit_roots", self.explicit_roots - def _build_execute_task_command(self, *, what): self.recorded_roots.append(list(self._get_scan_roots())) return list(self.command), self.schema_version +def _make_bundle(path: pathlib.Path, *, version: str | None = None, name: str = "dags") -> MagicMock: + bundle = MagicMock(spec=BaseDagBundle, path=path, version=version) + bundle.name = name + return bundle + + +@pytest.fixture +def artifact_bundle(tmp_path): + """Resolve every task's artifacts to an unversioned bundle at *tmp_path*, without the bundle machinery.""" + with patch.object( + SubprocessCoordinator, + "_resolve_artifact_bundle", + autospec=True, + return_value=_make_bundle(tmp_path), + ) as mock_resolve: + yield mock_resolve + + @pytest.fixture def mock_client(make_ti_context): client = MagicMock(spec=Client) @@ -618,6 +639,7 @@ class _Plain(SubprocessCoordinator): _Plain()._build_execute_task_command(what=_make_ti()) +@pytest.mark.usefixtures("artifact_bundle") class TestSubprocessCoordinatorExecuteTask: def _captured_popen_cmd( self, @@ -863,79 +885,20 @@ def test_register_pipe_readers_called_with_four_sockets(self, mock_client): assert mock_register.mock_calls == [call(ANY, ANY, ANY, ANY, data=ANY)] -class TestClassifyArtifactSource: - """Construction-time classification and per-mode validation. - - Classification runs from the base ``__attrs_post_init__``, so it is exercised - through construction rather than by calling the helper directly. - """ - - def test_explicit_root_classified_and_recorded(self): - configured = [pathlib.Path("/artifacts")] - coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=configured) - assert coordinator._artifact_source is _ArtifactSource.EXPLICIT_ROOT - assert coordinator._configured_roots == configured - - def test_task_bundle_classified_when_neither_set(self): - coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) - assert coordinator._artifact_source is _ArtifactSource.TASK_BUNDLE - - def test_subclass_without_overrides_defaults_to_task_bundle(self): - """A subclass that overrides neither hook is classified, not left to error at execute time.""" - - @attrs.define(kw_only=True) - class _Bare(SubprocessCoordinator): - def _build_execute_task_command(self, *, what): - return [], None - - assert _Bare()._artifact_source is _ArtifactSource.TASK_BUNDLE - - def test_rejects_explicit_root_and_dag_bundle_name_together(self): - with pytest.raises(ValueError, match="at most one of 'explicit_roots' or 'dag_bundle_name'"): - _StubSubprocessCoordinator( - command=["x"], explicit_roots=[pathlib.Path("/artifacts")], dag_bundle_name="artifacts" - ) - - @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") - def test_rejects_unconfigured_dag_bundle_name(self, mock_manager): - mock_manager.is_bundle_configured.return_value = False - with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): - _StubSubprocessCoordinator(command=["x"], explicit_roots=[], dag_bundle_name="ghost") - - @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") - def test_accepts_configured_dag_bundle_name(self, mock_manager): - mock_manager.is_bundle_configured.return_value = True - coordinator = _StubSubprocessCoordinator( - command=["x"], explicit_roots=[], dag_bundle_name="artifacts" - ) - assert coordinator._artifact_source is _ArtifactSource.NAMED_BUNDLE - mock_manager.is_bundle_configured.assert_called_once_with("artifacts") - - -class TestInitRootSource: - """Execute-time root resolution, dispatched on the classified mode.""" +class TestResolveArtifactBundle: + """Execute-time resolution of the Dag bundle holding the artifacts.""" - def test_returns_configured_root_and_no_bundle_in_explicit_mode(self): - configured = [pathlib.Path("/a"), pathlib.Path("/b")] - coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=configured) - roots, bundle = coordinator._init_root_source(MagicMock(), log) - assert roots == configured - assert bundle is None - - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") - def test_task_bundle_mode_uses_passed_task_bundle(self, mock_initialize, tmp_path): - resolved = MagicMock(path=tmp_path, version="v3") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) + def test_unset_name_uses_the_task_bundle(self, mock_initialize, tmp_path): + resolved = _make_bundle(tmp_path, version="v3") mock_initialize.return_value = resolved - coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) + coordinator = _StubSubprocessCoordinator(command=["x"]) bundle_info = BundleInfo(name="dags", version="v3", version_data={"k": "v"}) - roots, bundle = coordinator._init_root_source(bundle_info, log) - - assert roots == [tmp_path] - assert bundle is resolved + assert coordinator._resolve_artifact_bundle(bundle_info, log) is resolved mock_initialize.assert_called_once_with(bundle_info) - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) def test_version_less_bundle_is_rematerialized_at_its_current_version(self, mock_initialize, tmp_path): """A bundle resolved without a version is re-resolved at the version current now. @@ -947,67 +910,52 @@ def test_version_less_bundle_is_rematerialized_at_its_current_version(self, mock pinned_tree = tmp_path / "versions" / "sha-abc" pinned_tree.mkdir(parents=True) - unpinned = MagicMock(path=shared_checkout, version=None) + unpinned = _make_bundle(shared_checkout) unpinned.get_current_version.return_value = BundleVersion(version="sha-abc", data={"k": "v"}) - pinned = MagicMock(path=pinned_tree, version="sha-abc") + pinned = _make_bundle(pinned_tree, version="sha-abc") mock_initialize.side_effect = [unpinned, pinned] - coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) - coordinator.dag_bundle_name = "artifacts" - coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE + coordinator = _StubSubprocessCoordinator(command=["x"], task_handler_bundle_name="artifacts") - roots, bundle = coordinator._init_root_source(MagicMock(), log) - - assert roots == [pinned_tree] - assert bundle is pinned + assert coordinator._resolve_artifact_bundle(MagicMock(spec=BundleInfo), log) is pinned assert mock_initialize.call_args_list == [ call(BundleInfo(name="artifacts")), call(BundleInfo(name="artifacts", version="sha-abc", version_data={"k": "v"})), ] - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) def test_unversioned_bundle_keeps_its_single_path(self, mock_initialize, tmp_path): - resolved = MagicMock(path=tmp_path, version=None) + resolved = _make_bundle(tmp_path) resolved.get_current_version.return_value = None mock_initialize.return_value = resolved - coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) - coordinator.dag_bundle_name = "artifacts" - coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE - - roots, bundle = coordinator._init_root_source(MagicMock(), log) + coordinator = _StubSubprocessCoordinator(command=["x"], task_handler_bundle_name="artifacts") - assert roots == [tmp_path] - assert bundle is resolved + assert coordinator._resolve_artifact_bundle(MagicMock(spec=BundleInfo), log) is resolved mock_initialize.assert_called_once_with(BundleInfo(name="artifacts")) - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) def test_missing_resolved_path_raises(self, mock_initialize, tmp_path): - missing = tmp_path / "nope" - mock_initialize.return_value = MagicMock(path=missing, version="v1") - coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) - coordinator.dag_bundle_name = "artifacts" - coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE + mock_initialize.return_value = _make_bundle(tmp_path / "nope", version="v1") + coordinator = _StubSubprocessCoordinator(command=["x"], task_handler_bundle_name="artifacts") with pytest.raises(FileNotFoundError, match="does not exist"): - coordinator._init_root_source(MagicMock(), log) + coordinator._resolve_artifact_bundle(MagicMock(spec=BundleInfo), log) class TestExecuteTaskBundleWiring: """execute_task passes the task bundle, forwards resolved roots, and holds the version lock.""" - @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock") - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") - @patch.object(_PopenActivitySubprocess, "start") + @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) + @patch.object(_PopenActivitySubprocess, "start", autospec=True) def test_task_bundle_mode_binds_forwards_roots_and_locks( self, mock_start, mock_initialize, mock_lock, mock_client, tmp_path ): - resolved = MagicMock(path=tmp_path, version="v9") - resolved.name = "dags" - mock_initialize.return_value = resolved + mock_initialize.return_value = _make_bundle(tmp_path, version="v9") mock_start.return_value.wait.return_value = 0 - coordinator = _StubSubprocessCoordinator(command=["/runtime"], explicit_roots=[]) + coordinator = _StubSubprocessCoordinator(command=["/runtime"]) bundle_info = BundleInfo(name="dags", version="v9") coordinator.execute_task( @@ -1024,28 +972,9 @@ def test_task_bundle_mode_binds_forwards_roots_and_locks( mock_lock.return_value.__enter__.assert_called_once() mock_lock.return_value.__exit__.assert_called_once() - @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock") - @patch.object(_PopenActivitySubprocess, "start") - def test_explicit_root_mode_forwards_roots_without_locking( - self, mock_start, mock_lock, mock_client, tmp_path - ): - mock_start.return_value.wait.return_value = 0 - coordinator = _StubSubprocessCoordinator(command=["/runtime"], explicit_roots=[tmp_path]) - - coordinator.execute_task( - what=_make_ti(), - dag_rel_path="dag.py", - bundle_info=MagicMock(), - client=mock_client, - subprocess_logs_to_stdout=False, - ) - - assert coordinator.recorded_roots == [[tmp_path]] - mock_lock.assert_not_called() - - @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock") - @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") - @patch.object(_PopenActivitySubprocess, "start") + @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True) + @patch.object(_PopenActivitySubprocess, "start", autospec=True) def test_named_bundle_mode_locks_the_tree_it_scans( self, mock_start, mock_initialize, mock_lock, mock_client, tmp_path ): @@ -1053,21 +982,18 @@ def test_named_bundle_mode_locks_the_tree_it_scans( pinned_tree = tmp_path / "versions" / "sha-abc" pinned_tree.mkdir(parents=True) - unpinned = MagicMock(path=tmp_path, version=None) + unpinned = _make_bundle(tmp_path, name="artifacts") unpinned.get_current_version.return_value = BundleVersion(version="sha-abc", data=None) - pinned = MagicMock(path=pinned_tree, version="sha-abc") - pinned.name = "artifacts" + pinned = _make_bundle(pinned_tree, version="sha-abc", name="artifacts") mock_initialize.side_effect = [unpinned, pinned] mock_start.return_value.wait.return_value = 0 - coordinator = _StubSubprocessCoordinator(command=["/runtime"], explicit_roots=[]) - coordinator.dag_bundle_name = "artifacts" - coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE + coordinator = _StubSubprocessCoordinator(command=["/runtime"], task_handler_bundle_name="artifacts") coordinator.execute_task( what=_make_ti(), dag_rel_path="dag.py", - bundle_info=MagicMock(), + bundle_info=BundleInfo(name="dags", version="v9"), client=mock_client, subprocess_logs_to_stdout=False, ) @@ -1093,6 +1019,7 @@ def test_rejects_nested_reentry(self, tmp_path): with coordinator._set_scan_roots([tmp_path]): pass + @pytest.mark.usefixtures("artifact_bundle") def test_execute_task_is_not_reentrant(self, mock_client, tmp_path): coordinator = _StubSubprocessCoordinator(command=["/bin/true"]) with coordinator._set_scan_roots([tmp_path]): @@ -1104,3 +1031,137 @@ def test_execute_task_is_not_reentrant(self, mock_client, tmp_path): client=mock_client, subprocess_logs_to_stdout=False, ) + + +@attrs.define(kw_only=True) +class _TaskHandlerParsingCoordinator(_StubSubprocessCoordinator): + def _build_parse_task_handler_command(self, *, path): + self.recorded_roots.append(list(self._get_scan_roots())) + return [*self.command, os.fspath(path)], self.schema_version + + +def _parse_task_handler(coordinator: SubprocessCoordinator, tmp_path: pathlib.Path, reported: list) -> None: + coordinator.parse_task_handler( + path=tmp_path / "handlers.artifact", + bundle_path=tmp_path, + comm_address=("127.0.0.1", 1001), + logs_address=("127.0.0.1", 1002), + report_schema_version=reported.append, + ) + + +def _is_running(process: psutil.Process) -> bool: + try: + return process.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +class TestParseTaskHandler: + def test_build_parse_task_handler_command_default_raises(self, tmp_path): + coordinator = _StubSubprocessCoordinator(command=["x"]) + + with pytest.raises( + NotImplementedError, match="_StubSubprocessCoordinator does not parse task handlers" + ): + coordinator._build_parse_task_handler_command(path=tmp_path / "handlers.artifact") + + @patch("airflow.sdk.coordinators._subprocess._set_close_on_exec_above_stderr", autospec=True) + @patch("airflow.sdk.coordinators._subprocess._set_parent_death_signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.signal.signal", autospec=True) + @patch( + "airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True, side_effect=OSError("exec failed") + ) + def test_resolves_the_command_under_the_bundle_root_then_execs_it( + self, mock_execvpe, mock_signal, mock_death_signal, mock_close_on_exec, tmp_path + ): + coordinator = _TaskHandlerParsingCoordinator(command=["runtime"], schema_version="2026-06-16") + reported: list[str | None] = [] + + with pytest.raises(OSError, match="exec failed"): + _parse_task_handler(coordinator, tmp_path, reported) + + assert coordinator.recorded_roots == [[tmp_path]] + assert coordinator._active_scan_roots is None + assert reported == ["2026-06-16"] + argv = [ + "runtime", + os.fspath(tmp_path / "handlers.artifact"), + "--comm=127.0.0.1:1001", + "--logs=127.0.0.1:1002", + ] + mock_execvpe.assert_called_once_with("runtime", argv, ANY) + assert "AIRFLOW__LOGGING__LOGGING_LEVEL" in mock_execvpe.call_args.args[2] + restored = {c.args[0] for c in mock_signal.call_args_list if c.args[1] == signal.SIG_DFL} + assert {signal.SIGPIPE, signal.SIGXFSZ} <= restored + mock_death_signal.assert_called_once_with() + mock_close_on_exec.assert_called_once_with() + + @patch("airflow.sdk.coordinators._subprocess._set_close_on_exec_above_stderr", autospec=True) + @patch("airflow.sdk.coordinators._subprocess._set_parent_death_signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.signal.signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.os.getppid", autospec=True, side_effect=[4242, 1]) + @patch("airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True) + def test_does_not_exec_once_the_parent_has_exited( + self, mock_execvpe, mock_getppid, mock_signal, mock_death_signal, mock_close_on_exec, tmp_path + ): + coordinator = _TaskHandlerParsingCoordinator(command=["runtime"]) + + with pytest.raises(RuntimeError, match="The process that started the runtime has exited"): + _parse_task_handler(coordinator, tmp_path, []) + + mock_death_signal.assert_called_once_with() + mock_execvpe.assert_not_called() + + @patch("airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True) + def test_rejects_an_unknown_schema_version_before_reporting(self, mock_execvpe, tmp_path): + coordinator = _TaskHandlerParsingCoordinator(command=["runtime"], schema_version="1999-01-01") + reported: list[str | None] = [] + + with pytest.raises(ValueError, match="'1999-01-01' not found"): + _parse_task_handler(coordinator, tmp_path, reported) + + assert reported == [] + mock_execvpe.assert_not_called() + + @pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only") + def test_the_runtime_dies_with_the_process_that_started_it(self, tmp_path): + coordinator = _TaskHandlerParsingCoordinator(command=["/bin/sh", "-c", "exec sleep 60"]) + read_fd, write_fd = os.pipe() + parent_pid = os.fork() + if parent_pid == 0: + # The process reading the runtime's sockets: it starts the runtime, reports its pid and waits. + try: + os.close(read_fd) + runtime_pid = os.fork() + if runtime_pid == 0: + try: + _parse_task_handler(coordinator, tmp_path, []) + finally: + os._exit(1) + os.write(write_fd, str(runtime_pid).encode()) + os.close(write_fd) + os.waitpid(runtime_pid, 0) + finally: + os._exit(0) + + os.close(write_fd) + with os.fdopen(read_fd) as reader: + runtime = psutil.Process(int(reader.read())) + try: + # Polls: the runtime is not a child to wait for, and time_machine cannot reach another process. + deadline = time.monotonic() + 30 + while runtime.name() != "sleep": + assert time.monotonic() < deadline, "the runtime did not start" + time.sleep(0.05) + + os.kill(parent_pid, signal.SIGKILL) + os.waitpid(parent_pid, 0) + + deadline = time.monotonic() + 10 + while _is_running(runtime): + assert time.monotonic() < deadline, "the runtime outlived the process that started it" + time.sleep(0.05) + finally: + with contextlib.suppress(psutil.NoSuchProcess): + runtime.kill() diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index 3befcf14b421b..d837181c6f6d4 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -30,6 +30,7 @@ import subprocess import sys +from pathlib import Path from typing import Literal import pytest @@ -48,6 +49,13 @@ _SupervisorResponse, ) +from airflow.dag_processing.processor import ( + DagFileParseRequest, + TaskHandlerArtifact, + TaskHandlerDeclaration, + TaskHandlerParam, + TaskHandlerParseRequest, +) from airflow.sdk import TaskInstanceState from airflow.sdk.execution_time.comms import TaskState from airflow.sdk.execution_time.schema import ( @@ -337,6 +345,11 @@ def test_extra_fields_in_body_do_not_affect_resolution(self, mock_registry): assert resolve_body_class(body) is _LangSdkRequest +class TestRealRegistry: + def test_resolves_task_handler_parse_request(self): + assert resolve_body_class({"type": "TaskHandlerParseRequest"}) is TaskHandlerParseRequest + + class TestLazyCadwynImport: """``cadwyn`` (which imports FastAPI/Starlette/Jinja2) must stay off the worker import path. @@ -501,3 +514,67 @@ def test_upgrade_fills_missing_retry_reason_with_none(self, real_migrator): body = {"type": "TaskState", "state": "failed", "end_date": None, "rendered_map_index": None} out = real_migrator.upgrade(body, TaskState, "2026-06-16") assert out["retry_reason"] is None + + +class TestRealBundleKnownArtifacts: + """ + Drive the *real* supervisor bundle through the ``known_artifacts`` migration. + + ``DagFileParseRequest`` flows supervisor -> runtime, so ``downgrade`` is the direction a + pinned runtime travels. + """ + + @pytest.fixture + def request_with_known_artifacts(self) -> DagFileParseRequest: + return DagFileParseRequest( + file="/files/dags/etl.py", + bundle_path=Path("/files/dags"), + bundle_name="dags-folder", + known_artifacts=[ + TaskHandlerArtifact( + bundle_name="java-task-handlers", + relative_fileloc="etl.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={ + "etl": [ + TaskHandlerDeclaration( + task_id="extract", + binding="positional", + params=[TaskHandlerParam(name=None)], + ) + ] + }, + ) + ], + ) + + @pytest.fixture + def real_migrator(self) -> SchemaVersionMigrator: + return get_schema_version_migrator() + + def test_downgrade_strips_known_artifacts_for_previous_version( + self, real_migrator, request_with_known_artifacts + ): + out = real_migrator.downgrade(request_with_known_artifacts, "2026-06-16").model_dump() + assert "known_artifacts" not in out + + def test_downgrade_keeps_known_artifacts_at_head(self, real_migrator, request_with_known_artifacts): + out = real_migrator.downgrade(request_with_known_artifacts, "2026-10-30").model_dump() + assert out["known_artifacts"] == [ + { + "bundle_name": "java-task-handlers", + "relative_fileloc": "etl.jar", + "size_bytes": 1024, + "cache_digest": "ab12", + "task_handlers": { + "etl": [ + { + "task_id": "extract", + "binding": "positional", + "params": [{"name": None, "value_schema": None, "exact_name": False}], + } + ] + }, + } + ] diff --git a/task-sdk/tests/task_sdk/execution_time/test_coordinator.py b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py index c4185c1939092..cfca1e95baaba 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_coordinator.py +++ b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py @@ -37,6 +37,8 @@ reset_coordinator_manager, ) +from tests_common.test_utils.config import conf_vars + class _CoordinatorA(BaseCoordinator): def __init__(self, *, label: str = "a"): @@ -134,6 +136,76 @@ def test_from_config_rejects_invalid_queue_mapping(self, sdk_config): ): CoordinatorManager.from_config() + @pytest.mark.parametrize("bundle_name", ["ghost", ["ghost"]], ids=["unconfigured", "not-a-string"]) + def test_from_config_rejects_invalid_task_handler_bundle_name(self, sdk_config, bundle_name): + """Validated without importing or constructing the coordinator.""" + sdk_config( + coordinators=json.dumps( + { + "boom": { + "classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator", + "kwargs": {"task_handler_bundle_name": bundle_name}, + } + } + ), + queue_to_coordinator=json.dumps({"queue-boom": "boom"}), + ) + with pytest.raises(InvalidCoordinatorError, match=r"'boom' sets task_handler_bundle_name="): + CoordinatorManager.from_config() + + def test_from_config_accepts_configured_task_handler_bundle_name(self, sdk_config): + sdk_config( + coordinators=json.dumps( + { + "boom": { + "classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator", + "kwargs": {"task_handler_bundle_name": "handlers"}, + } + } + ), + queue_to_coordinator=json.dumps({"queue-boom": "boom"}), + ) + bundles = [ + { + "name": "handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {}, + } + ] + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(bundles)}): + manager = CoordinatorManager.from_config() + assert manager._queue_to_coordinator == {"queue-boom": "boom"} + assert manager._created_coordinators == {} + + def test_from_config_ignores_task_handler_bundle_name_of_unrouted_coordinator(self, sdk_config): + sdk_config( + coordinators=json.dumps( + { + "boom": { + "classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator", + "kwargs": {"task_handler_bundle_name": "ghost"}, + } + } + ), + ) + manager = CoordinatorManager.from_config() + assert set(manager._coordinator_specs) == {"boom"} + + @mock.patch( + "airflow.sdk.execution_time.coordinator.DagBundlesManager.is_bundle_configured", autospec=True + ) + def test_from_config_skips_bundle_lookup_without_task_handler_bundle_name( + self, is_bundle_configured, sdk_config + ): + sdk_config( + coordinators=json.dumps( + {"alpha": {"classpath": f"{_CoordinatorA.__module__}._CoordinatorA", "kwargs": {}}} + ), + queue_to_coordinator=json.dumps({"queue-a": "alpha"}), + ) + CoordinatorManager.from_config() + is_bundle_configured.assert_not_called() + @pytest.mark.parametrize( ("coordinator_spec", "expected_match"), [ @@ -172,6 +244,53 @@ def test_for_queue_raises_when_spec_missing(self): ): manager.for_queue("queue-bad") + def test_get_coordinator_builds_the_keyed_coordinator_once(self, sdk_config): + sdk_config( + coordinators=json.dumps( + { + "alpha": { + "classpath": f"{_CoordinatorA.__module__}._CoordinatorA", + "kwargs": {"label": "x"}, + } + } + ), + ) + manager = CoordinatorManager.from_config() + + coordinator = manager.get_coordinator("alpha") + + assert isinstance(coordinator, _CoordinatorA) + assert coordinator.label == "x" + assert manager.get_coordinator("alpha") is coordinator + + @pytest.mark.parametrize( + ("coordinators", "expected_match"), + [ + pytest.param({}, r"No coordinator 'bad-coord' in \[sdk\] coordinators", id="not-configured"), + pytest.param( + {"bad-coord": {"classpath": "nonexistent.module.MissingCoordinator"}}, + r"Cannot import coordinator 'bad-coord'", + id="classpath-cannot-be-imported", + ), + pytest.param( + { + "bad-coord": { + "classpath": f"{_CoordinatorA.__module__}._CoordinatorA", + "kwargs": {"unknown_param": "x"}, + } + }, + r"Cannot instantiate coordinator 'bad-coord'", + id="kwargs-dont-match-constructor", + ), + ], + ) + def test_get_coordinator_raises_invalid_coordinator_error(self, sdk_config, coordinators, expected_match): + sdk_config(coordinators=json.dumps(coordinators)) + manager = CoordinatorManager.from_config() + + with pytest.raises(InvalidCoordinatorError, match=expected_match): + manager.get_coordinator("bad-coord") + def test_get_coordinator_manager_is_cached(self, monkeypatch): monkeypatch.delenv("AIRFLOW__SDK__COORDINATORS", raising=False) @@ -257,12 +376,62 @@ def test_extra_for_queue_does_not_instantiate_coordinator(self, sdk_config): } assert manager._created_coordinators == {} + def test_get_task_handler_bundle_names_does_not_instantiate_coordinator(self, sdk_config): + sdk_config( + coordinators=json.dumps( + { + "named": { + "classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator", + "kwargs": {"task_handler_bundle_name": "java-task-handlers"}, + }, + "own-bundle": {"classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator"}, + } + ), + queue_to_coordinator=json.dumps({"queue-named": "named", "queue-own-bundle": "own-bundle"}), + ) + bundles = [ + { + "name": "java-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {}, + } + ] + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(bundles)}): + manager = CoordinatorManager.from_config() + + assert manager.get_task_handler_bundle_names() == { + "named": "java-task-handlers", + "own-bundle": None, + } + assert manager._created_coordinators == {} + + @pytest.mark.parametrize( + "kwargs", + [{"task_handler_bundle_name": "ghost"}, {"task_handler_bundle_name": ["ghost"]}, {}], + ids=["unconfigured", "not-a-string", "own-bundle"], + ) + def test_get_task_handler_bundle_names_leaves_out_unrouted_coordinator(self, sdk_config, kwargs): + sdk_config( + coordinators=json.dumps( + { + "routed": {"classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator"}, + "unrouted": { + "classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator", + "kwargs": kwargs, + }, + } + ), + queue_to_coordinator=json.dumps({"queue-routed": "routed"}), + ) + + assert CoordinatorManager.from_config().get_task_handler_bundle_names() == {"routed": None} + class TestConfigYamlCoordinatorsExample: """Guard the ``[sdk] coordinators`` example in ``config.yml`` against drift. - Nothing else exercises the example, so a broken one (e.g. dropping the - required ``jars_root`` kwarg) can ship unnoticed. Loading it through + Nothing else exercises the example, so a broken one (e.g. a kwarg the + coordinator does not accept) can ship unnoticed. Loading it through CoordinatorManager and constructing every entry keeps the example honest. """ @@ -279,7 +448,14 @@ def test_every_example_coordinator_constructs(self, sdk_config): coordinators=coordinators_example, queue_to_coordinator=json.dumps(queue_to_coordinator), ) - manager = CoordinatorManager.from_config() + # The deployment registers the Dag bundles the example names. + artifact_bundles = [ + {"name": name, "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}} + for spec in specs.values() + if (name := spec.get("kwargs", {}).get("task_handler_bundle_name")) + ] + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(artifact_bundles)}): + manager = CoordinatorManager.from_config() assert set(manager._coordinator_specs) == set(specs) for queue, key in queue_to_coordinator.items(): @@ -371,6 +547,17 @@ def test_sigterm_inside_context_does_not_kill(self): assert signal.getsignal(signal.SIGTERM) is signal.SIG_DFL +def test_a_coordinator_parses_no_task_handlers_by_default(tmp_path): + with pytest.raises(NotImplementedError, match="_CoordinatorB does not parse task handlers"): + _CoordinatorB().parse_task_handler( + path=tmp_path / "handlers.artifact", + bundle_path=tmp_path, + comm_address=("127.0.0.1", 1001), + logs_address=("127.0.0.1", 1002), + report_schema_version=lambda schema_version: None, + ) + + class TestPythonCoordinatorWarmShutdown: """The Python coordinator must wrap start() and wait() in the warm-shutdown handlers.""" diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 12daa4173d58f..7b931734bb566 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -172,6 +172,7 @@ WatchedSubprocess, _make_process_nondumpable, _remote_logging_conn, + _reopen_std_io_handles, forward_to_log, in_process_api_server, make_buffered_socket_reader, @@ -5060,6 +5061,25 @@ def on_close(sock): log_write.close() +def test_a_nested_child_keeps_its_standard_streams(): + """A child of a supervised child reopens handles its parent already reopened, and keeps them open.""" + pid = os.fork() + if pid == 0: + code = 1 + try: + for _ in ("supervised child", "its child"): + pairs = [socket.socketpair() for _ in range(3)] + _reopen_std_io_handles(*(child_end for child_end, _ in pairs)) + for fd in (0, 1, 2): + os.fstat(fd) + code = 0 + finally: + os._exit(code) + _, status = os.waitpid(pid, 0) + + assert os.waitstatus_to_exitcode(status) == 0 + + def test_reinit_supervisor_comms(monkeypatch, client_with_ti_start, caplog): def subprocess_main(): # This is run in the subprocess! diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor_schemas_name_type_sync.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor_schemas_name_type_sync.py index e1c3ef1311860..3ad92b430542e 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor_schemas_name_type_sync.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor_schemas_name_type_sync.py @@ -33,7 +33,7 @@ class by ``__name__``. If the two strings drift, a frame decodes against import pytest from pydantic import BaseModel -from airflow.dag_processing.processor import ToDagProcessor, ToManager +from airflow.dag_processing.processor import ToDagProcessor, ToManager, ToSDKTaskHandlerProcessor from airflow.jobs.triggerer_job_runner import ToTriggerRunner, ToTriggerSupervisor from airflow.sdk.execution_time.comms import ToSupervisor, ToTask @@ -45,7 +45,7 @@ def _members_of_union(union: object) -> tuple[type[BaseModel], ...]: return tuple(m for m in get_args(union) if isinstance(m, type) and issubclass(m, BaseModel)) -# All six supervisor discriminated unions. Triggerer's two unions are +# All seven supervisor discriminated unions. Triggerer's two unions are # not part of the lang-SDK-facing registry, but the same name/type # invariant is required for ``CommsDecoder`` to round-trip them. SUPERVISOR_UNIONS = [ @@ -53,6 +53,7 @@ def _members_of_union(union: object) -> tuple[type[BaseModel], ...]: pytest.param(ToSupervisor, id="ToSupervisor"), pytest.param(ToManager, id="ToManager"), pytest.param(ToDagProcessor, id="ToDagProcessor"), + pytest.param(ToSDKTaskHandlerProcessor, id="ToSDKTaskHandlerProcessor"), pytest.param(ToTriggerRunner, id="ToTriggerRunner"), pytest.param(ToTriggerSupervisor, id="ToTriggerSupervisor"), ] diff --git a/ts-sdk/README.md b/ts-sdk/README.md index 361e987fa3116..93fc2db0f8007 100644 --- a/ts-sdk/README.md +++ b/ts-sdk/README.md @@ -88,18 +88,35 @@ sales_pipeline() Airflow coordinator config: ```ini +[dag_processor] +dag_bundle_config_list = [ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + { + "name": "ts-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/opt/airflow/ts-bundles"} + } + ] + [sdk] coordinators = { "ts": { "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", - "kwargs": {"bundles_root": ["/opt/airflow/ts-bundles"]} + "kwargs": {"task_handler_bundle_name": "ts-task-handlers"} } } queue_to_coordinator = {"typescript": "ts"} ``` -Each configured bundle directory is searched recursively for `*.min.mjs` bundles built with `airflow-ts-pack` -(see [Packing bundles](#packing-bundles)), which embeds the Airflow metadata in the bundle itself. +The Dag bundle named by `task_handler_bundle_name` (the task's own Dag bundle when unset) is searched +recursively for `*.min.mjs` bundles built with `airflow-ts-pack` (see [Packing bundles](#packing-bundles)), +which embeds the Airflow metadata in the bundle itself. `task_handler_bundle_name` is used only by +mixed-language Dags, to locate the task handlers for the `@task.stub` tasks of a Python Dag; Dags defined +natively in a language SDK do not use it. + +This config, the packed bundles and Node.js 22 or later must be present wherever tasks execute and on the +Dag processor, which checks the stub tasks of each Python Dag against the task handlers the bundles +register. The API server does not need them. TypeScript entrypoint: @@ -270,7 +287,9 @@ Its values are Python names, which `tsc` cannot see and does not check. Airflow launches the bundled entrypoint with `--comm=host:port` and `--logs=host:port`. `bundle.serve()` connects to those sockets, receives the task startup message, finds the registered handler for the Dag/task pair, and -reports the terminal task state back to Airflow. +reports the terminal task state back to Airflow. Airflow can also start the bundle +to ask which task handlers it registers; `bundle.serve()` then answers with each +handler's task ID, without running any handler. See [`example/`](https://github.com/apache/airflow/tree/main/ts-sdk/example) for a coordinator-runtime example that packs a bundle with `airflow-ts-pack` and diff --git a/ts-sdk/example/README.md b/ts-sdk/example/README.md index 94c2ead395110..85b5c74099f80 100644 --- a/ts-sdk/example/README.md +++ b/ts-sdk/example/README.md @@ -58,14 +58,22 @@ ts-sdk/example/dist/ ## Airflow Configuration -Configure Airflow to route the `typescript` queue to the Node coordinator and -point it at the example bundle directory: +Register the example bundle directory as a Dag bundle, and configure Airflow to +route the `typescript` queue to the Node coordinator that reads it: ```bash +export AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='[ + {"name": "dags-folder", "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}}, + { + "name": "ts-task-handlers", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/absolute/path/to/airflow/ts-sdk/example/dist"} + } +]' export AIRFLOW__SDK__COORDINATORS='{ "ts": { "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", - "kwargs": {"bundles_root": ["/absolute/path/to/airflow/ts-sdk/example/dist"]} + "kwargs": {"task_handler_bundle_name": "ts-task-handlers"} } }' export AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{"typescript": "ts"}' diff --git a/ts-sdk/src/coordinator/comm-channel.ts b/ts-sdk/src/coordinator/comm-channel.ts index 2c8eeb0a908b3..cc0f629fc4ea3 100644 --- a/ts-sdk/src/coordinator/comm-channel.ts +++ b/ts-sdk/src/coordinator/comm-channel.ts @@ -22,8 +22,8 @@ // // The channel is the sole reader on the socket. The task sends // requests and awaits id-correlated replies. The supervisor's only -// unprompted frame is the greeting (StartupDetails / -// DagFileParseRequest), pre-caught into the `greeting` promise that +// unprompted frame is the greeting (StartupDetails, DagFileParseRequest +// or TaskHandlerParseRequest), pre-caught into the `greeting` promise that // `connect()` awaits — the protocol sends nothing else // supervisor-initiated (comms.py: "No messages are sent to task // process except in response to a request"). @@ -35,7 +35,8 @@ import { Deferred } from "./deferred.js"; import type { LogChannel } from "./log-channel.js"; /** What `CommChannel.connect` resolves to: the live channel plus the - * supervisor's first frame (StartupDetails / DagFileParseRequest), + * supervisor's first frame (StartupDetails, DagFileParseRequest or + * TaskHandlerParseRequest), * already in hand so the caller never has to manage a "frame arrived * with no consumer" window. */ export interface CommConnection { diff --git a/ts-sdk/src/coordinator/protocol.ts b/ts-sdk/src/coordinator/protocol.ts index 323765036ea19..312b8ba34997e 100644 --- a/ts-sdk/src/coordinator/protocol.ts +++ b/ts-sdk/src/coordinator/protocol.ts @@ -32,6 +32,7 @@ import type { SucceedTask, TaskState, } from "../generated/supervisor.js"; +import type { TaskHandlerParseRequest, TaskHandlerParsingResult } from "../generated/supervisor.js"; export { SUPERVISOR_API_VERSION } from "../generated/supervisor.js"; @@ -75,6 +76,7 @@ type WithRequiredType = Omit & { export type MsgFromSupervisor = | WithRequiredType | WithRequiredType + | WithRequiredType | WithRequiredType; // -------- Frames to supervisor -------- @@ -86,6 +88,7 @@ export type RuntimeSucceedTask = WithRequiredType & { outlet_events: NonNullable; }; export type RuntimeDagFileParsingResult = WithRequiredType; +export type RuntimeTaskHandlerParsingResult = WithRequiredType; // -------- Decoder: raw map → typed message -------- @@ -94,6 +97,7 @@ export function asMsgFromSupervisor(raw: unknown): MsgFromSupervisor { switch (body.type) { case "StartupDetails": case "DagFileParseRequest": + case "TaskHandlerParseRequest": case "ErrorResponse": return body as unknown as MsgFromSupervisor; default: diff --git a/ts-sdk/src/coordinator/runtime.ts b/ts-sdk/src/coordinator/runtime.ts index bf9764f3e6618..455cf3ab19b21 100644 --- a/ts-sdk/src/coordinator/runtime.ts +++ b/ts-sdk/src/coordinator/runtime.ts @@ -33,6 +33,7 @@ // 1. Parse --comm / --logs from argv // 2. Connect both TCP sockets // 3. Read the first frame from comm: +// - TaskHandlerParseRequest → send TaskHandlerParsingResult, await its ack, exit // - DagFileParseRequest → respond with DagFileParsingResult, exit // - StartupDetails → run task, respond Succeed or Fail, exit // @@ -40,6 +41,7 @@ import { resolveArgs, type BoundArgs } from "./arg-binding.js"; import { createCoordinatorClient } from "./client.js"; import { CommChannel } from "./comm-channel.js"; import { LogChannel } from "./log-channel.js"; +import { declareTaskHandlers } from "./task-handler-parse.js"; import { AIRFLOW_METADATA_FLAG, AIRFLOW_METADATA_SENTINEL, @@ -165,7 +167,14 @@ export async function startCoordinator( const body = asMsgFromSupervisor(firstFrame.body); - if (body.type === "DagFileParseRequest") { + if (body.type === "TaskHandlerParseRequest") { + runtimeLogs.info("Received task handler parse request", { file: body.file }); + const result = declareTaskHandlers(bundle, body); + await sendParseResult(result, comm, runtimeLogs); + runtimeLogs.info("Declared task handlers", { + dag_ids: Object.keys(result.task_handlers), + }); + } else if (body.type === "DagFileParseRequest") { runtimeLogs.info("Received Dag parse request", { file: body.file, bundle_path: body.bundle_path, @@ -194,7 +203,7 @@ export async function startCoordinator( runtimeLogs.info("Task succeeded", { task_id: body.ti.task_id }); } } else { - const errMsg = `First frame must be DagFileParseRequest or StartupDetails, got ${body.type}`; + const errMsg = `First frame must be TaskHandlerParseRequest, DagFileParseRequest or StartupDetails, got ${body.type}`; runtimeLogs.error("Unexpected first frame", { type: body.type }); await sendSupervisorResponse(firstFrame.id, null, comm, runtimeLogs, { error: "protocol_error", @@ -382,6 +391,26 @@ async function sendSupervisorResponse( } } +/** + * Send a parse result as a request and wait for Airflow's acknowledgement, so + * the process never exits before the result is recorded. An acknowledgement + * carrying an error, including the socket closing first, throws. + * + * Unlike {@link sendSupervisorResponse}, which answers the first frame's id + * and expects no reply, this sends a new request and awaits its answer. + */ +async function sendParseResult(body: unknown, comm: CommChannel, logs: LogChannel): Promise { + const ack = await comm.request(body, { timeoutMs: COORDINATOR_RESPONSE_TIMEOUT_MS }); + if (ack.error == null) return; + logs.error("Parse result was not acknowledged", { + response_type: responseType(body), + error: ack.error, + }); + throw new Error( + `Airflow did not acknowledge the ${responseType(body)}: ${JSON.stringify(ack.error)}`, + ); +} + function buildFailureResponse( details: StartupDetails, message: string, diff --git a/ts-sdk/src/coordinator/task-handler-parse.ts b/ts-sdk/src/coordinator/task-handler-parse.ts new file mode 100644 index 0000000000000..dfc5320c467db --- /dev/null +++ b/ts-sdk/src/coordinator/task-handler-parse.ts @@ -0,0 +1,55 @@ +/*! + * 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. + */ + +// The answer to a TaskHandlerParseRequest: every task handler this bundle +// registers, as the Dag processor checks them against a Python file's +// `@task.stub` calls. + +import type { TaskHandlerDeclaration, TaskHandlerParseRequest } from "../generated/supervisor.js"; +import type { RuntimeTaskHandlerParsingResult } from "./protocol.js"; +import { bundleTaskHandlers, type Bundle } from "../sdk/bundle.js"; + +/** + * Declare every registered task handler, keyed by Dag id, each Dag's in + * registration order. A Dag declared in TypeScript is never declared. No + * handler runs. + */ +export function declareTaskHandlers( + bundle: Bundle, + request: Pick, +): RuntimeTaskHandlerParsingResult { + const declared: [string, TaskHandlerDeclaration[]][] = []; + for (const [dagId, handlers] of bundleTaskHandlers(bundle)) { + declared.push([dagId, [...handlers.keys()].map((taskId) => declareTaskHandler(taskId))]); + } + return { + type: "TaskHandlerParsingResult", + fileloc: request.file, + // Defined rather than assigned, so a Dag named `__proto__` stays a key. + task_handlers: Object.fromEntries(declared), + }; +} + +/** + * Types are erased, so the SDK cannot list the params a handler takes, and the + * Dag processor checks only that the handler exists. + */ +function declareTaskHandler(taskId: string): TaskHandlerDeclaration { + return { task_id: taskId, binding: "named", params: null }; +} diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 76e9e727815cf..af0c907a5ed05 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -269,6 +269,16 @@ export type Msg2 = string | null; export type EmailType = "failure" | "retry"; export type Type14 = "EmailRequest"; export type CallbackRequests = (DagCallbackRequest | TaskCallbackRequest | EmailRequest)[]; +export type BundleName4 = string; +export type RelativeFileloc = string; +export type SizeBytes = number; +export type CacheDigest = string | null; +export type TaskId2 = string; +export type Binding = "positional" | "named"; +export type TaskHandlerParams = TaskHandlerParam[] | null; +export type Name10 = string | null; +export type ExactName = boolean; +export type KnownArtifacts = TaskHandlerArtifact[]; export type Type15 = "DagFileParseRequest"; export type Fileloc = string; export type LastLoaded = string | null; @@ -283,9 +293,9 @@ export type Language = string; export type Type16 = "DagFileParsingResult"; export type DagId4 = string; export type IsPaused = boolean; -export type BundleName4 = string | null; +export type BundleName5 = string | null; export type BundleVersion3 = string | null; -export type RelativeFileloc = string | null; +export type RelativeFileloc1 = string | null; export type Owners = string | null; export type Tags = string[]; export type NextDagrun = string | null; @@ -326,7 +336,7 @@ export type NextKwargs2 = { } | null; export type RenderedMapIndex1 = string | null; export type Type20 = "DeferTask"; -export type Name10 = string; +export type Name11 = string; export type Key1 = string; export type Type21 = "DeleteAssetStateStoreByName"; export type Uri5 = string; @@ -340,7 +350,7 @@ export type Type24 = "DeleteVariable"; export type Key5 = string; export type DagId6 = string; export type RunId5 = string; -export type TaskId2 = string; +export type TaskId3 = string; export type MapIndex1 = number | null; export type Type25 = "DeleteXCom"; /** @@ -378,11 +388,11 @@ export type ErrorType1 = | "PERMISSION_DENIED" | "GENERIC_ERROR" | "API_SERVER_ERROR"; -export type Name11 = string; +export type Name12 = string; export type Type27 = "GetAssetByName"; export type Uri6 = string; export type Type28 = "GetAssetByUri"; -export type Name12 = string | null; +export type Name13 = string | null; export type Uri7 = string | null; export type After = string | null; export type Before = string | null; @@ -405,7 +415,7 @@ export type Extra8 = { [k: string]: string; } | null; export type Type30 = "GetAssetEventByAssetAlias"; -export type Name13 = string; +export type Name14 = string; export type Key6 = string; export type Type31 = "GetAssetStateStoreByName"; export type Uri8 = string; @@ -437,7 +447,7 @@ export type LogicalDate3 = string; export type State3 = string | null; export type Type41 = "GetPreviousDagRun"; export type DagId12 = string; -export type TaskId3 = string; +export type TaskId4 = string; export type LogicalDate4 = string | null; export type MapIndex2 = number; export type Type42 = "GetPreviousTI"; @@ -474,25 +484,25 @@ export type Type49 = "GetVariableKeys"; export type Key10 = string; export type DagId16 = string; export type RunId9 = string; -export type TaskId4 = string; +export type TaskId5 = string; export type MapIndex5 = number | null; export type IncludePriorDates = boolean; export type Type50 = "GetXCom"; export type Key11 = string; export type DagId17 = string; export type RunId10 = string; -export type TaskId5 = string; +export type TaskId6 = string; export type Type51 = "GetXComCount"; export type Key12 = string; export type DagId18 = string; export type RunId11 = string; -export type TaskId6 = string; +export type TaskId7 = string; export type Offset1 = number; export type Type52 = "GetXComSequenceItem"; export type Key13 = string; export type DagId19 = string; export type RunId12 = string; -export type TaskId7 = string; +export type TaskId8 = string; export type Start = number | null; export type Stop = number | null; export type Step = number | null; @@ -514,7 +524,7 @@ export type AssignedUsers1 = HITLUser[] | null; export type Type54 = "HITLDetailRequestResult"; export type InactiveAssets = AssetProfile[] | null; export type Type55 = "InactiveAssetsResult"; -export type Name14 = string | null; +export type Name15 = string | null; export type Type56 = "MaskSecret"; export type Ok = boolean; export type Type57 = "OKResponse"; @@ -524,7 +534,7 @@ export type StartDate4 = string | null; export type EndDate3 = string | null; export type Type58 = "PrevSuccessfulDagRunResult"; export type Type59 = "PreviousDagRunResult"; -export type TaskId8 = string; +export type TaskId9 = string; export type DagId20 = string; export type RunId13 = string; export type LogicalDate5 = string | null; @@ -552,7 +562,7 @@ export type RetryReason = string | null; export type Type64 = "RetryTask"; export type Type65 = "SentFDs"; export type Fds = number[]; -export type Name15 = string; +export type Name16 = string; export type Key15 = string; export type Type66 = "SetAssetStateStoreByName"; export type Uri9 = string; @@ -568,7 +578,7 @@ export type Type70 = "SetTaskStateStore"; export type Key18 = string; export type DagId21 = string; export type RunId14 = string; -export type TaskId9 = string; +export type TaskId10 = string; export type MapIndex7 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; @@ -595,15 +605,25 @@ export type Breadcrumbs = { [k: string]: unknown; }[]; export type Type76 = "TaskBreadcrumbsResult"; +export type File1 = string; +export type BundlePath1 = string; +export type BundleName6 = string; +export type Type77 = "TaskHandlerParseRequest"; +export type Fileloc1 = string; +export type ImportErrors1 = { + [k: string]: string; +} | null; +export type Warnings1 = unknown[] | null; +export type Type78 = "TaskHandlerParsingResult"; export type StartDate7 = string | null; -export type Type77 = "TaskRescheduleStartDate"; +export type Type79 = "TaskRescheduleStartDate"; export type State8 = "failed" | "skipped" | "removed"; export type EndDate8 = string | null; -export type Type78 = "TaskState"; +export type Type80 = "TaskState"; export type RenderedMapIndex5 = string | null; export type RetryReason1 = string | null; -export type Type79 = "TaskStateStoreResult"; -export type Type80 = "TaskStatesResult"; +export type Type81 = "TaskStateStoreResult"; +export type Type82 = "TaskStatesResult"; export type LogicalDate6 = string | null; export type RunAfter2 = string | null; export type Conf2 = { @@ -614,10 +634,10 @@ export type PartitionKey7 = string | null; export type Note2 = string | null; export type DagId22 = string; export type DagRunId = string; -export type Type81 = "TriggerDagRun"; +export type Type83 = "TriggerDagRun"; export type TiId9 = string; export type Note3 = string | null; -export type Type82 = "UpdateDagRunNote"; +export type Type84 = "UpdateDagRunNote"; export type TiId10 = string; /** * @minItems 1 @@ -626,22 +646,22 @@ export type ChosenOptions = [string, ...string[]]; export type ParamsInput = { [k: string]: unknown; } | null; -export type Type83 = "UpdateHITLDetail"; +export type Type85 = "UpdateHITLDetail"; export type TiId11 = string; -export type Type84 = "ValidateInletsAndOutlets"; +export type Type86 = "ValidateInletsAndOutlets"; export type Keys = string[]; export type TotalEntries = number; -export type Type85 = "VariableKeysResult"; +export type Type87 = "VariableKeysResult"; export type Key19 = string; export type Value2 = string | null; -export type Type86 = "VariableResult"; +export type Type88 = "VariableResult"; export type Len = number; -export type Type87 = "XComCountResponse"; +export type Type89 = "XComCountResponse"; export type Key20 = string; -export type Type88 = "XComResult"; -export type Type89 = "XComSequenceIndexResult"; +export type Type90 = "XComResult"; +export type Type91 = "XComSequenceIndexResult"; export type Root = JsonValue[]; -export type Type90 = "XComSequenceSliceResult"; +export type Type92 = "XComSequenceSliceResult"; export interface SupervisorWireSchema {} /** @@ -1000,6 +1020,7 @@ export interface DagFileParseRequest { bundle_path: BundlePath; bundle_name: BundleName1; callback_requests?: CallbackRequests; + known_artifacts?: KnownArtifacts; type?: Type15; } /** @@ -1101,6 +1122,44 @@ export interface EmailRequest { context_from_server: TIRunContext; type?: Type14; } +/** + * A Lang-SDK artifact and every task handler it registers. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerArtifact". + */ +export interface TaskHandlerArtifact { + bundle_name: BundleName4; + relative_fileloc: RelativeFileloc; + size_bytes: SizeBytes; + cache_digest: CacheDigest; + task_handlers: TaskHandlers; +} +export interface TaskHandlers { + [k: string]: TaskHandlerDeclaration[]; +} +/** + * A task handler that a Lang-SDK artifact registers for one task. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerDeclaration". + */ +export interface TaskHandlerDeclaration { + task_id: TaskId2; + binding: Binding; + params: TaskHandlerParams; +} +/** + * One parameter of a task handler. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerParam". + */ +export interface TaskHandlerParam { + name: Name10; + value_schema?: ArgValueSchema | null; + exact_name?: ExactName; +} /** * Result of DAG File Parsing. * @@ -1155,9 +1214,9 @@ export interface DagSourceCode { export interface DagResult { dag_id: DagId4; is_paused: IsPaused; - bundle_name: BundleName4; + bundle_name: BundleName5; bundle_version: BundleVersion3; - relative_fileloc: RelativeFileloc; + relative_fileloc: RelativeFileloc1; owners: Owners; tags: Tags; next_dagrun: NextDagrun; @@ -1218,7 +1277,7 @@ export interface DeferTask { * via the `definition` "DeleteAssetStateStoreByName". */ export interface DeleteAssetStateStoreByName { - name: Name10; + name: Name11; key: Key1; type?: Type21; } @@ -1256,7 +1315,7 @@ export interface DeleteXCom { key: Key5; dag_id: DagId6; run_id: RunId5; - task_id: TaskId2; + task_id: TaskId3; map_index?: MapIndex1; type?: Type25; } @@ -1274,7 +1333,7 @@ export interface ErrorResponse { * via the `definition` "GetAssetByName". */ export interface GetAssetByName { - name: Name11; + name: Name12; type?: Type27; } /** @@ -1290,7 +1349,7 @@ export interface GetAssetByUri { * via the `definition` "GetAssetEventByAsset". */ export interface GetAssetEventByAsset { - name: Name12; + name: Name13; uri: Uri7; after?: After; before?: Before; @@ -1321,7 +1380,7 @@ export interface GetAssetEventByAssetAlias { * via the `definition` "GetAssetStateStoreByName". */ export interface GetAssetStateStoreByName { - name: Name13; + name: Name14; key: Key6; type?: Type31; } @@ -1423,7 +1482,7 @@ export interface GetPreviousDagRun { */ export interface GetPreviousTI { dag_id: DagId12; - task_id: TaskId3; + task_id: TaskId4; logical_date?: LogicalDate4; map_index?: MapIndex2; state?: TaskInstanceState | null; @@ -1509,7 +1568,7 @@ export interface GetXCom { key: Key10; dag_id: DagId16; run_id: RunId9; - task_id: TaskId4; + task_id: TaskId5; map_index?: MapIndex5; include_prior_dates?: IncludePriorDates; type?: Type50; @@ -1524,7 +1583,7 @@ export interface GetXComCount { key: Key11; dag_id: DagId17; run_id: RunId10; - task_id: TaskId5; + task_id: TaskId6; type?: Type51; } /** @@ -1535,7 +1594,7 @@ export interface GetXComSequenceItem { key: Key12; dag_id: DagId18; run_id: RunId11; - task_id: TaskId6; + task_id: TaskId7; offset: Offset1; type?: Type52; } @@ -1547,7 +1606,7 @@ export interface GetXComSequenceSlice { key: Key13; dag_id: DagId19; run_id: RunId12; - task_id: TaskId7; + task_id: TaskId8; start: Start; stop: Stop; step: Step; @@ -1589,7 +1648,7 @@ export interface InactiveAssetsResult { */ export interface MaskSecret { value: JsonValue; - name?: Name14; + name?: Name15; type?: Type56; } /** @@ -1628,7 +1687,7 @@ export interface PreviousDagRunResult { * via the `definition` "PreviousTIResponse". */ export interface PreviousTIResponse { - task_id: TaskId8; + task_id: TaskId9; dag_id: DagId20; run_id: RunId13; logical_date?: LogicalDate5; @@ -1705,7 +1764,7 @@ export interface SentFDs { * via the `definition` "SetAssetStateStoreByName". */ export interface SetAssetStateStoreByName { - name: Name15; + name: Name16; key: Key15; value: JsonValue; type?: Type66; @@ -1763,7 +1822,7 @@ export interface SetXCom { value: JsonValue; dag_id: DagId21; run_id: RunId14; - task_id: TaskId9; + task_id: TaskId10; map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; @@ -1824,6 +1883,40 @@ export interface TaskBreadcrumbsResult { breadcrumbs: Breadcrumbs; type?: Type76; } +/** + * Request for Task Handler Parsing. + * + * Asks a Lang-SDK runtime for every task handler an artifact registers. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerParseRequest". + */ +export interface TaskHandlerParseRequest { + file: File1; + bundle_path: BundlePath1; + bundle_name: BundleName6; + type?: Type77; +} +/** + * Result of Task Handler Parsing. + * + * Every task handler a Lang-SDK artifact registers, keyed by Dag id. + * + * The answer depends only on the artifact, never on the request. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerParsingResult". + */ +export interface TaskHandlerParsingResult { + fileloc: Fileloc1; + task_handlers: TaskHandlers1; + import_errors?: ImportErrors1; + warnings?: Warnings1; + type?: Type78; +} +export interface TaskHandlers1 { + [k: string]: TaskHandlerDeclaration[]; +} /** * Response containing the first reschedule date for a task instance. * @@ -1832,7 +1925,7 @@ export interface TaskBreadcrumbsResult { */ export interface TaskRescheduleStartDate { start_date: StartDate7; - type?: Type77; + type?: Type79; } /** * Update a task's state. @@ -1847,7 +1940,7 @@ export interface TaskRescheduleStartDate { export interface TaskState { state: State8; end_date?: EndDate8; - type?: Type78; + type?: Type80; rendered_map_index?: RenderedMapIndex5; retry_reason?: RetryReason1; } @@ -1859,7 +1952,7 @@ export interface TaskState { */ export interface TaskStateStoreResult { value: JsonValue | null; - type?: Type79; + type?: Type81; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1867,7 +1960,7 @@ export interface TaskStateStoreResult { */ export interface TaskStatesResult { task_states: TaskStates; - type?: Type80; + type?: Type82; } export interface TaskStates { [k: string]: unknown; @@ -1885,7 +1978,7 @@ export interface TriggerDagRun { note?: Note2; dag_id: DagId22; run_id: DagRunId; - type?: Type81; + type?: Type83; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1894,7 +1987,7 @@ export interface TriggerDagRun { export interface UpdateDagRunNote { ti_id: TiId9; note: Note3; - type?: Type82; + type?: Type84; } /** * Update the response content part of an existing Human-in-the-loop response. @@ -1906,7 +1999,7 @@ export interface UpdateHITLDetail { ti_id: TiId10; chosen_options: ChosenOptions; params_input?: ParamsInput; - type?: Type83; + type?: Type85; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1914,7 +2007,7 @@ export interface UpdateHITLDetail { */ export interface ValidateInletsAndOutlets { ti_id: TiId11; - type?: Type84; + type?: Type86; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1923,7 +2016,7 @@ export interface ValidateInletsAndOutlets { export interface VariableKeysResult { keys: Keys; total_entries: TotalEntries; - type?: Type85; + type?: Type87; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1932,7 +2025,7 @@ export interface VariableKeysResult { export interface VariableResult { key: Key19; value: Value2; - type?: Type86; + type?: Type88; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1940,7 +2033,7 @@ export interface VariableResult { */ export interface XComCountResponse { len: Len; - type?: Type87; + type?: Type89; } /** * Response to ReadXCom request. @@ -1951,7 +2044,7 @@ export interface XComCountResponse { export interface XComResult { key: Key20; value: JsonValue | null; - type?: Type88; + type?: Type90; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1959,7 +2052,7 @@ export interface XComResult { */ export interface XComSequenceIndexResult { root: JsonValue; - type?: Type89; + type?: Type91; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1967,7 +2060,7 @@ export interface XComSequenceIndexResult { */ export interface XComSequenceSliceResult { root: Root; - type?: Type90; + type?: Type92; } /** Cadwyn schema version this SDK was generated against. diff --git a/ts-sdk/src/sdk/bundle.ts b/ts-sdk/src/sdk/bundle.ts index a3e54f1a23172..4ae41cfe16ed4 100644 --- a/ts-sdk/src/sdk/bundle.ts +++ b/ts-sdk/src/sdk/bundle.ts @@ -226,6 +226,15 @@ export function bundleDags(bundle: Bundle): ReadonlyMap { return dagsOf(bundle); } +/** Internal: the handlers this bundle registered for tasks a Python Dag + * declares, keyed by dag_id and then task_id, both in registration order. + * Excludes the Dags declared in TypeScript. */ +export function bundleTaskHandlers( + bundle: Bundle, +): ReadonlyMap> { + return taskHandlersOf(bundle); +} + /** Internal: the task IDs this bundle can dispatch, per Dag: the Dags declared * in TypeScript first, then the Dags its task handlers name, each in * registration order. A Dag declared in TypeScript with no tasks is included. */ diff --git a/ts-sdk/tests/coordinator/integration.test.ts b/ts-sdk/tests/coordinator/integration.test.ts index d0b6d53f83fff..5aa1894594436 100644 --- a/ts-sdk/tests/coordinator/integration.test.ts +++ b/ts-sdk/tests/coordinator/integration.test.ts @@ -81,8 +81,8 @@ type Responder = ( body: Record, ) => { body: unknown; error?: unknown } | null; -function frameBytes(id: number, body: unknown, isResponse: boolean): Buffer { - const arr = isResponse ? [id, body, null] : [id, body]; +function frameBytes(id: number, body: unknown, isResponse: boolean, error: unknown = null): Buffer { + const arr = isResponse ? [id, body, error] : [id, body]; const payload = Buffer.from(encode(arr)); const header = Buffer.alloc(4); header.writeUInt32BE(payload.length, 0); @@ -203,7 +203,7 @@ async function driveSupervisor(initialFrame: unknown, responder?: Responder): Pr // runtime never hangs on auto-generated RPCs (e.g. the // return_value XCom push). { body: null }; - commSock.write(frameBytes(f.id, reply.body, true)); + commSock.write(frameBytes(f.id, reply.body, true, reply.error)); continue; } // Arity-3: a response from the runtime. The terminal frame @@ -218,10 +218,12 @@ async function driveSupervisor(initialFrame: unknown, responder?: Responder): Pr // Wait for the runtime to finish AND for the comm socket to deliver // its final bytes (FIN signals that all preceding frames are flushed). const commSocketEnded = new Promise((resolve) => commSock.on("end", () => resolve())); - await Promise.all([runtimeDone, commSocketEnded]); - - comm.server.close(); - logs.server.close(); + try { + await Promise.all([runtimeDone, commSocketEnded]); + } finally { + comm.server.close(); + logs.server.close(); + } const lines = Buffer.concat(logChunks).toString("utf8").split("\n").filter(Boolean); const logRecords = lines.map((l) => JSON.parse(l) as Record); @@ -611,6 +613,55 @@ describe("coordinator runtime integration", () => { ).toBe(true); }); + describe("TaskHandlerParseRequest", () => { + const parseRequest = { + type: "TaskHandlerParseRequest", + file: "/bundles/etl.min.mjs", + bundle_path: "/bundles", + bundle_name: "ts-task-handlers", + }; + + it("sends the declarations as a request once, and runs no handler", async () => { + let called = false; + bundle.register( + new TaskHandler("py_dag", "transform", async () => { + called = true; + }), + ); + + const result = await driveSupervisor(parseRequest); + + expect(result.firstResponse).toBeNull(); + expect(result.runtimeRequests).toEqual([ + { + type: "TaskHandlerParsingResult", + body: { + type: "TaskHandlerParsingResult", + fileloc: "/bundles/etl.min.mjs", + task_handlers: { + py_dag: [{ task_id: "transform", binding: "named", params: null }], + }, + }, + }, + ]); + expect(called).toBe(false); + }); + + it("fails when Airflow answers the result with an error", async () => { + bundle.register(new TaskHandler("py_dag", "transform", async () => undefined)); + + await expect( + driveSupervisor(parseRequest, () => ({ + body: null, + error: { + error: "generic_error", + detail: { message: "A parse result was already received" }, + }, + })), + ).rejects.toThrow(/did not acknowledge the TaskHandlerParsingResult/); + }); + }); + it("aborts the context signal on SIGTERM and reports a thrown task error", async () => { let sawAbort = false; testDag.task("aborted_then_failed", async () => { diff --git a/ts-sdk/tests/coordinator/protocol.test.ts b/ts-sdk/tests/coordinator/protocol.test.ts index 12ba947eeb3a2..0da421eb55b6c 100644 --- a/ts-sdk/tests/coordinator/protocol.test.ts +++ b/ts-sdk/tests/coordinator/protocol.test.ts @@ -54,6 +54,16 @@ describe("protocol decode", () => { expect(asMsgFromSupervisor(raw).type).toBe("DagFileParseRequest"); }); + it("accepts TaskHandlerParseRequest", () => { + const raw = { + type: "TaskHandlerParseRequest", + file: "/x.min.mjs", + bundle_path: "/", + bundle_name: "ts-task-handlers", + }; + expect(asMsgFromSupervisor(raw).type).toBe("TaskHandlerParseRequest"); + }); + it("rejects missing type", () => { expect(() => asMsgFromSupervisor({ ti: {} })).toThrow(/missing string 'type'/); }); diff --git a/ts-sdk/tests/coordinator/task-handler-parse.test.ts b/ts-sdk/tests/coordinator/task-handler-parse.test.ts new file mode 100644 index 0000000000000..f8feab803eee0 --- /dev/null +++ b/ts-sdk/tests/coordinator/task-handler-parse.test.ts @@ -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. + */ + +import { describe, expect, it } from "vitest"; + +import { declareTaskHandlers } from "../../src/coordinator/task-handler-parse.js"; +import { withArgNames } from "../../src/sdk/arg-names.js"; +import { Bundle } from "../../src/sdk/bundle.js"; +import { Dag } from "../../src/sdk/dag.js"; +import { TaskHandler } from "../../src/sdk/task-handler.js"; + +const FILE = "/bundles/etl.min.mjs"; + +async function noop(): Promise {} + +function declare(bundle: Bundle) { + return declareTaskHandlers(bundle, { file: FILE }); +} + +describe("declareTaskHandlers", () => { + it("declares every Dag's handlers in registration order", () => { + const bundle = new Bundle( + new TaskHandler("etl", "extract", noop), + new TaskHandler("reports", "send", noop), + new TaskHandler("etl", "load", noop), + ); + + const result = declare(bundle); + + expect(result).toStrictEqual({ + type: "TaskHandlerParsingResult", + fileloc: FILE, + task_handlers: { + etl: [ + { task_id: "extract", binding: "named", params: null }, + { task_id: "load", binding: "named", params: null }, + ], + reports: [{ task_id: "send", binding: "named", params: null }], + }, + }); + expect(Object.keys(result.task_handlers)).toEqual(["etl", "reports"]); + }); + + it("lists no params for a handler with withArgNames renames", () => { + interface ReportArgs { + label: string; + owner: string; + } + const report = withArgNames( + { label: "run_label", owner: "team" }, + async ({ label }: ReportArgs) => label, + ); + + const result = declare(new Bundle(new TaskHandler("etl", "report", report))); + + expect(result.task_handlers).toStrictEqual({ + etl: [{ task_id: "report", binding: "named", params: null }], + }); + }); + + it("does not declare a Dag declared in TypeScript", () => { + const native = new Dag("native"); + native.task("run", noop)(); + + const mixed = new Bundle(native, new TaskHandler("etl", "extract", noop)); + + expect(declare(new Bundle(native)).task_handlers).toEqual({}); + expect(Object.keys(declare(mixed).task_handlers)).toEqual(["etl"]); + }); + + it("keeps a Dag named __proto__ as a key", () => { + const result = declare(new Bundle(new TaskHandler("__proto__", "run", noop))); + + expect(Object.keys(result.task_handlers)).toEqual(["__proto__"]); + expect(Object.getPrototypeOf(result.task_handlers)).toBe(Object.prototype); + }); +});