wehub-resource-sync
292 lines
9.0 KiB
Python
292 lines
9.0 KiB
Python
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()
|