wehub-resource-sync
543 lines
22 KiB
Python
543 lines
22 KiB
Python
import copy
|
|
import dataclasses
|
|
import logging
|
|
from typing import Any, AsyncIterator, Dict, List, Optional, Union
|
|
|
|
import packaging.version
|
|
|
|
import ray
|
|
from ray.dashboard.modules.dashboard_sdk import SubmissionClient
|
|
from ray.dashboard.modules.job.common import (
|
|
JobDeleteResponse,
|
|
JobLogsResponse,
|
|
JobStatus,
|
|
JobStopResponse,
|
|
JobSubmitRequest,
|
|
JobSubmitResponse,
|
|
)
|
|
from ray.dashboard.modules.job.pydantic_models import JobDetails
|
|
from ray.dashboard.modules.job.utils import strip_keys_with_value_none
|
|
from ray.dashboard.utils import get_address_for_submission_client
|
|
from ray.runtime_env import RuntimeEnv
|
|
from ray.runtime_env.runtime_env import _validate_no_local_paths
|
|
from ray.util.annotations import PublicAPI
|
|
|
|
try:
|
|
import aiohttp
|
|
import requests
|
|
except ImportError:
|
|
aiohttp = None
|
|
requests = None
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(logging.INFO)
|
|
|
|
|
|
class JobSubmissionClient(SubmissionClient):
|
|
"""A local client for submitting and interacting with jobs on a remote cluster.
|
|
|
|
Submits requests over HTTP to the job server on the cluster using the REST API.
|
|
|
|
|
|
Args:
|
|
address: Either (1) the address of the Ray cluster, or (2) the HTTP address
|
|
of the dashboard server on the head node, e.g. "http://<head-node-ip>:8265".
|
|
In case (1) it must be specified as an address that can be passed to
|
|
ray.init(), e.g. a Ray Client address (ray://<head_node_host>:10001),
|
|
or "auto", or "localhost:<port>". If unspecified, will try to connect to
|
|
a running local Ray cluster. This argument is always overridden by the
|
|
RAY_API_SERVER_ADDRESS or RAY_ADDRESS environment variable.
|
|
create_cluster_if_needed: Indicates whether the cluster at the specified
|
|
address needs to already be running. Ray doesn't start a cluster
|
|
before interacting with jobs, but third-party job managers may do so.
|
|
cookies: Cookies to use when sending requests to the HTTP job server.
|
|
metadata: Arbitrary metadata to store along with all jobs. New metadata
|
|
specified per job will be merged with the global metadata provided here
|
|
via a simple dict update.
|
|
headers: Headers to use when sending requests to the HTTP job server, used
|
|
for cases like authentication to a remote cluster.
|
|
verify: Boolean indication to verify the server's TLS certificate or a path to
|
|
a file or directory of trusted certificates. Default: True.
|
|
**kwargs: Additional keyword arguments forwarded to the cluster info
|
|
resolution function. For external module addresses (e.g.,
|
|
``anyscale://``), these are passed through to the module's
|
|
``get_job_submission_client_cluster_info()`` implementation.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
address: Optional[str] = None,
|
|
create_cluster_if_needed: bool = False,
|
|
cookies: Optional[Dict[str, Any]] = None,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
headers: Optional[Dict[str, Any]] = None,
|
|
verify: Optional[Union[str, bool]] = True,
|
|
**kwargs,
|
|
):
|
|
self._client_ray_version = ray.__version__
|
|
"""Initialize a JobSubmissionClient and check the connection to the cluster."""
|
|
if requests is None:
|
|
raise RuntimeError(
|
|
"The Ray jobs CLI & SDK require the ray[default] "
|
|
"installation: `pip install 'ray[default]'`"
|
|
)
|
|
# Check types of arguments
|
|
if address is not None and not isinstance(address, str):
|
|
raise TypeError(f"address must be a string, got {type(address)}")
|
|
if not isinstance(create_cluster_if_needed, bool):
|
|
raise TypeError(
|
|
f"create_cluster_if_needed must be a bool, got"
|
|
f" {type(create_cluster_if_needed)}"
|
|
)
|
|
if cookies is not None and not isinstance(cookies, dict):
|
|
raise TypeError(f"cookies must be a dict, got {type(cookies)}")
|
|
if metadata is not None and not isinstance(metadata, dict):
|
|
raise TypeError(f"metadata must be a dict, got {type(metadata)}")
|
|
if headers is not None and not isinstance(headers, dict):
|
|
raise TypeError(f"headers must be a dict, got {type(headers)}")
|
|
if not (isinstance(verify, str) or isinstance(verify, bool)):
|
|
raise TypeError(f"verify must be a str or bool, got {type(verify)}")
|
|
|
|
api_server_url = get_address_for_submission_client(address)
|
|
|
|
super().__init__(
|
|
address=api_server_url,
|
|
create_cluster_if_needed=create_cluster_if_needed,
|
|
cookies=cookies,
|
|
metadata=metadata,
|
|
headers=headers,
|
|
verify=verify,
|
|
**kwargs,
|
|
)
|
|
self._check_connection_and_version(
|
|
min_version="1.9",
|
|
version_error_message="Jobs API is not supported on the Ray "
|
|
"cluster. Please ensure the cluster is "
|
|
"running Ray 1.9 or higher.",
|
|
)
|
|
|
|
# In ray>=2.0, the client sends the new kwarg `submission_id` to the server
|
|
# upon every job submission, which causes servers with ray<2.0 to error.
|
|
if packaging.version.parse(self._client_ray_version) > packaging.version.parse(
|
|
"2.0"
|
|
):
|
|
self._check_connection_and_version(
|
|
min_version="2.0",
|
|
version_error_message=f"Client Ray version {self._client_ray_version} "
|
|
"is not compatible with the Ray cluster. Please ensure the cluster is "
|
|
"running Ray 2.0 or higher or downgrade the client Ray version.",
|
|
)
|
|
|
|
@PublicAPI(stability="stable")
|
|
def submit_job(
|
|
self,
|
|
*,
|
|
entrypoint: str,
|
|
job_id: Optional[str] = None,
|
|
runtime_env: Optional[Dict[str, Any]] = None,
|
|
metadata: Optional[Dict[str, str]] = None,
|
|
submission_id: Optional[str] = None,
|
|
entrypoint_num_cpus: Optional[Union[int, float]] = None,
|
|
entrypoint_num_gpus: Optional[Union[int, float]] = None,
|
|
entrypoint_memory: Optional[int] = None,
|
|
entrypoint_resources: Optional[Dict[str, float]] = None,
|
|
entrypoint_label_selector: Optional[Dict[str, str]] = None,
|
|
) -> str:
|
|
"""Submit and execute a job asynchronously.
|
|
|
|
When a job is submitted, it runs once to completion or failure. Retries or
|
|
different runs with different parameters should be handled by the
|
|
submitter. Jobs are bound to the lifetime of a Ray cluster, so if the
|
|
cluster goes down, all running jobs on that cluster will be terminated.
|
|
|
|
Example:
|
|
>>> from ray.job_submission import JobSubmissionClient
|
|
>>> client = JobSubmissionClient("http://127.0.0.1:8265") # doctest: +SKIP
|
|
>>> client.submit_job( # doctest: +SKIP
|
|
... entrypoint="python script.py",
|
|
... runtime_env={
|
|
... "working_dir": "./",
|
|
... "pip": ["requests==2.26.0"]
|
|
... }
|
|
... ) # doctest: +SKIP
|
|
'raysubmit_4LamXRuQpYdSMg7J'
|
|
|
|
Args:
|
|
entrypoint: The shell command to run for this job.
|
|
job_id: DEPRECATED. This has been renamed to submission_id.
|
|
runtime_env: The runtime environment to install and run this job in.
|
|
metadata: Arbitrary data to store along with this job.
|
|
submission_id: A unique ID for this job.
|
|
entrypoint_num_cpus: The quantity of CPU cores to reserve for the execution
|
|
of the entrypoint command, separately from any tasks or actors launched
|
|
by it. Defaults to 0.
|
|
entrypoint_num_gpus: The quantity of GPUs to reserve for the execution
|
|
of the entrypoint command, separately from any tasks or actors launched
|
|
by it. Defaults to 0.
|
|
entrypoint_memory: The quantity of memory to reserve for the
|
|
execution of the entrypoint command, separately from any tasks or
|
|
actors launched by it. Defaults to 0.
|
|
entrypoint_resources: The quantity of custom resources to reserve for the
|
|
execution of the entrypoint command, separately from any tasks or
|
|
actors launched by it.
|
|
entrypoint_label_selector: Label selector for the entrypoint command.
|
|
|
|
Returns:
|
|
The submission ID of the submitted job. If not specified,
|
|
this is a randomly generated unique ID.
|
|
|
|
Raises:
|
|
RuntimeError: If the request to the job server fails, or if the specified
|
|
submission_id has already been used by a job on this cluster.
|
|
"""
|
|
if job_id:
|
|
logger.warning(
|
|
"job_id kwarg is deprecated. Please use submission_id instead."
|
|
)
|
|
|
|
if (
|
|
entrypoint_num_cpus
|
|
or entrypoint_num_gpus
|
|
or entrypoint_resources
|
|
or entrypoint_label_selector
|
|
):
|
|
self._check_connection_and_version(
|
|
min_version="2.2",
|
|
version_error_message="`entrypoint_num_cpus`, `entrypoint_num_gpus`, "
|
|
"`entrypoint_resources`, and `entrypoint_label_selector` kwargs "
|
|
"are not supported on the Ray cluster. Please ensure the cluster is "
|
|
"running Ray 2.2 or higher.",
|
|
)
|
|
|
|
if entrypoint_memory:
|
|
self._check_connection_and_version(
|
|
min_version="2.8",
|
|
version_error_message="`entrypoint_memory` kwarg "
|
|
"is not supported on the Ray cluster. Please ensure the cluster is "
|
|
"running Ray 2.8 or higher.",
|
|
)
|
|
|
|
runtime_env = copy.deepcopy(runtime_env or {})
|
|
metadata = metadata or {}
|
|
metadata.update(self._default_metadata)
|
|
|
|
self._upload_working_dir_if_needed(runtime_env)
|
|
self._upload_py_modules_if_needed(runtime_env)
|
|
|
|
# Verify worker_process_setup_hook type.
|
|
setup_hook = runtime_env.get("worker_process_setup_hook")
|
|
if setup_hook and not isinstance(setup_hook, str):
|
|
raise ValueError(
|
|
f"Invalid type {type(setup_hook)} for `worker_process_setup_hook`. "
|
|
"When a job submission API is used, `worker_process_setup_hook` "
|
|
"only allows a string type (module name). "
|
|
"Specify `worker_process_setup_hook` via "
|
|
"ray.init within a driver to use a `Callable` type. "
|
|
)
|
|
|
|
# Run the RuntimeEnv constructor to parse local pip/conda requirements files.
|
|
runtime_env = RuntimeEnv(**runtime_env)
|
|
_validate_no_local_paths(runtime_env)
|
|
runtime_env = runtime_env.to_dict()
|
|
|
|
submission_id = submission_id or job_id
|
|
req = JobSubmitRequest(
|
|
entrypoint=entrypoint,
|
|
submission_id=submission_id,
|
|
runtime_env=runtime_env,
|
|
metadata=metadata,
|
|
entrypoint_num_cpus=entrypoint_num_cpus,
|
|
entrypoint_num_gpus=entrypoint_num_gpus,
|
|
entrypoint_memory=entrypoint_memory,
|
|
entrypoint_resources=entrypoint_resources,
|
|
entrypoint_label_selector=entrypoint_label_selector,
|
|
)
|
|
|
|
# Remove keys with value None so that new clients with new optional fields
|
|
# are still compatible with older servers. This is also done on the server,
|
|
# but we do it here as well to be extra defensive.
|
|
json_data = strip_keys_with_value_none(dataclasses.asdict(req))
|
|
|
|
logger.debug(f"Submitting job with submission_id={submission_id}.")
|
|
r = self._do_request("POST", "/api/jobs/", json_data=json_data)
|
|
|
|
if r.status_code == 200:
|
|
return JobSubmitResponse(**r.json()).submission_id
|
|
else:
|
|
self._raise_error(r)
|
|
|
|
@PublicAPI(stability="stable")
|
|
def stop_job(
|
|
self,
|
|
job_id: str,
|
|
) -> bool:
|
|
"""Request a job to exit asynchronously.
|
|
|
|
Attempts to terminate process first, then kills process after timeout.
|
|
|
|
Example:
|
|
>>> from ray.job_submission import JobSubmissionClient
|
|
>>> client = JobSubmissionClient("http://127.0.0.1:8265") # doctest: +SKIP
|
|
>>> sub_id = client.submit_job(entrypoint="sleep 10") # doctest: +SKIP
|
|
>>> client.stop_job(sub_id) # doctest: +SKIP
|
|
True
|
|
|
|
Args:
|
|
job_id: The job ID or submission ID for the job to be stopped.
|
|
|
|
Returns:
|
|
True if the job was running, otherwise False.
|
|
|
|
Raises:
|
|
RuntimeError: If the job does not exist or if the request to the
|
|
job server fails.
|
|
"""
|
|
logger.debug(f"Stopping job with job_id={job_id}.")
|
|
r = self._do_request("POST", f"/api/jobs/{job_id}/stop")
|
|
|
|
if r.status_code == 200:
|
|
return JobStopResponse(**r.json()).stopped
|
|
else:
|
|
self._raise_error(r)
|
|
|
|
@PublicAPI(stability="stable")
|
|
def delete_job(
|
|
self,
|
|
job_id: str,
|
|
) -> bool:
|
|
"""Delete a job in a terminal state and all of its associated data.
|
|
|
|
If the job is not already in a terminal state, raises an error.
|
|
This does not delete the job logs from disk.
|
|
Submitting a job with the same submission ID as a previously
|
|
deleted job is not supported and may lead to unexpected behavior.
|
|
|
|
Example:
|
|
>>> from ray.job_submission import JobSubmissionClient
|
|
>>> client = JobSubmissionClient() # doctest: +SKIP
|
|
>>> job_id = client.submit_job(entrypoint="echo hello") # doctest: +SKIP
|
|
>>> client.delete_job(job_id) # doctest: +SKIP
|
|
True
|
|
|
|
Args:
|
|
job_id: submission ID for the job to be deleted.
|
|
|
|
Returns:
|
|
True if the job was deleted, otherwise False.
|
|
|
|
Raises:
|
|
RuntimeError: If the job does not exist, if the request to the
|
|
job server fails, or if the job is not in a terminal state.
|
|
"""
|
|
logger.debug(f"Deleting job with job_id={job_id}.")
|
|
r = self._do_request("DELETE", f"/api/jobs/{job_id}")
|
|
|
|
if r.status_code == 200:
|
|
return JobDeleteResponse(**r.json()).deleted
|
|
else:
|
|
self._raise_error(r)
|
|
|
|
@PublicAPI(stability="stable")
|
|
def get_job_info(
|
|
self,
|
|
job_id: str,
|
|
) -> JobDetails:
|
|
"""Get the latest status and other information associated with a job.
|
|
|
|
Example:
|
|
>>> from ray.job_submission import JobSubmissionClient
|
|
>>> client = JobSubmissionClient("http://127.0.0.1:8265") # doctest: +SKIP
|
|
>>> submission_id = client.submit_job(entrypoint="sleep 1") # doctest: +SKIP
|
|
>>> client.get_job_info(submission_id) # doctest: +SKIP
|
|
JobDetails(status='SUCCEEDED',
|
|
job_id='03000000', type='submission',
|
|
submission_id='raysubmit_4LamXRuQpYdSMg7J',
|
|
message='Job finished successfully.', error_type=None,
|
|
start_time=1647388711, end_time=1647388712, metadata={}, runtime_env={})
|
|
|
|
Args:
|
|
job_id: The job ID or submission ID of the job whose information
|
|
is being requested.
|
|
|
|
Returns:
|
|
The JobDetails for the job.
|
|
|
|
Raises:
|
|
RuntimeError: If the job does not exist or if the request to the
|
|
job server fails.
|
|
"""
|
|
r = self._do_request("GET", f"/api/jobs/{job_id}")
|
|
|
|
if r.status_code == 200:
|
|
if JobDetails is None:
|
|
raise RuntimeError(
|
|
"The Ray jobs CLI & SDK require the ray[default] "
|
|
"installation: `pip install 'ray[default]'`"
|
|
)
|
|
else:
|
|
return JobDetails(**r.json())
|
|
else:
|
|
self._raise_error(r)
|
|
|
|
@PublicAPI(stability="stable")
|
|
def list_jobs(self) -> List[JobDetails]:
|
|
"""List all jobs along with their status and other information.
|
|
|
|
Lists all jobs that have ever run on the cluster, including jobs that are
|
|
currently running and jobs that are no longer running.
|
|
|
|
Example:
|
|
>>> from ray.job_submission import JobSubmissionClient
|
|
>>> client = JobSubmissionClient("http://127.0.0.1:8265") # doctest: +SKIP
|
|
>>> client.submit_job(entrypoint="echo hello") # doctest: +SKIP
|
|
>>> client.submit_job(entrypoint="sleep 2") # doctest: +SKIP
|
|
>>> client.list_jobs() # doctest: +SKIP
|
|
[JobDetails(status='SUCCEEDED',
|
|
job_id='03000000', type='submission',
|
|
submission_id='raysubmit_4LamXRuQpYdSMg7J',
|
|
message='Job finished successfully.', error_type=None,
|
|
start_time=1647388711, end_time=1647388712, metadata={}, runtime_env={}),
|
|
JobDetails(status='RUNNING',
|
|
job_id='04000000', type='submission',
|
|
submission_id='raysubmit_1dxCeNvG1fCMVNHG',
|
|
message='Job is currently running.', error_type=None,
|
|
start_time=1647454832, end_time=None, metadata={}, runtime_env={})]
|
|
|
|
Returns:
|
|
A list of JobDetails containing the job status and other information.
|
|
|
|
Raises:
|
|
RuntimeError: If the request to the job server fails.
|
|
"""
|
|
r = self._do_request("GET", "/api/jobs/")
|
|
|
|
if r.status_code == 200:
|
|
jobs_info_json = r.json()
|
|
jobs_info = [
|
|
JobDetails(**job_info_json) for job_info_json in jobs_info_json
|
|
]
|
|
return jobs_info
|
|
else:
|
|
self._raise_error(r)
|
|
|
|
@PublicAPI(stability="stable")
|
|
def get_job_status(self, job_id: str) -> JobStatus:
|
|
"""Get the most recent status of a job.
|
|
|
|
Example:
|
|
>>> from ray.job_submission import JobSubmissionClient
|
|
>>> client = JobSubmissionClient("http://127.0.0.1:8265") # doctest: +SKIP
|
|
>>> client.submit_job(entrypoint="echo hello") # doctest: +SKIP
|
|
>>> client.get_job_status("raysubmit_4LamXRuQpYdSMg7J") # doctest: +SKIP
|
|
'SUCCEEDED'
|
|
|
|
Args:
|
|
job_id: The job ID or submission ID of the job whose status is being
|
|
requested.
|
|
|
|
Returns:
|
|
The JobStatus of the job.
|
|
|
|
Raises:
|
|
RuntimeError: If the job does not exist or if the request to the
|
|
job server fails.
|
|
"""
|
|
return self.get_job_info(job_id).status
|
|
|
|
@PublicAPI(stability="stable")
|
|
def get_job_logs(self, job_id: str) -> str:
|
|
"""Get all logs produced by a job.
|
|
|
|
Example:
|
|
>>> from ray.job_submission import JobSubmissionClient
|
|
>>> client = JobSubmissionClient("http://127.0.0.1:8265") # doctest: +SKIP
|
|
>>> sub_id = client.submit_job(entrypoint="echo hello") # doctest: +SKIP
|
|
>>> client.get_job_logs(sub_id) # doctest: +SKIP
|
|
'hello\\n'
|
|
|
|
Args:
|
|
job_id: The job ID or submission ID of the job whose logs are being
|
|
requested.
|
|
|
|
Returns:
|
|
A string containing the full logs of the job.
|
|
|
|
Raises:
|
|
RuntimeError: If the job does not exist or if the request to the
|
|
job server fails.
|
|
"""
|
|
r = self._do_request("GET", f"/api/jobs/{job_id}/logs")
|
|
|
|
if r.status_code == 200:
|
|
return JobLogsResponse(**r.json()).logs
|
|
else:
|
|
self._raise_error(r)
|
|
|
|
@PublicAPI(stability="stable")
|
|
async def tail_job_logs(self, job_id: str) -> AsyncIterator[str]:
|
|
"""Get an iterator that follows the logs of a job.
|
|
|
|
Example:
|
|
>>> from ray.job_submission import JobSubmissionClient
|
|
>>> client = JobSubmissionClient("http://127.0.0.1:8265") # doctest: +SKIP
|
|
>>> submission_id = client.submit_job( # doctest: +SKIP
|
|
... entrypoint="echo hi && sleep 5 && echo hi2")
|
|
>>> async for lines in client.tail_job_logs( # doctest: +SKIP
|
|
... 'raysubmit_Xe7cvjyGJCyuCvm2'):
|
|
... print(lines, end="") # doctest: +SKIP
|
|
hi
|
|
hi2
|
|
|
|
Args:
|
|
job_id: The job ID or submission ID of the job whose logs are being
|
|
requested.
|
|
|
|
Yields:
|
|
str: Successive chunks of the job's stdout/stderr as the driver
|
|
process produces them.
|
|
|
|
Raises:
|
|
RuntimeError: If the job does not exist, if the request to the
|
|
job server fails, or if the connection closes unexpectedly
|
|
before the job reaches a terminal state.
|
|
"""
|
|
async with aiohttp.ClientSession(
|
|
cookies=self._cookies, headers=self._headers
|
|
) as session:
|
|
ws = await session.ws_connect(
|
|
f"{self._address}/api/jobs/{job_id}/logs/tail",
|
|
headers=self._headers,
|
|
ssl=self._ssl_context,
|
|
)
|
|
|
|
while True:
|
|
msg = await ws.receive()
|
|
|
|
if msg.type == aiohttp.WSMsgType.TEXT:
|
|
yield msg.data
|
|
elif msg.type == aiohttp.WSMsgType.CLOSED:
|
|
logger.info(
|
|
f"WebSocket closed for job {job_id} with close code "
|
|
f"{ws.close_code}"
|
|
)
|
|
if ws.close_code == aiohttp.WSCloseCode.ABNORMAL_CLOSURE:
|
|
raise RuntimeError(
|
|
f"WebSocket connection closed unexpectedly with close code {ws.close_code}"
|
|
)
|
|
break
|
|
elif msg.type == aiohttp.WSMsgType.ERROR:
|
|
# Old Ray versions (<=2.0.1) may send ERROR on connection close
|
|
if self._server_ray_version is not None and packaging.version.parse(
|
|
self._server_ray_version
|
|
) > packaging.version.parse("2.0.1"):
|
|
raise RuntimeError(
|
|
f"WebSocket error for job {job_id}: {ws.exception()}"
|
|
)
|
|
else:
|
|
logger.warning(
|
|
f"WebSocket error for job {job_id}, treating as "
|
|
f"normal close. Err: {ws.exception()!r}"
|
|
)
|
|
break
|