image-match/python/py/Lib/site-packages/ray/dashboard/subprocesses/handle.py

import asyncio
import logging
import multiprocessing
import os
from typing import Optional, Union

import multidict

import ray.dashboard.consts as dashboard_consts
from ray.dashboard.optional_deps import aiohttp
from ray.dashboard.subprocesses.module import (
    SubprocessModule,
    SubprocessModuleConfig,
    run_module,
)
from ray.dashboard.subprocesses.utils import (
    ResponseType,
    get_http_session_to_module,
    module_logging_filename,
)

"""
This file contains code run in the parent process. It can start a subprocess and send
messages to it. Requires non-minimal Ray.
"""

logger = logging.getLogger(__name__)


def filter_hop_by_hop_headers(
    headers: Union[dict[str, str], multidict.CIMultiDictProxy[str]],
) -> dict[str, str]:
    """
    Filter out hop-by-hop headers from the headers dict.
    """
    HOP_BY_HOP_HEADERS = {
        "connection",
        "keep-alive",
        "proxy-authenticate",
        "proxy-authorization",
        "te",
        "trailers",
        "transfer-encoding",
        "upgrade",
    }
    if isinstance(headers, multidict.CIMultiDictProxy):
        headers = dict(headers)
    filtered_headers = {
        key: value
        for key, value in headers.items()
        if key.lower() not in HOP_BY_HOP_HEADERS
    }
    return filtered_headers


class SubprocessModuleHandle:
    """
    A handle to a module created as a subprocess. Can send messages to the module and
    receive responses. It only acts as a proxy to the aiohttp server running in the
    subprocess. On destruction, the subprocess is terminated.

    Lifecycle:
    1. In SubprocessModuleHandle creation, the subprocess is started and runs an aiohttp
       server.
    2. User must call start_module() and wait_for_module_ready() first.
    3. SubprocessRouteTable.bind(handle)
    4. app.add_routes(routes=SubprocessRouteTable.bound_routes())
    5. Run the app.

    Health check (_do_periodic_health_check):
    Every 1s, do a health check by _do_once_health_check. If the module is
    unhealthy:
      1. log the exception
      2. log the last N lines of the log file
      3. fail all active requests
      4. restart the module

    TODO(ryw): define policy for health check:
    - check period (Now: 1s)
    - define unhealthy. (Now: process exits. TODO: check_health() for event loop hang)
    - check number of failures in a row before we deem it unhealthy (Now: N/A)
    - "max number of restarts"? (Now: infinite)
    """

    # Class variable. Force using spawn because Ray C bindings have static variables
    # that need to be re-initialized for a new process.
    mp_context = multiprocessing.get_context("spawn")

    def __init__(
        self,
        loop: asyncio.AbstractEventLoop,
        module_cls: type[SubprocessModule],
        config: SubprocessModuleConfig,
    ):
        self.loop = loop
        self.module_cls = module_cls
        self.config = config

        # Increment this when the module is restarted.
        self.incarnation = 0
        # Runtime states, set by start_module() and wait_for_module_ready(),
        # reset by destroy_module().
        self.parent_conn = None
        self.process = None
        self.http_client_session: Optional[aiohttp.ClientSession] = None
        self.health_check_task = None

    def str_for_state(self, incarnation: int, pid: Optional[int]):
        return f"SubprocessModuleHandle(module_cls={self.module_cls.__name__}, incarnation={incarnation}, pid={pid})"

    def __str__(self):
        return self.str_for_state(
            self.incarnation, self.process.pid if self.process else None
        )

    def start_module(self):
        """
        Start the module. Should be non-blocking.
        """
        self.parent_conn, child_conn = self.mp_context.Pipe()
        if not os.path.exists(self.config.socket_dir):
            os.makedirs(self.config.socket_dir)
        self.process = self.mp_context.Process(
            target=run_module,
            args=(
                self.module_cls,
                self.config,
                self.incarnation,
                child_conn,
            ),
            daemon=True,
            name=f"{self.module_cls.__name__}-{self.incarnation}",
        )
        self.process.start()
        child_conn.close()

    def wait_for_module_ready(self):
        """
        Wait for the module to be ready. This is called after start_module()
        and can be blocking.
        """
        if self.parent_conn.poll(dashboard_consts.SUBPROCESS_MODULE_WAIT_READY_TIMEOUT):
            try:
                self.parent_conn.recv()
            except EOFError:
                raise RuntimeError(
                    f"Module {self.module_cls.__name__} failed to start. "
                    "Received EOF from pipe."
                )
            self.parent_conn.close()
            self.parent_conn = None
        else:
            raise RuntimeError(
                f"Module {self.module_cls.__name__} failed to start. "
                f"Timeout after {dashboard_consts.SUBPROCESS_MODULE_WAIT_READY_TIMEOUT} seconds."
            )

        module_name = self.module_cls.__name__
        self.http_client_session = get_http_session_to_module(
            module_name, self.config.socket_dir, self.config.session_name
        )

        self.health_check_task = self.loop.create_task(self._do_periodic_health_check())

    async def destroy_module(self):
        """
        Destroy the module with complete resource cleanup.
        This is called when the module is unhealthy or being shut down.
        """
        self.incarnation += 1

        # 1. Cancel health check task first to avoid race conditions
        if self.health_check_task:
            # NOTE: destroy_module() can be invoked from within the periodic health
            # check task itself (see _do_periodic_health_check()).
            # Cancelling the *current* task would raise CancelledError at the next
            # await and prevent cleanup + restart from completing.
            current_task = asyncio.current_task()
            if current_task is None or self.health_check_task is not current_task:
                self.health_check_task.cancel()
            self.health_check_task = None

        # 2. Close parent connection
        if self.parent_conn:
            self.parent_conn.close()
            self.parent_conn = None

        # 3. Terminate process gracefully, then forcefully if needed
        if self.process:
            try:
                # First, try graceful termination
                if self.process.is_alive():
                    self.process.terminate()
                    logger.debug(
                        f"Terminated process {self.process.pid}, waiting for exit..."
                    )

                    # Wait for process to exit (with timeout)
                    self.process.join(
                        timeout=dashboard_consts.SUBPROCESS_MODULE_GRACEFUL_SHUTDOWN_TIMEOUT
                    )

                    # Force kill if still alive
                    if self.process.is_alive():
                        logger.warning(
                            f"Process {self.process.pid} did not exit gracefully, "
                            "force killing..."
                        )
                        self.process.kill()
                        self.process.join(
                            timeout=dashboard_consts.SUBPROCESS_MODULE_JOIN_TIMEOUT
                        )
                else:
                    # Process already dead, just wait for it
                    self.process.join(
                        timeout=dashboard_consts.SUBPROCESS_MODULE_JOIN_TIMEOUT
                    )

                logger.debug(f"Process {self.process.pid} terminated successfully")
            except Exception as e:
                logger.warning(f"Error terminating process: {e}")
            finally:
                self.process = None

        # 4. Close HTTP client session with proper cleanup
        if self.http_client_session:
            await self.http_client_session.close()
            self.http_client_session = None

    async def _health_check(self) -> aiohttp.web.Response:
        """
        Do internal health check. The module should respond immediately with a 200 OK.
        This can be used to measure module responsiveness in RTT, it also indicates
        subprocess event loop lag.

        Currently you get a 200 OK with body = b'success'. Later if we want we can add more
        observability payloads.
        """
        resp = await self.http_client_session.get("http://localhost/api/healthz")
        return aiohttp.web.Response(
            status=resp.status,
            headers=filter_hop_by_hop_headers(resp.headers),
            body=await resp.read(),
        )

    async def _do_once_health_check(self):
        """
        Do a health check once. We check for:
        1. if the process exits, it's considered died.
        2. if the health check endpoint returns non-200, it's considered unhealthy.

        """
        if self.process.exitcode is not None:
            raise RuntimeError(f"Process exited with code {self.process.exitcode}")
        resp = await self._health_check()
        if resp.status != 200:
            raise RuntimeError(f"Health check failed: status code is {resp.status}")

    async def _do_periodic_health_check(self):
        """
        Every 1s, do a health check. If the module is unhealthy:
        1. log the exception
        2. log the last N lines of the log file
        3. restart the module
        """
        while True:
            try:
                await self._do_once_health_check()
            except Exception:
                filename = module_logging_filename(
                    self.module_cls.__name__, self.config.logging_filename
                )
                logger.exception(
                    f"Module {self.module_cls.__name__} is unhealthy. Please refer to "
                    f"{self.config.log_dir}/{filename} for more details. Failing all "
                    "active requests."
                )
                await self.destroy_module()
                self.start_module()
                self.wait_for_module_ready()
                return
            await asyncio.sleep(1)

    async def proxy_request(
        self, request: aiohttp.web.Request, resp_type: ResponseType = ResponseType.HTTP
    ) -> aiohttp.web.StreamResponse:
        """
        Sends a new request to the subprocess and returns the response.
        """
        if resp_type == ResponseType.HTTP:
            return await self.proxy_http(request)
        if resp_type == ResponseType.STREAM:
            return await self.proxy_stream(request)
        if resp_type == ResponseType.WEBSOCKET:
            return await self.proxy_websocket(request)
        raise ValueError(f"Unknown response type: {resp_type}")

    async def proxy_http(self, request: aiohttp.web.Request) -> aiohttp.web.Response:
        """
        Proxy handler for non-streaming HTTP API
        It forwards the method, query string, headers, and body to the backend.
        """
        url = f"http://localhost{request.path_qs}"
        body = await request.read()

        async with self.http_client_session.request(
            request.method,
            url,
            data=body,
            headers=filter_hop_by_hop_headers(request.headers),
            allow_redirects=False,
        ) as backend_resp:
            resp_body = await backend_resp.read()
            return aiohttp.web.Response(
                status=backend_resp.status,
                headers=filter_hop_by_hop_headers(backend_resp.headers),
                body=resp_body,
            )

    async def proxy_stream(
        self, request: aiohttp.web.Request
    ) -> aiohttp.web.StreamResponse:
        """
        Proxy handler for streaming HTTP API.
        It forwards the method, query string, and body to the backend.
        """
        url = f"http://localhost{request.path_qs}"
        body = await request.read()
        async with self.http_client_session.request(
            request.method,
            url,
            data=body,
            headers=filter_hop_by_hop_headers(request.headers),
        ) as backend_resp:
            proxy_resp = aiohttp.web.StreamResponse(
                status=backend_resp.status,
                headers=filter_hop_by_hop_headers(backend_resp.headers),
            )
            await proxy_resp.prepare(request)

            async for chunk, _ in backend_resp.content.iter_chunks():
                await proxy_resp.write(chunk)
            await proxy_resp.write_eof()
            return proxy_resp

    async def proxy_websocket(
        self, request: aiohttp.web.Request
    ) -> aiohttp.web.StreamResponse:
        """
        Proxy handler for WebSocket API
        It establishes a WebSocket connection with the client and simultaneously connects
        to the backend server's WebSocket endpoint. Messages are forwarded in single
        direction from the backend to the client.
        If the backend responds with normal HTTP response, then try to treat it as a normal
        HTTP request and calls proxy_http instead.

        TODO: Support bidirectional communication if needed. We only support one direction
              because it's sufficient for the current use case.
        """
        url = f"http://localhost{request.path_qs}"

        try:
            async with self.http_client_session.ws_connect(
                url, headers=filter_hop_by_hop_headers(request.headers)
            ) as ws_to_backend:
                ws_from_client = aiohttp.web.WebSocketResponse()
                await ws_from_client.prepare(request)
                async for msg in ws_to_backend:
                    if msg.type == aiohttp.WSMsgType.TEXT:
                        await ws_from_client.send_str(msg.data)
                    elif msg.type == aiohttp.WSMsgType.BINARY:
                        await ws_from_client.send_bytes(msg.data)
                    else:
                        logger.error(f"Unknown msg type: {msg.type}")
                await ws_from_client.close()
                return ws_from_client
        except aiohttp.WSServerHandshakeError as e:
            logger.warning(f"WebSocket handshake error: {repr(e)}")
            # Try to treat it as a normal HTTP request
            return await self.proxy_http(request)
        except Exception as e:
            logger.error(f"WebSocket proxy error: {repr(e)}")
            raise aiohttp.web.HTTPInternalServerError(reason="WebSocket proxy error")