Align the latest coordinator interface for executable coordinator

Author: jason810496

task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py

@@ -15,24 +15,74 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""Native executable coordinator for DAG file processing and task execution."""
+"""Native executable coordinator that launches a binary subprocess for Dag file processing and task execution."""
 
 from __future__ import annotations
 
 import os
 import pathlib
-from typing import TYPE_CHECKING
+import selectors
+import socket
+import subprocess
+import time
+from typing import TYPE_CHECKING, cast
 
 import attrs
-from structlog import get_logger
+import psutil
+import structlog
 
 from airflow.sdk.coordinators.executable.bundle_scanner import BundleScanner, read_source_code
 from airflow.sdk.execution_time.coordinator import BaseCoordinator
+from airflow.sdk.execution_time.supervisor import ActivitySubprocess
 
 if TYPE_CHECKING:
-    from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance
-
-log = get_logger(__name__)
+    from structlog.typing import FilteringBoundLogger
+    from typing_extensions import Self
+
+    from airflow.sdk.api.client import Client
+    from airflow.sdk.api.datamodels._generated import BundleInfo
+    from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO
+
+log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.executable")
+
+
+# TODO: extract the duplicated `_start_server`, `_accept_connections`, and the
+# ``ActivitySubprocess`` subclass below into a shared ``IPCActivitySubprocess``
+# base in ``airflow.sdk.execution_time.supervisor`` so JavaCoordinator and
+# ExecutableCoordinator (and future coordinator subclasses) can share the
+# same comm/logs server bootstrap and socket cleanup.
+def _start_server() -> socket.socket:
+    server = socket.socket()
+    server.bind(("", 0))
+    server.setblocking(True)
+    server.listen(1)  # Just need to listen to the child process.
+    return server
+
+
+def _accept_connections(
+    servers: dict[str, socket.socket],
+    proc: subprocess.Popen,
+    *,
+    max_wait: float = 10.0,
+) -> dict[str, socket.socket]:
+    """Block until the executable process connects to servers."""
+    accepted: dict[str, socket.socket] = {}
+    with selectors.DefaultSelector() as sel:
+        for key, soc in servers.items():
+            sel.register(soc, selectors.EVENT_READ, data=key)
+        deadline = time.monotonic() + max_wait
+        while len(accepted) < len(servers):
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise TimeoutError("process did not connect within timeout")
+            if proc.poll() is not None:
+                raise RuntimeError(f"process exited with {proc.returncode} before connecting")
+            for event, _ in sel.select(timeout=min(remaining, 1.0)):
+                log.debug("Accepting child process connection", key=(key := event.data))
+                conn, _ = cast("socket.socket", event.fileobj).accept()
+                sel.unregister(servers[key])
+                accepted[key] = conn
+    return accepted
 
 
 def _optional_path(value: str | os.PathLike[str] | None) -> pathlib.Path | None:
@@ -41,6 +91,82 @@ def _optional_path(value: str | os.PathLike[str] | None) -> pathlib.Path | None:
     return pathlib.Path(value)
 
 
+@attrs.define(kw_only=True)
+class _ExecutableActivitySubprocess(ActivitySubprocess):
+    """Native executable task runner process."""
+
+    _comm_server: socket.socket
+    _logs_server: socket.socket
+    _child_process: subprocess.Popen
+
+    # Keep track of channels used to pipe subprocess stdout and stderr so we can
+    # close them on exit. The "read" side is handled by _register_pipe_readers
+    # callbacks so we don't need to worry about them.
+    _stdout_w: socket.socket
+    _stderr_w: socket.socket
+
+    @classmethod
+    def start(  # type: ignore[override]
+        cls,
+        *,
+        what: TaskInstanceDTO,
+        dag_rel_path: str | os.PathLike[str],
+        bundle_info,
+        logger: FilteringBoundLogger | None = None,
+        sentry_integration: str = "",
+        executable_path: str,
+        **kwargs,
+    ) -> Self:
+        comm_server = _start_server()
+        logs_server = _start_server()
+
+        stdout_r, stdout_w = socket.socketpair()
+        stderr_r, stderr_w = socket.socketpair()
+
+        comm_host, comm_port = comm_server.getsockname()
+        logs_host, logs_port = logs_server.getsockname()
+
+        proc = subprocess.Popen(
+            [
+                executable_path,
+                f"--comm={comm_host}:{comm_port}",
+                f"--logs={logs_host}:{logs_port}",
+            ],
+            stdout=stdout_w.makefile("wb", buffering=0).fileno(),
+            stderr=stderr_w.makefile("wb", buffering=0).fileno(),
+        )
+        log.info("Starting subprocess", pid=proc.pid, executable=executable_path)
+        socks = _accept_connections({"comm": comm_server, "logs": logs_server}, proc)
+
+        self = cls(
+            id=what.id,
+            pid=proc.pid,
+            process=psutil.Process(proc.pid),
+            process_log=logger or structlog.get_logger(logger_name="task").bind(),
+            start_time=time.monotonic(),
+            stdin=socks["comm"],
+            child_process=proc,
+            comm_server=comm_server,
+            logs_server=logs_server,
+            stdout_w=stdout_w,
+            stderr_w=stderr_w,
+            **kwargs,
+        )
+        self._register_pipe_readers(stdout_r, stderr_r, socks["comm"], socks["logs"])
+        self._on_child_started(
+            ti=what,
+            dag_rel_path=dag_rel_path,
+            bundle_info=bundle_info,
+            sentry_integration=sentry_integration,
+        )
+        return self
+
+    def wait(self) -> int:
+        code = super().wait()
+        self._close_unused_sockets(self._comm_server, self._logs_server, self._stdout_w, self._stderr_w)
+        return code
+
+
 @attrs.define(kw_only=True)
 class ExecutableCoordinator(BaseCoordinator):
     """
@@ -87,71 +213,42 @@ def get_code_from_file(cls, fileloc: str) -> str:
             raise FileNotFoundError(f"No source code found for executable: {fileloc}")
         return code
 
-    @classmethod
-    def dag_parsing_cmd(
-        cls,
-        *,
-        dag_file_path: str,
-        bundle_name: str,
-        bundle_path: str,
-        comm_addr: str,
-        logs_addr: str,
-    ) -> list[str]:
-        """Build the subprocess command for parsing a native executable bundle."""
-        return [
-            dag_file_path,
-            f"--comm={comm_addr}",
-            f"--logs={logs_addr}",
-        ]
-
-    def task_execution_cmd(
+    def _resolve_executable(self, *, what: TaskInstanceDTO) -> str:
+        """
+        Resolve the executable path for *what*.
+
+        Looks up the bundle whose embedded manifest declares ``what.dag_id``
+        in the configured ``bundles_folder``.
+        """
+        if self.bundles_folder is None:
+            raise ValueError(
+                "The bundles_folder kwarg must be set on the ExecutableCoordinator "
+                "to resolve the executable for task execution."
+            )
+        return BundleScanner(self.bundles_folder).resolve(dag_id=what.dag_id)
+
+    def execute_task(
         self,
         *,
-        what: TaskInstance,
-        dag_file_path: str,
-        bundle_path: str,
+        what: TaskInstanceDTO,
+        dag_rel_path: str | os.PathLike[str],
         bundle_info: BundleInfo,
-        comm_addr: str,
-        logs_addr: str,
-    ) -> list[str]:
-        """Build the subprocess command for executing a task in a native executable bundle."""
-        # Case 1: Pure executable DAG — dag_file_path is itself a bundle.
-        # A ``.py`` extension always indicates a Python stub, so short-circuit
-        # there before validating the AFBNDL01 trailer. Both checks are needed:
-        # ``os.access(_, X_OK)`` alone is not reliable because bind-mounted
-        # Python stubs can satisfy X_OK (e.g. when running as root inside
-        # Breeze) yet are not real bundles.
-        log.debug(
-            "Resolving executable for task execution",
+        client: Client,
+        logger: FilteringBoundLogger | None = None,
+        sentry_integration: str = "",
+        subprocess_logs_to_stdout: bool,
+        **kwargs,
+    ) -> BaseCoordinator.ExecutionResult:
+        executable_path = self._resolve_executable(what=what)
+        process = _ExecutableActivitySubprocess.start(
             what=what,
-            dag_file_path=dag_file_path,
-            bundle_path=bundle_path,
+            dag_rel_path=dag_rel_path,
             bundle_info=bundle_info,
-            comm_addr=comm_addr,
-            logs_addr=logs_addr,
+            client=client,
+            logger=logger,
+            subprocess_logs_to_stdout=subprocess_logs_to_stdout,
+            sentry_integration=sentry_integration,
+            executable_path=executable_path,
         )
-        if (
-            not dag_file_path.endswith(".py")
-            and (resolved := BundleScanner.resolve_executable(pathlib.Path(dag_file_path))) is not None
-        ):
-            return [
-                resolved,
-                f"--comm={comm_addr}",
-                f"--logs={logs_addr}",
-            ]
-
-        # Case 2: Python Stub DAG — the dag_file_path is a Python file but
-        # the task delegates to a native runtime.  The actual binary lives
-        # in the coordinator's configured ``bundles_folder`` kwarg.
-        if self.bundles_folder is None:
-            raise ValueError(
-                "The bundles_folder kwarg must be set on the ExecutableCoordinator "
-                "for Python stub DAGs that delegate to native executable task execution."
-            )
-
-        executable_path = BundleScanner(self.bundles_folder).resolve(dag_id=what.dag_id)
-        return [
-            executable_path,
-            f"--comm={comm_addr}",
-            f"--logs={logs_addr}",
-        ]
+        exit_code = process.wait()
+        return self.ExecutionResult(exit_code, process.final_state)