kaxil commented on code in PR #65587:
URL: https://github.com/apache/airflow/pull/65587#discussion_r4008108578
##########
airflow-core/src/airflow/models/variable.py:
##########
@@ -171,7 +171,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():
Review Comment:
This needs a rebase: #71968 (2c9088e078) landed on 14 Sep and added a
`session is not None` guard as the first statement inside this same block, so
`git merge-tree apache/main HEAD` reports this file as the only conflict. The
resolution should be mechanical, keeping that guard inside the block and now
gated on `should_use_task_sdk_api_path()`: the guard is there to reject a
metastore session taken in an execution context, and in the in-process server
context a session is legitimate (no Execution API route passes one today
anyway). Worth noting the last green CI run is from a base 304 commits old, so
it says nothing about the merged result.
##########
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* (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
Review Comment:
The figure referenced three lines below still asserts what this sentence now
denies: the generator at
`contributing-docs/images/diagram_task_sdk_execution_architecture.py:234`
renders the cluster label "ONE Python process / no fork / no sockets / no
HTTP", and `:256` labels the `InProcessSupervisorComms` node "in-memory deques,
not sockets", which is the deque this PR removes. The
`generate-airflow-diagrams` prek hook rebuilds the PNG but only fires on
`^contributing-docs/images/diagram_[^/]*\.py$`, so editing those two strings is
what gets the image regenerated. Line 44 also still says "no fork" while the
next sentence describes an operator spawning its own child process, so
something like "one Python process for the task itself, no HTTP" would settle
it.
##########
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]:
Review Comment:
`should_use_task_sdk_api_path()` only tests `== "server"` at line 53, so
`"client"` here is indistinguishable from not using the context manager at all,
and nothing in the repo passes it. The failure is quiet and in the unsafe
direction: someone writing `with override_process_context("client")` to force
the SDK path gets server behaviour and no error. Narrowing to
`Literal["server"]`, or a no-argument `force_server_context()`, would make the
signature match what is implemented; `_PROCESS_CONTEXT_OVERRIDE` is also
declared `ContextVar[str | None]`, wider than the Literal and unvalidated.
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2014,26 +2015,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
Review Comment:
This says the socket thread never has a sink set, but every child
`GetVariable`/`GetConnection` handled on that thread calls `mask_secret`
(`request_handlers.py:106` and `:94`), which re-enters
`InProcessSupervisorComms.send` at `sdk/log.py:266` and sets a sink on that
very thread. The code survives it, since the nested token resets before the
outer `send_msg` runs and the nested `MaskSecret` reply lands in the sink
rather than on the child's socket, but this docstring is the only place the
design is written down and it is the sentence someone would rely on to replace
the sink with a thread-identity check. The class docstring just below at line
2032 also still says this handler "uses deques instead of sockets", which this
diff makes untrue in both halves.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -3782,6 +3783,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.
Review Comment:
This docstring describes the socket being serviced on its own thread while
the in-process task has a request outstanding, but there is no socket in this
test: `ConcurrentSupervisor` is built without `stdin` and both requests go
through `comms.send` as ordinary in-process callers. It is a real test of the
per-call sink, just not of the scenario named, and the interleaving that does
happen in production has no test at all: on the socket thread `mask_secret`
re-enters `comms.send` inside the child's `GetVariable`, which is the one case
where a live sink and the `sink is None` branch coexist (all three socket tests
patch `mask_secret` out). Also `first.join(10)` below has no `assert not
first.is_alive()`, so a regression that hangs the first caller surfaces 20s
later as an opaque dict comparison instead of the clear message the sibling
test gives at line 3906.
--
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]
