wehub-resource-sync
485 lines
21 KiB
Python
485 lines
21 KiB
Python
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import traceback
|
|
from asyncio.tasks import FIRST_COMPLETED
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import ray
|
|
import ray._private.ray_constants as ray_constants
|
|
from ray._common.filters import CoreContextFilter
|
|
from ray._common.formatters import JSONFormatter, TextFormatter
|
|
from ray._common.network_utils import build_address
|
|
from ray._private.accelerators.npu import NOSET_ASCEND_RT_VISIBLE_DEVICES_ENV_VAR
|
|
from ray._private.accelerators.nvidia_gpu import NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR
|
|
from ray._private.runtime_env.constants import RAY_JOB_CONFIG_JSON_ENV_VAR
|
|
from ray._private.utils import remove_ray_internal_flags_from_env
|
|
from ray._raylet import GcsClient
|
|
from ray.actor import ActorHandle
|
|
from ray.dashboard.modules.job.common import (
|
|
JOB_ID_METADATA_KEY,
|
|
JOB_NAME_METADATA_KEY,
|
|
JobInfoStorageClient,
|
|
)
|
|
from ray.dashboard.modules.job.job_log_storage_client import JobLogStorageClient
|
|
from ray.job_submission import JobErrorType, JobStatus
|
|
|
|
import psutil
|
|
|
|
# asyncio python version compatibility
|
|
try:
|
|
create_task = asyncio.create_task
|
|
except AttributeError:
|
|
create_task = asyncio.ensure_future
|
|
|
|
# Windows requires additional packages for proper process control.
|
|
if sys.platform == "win32":
|
|
try:
|
|
import win32api
|
|
import win32con
|
|
import win32job
|
|
except (ModuleNotFoundError, ImportError) as e:
|
|
win32api = None
|
|
win32con = None
|
|
win32job = None
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logger.warning(
|
|
"Failed to Import win32api. For best usage experience run "
|
|
f"'conda install pywin32'. Import error: {e}"
|
|
)
|
|
|
|
|
|
class JobSupervisor:
|
|
"""
|
|
Ray actor created by JobManager for each submitted job, responsible to
|
|
setup runtime_env, execute given shell command in subprocess, update job
|
|
status, persist job logs and manage subprocess group cleaning.
|
|
|
|
One job supervisor actor maps to one subprocess, for one job_id.
|
|
Job supervisor actor should fate share with subprocess it created.
|
|
"""
|
|
|
|
DEFAULT_RAY_JOB_STOP_WAIT_TIME_S = 3
|
|
SUBPROCESS_POLL_PERIOD_S = 0.1
|
|
VALID_STOP_SIGNALS = ["SIGINT", "SIGTERM"]
|
|
|
|
def __init__(
|
|
self,
|
|
job_id: str,
|
|
entrypoint: str,
|
|
user_metadata: Dict[str, str],
|
|
gcs_address: str,
|
|
cluster_id_hex: str,
|
|
logs_dir: Optional[str] = None,
|
|
):
|
|
self._job_id = job_id
|
|
gcs_client = GcsClient(address=gcs_address, cluster_id=cluster_id_hex)
|
|
self._job_info_client = JobInfoStorageClient(gcs_client, logs_dir)
|
|
self._log_client = JobLogStorageClient()
|
|
self._entrypoint = entrypoint
|
|
|
|
# Default metadata if not passed by the user.
|
|
self._metadata = {JOB_ID_METADATA_KEY: job_id, JOB_NAME_METADATA_KEY: job_id}
|
|
self._metadata.update(user_metadata)
|
|
|
|
# Event used to signal that a job should be stopped.
|
|
# Set in the `stop_job` method.
|
|
self._stop_event = asyncio.Event()
|
|
|
|
# Windows Job Object used to handle stopping the child processes.
|
|
self._win32_job_object = None
|
|
|
|
# Logger object to persist JobSupervisor logs in separate file.
|
|
self._logger = logging.getLogger(f"{__name__}.supervisor-{job_id}")
|
|
self._configure_logger()
|
|
|
|
def _configure_logger(self) -> None:
|
|
"""
|
|
Configure self._logger object to write logs to file based on job
|
|
submission ID and to console.
|
|
"""
|
|
supervisor_log_file_name = os.path.join(
|
|
ray._private.worker._global_node.get_logs_dir_path(),
|
|
f"jobs/supervisor-{self._job_id}.log",
|
|
)
|
|
os.makedirs(os.path.dirname(supervisor_log_file_name), exist_ok=True)
|
|
self._logger.addFilter(CoreContextFilter())
|
|
stream_handler = logging.StreamHandler()
|
|
file_handler = logging.FileHandler(supervisor_log_file_name)
|
|
formatter = TextFormatter()
|
|
if ray_constants.env_bool(ray_constants.RAY_BACKEND_LOG_JSON_ENV_VAR, False):
|
|
formatter = JSONFormatter()
|
|
stream_handler.setFormatter(formatter)
|
|
file_handler.setFormatter(formatter)
|
|
self._logger.addHandler(stream_handler)
|
|
self._logger.addHandler(file_handler)
|
|
self._logger.propagate = False
|
|
|
|
def _get_driver_runtime_env(
|
|
self, resources_specified: bool = False
|
|
) -> Dict[str, Any]:
|
|
"""Get the runtime env that should be set in the job driver.
|
|
|
|
Args:
|
|
resources_specified: Whether the user specified resources (CPUs, GPUs,
|
|
custom resources) in the submit_job request. If so, we will skip
|
|
the workaround for GPU detection introduced in #24546, so that the
|
|
behavior matches that of the user specifying resources for any
|
|
other actor.
|
|
|
|
Returns:
|
|
The runtime env that should be set in the job driver.
|
|
"""
|
|
# Get the runtime_env set for the supervisor actor.
|
|
curr_runtime_env = dict(ray.get_runtime_context().runtime_env)
|
|
if resources_specified:
|
|
return curr_runtime_env
|
|
# Allow CUDA_VISIBLE_DEVICES to be set normally for the driver's tasks
|
|
# & actors.
|
|
env_vars = curr_runtime_env.get("env_vars", {})
|
|
env_vars.pop(NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR)
|
|
env_vars.pop(NOSET_ASCEND_RT_VISIBLE_DEVICES_ENV_VAR)
|
|
env_vars.pop(ray_constants.RAY_WORKER_NICENESS)
|
|
curr_runtime_env["env_vars"] = env_vars
|
|
return curr_runtime_env
|
|
|
|
def ping(self):
|
|
"""Used to check the health of the actor."""
|
|
pass
|
|
|
|
def _exec_entrypoint(self, env: dict, logs_path: str) -> subprocess.Popen:
|
|
"""
|
|
Runs the entrypoint command as a child process, streaming stderr &
|
|
stdout to given log files.
|
|
|
|
Unix systems:
|
|
Meanwhile we start a demon process and group driver
|
|
subprocess in same pgid, such that if job actor dies, entire process
|
|
group also fate share with it.
|
|
|
|
Windows systems:
|
|
A jobObject is created to enable fate sharing for the entire process group.
|
|
|
|
Args:
|
|
env: Environment variables passed through to the driver subprocess.
|
|
logs_path: File path on head node's local disk to store driver
|
|
command's stdout & stderr.
|
|
Returns:
|
|
child_process: Child process that runs the driver command. Can be
|
|
terminated or killed upon user calling stop().
|
|
"""
|
|
# Open in append mode to avoid overwriting runtime_env setup logs for the
|
|
# supervisor actor, which are also written to the same file.
|
|
with open(logs_path, "a") as logs_file:
|
|
logs_file.write(
|
|
f"Running entrypoint for job {self._job_id}: {self._entrypoint}\n"
|
|
)
|
|
child_process = subprocess.Popen(
|
|
self._entrypoint,
|
|
shell=True,
|
|
start_new_session=True,
|
|
stdout=logs_file,
|
|
stderr=subprocess.STDOUT,
|
|
env=env,
|
|
# Ray intentionally blocks SIGINT in all processes, so if the user wants
|
|
# to stop job through SIGINT, we need to unblock it in the child process
|
|
preexec_fn=(
|
|
(
|
|
lambda: signal.pthread_sigmask(
|
|
signal.SIG_UNBLOCK, {signal.SIGINT}
|
|
)
|
|
)
|
|
if sys.platform != "win32"
|
|
and os.environ.get("RAY_JOB_STOP_SIGNAL") == "SIGINT"
|
|
else None
|
|
),
|
|
)
|
|
parent_pid = os.getpid()
|
|
child_pid = child_process.pid
|
|
# Create new pgid with new subprocess to execute driver command
|
|
|
|
if sys.platform != "win32":
|
|
try:
|
|
child_pgid = os.getpgid(child_pid)
|
|
except ProcessLookupError:
|
|
# Process died before we could get its pgid.
|
|
return child_process
|
|
|
|
# Open a new subprocess to kill the child process when the parent
|
|
# process dies kill -s 0 parent_pid will succeed if the parent is
|
|
# alive. If it fails, SIGKILL the child process group and exit
|
|
subprocess.Popen(
|
|
f"while kill -s 0 {parent_pid}; do sleep 1; done; kill -9 -{child_pgid}", # noqa: E501
|
|
shell=True,
|
|
# Suppress output
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
)
|
|
|
|
elif sys.platform == "win32" and win32api:
|
|
# Create a JobObject to which the child process (and its children)
|
|
# will be connected. This job object can be used to kill the child
|
|
# processes explicitly or when the jobObject gets deleted during
|
|
# garbage collection.
|
|
self._win32_job_object = win32job.CreateJobObject(None, "")
|
|
win32_job_info = win32job.QueryInformationJobObject(
|
|
self._win32_job_object, win32job.JobObjectExtendedLimitInformation
|
|
)
|
|
win32_job_info["BasicLimitInformation"][
|
|
"LimitFlags"
|
|
] = win32job.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
|
|
win32job.SetInformationJobObject(
|
|
self._win32_job_object,
|
|
win32job.JobObjectExtendedLimitInformation,
|
|
win32_job_info,
|
|
)
|
|
child_handle = win32api.OpenProcess(
|
|
win32con.PROCESS_TERMINATE | win32con.PROCESS_SET_QUOTA,
|
|
False,
|
|
child_pid,
|
|
)
|
|
win32job.AssignProcessToJobObject(self._win32_job_object, child_handle)
|
|
|
|
return child_process
|
|
|
|
def _get_driver_env_vars(self, resources_specified: bool) -> Dict[str, str]:
|
|
"""Returns environment variables that should be set in the driver."""
|
|
# RAY_ADDRESS may be the dashboard URL but not the gcs address,
|
|
# so when the environment variable is not empty, we force set RAY_ADDRESS
|
|
# to "auto" to avoid function `canonicalize_bootstrap_address_or_die` returning
|
|
# the wrong GCS address.
|
|
# TODO(Jialing He, Archit Kulkarni): Definition of Specification RAY_ADDRESS
|
|
if ray_constants.RAY_ADDRESS_ENVIRONMENT_VARIABLE in os.environ:
|
|
os.environ[ray_constants.RAY_ADDRESS_ENVIRONMENT_VARIABLE] = "auto"
|
|
ray_addr = ray._private.services.canonicalize_bootstrap_address_or_die(
|
|
"auto", ray._private.worker._global_node._ray_params.temp_dir
|
|
)
|
|
assert ray_addr is not None
|
|
return {
|
|
# Set JobConfig for the child process (runtime_env, metadata).
|
|
RAY_JOB_CONFIG_JSON_ENV_VAR: json.dumps(
|
|
{
|
|
"runtime_env": self._get_driver_runtime_env(resources_specified),
|
|
"metadata": self._metadata,
|
|
}
|
|
),
|
|
# Always set RAY_ADDRESS as find_bootstrap_address address for
|
|
# job submission. In case of local development, prevent user from
|
|
# re-using http://{address}:{dashboard_port} to interact with
|
|
# jobs SDK.
|
|
# TODO:(mwtian) Check why "auto" does not work in entrypoint script
|
|
ray_constants.RAY_ADDRESS_ENVIRONMENT_VARIABLE: ray_addr,
|
|
# Set PYTHONUNBUFFERED=1 to stream logs during the job instead of
|
|
# only streaming them upon completion of the job.
|
|
"PYTHONUNBUFFERED": "1",
|
|
}
|
|
|
|
async def _polling(self, child_process: subprocess.Popen) -> int:
|
|
while child_process is not None:
|
|
return_code = child_process.poll()
|
|
if return_code is not None:
|
|
# subprocess finished with return code
|
|
return return_code
|
|
else:
|
|
# still running, yield control, 0.1s by default
|
|
await asyncio.sleep(self.SUBPROCESS_POLL_PERIOD_S)
|
|
|
|
async def _poll_all(self, processes: List[psutil.Process]):
|
|
"""Poll processes until all are completed."""
|
|
while True:
|
|
(_, alive) = psutil.wait_procs(processes, timeout=0)
|
|
if len(alive) == 0:
|
|
return
|
|
else:
|
|
await asyncio.sleep(self.SUBPROCESS_POLL_PERIOD_S)
|
|
|
|
def _kill_processes(self, processes: List[psutil.Process], sig: signal.Signals):
|
|
"""Ensure each process is already finished or send a kill signal."""
|
|
for proc in processes:
|
|
try:
|
|
os.kill(proc.pid, sig)
|
|
except ProcessLookupError:
|
|
# Process is already dead
|
|
pass
|
|
|
|
async def run(
|
|
self,
|
|
# Signal actor used in testing to capture PENDING -> RUNNING cases
|
|
_start_signal_actor: Optional[ActorHandle] = None,
|
|
resources_specified: bool = False,
|
|
):
|
|
"""
|
|
Stop and start both happen asynchronously, coordinated by asyncio event
|
|
and coroutine, respectively.
|
|
|
|
1) Sets job status as running
|
|
2) Pass runtime env and metadata to subprocess as serialized env
|
|
variables.
|
|
3) Handle concurrent events of driver execution and
|
|
"""
|
|
curr_info = await self._job_info_client.get_info(self._job_id)
|
|
if curr_info is None:
|
|
raise RuntimeError(f"Status could not be retrieved for job {self._job_id}.")
|
|
curr_status = curr_info.status
|
|
curr_message = curr_info.message
|
|
if curr_status == JobStatus.RUNNING:
|
|
raise RuntimeError(
|
|
f"Job {self._job_id} is already in RUNNING state. "
|
|
f"JobSupervisor.run() should only be called once. "
|
|
)
|
|
if curr_status != JobStatus.PENDING:
|
|
raise RuntimeError(
|
|
f"Job {self._job_id} is not in PENDING state. "
|
|
f"Current status is {curr_status} with message {curr_message}."
|
|
)
|
|
|
|
if _start_signal_actor:
|
|
# Block in PENDING state until start signal received.
|
|
await _start_signal_actor.wait.remote()
|
|
|
|
node = ray._private.worker.global_worker.node
|
|
driver_agent_http_address = f"http://{build_address(node.node_ip_address, node.dashboard_agent_listen_port)}"
|
|
driver_node_id = ray.get_runtime_context().get_node_id()
|
|
|
|
await self._job_info_client.put_status(
|
|
self._job_id,
|
|
JobStatus.RUNNING,
|
|
jobinfo_replace_kwargs={
|
|
"driver_agent_http_address": driver_agent_http_address,
|
|
"driver_node_id": driver_node_id,
|
|
},
|
|
)
|
|
|
|
try:
|
|
# Configure environment variables for the child process.
|
|
env = os.environ.copy()
|
|
# Remove internal Ray flags. They present because JobSuperVisor itself is
|
|
# a Ray worker process but we don't want to pass them to the driver.
|
|
remove_ray_internal_flags_from_env(env)
|
|
# These will *not* be set in the runtime_env, so they apply to the driver
|
|
# only, not its tasks & actors.
|
|
env.update(self._get_driver_env_vars(resources_specified))
|
|
|
|
self._logger.info(
|
|
"Submitting job with RAY_ADDRESS = "
|
|
f"{env[ray_constants.RAY_ADDRESS_ENVIRONMENT_VARIABLE]}"
|
|
)
|
|
log_path = self._log_client.get_log_file_path(self._job_id)
|
|
child_process = self._exec_entrypoint(env, log_path)
|
|
child_pid = child_process.pid
|
|
|
|
polling_task = create_task(self._polling(child_process))
|
|
finished, _ = await asyncio.wait(
|
|
[polling_task, create_task(self._stop_event.wait())],
|
|
return_when=FIRST_COMPLETED,
|
|
)
|
|
|
|
if self._stop_event.is_set():
|
|
polling_task.cancel()
|
|
if sys.platform == "win32" and self._win32_job_object:
|
|
win32job.TerminateJobObject(self._win32_job_object, -1)
|
|
elif sys.platform != "win32":
|
|
stop_signal = os.environ.get("RAY_JOB_STOP_SIGNAL", "SIGTERM")
|
|
if stop_signal not in self.VALID_STOP_SIGNALS:
|
|
self._logger.warning(
|
|
f"{stop_signal} not a valid stop signal. Terminating "
|
|
"job with SIGTERM."
|
|
)
|
|
stop_signal = "SIGTERM"
|
|
|
|
job_process = psutil.Process(child_pid)
|
|
proc_to_kill = [job_process] + job_process.children(recursive=True)
|
|
|
|
# Send stop signal and wait for job to terminate gracefully,
|
|
# otherwise SIGKILL job forcefully after timeout.
|
|
self._kill_processes(proc_to_kill, getattr(signal, stop_signal))
|
|
try:
|
|
stop_job_wait_time = int(
|
|
os.environ.get(
|
|
"RAY_JOB_STOP_WAIT_TIME_S",
|
|
self.DEFAULT_RAY_JOB_STOP_WAIT_TIME_S,
|
|
)
|
|
)
|
|
poll_job_stop_task = create_task(self._poll_all(proc_to_kill))
|
|
await asyncio.wait_for(poll_job_stop_task, stop_job_wait_time)
|
|
self._logger.info(
|
|
f"Job {self._job_id} has been terminated gracefully "
|
|
f"with {stop_signal}."
|
|
)
|
|
except asyncio.TimeoutError:
|
|
self._logger.warning(
|
|
f"Attempt to gracefully terminate job {self._job_id} "
|
|
f"through {stop_signal} has timed out after "
|
|
f"{stop_job_wait_time} seconds. Job is now being "
|
|
"force-killed with SIGKILL."
|
|
)
|
|
self._kill_processes(proc_to_kill, signal.SIGKILL)
|
|
|
|
await self._job_info_client.put_status(self._job_id, JobStatus.STOPPED)
|
|
else:
|
|
# Child process finished execution and no stop event is set
|
|
# at the same time
|
|
assert len(finished) == 1, "Should have only one coroutine done"
|
|
[child_process_task] = finished
|
|
return_code = child_process_task.result()
|
|
self._logger.info(
|
|
f"Job {self._job_id} entrypoint command "
|
|
f"exited with code {return_code}"
|
|
)
|
|
if return_code == 0:
|
|
await self._job_info_client.put_status(
|
|
self._job_id,
|
|
JobStatus.SUCCEEDED,
|
|
driver_exit_code=return_code,
|
|
)
|
|
else:
|
|
log_tail = await self._log_client.get_last_n_log_lines(self._job_id)
|
|
if log_tail is not None and log_tail != "":
|
|
message = (
|
|
"Job entrypoint command "
|
|
f"failed with exit code {return_code}, "
|
|
"last available logs (truncated to 20,000 chars):\n"
|
|
+ log_tail
|
|
)
|
|
else:
|
|
message = (
|
|
"Job entrypoint command "
|
|
f"failed with exit code {return_code}. No logs available."
|
|
)
|
|
await self._job_info_client.put_status(
|
|
self._job_id,
|
|
JobStatus.FAILED,
|
|
message=message,
|
|
driver_exit_code=return_code,
|
|
error_type=JobErrorType.JOB_ENTRYPOINT_COMMAND_ERROR,
|
|
)
|
|
except Exception:
|
|
self._logger.error(
|
|
"Got unexpected exception while trying to execute driver "
|
|
f"command. {traceback.format_exc()}"
|
|
)
|
|
try:
|
|
await self._job_info_client.put_status(
|
|
self._job_id,
|
|
JobStatus.FAILED,
|
|
message=traceback.format_exc(),
|
|
error_type=JobErrorType.JOB_ENTRYPOINT_COMMAND_START_ERROR,
|
|
)
|
|
except Exception:
|
|
self._logger.error(
|
|
"Failed to update job status to FAILED. "
|
|
f"Exception: {traceback.format_exc()}"
|
|
)
|
|
finally:
|
|
# clean up actor after tasks are finished
|
|
ray.actor.exit_actor()
|
|
|
|
def stop(self):
|
|
"""Set step_event and let run() handle the rest in its asyncio.wait()."""
|
|
self._stop_event.set()
|