amoghrajesh commented on code in PR #74074: URL: https://github.com/apache/airflow/pull/74074#discussion_r4559979745
########## 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: nit: rename to `get_supervisor_comms()`. ########## 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: This description shows up in mutliple places. Can we only keep it here and also shorten it to someting like: ``` # True while code serves an in-process supervisor request (dag.test()). That code acts as # the server side, so supervisor_comms() returns None there. Other task threads are unaffected. ``` ########## airflow-core/tests/unit/core/test_process_context.py: ########## @@ -0,0 +1,61 @@ +# +# 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 types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from airflow.process_context import force_server_context, should_use_task_sdk_api_path + +TASK_RUNNER = "airflow.sdk.execution_time.task_runner" + + +class TestShouldUseTaskSdkApiPath: + """Whether models.Variable/Connection take the Execution API path, per thread.""" + + def test_false_when_the_task_runner_was_never_imported(self, monkeypatch): Review Comment: This test and `test_forced_server_context_wins_over_the_comms` both pass on main. Can we drop them? ########## task-sdk/tests/task_sdk/execution_time/test_supervisor.py: ########## @@ -5319,6 +5320,91 @@ def fake_api_client(dag, attempt_id): assert isinstance(collected[0], _Failure) +class TestInProcessSupervisorCommsAcrossThreads: + """ + dag.test() serves a task's requests in the task's own process, possibly for several of the task's + threads at once. While a request is served, the code serving it stops seeing the comms, and only + that code: the task's other threads keep theirs. The in-process API server's side is covered in the + execution API's test_app.py. + """ + + @staticmethod + def _comms(serve): + """In-process comms whose supervisor answers each request with ``serve(msg)``.""" + comms = InProcessSupervisorComms(supervisor=mock.Mock(spec=InProcessTestSupervisor)) + + def handle_request(msg, log, req_id): + sink = _IN_PROCESS_RESPONSE_SINK.get() + assert sink is not None + sink.append(serve(msg)) + + comms.supervisor._handle_request.side_effect = handle_request + return comms + + def test_only_the_code_serving_a_request_stops_seeing_the_comms(self): + """ + While one thread's request is served, the task's other threads must keep their comms: seeing + none, they would take their Variable and Connection lookups for the server side and find nothing. + """ + import threading Review Comment: Already imported at top ########## 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: This version of task sdk might run with an older core (in practice we should do it). Airflow core 3.4 checks `hasattr(task_runner, "SUPERVISOR_COMMS")` in `models.Variable` and `models.Connection` to decide whether it's the server comms. Now that the attribute stays set while a request is served, the inprocess API in 3.4 will route lookups back through the comms. I think that fixes the dag.test() hang in #65587. Can we keep deleting the attribute for older airflow cores please? ########## 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: If `serving_supervisor_request` is ever renamed, this `getattr` quietly falls back to `nullcontext` and the fix stops working. Can we add a test that checks the SDK still has `serving_supervisor_request` and `supervisor_comms`? ########## airflow-core/tests/unit/models/test_taskinstance.py: ########## @@ -5668,3 +5668,86 @@ def test_mapped_archival_ack_uses_fork_facts_when_definition_is_missing(loop_run archived = session.get(TaskInstance, archiving_id) assert (archived.working_set, archived.archived_reason) == (None, "superseded") + + +class TestInProcessSupervisorFromThreads: + """ + Under the in-process supervisor (dag.test(), dag_maker) a task's own threads make SDK calls + while others are being served: every call must go through and get its own answer. + """ + + @staticmethod + def _run(dag_maker, session): + dag_run = dag_maker.create_dagrun() + (ti,) = dag_run.task_instance_scheduling_decisions(session=session).schedulable_tis + with contextlib.suppress(BaseException): Review Comment: `suppress(BaseException)` hides the task's error, so a failure only shows `state != SUCCESS`. Can we drop it? ########## 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: nit: rename to `hide_supervisor_comms()`, to match `force_server_context`. ########## airflow-core/tests/unit/models/test_taskinstance.py: ########## @@ -5668,3 +5668,86 @@ def test_mapped_archival_ack_uses_fork_facts_when_definition_is_missing(loop_run archived = session.get(TaskInstance, archiving_id) assert (archived.working_set, archived.archived_reason) == (None, "superseded") + + +class TestInProcessSupervisorFromThreads: Review Comment: This tests the inprocess supervisor, not `TaskInstance`. It also only fails on main when threads happen to overlap. Can it move next to the `dag.test()` tests? -- 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]
