diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index 83bb2e32358bf..01f16537f59f0 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -18,6 +18,7 @@ import asyncio import functools +import io import logging import os import selectors @@ -182,6 +183,7 @@ class StartTriggerer(BaseModel): """Tell the async trigger runner process to start, and where to send status update messages.""" requests_fd: int + trigger_requests_fd: int type: Literal["StartTriggerer"] = "StartTriggerer" class CancelTriggers(BaseModel): @@ -322,8 +324,11 @@ def start( # type: ignore[override] ): proc = super().start(id=job.id, job=job, target=cls.run_in_process, logger=logger, **kwargs) - msg = messages.StartTriggerer(requests_fd=proc._requests_fd) - proc._send(msg) + msg = messages.StartTriggerer( + requests_fd=proc._requests_fd, + trigger_requests_fd=proc._trigger_requests_fd, # type: ignore[arg-type] + ) + proc._send_first_message(msg) return proc @functools.cached_property @@ -463,6 +468,9 @@ def emit_metrics(self): ) def _send(self, msg: BaseModel): + self.trigger_stdin.write(msg.model_dump_json().encode("utf-8") + b"\n") # type: ignore[union-attr] + + def _send_first_message(self, msg: BaseModel): self.stdin.write(msg.model_dump_json().encode("utf-8") + b"\n") def update_triggers(self, requested_trigger_ids: set[int]): @@ -678,8 +686,12 @@ def init_comms(self): msg = comms_decoder.get_message() if not isinstance(msg, messages.StartTriggerer): raise RuntimeError(f"Required first message to be a messages.StartTriggerer, it was {msg}") + comms_decoder.request_socket = os.fdopen(msg.requests_fd, "wb", buffering=0) + binary = os.fdopen(msg.trigger_requests_fd, "rb") + handle = io.TextIOWrapper(binary, line_buffering=True) + comms_decoder.trigger_input = handle task_runner.SUPERVISOR_COMMS = comms_decoder def run(self): @@ -750,7 +762,7 @@ async def read_workloads(self, ready_event: asyncio.Event): async def connect_stdin() -> asyncio.StreamReader: reader = asyncio.StreamReader() protocol = asyncio.StreamReaderProtocol(reader) - await loop.connect_read_pipe(lambda: protocol, sys.stdin) + await loop.connect_read_pipe(lambda: protocol, task_runner.SUPERVISOR_COMMS.trigger_input) return reader stdin = await connect_stdin() diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index e53b893178df5..46aa6c810dbdc 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -172,8 +172,10 @@ def builder(job=None): job=job, pid=process.pid, stdin=mocker.Mock(), + trigger_stdin=mocker.Mock(), process=process, requests_fd=-1, + trigger_requests_fd=-1, capacity=10, ) # Mock the selector @@ -447,7 +449,7 @@ def is_alive(self): return True supervisor.run() - assert supervisor.stdin.write.call_count == 1 + assert supervisor.trigger_stdin.write.call_count == 1 @pytest.mark.execution_timeout(5) @@ -624,10 +626,10 @@ async def run(self, **args) -> AsyncIterator[TriggerEvent]: from airflow.sdk import Variable from airflow.sdk.execution_time.xcom import XCom - conn = await sync_to_async(BaseHook.get_connection)("test_connection") + conn = await sync_to_async(BaseHook.get_connection)(f"test_connection_{self.task_id}") self.log.info("Loaded conn %s", conn.conn_id) - variable = await sync_to_async(Variable.get)("test_variable") + variable = await sync_to_async(Variable.get)(f"test_variable_{self.task_id}") self.log.info("Loaded variable %s", variable) xcom = await sync_to_async(XCom.get_one)( @@ -693,8 +695,8 @@ async def test_trigger_can_access_variables_connections_and_xcoms(session, dag_m task_instance.trigger_id = trigger_orm.id # Create the appropriate Connection, Variable and XCom - connection = Connection(conn_id="test_connection", conn_type="http") - variable = Variable(key="test_variable", val="some_variable_value") + connection = Connection(conn_id=f"test_connection_{task_instance.task_id}", conn_type="http") + variable = Variable(key=f"test_variable_{task_instance.task_id}", val="some_variable_value") XComModel.set( key="test_xcom", value="some_xcom_value", @@ -720,7 +722,7 @@ async def test_trigger_can_access_variables_connections_and_xcoms(session, dag_m assert task_instance.next_kwargs == { "event": { "connection": { - "conn_id": "test_connection", + "conn_id": f"test_connection_{task_instance.task_id}", "conn_type": "http", "description": None, "host": None, @@ -734,3 +736,80 @@ async def test_trigger_can_access_variables_connections_and_xcoms(session, dag_m "xcom": '"some_xcom_value"', } } + + +@pytest.mark.flaky(reruns=2, reruns_delay=10) +@pytest.mark.execution_timeout(30) +@pytest.mark.asyncio +def test_trigger_multiple_task_instances_correctly_gets_messages_running_in_trigger(session, dag_maker): + """ + Check that the trigger will successfully access Variables, Connections and XComs. + + Running multiple tasks in triggerer to get the correct messages. + Create 3 connections, variables, and XComs for 3 tasks. the key will have task_id to + assert the correct message received + """ + + with dag_maker(dag_id="trigger_accessing_variable_connection_and_xcom", session=session): + EmptyOperator(task_id="task_one") + EmptyOperator(task_id="task_two") + EmptyOperator(task_id="task_three") + dr = dag_maker.create_dagrun() + + for trigger_id, ti in enumerate(dr.task_instances, start=1): + ti.state = TaskInstanceState.DEFERRED + + # Create a Trigger + trigger = CustomTrigger(dag_id=dr.dag_id, run_id=dr.run_id, task_id=ti.task_id, map_index=-1) + trigger_orm = Trigger( + classpath=trigger.serialize()[0], + kwargs={"dag_id": dr.dag_id, "run_id": dr.run_id, "task_id": ti.task_id, "map_index": -1}, + ) + trigger_orm.id = trigger_id + session.add(trigger_orm) + session.commit() + ti.trigger_id = trigger_orm.id + + # Create the appropriate Connection, Variable and XCom + connection = Connection(conn_id=f"test_connection_{ti.task_id}", conn_type="http") + variable = Variable(key=f"test_variable_{ti.task_id}", val=f"some_variable_value {trigger_id}") + XComModel.set( + key="test_xcom", + value=f"some_xcom_value {trigger_id}", + task_id=ti.task_id, + dag_id=dr.dag_id, + run_id=dr.run_id, + map_index=-1, + session=session, + ) + session.add(connection) + session.add(variable) + session.commit() + + job = Job() + session.add(job) + session.commit() + + supervisor = DummyTriggerRunnerSupervisor.start(job=job, capacity=4, logger=None) + supervisor.run() + for trigger_id, ti in enumerate(dr.task_instances, start=1): + ti.refresh_from_db() + assert ti.state == TaskInstanceState.SCHEDULED + assert ti.next_method != "__fail__" + assert ti.next_kwargs == { + "event": { + "connection": { + "conn_id": f"test_connection_{ti.task_id}", + "conn_type": "http", + "description": None, + "host": None, + "schema": None, + "login": None, + "password": None, + "port": None, + "extra": None, + }, + "variable": f"some_variable_value {trigger_id}", + "xcom": f'"some_xcom_value {trigger_id}"', + } + } diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index c924e11322cd7..a7b13bcba9454 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -385,6 +385,16 @@ class WatchedSubprocess: _requests_fd: int """File descriptor for request handling.""" + """" + It require special case to handle the workloads and api calls to api-server in Triggerer, due to mixing up messages + a separate channel is used to send the workloads from parent process to child process the child process. + connect_stdin will use this channel to read the workloads in read_workload method. + trigger_stdin: trigger parent process uses this to send the workloads to child process. + _trigger_requests_fd: trigger child process uses this to read the workloads. + """ + _trigger_requests_fd: int | None = None + trigger_stdin: BinaryIO | None = None + _num_open_sockets: int = 4 _exit_code: int | None = attrs.field(default=None, init=False) @@ -413,11 +423,19 @@ def start( child_comms, read_msgs = mkpipe() child_logs, read_logs = mkpipe() + if "TriggerRunnerSupervisor" in cls.__name__: + trigger_child_stdin, trigger_feed_stdin = mkpipe(remote_read=True) + else: + trigger_child_stdin, trigger_feed_stdin = None, None + pid = os.fork() if pid == 0: # Close and delete of the parent end of the sockets. cls._close_unused_sockets(feed_stdin, read_stdout, read_stderr, read_msgs, read_logs) + if trigger_feed_stdin: + cls._close_unused_sockets(trigger_feed_stdin) + # Python GC should delete these for us, but lets make double sure that we don't keep anything # around in the forked processes, especially things that might involve open files or sockets! del constructor_kwargs @@ -439,6 +457,15 @@ def start( requests_fd = child_comms.fileno() + if "TriggerRunnerSupervisor" in cls.__name__: + additional_trigger_fds = { + "trigger_stdin": trigger_feed_stdin, + "trigger_requests_fd": trigger_child_stdin.fileno(), # type: ignore[union-attr] + } + cls._close_unused_sockets(trigger_child_stdin) + else: + additional_trigger_fds = {} + # Close the remaining parent-end of the sockets we've passed to the child via fork. We still have the # other end of the pair open cls._close_unused_sockets(child_stdin, child_stdout, child_stderr, child_comms, child_logs) @@ -450,6 +477,7 @@ def start( process=psutil.Process(pid), requests_fd=requests_fd, process_log=logger, + **additional_trigger_fds, # type: ignore[arg-type] **constructor_kwargs, ) @@ -668,6 +696,8 @@ def _check_subprocess_exit(self, raise_on_timeout: bool = False) -> int | None: self._exit_code = self._process.wait(timeout=0) log.debug("%s process exited", type(self).__name__, exit_code=self._exit_code) self._close_unused_sockets(self.stdin) + if self.trigger_stdin: + self._close_unused_sockets(self.trigger_stdin) except psutil.TimeoutExpired: if raise_on_timeout: raise diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 04165384fef5f..830ecfb016b5f 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -540,6 +540,13 @@ class CommsDecoder(Generic[ReceiveMsgType, SendMsgType]): input: TextIO + """" + It require special case to handle the workloads and api calls to api-server, due to mixing up messages + a separate channel is used to send the workloads from parent process to child process the child process. + connect_stdin will use this channel to read the workloads in read_workload method. + """ + trigger_input: TextIO = attrs.field(init=False, default=None) + request_socket: FileIO = attrs.field(init=False, default=None) # We could be "clever" here and set the default to this based type parameters and a custom