chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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()
|
||||
Reference in New Issue
Block a user