Skip to content

Commit cb9baf8

Browse files
committed
Report a Lang-SDK parse past its import timeout as an import error
The parse child resolves the get_dagbag_import_timeout policy for the file and sends the timeout with the runtime's schema version. The policy thus runs where it does for a Python Dag file, and a failure is only that file's import error. The timeout counts from the start of the parse. Past it, with no result yet, the runtime is killed and the file gets an import error. LangSDKDagFileProcessorProcess.run applies the same timeout, so it no longer takes one. Until the child reports it, dag_file_processor_timeout bounds the parse, as it does in the Dag processor. At the deadline run stops waiting even if an exited runtime left a process that holds its output open.
1 parent 789dc06 commit cb9baf8

5 files changed

Lines changed: 174 additions & 87 deletions

File tree

‎airflow-core/src/airflow/dag_processing/lang_sdk_processor.py‎

Lines changed: 43 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@
3232
from pydantic import BaseModel, Field, TypeAdapter
3333
from uuid6 import uuid7
3434

35+
from airflow import settings
36+
from airflow.configuration import conf
3537
from airflow.dag_processing.importer_routing import get_claiming_coordinator
3638
from airflow.dag_processing.processor import (
3739
BaseDagFileProcessorProcess,
@@ -78,12 +80,22 @@ class StartLangSDKRuntime(BaseModel):
7880

7981

8082
class LangSDKRuntimeSchemaVersion(BaseModel):
81-
"""The supervisor schema version of the runtime the parse child is about to exec."""
83+
"""The schema version and the import timeout of the runtime the parse child is about to exec."""
8284

8385
schema_version: str | None
86+
import_timeout: float | None = None
87+
"""Seconds from the start of the parse; ``None`` means no timeout."""
8488
type: Literal["LangSDKRuntimeSchemaVersion"] = "LangSDKRuntimeSchemaVersion"
8589

8690

91+
def _get_import_timeout(path: str) -> float | None:
92+
"""Return the ``get_dagbag_import_timeout`` policy's timeout for *path*; ``None`` means none."""
93+
timeout = settings.get_dagbag_import_timeout(path)
94+
if not isinstance(timeout, (int, float)):
95+
raise TypeError(f"Value ({timeout}) from get_dagbag_import_timeout must be int or float")
96+
return timeout if timeout > 0 else None
97+
98+
8799
def _start_runtime_entrypoint() -> None:
88100
"""Exec the runtime that parses the file named by the start request, or report why it cannot start."""
89101
os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "client"
@@ -101,9 +113,11 @@ def _start_runtime_entrypoint() -> None:
101113
raise RuntimeError(f"Required first message to be a StartLangSDKRuntime, it was {msg}")
102114

103115
def report_schema_version(schema_version: str | None) -> None:
104-
comms.send(LangSDKRuntimeSchemaVersion(schema_version=schema_version))
116+
comms.send(LangSDKRuntimeSchemaVersion(schema_version=schema_version, import_timeout=import_timeout))
105117

106118
try:
119+
# The policy is user code: it runs in this child, where a failure is only this file's import error.
120+
import_timeout = _get_import_timeout(msg.file)
107121
coordinator = get_claiming_coordinator(msg.file, msg.bundle_name)
108122
if coordinator is None:
109123
raise RuntimeError(f"No coordinator's Dag importer claims {msg.file}")
@@ -149,6 +163,7 @@ class LangSDKDagFileProcessorProcess(BaseDagFileProcessorProcess):
149163
_listeners: dict[_Channel, socket]
150164
_parse_request: DagFileParseRequest
151165
_runtime_schema_version: str | None = attrs.field(default=None, init=False)
166+
_import_timeout: float | None = attrs.field(default=None, init=False)
152167
_schema_version_reported: bool = attrs.field(default=False, init=False)
153168
_parsing_result_monotonic: float | None = attrs.field(default=None, init=False)
154169
_unverified_connections: list[tuple[socket, _Channel]] = attrs.field(factory=list, init=False)
@@ -212,17 +227,16 @@ def run(
212227
bundle_path: Path,
213228
bundle_name: str,
214229
dag_file_rel_path: str,
215-
timeout: float | None,
216230
logger: FilteringBoundLogger,
217231
) -> DagFileParsingResult:
218232
"""
219233
Parse *path* outside the Dag processor and wait for the result.
220234
221-
There is no API client, so each request of the runtime that needs one gets an error.
222-
223-
:raises TimeoutError: if the parse does not finish within *timeout* seconds. The runtime is
224-
killed.
235+
There is no API client, so each request of the runtime that needs one gets an error. The file's import
236+
timeout bounds the parse, and ``[dag_processor] dag_file_processor_timeout`` until the parse child
237+
reports it.
225238
"""
239+
processor_timeout = conf.getfloat("dag_processor", "dag_file_processor_timeout")
226240
with selectors.DefaultSelector() as selector:
227241
proc = cls.start(
228242
id=uuid7(),
@@ -235,14 +249,13 @@ def run(
235249
)
236250
try:
237251
while not proc.is_ready:
238-
wait = 0.1
239-
if timeout is not None:
240-
if (remaining := proc.start_time + timeout - time.monotonic()) <= 0:
241-
raise TimeoutError(
242-
f"The Lang-SDK runtime did not parse {os.fspath(path)} within {timeout}s"
243-
)
244-
wait = min(wait, remaining)
245-
proc._service_subprocess(max_wait_time=wait)
252+
timeout = proc._import_timeout if proc._schema_version_reported else processor_timeout
253+
if timeout is not None and time.monotonic() - proc.start_time > timeout:
254+
# Unlike is_ready, this does not wait for an exited runtime's leftover processes,
255+
# which can hold its sockets open. close() closes them.
256+
proc._time_out(timeout)
257+
break
258+
proc._service_subprocess(max_wait_time=0.1)
246259
except BaseException:
247260
proc._kill_runtime()
248261
raise
@@ -376,6 +389,7 @@ def _handle_runtime_schema_version(
376389
self._reject_request(msg, log, req_id)
377390
return ResponseSent.ALREADY_SENT
378391
self._runtime_schema_version = msg.schema_version
392+
self._import_timeout = msg.import_timeout
379393
self._schema_version_reported = True
380394
return None, {}
381395

@@ -438,6 +452,13 @@ def is_ready(self) -> bool:
438452
):
439453
self.process_log.warning("The Lang-SDK runtime did not exit after its parse result; killing it")
440454
self._kill_runtime()
455+
if (
456+
self._import_timeout is not None
457+
and self.parsing_result is None
458+
and self._exit_code is None
459+
and time.monotonic() - self.start_time > self._import_timeout
460+
):
461+
self._time_out(self._import_timeout)
441462
if self._check_subprocess_exit() is None:
442463
return False
443464
self._close_listeners()
@@ -449,6 +470,13 @@ def is_ready(self) -> bool:
449470
)
450471
return True
451472

473+
def _time_out(self, timeout: float) -> None:
474+
if self.parsing_result is None:
475+
self._set_import_error(
476+
f"The Lang-SDK runtime did not parse {self._parse_request.file} within {timeout}s"
477+
)
478+
self._kill_runtime()
479+
452480
def _kill_runtime(self) -> None:
453481
"""Kill the runtime and wait for it, without servicing its sockets, whose handler may have failed."""
454482
if self._exit_code is not None:

‎airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py‎

Lines changed: 74 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
from airflow.dag_processing.lang_sdk_processor import (
3737
LangSDKDagFileProcessorProcess,
3838
LangSDKRuntimeSchemaVersion,
39+
_get_import_timeout,
3940
)
4041
from airflow.dag_processing.processor import DagFileParseRequest, DagFileParsingResult
4142
from airflow.sdk import DAG, BaseOperator
@@ -332,6 +333,27 @@ def reply(request, comms):
332333
assert proc._exit_code == -signal.SIGKILL
333334
assert "The Lang-SDK runtime did not exit after its parse result; killing it" in cap_structlog
334335

336+
@pytest.mark.parametrize(
337+
("policy", "error"),
338+
[
339+
pytest.param(
340+
{"side_effect": RuntimeError("policy bug")}, "RuntimeError: policy bug", id="raises"
341+
),
342+
pytest.param(
343+
{"return_value": "30"},
344+
"TypeError: Value (30) from get_dagbag_import_timeout must be int or float",
345+
id="not-a-number",
346+
),
347+
],
348+
)
349+
def test_a_failing_import_timeout_policy_is_an_import_error(self, parse, policy, error):
350+
with patch("airflow.settings.get_dagbag_import_timeout", autospec=True, **policy):
351+
proc = parse()
352+
353+
assert proc.parsing_result.import_errors == {
354+
"dag.native": f"Cannot start the Lang-SDK runtime: {error}"
355+
}
356+
335357
@pytest.mark.skipif(not Path("/proc/self/fd").is_dir(), reason="reads /proc")
336358
@pytest.mark.parametrize("use_exec", [False, True], ids=["fork", "spawn"])
337359
def test_the_runtime_inherits_only_its_standard_streams(self, monkeypatch, tmp_path, use_exec):
@@ -357,13 +379,12 @@ def test_the_runtime_inherits_only_its_standard_streams(self, monkeypatch, tmp_p
357379

358380
class TestRun:
359381
@staticmethod
360-
def _run(tmp_path, *, timeout: float | None = 30) -> DagFileParsingResult:
382+
def _run(tmp_path, **spec) -> DagFileParsingResult:
361383
return LangSDKDagFileProcessorProcess.run(
362-
path=write_native_file(tmp_path / "dag.native"),
384+
path=write_native_file(tmp_path / "dag.native", **spec),
363385
bundle_path=tmp_path,
364386
bundle_name="testing",
365387
dag_file_rel_path="dag.native",
366-
timeout=timeout,
367388
logger=structlog.get_logger(),
368389
)
369390

@@ -398,9 +419,9 @@ def reply(request, comms):
398419
mock_mask_secret.assert_called_once_with("native-secret", "native_conn")
399420

400421
@pytest.mark.parametrize("connected", [False, True], ids=["before-connecting", "after-connecting"])
401-
def test_a_parse_past_its_timeout_is_killed(self, tmp_path, connected):
422+
@patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1)
423+
def test_a_parse_past_the_import_timeout_is_killed(self, mock_timeout, tmp_path, connected):
402424
# A runtime that never connects leaves both listeners open when it is killed.
403-
write_native_file(tmp_path / "dag.native", argv=["/bin/sh", "-c", "exec sleep 60"])
404425
fds_before = _get_open_fds()
405426

406427
with (
@@ -418,22 +439,61 @@ def test_a_parse_past_its_timeout_is_killed(self, tmp_path, connected):
418439
autospec=True,
419440
side_effect=LangSDKDagFileProcessorProcess.close,
420441
) as mock_close,
421-
pytest.raises(TimeoutError, match=r"did not parse .*dag\.native within 1s"),
422442
):
423-
LangSDKDagFileProcessorProcess.run(
424-
path=tmp_path / "dag.native",
425-
bundle_path=tmp_path,
426-
bundle_name="testing",
427-
dag_file_rel_path="dag.native",
428-
timeout=1,
429-
logger=structlog.get_logger(),
430-
)
443+
result = self._run(tmp_path, argv=["/bin/sh", "-c", "exec sleep 60"])
431444

445+
assert result.import_errors == {
446+
"dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} within 1.0s"
447+
}
432448
[proc] = [c.args[0] for c in mock_close.call_args_list]
433449
assert proc._exit_code == -9
434450
assert not proc._open_sockets
435451
assert _get_open_fds() <= fds_before
436452

453+
@patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1)
454+
def test_the_import_timeout_holds_after_the_runtime_exits(self, mock_timeout, tmp_path):
455+
with patch.object(
456+
LangSDKDagFileProcessorProcess,
457+
"close",
458+
autospec=True,
459+
side_effect=LangSDKDagFileProcessorProcess.close,
460+
) as mock_close:
461+
# The runtime exits, and the process it leaves behind keeps its output open.
462+
result = self._run(tmp_path, argv=["/bin/sh", "-c", "sleep 30 & exit 0"])
463+
[proc] = [c.args[0] for c in mock_close.call_args_list]
464+
os.killpg(proc.pid, signal.SIGKILL)
465+
466+
assert result.import_errors == {
467+
"dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} within 1.0s"
468+
}
469+
assert proc._exit_code == 0
470+
assert not proc._open_sockets
471+
472+
@conf_vars({("dag_processor", "dag_file_processor_timeout"): "1"})
473+
@patch.object(
474+
FakeCoordinator,
475+
"_build_parse_dag_command",
476+
autospec=True,
477+
side_effect=lambda self, *, path: time.sleep(60),
478+
)
479+
def test_the_dag_file_processor_timeout_applies_until_the_import_timeout_is_reported(
480+
self, mock_build_parse_dag_command, tmp_path
481+
):
482+
result = self._run(tmp_path)
483+
484+
assert result.import_errors == {
485+
"dag.native": f"The Lang-SDK runtime did not parse {tmp_path / 'dag.native'} within 1.0s"
486+
}
487+
488+
489+
@pytest.mark.parametrize(("configured", "expected"), [(30, 30), (0.5, 0.5), (0, None), (-1, None)])
490+
@patch("airflow.settings.get_dagbag_import_timeout", autospec=True)
491+
def test_only_a_positive_import_timeout_applies(mock_timeout, configured, expected):
492+
mock_timeout.return_value = configured
493+
494+
assert _get_import_timeout("/b/dag.native") == expected
495+
mock_timeout.assert_called_once_with("/b/dag.native")
496+
437497

438498
def _make_process(**kwargs) -> LangSDKDagFileProcessorProcess:
439499
return LangSDKDagFileProcessorProcess(

‎airflow-core/tests/unit/dag_processing/test_manager.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1768,6 +1768,54 @@ def reply(request, comms):
17681768
assert import_error.stacktrace.startswith("The Lang-SDK runtime sent an invalid frame: ")
17691769
assert manager.selector.get_map() == {}
17701770

1771+
@mock.patch("airflow.settings.get_dagbag_import_timeout", autospec=True, return_value=1)
1772+
@mock.patch.object(FakeCoordinator, "parse_dag", autospec=True)
1773+
def test_a_coordinator_file_past_the_import_timeout_is_an_import_error(
1774+
self, mock_parse_dag, mock_timeout, tmp_path, configure_testing_dag_bundle
1775+
):
1776+
mock_parse_dag.side_effect = play_runtime(lambda request, comms: time.sleep(60))
1777+
slow = write_native_file(tmp_path / "slow.native")
1778+
1779+
with fake_coordinator(), configure_testing_dag_bundle(tmp_path):
1780+
manager = DagFileProcessorManager(max_runs=1, processor_timeout=60)
1781+
manager.run()
1782+
1783+
with create_session() as session:
1784+
[import_error] = session.scalars(select(ParseImportError)).all()
1785+
1786+
assert import_error.filename == "slow.native"
1787+
assert import_error.stacktrace == f"The Lang-SDK runtime did not parse {slow} within 1.0s"
1788+
assert manager.selector.get_map() == {}
1789+
1790+
@mock.patch("airflow.settings.get_dagbag_import_timeout", autospec=True)
1791+
def test_a_failing_import_timeout_policy_fails_only_its_file(
1792+
self, mock_timeout, tmp_path, configure_testing_dag_bundle
1793+
):
1794+
def get_dagbag_import_timeout(dag_file_path):
1795+
if dag_file_path.endswith(".native"):
1796+
raise RuntimeError("policy bug")
1797+
return 30
1798+
1799+
mock_timeout.side_effect = get_dagbag_import_timeout
1800+
write_native_file(tmp_path / "broken.native")
1801+
(tmp_path / "python_dag.py").write_text(
1802+
"from airflow.sdk import DAG\nfrom airflow.sdk.bases.operator import BaseOperator\n\n"
1803+
'with DAG("python_dag", schedule=None):\n BaseOperator(task_id="task")\n'
1804+
)
1805+
1806+
with fake_coordinator(), configure_testing_dag_bundle(tmp_path):
1807+
DagFileProcessorManager(max_runs=1, processor_timeout=60).run()
1808+
1809+
with create_session() as session:
1810+
dag_ids = session.scalars(select(SerializedDagModel.dag_id)).all()
1811+
[import_error] = session.scalars(select(ParseImportError)).all()
1812+
1813+
assert dag_ids == ["python_dag"]
1814+
assert (import_error.filename, import_error.stacktrace) == (
1815+
"broken.native",
1816+
"Cannot start the Lang-SDK runtime: RuntimeError: policy bug",
1817+
)
1818+
17711819
def test_terminate_orphan_processes_kills_then_closes_processor(self):
17721820
manager = DagFileProcessorManager(max_runs=1)
17731821
processor, _ = self.mock_processor()

‎task-sdk/src/airflow/sdk/coordinators/_dag_importer.py‎

Lines changed: 8 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323

2424
import structlog
2525

26-
from airflow.sdk.exceptions import AirflowConfigException
2726
from airflow.sdk.importers.base import (
2827
AbstractDagImporter,
2928
DagImportError,
@@ -45,19 +44,6 @@
4544
log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.dag_importer")
4645

4746

48-
def _get_import_timeout(definition: FilesystemDagDefinition) -> float | None:
49-
"""Return the parse timeout for *definition*, as ``PythonDagImporter`` does; ``None`` means none."""
50-
try:
51-
from airflow import settings # noqa: SDK002
52-
53-
timeout = settings.get_dagbag_import_timeout(repr(definition))
54-
except (ImportError, AttributeError):
55-
timeout = 30.0
56-
if not isinstance(timeout, (int, float)):
57-
raise AirflowConfigException(f"Value ({timeout}) from get_dagbag_import_timeout must be int or float")
58-
return timeout if timeout > 0 else None
59-
60-
6147
class CoordinatorDagImporter(AbstractDagImporter[FilesystemDagDefinition]):
6248
"""
6349
Import the native Dags of a coordinator's artifacts by running the coordinator's runtime.
@@ -104,23 +90,15 @@ def import_definition(
10490

10591
source_reference = repr(definition)
10692
bundle_path = bundle.path or definition.path.parent
107-
result = DagImportResult(definition=definition)
108-
try:
109-
parsing_result = LangSDKDagFileProcessorProcess.run(
110-
path=definition.path,
111-
bundle_path=bundle_path,
112-
bundle_name=bundle.name,
113-
dag_file_rel_path=definition.get_relative_loc(bundle_path),
114-
timeout=_get_import_timeout(definition),
115-
logger=log,
116-
)
117-
except TimeoutError as e:
118-
result.errors.append(
119-
DagImportError(source_reference=source_reference, message=str(e), error_type="timeout")
120-
)
121-
return result
122-
12393
relative_loc = definition.get_relative_loc(bundle_path)
94+
result = DagImportResult(definition=definition)
95+
parsing_result = LangSDKDagFileProcessorProcess.run(
96+
path=definition.path,
97+
bundle_path=bundle_path,
98+
bundle_name=bundle.name,
99+
dag_file_rel_path=relative_loc,
100+
logger=log,
101+
)
124102
for key, message in (parsing_result.import_errors or {}).items():
125103
result.errors.append(
126104
DagImportError(

0 commit comments

Comments
 (0)