From bfc85b0a57c4c865f478666276e8e67b20184a91 Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Wed, 7 Oct 2026 13:16:49 +0800 Subject: [PATCH 1/6] Go SDK: answer the Dag-parsing request from Serve Signed-off-by: PoAn Yang --- .../dag_processing/lang_sdk_processor.py | 67 +++++- .../dag_processing/test_lang_sdk_processor.py | 119 +++++++++++ go-sdk/README.md | 5 + go-sdk/airflow/bundle.go | 48 ++++- go-sdk/airflow/bundle_test.go | 42 ++++ go-sdk/airflow/dag.go | 5 +- go-sdk/airflow/serve.go | 8 +- go-sdk/airflow/serve_test.go | 103 ++++++++-- go-sdk/internal/bundle/doc.go | 9 +- go-sdk/internal/bundle/task.go | 24 +++ go-sdk/pkg/execution/dag_parse.go | 78 +++++++ go-sdk/pkg/execution/dag_parse_test.go | 194 ++++++++++++++++++ go-sdk/pkg/execution/integration_test.go | 102 ++++++++- go-sdk/pkg/execution/server.go | 46 ++++- .../airflow/sdk/execution_time/supervisor.py | 8 +- 15 files changed, 818 insertions(+), 40 deletions(-) create mode 100644 go-sdk/pkg/execution/dag_parse.go create mode 100644 go-sdk/pkg/execution/dag_parse_test.go diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py index 9a0998dd05109..57d0aa662c1e8 100644 --- a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -23,6 +23,7 @@ import selectors import signal import time +from collections import deque from contextlib import suppress from pathlib import Path from socket import socket @@ -46,6 +47,7 @@ ) from airflow.exceptions import DeserializationError from airflow.models.pool import Pool +from airflow.sdk._shared.secrets_masker import redact # noqa: SDK001 from airflow.sdk.coordinators._subprocess import _is_connection_from_pid, _start_server from airflow.sdk.execution_time import supervisor from airflow.sdk.execution_time.comms import CommsDecoder, ErrorResponse, _RequestFrame @@ -75,6 +77,15 @@ _IMPORT_TIMEOUT_SETTING = "[core] dagbag_import_timeout or the get_dagbag_import_timeout policy" _PROCESSOR_TIMEOUT_SETTING = "[dag_processor] dag_file_processor_timeout" +# When a runtime exits without a parse result, its import error shows the first and the last lines +# that the runtime wrote to stderr. A Go panic and a Java exception print their message on the first +# line, and a Python traceback prints its message on the last line. Each line is cut to a maximum +# length. Node writes the whole source line that threw an error to stderr, and in a minified +# JavaScript bundle one source line can hold the whole bundle. +_STDERR_HEAD_LINES = 5 +_STDERR_TAIL_LINES = 15 +_STDERR_LINE_LENGTH = 1000 + # StartLangSDKRuntime and LangSDKRuntimeSchemaVersion pass only between the manager and its forked # child before the exec, so they are not part of the supervisor schema the runtimes speak. @@ -155,6 +166,48 @@ def report_schema_version(schema_version: str | None) -> None: _Channel = Literal["comm", "logs"] +@attrs.define +class _StderrExcerpt: + """Keeps the first and the last lines of the runtime's stderr for the import error.""" + + head: list[str] = attrs.field(factory=list) + tail: deque[str] = attrs.field(factory=functools.partial(deque, maxlen=_STDERR_TAIL_LINES)) + omitted: int = 0 + + def add(self, line: bytes | bytearray) -> None: + # The masker finds a secret only when the whole secret is in the text, so mask the line before + # cutting it. + text = cast("str", redact(line.rstrip().decode("utf-8", errors="replace"))) + # Postgres cannot store a NUL character in a text column. + text = text.replace("\x00", "\N{REPLACEMENT CHARACTER}") + if len(text) > _STDERR_LINE_LENGTH: + text = text[:_STDERR_LINE_LENGTH] + "\N{HORIZONTAL ELLIPSIS}" + if len(self.head) < _STDERR_HEAD_LINES: + self.head.append(text) + return + if len(self.tail) == _STDERR_TAIL_LINES: + self.omitted += 1 + self.tail.append(text) + + def render(self) -> str: + lines = "line" if self.omitted == 1 else "lines" + omitted = [f"... {self.omitted} {lines} omitted ..."] if self.omitted else [] + # The runtime can ask the Dag processor to mask a secret after printing the secret. A secret + # can also span several lines. + return cast("str", redact("\n".join([*self.head, *omitted, *self.tail]).strip("\n"))) + + +def _record_lines( + lines: Generator[None, bytes | bytearray, None], excerpt: _StderrExcerpt +) -> Generator[None, bytes | bytearray, None]: + """Pass each line on to *lines* and add the line to *excerpt*.""" + next(lines) + while True: + line = yield + excerpt.add(line) + lines.send(line) + + @attrs.define(kw_only=True) class LangSDKDagFileProcessorProcess(BaseDagFileProcessorProcess): """ @@ -177,6 +230,7 @@ class LangSDKDagFileProcessorProcess(BaseDagFileProcessorProcess): _parsing_result_monotonic: float | None = attrs.field(default=None, init=False) _unverified_connections: list[tuple[socket, _Channel]] = attrs.field(factory=list, init=False) _group_killed: bool = attrs.field(default=False, init=False) + _stderr_excerpt: _StderrExcerpt = attrs.field(factory=_StderrExcerpt, init=False) @classmethod def start( # type: ignore[override] @@ -394,6 +448,12 @@ def _register_logs(self, conn: socket) -> None: ), ) + def _create_line_forwarder( + self, loggers: tuple[FilteringBoundLogger, ...], name: str, level: int + ) -> Generator[None, bytes | bytearray, None]: + lines = super()._create_line_forwarder(loggers, name, level) + return _record_lines(lines, self._stderr_excerpt) if name.endswith(".stderr") else lines + def _set_import_error(self, message: str) -> None: self.parsing_result = DagFileParsingResult( fileloc=self._parse_request.file, @@ -538,9 +598,10 @@ def is_ready(self) -> bool: 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" - ) + message = f"The Lang-SDK runtime exited with code {self._exit_code} without a parse result" + if stderr := self._stderr_excerpt.render(): + message += f". Its stderr output:\n{stderr}" + self._set_import_error(message) return True def _time_out(self, timeout: float, setting: str) -> None: diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index 8be3221b37874..1dfbcc394511b 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -43,11 +43,13 @@ LangSDKDagFileProcessorProcess, LangSDKRuntimeSchemaVersion, _get_import_timeout, + _StderrExcerpt, ) from airflow.dag_processing.processor import DagFileParseRequest, DagFileParsingResult from airflow.exceptions import UnknownExecutorException from airflow.executors.executor_loader import ExecutorLoader from airflow.sdk import DAG, BaseOperator, task +from airflow.sdk._shared.secrets_masker import _secrets_masker as sdk_secrets_masker from airflow.sdk.api.client import Client from airflow.sdk.api.datamodels._generated import VariableResponse from airflow.sdk.exceptions import AirflowRuntimeError @@ -115,6 +117,18 @@ def _stops_running(pid: int, timeout: float = 10) -> bool: return True +@contextlib.contextmanager +def _masked_secret(secret: str): + """Register *secret* with the masker of the parse log, and restore the masker's state afterwards.""" + masker = sdk_secrets_masker() + patterns, replacer = set(masker.patterns), masker.replacer + masker.add_mask(secret) + try: + yield + finally: + masker.patterns, masker.replacer = patterns, replacer + + @pytest.fixture(autouse=True) def _coordinator(): with fake_coordinator(): @@ -350,6 +364,32 @@ def test_a_failed_parse_is_an_import_error(self, parse, spec, reply, error): assert proc.parsing_result.serialized_dags == [] assert proc.parsing_result.import_errors == {"dag.native": error} + def test_an_exit_without_a_result_shows_the_stderr_output(self, parse): + runtime = ( + "import sys; sys.stdout.write('not on stderr\\n'); " + "sys.stderr.write('panic: Dag \"etl\" is already registered\\n\\ngoroutine 1 [running]:\\n'); " + "sys.exit(2)" + ) + + proc = parse(argv=[sys.executable, "-c", runtime]) + + assert proc.parsing_result.import_errors == { + "dag.native": "The Lang-SDK runtime exited with code 2 without a parse result. " + 'Its stderr output:\npanic: Dag "etl" is already registered\n\ngoroutine 1 [running]:' + } + + @pytest.mark.enable_redact + def test_an_exit_without_a_result_masks_the_secrets_in_stderr(self, parse): + runtime = "import sys; sys.stderr.write('panic: native-secret-value\\n'); sys.exit(2)" + + with _masked_secret("native-secret-value"): + proc = parse(argv=[sys.executable, "-c", runtime]) + + assert proc.parsing_result.import_errors == { + "dag.native": "The Lang-SDK runtime exited with code 2 without a parse result. " + "Its stderr output:\npanic: ***" + } + @pytest.mark.parametrize( ("mapping", "error"), [ @@ -540,6 +580,85 @@ def test_the_runtime_inherits_only_its_standard_streams(self, monkeypatch, tmp_p assert fds["0"] == "/dev/null" +def _render_excerpt(stderr: bytes) -> str: + excerpt = _StderrExcerpt() + for line in stderr.splitlines(keepends=True): + excerpt.add(line) + return excerpt.render() + + +class TestStderrExcerpt: + @pytest.mark.parametrize( + ("stderr", "shown"), + [ + pytest.param( + b'panic: Dag "etl" is already registered\n\ngoroutine 1 [running]:\n', + 'panic: Dag "etl" is already registered\n\ngoroutine 1 [running]:', + id="all-lines", + ), + pytest.param( + b"".join(b"line %d\n" % i for i in range(25)), + "\n".join( + [ + *(f"line {i}" for i in range(5)), + "... 5 lines omitted ...", + *(f"line {i}" for i in range(10, 25)), + ] + ), + id="first-and-last-lines", + ), + pytest.param( + b"".join(b"line %d\n" % i for i in range(21)), + "\n".join( + [ + *(f"line {i}" for i in range(5)), + "... 1 line omitted ...", + *(f"line {i}" for i in range(6, 21)), + ] + ), + id="one-line-omitted", + ), + pytest.param(b"a" * 1200 + b"\n", "a" * 1000 + "\N{HORIZONTAL ELLIPSIS}", id="long-line"), + pytest.param(b"\n\npanic: x \r\n\n", "panic: x", id="blank-lines-and-trailing-spaces"), + pytest.param(b"\tat frame\n", "\tat frame", id="indented-first-line"), + pytest.param( + b"bad \xff \xce\xbb\n", + "bad \N{REPLACEMENT CHARACTER} \u03bb", + id="not-utf-8", + ), + pytest.param(b"nul\x00byte\n", "nul\N{REPLACEMENT CHARACTER}byte", id="nul"), + pytest.param(b"\n\n", "", id="only-blank-lines"), + ], + ) + def test_render(self, stderr, shown): + assert _render_excerpt(stderr) == shown + + @pytest.mark.enable_redact + @pytest.mark.parametrize( + ("secret", "stderr", "shown"), + [ + pytest.param( + "native-secret-value", + b"x" * 990 + b"native-secret-value\n", + "x" * 990 + "***", + id="across-the-cut", + ), + pytest.param("k" * 1500, b"key " + b"k" * 1500 + b"\n", "key ***", id="longer-than-the-cut"), + ], + ) + def test_render_hides_a_secret_before_it_cuts_the_line(self, secret, stderr, shown): + with _masked_secret(secret): + assert _render_excerpt(stderr) == shown + + @pytest.mark.enable_redact + def test_render_hides_a_secret_masked_after_its_line(self): + excerpt = _StderrExcerpt() + excerpt.add(b"panic: late-secret-value\n") + + with _masked_secret("late-secret-value"): + assert excerpt.render() == "panic: ***" + + class TestRun: @staticmethod def _run(tmp_path, **spec) -> DagFileParsingResult: diff --git a/go-sdk/README.md b/go-sdk/README.md index ad344112abebf..759d683923a23 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -389,6 +389,11 @@ 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. +- When the first frame on the comm socket is a `DagFileParseRequest` from the Dag processor, the + binary runs no task. It answers with one `DagFileParsingResult` that holds the Dags from + `airflow.Dag` that the binary registered. The Dags are serialized as + [Serializing a native Dag](#serializing-a-native-dag) describes. The binary exits once the Dag + processor acknowledges the result. 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 9abb02717a085..54cbb1c04d1c5 100644 --- a/go-sdk/airflow/bundle.go +++ b/go-sdk/airflow/bundle.go @@ -193,10 +193,11 @@ func (m *taskHandlerMap) ListTaskHandlers() []bundle.TaskHandlerInfo { return slices.Clone(m.order) } -// dagMap holds the registered Dags by dag_id. +// dagMap holds the registered Dags by dag_id and keeps the order in which they were registered. type dagMap struct { - mu sync.Mutex - dags map[string]*DagRef + mu sync.Mutex + dags map[string]*DagRef + order []*DagRef } func (m *dagMap) add(dag *DagRef) { @@ -211,6 +212,47 @@ func (m *dagMap) add(dag *DagRef) { m.dags = make(map[string]*DagRef) } m.dags[dag.dagID] = dag + m.order = append(m.order, dag) +} + +// serialize serializes the registered Dags in the order they were registered. If serializing a Dag +// panics, the panic becomes the error of that Dag instead of ending the parse. The import error of +// the parse then names the Dag. +func (m *dagMap) serialize(fileloc, relativeFileloc string) []bundle.SerializedDag { + m.mu.Lock() + dags := slices.Clone(m.order) + m.mu.Unlock() + + serialized := make([]bundle.SerializedDag, len(dags)) + for i, dag := range dags { + serialized[i] = serializeRecovering(dag, fileloc, relativeFileloc) + } + return serialized +} + +func serializeRecovering(dag *DagRef, fileloc, relativeFileloc string) (s bundle.SerializedDag) { + s.DagID = dag.dagID + defer func() { + if r := recover(); r != nil { + s.Data, s.Err = nil, fmt.Errorf("%v", r) + } + }() + s.Data = dag.serialize(fileloc, relativeFileloc) + return s +} + +// coordinatorBundle is what Serve passes to execution.Serve. It holds the task handlers that a task +// run looks up and the Dags that a Dag parse serializes. +type coordinatorBundle struct{ b *BundleRef } + +var _ bundle.Registry = coordinatorBundle{} + +func (c coordinatorBundle) LookupTask(dagID, taskID string) (bundle.Task, bool) { + return c.b.taskHandlers.LookupTask(dagID, taskID) +} + +func (c coordinatorBundle) SerializeDags(fileloc, relativeFileloc string) []bundle.SerializedDag { + return c.b.dags.serialize(fileloc, relativeFileloc) } func (m *dagMap) has(dagID string) bool { diff --git a/go-sdk/airflow/bundle_test.go b/go-sdk/airflow/bundle_test.go index 7a5874e29b474..45717e577b36b 100644 --- a/go-sdk/airflow/bundle_test.go +++ b/go-sdk/airflow/bundle_test.go @@ -412,3 +412,45 @@ func TestRegisterableRejectsForeignTypes(t *testing.T) { assert.Contains(t, string(out), "foreignItem does not implement airflow.Registerable") assert.Contains(t, string(out), "unexported method registerable") } + +func TestSerializeDagsKeepsTheOtherDagsWhenADagCannotBeSerialized(t *testing.T) { + b := Bundle() + etl := Dag("etl") + etl.Task(noop) + b.Register(etl) + // Register checks and expands every Dag it takes. A Dag that skipped Register therefore stands + // in for a Dag that the serializer fails on. + broken := Dag("broken") + b.dags.dags["broken"] = broken + b.dags.order = append(b.dags.order, broken) + reports := Dag("reports") + reports.Task(noop) + b.Register(reports) + + serialized := coordinatorBundle{b}.SerializeDags("/bundles/go/etl", "etl") + + require.Len(t, serialized, 3) + assert.Equal(t, bundle.SerializedDag{ + DagID: "etl", + Data: etl.serialize("/bundles/go/etl", "etl"), + }, serialized[0]) + assert.Equal(t, "broken", serialized[1].DagID) + assert.Nil(t, serialized[1].Data) + assert.ErrorContains(t, serialized[1].Err, `Dag "broken" is not registered`) + assert.Equal(t, bundle.SerializedDag{ + DagID: "reports", + Data: reports.serialize("/bundles/go/etl", "etl"), + }, serialized[2]) +} + +func TestSerializeDagsLeavesOutADagThatRegisterRejected(t *testing.T) { + b := Bundle() + cyclic := Dag("cyclic") + extracted := orderedTask(t, cyclic, "extract") + loaded := orderedTask(t, cyclic, "load") + extracted.Before(loaded) + loaded.Before(extracted) + require.Panics(t, func() { b.Register(cyclic) }) + + assert.Empty(t, coordinatorBundle{b}.SerializeDags("/bundles/go/etl", "etl")) +} diff --git a/go-sdk/airflow/dag.go b/go-sdk/airflow/dag.go index e95505de0ca4d..11b83c0dc983e 100644 --- a/go-sdk/airflow/dag.go +++ b/go-sdk/airflow/dag.go @@ -71,8 +71,9 @@ type DagRef struct { // [IfRef.Then], [IfRef.Else], [SwitchRef.Case] and the methods of [TaskGroupRef] panic once the Dag // is registered. // -// [BundleRef.Serve] does not yet serve the Dags that Dag returns. It leaves them out of the -// --airflow-metadata manifest and cannot run their tasks. +// When the Dag processor asks to parse the bundle, [BundleRef.Serve] answers with the registered +// Dags in serialized form. Serve does not yet list these Dags in the --airflow-metadata manifest or +// run their tasks. // // Dag panics if it gets more than one DagSpec, or if the DagSpec has a value that Python rejects // when it builds or validates a Dag: diff --git a/go-sdk/airflow/serve.go b/go-sdk/airflow/serve.go index bcfb33ab934d1..c0249f0f2a466 100644 --- a/go-sdk/airflow/serve.go +++ b/go-sdk/airflow/serve.go @@ -57,8 +57,10 @@ 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. +// Airflow starts the executable with --comm and --logs. With these flags, Serve speaks the +// coordinator protocol. It either runs one task or answers the Dag processor's request to parse +// the bundle. To answer a parse request, Serve sends the Dags from [Dag] that the bundle registered +// and runs no task. // // 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: @@ -141,7 +143,7 @@ func (b *BundleRef) serve(args []string, stdout io.Writer) error { } return execution.DumpAirflowMetadata(stdout, &b.taskHandlers, format) case modeCoordinator: - return execution.Serve(&b.taskHandlers, *commAddr, *logsAddr) + return execution.Serve(coordinatorBundle{b}, *commAddr, *logsAddr) case modeCoordinatorUsageError: return errCoordinatorFlagsRequired } diff --git a/go-sdk/airflow/serve_test.go b/go-sdk/airflow/serve_test.go index 93de70af0231c..608cc70e9e953 100644 --- a/go-sdk/airflow/serve_test.go +++ b/go-sdk/airflow/serve_test.go @@ -33,6 +33,7 @@ import ( "gopkg.in/yaml.v3" "github.com/apache/airflow/go-sdk/pkg/execution" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" ) func TestDecideMode(t *testing.T) { @@ -219,28 +220,23 @@ func TestServeHelpIsNotAnError(t *testing.T) { assert.Contains(t, stdout.String(), "--airflow-metadata") } -// A fake supervisor sends StartupDetails over the comm socket, as the Python -// ExecutableCoordinator does after it starts the bundle with --comm and --logs. -func TestServeRunsTaskForSupervisor(t *testing.T) { +// serveForSupervisor runs b.serve with --comm and --logs, the same flags that Airflow starts the +// bundle with. It returns the supervisor's end of the comm socket and the channel that gets what +// serve returns. +func serveForSupervisor(t *testing.T, b *BundleRef) (*execution.CoordinatorComm, <-chan error) { + t.Helper() commLn, err := net.Listen("tcp", "127.0.0.1:0") require.NoError(t, err) - defer commLn.Close() + t.Cleanup(func() { commLn.Close() }) logsLn, err := net.Listen("tcp", "127.0.0.1:0") require.NoError(t, err) - defer logsLn.Close() + t.Cleanup(func() { logsLn.Close() }) // Without a deadline, a Serve that never dials would leave Accept blocked until the test // binary times out. deadline := time.Now().Add(10 * time.Second) require.NoError(t, commLn.(*net.TCPListener).SetDeadline(deadline)) require.NoError(t, logsLn.(*net.TCPListener).SetDeadline(deadline)) - ran := false - b := Bundle() - b.Register(TaskHandler("py_etl", "transform", func(Context) error { - ran = true - return nil - })) - done := make(chan error, 1) go func() { done <- b.serve( @@ -251,13 +247,26 @@ func TestServeRunsTaskForSupervisor(t *testing.T) { commConn, err := commLn.Accept() require.NoError(t, err) - defer commConn.Close() + t.Cleanup(func() { commConn.Close() }) logsConn, err := logsLn.Accept() require.NoError(t, err) - defer logsConn.Close() + t.Cleanup(func() { logsConn.Close() }) require.NoError(t, commConn.SetDeadline(deadline)) - supervisor := execution.NewCoordinatorComm(commConn, commConn, discardLogger()) + return execution.NewCoordinatorComm(commConn, commConn, discardLogger()), done +} + +// A fake supervisor sends StartupDetails over the comm socket, as the Python +// ExecutableCoordinator does after it starts the bundle with --comm and --logs. +func TestServeRunsTaskForSupervisor(t *testing.T) { + ran := false + b := Bundle() + b.Register(TaskHandler("py_etl", "transform", func(Context) error { + ran = true + return nil + })) + + supervisor, done := serveForSupervisor(t, b) require.NoError(t, supervisor.SendRequest(0, map[string]any{ "type": "StartupDetails", "ti": map[string]any{ @@ -284,3 +293,67 @@ func TestServeRunsTaskForSupervisor(t *testing.T) { } assert.True(t, ran) } + +// To parse the bundle, the Dag processor starts it with --comm and --logs and then sends a +// DagFileParseRequest over the comm socket. +func TestServeAnswersTheDagParseRequestWithoutRunningATask(t *testing.T) { + var ran []string + record := func(name string) func(Context) error { + return func(Context) error { + ran = append(ran, name) + return nil + } + } + etl := Dag("etl") + extracted := etl.Task(func(Context) (int, error) { + ran = append(ran, "extract") + return 3, nil + }, TaskSpec{TaskID: "extract"}) + loaded := etl.Task(record("load"), TaskSpec{TaskID: "load"}) + etl.If(func(_ Context, rows int) (bool, error) { + ran = append(ran, "has_rows") + return rows > 0, nil + }, Inputs(extracted), TaskSpec{TaskID: "has_rows"}).Then(loaded) + reports := Dag("reports") + reports.Task(record("publish"), TaskSpec{TaskID: "publish"}) + + b := Bundle() + b.Register(etl, TaskHandler("py_etl", "transform", record("transform")), reports) + + supervisor, done := serveForSupervisor(t, b) + require.NoError(t, supervisor.SendRequest(0, map[string]any{ + "type": "DagFileParseRequest", + "file": "/bundles/go/etl", + "bundle_path": "/bundles/go", + "bundle_name": "go", + })) + + frame, err := supervisor.ReadMessage() + require.NoError(t, err) + var result genmodels.DagFileParsingResult + require.NoError(t, msgpack.Unmarshal(frame.Body, &result)) + assert.Equal(t, "DagFileParsingResult", result.Type) + assert.Equal(t, "/bundles/go/etl", result.Fileloc) + assert.Nil(t, result.ImportErrors) + require.Len(t, result.SerializedDags, 2) + for i, dag := range []*DagRef{etl, reports} { + want := jsonOf(t, dag.serialize("/bundles/go/etl", "etl")) + assertJSON(t, want, result.SerializedDags[i].Data) + } + require.NoError(t, supervisor.SendRequest(frame.ID, nil)) + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return after it answered the Dag parse request") + } + assert.Empty(t, ran) +} + +func jsonOf(t *testing.T, value any) string { + t.Helper() + raw, err := json.Marshal(value) + require.NoError(t, err) + return string(raw) +} diff --git a/go-sdk/internal/bundle/doc.go b/go-sdk/internal/bundle/doc.go index dda8b6ed6d4c4..af14047d0f0a8 100644 --- a/go-sdk/internal/bundle/doc.go +++ b/go-sdk/internal/bundle/doc.go @@ -15,8 +15,11 @@ // specific language governing permissions and limitations // under the License. -// Package bundle defines what the coordinator runtime needs from a bundle: -// the tasks it looks up and runs, and the Dag and task ids it lists in the manifest. +// Package bundle defines what the coordinator runtime needs from a bundle. The +// runtime needs the tasks to look up and run, and the Dag and task ids to list +// in the manifest. It also needs the serialized Dags to send in answer to a Dag +// parse request. // -// Package airflow builds both from the task handlers a bundle registers. +// Package airflow builds the tasks, the ids, and the serialized Dags from the +// task handlers and the Dags that a bundle registers. package bundle diff --git a/go-sdk/internal/bundle/task.go b/go-sdk/internal/bundle/task.go index c0cccb95697be..2b91762250270 100644 --- a/go-sdk/internal/bundle/task.go +++ b/go-sdk/internal/bundle/task.go @@ -57,6 +57,30 @@ type EnumerableBundle interface { ListTaskHandlers() []TaskHandlerInfo } +// SerializedDag is a Dag authored with airflow.Dag, serialized for the Dag processor. +type SerializedDag struct { + DagID string + // Data is the serialized Dag that Airflow stores. It is nil when Err is set. + Data map[string]any + // Err says why the Dag could not be serialized. + Err error +} + +// DagSerializer serializes the Dags from airflow.Dag that a bundle registered, in the order they +// were registered. The coordinator runtime answers a DagFileParseRequest with the serialized Dags. +// fileloc is the path of the bundle binary, and relativeFileloc is that path relative to the root +// of its Dag bundle. +type DagSerializer interface { + SerializeDags(fileloc, relativeFileloc string) []SerializedDag +} + +// Registry is what a bundle binary serves. A task run looks up its task in the Registry, and a Dag +// parse serializes the Dags that the Registry holds. +type Registry interface { + Bundle + DagSerializer +} + type taskFunction struct { fn reflect.Value fullName string diff --git a/go-sdk/pkg/execution/dag_parse.go b/go-sdk/pkg/execution/dag_parse.go new file mode 100644 index 0000000000000..9d876db0bc0ba --- /dev/null +++ b/go-sdk/pkg/execution/dag_parse.go @@ -0,0 +1,78 @@ +// 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 ( + "fmt" + "log/slog" + "path/filepath" + "strings" + + "github.com/apache/airflow/go-sdk/internal/bundle" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +// parseDags answers a DagFileParseRequest with the Dags from airflow.Dag that the bundle +// registered. The answer does not include the Dag of a task handler, because the Python Dag file +// that declares that Dag is parsed separately. +// +// If a Dag cannot be serialized, parseDags turns the failure into an import error that names the +// Dag. parseDags still sends the other Dags, and the Dag processor stores them. But while the +// bundle binary has an import error, the Dag processor deactivates every Dag of the bundle binary. +// The Dag processor treats a Python Dag file with an import error in the same way. The Dag +// processor stores each import error under the path of the file relative to its Dag bundle, and it +// keeps only one import error for each file. So the messages for all the Dags that cannot be +// serialized share one import error. +func parseDags( + b bundle.DagSerializer, + req *genmodels.DagFileParseRequest, + logger *slog.Logger, +) genmodels.DagFileParsingResult { + relative := computeRelativeFileloc(req.File, req.BundlePath) + // A nil slice would be sent as null, which the Dag processor rejects. + result := genmodels.DagFileParsingResult{ + Fileloc: req.File, + SerializedDags: []genmodels.LazyDeserializedDAG{}, + } + var failures []string + for _, dag := range b.SerializeDags(req.File, relative) { + if dag.Err != nil { + logger.Error("Dag could not be serialized", "dag_id", dag.DagID, "error", dag.Err) + failures = append(failures, fmt.Sprintf("Dag %q: %v", dag.DagID, dag.Err)) + continue + } + result.SerializedDags = append( + result.SerializedDags, genmodels.LazyDeserializedDAG{Data: dag.Data}, + ) + } + if len(failures) > 0 { + result.ImportErrors = &genmodels.ImportErrors{relative: strings.Join(failures, "\n")} + } + return result +} + +// computeRelativeFileloc returns file relative to bundlePath, in the same way that Python's DagBag +// computes the relative_fileloc of a Dag file. Like DagBag, it returns file unchanged when file is +// not inside bundlePath. +func computeRelativeFileloc(file, bundlePath string) string { + rel, err := filepath.Rel(bundlePath, file) + if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) { + return file + } + return rel +} diff --git a/go-sdk/pkg/execution/dag_parse_test.go b/go-sdk/pkg/execution/dag_parse_test.go new file mode 100644 index 0000000000000..86f4595d8ca38 --- /dev/null +++ b/go-sdk/pkg/execution/dag_parse_test.go @@ -0,0 +1,194 @@ +// 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 ( + "bytes" + "encoding/json" + "errors" + "io" + "log/slog" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/vmihailenco/msgpack/v5" + + "github.com/apache/airflow/go-sdk/internal/bundle" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" +) + +// serializedDags serializes the Dags it holds. It also records the fileloc and relativeFileloc +// that SerializeDags was called with. +type serializedDags struct { + dags []bundle.SerializedDag + fileloc, relative string +} + +func (s *serializedDags) LookupTask(string, string) (bundle.Task, bool) { return nil, false } + +func (s *serializedDags) SerializeDags(fileloc, relativeFileloc string) []bundle.SerializedDag { + s.fileloc, s.relative = fileloc, relativeFileloc + return s.dags +} + +func serializedDag(dagID string) bundle.SerializedDag { + return bundle.SerializedDag{ + DagID: dagID, + Data: map[string]any{"__version": 3, "dag": map[string]any{"dag_id": dagID}}, + } +} + +var etlParseRequest = &genmodels.DagFileParseRequest{ + File: "/bundles/go/etl", + BundlePath: "/bundles/go", + BundleName: "go", +} + +func discardLogger() *slog.Logger { return slog.New(slog.NewTextHandler(io.Discard, nil)) } + +// wireJSON passes body through the frame encoding that SendRequest uses and returns the result as +// JSON. A test then sees the keys, nulls and empty lists that the Dag processor receives. +func wireJSON(t *testing.T, body any) string { + t.Helper() + payload, err := encodeRequest(0, body) + require.NoError(t, err) + frame, err := decodeFrame(payload) + require.NoError(t, err) + dec := msgpack.NewDecoder(bytes.NewReader(frame.Body)) + dec.UseLooseInterfaceDecoding(true) + var decoded any + require.NoError(t, dec.Decode(&decoded)) + raw, err := json.Marshal(decoded) + require.NoError(t, err) + return string(raw) +} + +func TestParseDagsAnswersWithTheSerializedDags(t *testing.T) { + dags := &serializedDags{dags: []bundle.SerializedDag{ + serializedDag("etl"), + serializedDag("reports"), + }} + + result := parseDags(dags, etlParseRequest, discardLogger()) + + assert.Equal(t, "/bundles/go/etl", dags.fileloc) + assert.Equal(t, "etl", dags.relative) + assert.JSONEq(t, `{ + "type": "DagFileParsingResult", + "fileloc": "/bundles/go/etl", + "serialized_dags": [ + {"data": {"__version": 3, "dag": {"dag_id": "etl"}}}, + {"data": {"__version": 3, "dag": {"dag_id": "reports"}}} + ] + }`, wireJSON(t, result)) +} + +func TestParseDagsSendsAnEmptyListForABundleWithoutDags(t *testing.T) { + result := parseDags(&serializedDags{}, etlParseRequest, discardLogger()) + + assert.JSONEq(t, `{ + "type": "DagFileParsingResult", + "fileloc": "/bundles/go/etl", + "serialized_dags": [] + }`, wireJSON(t, result)) +} + +func TestParseDagsReportsEachDagThatCannotBeSerialized(t *testing.T) { + dags := &serializedDags{dags: []bundle.SerializedDag{ + serializedDag("etl"), + {DagID: "reports", Err: errors.New("no schema field")}, + serializedDag("cleanup"), + {DagID: "audit", Err: errors.New("bad field type")}, + }} + var logs bytes.Buffer + logger := slog.New(slog.NewJSONHandler(&logs, nil)) + + result := parseDags(dags, etlParseRequest, logger) + + assert.JSONEq(t, `{ + "type": "DagFileParsingResult", + "fileloc": "/bundles/go/etl", + "serialized_dags": [ + {"data": {"__version": 3, "dag": {"dag_id": "etl"}}}, + {"data": {"__version": 3, "dag": {"dag_id": "cleanup"}}} + ], + "import_errors": { + "etl": "Dag \"reports\": no schema field\nDag \"audit\": bad field type" + } + }`, wireJSON(t, result)) + + var logged []map[string]any + for line := range bytes.Lines(logs.Bytes()) { + var record map[string]any + require.NoError(t, json.Unmarshal(line, &record)) + logged = append(logged, record) + } + require.Len(t, logged, 2) + for i, dagID := range []string{"reports", "audit"} { + assert.Equal(t, "ERROR", logged[i]["level"]) + assert.Equal(t, dagID, logged[i]["dag_id"]) + } +} + +func TestComputeRelativeFileloc(t *testing.T) { + tests := []struct { + name string + file string + bundlePath string + want string + }{ + {name: "in the bundle", file: "/bundles/go/etl", bundlePath: "/bundles/go", want: "etl"}, + { + name: "in a directory of the bundle", + file: "/bundles/go/dags/etl", + bundlePath: "/bundles/go", + want: "dags/etl", + }, + { + name: "the bundle itself", + file: "/bundles/go/etl", + bundlePath: "/bundles/go/etl", + want: ".", + }, + { + name: "a name that starts with two dots", + file: "/bundles/go/..etl", + bundlePath: "/bundles/go", + want: "..etl", + }, + { + name: "outside the bundle", + file: "/bundles/gopher/etl", + bundlePath: "/bundles/go", + want: "/bundles/gopher/etl", + }, + { + name: "the directory that holds the bundle", + file: "/bundles", + bundlePath: "/bundles/go", + want: "/bundles", + }, + {name: "a relative file", file: "etl", bundlePath: "/bundles/go", want: "etl"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, computeRelativeFileloc(tt.file, tt.bundlePath)) + }) + } +} diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 664563511d472..affd41c10c3c6 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -103,7 +103,11 @@ func (b testBundle) LookupTask(dagID, taskID string) (bundle.Task, bool) { return task, ok } -func buildBundle(t *testing.T, register func(testBundle)) bundle.Bundle { +// SerializeDags serializes nothing because the Dags of a testBundle only hold tasks to run. The Dag +// parse tests use serializedDags instead. +func (b testBundle) SerializeDags(string, string) []bundle.SerializedDag { return nil } + +func buildBundle(t *testing.T, register func(testBundle)) bundle.Registry { t.Helper() b := testBundle{} register(b) @@ -1361,7 +1365,8 @@ func TestServeFailureAfterConnectClosesComm(t *testing.T) { logsConn := <-logsCh defer logsConn.Close() - // Serve expects StartupDetails as the first frame, so it fails to decode a VariableResult. + // Serve expects StartupDetails or DagFileParseRequest as the first frame, so it fails to + // decode a VariableResult. payload, err := encodeRequest( 0, map[string]any{"type": "VariableResult", "key": "k", "value": "v"}, @@ -1382,3 +1387,96 @@ func TestServeFailureAfterConnectClosesComm(t *testing.T) { _, err = readFrame(commConn) require.Error(t, err) } + +// startDagParse runs Serve for dags and sends it a DagFileParseRequest. It returns the comm and +// logs connections, the frame that Serve answers with, and the channel that gets what Serve +// returns. +func startDagParse( + t *testing.T, + dags bundle.Registry, +) (commConn, logsConn net.Conn, frame IncomingFrame, done <-chan error) { + t.Helper() + commAddr, logsAddr, commCh, logsCh, cleanup := startSupervisor(t) + t.Cleanup(cleanup) + + served := make(chan error, 1) + go func() { served <- Serve(dags, commAddr, logsAddr) }() + + commConn = <-commCh + t.Cleanup(func() { commConn.Close() }) + logsConn = <-logsCh + t.Cleanup(func() { logsConn.Close() }) + deadline := time.Now().Add(10 * time.Second) + require.NoError(t, commConn.SetDeadline(deadline)) + require.NoError(t, logsConn.SetDeadline(deadline)) + + payload, err := encodeRequest(0, map[string]any{ + "type": "DagFileParseRequest", + "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, logsConn, frame, served +} + +func TestServeDagFileParseRequestEndToEnd(t *testing.T) { + dags := &serializedDags{dags: []bundle.SerializedDag{serializedDag("etl")}} + commConn, _, frame, done := startDagParse(t, dags) + + var result genmodels.DagFileParsingResult + require.NoError(t, decodeBody(frame.Body, &result)) + assert.Equal(t, "DagFileParsingResult", result.Type) + assert.Equal(t, "/bundles/go/etl", result.Fileloc) + require.Len(t, result.SerializedDags, 1) + assert.Equal(t, "etl", result.SerializedDags[0].Data["dag"].(map[string]any)["dag_id"]) + assert.Equal(t, "etl", dags.relative) + 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 Dag processor acknowledged the result") + } +} + +func TestServeFailsWhenTheDagProcessorRejectsTheParseResult(t *testing.T) { + commConn, logsConn, frame, done := startDagParse(t, &serializedDags{}) + + rejection := map[string]any{ + "type": "ErrorResponse", + "error": "generic_error", + "detail": map[string]any{"message": "A parse result was already received"}, + } + require.NoError(t, writeFrame(commConn, encodeResponseFrame(t, frame.ID, nil, rejection))) + + select { + case err := <-done: + require.ErrorContains(t, err, "sending Dag parsing result") + require.ErrorContains(t, err, "A parse result was already received") + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return after the Dag processor rejected the result") + } + + // main's log.Fatal writes to the logs socket after Serve has closed that socket. The error + // therefore shows up only if Serve logs it first. + logs, err := io.ReadAll(logsConn) + require.NoError(t, err) + var failure map[string]any + for line := range bytes.Lines(logs) { + var record map[string]any + require.NoError(t, json.Unmarshal(line, &record)) + if record["event"] == "Failed to send the Dag parsing result" { + failure = record + } + } + require.NotNil(t, failure, "Serve did not log why it failed: %s", logs) + assert.Equal(t, "error", failure["level"]) + assert.Contains(t, failure["error"], "A parse result was already received") +} diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index 0f56e005315ca..c570cfb7d7e52 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -20,8 +20,11 @@ // 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 frame that the runtime reads on the comm socket decides what the +// runtime does. After a StartupDetails message, the runtime runs a task over +// several rounds of messages. After a DagFileParseRequest from the Dag +// processor, the runtime runs no task. It answers with one DagFileParsingResult +// that holds the Dags from airflow.Dag that the bundle registered. // // See go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md. package execution @@ -52,17 +55,21 @@ const dialTimeout = 30 * time.Second // 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. +// terminalSendTimeout also bounds the write of a DagFileParsingResult. 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 then exit 0. A +// clean shutdown means one of two things. Either the task ran and its terminal +// TaskState/SucceedTask frame was delivered, or the DagFileParsingResult was +// delivered and the Dag processor acknowledged it. A non-nil error indicates a +// protocol-level failure (connection loss, malformed frames, unknown first +// message type). In a task run, the failure happens before or instead of the +// delivery of a terminal frame. In a Dag parse, the failure happens before the +// Dag processor acknowledges the DagFileParsingResult. // // 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 +80,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 +174,29 @@ func Serve(b bundle.Bundle, commAddr, logsAddr string) error { } logger.Debug("Task execution complete") + case *genmodels.DagFileParseRequest: + logger.Debug("Dag parsing mode", "file", msg.File) + result := parseDags(b, msg, logger) + // The Dag processor reads from a connection only after it checks that the connection + // belongs to this process. The Dag processor cannot make that check after this process + // exits. So a process that exits right after it sends the result often loses the logs of + // its parse. Waiting for the acknowledgement keeps this process running until the Dag + // processor has read the result. The wait also gives the Dag processor time to check the + // logs connection. + // + // The write has a deadline. The wait for the acknowledgement has no deadline. The Dag + // processor validates the Dags in the result before it sends the acknowledgement. A larger + // Dag takes longer to validate. So a deadline on the wait could expire while the + // validation is still running. Even without a deadline, the wait cannot hang. The Dag + // processor kills a runtime that keeps running too long after the Dag processor receives + // the parse result from that runtime. + _ = commConn.SetWriteDeadline(time.Now().Add(terminalSendTimeout)) + if _, err := comm.Communicate(ctx, result); err != nil { + logger.Error("Failed to send the Dag parsing result", "error", err) + return fmt.Errorf("sending Dag parsing result: %w", err) + } + logger.Debug("Dag parsing complete", "dags", len(result.SerializedDags)) + default: logger.Error("Unexpected initial message type", "type", fmt.Sprintf("%T", body)) return fmt.Errorf("unexpected initial message type: %T", body) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 372d12bc62102..332c316f8e6b6 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1071,11 +1071,17 @@ def _create_log_forwarder( for log in loggers ) return make_buffered_socket_reader( - forward_to_log(loggers, logger=name, level=log_level), + self._create_line_forwarder(loggers, name, log_level), data=data, on_close=self._on_socket_closed, ) + def _create_line_forwarder( + self, loggers: tuple[FilteringBoundLogger, ...], name: str, level: int + ) -> Generator[None, bytes | bytearray, None]: + """Create the generator that logs each line of the child's output stream named *name*.""" + return forward_to_log(loggers, logger=name, level=level) + def _on_socket_closed(self, sock: socket): # We want to keep servicing this process until we've read up to EOF from all the sockets. with suppress(KeyError): From 79426c636119ab685f722527bac70c815b8c468e Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Wed, 7 Oct 2026 07:49:59 +0000 Subject: [PATCH 2/6] Drop the stderr excerpt from the Lang-SDK Dag processor The stderr excerpt is unrelated to Go Dag parsing, so this PR leaves it out. --- .../dag_processing/lang_sdk_processor.py | 67 +--------- .../dag_processing/test_lang_sdk_processor.py | 119 ------------------ .../airflow/sdk/execution_time/supervisor.py | 8 +- 3 files changed, 4 insertions(+), 190 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py index 57d0aa662c1e8..9a0998dd05109 100644 --- a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -23,7 +23,6 @@ import selectors import signal import time -from collections import deque from contextlib import suppress from pathlib import Path from socket import socket @@ -47,7 +46,6 @@ ) from airflow.exceptions import DeserializationError from airflow.models.pool import Pool -from airflow.sdk._shared.secrets_masker import redact # noqa: SDK001 from airflow.sdk.coordinators._subprocess import _is_connection_from_pid, _start_server from airflow.sdk.execution_time import supervisor from airflow.sdk.execution_time.comms import CommsDecoder, ErrorResponse, _RequestFrame @@ -77,15 +75,6 @@ _IMPORT_TIMEOUT_SETTING = "[core] dagbag_import_timeout or the get_dagbag_import_timeout policy" _PROCESSOR_TIMEOUT_SETTING = "[dag_processor] dag_file_processor_timeout" -# When a runtime exits without a parse result, its import error shows the first and the last lines -# that the runtime wrote to stderr. A Go panic and a Java exception print their message on the first -# line, and a Python traceback prints its message on the last line. Each line is cut to a maximum -# length. Node writes the whole source line that threw an error to stderr, and in a minified -# JavaScript bundle one source line can hold the whole bundle. -_STDERR_HEAD_LINES = 5 -_STDERR_TAIL_LINES = 15 -_STDERR_LINE_LENGTH = 1000 - # StartLangSDKRuntime and LangSDKRuntimeSchemaVersion pass only between the manager and its forked # child before the exec, so they are not part of the supervisor schema the runtimes speak. @@ -166,48 +155,6 @@ def report_schema_version(schema_version: str | None) -> None: _Channel = Literal["comm", "logs"] -@attrs.define -class _StderrExcerpt: - """Keeps the first and the last lines of the runtime's stderr for the import error.""" - - head: list[str] = attrs.field(factory=list) - tail: deque[str] = attrs.field(factory=functools.partial(deque, maxlen=_STDERR_TAIL_LINES)) - omitted: int = 0 - - def add(self, line: bytes | bytearray) -> None: - # The masker finds a secret only when the whole secret is in the text, so mask the line before - # cutting it. - text = cast("str", redact(line.rstrip().decode("utf-8", errors="replace"))) - # Postgres cannot store a NUL character in a text column. - text = text.replace("\x00", "\N{REPLACEMENT CHARACTER}") - if len(text) > _STDERR_LINE_LENGTH: - text = text[:_STDERR_LINE_LENGTH] + "\N{HORIZONTAL ELLIPSIS}" - if len(self.head) < _STDERR_HEAD_LINES: - self.head.append(text) - return - if len(self.tail) == _STDERR_TAIL_LINES: - self.omitted += 1 - self.tail.append(text) - - def render(self) -> str: - lines = "line" if self.omitted == 1 else "lines" - omitted = [f"... {self.omitted} {lines} omitted ..."] if self.omitted else [] - # The runtime can ask the Dag processor to mask a secret after printing the secret. A secret - # can also span several lines. - return cast("str", redact("\n".join([*self.head, *omitted, *self.tail]).strip("\n"))) - - -def _record_lines( - lines: Generator[None, bytes | bytearray, None], excerpt: _StderrExcerpt -) -> Generator[None, bytes | bytearray, None]: - """Pass each line on to *lines* and add the line to *excerpt*.""" - next(lines) - while True: - line = yield - excerpt.add(line) - lines.send(line) - - @attrs.define(kw_only=True) class LangSDKDagFileProcessorProcess(BaseDagFileProcessorProcess): """ @@ -230,7 +177,6 @@ class LangSDKDagFileProcessorProcess(BaseDagFileProcessorProcess): _parsing_result_monotonic: float | None = attrs.field(default=None, init=False) _unverified_connections: list[tuple[socket, _Channel]] = attrs.field(factory=list, init=False) _group_killed: bool = attrs.field(default=False, init=False) - _stderr_excerpt: _StderrExcerpt = attrs.field(factory=_StderrExcerpt, init=False) @classmethod def start( # type: ignore[override] @@ -448,12 +394,6 @@ def _register_logs(self, conn: socket) -> None: ), ) - def _create_line_forwarder( - self, loggers: tuple[FilteringBoundLogger, ...], name: str, level: int - ) -> Generator[None, bytes | bytearray, None]: - lines = super()._create_line_forwarder(loggers, name, level) - return _record_lines(lines, self._stderr_excerpt) if name.endswith(".stderr") else lines - def _set_import_error(self, message: str) -> None: self.parsing_result = DagFileParsingResult( fileloc=self._parse_request.file, @@ -598,10 +538,9 @@ def is_ready(self) -> bool: if not super().is_ready: return False if self.parsing_result is None: - message = f"The Lang-SDK runtime exited with code {self._exit_code} without a parse result" - if stderr := self._stderr_excerpt.render(): - message += f". Its stderr output:\n{stderr}" - self._set_import_error(message) + 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, setting: str) -> None: diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index 1dfbcc394511b..8be3221b37874 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -43,13 +43,11 @@ LangSDKDagFileProcessorProcess, LangSDKRuntimeSchemaVersion, _get_import_timeout, - _StderrExcerpt, ) from airflow.dag_processing.processor import DagFileParseRequest, DagFileParsingResult from airflow.exceptions import UnknownExecutorException from airflow.executors.executor_loader import ExecutorLoader from airflow.sdk import DAG, BaseOperator, task -from airflow.sdk._shared.secrets_masker import _secrets_masker as sdk_secrets_masker from airflow.sdk.api.client import Client from airflow.sdk.api.datamodels._generated import VariableResponse from airflow.sdk.exceptions import AirflowRuntimeError @@ -117,18 +115,6 @@ def _stops_running(pid: int, timeout: float = 10) -> bool: return True -@contextlib.contextmanager -def _masked_secret(secret: str): - """Register *secret* with the masker of the parse log, and restore the masker's state afterwards.""" - masker = sdk_secrets_masker() - patterns, replacer = set(masker.patterns), masker.replacer - masker.add_mask(secret) - try: - yield - finally: - masker.patterns, masker.replacer = patterns, replacer - - @pytest.fixture(autouse=True) def _coordinator(): with fake_coordinator(): @@ -364,32 +350,6 @@ def test_a_failed_parse_is_an_import_error(self, parse, spec, reply, error): assert proc.parsing_result.serialized_dags == [] assert proc.parsing_result.import_errors == {"dag.native": error} - def test_an_exit_without_a_result_shows_the_stderr_output(self, parse): - runtime = ( - "import sys; sys.stdout.write('not on stderr\\n'); " - "sys.stderr.write('panic: Dag \"etl\" is already registered\\n\\ngoroutine 1 [running]:\\n'); " - "sys.exit(2)" - ) - - proc = parse(argv=[sys.executable, "-c", runtime]) - - assert proc.parsing_result.import_errors == { - "dag.native": "The Lang-SDK runtime exited with code 2 without a parse result. " - 'Its stderr output:\npanic: Dag "etl" is already registered\n\ngoroutine 1 [running]:' - } - - @pytest.mark.enable_redact - def test_an_exit_without_a_result_masks_the_secrets_in_stderr(self, parse): - runtime = "import sys; sys.stderr.write('panic: native-secret-value\\n'); sys.exit(2)" - - with _masked_secret("native-secret-value"): - proc = parse(argv=[sys.executable, "-c", runtime]) - - assert proc.parsing_result.import_errors == { - "dag.native": "The Lang-SDK runtime exited with code 2 without a parse result. " - "Its stderr output:\npanic: ***" - } - @pytest.mark.parametrize( ("mapping", "error"), [ @@ -580,85 +540,6 @@ def test_the_runtime_inherits_only_its_standard_streams(self, monkeypatch, tmp_p assert fds["0"] == "/dev/null" -def _render_excerpt(stderr: bytes) -> str: - excerpt = _StderrExcerpt() - for line in stderr.splitlines(keepends=True): - excerpt.add(line) - return excerpt.render() - - -class TestStderrExcerpt: - @pytest.mark.parametrize( - ("stderr", "shown"), - [ - pytest.param( - b'panic: Dag "etl" is already registered\n\ngoroutine 1 [running]:\n', - 'panic: Dag "etl" is already registered\n\ngoroutine 1 [running]:', - id="all-lines", - ), - pytest.param( - b"".join(b"line %d\n" % i for i in range(25)), - "\n".join( - [ - *(f"line {i}" for i in range(5)), - "... 5 lines omitted ...", - *(f"line {i}" for i in range(10, 25)), - ] - ), - id="first-and-last-lines", - ), - pytest.param( - b"".join(b"line %d\n" % i for i in range(21)), - "\n".join( - [ - *(f"line {i}" for i in range(5)), - "... 1 line omitted ...", - *(f"line {i}" for i in range(6, 21)), - ] - ), - id="one-line-omitted", - ), - pytest.param(b"a" * 1200 + b"\n", "a" * 1000 + "\N{HORIZONTAL ELLIPSIS}", id="long-line"), - pytest.param(b"\n\npanic: x \r\n\n", "panic: x", id="blank-lines-and-trailing-spaces"), - pytest.param(b"\tat frame\n", "\tat frame", id="indented-first-line"), - pytest.param( - b"bad \xff \xce\xbb\n", - "bad \N{REPLACEMENT CHARACTER} \u03bb", - id="not-utf-8", - ), - pytest.param(b"nul\x00byte\n", "nul\N{REPLACEMENT CHARACTER}byte", id="nul"), - pytest.param(b"\n\n", "", id="only-blank-lines"), - ], - ) - def test_render(self, stderr, shown): - assert _render_excerpt(stderr) == shown - - @pytest.mark.enable_redact - @pytest.mark.parametrize( - ("secret", "stderr", "shown"), - [ - pytest.param( - "native-secret-value", - b"x" * 990 + b"native-secret-value\n", - "x" * 990 + "***", - id="across-the-cut", - ), - pytest.param("k" * 1500, b"key " + b"k" * 1500 + b"\n", "key ***", id="longer-than-the-cut"), - ], - ) - def test_render_hides_a_secret_before_it_cuts_the_line(self, secret, stderr, shown): - with _masked_secret(secret): - assert _render_excerpt(stderr) == shown - - @pytest.mark.enable_redact - def test_render_hides_a_secret_masked_after_its_line(self): - excerpt = _StderrExcerpt() - excerpt.add(b"panic: late-secret-value\n") - - with _masked_secret("late-secret-value"): - assert excerpt.render() == "panic: ***" - - class TestRun: @staticmethod def _run(tmp_path, **spec) -> DagFileParsingResult: diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 332c316f8e6b6..372d12bc62102 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1071,17 +1071,11 @@ def _create_log_forwarder( for log in loggers ) return make_buffered_socket_reader( - self._create_line_forwarder(loggers, name, log_level), + forward_to_log(loggers, logger=name, level=log_level), data=data, on_close=self._on_socket_closed, ) - def _create_line_forwarder( - self, loggers: tuple[FilteringBoundLogger, ...], name: str, level: int - ) -> Generator[None, bytes | bytearray, None]: - """Create the generator that logs each line of the child's output stream named *name*.""" - return forward_to_log(loggers, logger=name, level=level) - def _on_socket_closed(self, sock: socket): # We want to keep servicing this process until we've read up to EOF from all the sockets. with suppress(KeyError): From 5c814635838e7e0af32a9fe118074ab8b5268b2c Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Wed, 7 Oct 2026 07:52:11 +0000 Subject: [PATCH 3/6] Go SDK: serve the Dag parse request from bundle.Bundle Serve takes a bundle.Bundle again, and the Dag serializer is an optional interface the bundle can implement. --- go-sdk/airflow/bundle.go | 36 ++++++++++---------- go-sdk/airflow/bundle_test.go | 4 +-- go-sdk/airflow/serve.go | 2 +- go-sdk/internal/bundle/task.go | 22 ++++-------- go-sdk/pkg/execution/dag_parse.go | 43 ++++++++++++------------ go-sdk/pkg/execution/dag_parse_test.go | 16 +++++++-- go-sdk/pkg/execution/integration_test.go | 16 +++++---- go-sdk/pkg/execution/server.go | 7 ++-- 8 files changed, 76 insertions(+), 70 deletions(-) diff --git a/go-sdk/airflow/bundle.go b/go-sdk/airflow/bundle.go index 54cbb1c04d1c5..788f59cdd4bdb 100644 --- a/go-sdk/airflow/bundle.go +++ b/go-sdk/airflow/bundle.go @@ -215,6 +215,14 @@ func (m *dagMap) add(dag *DagRef) { m.order = append(m.order, dag) } +func (m *dagMap) has(dagID string) bool { + m.mu.Lock() + defer m.mu.Unlock() + + _, exists := m.dags[dagID] + return exists +} + // serialize serializes the registered Dags in the order they were registered. If serializing a Dag // panics, the panic becomes the error of that Dag instead of ending the parse. The import error of // the parse then names the Dag. @@ -241,24 +249,18 @@ func serializeRecovering(dag *DagRef, fileloc, relativeFileloc string) (s bundle return s } -// coordinatorBundle is what Serve passes to execution.Serve. It holds the task handlers that a task -// run looks up and the Dags that a Dag parse serializes. -type coordinatorBundle struct{ b *BundleRef } - -var _ bundle.Registry = coordinatorBundle{} - -func (c coordinatorBundle) LookupTask(dagID, taskID string) (bundle.Task, bool) { - return c.b.taskHandlers.LookupTask(dagID, taskID) +// coordinatorSource is what Serve hands to execution.Serve: the task handlers and the Dags of one +// bundle. +type coordinatorSource struct { + *taskHandlerMap + dags *dagMap } -func (c coordinatorBundle) SerializeDags(fileloc, relativeFileloc string) []bundle.SerializedDag { - return c.b.dags.serialize(fileloc, relativeFileloc) -} - -func (m *dagMap) has(dagID string) bool { - m.mu.Lock() - defer m.mu.Unlock() +var ( + _ bundle.Bundle = coordinatorSource{} + _ bundle.DagSerializer = coordinatorSource{} +) - _, exists := m.dags[dagID] - return exists +func (s coordinatorSource) SerializeDags(fileloc, relativeFileloc string) []bundle.SerializedDag { + return s.dags.serialize(fileloc, relativeFileloc) } diff --git a/go-sdk/airflow/bundle_test.go b/go-sdk/airflow/bundle_test.go index 45717e577b36b..c49af008e8209 100644 --- a/go-sdk/airflow/bundle_test.go +++ b/go-sdk/airflow/bundle_test.go @@ -427,7 +427,7 @@ func TestSerializeDagsKeepsTheOtherDagsWhenADagCannotBeSerialized(t *testing.T) reports.Task(noop) b.Register(reports) - serialized := coordinatorBundle{b}.SerializeDags("/bundles/go/etl", "etl") + serialized := b.dags.serialize("/bundles/go/etl", "etl") require.Len(t, serialized, 3) assert.Equal(t, bundle.SerializedDag{ @@ -452,5 +452,5 @@ func TestSerializeDagsLeavesOutADagThatRegisterRejected(t *testing.T) { loaded.Before(extracted) require.Panics(t, func() { b.Register(cyclic) }) - assert.Empty(t, coordinatorBundle{b}.SerializeDags("/bundles/go/etl", "etl")) + assert.Empty(t, b.dags.serialize("/bundles/go/etl", "etl")) } diff --git a/go-sdk/airflow/serve.go b/go-sdk/airflow/serve.go index c0249f0f2a466..0985911a04fc8 100644 --- a/go-sdk/airflow/serve.go +++ b/go-sdk/airflow/serve.go @@ -143,7 +143,7 @@ func (b *BundleRef) serve(args []string, stdout io.Writer) error { } return execution.DumpAirflowMetadata(stdout, &b.taskHandlers, format) case modeCoordinator: - return execution.Serve(coordinatorBundle{b}, *commAddr, *logsAddr) + return execution.Serve(coordinatorSource{&b.taskHandlers, &b.dags}, *commAddr, *logsAddr) case modeCoordinatorUsageError: return errCoordinatorFlagsRequired } diff --git a/go-sdk/internal/bundle/task.go b/go-sdk/internal/bundle/task.go index 2b91762250270..c3cb8cc65570a 100644 --- a/go-sdk/internal/bundle/task.go +++ b/go-sdk/internal/bundle/task.go @@ -57,30 +57,20 @@ type EnumerableBundle interface { ListTaskHandlers() []TaskHandlerInfo } -// SerializedDag is a Dag authored with airflow.Dag, serialized for the Dag processor. +// SerializedDag is one Dag from airflow.Dag, serialized for the Dag processor. Data is nil when +// Err is set. type SerializedDag struct { DagID string - // Data is the serialized Dag that Airflow stores. It is nil when Err is set. - Data map[string]any - // Err says why the Dag could not be serialized. - Err error + Data map[string]any + Err error } -// DagSerializer serializes the Dags from airflow.Dag that a bundle registered, in the order they -// were registered. The coordinator runtime answers a DagFileParseRequest with the serialized Dags. -// fileloc is the path of the bundle binary, and relativeFileloc is that path relative to the root -// of its Dag bundle. +// DagSerializer serializes the Dags from airflow.Dag that a bundle registered, in registration +// order. type DagSerializer interface { SerializeDags(fileloc, relativeFileloc string) []SerializedDag } -// Registry is what a bundle binary serves. A task run looks up its task in the Registry, and a Dag -// parse serializes the Dags that the Registry holds. -type Registry interface { - Bundle - DagSerializer -} - type taskFunction struct { fn reflect.Value fullName string diff --git a/go-sdk/pkg/execution/dag_parse.go b/go-sdk/pkg/execution/dag_parse.go index 9d876db0bc0ba..96d33416ea79a 100644 --- a/go-sdk/pkg/execution/dag_parse.go +++ b/go-sdk/pkg/execution/dag_parse.go @@ -27,19 +27,11 @@ import ( "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" ) -// parseDags answers a DagFileParseRequest with the Dags from airflow.Dag that the bundle -// registered. The answer does not include the Dag of a task handler, because the Python Dag file -// that declares that Dag is parsed separately. -// -// If a Dag cannot be serialized, parseDags turns the failure into an import error that names the -// Dag. parseDags still sends the other Dags, and the Dag processor stores them. But while the -// bundle binary has an import error, the Dag processor deactivates every Dag of the bundle binary. -// The Dag processor treats a Python Dag file with an import error in the same way. The Dag -// processor stores each import error under the path of the file relative to its Dag bundle, and it -// keeps only one import error for each file. So the messages for all the Dags that cannot be -// serialized share one import error. +// parseDags answers a DagFileParseRequest with the bundle's Dags from airflow.Dag. A Dag that +// cannot be serialized adds a line to the import error of the file, and the other Dags are still +// sent. func parseDags( - b bundle.DagSerializer, + s bundle.DagSerializer, req *genmodels.DagFileParseRequest, logger *slog.Logger, ) genmodels.DagFileParsingResult { @@ -49,20 +41,29 @@ func parseDags( Fileloc: req.File, SerializedDags: []genmodels.LazyDeserializedDAG{}, } - var failures []string - for _, dag := range b.SerializeDags(req.File, relative) { - if dag.Err != nil { - logger.Error("Dag could not be serialized", "dag_id", dag.DagID, "error", dag.Err) - failures = append(failures, fmt.Sprintf("Dag %q: %v", dag.DagID, dag.Err)) - continue + var dagIDs, failures []string + if s != nil { + for _, dag := range s.SerializeDags(req.File, relative) { + dagIDs = append(dagIDs, dag.DagID) + if dag.Err != nil { + logger.Error("Dag could not be serialized", "dag_id", dag.DagID, "error", dag.Err) + failures = append(failures, fmt.Sprintf("Dag %q: %v", dag.DagID, dag.Err)) + continue + } + result.SerializedDags = append( + result.SerializedDags, genmodels.LazyDeserializedDAG{Data: dag.Data}, + ) } - result.SerializedDags = append( - result.SerializedDags, genmodels.LazyDeserializedDAG{Data: dag.Data}, - ) } if len(failures) > 0 { result.ImportErrors = &genmodels.ImportErrors{relative: strings.Join(failures, "\n")} } + logger.Info("Parse-mode response", + "fileloc", req.File, + "dag_ids", dagIDs, + "serialized", len(result.SerializedDags), + "import_errors", len(failures), + ) return result } diff --git a/go-sdk/pkg/execution/dag_parse_test.go b/go-sdk/pkg/execution/dag_parse_test.go index 86f4595d8ca38..e0d2078cba452 100644 --- a/go-sdk/pkg/execution/dag_parse_test.go +++ b/go-sdk/pkg/execution/dag_parse_test.go @@ -40,8 +40,6 @@ type serializedDags struct { fileloc, relative string } -func (s *serializedDags) LookupTask(string, string) (bundle.Task, bool) { return nil, false } - func (s *serializedDags) SerializeDags(fileloc, relativeFileloc string) []bundle.SerializedDag { s.fileloc, s.relative = fileloc, relativeFileloc return s.dags @@ -99,6 +97,14 @@ func TestParseDagsAnswersWithTheSerializedDags(t *testing.T) { }`, wireJSON(t, result)) } +func TestParseDagsAnswersWithNoDagsWhenTheBundleHasNoSerializer(t *testing.T) { + result := parseDags(nil, etlParseRequest, discardLogger()) + + assert.NotNil(t, result.SerializedDags) + assert.Empty(t, result.SerializedDags) + assert.Nil(t, result.ImportErrors) +} + func TestParseDagsSendsAnEmptyListForABundleWithoutDags(t *testing.T) { result := parseDags(&serializedDags{}, etlParseRequest, discardLogger()) @@ -139,11 +145,15 @@ func TestParseDagsReportsEachDagThatCannotBeSerialized(t *testing.T) { require.NoError(t, json.Unmarshal(line, &record)) logged = append(logged, record) } - require.Len(t, logged, 2) + require.Len(t, logged, 3) for i, dagID := range []string{"reports", "audit"} { assert.Equal(t, "ERROR", logged[i]["level"]) assert.Equal(t, dagID, logged[i]["dag_id"]) } + assert.Equal(t, "INFO", logged[2]["level"]) + assert.Equal(t, []any{"etl", "reports", "cleanup", "audit"}, logged[2]["dag_ids"]) + assert.EqualValues(t, 2, logged[2]["serialized"]) + assert.EqualValues(t, 2, logged[2]["import_errors"]) } func TestComputeRelativeFileloc(t *testing.T) { diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index affd41c10c3c6..4f3411014ce22 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -103,11 +103,7 @@ func (b testBundle) LookupTask(dagID, taskID string) (bundle.Task, bool) { return task, ok } -// SerializeDags serializes nothing because the Dags of a testBundle only hold tasks to run. The Dag -// parse tests use serializedDags instead. -func (b testBundle) SerializeDags(string, string) []bundle.SerializedDag { return nil } - -func buildBundle(t *testing.T, register func(testBundle)) bundle.Registry { +func buildBundle(t *testing.T, register func(testBundle)) bundle.Bundle { t.Helper() b := testBundle{} register(b) @@ -1388,19 +1384,25 @@ func TestServeFailureAfterConnectClosesComm(t *testing.T) { require.Error(t, err) } +// parseBundle is a bundle that serializes Dags and has no task to run. +type parseBundle struct { + testBundle + *serializedDags +} + // startDagParse runs Serve for dags and sends it a DagFileParseRequest. It returns the comm and // logs connections, the frame that Serve answers with, and the channel that gets what Serve // returns. func startDagParse( t *testing.T, - dags bundle.Registry, + dags *serializedDags, ) (commConn, logsConn net.Conn, frame IncomingFrame, done <-chan error) { t.Helper() commAddr, logsAddr, commCh, logsCh, cleanup := startSupervisor(t) t.Cleanup(cleanup) served := make(chan error, 1) - go func() { served <- Serve(dags, commAddr, logsAddr) }() + go func() { served <- Serve(parseBundle{testBundle{}, dags}, commAddr, logsAddr) }() commConn = <-commCh t.Cleanup(func() { commConn.Close() }) diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index c570cfb7d7e52..398ab7d945406 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -80,7 +80,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.Registry, commAddr, logsAddr string) error { +func Serve(b bundle.Bundle, commAddr, logsAddr string) error { if commAddr == "" { return fmt.Errorf("missing --comm=host:port argument") } @@ -175,8 +175,9 @@ func Serve(b bundle.Registry, commAddr, logsAddr string) error { logger.Debug("Task execution complete") case *genmodels.DagFileParseRequest: - logger.Debug("Dag parsing mode", "file", msg.File) - result := parseDags(b, msg, logger) + logger.Info("Received Dag parse request", "file", msg.File, "bundle_path", msg.BundlePath) + serializer, _ := b.(bundle.DagSerializer) + result := parseDags(serializer, msg, logger) // The Dag processor reads from a connection only after it checks that the connection // belongs to this process. The Dag processor cannot make that check after this process // exits. So a process that exits right after it sends the result often loses the logs of From 51fcfe68d8298a2ed54172f6becdb4b78fd34a57 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Wed, 7 Oct 2026 07:53:03 +0000 Subject: [PATCH 4/6] Go SDK: reply to the Dag parse request on its request id The runtime replies on the request id and exits without waiting for an acknowledgement, as the TS runtime and ADR 0003 do. --- go-sdk/README.md | 3 +- go-sdk/airflow/serve_test.go | 5 +- go-sdk/pkg/execution/integration_test.go | 61 +++++------------------- go-sdk/pkg/execution/server.go | 30 ++---------- 4 files changed, 21 insertions(+), 78 deletions(-) diff --git a/go-sdk/README.md b/go-sdk/README.md index 759d683923a23..f099bee045efe 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -392,8 +392,7 @@ Python supervisor / task runner - When the first frame on the comm socket is a `DagFileParseRequest` from the Dag processor, the binary runs no task. It answers with one `DagFileParsingResult` that holds the Dags from `airflow.Dag` that the binary registered. The Dags are serialized as - [Serializing a native Dag](#serializing-a-native-dag) describes. The binary exits once the Dag - processor acknowledges the result. + [Serializing a native Dag](#serializing-a-native-dag) describes. The binary then 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/serve_test.go b/go-sdk/airflow/serve_test.go index 608cc70e9e953..bdfb1ac590891 100644 --- a/go-sdk/airflow/serve_test.go +++ b/go-sdk/airflow/serve_test.go @@ -321,7 +321,8 @@ func TestServeAnswersTheDagParseRequestWithoutRunningATask(t *testing.T) { b.Register(etl, TaskHandler("py_etl", "transform", record("transform")), reports) supervisor, done := serveForSupervisor(t, b) - require.NoError(t, supervisor.SendRequest(0, map[string]any{ + const requestID = 7 + require.NoError(t, supervisor.SendRequest(requestID, map[string]any{ "type": "DagFileParseRequest", "file": "/bundles/go/etl", "bundle_path": "/bundles/go", @@ -330,6 +331,7 @@ func TestServeAnswersTheDagParseRequestWithoutRunningATask(t *testing.T) { frame, err := supervisor.ReadMessage() require.NoError(t, err) + assert.EqualValues(t, requestID, frame.ID) var result genmodels.DagFileParsingResult require.NoError(t, msgpack.Unmarshal(frame.Body, &result)) assert.Equal(t, "DagFileParsingResult", result.Type) @@ -340,7 +342,6 @@ func TestServeAnswersTheDagParseRequestWithoutRunningATask(t *testing.T) { want := jsonOf(t, dag.serialize("/bundles/go/etl", "etl")) assertJSON(t, want, result.SerializedDags[i].Data) } - require.NoError(t, supervisor.SendRequest(frame.ID, nil)) select { case err := <-done: diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 4f3411014ce22..3eb6cade4f8ec 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -1390,13 +1390,11 @@ type parseBundle struct { *serializedDags } -// startDagParse runs Serve for dags and sends it a DagFileParseRequest. It returns the comm and -// logs connections, the frame that Serve answers with, and the channel that gets what Serve -// returns. -func startDagParse( - t *testing.T, - dags *serializedDags, -) (commConn, logsConn net.Conn, frame IncomingFrame, done <-chan error) { +const dagParseRequestID = 7 + +// startDagParse runs Serve for dags and sends it a DagFileParseRequest. It returns the frame that +// Serve answers with and the channel that gets what Serve returns. +func startDagParse(t *testing.T, dags *serializedDags) (frame IncomingFrame, done <-chan error) { t.Helper() commAddr, logsAddr, commCh, logsCh, cleanup := startSupervisor(t) t.Cleanup(cleanup) @@ -1404,15 +1402,15 @@ func startDagParse( served := make(chan error, 1) go func() { served <- Serve(parseBundle{testBundle{}, dags}, commAddr, logsAddr) }() - commConn = <-commCh + commConn := <-commCh t.Cleanup(func() { commConn.Close() }) - logsConn = <-logsCh + logsConn := <-logsCh t.Cleanup(func() { logsConn.Close() }) deadline := time.Now().Add(10 * time.Second) require.NoError(t, commConn.SetDeadline(deadline)) require.NoError(t, logsConn.SetDeadline(deadline)) - payload, err := encodeRequest(0, map[string]any{ + payload, err := encodeRequest(dagParseRequestID, map[string]any{ "type": "DagFileParseRequest", "file": "/bundles/go/etl", "bundle_path": "/bundles/go", @@ -1424,13 +1422,14 @@ func startDagParse( frame, err = readFrame(commConn) require.NoError(t, err) require.True(t, isNilRaw(frame.Err)) - return commConn, logsConn, frame, served + return frame, served } func TestServeDagFileParseRequestEndToEnd(t *testing.T) { dags := &serializedDags{dags: []bundle.SerializedDag{serializedDag("etl")}} - commConn, _, frame, done := startDagParse(t, dags) + frame, done := startDagParse(t, dags) + assert.EqualValues(t, dagParseRequestID, frame.ID) var result genmodels.DagFileParsingResult require.NoError(t, decodeBody(frame.Body, &result)) assert.Equal(t, "DagFileParsingResult", result.Type) @@ -1438,47 +1437,11 @@ func TestServeDagFileParseRequestEndToEnd(t *testing.T) { require.Len(t, result.SerializedDags, 1) assert.Equal(t, "etl", result.SerializedDags[0].Data["dag"].(map[string]any)["dag_id"]) assert.Equal(t, "etl", dags.relative) - 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 Dag processor acknowledged the result") - } -} - -func TestServeFailsWhenTheDagProcessorRejectsTheParseResult(t *testing.T) { - commConn, logsConn, frame, done := startDagParse(t, &serializedDags{}) - - rejection := map[string]any{ - "type": "ErrorResponse", - "error": "generic_error", - "detail": map[string]any{"message": "A parse result was already received"}, - } - require.NoError(t, writeFrame(commConn, encodeResponseFrame(t, frame.ID, nil, rejection))) - - select { - case err := <-done: - require.ErrorContains(t, err, "sending Dag parsing result") - require.ErrorContains(t, err, "A parse result was already received") - case <-time.After(2 * time.Second): - t.Fatal("Serve did not return after the Dag processor rejected the result") - } - - // main's log.Fatal writes to the logs socket after Serve has closed that socket. The error - // therefore shows up only if Serve logs it first. - logs, err := io.ReadAll(logsConn) - require.NoError(t, err) - var failure map[string]any - for line := range bytes.Lines(logs) { - var record map[string]any - require.NoError(t, json.Unmarshal(line, &record)) - if record["event"] == "Failed to send the Dag parsing result" { - failure = record - } + t.Fatal("Serve did not return after it sent the Dag parsing result") } - require.NotNil(t, failure, "Serve did not log why it failed: %s", logs) - assert.Equal(t, "error", failure["level"]) - assert.Contains(t, failure["error"], "A parse result was already received") } diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index 398ab7d945406..f3ef70418fdd1 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -55,21 +55,15 @@ const dialTimeout = 30 * time.Second // 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. -// terminalSendTimeout also bounds the write of a DagFileParsingResult. 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, and the caller should then exit 0. A -// clean shutdown means one of two things. Either the task ran and its terminal -// TaskState/SucceedTask frame was delivered, or the DagFileParsingResult was -// delivered and the Dag processor acknowledged it. A non-nil error indicates a -// protocol-level failure (connection loss, malformed frames, unknown first -// message type). In a task run, the failure happens before or instead of the -// delivery of a terminal frame. In a Dag parse, the failure happens before the -// Dag processor acknowledges the DagFileParsingResult. +// Serve returns nil once it delivers the terminal frame of a task run or the DagFileParsingResult +// of a Dag parse, and the caller should then exit 0. A non-nil error indicates a protocol-level +// failure (connection loss, malformed frames, unknown first message type) before that. // // 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 @@ -178,25 +172,11 @@ func Serve(b bundle.Bundle, commAddr, logsAddr string) error { logger.Info("Received Dag parse request", "file", msg.File, "bundle_path", msg.BundlePath) serializer, _ := b.(bundle.DagSerializer) result := parseDags(serializer, msg, logger) - // The Dag processor reads from a connection only after it checks that the connection - // belongs to this process. The Dag processor cannot make that check after this process - // exits. So a process that exits right after it sends the result often loses the logs of - // its parse. Waiting for the acknowledgement keeps this process running until the Dag - // processor has read the result. The wait also gives the Dag processor time to check the - // logs connection. - // - // The write has a deadline. The wait for the acknowledgement has no deadline. The Dag - // processor validates the Dags in the result before it sends the acknowledgement. A larger - // Dag takes longer to validate. So a deadline on the wait could expire while the - // validation is still running. Even without a deadline, the wait cannot hang. The Dag - // processor kills a runtime that keeps running too long after the Dag processor receives - // the parse result from that runtime. + // Bound the write so a wedged socket cannot hang shutdown. _ = commConn.SetWriteDeadline(time.Now().Add(terminalSendTimeout)) - if _, err := comm.Communicate(ctx, result); err != nil { - logger.Error("Failed to send the Dag parsing result", "error", err) + if err := comm.SendRequest(frame.ID, result); err != nil { return fmt.Errorf("sending Dag parsing result: %w", err) } - logger.Debug("Dag parsing complete", "dags", len(result.SerializedDags)) default: logger.Error("Unexpected initial message type", "type", fmt.Sprintf("%T", body)) From bd64f0eed82947f5aebd31780253bf2db3dbef00 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Wed, 7 Oct 2026 07:53:14 +0000 Subject: [PATCH 5/6] Go SDK: log the stack when serializing a Dag panics The panic becomes the Dag's error, and the stack now reaches the Dag processor log through the default slog logger. --- go-sdk/airflow/bundle.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/go-sdk/airflow/bundle.go b/go-sdk/airflow/bundle.go index 788f59cdd4bdb..9c01808d8f727 100644 --- a/go-sdk/airflow/bundle.go +++ b/go-sdk/airflow/bundle.go @@ -19,6 +19,8 @@ package airflow import ( "fmt" + "log/slog" + "runtime/debug" "slices" "sync" "sync/atomic" @@ -242,6 +244,10 @@ func serializeRecovering(dag *DagRef, fileloc, relativeFileloc string) (s bundle s.DagID = dag.dagID defer func() { if r := recover(); r != nil { + slog.Error( + "Dag serialization panicked", + "dag_id", dag.dagID, "panic", r, "stack", string(debug.Stack()), + ) s.Data, s.Err = nil, fmt.Errorf("%v", r) } }() From e84f9b3bbf9918bd277add6af0f9d1667fb7161f Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Wed, 7 Oct 2026 07:53:48 +0000 Subject: [PATCH 6/6] Go SDK: trim the comments of the Dag parse path The comments now state the behavior in one or two lines. --- go-sdk/airflow/bundle.go | 7 +++---- go-sdk/airflow/bundle_test.go | 3 +-- go-sdk/airflow/dag.go | 5 ++--- go-sdk/airflow/serve.go | 7 +++---- go-sdk/airflow/serve_test.go | 8 +++----- go-sdk/internal/bundle/doc.go | 11 +++++------ go-sdk/pkg/execution/dag_parse_test.go | 6 ++---- go-sdk/pkg/execution/integration_test.go | 4 ++-- go-sdk/pkg/execution/server.go | 7 ++----- 9 files changed, 23 insertions(+), 35 deletions(-) diff --git a/go-sdk/airflow/bundle.go b/go-sdk/airflow/bundle.go index 9c01808d8f727..df10c5d5255fe 100644 --- a/go-sdk/airflow/bundle.go +++ b/go-sdk/airflow/bundle.go @@ -195,7 +195,7 @@ func (m *taskHandlerMap) ListTaskHandlers() []bundle.TaskHandlerInfo { return slices.Clone(m.order) } -// dagMap holds the registered Dags by dag_id and keeps the order in which they were registered. +// dagMap holds the registered Dags by dag_id, in registration order. type dagMap struct { mu sync.Mutex dags map[string]*DagRef @@ -225,9 +225,8 @@ func (m *dagMap) has(dagID string) bool { return exists } -// serialize serializes the registered Dags in the order they were registered. If serializing a Dag -// panics, the panic becomes the error of that Dag instead of ending the parse. The import error of -// the parse then names the Dag. +// serialize serializes the Dags in registration order. If a Dag panics, the panic becomes that +// Dag's Err. func (m *dagMap) serialize(fileloc, relativeFileloc string) []bundle.SerializedDag { m.mu.Lock() dags := slices.Clone(m.order) diff --git a/go-sdk/airflow/bundle_test.go b/go-sdk/airflow/bundle_test.go index c49af008e8209..d6c2da67f5120 100644 --- a/go-sdk/airflow/bundle_test.go +++ b/go-sdk/airflow/bundle_test.go @@ -418,8 +418,7 @@ func TestSerializeDagsKeepsTheOtherDagsWhenADagCannotBeSerialized(t *testing.T) etl := Dag("etl") etl.Task(noop) b.Register(etl) - // Register checks and expands every Dag it takes. A Dag that skipped Register therefore stands - // in for a Dag that the serializer fails on. + // A Dag that skipped Register stands in for one the serializer fails on. broken := Dag("broken") b.dags.dags["broken"] = broken b.dags.order = append(b.dags.order, broken) diff --git a/go-sdk/airflow/dag.go b/go-sdk/airflow/dag.go index 11b83c0dc983e..9415e514caf2c 100644 --- a/go-sdk/airflow/dag.go +++ b/go-sdk/airflow/dag.go @@ -71,9 +71,8 @@ type DagRef struct { // [IfRef.Then], [IfRef.Else], [SwitchRef.Case] and the methods of [TaskGroupRef] panic once the Dag // is registered. // -// When the Dag processor asks to parse the bundle, [BundleRef.Serve] answers with the registered -// Dags in serialized form. Serve does not yet list these Dags in the --airflow-metadata manifest or -// run their tasks. +// [BundleRef.Serve] sends the registered Dags to the Dag processor, but does not yet list them in +// the --airflow-metadata manifest or run their tasks. // // Dag panics if it gets more than one DagSpec, or if the DagSpec has a value that Python rejects // when it builds or validates a Dag: diff --git a/go-sdk/airflow/serve.go b/go-sdk/airflow/serve.go index 0985911a04fc8..7fafd79596529 100644 --- a/go-sdk/airflow/serve.go +++ b/go-sdk/airflow/serve.go @@ -57,10 +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. -// Airflow starts the executable with --comm and --logs. With these flags, Serve speaks the -// coordinator protocol. It either runs one task or answers the Dag processor's request to parse -// the bundle. To answer a parse request, Serve sends the Dags from [Dag] that the bundle registered -// and runs no task. +// Airflow starts the executable with --comm and --logs. Serve then speaks the coordinator +// protocol and either runs one task or answers the Dag processor's parse request with the Dags +// from [Dag]. // // 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/airflow/serve_test.go b/go-sdk/airflow/serve_test.go index bdfb1ac590891..bca506bc14c6a 100644 --- a/go-sdk/airflow/serve_test.go +++ b/go-sdk/airflow/serve_test.go @@ -220,9 +220,8 @@ func TestServeHelpIsNotAnError(t *testing.T) { assert.Contains(t, stdout.String(), "--airflow-metadata") } -// serveForSupervisor runs b.serve with --comm and --logs, the same flags that Airflow starts the -// bundle with. It returns the supervisor's end of the comm socket and the channel that gets what -// serve returns. +// serveForSupervisor runs b.serve with --comm and --logs and returns the supervisor's end of the +// comm socket and the channel that gets what serve returns. func serveForSupervisor(t *testing.T, b *BundleRef) (*execution.CoordinatorComm, <-chan error) { t.Helper() commLn, err := net.Listen("tcp", "127.0.0.1:0") @@ -294,8 +293,7 @@ func TestServeRunsTaskForSupervisor(t *testing.T) { assert.True(t, ran) } -// To parse the bundle, the Dag processor starts it with --comm and --logs and then sends a -// DagFileParseRequest over the comm socket. +// The Dag processor sends a DagFileParseRequest instead of StartupDetails. func TestServeAnswersTheDagParseRequestWithoutRunningATask(t *testing.T) { var ran []string record := func(name string) func(Context) error { diff --git a/go-sdk/internal/bundle/doc.go b/go-sdk/internal/bundle/doc.go index af14047d0f0a8..8a77698fe70a8 100644 --- a/go-sdk/internal/bundle/doc.go +++ b/go-sdk/internal/bundle/doc.go @@ -15,11 +15,10 @@ // specific language governing permissions and limitations // under the License. -// Package bundle defines what the coordinator runtime needs from a bundle. The -// runtime needs the tasks to look up and run, and the Dag and task ids to list -// in the manifest. It also needs the serialized Dags to send in answer to a Dag -// parse request. +// Package bundle defines what the coordinator runtime needs from a bundle: the +// tasks it looks up and runs, the Dag and task ids it lists in the manifest, and +// the serialized Dags it sends to the Dag processor. // -// Package airflow builds the tasks, the ids, and the serialized Dags from the -// task handlers and the Dags that a bundle registers. +// Package airflow builds the tasks and ids from the task handlers a bundle +// registers, and the serialized Dags from its Dags. package bundle diff --git a/go-sdk/pkg/execution/dag_parse_test.go b/go-sdk/pkg/execution/dag_parse_test.go index e0d2078cba452..27f4d8a95d417 100644 --- a/go-sdk/pkg/execution/dag_parse_test.go +++ b/go-sdk/pkg/execution/dag_parse_test.go @@ -33,8 +33,7 @@ import ( "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" ) -// serializedDags serializes the Dags it holds. It also records the fileloc and relativeFileloc -// that SerializeDags was called with. +// serializedDags is a DagSerializer that records the paths it was called with. type serializedDags struct { dags []bundle.SerializedDag fileloc, relative string @@ -60,8 +59,7 @@ var etlParseRequest = &genmodels.DagFileParseRequest{ func discardLogger() *slog.Logger { return slog.New(slog.NewTextHandler(io.Discard, nil)) } -// wireJSON passes body through the frame encoding that SendRequest uses and returns the result as -// JSON. A test then sees the keys, nulls and empty lists that the Dag processor receives. +// wireJSON returns body as the Dag processor receives it, after the frame encoding of SendRequest. func wireJSON(t *testing.T, body any) string { t.Helper() payload, err := encodeRequest(0, body) diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 3eb6cade4f8ec..8e4b24358c83b 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -1361,8 +1361,8 @@ func TestServeFailureAfterConnectClosesComm(t *testing.T) { logsConn := <-logsCh defer logsConn.Close() - // Serve expects StartupDetails or DagFileParseRequest as the first frame, so it fails to - // decode a VariableResult. + // Serve expects StartupDetails or DagFileParseRequest first, so it fails to decode a + // VariableResult. payload, err := encodeRequest( 0, map[string]any{"type": "VariableResult", "key": "k", "value": "v"}, diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index f3ef70418fdd1..062658930d0ba 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -20,11 +20,8 @@ // the Airflow supervisor (Python ExecutableCoordinator), the Serve method of // airflow.BundleRef dispatches here. // -// The first frame that the runtime reads on the comm socket decides what the -// runtime does. After a StartupDetails message, the runtime runs a task over -// several rounds of messages. After a DagFileParseRequest from the Dag -// processor, the runtime runs no task. It answers with one DagFileParsingResult -// that holds the Dags from airflow.Dag that the bundle registered. +// The first frame on the comm socket picks the mode: StartupDetails runs one +// task, and DagFileParseRequest is answered with the bundle's serialized Dags. // // See go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md. package execution