dabla commented on code in PR #74074:
URL: https://github.com/apache/airflow/pull/74074#discussion_r4562044716
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2555,13 +2555,22 @@ class InProcessSupervisorComms:
supervisor: InProcessTestSupervisor
def send(self, msg: BaseModel) -> BaseModel | None:
- """Send a request to the supervisor and return its response."""
+ """
+ Send a request to the supervisor and return its response.
+
+ Safe to call from several threads: each call gets its own response
sink, and only the code
+ serving the request stops seeing the comms meanwhile
(``task_runner.serving_supervisor_request``,
+ also entered by the in-process API server for its own threads), so the
task's other threads
+ keep theirs.
+ """
+ from airflow.sdk.execution_time.task_runner import
serving_supervisor_request
+
self.log.debug("Sending request", msg=msg)
responses: deque[BaseModel | None] = deque()
token = _IN_PROCESS_RESPONSE_SINK.set(responses)
try:
- with set_supervisor_comms(None):
+ with serving_supervisor_request():
Review Comment:
Good catch, and confirmed on the `3.4.0` tag: `models/variable.py` and
`models/connection.py` decide with `hasattr(task_runner, "SUPERVISOR_COMMS")`,
and there is no `airflow.process_context` there, so the route's lookup would go
back through the comms. In 82a0f4b15e `send()` keeps removing the attribute
when the installed core has no `airflow.process_context`
(`_core_decides_server_side_per_request()`), and hides the comms per request
only otherwise. On such a core the behaviour is what it is today, threads
included. Covered by
`test_with_an_older_core_the_comms_are_removed_while_a_request_is_served`; I
read the chain from the tag and did not run it against a real 3.4.0 core.
##########
airflow-core/src/airflow/api_fastapi/execution_api/in_process.py:
##########
@@ -94,13 +95,24 @@ async def always_allow(request: Request):
class _RequestScopedServerContextApp:
- """Wrap an ASGI app so in-process requests behave like server-side API
handling."""
+ """
+ Wrap an ASGI app so in-process requests behave like server-side API
handling.
+
+ The flags are set in the request's context, which the event loop task and
the worker threads of
+ sync routes inherit, so the task's other threads keep their comms
meanwhile:
+ ``force_server_context`` for airflow-core (``models.Variable``,
``models.Connection``), and the
+ Task SDK's ``serving_supervisor_request`` for its own readers (mask
forwarding, the secrets
+ backend choice).
+ """
def __init__(self, app: FastAPI) -> None:
self.app = app
async def __call__(self, scope: Scope, receive: Receive, send: Send) ->
None:
- with force_server_context():
+ task_runner = sys.modules.get("airflow.sdk.execution_time.task_runner")
+ # No task runs in this process, or a Task SDK from before the flag:
nothing to hide from.
+ serving_supervisor_request = getattr(task_runner,
"serving_supervisor_request", nullcontext)
Review Comment:
Added `test_the_task_sdk_has_the_names_looked_up_with_getattr` in
`test_process_context.py` for both names (82a0f4b15e).
`test_in_process_execution_api_serves_requests_as_the_server_side` would also
fail on a rename, since the route would then see the comms, but it did not say
so.
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -1245,6 +1245,41 @@ def parse(what: StartupDetails, log: Logger) ->
RuntimeTaskInstance:
# accessible wherever needed during task execution without modifying every
layer of the call stack.
SUPERVISOR_COMMS: CommsDecoder[ToTask, ToSupervisor]
+# The in-process supervisor of dag.test() serves a task's requests in the
task's own process. The
+# code serving a request (models.Variable, models.Connection, the secrets
backends, mask forwarding)
+# has to act as the server side there, so it must not see the comms. That code
runs in the sender's
+# thread (the supervisor's _handle_request) and in the in-process API server,
on its event loop and
+# its worker threads for sync routes. Both set this flag around the request:
the sender here, the
+# server around each request it handles (InProcessExecutionAPI), and the
server's threads inherit it
+# from there. Being a ContextVar, it hides the comms from the code serving the
request only: the
+# task's other threads keep theirs while it is served.
Review Comment:
Done in 82a0f4b15e: your two lines are the comment at the ContextVar, and
the copies in `send()`, `_RequestScopedServerContextApp`,
`should_use_task_sdk_api_path()` and `sdk/log.py` are cut down to a pointer.
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -1245,6 +1245,41 @@ def parse(what: StartupDetails, log: Logger) ->
RuntimeTaskInstance:
# accessible wherever needed during task execution without modifying every
layer of the call stack.
SUPERVISOR_COMMS: CommsDecoder[ToTask, ToSupervisor]
+# The in-process supervisor of dag.test() serves a task's requests in the
task's own process. The
+# code serving a request (models.Variable, models.Connection, the secrets
backends, mask forwarding)
+# has to act as the server side there, so it must not see the comms. That code
runs in the sender's
+# thread (the supervisor's _handle_request) and in the in-process API server,
on its event loop and
+# its worker threads for sync routes. Both set this flag around the request:
the sender here, the
+# server around each request it handles (InProcessExecutionAPI), and the
server's threads inherit it
+# from there. Being a ContextVar, it hides the comms from the code serving the
request only: the
+# task's other threads keep theirs while it is served.
+_serving_supervisor_request: contextvars.ContextVar[bool] =
contextvars.ContextVar(
+ "_serving_supervisor_request", default=False
+)
+
+
+def supervisor_comms() -> CommsDecoder[ToTask, ToSupervisor] | None:
Review Comment:
Renamed to `get_supervisor_comms()` in 82a0f4b15e.
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -1245,6 +1245,41 @@ def parse(what: StartupDetails, log: Logger) ->
RuntimeTaskInstance:
# accessible wherever needed during task execution without modifying every
layer of the call stack.
SUPERVISOR_COMMS: CommsDecoder[ToTask, ToSupervisor]
+# The in-process supervisor of dag.test() serves a task's requests in the
task's own process. The
+# code serving a request (models.Variable, models.Connection, the secrets
backends, mask forwarding)
+# has to act as the server side there, so it must not see the comms. That code
runs in the sender's
+# thread (the supervisor's _handle_request) and in the in-process API server,
on its event loop and
+# its worker threads for sync routes. Both set this flag around the request:
the sender here, the
+# server around each request it handles (InProcessExecutionAPI), and the
server's threads inherit it
+# from there. Being a ContextVar, it hides the comms from the code serving the
request only: the
+# task's other threads keep theirs while it is served.
+_serving_supervisor_request: contextvars.ContextVar[bool] =
contextvars.ContextVar(
+ "_serving_supervisor_request", default=False
+)
+
+
+def supervisor_comms() -> CommsDecoder[ToTask, ToSupervisor] | None:
+ """
+ Return the comms to the supervisor, or ``None``.
+
+ ``None`` outside a task execution context, and in the code serving a
request of the in-process
+ supervisor (see :func:`serving_supervisor_request`). Code that decides
whether it runs inside a
+ task should ask this instead of checking for ``SUPERVISOR_COMMS``.
+ """
+ if _serving_supervisor_request.get():
+ return None
+ return globals().get("SUPERVISOR_COMMS")
+
+
+@contextmanager
+def serving_supervisor_request() -> Iterator[None]:
Review Comment:
Renamed to `hide_supervisor_comms()` in 82a0f4b15e; the ContextVar is
`_supervisor_comms_hidden` now, so there is one term for it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
