From 0a6a3ade21f94d6cf0a7c3c8b0eaa31006b9dd2b Mon Sep 17 00:00:00 2001 From: henry3260 Date: Wed, 5 Aug 2026 23:56:17 +0800 Subject: [PATCH] Fix dag.test() hanging when a task reads a Variable in a subprocess MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Under dag.test() the Execution API server and the task share one process, so the API's own route handlers can see SUPERVISOR_COMMS. They took that as proof they were client-side and issued a fresh Execution API request, calling back into themselves instead of reading the metastore. Operators that spawn their own child process — PythonVirtualenvOperator and ExternalPythonOperator — have that child reconnect to the supervisor over a socket. The in-process supervisor answered every request on an in-memory queue the child cannot read, so it stayed blocked on a response frame that was never written: no error, no timeout, no traceback. --- airflow-core/docs/core-concepts/overview.rst | 7 +- .../airflow/api_fastapi/execution_api/app.py | 24 +++- airflow-core/src/airflow/models/connection.py | 6 +- airflow-core/src/airflow/models/variable.py | 10 +- airflow-core/src/airflow/process_context.py | 57 ++++++++ .../api_fastapi/execution_api/test_app.py | 33 +++++ .../versions/head/test_connections.py | 34 +++++ .../versions/head/test_variables.py | 63 +++++++++ .../tests/unit/models/test_connection.py | 21 +++ .../airflow/sdk/execution_time/supervisor.py | 44 ++++-- .../execution_time/test_supervisor.py | 125 ++++++++++++++++++ 11 files changed, 398 insertions(+), 26 deletions(-) create mode 100644 airflow-core/src/airflow/process_context.py diff --git a/airflow-core/docs/core-concepts/overview.rst b/airflow-core/docs/core-concepts/overview.rst index 1f2b1ad8a11bf..fc53fa02b3c79 100644 --- a/airflow-core/docs/core-concepts/overview.rst +++ b/airflow-core/docs/core-concepts/overview.rst @@ -212,9 +212,10 @@ The two processes talk over a socket, and the Supervisor is the only side that e task JWT or talks to the *Execution API* — the user's code never sees the token and never touches the database. -The same runtime can also run *in-process* (a single Python process, no fork, no sockets, no HTTP) for -``dag.test()`` and local runs. The diagram below contrasts the two paths and marks where each Python process -lives: +The same runtime can also run *in-process* (a single Python process, no fork, no HTTP) for +``dag.test()`` and local runs. A supervisor socket is still set up, because operators such as +``PythonVirtualenvOperator`` spawn their own child process that has to reconnect to ask for Connections +and Variables. The diagram below contrasts the two paths and marks where each Python process lives: .. image:: ../img/diagram_task_sdk_execution_architecture.png diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index f9e7cb1725000..27bd70b2be4d2 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -44,9 +44,11 @@ get_sig_validation_args, get_signing_args, ) +from airflow.process_context import override_process_context if TYPE_CHECKING: import httpx + from starlette.types import Receive, Scope, Send import structlog from structlog.contextvars import bind_contextvars @@ -372,6 +374,17 @@ def _shutdown_loop( thread.join(timeout=5) +class _RequestScopedServerContextApp: + """Wrap an ASGI app so in-process requests behave like server-side API handling.""" + + def __init__(self, app: FastAPI) -> None: + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + with override_process_context("server"): + await self.app(scope, receive, send) + + @attrs.define() class InProcessExecutionAPI: """ @@ -379,6 +392,10 @@ class InProcessExecutionAPI: The sync version of this makes use of a2wsgi which runs the async loop in a separate thread. This is needed so that we can use the sync httpx client + + Requests are always dispatched in a server process context, with no way to opt out: this app *is* + the server side of the Execution API, and in a single process its route handlers would otherwise + read the caller's ``SUPERVISOR_COMMS`` and re-enter the API through the Task SDK. """ _app: FastAPI | None = None @@ -433,7 +450,8 @@ def transport(self) -> httpx.WSGITransport: thread = threading.Thread(target=loop.run_forever, name="InProcessExecutionAPI-loop", daemon=True) thread.start() - middleware = ASGIMiddleware(self.app, loop=loop) + app = self.app + middleware = ASGIMiddleware(cast("Any", _RequestScopedServerContextApp(app)), loop=loop) # https://github.com/abersheeran/a2wsgi/discussions/64 async def start_lifespan(cm: AsyncExitStack, app: FastAPI): @@ -443,7 +461,7 @@ async def start_lifespan(cm: AsyncExitStack, app: FastAPI): # Wait for lifespan startup to complete so callers see a ready app and so the finalizer can # safely aclose() a context whose __aenter__ has actually run. - asyncio.run_coroutine_threadsafe(start_lifespan(cm, self.app), loop).result() + asyncio.run_coroutine_threadsafe(start_lifespan(cm, app), loop).result() transport = httpx.WSGITransport(app=middleware) # type: ignore[arg-type] @@ -460,4 +478,4 @@ async def start_lifespan(cm: AsyncExitStack, app: FastAPI): def atransport(self) -> httpx.ASGITransport: import httpx - return httpx.ASGITransport(app=self.app) + return httpx.ASGITransport(app=_RequestScopedServerContextApp(self.app)) diff --git a/airflow-core/src/airflow/models/connection.py b/airflow-core/src/airflow/models/connection.py index 1b4b0f8f86768..b1ec6beaee562 100644 --- a/airflow-core/src/airflow/models/connection.py +++ b/airflow-core/src/airflow/models/connection.py @@ -20,7 +20,6 @@ import json import logging import re -import sys import warnings from contextlib import suppress from json import JSONDecodeError @@ -50,6 +49,7 @@ class AirflowSecretsBackendAccessDenied(PermissionError): # type: ignore[no-red """Compat stub — never raised by task-sdk <1.2.2.""" +from airflow.process_context import should_use_task_sdk_api_path from airflow.utils.helpers import prune_dict from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.session import NEW_SESSION, provide_session @@ -475,7 +475,7 @@ def get_connection_from_secrets(cls, conn_id: str, team_name: str | None = None) # If this is set it means are in some kind of execution context (Task, Dag Parse or Triggerer perhaps) # and should use the Task SDK API server path - if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"): + if should_use_task_sdk_api_path(): from airflow.sdk import Connection as TaskSDKConnection from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType @@ -566,7 +566,7 @@ def to_dict(self, *, prune_empty: bool = False, validate: bool = True) -> dict[s @classmethod def from_json(cls, value, conn_id=None) -> Connection: - if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"): + if should_use_task_sdk_api_path(): from airflow.sdk import Connection as TaskSDKConnection warnings.warn( diff --git a/airflow-core/src/airflow/models/variable.py b/airflow-core/src/airflow/models/variable.py index b06e73cd5f50a..f29bdd55fc7eb 100644 --- a/airflow-core/src/airflow/models/variable.py +++ b/airflow-core/src/airflow/models/variable.py @@ -20,7 +20,6 @@ import contextlib import json import logging -import sys import warnings from typing import TYPE_CHECKING, Any @@ -47,6 +46,7 @@ class AirflowSecretsBackendAccessDenied(PermissionError): # type: ignore[no-red """Compat stub — never raised by task-sdk <1.2.2.""" +from airflow.process_context import should_use_task_sdk_api_path from airflow.secrets.metastore import MetastoreBackend from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.session import NEW_SESSION, create_session, provide_session @@ -166,7 +166,7 @@ def get( # If this is set it means we are in some kind of execution context (Task, Dag Parse or Triggerer perhaps) # and should use the Task SDK API server path - if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"): + if should_use_task_sdk_api_path(): warnings.warn( "Using Variable.get from `airflow.models` is deprecated." "Please use `get` on Variable from sdk(`airflow.sdk.Variable`) instead", @@ -226,7 +226,7 @@ def set( # If this is set it means we are in some kind of execution context (Task, Dag Parse or Triggerer perhaps) # and should use the Task SDK API server path - if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"): + if should_use_task_sdk_api_path(): warnings.warn( "Using Variable.set from `airflow.models` is deprecated." "Please use `set` on Variable from sdk(`airflow.sdk.Variable`) instead", @@ -314,7 +314,7 @@ def update( # If this is set it means are in some kind of execution context (Task, Dag Parse or Triggerer perhaps) # and should use the Task SDK API server path - if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"): + if should_use_task_sdk_api_path(): warnings.warn( "Using Variable.update from `airflow.models` is deprecated." "Please use `set` on Variable from sdk(`airflow.sdk.Variable`) instead as it is an upsert.", @@ -380,7 +380,7 @@ def delete(key: str, team_name: str | None = None, session: Session | None = Non # If this is set it means are in some kind of execution context (Task, Dag Parse or Triggerer perhaps) # and should use the Task SDK API server path - if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"): + if should_use_task_sdk_api_path(): warnings.warn( "Using Variable.delete from `airflow.models` is deprecated." "Please use `delete` on Variable from sdk(`airflow.sdk.Variable`) instead", diff --git a/airflow-core/src/airflow/process_context.py b/airflow-core/src/airflow/process_context.py new file mode 100644 index 0000000000000..7c97948977b04 --- /dev/null +++ b/airflow-core/src/airflow/process_context.py @@ -0,0 +1,57 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import sys +from collections.abc import Generator +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Literal + +__all__ = [ + "override_process_context", + "should_use_task_sdk_api_path", +] + +_PROCESS_CONTEXT_OVERRIDE: ContextVar[str | None] = ContextVar( + "_AIRFLOW_PROCESS_CONTEXT_OVERRIDE", + default=None, +) + + +@contextmanager +def override_process_context(context: Literal["server", "client"]) -> Generator[None, None, None]: + """Temporarily override the current process context for the active execution flow.""" + token = _PROCESS_CONTEXT_OVERRIDE.set(context) + try: + yield + finally: + _PROCESS_CONTEXT_OVERRIDE.reset(token) + + +def should_use_task_sdk_api_path() -> bool: + """Return True when execution-context helpers should route through Task SDK APIs.""" + # Only the ContextVar, never the ``_AIRFLOW_PROCESS_CONTEXT`` env var: that env var is + # process-wide and inherited by children (``action_cli`` sets it around the whole + # ``airflow dags test`` body, and PythonVirtualenvOperator passes it to the venv child), so + # letting it win here would send worker-side code straight to the metastore. ``SUPERVISOR_COMMS`` + # keeps precedence over it, matching ``ensure_secrets_backend_loaded()`` in the Task SDK. + if _PROCESS_CONTEXT_OVERRIDE.get() == "server": + return False + + task_runner_module = sys.modules.get("airflow.sdk.execution_time.task_runner") + return bool(getattr(task_runner_module, "SUPERVISOR_COMMS", None)) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py index bb2d2d557dc36..6cb865448a02d 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py @@ -163,6 +163,39 @@ def test_in_process_execution_api_runs_without_jwt_secret(): assert response.status_code == 200 +def test_in_process_execution_api_does_not_reenter_task_sdk(session): + """A request served in-process must not route back out through the Task SDK. + + Under ``dag.test()`` the API server and the task share one process, so ``SUPERVISOR_COMMS`` is + visible to the route handler too. Without a request-scoped server context the handler reads it, + decides it is client-side, and issues another Execution API request -- looping until the caller + is killed. No env var is set here on purpose: ``dag.test()`` from a script sets none, so the + ContextVar is the only thing preventing re-entry. + """ + from airflow.models.variable import Variable + + # Patched below, not by decorator: this setup must run before ``SUPERVISOR_COMMS`` is visible. + Variable.set(key="inproc_key", value="VALUE", session=session) + session.commit() + + api = InProcessExecutionAPI() + with ( + mock.patch.dict( + "sys.modules", + {"airflow.sdk.execution_time.task_runner": mock.Mock(spec=["SUPERVISOR_COMMS"])}, + ), + mock.patch( + "airflow.sdk.Variable.get", + side_effect=AssertionError("in-process Execution API re-entered the Task SDK path"), + ), + httpx.Client(transport=api.transport) as client, + ): + response = client.get("http://localhost/variables/inproc_key") + + assert response.status_code == 200 + assert response.json() == {"key": "inproc_key", "value": "VALUE"} + + def test_in_process_execution_api_transport_lifecycle(): """The background loop + thread lifecycle is tied to the ``.transport``, not the factory instance. diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connections.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connections.py index a2e3cb51fab32..51fe3959d33f3 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connections.py @@ -17,12 +17,14 @@ from __future__ import annotations +import sys from unittest import mock import pytest from fastapi import FastAPI, HTTPException, status from airflow.models.connection import Connection +from airflow.process_context import override_process_context pytestmark = pytest.mark.db_test @@ -105,6 +107,38 @@ def test_connection_get_from_env_var(self, client, session): "extra": '{"headers": "header"}', } + @mock.patch.dict( + "os.environ", + { + "AIRFLOW_CONN_TEST_CONN_SERVER": '{"uri": "http://root:admin@localhost:8080/https?headers=header"}', + }, + ) + @mock.patch.dict( + sys.modules, + {"airflow.sdk.execution_time.task_runner": mock.Mock(spec=["SUPERVISOR_COMMS"])}, + ) + @mock.patch( + "airflow.sdk.Connection.get", + side_effect=AssertionError( + "Execution API should not route through Task SDK Connection.get in server context" + ), + ) + def test_connection_get_uses_server_path_when_supervisor_comms_exists(self, mock_sdk_get, client): + with override_process_context("server"): + response = client.get("/execution/connections/test_conn_server") + + assert response.status_code == 200 + assert response.json() == { + "conn_id": "test_conn_server", + "conn_type": "http", + "host": "localhost", + "login": "root", + "password": "admin", + "schema": "https", + "port": 8080, + "extra": '{"headers": "header"}', + } + def test_connection_get_not_found(self, client): response = client.get("/execution/connections/non_existent_test_conn") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_variables.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_variables.py index f078d6c2fe06f..698d324bc6ced 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_variables.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_variables.py @@ -26,6 +26,7 @@ from sqlalchemy import select from airflow.models.variable import Variable +from airflow.process_context import override_process_context from tests_common.test_utils.db import clear_db_variables @@ -96,6 +97,24 @@ def test_variable_get_from_env_var(self, client, session): assert response.status_code == 200 assert response.json() == {"key": "key1", "value": "VALUE"} + @mock.patch.dict("os.environ", {"AIRFLOW_VAR_KEY1": "VALUE"}) + @mock.patch.dict( + "sys.modules", + {"airflow.sdk.execution_time.task_runner": mock.Mock(spec=["SUPERVISOR_COMMS"])}, + ) + @mock.patch( + "airflow.sdk.Variable.get", + side_effect=AssertionError( + "Execution API should not route through Task SDK Variable.get in server context" + ), + ) + def test_variable_get_uses_server_path_when_supervisor_comms_exists(self, mock_sdk_get, client): + with override_process_context("server"): + response = client.get("/execution/variables/key1") + + assert response.status_code == 200 + assert response.json() == {"key": "key1", "value": "VALUE"} + @pytest.mark.parametrize( "key", [ @@ -158,6 +177,26 @@ def test_should_create_variable(self, client, key, payload, session): if "description" in payload: assert var_from_db.description == payload["description"] + @mock.patch.dict( + "sys.modules", + {"airflow.sdk.execution_time.task_runner": mock.Mock(spec=["SUPERVISOR_COMMS"])}, + ) + @mock.patch( + "airflow.sdk.Variable.set", + side_effect=AssertionError( + "Execution API should not route through Task SDK Variable.set in server context" + ), + ) + def test_variable_put_uses_server_path_when_supervisor_comms_exists(self, mock_sdk_set, client, session): + with override_process_context("server"): + response = client.put("/execution/variables/var_server_only", json={"value": "server_value"}) + + assert response.status_code == 201 + assert response.json()["message"] == "Variable successfully set" + var_from_db = session.scalars(select(Variable).where(Variable.key == "var_server_only")).first() + assert var_from_db is not None + assert var_from_db.val == "server_value" + @pytest.mark.parametrize( ("key", "payload", "error_type"), [ @@ -342,3 +381,27 @@ def test_should_not_delete_variable(self, client, session): vars = session.scalars(select(Variable)).all() assert len(vars) == 1 + + def test_variable_delete_uses_server_path_when_supervisor_comms_exists(self, client, session): + # Patched below, not by decorator: this setup must run before ``SUPERVISOR_COMMS`` is visible. + Variable.set(key="var_server_delete", value="to_delete", session=session) + session.commit() + + with ( + override_process_context("server"), + mock.patch.dict( + "sys.modules", + {"airflow.sdk.execution_time.task_runner": mock.Mock(spec=["SUPERVISOR_COMMS"])}, + ), + mock.patch( + "airflow.sdk.Variable.delete", + side_effect=AssertionError( + "Execution API should not route through Task SDK Variable.delete in server context" + ), + ), + ): + response = client.delete("/execution/variables/var_server_delete") + + assert response.status_code == 204 + session.expire_all() + assert session.scalar(select(Variable).where(Variable.key == "var_server_delete")) is None diff --git a/airflow-core/tests/unit/models/test_connection.py b/airflow-core/tests/unit/models/test_connection.py index 94cabe5e4daf4..ffd2af1785ef9 100644 --- a/airflow-core/tests/unit/models/test_connection.py +++ b/airflow-core/tests/unit/models/test_connection.py @@ -27,6 +27,7 @@ from airflow.exceptions import AirflowException, AirflowNotFoundException from airflow.models import Connection +from airflow.process_context import override_process_context from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType from airflow.sdk.execution_time.comms import ErrorResponse @@ -455,6 +456,26 @@ def test_get_connection_from_secrets_task_sdk_not_found(self, mock_task_sdk_conn with pytest.raises(AirflowNotFoundException): Connection.get_connection_from_secrets("test_conn") + @mock.patch.dict( + sys.modules, + {"airflow.sdk.execution_time.task_runner": mock.Mock(spec=["SUPERVISOR_COMMS"])}, + ) + @mock.patch( + "airflow.sdk.Connection.from_json", + side_effect=AssertionError( + "Connection.from_json should not route through Task SDK in server context" + ), + ) + def test_connection_from_json_uses_core_path_when_server_context(self, mock_sdk_from_json): + """Server context should prefer core Connection.from_json even if comms exist.""" + with override_process_context("server"): + result = Connection.from_json('{"conn_type": "http", "host": "localhost"}', conn_id="test_conn") + + assert isinstance(result, Connection) + assert result.conn_id == "test_conn" + assert result.conn_type == "http" + assert result.host == "localhost" + @mock.patch.dict(sys.modules, {"airflow.sdk.execution_time.task_runner": None}) @mock.patch("airflow.sdk.Connection") @mock.patch("airflow.secrets.environment_variables.EnvironmentVariablesBackend.get_connection") diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 87311f02da7a1..99ee75db99e7a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -35,6 +35,7 @@ from collections import deque from collections.abc import Callable, Generator from contextlib import contextmanager, suppress +from contextvars import ContextVar from datetime import datetime, timezone from http import HTTPStatus from socket import socket, socketpair @@ -1951,26 +1952,38 @@ def in_process_api_server(): return api +_IN_PROCESS_RESPONSE_SINK: ContextVar[deque[BaseModel | None] | None] = ContextVar( + "in_process_response_sink", default=None +) +"""Where :meth:`InProcessTestSupervisor.send_msg` must deliver the response it is about to send. + +Only :meth:`InProcessSupervisorComms.send` sets it, and it gets a fresh sink per call, so a response +can never reach a caller other than the one that is waiting for it. The socket is read by the raw +thread started in ``_setup_subprocess_socket``, which never has a sink set: requests from a child +process are answered on the socket instead. +""" + + @attrs.define(kw_only=True) class InProcessSupervisorComms: """In-process communication handler that uses deques instead of sockets.""" log: FilteringBoundLogger = attrs.field(repr=False, factory=structlog.get_logger) supervisor: InProcessTestSupervisor - messages: deque[BaseModel | None] = attrs.field(factory=deque) - def _get_response(self) -> BaseModel | None: - """Get a message from the supervisor. Blocks until a message is available.""" - return self.messages.popleft() - - def send(self, msg: BaseModel): - """Send a request to the supervisor.""" + def send(self, msg: BaseModel) -> BaseModel | None: + """Send a request to the supervisor and return its response.""" self.log.debug("Sending request", msg=msg) - with set_supervisor_comms(None): - self.supervisor._handle_request(msg, log, 0) # type: ignore[arg-type] + responses: deque[BaseModel | None] = deque() + token = _IN_PROCESS_RESPONSE_SINK.set(responses) + try: + with set_supervisor_comms(None): + self.supervisor._handle_request(msg, log, 0) # type: ignore[arg-type] + finally: + _IN_PROCESS_RESPONSE_SINK.reset(token) - return self._get_response() + return responses.popleft() if responses else None @attrs.define @@ -2133,8 +2146,15 @@ def _api_client(dag=None): def send_msg( self, msg: BaseModel | None, request_id: int, error: ErrorResponse | None = None, **dump_opts ): - """Override to use in-process comms.""" - self.comms.messages.append(msg) + """Deliver the response in-process, or over the socket when the request came from there.""" + sink = _IN_PROCESS_RESPONSE_SINK.get() + if sink is None: + # A real child process (a virtualenv operator's, say) sent this request over the socket + # set up by `_setup_subprocess_socket`. It is blocked reading a response frame, so + # anything queued in-process would never reach it. + super().send_msg(msg, request_id, error=error, **dump_opts) + return + sink.append(msg) @classmethod def run_trigger_in_process(cls, *, trigger, ti): diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index f777b2d5a8a90..14091f0c46c81 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -27,6 +27,7 @@ import socket import subprocess import sys +import threading import time from contextlib import nullcontext from dataclasses import dataclass, field @@ -3628,6 +3629,130 @@ def _handle_request(self, msg, log, req_id): assert isinstance(response, VariableResult) assert response.value == "value" + @pytest.fixture + def socket_supervisor(self, mocker, socket_pair): + """An in-process supervisor wired to a socket, as ``_setup_subprocess_socket`` leaves it.""" + read_end, write_end = socket_pair + + supervisor = InProcessTestSupervisor( + id=TI_ID, + pid=12345, + process=mocker.Mock(), + process_log=mocker.MagicMock(), + client=mocker.MagicMock(spec=sdk_client.Client), + ) + supervisor.comms = InProcessSupervisorComms(supervisor=supervisor) + supervisor.stdin = write_end + supervisor.client.variables.get.return_value = VariableResult(key="test_key", value="test_value") + + return supervisor, read_end + + @patch("airflow.sdk.execution_time.request_handlers.mask_secret") + @pytest.mark.parametrize("req_id", [0, 42], ids=["first_request", "later_request"]) + def test_socket_request_is_answered_over_the_socket( + self, mock_mask_secret, socket_supervisor, mocker, req_id + ): + """A virtualenv operator under ``dag.test()`` runs in a real child process that reconnects to + the supervisor over ``__AIRFLOW_SUPERVISOR_FD``, so its requests can only be answered with a + response frame on that socket. ``req_id=0`` is deliberate: the child's ``CommsDecoder`` + numbers its requests from 0, so the id cannot tell the two paths apart. + """ + supervisor, read_end = socket_supervisor + + generator = supervisor.handle_requests(log=mocker.Mock()) + next(generator) + generator.send(_RequestFrame(id=req_id, body=GetVariable(key="test_key").model_dump())) + + read_end.settimeout(1) + frame_len = int.from_bytes(read_end.recv(4), "big") + frame = msgspec.msgpack.Decoder(_ResponseFrame).decode(read_end.recv(frame_len)) + + assert frame.id == req_id + assert frame.body == {"key": "test_key", "value": "test_value", "type": "VariableResult"} + + @patch("airflow.sdk.execution_time.request_handlers.mask_secret") + def test_in_process_request_is_not_written_to_the_socket(self, mock_mask_secret, socket_supervisor): + """The task running in this process reads its response from the queue, not the socket.""" + supervisor, read_end = socket_supervisor + + response = supervisor.comms.send(GetVariable(key="test_key")) + + assert response == VariableResult(key="test_key", value="test_value") + read_end.settimeout(0.1) + with pytest.raises(TimeoutError): + read_end.recv(1) + + def test_concurrent_in_process_requests_get_their_own_response(self, mocker): + """Requests in flight at the same time must not be answered with each other's response. + + The socket is serviced on its own thread, so the supervisor answers child-process requests + while the task in this process has a request of its own outstanding. + """ + first_queued = threading.Event() + second_answered = threading.Event() + + class ConcurrentSupervisor(InProcessTestSupervisor): + def _handle_request(self, msg, log, req_id): + self.send_msg(VariableResult(key=msg.key, value=msg.key), req_id) + if msg.key == "first": + first_queued.set() + second_answered.wait(10) + + supervisor = ConcurrentSupervisor( + id=TI_ID, + pid=12345, + process=mocker.Mock(), + process_log=mocker.MagicMock(), + client=mocker.MagicMock(spec=sdk_client.Client), + ) + supervisor.comms = InProcessSupervisorComms(supervisor=supervisor) + + answers: dict[str, Any] = {} + + def ask(key): + answers[key] = supervisor.comms.send(GetVariable(key=key)) + + first = threading.Thread(target=ask, args=("first",), daemon=True) + first.start() + assert first_queued.wait(10) + + ask("second") + second_answered.set() + first.join(10) + + assert answers == { + "first": VariableResult(key="first", value="first"), + "second": VariableResult(key="second", value="second"), + } + + @patch("airflow.sdk.execution_time.request_handlers.mask_secret") + def test_child_process_gets_a_response_through_the_real_socket(self, mock_mask_secret, mocker): + """The same round trip, driven through the socket machinery a real child process uses.""" + supervisor = InProcessTestSupervisor( + id=TI_ID, + pid=12345, + process=mocker.Mock(), + process_log=mocker.MagicMock(), + client=mocker.MagicMock(spec=sdk_client.Client), + ) + supervisor.comms = InProcessSupervisorComms(supervisor=supervisor) + supervisor.client.variables.get.return_value = VariableResult(key="test_key", value="test_value") + + received: list[Any] = [] + with supervisor._setup_subprocess_socket() as child_sock: + comms = CommsDecoder(socket=child_sock) + # `CommsDecoder._read_frame` forces the socket back to blocking mode, so a lost response + # hangs forever -- read it from a thread we can give up on to get a failure instead. + reader = threading.Thread( + target=lambda: received.append(comms.send(GetVariable(key="test_key"))), daemon=True + ) + reader.start() + reader.join(10) + + assert not reader.is_alive(), "supervisor never answered the request sent on its socket" + + assert received == [VariableResult(key="test_key", value="test_value")] + def test_inprocess_failure_callback_receives_exception( self, monkeypatch,