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