This is an automated email from the ASF dual-hosted git repository.
henry3260 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 79a1dd12ee5 Fix dag.test() hanging when a task reads a Variable in a
subprocess (#65587)
79a1dd12ee5 is described below
commit 79a1dd12ee557614a7692120f77d165ddbd887a6
Author: Henry Chen <[email protected]>
AuthorDate: Sat Nov 7 20:06:02 2026 +0800
Fix dag.test() hanging when a task reads a Variable in a subprocess (#65587)
* Fix dag.test() hanging when a task reads a Variable in a subprocess
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.
* Keep in-process comms docs and server-context API in step with behavior
The response-sink docstring claimed the socket thread never sets a sink,
yet mask_secret re-enters the comms on that thread; a reader trusting it
could swap the sink for a thread check and misroute MaskSecret replies.
The architecture doc and diagram still described deques and no sockets,
and override_process_context accepted a "client" value that did nothing,
silently falling back to server behavior for anyone relying on it.
* Cover secret masking on the socket thread in in-process supervisor tests
Every socket test patched mask_secret out, so the one path where a child's
request re-enters the in-process comms on the socket thread was never
exercised, and a regression there would silently hand the child the wrong
frame. The concurrent test's docstring also described a socket scenario it
never ran, and a hung first request only surfaced as an opaque dict diff.
* Narrow the a2wsgi middleware cast to a2wsgi's ASGIApp
a2wsgi types ASGI scopes and messages as TypedDict unions and Starlette
as mutable mappings, so the Starlette-typed wrapper never matches the
ASGIApp that ASGIMiddleware declares. Casting to that type instead of
Any ties the call to a2wsgi's signature, so a change there shows up in
mypy. The wrapper's own signature stays checked where atransport passes
it to httpx2.ASGITransport.
---
.../api_fastapi/execution_api/in_process.py | 30 +++-
airflow-core/src/airflow/models/connection.py | 6 +-
airflow-core/src/airflow/models/variable.py | 10 +-
airflow-core/src/airflow/process_context.py | 53 +++++++
.../unit/api_fastapi/execution_api/test_app.py | 33 ++++
.../versions/head/test_connections.py | 34 +++++
.../execution_api/versions/head/test_variables.py | 63 ++++++++
airflow-core/tests/unit/models/test_connection.py | 21 +++
.../31_task_execution_architecture.rst | 7 +-
.../diagram_task_sdk_execution_architecture.md5sum | 2 +-
.../diagram_task_sdk_execution_architecture.png | Bin 184731 -> 170906 bytes
.../diagram_task_sdk_execution_architecture.py | 8 +-
.../src/airflow/sdk/execution_time/supervisor.py | 52 +++++--
.../task_sdk/execution_time/test_supervisor.py | 166 +++++++++++++++++++++
14 files changed, 452 insertions(+), 33 deletions(-)
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py
b/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py
index eb428196d43..4405e387b76 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py
@@ -23,15 +23,19 @@ import threading
import weakref
from contextlib import AsyncExitStack
from functools import cached_property
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
import attrs
import structlog
from starlette.requests import Request
+from airflow.process_context import force_server_context
+
if TYPE_CHECKING:
import httpx2
+ from a2wsgi.asgi_typing import ASGIApp
from fastapi import FastAPI
+ from starlette.types import Receive, Scope, Send
logger = structlog.get_logger(logger_name=__name__)
@@ -89,6 +93,17 @@ def _configure_in_process_auth(app: FastAPI, *,
triggerer_id: int | None = None)
app.dependency_overrides[has_xcom_access] = always_allow
+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 force_server_context():
+ await self.app(scope, receive, send)
+
+
@attrs.define()
class InProcessExecutionAPI:
"""
@@ -96,6 +111,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 httpx2 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
@@ -131,7 +150,10 @@ class InProcessExecutionAPI:
thread = threading.Thread(target=loop.run_forever,
name="InProcessExecutionAPI-loop", daemon=True)
thread.start()
- middleware = ASGIMiddleware(self.app, loop=loop)
+ app = self.app
+ # a2wsgi types ASGI scopes and messages as TypedDict unions and
Starlette as mutable mappings,
+ # so no app typed against Starlette matches a2wsgi's ``ASGIApp``.
+ middleware = ASGIMiddleware(cast("ASGIApp",
_RequestScopedServerContextApp(app)), loop=loop)
# https://github.com/abersheeran/a2wsgi/discussions/64
async def start_lifespan(cm: AsyncExitStack, app: FastAPI):
@@ -141,7 +163,7 @@ class InProcessExecutionAPI:
cm = AsyncExitStack()
# Wait for startup so callers see a ready app and the finalizer can
close an entered context.
- asyncio.run_coroutine_threadsafe(start_lifespan(cm, self.app),
loop).result()
+ asyncio.run_coroutine_threadsafe(start_lifespan(cm, app),
loop).result()
transport = httpx2.WSGITransport(app=middleware) # type:
ignore[arg-type]
@@ -155,4 +177,4 @@ class InProcessExecutionAPI:
def atransport(self) -> httpx2.ASGITransport:
import httpx2
- return httpx2.ASGITransport(app=self.app)
+ return
httpx2.ASGITransport(app=_RequestScopedServerContextApp(self.app))
diff --git a/airflow-core/src/airflow/models/connection.py
b/airflow-core/src/airflow/models/connection.py
index 52573a750e9..413902030f9 100644
--- a/airflow-core/src/airflow/models/connection.py
+++ b/airflow-core/src/airflow/models/connection.py
@@ -20,7 +20,6 @@ from __future__ import annotations
import json
import logging
import re
-import sys
import warnings
from contextlib import suppress
from json import JSONDecodeError
@@ -36,6 +35,7 @@ from airflow._shared.secrets_masker import mask_secret
from airflow.exceptions import AirflowException, AirflowNotFoundException
from airflow.models.base import ID_LEN, Base
from airflow.models.crypto import FernetFieldsMixin, get_fernet
+from airflow.process_context import should_use_task_sdk_api_path
# AirflowSecretsBackendAccessDenied was added to task-sdk in 1.2.2. When
# airflow-core is installed alongside an older published task-sdk (e.g. 1.2.1
or earlier),
@@ -535,7 +535,7 @@ class Connection(Base, FernetFieldsMixin, LoggingMixin):
# 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
@@ -626,7 +626,7 @@ class Connection(Base, FernetFieldsMixin, LoggingMixin):
@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 1c53237aa16..bf7138545f3 100644
--- a/airflow-core/src/airflow/models/variable.py
+++ b/airflow-core/src/airflow/models/variable.py
@@ -20,7 +20,6 @@ from __future__ import annotations
import contextlib
import json
import logging
-import sys
import warnings
from typing import TYPE_CHECKING, Any
@@ -33,6 +32,7 @@ from airflow._shared.secrets_masker import mask_secret
from airflow.configuration import conf, ensure_secrets_loaded
from airflow.models.base import ID_LEN, Base
from airflow.models.crypto import get_fernet
+from airflow.process_context import should_use_task_sdk_api_path
# AirflowSecretsBackendAccessDenied was added to task-sdk in 1.2.2. When
# airflow-core is installed alongside an older published task-sdk (e.g. 1.2.1
or earlier),
@@ -188,7 +188,7 @@ class Variable(Base, LoggingMixin):
# 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():
if session is not None:
raise ValueError(
"Variable.get() cannot use a metadata database session
from an execution context; "
@@ -255,7 +255,7 @@ class Variable(Base, LoggingMixin):
# 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",
@@ -342,7 +342,7 @@ class Variable(Base, LoggingMixin):
# 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.",
@@ -408,7 +408,7 @@ class Variable(Base, LoggingMixin):
# 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 00000000000..af506bc00a2
--- /dev/null
+++ b/airflow-core/src/airflow/process_context.py
@@ -0,0 +1,53 @@
+# 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
+
+__all__ = [
+ "force_server_context",
+ "should_use_task_sdk_api_path",
+]
+
+_FORCE_SERVER_CONTEXT: ContextVar[bool] =
ContextVar("_AIRFLOW_FORCE_SERVER_CONTEXT", default=False)
+
+
+@contextmanager
+def force_server_context() -> Generator[None, None, None]:
+ """Handle the active execution flow as server-side, even if
``SUPERVISOR_COMMS`` is set."""
+ token = _FORCE_SERVER_CONTEXT.set(True)
+ try:
+ yield
+ finally:
+ _FORCE_SERVER_CONTEXT.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 _FORCE_SERVER_CONTEXT.get():
+ 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 5f708651294..e4c9199e3b6 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
@@ -231,6 +231,39 @@ def
test_in_process_task_xcom_write_uses_its_attempt_id(create_task_instance, se
assert stored.value == "correct"
+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"),
+ ),
+ httpx2.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 8d5da4fa3ae..d5cb7fca7f5 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 force_server_context
pytestmark = pytest.mark.db_test
@@ -137,6 +139,38 @@ class TestGetConnection:
"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 force_server_context():
+ 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 f078d6c2fe0..e7e8fc0e717 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 fastapi.routing import Mount
from sqlalchemy import select
from airflow.models.variable import Variable
+from airflow.process_context import force_server_context
from tests_common.test_utils.db import clear_db_variables
@@ -96,6 +97,24 @@ class TestGetVariable:
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 force_server_context():
+ 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 @@ class TestPutVariable:
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 force_server_context():
+ 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 @@ class TestDeleteVariable:
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 (
+ force_server_context(),
+ 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 3e83a20cf2c..c4b7a728310 100644
--- a/airflow-core/tests/unit/models/test_connection.py
+++ b/airflow-core/tests/unit/models/test_connection.py
@@ -27,6 +27,7 @@ from structlog.testing import capture_logs
from airflow.exceptions import AirflowException, AirflowNotFoundException
from airflow.models import Connection
+from airflow.process_context import force_server_context
from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType
from airflow.sdk.execution_time.comms import ErrorResponse
@@ -565,6 +566,26 @@ class TestConnection:
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 force_server_context():
+ 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/contributing-docs/31_task_execution_architecture.rst
b/contributing-docs/31_task_execution_architecture.rst
index 6b12ddd0a14..14d4670dc66 100644
--- a/contributing-docs/31_task_execution_architecture.rst
+++ b/contributing-docs/31_task_execution_architecture.rst
@@ -41,9 +41,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* (one Python process for the task
itself, 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:: images/diagram_task_sdk_execution_architecture.png
diff --git
a/contributing-docs/images/diagram_task_sdk_execution_architecture.md5sum
b/contributing-docs/images/diagram_task_sdk_execution_architecture.md5sum
index f5ddfd6b4c9..c6e01eb7524 100644
--- a/contributing-docs/images/diagram_task_sdk_execution_architecture.md5sum
+++ b/contributing-docs/images/diagram_task_sdk_execution_architecture.md5sum
@@ -1 +1 @@
-3347c12823e4feee7dd92e378b5b76e8
+24a90b983c117809b10b03ae2db45159
diff --git
a/contributing-docs/images/diagram_task_sdk_execution_architecture.png
b/contributing-docs/images/diagram_task_sdk_execution_architecture.png
index 7550e570489..0c97f44c153 100644
Binary files
a/contributing-docs/images/diagram_task_sdk_execution_architecture.png and
b/contributing-docs/images/diagram_task_sdk_execution_architecture.png differ
diff --git
a/contributing-docs/images/diagram_task_sdk_execution_architecture.py
b/contributing-docs/images/diagram_task_sdk_execution_architecture.py
index cc483cc03d7..bdb5f3e5475 100644
--- a/contributing-docs/images/diagram_task_sdk_execution_architecture.py
+++ b/contributing-docs/images/diagram_task_sdk_execution_architecture.py
@@ -226,12 +226,12 @@ def generate_task_sdk_execution_architecture_diagram():
)
# ------------------------------------------------------------------ #
- # Native / in-process path — one Python process, no fork, no sockets.
+ # Native / in-process path — one Python process for the task itself, no
HTTP.
# ------------------------------------------------------------------ #
with g.subgraph(name="cluster_native") as nat:
nat.attr(
label="Native (in-process) path — dag.test() / local run"
- " · ONE Python process · no fork · no sockets · no HTTP",
+ " · ONE Python process for the task itself · no HTTP",
labelloc="t",
style="rounded,filled",
fillcolor="#faf3fc",
@@ -253,7 +253,7 @@ def generate_task_sdk_execution_architecture_diagram():
nat,
"in_comms",
"InProcessSupervisorComms",
- "in-memory deques, not sockets",
+ "direct call for the task · supervisor socket<br/>for child
processes (e.g. virtualenv)",
shape="box",
theme=NATIVE,
)
@@ -269,7 +269,7 @@ def generate_task_sdk_execution_architecture_diagram():
nat.edge(
"in_runner",
"in_comms",
- label="ToTask / ToSupervisor messages\nvia deque.append / popleft",
+ label="ToTask / ToSupervisor messages\nresponse returned to the
caller",
color=NATIVE[1],
fontcolor=NATIVE[1],
dir="both",
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index 38c14782351..3b1219a4f14 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -36,6 +36,7 @@ import weakref
from collections import deque
from collections.abc import Callable, Generator
from contextlib import contextmanager, nullcontext, suppress
+from contextvars import ContextVar
from datetime import UTC, datetime
from enum import Enum, auto
from http import HTTPStatus
@@ -2528,26 +2529,44 @@ 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. When no
sink is set, the request
+came over the socket from a child process (read by the thread started in
+``_setup_subprocess_socket``) and is answered on the socket instead.
+
+Whether a sink is set depends on the call, not the thread, so the sink must
not be replaced with a
+thread-identity check: handling a child's ``GetVariable``/``GetConnection`` on
the socket thread calls
+``mask_secret``, which re-enters :meth:`InProcessSupervisorComms.send` on that
same thread. The nested
+``MaskSecret`` response lands in the nested sink, which is reset before the
outer response is written
+to the socket.
+"""
+
+
@attrs.define(kw_only=True)
class InProcessSupervisorComms:
- """In-process communication handler that uses deques instead of sockets."""
+ """In-process communication handler that returns each response directly to
its caller."""
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
@@ -2722,8 +2741,15 @@ class InProcessTestSupervisor(ActivitySubprocess):
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 d2381dc6046..e8621d625ce 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 signal
import socket
import subprocess
import sys
+import threading
import time
from contextlib import nullcontext
from dataclasses import dataclass, field
@@ -74,6 +75,7 @@ from airflow.sdk.api.datamodels._generated import (
TaskInstance,
TaskInstanceState,
TIRunContext,
+ VariableResponse,
)
from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType,
TaskAlreadyRunningError
@@ -5098,6 +5100,170 @@ class TestInProcessTestSupervisor:
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):
+ """In-process requests in flight at the same time must not be answered
with each other's response.
+
+ Both requests go through ``comms.send`` from different threads, with
the first one still
+ being handled when the second arrives, so each must read its response
from its own sink.
+ """
+ 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 not first.is_alive(), "first in-process request never got its
response"
+
+ 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")]
+
+ @patch("airflow.sdk.log._secrets_masker")
+ def
test_child_request_masking_secret_on_socket_thread_gets_only_its_response(
+ self, mock_secrets_masker, mocker
+ ):
+ """Handling a child's ``GetVariable`` on the socket thread calls
``mask_secret``, which
+ re-enters ``comms.send`` on that same thread with a sink set; the
nested ``MaskSecret`` reply
+ must stay in that sink and only the ``VariableResult`` may reach the
child's socket.
+ """
+ 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 =
VariableResponse(key="test_key", value="test_value")
+ comms_send = mocker.spy(InProcessSupervisorComms, "send")
+
+ received: list[Any] = []
+ with (
+ set_supervisor_comms(supervisor.comms),
+ supervisor._setup_subprocess_socket() as child_sock,
+ ):
+ comms = CommsDecoder(socket=child_sock)
+ 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"
+ child_sock.settimeout(0.1)
+ with pytest.raises(TimeoutError):
+ child_sock.recv(1)
+
+ assert received == [VariableResult(key="test_key", value="test_value")]
+ comms_send.assert_called_once_with(supervisor.comms,
MaskSecret(value="test_value", name="test_key"))
+
def test_inprocess_failure_callback_receives_exception(
self,
monkeypatch,
