chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,383 @@
|
||||
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")
|
||||
@@ -0,0 +1,291 @@
|
||||
import abc
|
||||
import asyncio
|
||||
import inspect
|
||||
import logging
|
||||
import multiprocessing
|
||||
import multiprocessing.connection
|
||||
import os
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
|
||||
import aiohttp
|
||||
|
||||
import ray
|
||||
from ray import ray_constants
|
||||
from ray._private import logging_utils
|
||||
from ray._private.gcs_utils import GcsChannel
|
||||
from ray._private.ray_logging import setup_component_logger
|
||||
from ray._raylet import GcsClient
|
||||
from ray.dashboard.subprocesses.utils import (
|
||||
get_named_pipe_path,
|
||||
get_socket_path,
|
||||
module_logging_filename,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SubprocessModuleConfig:
|
||||
"""
|
||||
Configuration for a SubprocessModule.
|
||||
Pickleable.
|
||||
"""
|
||||
|
||||
cluster_id_hex: str
|
||||
gcs_address: str
|
||||
session_name: str
|
||||
temp_dir: str
|
||||
session_dir: str
|
||||
# Logger configs. Will be set up in subprocess entrypoint `run_module`.
|
||||
logging_level: str
|
||||
logging_format: str
|
||||
log_dir: str
|
||||
# Name of the "base" log file. Its stem is appended with the Module.__name__.
|
||||
# e.g. when logging_filename = "dashboard.log", and Module is JobHead,
|
||||
# we will set up logger with name "dashboard_JobHead.log". This name will again be
|
||||
# appended with .1 and .2 for rotation.
|
||||
logging_filename: str
|
||||
logging_rotate_bytes: int
|
||||
logging_rotate_backup_count: int
|
||||
# The directory where the socket file will be created.
|
||||
socket_dir: str
|
||||
|
||||
|
||||
class SubprocessModule(abc.ABC):
|
||||
"""
|
||||
A Dashboard Head Module that runs in a subprocess as a standalone aiohttp server.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: SubprocessModuleConfig,
|
||||
):
|
||||
"""Initialize current module when DashboardHead loading modules.
|
||||
|
||||
Args:
|
||||
config: The SubprocessModuleConfig instance.
|
||||
"""
|
||||
self._config = config
|
||||
self._parent_process = multiprocessing.parent_process()
|
||||
# Lazy init
|
||||
self._gcs_client = None
|
||||
self._aiogrpc_gcs_channel = None
|
||||
self._parent_process_death_detection_task = None
|
||||
self._http_session = None
|
||||
|
||||
async def _detect_parent_process_death(self):
|
||||
"""
|
||||
Detect parent process liveness. Only returns when parent process is dead.
|
||||
"""
|
||||
while True:
|
||||
if not self._parent_process.is_alive():
|
||||
logger.warning(
|
||||
f"Parent process {self._parent_process.pid} died. Exiting..."
|
||||
)
|
||||
return
|
||||
await asyncio.sleep(1)
|
||||
|
||||
@staticmethod
|
||||
def is_minimal_module():
|
||||
"""
|
||||
Currently all SubprocessModule classes should be non-minimal.
|
||||
|
||||
We require this because SubprocessModuleHandle tracks aiohttp requests and
|
||||
responses. To ease this, we can define another SubprocessModuleMinimalHandle
|
||||
that doesn't track requests and responses, but still provides Queue interface
|
||||
and health check.
|
||||
TODO(ryw): If needed, create SubprocessModuleMinimalHandle.
|
||||
"""
|
||||
return False
|
||||
|
||||
async def run(self):
|
||||
"""
|
||||
Start running the module.
|
||||
This method should be called first before the module starts receiving requests.
|
||||
"""
|
||||
app = aiohttp.web.Application(
|
||||
client_max_size=ray_constants.DASHBOARD_CLIENT_MAX_SIZE,
|
||||
)
|
||||
routes: list[aiohttp.web.RouteDef] = [
|
||||
aiohttp.web.get("/api/healthz", self._internal_module_health_check)
|
||||
]
|
||||
handlers = inspect.getmembers(
|
||||
self,
|
||||
lambda x: (
|
||||
inspect.ismethod(x)
|
||||
and hasattr(x, "__route_method__")
|
||||
and hasattr(x, "__route_path__")
|
||||
),
|
||||
)
|
||||
for _, handler in handlers:
|
||||
routes.append(
|
||||
aiohttp.web.route(
|
||||
handler.__route_method__,
|
||||
handler.__route_path__,
|
||||
handler,
|
||||
)
|
||||
)
|
||||
app.add_routes(routes)
|
||||
runner = aiohttp.web.AppRunner(app, access_log=None)
|
||||
await runner.setup()
|
||||
|
||||
module_name = self.__class__.__name__
|
||||
if sys.platform == "win32":
|
||||
named_pipe_path = get_named_pipe_path(
|
||||
module_name, self._config.session_name
|
||||
)
|
||||
site = aiohttp.web.NamedPipeSite(runner, named_pipe_path)
|
||||
logger.info(f"Started aiohttp server over {named_pipe_path}.")
|
||||
else:
|
||||
socket_path = get_socket_path(self._config.socket_dir, module_name)
|
||||
site = aiohttp.web.UnixSite(runner, socket_path)
|
||||
logger.info(f"Started aiohttp server over {socket_path}.")
|
||||
await site.start()
|
||||
|
||||
@property
|
||||
def gcs_client(self):
|
||||
if self._gcs_client is None:
|
||||
if not ray.experimental.internal_kv._internal_kv_initialized():
|
||||
gcs_client = GcsClient(
|
||||
address=self._config.gcs_address,
|
||||
cluster_id=self._config.cluster_id_hex,
|
||||
)
|
||||
ray.experimental.internal_kv._initialize_internal_kv(gcs_client)
|
||||
self._gcs_client = ray.experimental.internal_kv.internal_kv_get_gcs_client()
|
||||
return self._gcs_client
|
||||
|
||||
@property
|
||||
def aiogrpc_gcs_channel(self):
|
||||
if self._aiogrpc_gcs_channel is None:
|
||||
gcs_channel = GcsChannel(gcs_address=self._config.gcs_address, aio=True)
|
||||
gcs_channel.connect()
|
||||
self._aiogrpc_gcs_channel = gcs_channel.channel()
|
||||
return self._aiogrpc_gcs_channel
|
||||
|
||||
@property
|
||||
def session_name(self):
|
||||
"""
|
||||
Return the Ray session name. It's not related to the aiohttp session.
|
||||
"""
|
||||
return self._config.session_name
|
||||
|
||||
@property
|
||||
def temp_dir(self):
|
||||
return self._config.temp_dir
|
||||
|
||||
@property
|
||||
def session_dir(self):
|
||||
return self._config.session_dir
|
||||
|
||||
@property
|
||||
def log_dir(self):
|
||||
return self._config.log_dir
|
||||
|
||||
@property
|
||||
def http_session(self):
|
||||
if self._http_session is None:
|
||||
self._http_session = aiohttp.ClientSession()
|
||||
return self._http_session
|
||||
|
||||
@property
|
||||
def gcs_address(self):
|
||||
return self._config.gcs_address
|
||||
|
||||
async def _internal_module_health_check(self, request):
|
||||
return aiohttp.web.Response(
|
||||
text="success",
|
||||
content_type="application/text",
|
||||
)
|
||||
|
||||
|
||||
async def run_module_inner(
|
||||
cls: type[SubprocessModule],
|
||||
config: SubprocessModuleConfig,
|
||||
incarnation: int,
|
||||
child_conn: multiprocessing.connection.Connection,
|
||||
):
|
||||
|
||||
module_name = cls.__name__
|
||||
|
||||
logger.info(
|
||||
f"Starting module {module_name} with incarnation {incarnation} and config {config}"
|
||||
)
|
||||
|
||||
try:
|
||||
module = cls(config)
|
||||
module._parent_process_death_detection_task = asyncio.create_task(
|
||||
module._detect_parent_process_death()
|
||||
)
|
||||
module._parent_process_death_detection_task.add_done_callback(
|
||||
lambda _: sys.exit()
|
||||
)
|
||||
await module.run()
|
||||
child_conn.send(None)
|
||||
child_conn.close()
|
||||
logger.info(f"Module {module_name} initialized, receiving messages...")
|
||||
except Exception as e:
|
||||
logger.exception(f"Error creating module {module_name}")
|
||||
raise e
|
||||
|
||||
|
||||
def run_module(
|
||||
cls: type[SubprocessModule],
|
||||
config: SubprocessModuleConfig,
|
||||
incarnation: int,
|
||||
child_conn: multiprocessing.connection.Connection,
|
||||
):
|
||||
"""
|
||||
Entrypoint for a subprocess module.
|
||||
"""
|
||||
module_name = cls.__name__
|
||||
current_proctitle = ray._raylet.getproctitle()
|
||||
ray._raylet.setproctitle(
|
||||
f"ray-dashboard-{module_name}-{incarnation} ({current_proctitle})"
|
||||
)
|
||||
logging_filename = module_logging_filename(module_name, config.logging_filename)
|
||||
setup_component_logger(
|
||||
logging_level=config.logging_level,
|
||||
logging_format=config.logging_format,
|
||||
log_dir=config.log_dir,
|
||||
filename=logging_filename,
|
||||
max_bytes=config.logging_rotate_bytes,
|
||||
backup_count=config.logging_rotate_backup_count,
|
||||
)
|
||||
|
||||
if config.logging_filename:
|
||||
stdout_filename = module_logging_filename(
|
||||
module_name, config.logging_filename, extension=".out"
|
||||
)
|
||||
stderr_filename = module_logging_filename(
|
||||
module_name, config.logging_filename, extension=".err"
|
||||
)
|
||||
logging_utils.redirect_stdout_stderr_if_needed(
|
||||
os.path.join(config.log_dir, stdout_filename),
|
||||
os.path.join(config.log_dir, stderr_filename),
|
||||
config.logging_rotate_bytes,
|
||||
config.logging_rotate_backup_count,
|
||||
)
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
task = loop.create_task(
|
||||
run_module_inner(
|
||||
cls,
|
||||
config,
|
||||
incarnation,
|
||||
child_conn,
|
||||
)
|
||||
)
|
||||
# TODO: do graceful shutdown.
|
||||
# 1. define a stop token.
|
||||
# 2. join the loop to wait for all pending tasks to finish, up until a timeout.
|
||||
# 3. close the loop and exit.
|
||||
|
||||
def sigterm_handler(signum, frame):
|
||||
logger.warning(f"Exiting with signal {signum} immediately...")
|
||||
sys.exit(signum)
|
||||
|
||||
ray._private.utils.set_sigterm_handler(sigterm_handler)
|
||||
|
||||
loop.run_until_complete(task)
|
||||
loop.run_forever()
|
||||
@@ -0,0 +1,96 @@
|
||||
import collections
|
||||
import inspect
|
||||
|
||||
from ray.dashboard.optional_deps import aiohttp
|
||||
from ray.dashboard.routes import BaseRouteTable
|
||||
from ray.dashboard.subprocesses.handle import SubprocessModuleHandle
|
||||
from ray.dashboard.subprocesses.utils import ResponseType
|
||||
|
||||
|
||||
class SubprocessRouteTable(BaseRouteTable):
|
||||
"""
|
||||
A route table to bind http route to SubprocessModuleHandle and SubprocessModule.
|
||||
|
||||
This class is used in cls object: all the decorator methods are @classmethod, and
|
||||
the routes are binded to the cls object.
|
||||
|
||||
Note we have 2 handlers:
|
||||
1. the child side handler, that is `handler` that contains real logic and
|
||||
is executed in the child process. It's added with __route_method__ and
|
||||
__route_path__ attributes. The child process runs a standalone aiohttp
|
||||
server.
|
||||
2. the parent side handler, that just sends the request to the
|
||||
SubprocessModuleHandle at cls._bind_map[method][path].instance.
|
||||
|
||||
With modifications:
|
||||
- __route_method__ and __route_path__ are added to both side's handlers.
|
||||
- method and path are added to self._bind_map.
|
||||
|
||||
Lifecycle of a request:
|
||||
1. Parent receives an aiohttp request.
|
||||
2. Router finds by [method][path] and calls parent_side_handler.
|
||||
3. `parent_side_handler` bookkeeps the request with a Future and sends a
|
||||
request to the subprocess.
|
||||
4. Subprocesses receives the response and sends it back to the parent.
|
||||
5. Parent responds to the aiohttp request with the response from the subprocess.
|
||||
"""
|
||||
|
||||
_bind_map = collections.defaultdict(dict)
|
||||
_routes = aiohttp.web.RouteTableDef()
|
||||
|
||||
@classmethod
|
||||
def bind(cls, instance: "SubprocessModuleHandle"):
|
||||
# __route_method__ and __route_path__ are added to SubprocessModule's methods,
|
||||
# not the SubprocessModuleHandle's methods.
|
||||
def predicate(o):
|
||||
if inspect.isfunction(o):
|
||||
return hasattr(o, "__route_method__") and hasattr(o, "__route_path__")
|
||||
return False
|
||||
|
||||
handler_routes = inspect.getmembers(instance.module_cls, predicate)
|
||||
for _, h in handler_routes:
|
||||
cls._bind_map[h.__route_method__][h.__route_path__].instance = instance
|
||||
|
||||
@classmethod
|
||||
def _register_route(
|
||||
cls, method, path, resp_type: ResponseType = ResponseType.HTTP, **kwargs
|
||||
):
|
||||
"""
|
||||
Register a route to the module and return the decorated handler.
|
||||
"""
|
||||
|
||||
def _wrapper(handler):
|
||||
if path in cls._bind_map[method]:
|
||||
bind_info = cls._bind_map[method][path]
|
||||
raise Exception(
|
||||
f"Duplicated route path: {path}, "
|
||||
f"previous one registered at "
|
||||
f"{bind_info.filename}:{bind_info.lineno}"
|
||||
)
|
||||
|
||||
bind_info = cls._BindInfo(
|
||||
handler.__code__.co_filename, handler.__code__.co_firstlineno, None
|
||||
)
|
||||
|
||||
cls._bind_map[method][path] = bind_info
|
||||
|
||||
async def parent_side_handler(
|
||||
request: aiohttp.web.Request,
|
||||
) -> aiohttp.web.Response:
|
||||
bind_info = cls._bind_map[method][path]
|
||||
subprocess_module_handle = bind_info.instance
|
||||
return await subprocess_module_handle.proxy_request(request, resp_type)
|
||||
|
||||
# Used in bind().
|
||||
handler.__route_method__ = method
|
||||
handler.__route_path__ = path
|
||||
# Used in bound_routes().
|
||||
parent_side_handler.__route_method__ = method
|
||||
parent_side_handler.__route_path__ = path
|
||||
parent_side_handler.__name__ = handler.__name__
|
||||
|
||||
cls._routes.route(method, path)(parent_side_handler)
|
||||
|
||||
return handler
|
||||
|
||||
return _wrapper
|
||||
@@ -0,0 +1,489 @@
|
||||
import asyncio
|
||||
import pathlib
|
||||
import re
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
from typing import List
|
||||
|
||||
import pytest
|
||||
|
||||
import ray._private.ray_constants as ray_constants
|
||||
import ray.dashboard.consts as dashboard_consts
|
||||
from ray._common.ray_constants import (
|
||||
LOGGING_ROTATE_BACKUP_COUNT,
|
||||
LOGGING_ROTATE_BYTES,
|
||||
)
|
||||
from ray._common.test_utils import async_wait_for_condition, wait_for_condition
|
||||
from ray.dashboard.optional_deps import aiohttp
|
||||
from ray.dashboard.subprocesses.handle import SubprocessModuleHandle
|
||||
from ray.dashboard.subprocesses.module import SubprocessModule, SubprocessModuleConfig
|
||||
from ray.dashboard.subprocesses.routes import SubprocessRouteTable
|
||||
from ray.dashboard.subprocesses.tests.utils import TestModule, TestModule1
|
||||
|
||||
# This test requires non-minimal Ray.
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def default_module_config(tmp_path) -> SubprocessModuleConfig:
|
||||
"""
|
||||
Creates a tmpdir to hold the logs.
|
||||
"""
|
||||
# macOS caps AF_UNIX sun_path at 104 bytes. pytest's tmp_path on macOS is
|
||||
# rooted at /private/var/folders/... which is already long enough that
|
||||
# appending the socket filename exceeds the limit, so use a short directory
|
||||
# under /tmp for the socket files specifically.
|
||||
socket_dir = tempfile.mkdtemp(
|
||||
prefix="ray_e2e_", dir="/tmp" if sys.platform != "win32" else None
|
||||
)
|
||||
try:
|
||||
yield SubprocessModuleConfig(
|
||||
cluster_id_hex="test_cluster_id",
|
||||
gcs_address="",
|
||||
session_name="test_session",
|
||||
temp_dir=str(tmp_path),
|
||||
session_dir=str(tmp_path),
|
||||
logging_level=ray_constants.LOGGER_LEVEL,
|
||||
logging_format=ray_constants.LOGGER_FORMAT,
|
||||
log_dir=str(tmp_path),
|
||||
logging_filename=dashboard_consts.DASHBOARD_LOG_FILENAME,
|
||||
logging_rotate_bytes=LOGGING_ROTATE_BYTES,
|
||||
logging_rotate_backup_count=LOGGING_ROTATE_BACKUP_COUNT,
|
||||
socket_dir=socket_dir,
|
||||
)
|
||||
finally:
|
||||
shutil.rmtree(socket_dir, ignore_errors=True)
|
||||
|
||||
|
||||
class _DummyConn:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _DummyProcess:
|
||||
def __init__(self, alive: bool = True, pid: int = 12345):
|
||||
self._alive = alive
|
||||
self.pid = pid
|
||||
self.terminate_called = False
|
||||
self.kill_called = False
|
||||
self.join_called = False
|
||||
self.join_timeout = None
|
||||
|
||||
def is_alive(self) -> bool:
|
||||
return self._alive
|
||||
|
||||
def terminate(self):
|
||||
self.terminate_called = True
|
||||
# Simulate graceful exit.
|
||||
self._alive = False
|
||||
|
||||
def kill(self):
|
||||
self.kill_called = True
|
||||
self._alive = False
|
||||
|
||||
def join(self, timeout: float | None = None):
|
||||
self.join_called = True
|
||||
self.join_timeout = timeout
|
||||
|
||||
|
||||
class _DummySession:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_can_health_check(default_module_config):
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
subprocess_handle = SubprocessModuleHandle(loop, TestModule, default_module_config)
|
||||
subprocess_handle.start_module()
|
||||
subprocess_handle.wait_for_module_ready()
|
||||
response = await subprocess_handle._health_check()
|
||||
assert response.status == 200
|
||||
assert response.body == b"success"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_destroy_module_cleans_up_resources(default_module_config, monkeypatch):
|
||||
"""Ensure destroy_module closes connections, terminates process, and cancels tasks."""
|
||||
loop = asyncio.get_event_loop()
|
||||
handle = SubprocessModuleHandle(loop, TestModule, default_module_config)
|
||||
|
||||
# Inject dummy resources.
|
||||
parent_conn = _DummyConn()
|
||||
process = _DummyProcess(alive=True)
|
||||
http_session = _DummySession()
|
||||
health_task = asyncio.create_task(asyncio.sleep(1000))
|
||||
|
||||
handle.parent_conn = parent_conn
|
||||
handle.process = process
|
||||
handle.http_client_session = http_session
|
||||
handle.health_check_task = health_task
|
||||
|
||||
await handle.destroy_module()
|
||||
|
||||
# Parent connection is closed and cleared.
|
||||
assert parent_conn.closed
|
||||
assert handle.parent_conn is None
|
||||
|
||||
# Process was terminated gracefully and cleared.
|
||||
assert process.terminate_called
|
||||
assert process.join_called
|
||||
assert not process.kill_called
|
||||
assert handle.process is None
|
||||
|
||||
# HTTP client session is closed and cleared.
|
||||
assert http_session.closed
|
||||
assert handle.http_client_session is None
|
||||
|
||||
# Health check task is cancelled and cleared.
|
||||
# Note: destroy_module() may call cancel() without awaiting the task, so the
|
||||
# cancelled state might not be observable immediately. Wait briefly for the
|
||||
# cancellation to be delivered.
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(health_task, timeout=0.2)
|
||||
assert health_task.cancelled()
|
||||
assert handle.health_check_task is None
|
||||
|
||||
|
||||
async def start_http_server_app(
|
||||
default_module_config, modules: List[type(SubprocessModule)]
|
||||
):
|
||||
loop = asyncio.get_event_loop()
|
||||
handles = [
|
||||
SubprocessModuleHandle(loop, module, default_module_config)
|
||||
for module in modules
|
||||
]
|
||||
# Parallel start all modules.
|
||||
for handle in handles:
|
||||
handle.start_module()
|
||||
# Wait for all modules to be ready.
|
||||
for handle in handles:
|
||||
handle.wait_for_module_ready()
|
||||
SubprocessRouteTable.bind(handle)
|
||||
|
||||
app = aiohttp.web.Application()
|
||||
app.add_routes(routes=SubprocessRouteTable.bound_routes())
|
||||
|
||||
async def _cleanup_app(app):
|
||||
# Best-effort cleanup to avoid leaking subprocesses across the pytest session.
|
||||
# If subprocesses leak, later tests (e.g. starting a Ray cluster) can fail due
|
||||
# to process/FD exhaustion and manifest as unrelated startup errors.
|
||||
for handle in handles:
|
||||
try:
|
||||
await handle.destroy_module()
|
||||
except Exception:
|
||||
# Don't mask original test failures.
|
||||
pass
|
||||
|
||||
# SubprocessRouteTable is a global singleton route table. Ensure we drop strong
|
||||
# references to handles so they can be GC'd.
|
||||
for _, path_map in list(SubprocessRouteTable._bind_map.items()):
|
||||
for _, bind_info in list(path_map.items()):
|
||||
try:
|
||||
bind_info.instance = None
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
app.on_cleanup.append(_cleanup_app)
|
||||
return app
|
||||
|
||||
|
||||
# @pytest.mark.asyncio is not compatible with aiohttp_client.
|
||||
async def test_http_server(aiohttp_client, default_module_config):
|
||||
"""
|
||||
Tests that the http server works. It must
|
||||
1. bind a SubprocessModuleHandle
|
||||
2. add_routes
|
||||
3. run
|
||||
"""
|
||||
app = await start_http_server_app(default_module_config, [TestModule])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
# Test HTTP
|
||||
response = await client.get("/test")
|
||||
assert response.status == 200
|
||||
assert await response.text() == "Hello, World from GET /test, run_finished: True"
|
||||
|
||||
response = await client.post("/echo", data="a new dashboard")
|
||||
assert response.status == 200
|
||||
assert await response.text() == "Hello, World from POST /echo from a new dashboard"
|
||||
|
||||
response = await client.put("/error")
|
||||
assert response.status == 500
|
||||
assert "Internal Server Error" in await response.text()
|
||||
|
||||
response = await client.put("/error_403")
|
||||
assert response.status == 403
|
||||
assert "you shall not pass" in await response.text()
|
||||
|
||||
|
||||
async def test_load_multiple_modules(aiohttp_client, default_module_config):
|
||||
"""
|
||||
Tests multiple modules can be loaded.
|
||||
"""
|
||||
app = await start_http_server_app(default_module_config, [TestModule, TestModule1])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
response = await client.get("/test")
|
||||
assert response.status == 200
|
||||
assert await response.text() == "Hello, World from GET /test, run_finished: True"
|
||||
|
||||
response = await client.get("/test1")
|
||||
assert response.status == 200
|
||||
assert await response.text() == "Hello from TestModule1"
|
||||
|
||||
|
||||
async def test_redirect_between_modules(aiohttp_client, default_module_config):
|
||||
"""Tests that a redirect can be handled between modules."""
|
||||
app = await start_http_server_app(default_module_config, [TestModule, TestModule1])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
# Allow redirects to be handled between modules.
|
||||
# NOTE: If redirects were followed at the module level,
|
||||
# the test would error, since following /test in TestModule1 would
|
||||
# result in a 404.
|
||||
# Instead, the redirect should be handled at the subprocess proxy level,
|
||||
# where the redirect request is forwarded to the correct module.
|
||||
response = await client.get("/redirect_between_modules", allow_redirects=True)
|
||||
assert response.status == 200
|
||||
assert await response.text() == "Hello, World from GET /test, run_finished: True"
|
||||
|
||||
response = await client.get("/redirect_within_module", allow_redirects=True)
|
||||
assert response.status == 200
|
||||
assert await response.text() == "Hello from TestModule1"
|
||||
|
||||
|
||||
async def test_cached_endpoint(aiohttp_client, default_module_config):
|
||||
"""
|
||||
Test whether the ray.dashboard.optional_utils.aiohttp_cache decorator works.
|
||||
"""
|
||||
app = await start_http_server_app(default_module_config, [TestModule])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
response = await client.get("/not_cached")
|
||||
assert response.status == 200
|
||||
assert await response.text() == "Hello, World from GET /not_cached, count: 1"
|
||||
|
||||
# Call again, count should increase.
|
||||
response = await client.get("/not_cached")
|
||||
assert response.status == 200
|
||||
assert await response.text() == "Hello, World from GET /not_cached, count: 2"
|
||||
|
||||
response = await client.get("/cached")
|
||||
assert response.status == 200
|
||||
assert await response.text() == "Hello, World from GET /cached, count: 1"
|
||||
|
||||
# Call again, count should NOT increase.
|
||||
response = await client.get("/cached")
|
||||
assert response.status == 200
|
||||
assert await response.text() == "Hello, World from GET /cached, count: 1"
|
||||
|
||||
|
||||
async def test_streamed_iota(aiohttp_client, default_module_config):
|
||||
# TODO(ryw): also test streams that raise exceptions.
|
||||
app = await start_http_server_app(default_module_config, [TestModule])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
response = await client.post("/streamed_iota", data=b"10")
|
||||
assert response.status == 200
|
||||
assert await response.text() == "0\n1\n2\n3\n4\n5\n6\n7\n8\n9\n"
|
||||
|
||||
|
||||
async def test_streamed_error(aiohttp_client, default_module_config):
|
||||
app = await start_http_server_app(default_module_config, [TestModule])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
response = await client.post("/streamed_401", data=b"")
|
||||
assert response.status == 401
|
||||
assert await response.text() == "401: Unauthorized although I am not a teapot"
|
||||
|
||||
|
||||
async def test_websocket_bytes_res(aiohttp_client, default_module_config):
|
||||
app = await start_http_server_app(default_module_config, [TestModule])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
res = []
|
||||
async with client.ws_connect("/websocket_one_to_five_bytes") as ws:
|
||||
async for msg in ws:
|
||||
assert msg.type == aiohttp.WSMsgType.BINARY
|
||||
res.append(msg.data)
|
||||
assert res == [b"1\n", b"2\n", b"3\n", b"4\n", b"5\n"]
|
||||
|
||||
|
||||
async def test_websocket_bytes_str(aiohttp_client, default_module_config):
|
||||
app = await start_http_server_app(default_module_config, [TestModule])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
res = []
|
||||
async with client.ws_connect("/websocket_one_to_five_strs") as ws:
|
||||
async for msg in ws:
|
||||
assert msg.type == aiohttp.WSMsgType.TEXT
|
||||
res.append(msg.data)
|
||||
assert res == ["1\n", "2\n", "3\n", "4\n", "5\n"]
|
||||
|
||||
|
||||
async def test_websocket_raise_http_error(aiohttp_client, default_module_config):
|
||||
app = await start_http_server_app(default_module_config, [TestModule])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
response = await client.get("/websocket_raise_http_error")
|
||||
assert response.status == 400
|
||||
assert await response.text() == "400: Hello this is a bad request"
|
||||
|
||||
|
||||
async def test_websocket_raise_non_http_error(aiohttp_client, default_module_config):
|
||||
app = await start_http_server_app(default_module_config, [TestModule])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
response = await client.get("/websocket_raise_non_http_error")
|
||||
assert response.status == 500
|
||||
|
||||
|
||||
async def test_kill_self(aiohttp_client, default_module_config):
|
||||
"""
|
||||
If a module died, all pending requests should be failed, and the module should be
|
||||
restarted. After the restart, subsequent requests should be successful.
|
||||
"""
|
||||
app = await start_http_server_app(default_module_config, [TestModule])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
long_running_task = asyncio.create_task(client.post("/run_forever", data=b""))
|
||||
# Wait for 1s for the long running request to start.
|
||||
await asyncio.sleep(1)
|
||||
|
||||
response = await client.post("/kill_self", data=b"")
|
||||
assert response.status == 500
|
||||
assert (
|
||||
await response.text()
|
||||
== "500 Internal Server Error\n\nServer got itself in trouble"
|
||||
)
|
||||
|
||||
# Long running request should get a 500.
|
||||
long_running_response = await long_running_task
|
||||
assert long_running_response.status == 500
|
||||
assert (
|
||||
await long_running_response.text()
|
||||
== "500 Internal Server Error\n\nServer got itself in trouble"
|
||||
)
|
||||
|
||||
async def verify():
|
||||
response = await client.post(
|
||||
"/echo", data=b"a restarted dashboard", timeout=0.5
|
||||
)
|
||||
assert response.status == 200
|
||||
assert (
|
||||
await response.text()
|
||||
== "Hello, World from POST /echo from a restarted dashboard"
|
||||
)
|
||||
return True
|
||||
|
||||
await async_wait_for_condition(verify)
|
||||
|
||||
|
||||
async def test_logging_in_module(aiohttp_client, default_module_config):
|
||||
app = await start_http_server_app(default_module_config, [TestModule])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
def read_file_content(file_path):
|
||||
with file_path.open("r") as f:
|
||||
return f.read()
|
||||
|
||||
response = await client.post(
|
||||
"/logging_in_module", data=b"Not all those who wander are lost"
|
||||
)
|
||||
assert response.status == 200
|
||||
assert await response.text() == "done!"
|
||||
|
||||
# Assert the log file name and read the log file
|
||||
log_file_path = (
|
||||
pathlib.Path(default_module_config.log_dir) / "dashboard_TestModule.log"
|
||||
)
|
||||
|
||||
def verify():
|
||||
with log_file_path.open("r") as f:
|
||||
log_file_content = f.read()
|
||||
|
||||
# Assert on the log format and the content.
|
||||
log_pattern = (
|
||||
r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2},\d{3}\tINFO ([\w\.]+):\d+ -- (.*)"
|
||||
)
|
||||
matches = re.findall(log_pattern, log_file_content)
|
||||
|
||||
# Expected format: [(file_name, content), ...]
|
||||
expected_logs = [
|
||||
("utils.py", "TestModule is initing"),
|
||||
("utils.py", "TestModule is done initing"),
|
||||
("utils.py", "In /logging_in_module, Not all those who wander are lost."),
|
||||
]
|
||||
return all(
|
||||
(file_name, content) in matches for (file_name, content) in expected_logs
|
||||
)
|
||||
|
||||
wait_for_condition(verify)
|
||||
|
||||
# Assert that stdout is logged to "dashboard_TestModule.out"
|
||||
out_log_file_path = (
|
||||
pathlib.Path(default_module_config.log_dir) / "dashboard_TestModule.out"
|
||||
)
|
||||
wait_for_condition(
|
||||
lambda: read_file_content(out_log_file_path)
|
||||
== "In /logging_in_module, stdout\n"
|
||||
)
|
||||
|
||||
# Assert that stderr is logged to "dashboard_TestModule.err"
|
||||
err_log_file_path = (
|
||||
pathlib.Path(default_module_config.log_dir) / "dashboard_TestModule.err"
|
||||
)
|
||||
wait_for_condition(
|
||||
lambda: read_file_content(err_log_file_path)
|
||||
== "In /logging_in_module, stderr\n"
|
||||
)
|
||||
|
||||
|
||||
async def test_logging_in_module_with_multiple_incarnations(
|
||||
aiohttp_client, default_module_config
|
||||
):
|
||||
app = await start_http_server_app(default_module_config, [TestModule])
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
response = await client.post(
|
||||
"/logging_in_module", data=b"this is from incarnation 0"
|
||||
)
|
||||
assert response.status == 200
|
||||
assert await response.text() == "done!"
|
||||
|
||||
response = await client.post("/kill_self", data=b"")
|
||||
assert response.status == 500
|
||||
assert (
|
||||
await response.text()
|
||||
== "500 Internal Server Error\n\nServer got itself in trouble"
|
||||
)
|
||||
|
||||
async def verify():
|
||||
response = await client.post(
|
||||
"/logging_in_module", data=b"and this is from incarnation 1"
|
||||
)
|
||||
assert response.status == 200
|
||||
assert await response.text() == "done!"
|
||||
return True
|
||||
|
||||
await async_wait_for_condition(verify)
|
||||
|
||||
log_file_path = (
|
||||
pathlib.Path(default_module_config.log_dir) / "dashboard_TestModule.log"
|
||||
)
|
||||
with log_file_path.open("r") as f:
|
||||
log_file_content = f.read()
|
||||
assert "In /logging_in_module, this is from incarnation 0." in log_file_content
|
||||
assert "In /logging_in_module, and this is from incarnation 1." in log_file_content
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,179 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
from typing import AsyncIterator
|
||||
|
||||
from ray.dashboard import optional_utils
|
||||
from ray.dashboard.optional_deps import aiohttp
|
||||
from ray.dashboard.subprocesses.module import SubprocessModule
|
||||
from ray.dashboard.subprocesses.routes import SubprocessRouteTable as routes
|
||||
from ray.dashboard.subprocesses.utils import ResponseType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BaseTestModule(SubprocessModule):
|
||||
@property
|
||||
def gcs_client(self):
|
||||
return None
|
||||
|
||||
@property
|
||||
def aiogrpc_gcs_channel(self):
|
||||
return None
|
||||
|
||||
|
||||
class TestModule(BaseTestModule):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.run_finished = False
|
||||
self.not_cached_count = 0
|
||||
self.cached_count = 0
|
||||
|
||||
async def run(self):
|
||||
await super().run()
|
||||
logger.info("TestModule is initing")
|
||||
self.run_finished = True
|
||||
await asyncio.sleep(0.1)
|
||||
logger.info("TestModule is done initing")
|
||||
|
||||
@routes.get("/test")
|
||||
async def test(self, req: aiohttp.web.Request) -> aiohttp.web.Response:
|
||||
return aiohttp.web.Response(
|
||||
text="Hello, World from GET /test, run_finished: " + str(self.run_finished)
|
||||
)
|
||||
|
||||
@routes.post("/echo")
|
||||
async def echo(self, req: aiohttp.web.Request) -> aiohttp.web.Response:
|
||||
# await works
|
||||
await asyncio.sleep(0.1)
|
||||
body = await req.text()
|
||||
return aiohttp.web.Response(text="Hello, World from POST /echo from " + body)
|
||||
|
||||
@routes.get("/not_cached")
|
||||
async def not_cached(self, req: aiohttp.web.Request) -> aiohttp.web.Response:
|
||||
self.not_cached_count += 1
|
||||
return aiohttp.web.Response(
|
||||
text=f"Hello, World from GET /not_cached, count: {self.not_cached_count}"
|
||||
)
|
||||
|
||||
@routes.get("/cached")
|
||||
@optional_utils.aiohttp_cache
|
||||
async def cached(self, req: aiohttp.web.Request) -> aiohttp.web.Response:
|
||||
self.cached_count += 1
|
||||
return aiohttp.web.Response(
|
||||
text=f"Hello, World from GET /cached, count: {self.cached_count}"
|
||||
)
|
||||
|
||||
@routes.put("/error")
|
||||
async def make_error(self, req: aiohttp.web.Request) -> aiohttp.web.Response:
|
||||
raise ValueError("This is an error")
|
||||
|
||||
@routes.put("/error_403")
|
||||
async def make_error_403(self, req: aiohttp.web.Request) -> aiohttp.web.Response:
|
||||
# For an ascii art of Gandalf the Grey, see:
|
||||
# https://github.com/ray-project/ray/pull/49732#discussion_r1919292428
|
||||
raise aiohttp.web.HTTPForbidden(reason="you shall not pass")
|
||||
|
||||
@routes.post("/streamed_iota", resp_type=ResponseType.STREAM)
|
||||
async def streamed_iota(
|
||||
self, req: aiohttp.web.Request
|
||||
) -> aiohttp.web.StreamResponse:
|
||||
"""
|
||||
Streams the numbers 0 to N.
|
||||
"""
|
||||
request_body = await req.text()
|
||||
n = int(request_body)
|
||||
resp = aiohttp.web.StreamResponse()
|
||||
await resp.prepare(req)
|
||||
for i in range(n):
|
||||
await asyncio.sleep(0.001)
|
||||
await resp.write(f"{i}\n".encode())
|
||||
await resp.write_eof()
|
||||
return resp
|
||||
|
||||
@routes.post("/streamed_401", resp_type=ResponseType.STREAM)
|
||||
async def streamed_401(self, req: aiohttp.web.Request) -> AsyncIterator[bytes]:
|
||||
"""
|
||||
Raise a 401 error instead of streaming.
|
||||
"""
|
||||
raise aiohttp.web.HTTPUnauthorized(
|
||||
reason="Unauthorized although I am not a teapot"
|
||||
)
|
||||
|
||||
@routes.get("/websocket_one_to_five_bytes", resp_type=ResponseType.WEBSOCKET)
|
||||
async def websocket_one_to_five_bytes(
|
||||
self, req: aiohttp.web.Request
|
||||
) -> aiohttp.web.WebSocketResponse:
|
||||
ws = aiohttp.web.WebSocketResponse()
|
||||
await ws.prepare(req)
|
||||
for i in range(1, 6):
|
||||
await asyncio.sleep(0.001)
|
||||
await ws.send_bytes(f"{i}\n".encode())
|
||||
await ws.close()
|
||||
return ws
|
||||
|
||||
@routes.get("/websocket_one_to_five_strs", resp_type=ResponseType.WEBSOCKET)
|
||||
async def websocket_one_to_five_strs(
|
||||
self, req: aiohttp.web.Request
|
||||
) -> aiohttp.web.WebSocketResponse:
|
||||
ws = aiohttp.web.WebSocketResponse()
|
||||
await ws.prepare(req)
|
||||
for i in range(1, 6):
|
||||
await asyncio.sleep(0.001)
|
||||
await ws.send_str(f"{i}\n")
|
||||
await ws.close()
|
||||
return ws
|
||||
|
||||
@routes.get("/websocket_raise_http_error", resp_type=ResponseType.WEBSOCKET)
|
||||
async def websocket_raise_http_error(
|
||||
self, req: aiohttp.web.Request
|
||||
) -> aiohttp.web.WebSocketResponse:
|
||||
raise aiohttp.web.HTTPBadRequest(reason="Hello this is a bad request")
|
||||
|
||||
@routes.get("/websocket_raise_non_http_error", resp_type=ResponseType.WEBSOCKET)
|
||||
async def websocket_raise_non_http_error(
|
||||
self, req: aiohttp.web.Request
|
||||
) -> aiohttp.web.WebSocketResponse:
|
||||
raise ValueError("Hello world")
|
||||
|
||||
@routes.post("/run_forever")
|
||||
async def run_forever(self, req: aiohttp.web.Request) -> aiohttp.web.Response:
|
||||
while True:
|
||||
await asyncio.sleep(1)
|
||||
return aiohttp.web.Response(text="done in the infinite future!")
|
||||
|
||||
@routes.post("/logging_in_module")
|
||||
async def logging_in_module(self, req: aiohttp.web.Request) -> aiohttp.web.Response:
|
||||
request_body_str = await req.text()
|
||||
logger.info(f"In /logging_in_module, {request_body_str}.")
|
||||
print("In /logging_in_module, stdout")
|
||||
print("In /logging_in_module, stderr", file=sys.stderr)
|
||||
return aiohttp.web.Response(text="done!")
|
||||
|
||||
@routes.post("/kill_self")
|
||||
async def kill_self(self, req: aiohttp.web.Request) -> aiohttp.web.Response:
|
||||
logger.error("Crashing by sending myself a sigkill")
|
||||
os.kill(os.getpid(), signal.SIGKILL)
|
||||
asyncio.sleep(1000)
|
||||
return aiohttp.web.Response(text="done!")
|
||||
|
||||
|
||||
class TestModule1(BaseTestModule):
|
||||
@routes.get("/test1")
|
||||
async def test(self, req: aiohttp.web.Request) -> aiohttp.web.Response:
|
||||
return aiohttp.web.Response(text="Hello from TestModule1")
|
||||
|
||||
@routes.get("/redirect_between_modules")
|
||||
async def redirect_between_modules(
|
||||
self, req: aiohttp.web.Request
|
||||
) -> aiohttp.web.Response:
|
||||
# Redirect to the /test route in TestModule
|
||||
raise aiohttp.web.HTTPFound(location="/test")
|
||||
|
||||
@routes.get("/redirect_within_module")
|
||||
async def redirect_within_module(
|
||||
self, req: aiohttp.web.Request
|
||||
) -> aiohttp.web.Response:
|
||||
raise aiohttp.web.HTTPFound(location="/test1")
|
||||
@@ -0,0 +1,75 @@
|
||||
import enum
|
||||
import os
|
||||
import sys
|
||||
from typing import TypeVar
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ray._private.utils import validate_socket_filepath
|
||||
|
||||
K = TypeVar("K")
|
||||
V = TypeVar("V")
|
||||
|
||||
|
||||
class ResponseType(enum.Enum):
|
||||
HTTP = "http"
|
||||
STREAM = "stream"
|
||||
WEBSOCKET = "websocket"
|
||||
|
||||
|
||||
def module_logging_filename(
|
||||
module_name: str, logging_filename: str, extension: str = ""
|
||||
) -> str:
|
||||
"""Parse logging_filename = STEM EXTENSION, return STEM _ MODULE_NAME _ EXTENSION.
|
||||
|
||||
If logging_filename is empty, return empty string.
|
||||
If extension is empty, use the extension from logging_filename.
|
||||
|
||||
Args:
|
||||
module_name: Name of the subprocess module, embedded in the output stem.
|
||||
logging_filename: Original log filename (e.g. ``dashboard.log``).
|
||||
extension: Override extension. Defaults to the extension of
|
||||
``logging_filename``.
|
||||
|
||||
Returns:
|
||||
The new filename with the module name embedded between the stem and
|
||||
extension, or an empty string when ``logging_filename`` is empty.
|
||||
|
||||
Example:
|
||||
module_name = "TestModule"
|
||||
logging_filename = "dashboard.log"
|
||||
STEM = "dashboard"
|
||||
EXTENSION = ".log"
|
||||
return "dashboard_TestModule.log"
|
||||
"""
|
||||
if not logging_filename:
|
||||
return ""
|
||||
stem, ext = os.path.splitext(logging_filename)
|
||||
if not extension:
|
||||
extension = ext
|
||||
return f"{stem}_{module_name}{extension}"
|
||||
|
||||
|
||||
def get_socket_path(socket_dir: str, module_name: str) -> str:
|
||||
socket_path = os.path.join(socket_dir, "dash_" + module_name)
|
||||
validate_socket_filepath(socket_path)
|
||||
return socket_path
|
||||
|
||||
|
||||
def get_named_pipe_path(module_name: str, session_name: str) -> str:
|
||||
return r"\\.\pipe\dash_" + module_name + "_" + session_name
|
||||
|
||||
|
||||
def get_http_session_to_module(
|
||||
module_name: str, socket_dir: str, session_name: str
|
||||
) -> aiohttp.ClientSession:
|
||||
"""
|
||||
Get the aiohttp http client session to the subprocess module.
|
||||
"""
|
||||
if sys.platform == "win32":
|
||||
named_pipe_path = get_named_pipe_path(module_name, session_name)
|
||||
connector = aiohttp.NamedPipeConnector(named_pipe_path)
|
||||
else:
|
||||
socket_path = get_socket_path(socket_dir, module_name)
|
||||
connector = aiohttp.UnixConnector(socket_path)
|
||||
return aiohttp.ClientSession(connector=connector)
|
||||
Reference in New Issue
Block a user