Skip to content

Commit 492518e

Browse files
ashbashb
authored
Switch the Supervisor/task process from line-based to length-prefixed (#51699)
* Switch the Supervisor/task process from line-based to length-prefixed The existing JSON Lines based approach had two major drawbacks 1. In the case of really large lines (in the region of 10 or 20MB) the python line buffering could _sometimes_ result in a partial read 2. The JSON based approach didn't have the ability to add any metadata (such as errors). 3. Not every message type/call-site waited for a response, which meant those client functions could never get told about an error One of the ways this line-based approach fell down was if you suddenly tried to run 100s of triggers at the same time you would get an error like this: ``` Traceback (most recent call last): File "/Users/ash/.local/share/uv/python/cpython-3.12.7-macos-aarch64-none/lib/python3.12/asyncio/streams.py", line 568, in readline line = await self.readuntil(sep) ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/ash/.local/share/uv/python/cpython-3.12.7-macos-aarch64-none/lib/python3.12/asyncio/streams.py", line 663, in readuntil raise exceptions.LimitOverrunError( asyncio.exceptions.LimitOverrunError: Separator is found, but chunk is longer than limit ``` The other way this caused problems was if you parse a large dag (as in one with 20k tasks or more) the DagFileProcessor could end up getting a partial read which would be invalid JSON. This changes the communications protocol in in a couple of ways. First off at the python level the separate send and receive methods in the client/task side have been removed and replaced with a single `send()` that sends the request, reads the response and raises an error if one is returned. (But note, right now almost nothing in the supervisor side sets the error, that will be a future PR.) Secondly the JSON Lines approach has been changed from a line-based protocol to a binary "frame" one. The protocol (which is the same for whichever side is sending) is length-prefixed, i.e. we first send the length of the data as a 4byte big-endian integer, followed by the data itself. This should remove the possibility of JSON parse errors due to reading incomplete lines Finally the last change made in this PR is to remove the "extra" requests socket/channel. Upon closer examination with this comms path I realised that this socket is unnecessary: Since we are in 100% control of the client side we can make use of the bi-directional nature of `socketpair` and save file handles. This also happens to help the `run_as_user` feature which is currently broken, as without extra config to `sudoers` file, `sudo` will close all filehandles other than stdin, stdout, and stderr -- so by introducing this change we make it easier to re-add run_as_user support. In order to support this in the DagFileProcessor (as the fact that the proc manager uses a single selector for multiple processes) means I have moved the `on_close` callback to be part of the object we store in the `selector` object in the supervisors, previoulsy it was the "on_read" callback, now we store a tuple of `(on_read, on_close)` and on_close is called once universally. This also changes the way comms are handled from the (async) TriggerRunner process. Previously we had a sync+async lock, but that made it possible to end up deadlocking things. The change now is to have `send` on `TriggerCommsDecoder` "go back" to the async even loop via `async_to_sync`, so that only async code deals with the socket, and we can use an async lock (rather than the hybrid sync and async lock we tried before). This seems to help the deadlock issue, but I'm not 100% sure it will remove it entirely, but it makes it much much harder to hit - I've not been able to reprouce it with this change * Deal with compat in tests This compat issue is only in tests, as nothing in the runtime of airflow-core imports/calls methods directly on SUPERVISOR_COMMS, we are only importing it in tests to mkae assertions about the behavour/to stub the return values.
1 parent af94629 commit 492518e

48 files changed

Lines changed: 1249 additions & 1218 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,8 @@
7777
from airflow.utils.sqlalchemy import prohibit_commit, with_row_locks
7878

7979
if TYPE_CHECKING:
80+
from socket import socket
81+
8082
from sqlalchemy.orm import Session
8183

8284
from airflow.callbacks.callback_requests import CallbackRequest
@@ -388,7 +390,7 @@ def _service_processor_sockets(self, timeout: float | None = 1.0):
388390
"""
389391
events = self.selector.select(timeout=timeout)
390392
for key, _ in events:
391-
socket_handler = key.data
393+
socket_handler, on_close = key.data
392394

393395
# BrokenPipeError should be caught and treated as if the handler returned false, similar
394396
# to EOF case
@@ -397,8 +399,9 @@ def _service_processor_sockets(self, timeout: float | None = 1.0):
397399
except BrokenPipeError:
398400
need_more = False
399401
if not need_more:
400-
self.selector.unregister(key.fileobj)
401-
key.fileobj.close() # type: ignore[union-attr]
402+
sock: socket = key.fileobj # type: ignore[assignment]
403+
on_close(sock)
404+
sock.close()
402405

403406
def _queue_requested_files_for_parsing(self) -> None:
404407
"""Queue any files requested for parsing as requested by users via UI/API."""

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

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,6 @@ class DagFileParseRequest(BaseModel):
6868
bundle_path: Path
6969
"""Passing bundle path around lets us figure out relative file path."""
7070

71-
requests_fd: int
7271
callback_requests: list[CallbackRequest] = Field(default_factory=list)
7372
type: Literal["DagFileParseRequest"] = "DagFileParseRequest"
7473

@@ -102,18 +101,16 @@ class DagFileParsingResult(BaseModel):
102101
def _parse_file_entrypoint():
103102
import structlog
104103

105-
from airflow.sdk.execution_time import task_runner
104+
from airflow.sdk.execution_time import comms, task_runner
106105

107106
# Parse DAG file, send JSON back up!
108-
comms_decoder = task_runner.CommsDecoder[ToDagProcessor, ToManager](
109-
input=sys.stdin,
110-
decoder=TypeAdapter[ToDagProcessor](ToDagProcessor),
107+
comms_decoder = comms.CommsDecoder[ToDagProcessor, ToManager](
108+
body_decoder=TypeAdapter[ToDagProcessor](ToDagProcessor),
111109
)
112110

113-
msg = comms_decoder.get_message()
111+
msg = comms_decoder._get_response()
114112
if not isinstance(msg, DagFileParseRequest):
115113
raise RuntimeError(f"Required first message to be a DagFileParseRequest, it was {msg}")
116-
comms_decoder.request_socket = os.fdopen(msg.requests_fd, "wb", buffering=0)
117114

118115
task_runner.SUPERVISOR_COMMS = comms_decoder
119116
log = structlog.get_logger(logger_name="task")
@@ -125,7 +122,7 @@ def _parse_file_entrypoint():
125122

126123
result = _parse_file(msg, log)
127124
if result is not None:
128-
comms_decoder.send_request(log, result)
125+
comms_decoder.send(result)
129126

130127

131128
def _parse_file(msg: DagFileParseRequest, log: FilteringBoundLogger) -> DagFileParsingResult | None:
@@ -266,20 +263,18 @@ def _on_child_started(
266263
msg = DagFileParseRequest(
267264
file=os.fspath(path),
268265
bundle_path=bundle_path,
269-
requests_fd=self._requests_fd,
270266
callback_requests=callbacks,
271267
)
272-
self.send_msg(msg)
268+
self.send_msg(msg, request_id=0)
273269

274-
def _handle_request(self, msg: ToManager, log: FilteringBoundLogger) -> None: # type: ignore[override]
270+
def _handle_request(self, msg: ToManager, log: FilteringBoundLogger, req_id: int) -> None: # type: ignore[override]
275271
from airflow.sdk.api.datamodels._generated import ConnectionResponse, VariableResponse
276272

277273
resp: BaseModel | None = None
278274
dump_opts = {}
279275
if isinstance(msg, DagFileParsingResult):
280276
self.parsing_result = msg
281-
return
282-
if isinstance(msg, GetConnection):
277+
elif isinstance(msg, GetConnection):
283278
conn = self.client.connections.get(msg.conn_id)
284279
if isinstance(conn, ConnectionResponse):
285280
conn_result = ConnectionResult.from_conn_response(conn)
@@ -301,18 +296,24 @@ def _handle_request(self, msg: ToManager, log: FilteringBoundLogger) -> None: #
301296
resp = self.client.variables.delete(msg.key)
302297
else:
303298
log.error("Unhandled request", msg=msg)
299+
self.send_msg(
300+
None,
301+
request_id=req_id,
302+
error=ErrorResponse(
303+
detail={"status_code": 400, "message": "Unhandled request"},
304+
),
305+
)
304306
return
305307

306-
if resp:
307-
self.send_msg(resp, **dump_opts)
308+
self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
308309

309310
@property
310311
def is_ready(self) -> bool:
311312
if self._check_subprocess_exit() is None:
312313
# Process still alive, def can't be finished yet
313314
return False
314315

315-
return self._num_open_sockets == 0
316+
return not self._open_sockets
316317

317318
def wait(self) -> int:
318319
raise NotImplementedError(f"Don't call wait on {type(self).__name__} objects")

‎airflow-core/src/airflow/jobs/triggerer_job_runner.py‎

Lines changed: 74 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from collections.abc import Generator, Iterable
2929
from contextlib import suppress
3030
from datetime import datetime
31+
from socket import socket
3132
from traceback import format_exception
3233
from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal, TypedDict, Union
3334

@@ -43,6 +44,7 @@
4344
from airflow.jobs.job import perform_heartbeat
4445
from airflow.models.trigger import Trigger
4546
from airflow.sdk.execution_time.comms import (
47+
CommsDecoder,
4648
ConnectionResult,
4749
DagRunStateResult,
4850
DRCount,
@@ -58,6 +60,7 @@
5860
TICount,
5961
VariableResult,
6062
XComResult,
63+
_RequestFrame,
6164
)
6265
from airflow.sdk.execution_time.supervisor import WatchedSubprocess, make_buffered_socket_reader
6366
from airflow.stats import Stats
@@ -70,8 +73,6 @@
7073
from airflow.utils.session import provide_session
7174

7275
if TYPE_CHECKING:
73-
from socket import socket
74-
7576
from sqlalchemy.orm import Session
7677
from structlog.typing import FilteringBoundLogger, WrappedLogger
7778

@@ -181,7 +182,6 @@ class messages:
181182
class StartTriggerer(BaseModel):
182183
"""Tell the async trigger runner process to start, and where to send status update messages."""
183184

184-
requests_fd: int
185185
type: Literal["StartTriggerer"] = "StartTriggerer"
186186

187187
class TriggerStateChanges(BaseModel):
@@ -295,7 +295,7 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
295295
"""
296296
TriggerRunnerSupervisor is responsible for monitoring the subprocess and marshalling DB access.
297297
298-
This class (which runs in the main process) is responsible for querying the DB, sending RunTrigger
298+
This class (which runs in the main/sync process) is responsible for querying the DB, sending RunTrigger
299299
workload messages to the subprocess, and collecting results and updating them in the DB.
300300
"""
301301

@@ -342,8 +342,8 @@ def start( # type: ignore[override]
342342
):
343343
proc = super().start(id=job.id, job=job, target=cls.run_in_process, logger=logger, **kwargs)
344344

345-
msg = messages.StartTriggerer(requests_fd=proc._requests_fd)
346-
proc.send_msg(msg)
345+
msg = messages.StartTriggerer()
346+
proc.send_msg(msg, request_id=0)
347347
return proc
348348

349349
@functools.cached_property
@@ -355,7 +355,7 @@ def client(self) -> Client:
355355
client.base_url = "http://in-process.invalid./" # type: ignore[assignment]
356356
return client
357357

358-
def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger) -> None: # type: ignore[override]
358+
def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger, req_id: int) -> None: # type: ignore[override]
359359
from airflow.sdk.api.datamodels._generated import (
360360
ConnectionResponse,
361361
TaskStatesResponse,
@@ -454,8 +454,7 @@ def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger) -
454454
else:
455455
raise ValueError(f"Unknown message type {type(msg)}")
456456

457-
if resp:
458-
self.send_msg(resp, **dump_opts)
457+
self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
459458

460459
def run(self) -> None:
461460
"""Run synchronously and handle all database reads/writes."""
@@ -628,7 +627,7 @@ def _register_pipe_readers(self, stdout: socket, stderr: socket, requests: socke
628627
),
629628
)
630629

631-
def _process_log_messages_from_subprocess(self) -> Generator[None, bytes, None]:
630+
def _process_log_messages_from_subprocess(self) -> Generator[None, bytes | bytearray, None]:
632631
import msgspec
633632
from structlog.stdlib import NAME_TO_LEVEL
634633

@@ -691,14 +690,60 @@ class TriggerDetails(TypedDict):
691690
events: int
692691

693692

693+
@attrs.define(kw_only=True)
694+
class TriggerCommsDecoder(CommsDecoder[ToTriggerRunner, ToTriggerSupervisor]):
695+
_async_writer: asyncio.StreamWriter = attrs.field(alias="async_writer")
696+
_async_reader: asyncio.StreamReader = attrs.field(alias="async_reader")
697+
698+
body_decoder: TypeAdapter[ToTriggerRunner] = attrs.field(
699+
factory=lambda: TypeAdapter(ToTriggerRunner), repr=False
700+
)
701+
702+
_lock: asyncio.Lock = attrs.field(factory=asyncio.Lock, repr=False)
703+
704+
def _read_frame(self):
705+
from asgiref.sync import async_to_sync
706+
707+
return async_to_sync(self._aread_frame)()
708+
709+
def send(self, msg: ToTriggerSupervisor) -> ToTriggerRunner | None:
710+
from asgiref.sync import async_to_sync
711+
712+
return async_to_sync(self.asend)(msg)
713+
714+
async def _aread_frame(self):
715+
len_bytes = await self._async_reader.readexactly(4)
716+
len = int.from_bytes(len_bytes, byteorder="big")
717+
if len >= 2**32:
718+
raise OverflowError(f"Refusing to receive messages larger than 4GiB {len=}")
719+
720+
buffer = await self._async_reader.readexactly(len)
721+
return self.resp_decoder.decode(buffer)
722+
723+
async def _aget_response(self, expect_id: int) -> ToTriggerRunner | None:
724+
frame = await self._aread_frame()
725+
if frame.id != expect_id:
726+
# Given the lock we take out in `asend`, this _shouldn't_ be possible, but I'd rather fail with
727+
# this explicit error return the wrong type of message back to a Trigger
728+
raise RuntimeError(f"Response read out of order! Got {frame.id=}, {expect_id=}")
729+
return self._from_frame(frame)
730+
731+
async def asend(self, msg: ToTriggerSupervisor) -> ToTriggerRunner | None:
732+
frame = _RequestFrame(id=next(self.id_counter), body=msg.model_dump())
733+
bytes = frame.as_bytes()
734+
735+
async with self._lock:
736+
self._async_writer.write(bytes)
737+
738+
return await self._aget_response(frame.id)
739+
740+
694741
class TriggerRunner:
695742
"""
696743
Runtime environment for all triggers.
697744
698-
Mainly runs inside its own thread, where it hands control off to an asyncio
699-
event loop, but is also sometimes interacted with from the main thread
700-
(where all the DB queries are done). All communication between threads is
701-
done via Deques.
745+
Mainly runs inside its own process, where it hands control off to an asyncio
746+
event loop. All communication between this and it's (sync) supervisor is done via sockets
702747
"""
703748

704749
# Maps trigger IDs to their running tasks and other info
@@ -726,10 +771,7 @@ class TriggerRunner:
726771
# TODO: connect this to the parent process
727772
log: FilteringBoundLogger = structlog.get_logger()
728773

729-
requests_sock: asyncio.StreamWriter
730-
response_sock: asyncio.StreamReader
731-
732-
decoder: TypeAdapter[ToTriggerRunner]
774+
comms_decoder: TriggerCommsDecoder
733775

734776
def __init__(self):
735777
super().__init__()
@@ -740,7 +782,6 @@ def __init__(self):
740782
self.events = deque()
741783
self.failed_triggers = deque()
742784
self.job_id = None
743-
self.decoder = TypeAdapter(ToTriggerRunner)
744785

745786
def run(self):
746787
"""Sync entrypoint - just run a run in an async loop."""
@@ -796,36 +837,21 @@ async def init_comms(self):
796837
"""
797838
from airflow.sdk.execution_time import task_runner
798839

799-
loop = asyncio.get_event_loop()
840+
# Yes, we read and write to stdin! It's a socket, not a normal stdin.
841+
reader, writer = await asyncio.open_connection(sock=socket(fileno=0))
800842

801-
comms_decoder = task_runner.CommsDecoder[ToTriggerRunner, ToTriggerSupervisor](
802-
input=sys.stdin,
803-
decoder=self.decoder,
843+
self.comms_decoder = TriggerCommsDecoder(
844+
async_writer=writer,
845+
async_reader=reader,
804846
)
805847

806-
task_runner.SUPERVISOR_COMMS = comms_decoder
807-
808-
async def connect_stdin() -> asyncio.StreamReader:
809-
reader = asyncio.StreamReader()
810-
protocol = asyncio.StreamReaderProtocol(reader)
811-
await loop.connect_read_pipe(lambda: protocol, sys.stdin)
812-
return reader
813-
814-
self.response_sock = await connect_stdin()
848+
task_runner.SUPERVISOR_COMMS = self.comms_decoder
815849

816-
line = await self.response_sock.readline()
850+
msg = await self.comms_decoder._aget_response(expect_id=0)
817851

818-
msg = self.decoder.validate_json(line)
819852
if not isinstance(msg, messages.StartTriggerer):
820853
raise RuntimeError(f"Required first message to be a messages.StartTriggerer, it was {msg}")
821854

822-
comms_decoder.request_socket = os.fdopen(msg.requests_fd, "wb", buffering=0)
823-
writer_transport, writer_protocol = await loop.connect_write_pipe(
824-
lambda: asyncio.streams.FlowControlMixin(loop=loop),
825-
comms_decoder.request_socket,
826-
)
827-
self.requests_sock = asyncio.streams.StreamWriter(writer_transport, writer_protocol, None, loop)
828-
829855
async def create_triggers(self):
830856
"""Drain the to_create queue and create all new triggers that have been requested in the DB."""
831857
while self.to_create:
@@ -934,8 +960,6 @@ async def cleanup_finished_triggers(self) -> list[int]:
934960
return finished_ids
935961

936962
async def sync_state_to_supervisor(self, finished_ids: list[int]):
937-
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
938-
939963
# Copy out of our deques in threadsafe manner to sync state with parent
940964
events_to_send = []
941965
while self.events:
@@ -961,19 +985,17 @@ async def sync_state_to_supervisor(self, finished_ids: list[int]):
961985
if not finished_ids:
962986
msg.finished = None
963987

964-
# Block triggers from making any requests for the duration of this
965-
async with SUPERVISOR_COMMS.lock:
966-
# Tell the monitor that we've finished triggers so it can update things
967-
self.requests_sock.write(msg.model_dump_json(exclude_none=True).encode() + b"\n")
968-
line = await self.response_sock.readline()
969-
970-
if line == b"": # EoF received!
988+
# Tell the monitor that we've finished triggers so it can update things
989+
try:
990+
resp = await self.comms_decoder.asend(msg)
991+
except asyncio.IncompleteReadError:
971992
if task := asyncio.current_task():
972993
task.cancel("EOF - shutting down")
994+
return
995+
raise
973996

974-
resp = self.decoder.validate_json(line)
975997
if not isinstance(resp, messages.TriggerStateSync):
976-
raise RuntimeError(f"Expected to get a TriggerStateSync message, instead we got f{type(msg)}")
998+
raise RuntimeError(f"Expected to get a TriggerStateSync message, instead we got {type(msg)}")
977999
self.to_create.extend(resp.to_create)
9781000
self.to_cancel.extend(resp.to_cancel)
9791001

0 commit comments

Comments
 (0)