chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,600 @@
|
||||
import json
|
||||
import os
|
||||
import pprint
|
||||
import shlex
|
||||
import sys
|
||||
import time
|
||||
from subprocess import list2cmdline
|
||||
from typing import Any, Dict, Optional, Tuple, Union
|
||||
|
||||
import click
|
||||
|
||||
import ray._private.ray_constants as ray_constants
|
||||
from ray._common.utils import (
|
||||
get_or_create_event_loop,
|
||||
load_class,
|
||||
)
|
||||
from ray._private.utils import (
|
||||
parse_metadata_json,
|
||||
parse_resources_json,
|
||||
)
|
||||
from ray.autoscaler._private.cli_logger import add_click_logging_options, cf, cli_logger
|
||||
from ray.dashboard.modules.dashboard_sdk import parse_runtime_env_args
|
||||
from ray.dashboard.modules.job.cli_utils import add_common_job_options
|
||||
from ray.dashboard.modules.job.utils import redact_url_password
|
||||
from ray.job_submission import JobStatus, JobSubmissionClient
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
|
||||
def _get_sdk_client(
|
||||
address: Optional[str],
|
||||
create_cluster_if_needed: bool = False,
|
||||
headers: Optional[str] = None,
|
||||
verify: Union[bool, str] = True,
|
||||
) -> JobSubmissionClient:
|
||||
client = JobSubmissionClient(
|
||||
address,
|
||||
create_cluster_if_needed,
|
||||
headers=_handle_headers(headers),
|
||||
verify=verify,
|
||||
)
|
||||
client_address = client.get_address()
|
||||
cli_logger.labeled_value(
|
||||
"Job submission server address", redact_url_password(client_address)
|
||||
)
|
||||
return client
|
||||
|
||||
|
||||
def _handle_headers(headers: Optional[str]) -> Optional[Dict[str, Any]]:
|
||||
if headers is None and "RAY_JOB_HEADERS" in os.environ:
|
||||
headers = os.environ["RAY_JOB_HEADERS"]
|
||||
if headers is not None:
|
||||
try:
|
||||
return json.loads(headers)
|
||||
except Exception as exc:
|
||||
raise ValueError(
|
||||
"""Failed to parse headers into JSON.
|
||||
Expected format: {{"KEY": "VALUE"}}, got {}, {}""".format(
|
||||
headers, exc
|
||||
)
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _log_big_success_msg(success_msg):
|
||||
cli_logger.newline()
|
||||
cli_logger.success("-" * len(success_msg))
|
||||
cli_logger.success(success_msg)
|
||||
cli_logger.success("-" * len(success_msg))
|
||||
cli_logger.newline()
|
||||
|
||||
|
||||
def _log_big_error_msg(success_msg):
|
||||
cli_logger.newline()
|
||||
cli_logger.error("-" * len(success_msg))
|
||||
cli_logger.error(success_msg)
|
||||
cli_logger.error("-" * len(success_msg))
|
||||
cli_logger.newline()
|
||||
|
||||
|
||||
def _log_job_status(client: JobSubmissionClient, job_id: str) -> JobStatus:
|
||||
info = client.get_job_info(job_id)
|
||||
if info.status == JobStatus.SUCCEEDED:
|
||||
_log_big_success_msg(f"Job '{job_id}' succeeded")
|
||||
elif info.status == JobStatus.STOPPED:
|
||||
cli_logger.warning(f"Job '{job_id}' was stopped")
|
||||
elif info.status == JobStatus.FAILED:
|
||||
_log_big_error_msg(f"Job '{job_id}' failed")
|
||||
if info.message is not None:
|
||||
cli_logger.print(f"Status message: {info.message}", no_format=True)
|
||||
else:
|
||||
# Catch-all.
|
||||
cli_logger.print(f"Status for job '{job_id}': {info.status}")
|
||||
if info.message is not None:
|
||||
cli_logger.print(f"Status message: {info.message}", no_format=True)
|
||||
return info.status
|
||||
|
||||
|
||||
async def _tail_logs(client: JobSubmissionClient, job_id: str) -> JobStatus:
|
||||
async for lines in client.tail_job_logs(job_id):
|
||||
print(lines, end="")
|
||||
|
||||
return _log_job_status(client, job_id)
|
||||
|
||||
|
||||
@click.group("job")
|
||||
def job_cli_group():
|
||||
"""Submit, stop, delete, or list Ray jobs."""
|
||||
pass
|
||||
|
||||
|
||||
@job_cli_group.command()
|
||||
@click.option(
|
||||
"--address",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help=(
|
||||
"Address of the Ray cluster to connect to. Can also be specified "
|
||||
"using the RAY_API_SERVER_ADDRESS environment variable (falls back to RAY_ADDRESS)."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--job-id",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help=("DEPRECATED: Use `--submission-id` instead."),
|
||||
)
|
||||
@click.option(
|
||||
"--submission-id",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help=(
|
||||
"Submission ID to specify for the job. If not provided, one will be generated."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--runtime-env",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help="Path to a local YAML file containing a runtime_env definition.",
|
||||
)
|
||||
@click.option(
|
||||
"--runtime-env-json",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help="JSON-serialized runtime_env dictionary.",
|
||||
)
|
||||
@click.option(
|
||||
"--working-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help=(
|
||||
"Directory containing files that your job will run in. Can be a "
|
||||
"local directory or a remote URI to a .zip file (S3, GS, HTTP). "
|
||||
"If specified, this overrides the option in `--runtime-env`."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--metadata-json",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help="JSON-serialized dictionary of metadata to attach to the job.",
|
||||
)
|
||||
@click.option(
|
||||
"--entrypoint-num-cpus",
|
||||
required=False,
|
||||
type=float,
|
||||
help="the quantity of CPU cores to reserve for the entrypoint command, "
|
||||
"separately from any tasks or actors that are launched by it",
|
||||
)
|
||||
@click.option(
|
||||
"--entrypoint-num-gpus",
|
||||
required=False,
|
||||
type=float,
|
||||
help="the quantity of GPUs to reserve for the entrypoint command, "
|
||||
"separately from any tasks or actors that are launched by it",
|
||||
)
|
||||
@click.option(
|
||||
"--entrypoint-memory",
|
||||
required=False,
|
||||
type=int,
|
||||
help="the amount of memory to reserve "
|
||||
"for the entrypoint command, separately from any tasks or actors that are "
|
||||
"launched by it",
|
||||
)
|
||||
@click.option(
|
||||
"--entrypoint-resources",
|
||||
required=False,
|
||||
type=str,
|
||||
help="a JSON-serialized dictionary mapping resource name to resource quantity "
|
||||
"describing resources to reserve for the entrypoint command, "
|
||||
"separately from any tasks or actors that are launched by it",
|
||||
)
|
||||
@click.option(
|
||||
"--entrypoint-label-selector",
|
||||
required=False,
|
||||
type=str,
|
||||
help="a JSON-serialized dictionary mapping label keys to selector strings "
|
||||
"describing placement constraints for the entrypoint command",
|
||||
)
|
||||
@click.option(
|
||||
"--no-wait",
|
||||
is_flag=True,
|
||||
type=bool,
|
||||
default=False,
|
||||
help="If set, will not stream logs and wait for the job to exit.",
|
||||
)
|
||||
@add_common_job_options
|
||||
@add_click_logging_options
|
||||
@click.argument("entrypoint", nargs=-1, required=True, type=click.UNPROCESSED)
|
||||
@PublicAPI
|
||||
def submit(
|
||||
address: Optional[str],
|
||||
job_id: Optional[str],
|
||||
submission_id: Optional[str],
|
||||
runtime_env: Optional[str],
|
||||
runtime_env_json: Optional[str],
|
||||
metadata_json: Optional[str],
|
||||
working_dir: Optional[str],
|
||||
entrypoint: Tuple[str],
|
||||
entrypoint_num_cpus: Optional[Union[int, float]],
|
||||
entrypoint_num_gpus: Optional[Union[int, float]],
|
||||
entrypoint_memory: Optional[int],
|
||||
entrypoint_resources: Optional[str],
|
||||
entrypoint_label_selector: Optional[str],
|
||||
no_wait: bool,
|
||||
verify: Union[bool, str],
|
||||
headers: Optional[str],
|
||||
):
|
||||
"""Submits a job to be run on the cluster.
|
||||
|
||||
By default (if --no-wait is not set), streams logs to stdout until the job finishes.
|
||||
If the job succeeded, exits with 0. If it failed, exits with 1.
|
||||
|
||||
Example:
|
||||
`ray job submit -- python my_script.py --arg=val`
|
||||
|
||||
Args:
|
||||
address: Job submission server address.
|
||||
job_id: DEPRECATED. Use submission_id instead.
|
||||
submission_id: Submission ID for the job.
|
||||
runtime_env: Path to a runtime_env YAML file.
|
||||
runtime_env_json: JSON-serialized runtime_env dictionary.
|
||||
metadata_json: JSON-serialized metadata dictionary.
|
||||
working_dir: Working directory for the job.
|
||||
entrypoint: Entrypoint command.
|
||||
entrypoint_num_cpus: CPU cores to reserve.
|
||||
entrypoint_num_gpus: GPUs to reserve.
|
||||
entrypoint_memory: Memory to reserve.
|
||||
entrypoint_resources: JSON-serialized custom resources dict.
|
||||
entrypoint_label_selector: JSON-serialized label selector dict.
|
||||
no_wait: Do not wait for job completion.
|
||||
verify: TLS verification flag or path.
|
||||
headers: JSON-serialized headers.
|
||||
"""
|
||||
if job_id:
|
||||
cli_logger.warning(
|
||||
"--job-id option is deprecated. Please use --submission-id instead."
|
||||
)
|
||||
if entrypoint_resources is not None:
|
||||
entrypoint_resources = parse_resources_json(
|
||||
entrypoint_resources, cli_logger, cf, command_arg="entrypoint-resources"
|
||||
)
|
||||
if entrypoint_label_selector is not None:
|
||||
entrypoint_label_selector = parse_resources_json(
|
||||
entrypoint_label_selector,
|
||||
cli_logger,
|
||||
cf,
|
||||
command_arg="entrypoint-label-selector",
|
||||
)
|
||||
if metadata_json is not None:
|
||||
metadata_json = parse_metadata_json(
|
||||
metadata_json, cli_logger, cf, command_arg="metadata-json"
|
||||
)
|
||||
|
||||
submission_id = submission_id or job_id
|
||||
|
||||
if ray_constants.RAY_JOB_SUBMIT_HOOK in os.environ:
|
||||
# Submit all args as **kwargs per the JOB_SUBMIT_HOOK contract.
|
||||
load_class(os.environ[ray_constants.RAY_JOB_SUBMIT_HOOK])(
|
||||
address=address,
|
||||
job_id=submission_id,
|
||||
submission_id=submission_id,
|
||||
runtime_env=runtime_env,
|
||||
runtime_env_json=runtime_env_json,
|
||||
metadata_json=metadata_json,
|
||||
working_dir=working_dir,
|
||||
entrypoint=entrypoint,
|
||||
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,
|
||||
no_wait=no_wait,
|
||||
)
|
||||
|
||||
client = _get_sdk_client(
|
||||
address, create_cluster_if_needed=True, headers=headers, verify=verify
|
||||
)
|
||||
|
||||
final_runtime_env = parse_runtime_env_args(
|
||||
runtime_env=runtime_env,
|
||||
runtime_env_json=runtime_env_json,
|
||||
working_dir=working_dir,
|
||||
)
|
||||
job_id = client.submit_job(
|
||||
entrypoint=(
|
||||
list2cmdline(entrypoint)
|
||||
if sys.platform == "win32"
|
||||
else shlex.join(entrypoint)
|
||||
),
|
||||
submission_id=submission_id,
|
||||
runtime_env=final_runtime_env,
|
||||
metadata=metadata_json,
|
||||
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,
|
||||
)
|
||||
|
||||
_log_big_success_msg(f"Job '{job_id}' submitted successfully")
|
||||
|
||||
with cli_logger.group("Next steps"):
|
||||
cli_logger.print("Query the logs of the job:")
|
||||
with cli_logger.indented():
|
||||
cli_logger.print(cf.bold(f"ray job logs {job_id}"))
|
||||
|
||||
cli_logger.print("Query the status of the job:")
|
||||
with cli_logger.indented():
|
||||
cli_logger.print(cf.bold(f"ray job status {job_id}"))
|
||||
|
||||
cli_logger.print("Request the job to be stopped:")
|
||||
with cli_logger.indented():
|
||||
cli_logger.print(cf.bold(f"ray job stop {job_id}"))
|
||||
|
||||
cli_logger.newline()
|
||||
# Flush stdout to ensure the Ray job ID is output immediately
|
||||
# for the kubectl plugin, ref PR #52780, Issue kuberay/#3508.
|
||||
cli_logger.flush()
|
||||
sdk_version = client.get_version()
|
||||
# sdk version 0 does not have log streaming
|
||||
if not no_wait:
|
||||
if int(sdk_version) > 0:
|
||||
cli_logger.print(
|
||||
"Tailing logs until the job exits (disable with --no-wait):"
|
||||
)
|
||||
job_status = get_or_create_event_loop().run_until_complete(
|
||||
_tail_logs(client, job_id)
|
||||
)
|
||||
if job_status == JobStatus.FAILED:
|
||||
sys.exit(1)
|
||||
else:
|
||||
cli_logger.warning(
|
||||
"Tailing logs is not enabled for job sdk client version "
|
||||
f"{sdk_version}. Please upgrade Ray to the latest version "
|
||||
"for this feature."
|
||||
)
|
||||
|
||||
|
||||
@job_cli_group.command()
|
||||
@click.option(
|
||||
"--address",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help=(
|
||||
"Address of the Ray cluster to connect to. Can also be specified "
|
||||
"using the RAY_API_SERVER_ADDRESS environment variable (falls back to RAY_ADDRESS)."
|
||||
),
|
||||
)
|
||||
@click.argument("job-id", type=str)
|
||||
@add_common_job_options
|
||||
@add_click_logging_options
|
||||
@PublicAPI(stability="stable")
|
||||
def status(
|
||||
address: Optional[str],
|
||||
job_id: str,
|
||||
headers: Optional[str],
|
||||
verify: Union[bool, str],
|
||||
):
|
||||
"""Queries for the current status of a job.
|
||||
|
||||
Example:
|
||||
`ray job status <my_job_id>`
|
||||
|
||||
Args:
|
||||
address: Address of the Ray cluster to connect to.
|
||||
job_id: The submission ID of the job to query.
|
||||
headers: JSON string of headers to attach to requests.
|
||||
verify: Path to a CA bundle, or boolean toggling TLS verification.
|
||||
"""
|
||||
client = _get_sdk_client(address, headers=headers, verify=verify)
|
||||
_log_job_status(client, job_id)
|
||||
|
||||
|
||||
@job_cli_group.command()
|
||||
@click.option(
|
||||
"--address",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help=(
|
||||
"Address of the Ray cluster to connect to. Can also be specified "
|
||||
"using the RAY_API_SERVER_ADDRESS environment variable (falls back to RAY_ADDRESS)."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--no-wait",
|
||||
is_flag=True,
|
||||
type=bool,
|
||||
default=False,
|
||||
help="If set, will not wait for the job to exit.",
|
||||
)
|
||||
@click.argument("job-id", type=str)
|
||||
@add_common_job_options
|
||||
@add_click_logging_options
|
||||
@PublicAPI(stability="stable")
|
||||
def stop(
|
||||
address: Optional[str],
|
||||
no_wait: bool,
|
||||
job_id: str,
|
||||
headers: Optional[str],
|
||||
verify: Union[bool, str],
|
||||
):
|
||||
"""Attempts to stop a job.
|
||||
|
||||
Example:
|
||||
`ray job stop <my_job_id>`
|
||||
|
||||
Args:
|
||||
address: Address of the Ray cluster to connect to.
|
||||
no_wait: If True, return immediately instead of waiting for the job to reach a terminal state.
|
||||
job_id: The submission ID of the job to stop.
|
||||
headers: JSON string of headers to attach to requests.
|
||||
verify: Path to a CA bundle, or boolean toggling TLS verification.
|
||||
|
||||
Returns:
|
||||
None. The function returns early when ``no_wait`` is True; otherwise it
|
||||
polls until the job reaches a terminal state.
|
||||
"""
|
||||
client = _get_sdk_client(address, headers=headers, verify=verify)
|
||||
cli_logger.print(f"Attempting to stop job '{job_id}'")
|
||||
client.stop_job(job_id)
|
||||
|
||||
if no_wait:
|
||||
return
|
||||
else:
|
||||
cli_logger.print(
|
||||
f"Waiting for job '{job_id}' to exit (disable with --no-wait):"
|
||||
)
|
||||
|
||||
while True:
|
||||
status = client.get_job_status(job_id)
|
||||
if status in {JobStatus.STOPPED, JobStatus.SUCCEEDED, JobStatus.FAILED}:
|
||||
_log_job_status(client, job_id)
|
||||
break
|
||||
else:
|
||||
cli_logger.print(f"Job has not exited yet. Status: {status}")
|
||||
time.sleep(1)
|
||||
|
||||
|
||||
@job_cli_group.command()
|
||||
@click.option(
|
||||
"--address",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help=(
|
||||
"Address of the Ray cluster to connect to. Can also be specified "
|
||||
"using the RAY_API_SERVER_ADDRESS environment variable (falls back to RAY_ADDRESS)."
|
||||
),
|
||||
)
|
||||
@click.argument("job-id", type=str)
|
||||
@add_common_job_options
|
||||
@add_click_logging_options
|
||||
@PublicAPI(stability="stable")
|
||||
def delete(
|
||||
address: Optional[str],
|
||||
job_id: str,
|
||||
headers: Optional[str],
|
||||
verify: Union[bool, str],
|
||||
):
|
||||
"""Deletes a stopped job and its associated data from memory.
|
||||
|
||||
Only supported for jobs that are already in a terminal state.
|
||||
Fails with exit code 1 if the job is not already stopped.
|
||||
Does not delete 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:
|
||||
ray job delete <my_job_id>
|
||||
|
||||
Args:
|
||||
address: Address of the Ray cluster to connect to.
|
||||
job_id: The submission ID of the job to delete.
|
||||
headers: JSON string of headers to attach to requests.
|
||||
verify: Path to a CA bundle, or boolean toggling TLS verification.
|
||||
"""
|
||||
client = _get_sdk_client(address, headers=headers, verify=verify)
|
||||
client.delete_job(job_id)
|
||||
cli_logger.print(f"Job '{job_id}' deleted successfully")
|
||||
|
||||
|
||||
@job_cli_group.command()
|
||||
@click.option(
|
||||
"--address",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help=(
|
||||
"Address of the Ray cluster to connect to. Can also be specified "
|
||||
"using the RAY_API_SERVER_ADDRESS environment variable (falls back to RAY_ADDRESS)."
|
||||
),
|
||||
)
|
||||
@click.argument("job-id", type=str)
|
||||
@click.option(
|
||||
"-f",
|
||||
"--follow",
|
||||
is_flag=True,
|
||||
type=bool,
|
||||
default=False,
|
||||
help="If set, follow the logs (like `tail -f`).",
|
||||
)
|
||||
@add_common_job_options
|
||||
@add_click_logging_options
|
||||
@PublicAPI(stability="stable")
|
||||
def logs(
|
||||
address: Optional[str],
|
||||
job_id: str,
|
||||
follow: bool,
|
||||
headers: Optional[str],
|
||||
verify: Union[bool, str],
|
||||
):
|
||||
"""Gets the logs of a job.
|
||||
|
||||
Example:
|
||||
`ray job logs <my_job_id>`
|
||||
|
||||
Args:
|
||||
address: Address of the Ray cluster to connect to.
|
||||
job_id: The submission ID of the job whose logs to fetch.
|
||||
follow: If True, stream the logs (``tail -f`` style) instead of printing them once.
|
||||
headers: JSON string of headers to attach to requests.
|
||||
verify: Path to a CA bundle, or boolean toggling TLS verification.
|
||||
"""
|
||||
client = _get_sdk_client(address, headers=headers, verify=verify)
|
||||
sdk_version = client.get_version()
|
||||
# sdk version 0 did not have log streaming
|
||||
if follow:
|
||||
if int(sdk_version) > 0:
|
||||
get_or_create_event_loop().run_until_complete(_tail_logs(client, job_id))
|
||||
else:
|
||||
cli_logger.warning(
|
||||
"Tailing logs is not enabled for the Jobs SDK client version "
|
||||
f"{sdk_version}. Please upgrade Ray to latest version "
|
||||
"for this feature."
|
||||
)
|
||||
else:
|
||||
# Set no_format to True because the logs may have unescaped "{" and "}"
|
||||
# and the CLILogger calls str.format().
|
||||
cli_logger.print(client.get_job_logs(job_id), end="", no_format=True)
|
||||
|
||||
|
||||
@job_cli_group.command()
|
||||
@click.option(
|
||||
"--address",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help=(
|
||||
"Address of the Ray cluster to connect to. Can also be specified "
|
||||
"using the RAY_API_SERVER_ADDRESS environment variable (falls back to RAY_ADDRESS)."
|
||||
),
|
||||
)
|
||||
@add_common_job_options
|
||||
@add_click_logging_options
|
||||
@PublicAPI(stability="stable")
|
||||
def list(address: Optional[str], headers: Optional[str], verify: Union[bool, str]):
|
||||
"""Lists all running jobs and their information.
|
||||
|
||||
Example:
|
||||
`ray job list`
|
||||
|
||||
Args:
|
||||
address: Address of the Ray cluster to connect to.
|
||||
headers: JSON string of headers to attach to requests.
|
||||
verify: Path to a CA bundle, or boolean toggling TLS verification.
|
||||
"""
|
||||
client = _get_sdk_client(address, headers=headers, verify=verify)
|
||||
# Set no_format to True because the logs may have unescaped "{" and "}"
|
||||
# and the CLILogger calls str.format().
|
||||
cli_logger.print(pprint.pformat(client.list_jobs()), no_format=True)
|
||||
@@ -0,0 +1,56 @@
|
||||
import functools
|
||||
from typing import Union
|
||||
|
||||
import click
|
||||
|
||||
|
||||
def bool_cast(string: str) -> Union[bool, str]:
|
||||
"""Cast a string to a boolean if possible, otherwise return the string."""
|
||||
if string.lower() == "true" or string == "1":
|
||||
return True
|
||||
elif string.lower() == "false" or string == "0":
|
||||
return False
|
||||
else:
|
||||
return string
|
||||
|
||||
|
||||
class BoolOrStringParam(click.ParamType):
|
||||
"""A click parameter that can be either a boolean or a string."""
|
||||
|
||||
name = "BOOL | TEXT"
|
||||
|
||||
def convert(self, value, param, ctx):
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
else:
|
||||
return bool_cast(value)
|
||||
|
||||
|
||||
def add_common_job_options(func):
|
||||
"""Decorator for adding CLI flags shared by all `ray job` commands."""
|
||||
|
||||
@click.option(
|
||||
"--verify",
|
||||
default=True,
|
||||
show_default=True,
|
||||
type=BoolOrStringParam(),
|
||||
help=(
|
||||
"Boolean indication to verify the server's TLS certificate or a path to"
|
||||
" a file or directory of trusted certificates."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--headers",
|
||||
required=False,
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Used to pass headers through http/s to the Ray Cluster."
|
||||
'please follow JSON formatting formatting {"key": "value"}'
|
||||
),
|
||||
)
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
@@ -0,0 +1,599 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, replace
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional, Tuple, Union
|
||||
|
||||
from ray._private import ray_constants
|
||||
from ray._private.event.export_event_logger import (
|
||||
EventLogType,
|
||||
check_export_api_enabled,
|
||||
get_export_event_logger,
|
||||
)
|
||||
from ray._private.runtime_env.packaging import parse_uri
|
||||
from ray._raylet import RAY_INTERNAL_NAMESPACE_PREFIX, GcsClient
|
||||
from ray.core.generated.export_event_pb2 import ExportEvent
|
||||
from ray.core.generated.export_submission_job_event_pb2 import (
|
||||
ExportSubmissionJobEventData,
|
||||
)
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
# NOTE(edoakes): these constants should be considered a public API because
|
||||
# they're exposed in the snapshot API.
|
||||
JOB_ID_METADATA_KEY = "job_submission_id"
|
||||
JOB_NAME_METADATA_KEY = "job_name"
|
||||
JOB_ACTOR_NAME_TEMPLATE = f"{RAY_INTERNAL_NAMESPACE_PREFIX}job_actor_" + "{job_id}"
|
||||
# In order to get information about SupervisorActors launched by different jobs,
|
||||
# they must be set to the same namespace.
|
||||
SUPERVISOR_ACTOR_RAY_NAMESPACE = "SUPERVISOR_ACTOR_RAY_NAMESPACE"
|
||||
JOB_LOGS_PATH_TEMPLATE = "job-driver-{submission_id}.log"
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
class JobStatus(str, Enum):
|
||||
"""An enumeration for describing the status of a job."""
|
||||
|
||||
#: The job has not started yet, likely waiting for the runtime_env to be set up.
|
||||
PENDING = "PENDING"
|
||||
#: The job is currently running.
|
||||
RUNNING = "RUNNING"
|
||||
#: The job was intentionally stopped by the user.
|
||||
STOPPED = "STOPPED"
|
||||
#: The job finished successfully.
|
||||
SUCCEEDED = "SUCCEEDED"
|
||||
#: The job failed.
|
||||
FAILED = "FAILED"
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.value}"
|
||||
|
||||
def is_terminal(self) -> bool:
|
||||
"""Return whether or not this status is terminal.
|
||||
|
||||
A terminal status is one that cannot transition to any other status.
|
||||
The terminal statuses are "STOPPED", "SUCCEEDED", and "FAILED".
|
||||
|
||||
Returns:
|
||||
True if this status is terminal, otherwise False.
|
||||
"""
|
||||
return self.value in {"STOPPED", "SUCCEEDED", "FAILED"}
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
class JobErrorType(str, Enum):
|
||||
"""An enumeration for describing the error type of a job."""
|
||||
|
||||
# Runtime environment failed to be set up
|
||||
RUNTIME_ENV_SETUP_FAILURE = "RUNTIME_ENV_SETUP_FAILURE"
|
||||
# Job supervisor actor launched, but job failed to start within timeout
|
||||
JOB_SUPERVISOR_ACTOR_START_TIMEOUT = "JOB_SUPERVISOR_ACTOR_START_TIMEOUT"
|
||||
# Job supervisor actor failed to start
|
||||
JOB_SUPERVISOR_ACTOR_START_FAILURE = "JOB_SUPERVISOR_ACTOR_START_FAILURE"
|
||||
# Job supervisor actor failed to be scheduled
|
||||
JOB_SUPERVISOR_ACTOR_UNSCHEDULABLE = "JOB_SUPERVISOR_ACTOR_UNSCHEDULABLE"
|
||||
# Job supervisor actor failed for unknown exception
|
||||
JOB_SUPERVISOR_ACTOR_UNKNOWN_FAILURE = "JOB_SUPERVISOR_ACTOR_UNKNOWN_FAILURE"
|
||||
# Job supervisor actor died
|
||||
JOB_SUPERVISOR_ACTOR_DIED = "JOB_SUPERVISOR_ACTOR_DIED"
|
||||
# Job driver script failed to start due to exception
|
||||
JOB_ENTRYPOINT_COMMAND_START_ERROR = "JOB_ENTRYPOINT_COMMAND_START_ERROR"
|
||||
# Job driver script failed due to non-zero exit code
|
||||
JOB_ENTRYPOINT_COMMAND_ERROR = "JOB_ENTRYPOINT_COMMAND_ERROR"
|
||||
|
||||
|
||||
# TODO(aguo): Convert to pydantic model
|
||||
@PublicAPI(stability="stable")
|
||||
@dataclass
|
||||
class JobInfo:
|
||||
"""A class for recording information associated with a job and its execution.
|
||||
|
||||
Please keep this in sync with the JobsAPIInfo proto in src/ray/protobuf/gcs.proto.
|
||||
"""
|
||||
|
||||
#: The status of the job.
|
||||
status: JobStatus
|
||||
#: The entrypoint command for this job.
|
||||
entrypoint: str
|
||||
#: A message describing the status in more detail.
|
||||
message: Optional[str] = None
|
||||
#: Internal error, user script error
|
||||
error_type: Optional[JobErrorType] = None
|
||||
#: The time when the job was started. A Unix timestamp in ms.
|
||||
start_time: Optional[int] = None
|
||||
#: The time when the job moved into a terminal state. A Unix timestamp in ms.
|
||||
end_time: Optional[int] = None
|
||||
#: Arbitrary user-provided metadata for the job.
|
||||
metadata: Optional[Dict[str, str]] = None
|
||||
#: The runtime environment for the job.
|
||||
runtime_env: Optional[Dict[str, Any]] = None
|
||||
#: The quantity of CPU cores to reserve for the entrypoint command.
|
||||
entrypoint_num_cpus: Optional[Union[int, float]] = None
|
||||
#: The number of GPUs to reserve for the entrypoint command.
|
||||
entrypoint_num_gpus: Optional[Union[int, float]] = None
|
||||
#: The amount of memory for workers requesting memory for the entrypoint command.
|
||||
entrypoint_memory: Optional[int] = None
|
||||
#: The quantity of various custom resources to reserve for the entrypoint command.
|
||||
entrypoint_resources: Optional[Dict[str, float]] = None
|
||||
#: Driver agent http address
|
||||
driver_agent_http_address: Optional[str] = None
|
||||
#: The node id that driver running on. It will be None only when the job status
|
||||
# is PENDING, and this field will not be deleted or modified even if the driver dies
|
||||
driver_node_id: Optional[str] = None
|
||||
#: The driver process exit code after the driver executed. Return None if driver
|
||||
#: doesn't finish executing
|
||||
driver_exit_code: Optional[int] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if isinstance(self.status, str):
|
||||
self.status = JobStatus(self.status)
|
||||
if self.message is None:
|
||||
if self.status == JobStatus.PENDING:
|
||||
self.message = "Job has not started yet."
|
||||
if any(
|
||||
[
|
||||
self.entrypoint_num_cpus is not None
|
||||
and self.entrypoint_num_cpus > 0,
|
||||
self.entrypoint_num_gpus is not None
|
||||
and self.entrypoint_num_gpus > 0,
|
||||
self.entrypoint_memory is not None
|
||||
and self.entrypoint_memory > 0,
|
||||
self.entrypoint_resources not in [None, {}],
|
||||
]
|
||||
):
|
||||
self.message += (
|
||||
" It may be waiting for resources "
|
||||
"(CPUs, GPUs, memory, custom resources) to become available."
|
||||
)
|
||||
if self.runtime_env not in [None, {}]:
|
||||
self.message += (
|
||||
" It may be waiting for the runtime environment to be set up."
|
||||
)
|
||||
elif self.status == JobStatus.RUNNING:
|
||||
self.message = "Job is currently running."
|
||||
elif self.status == JobStatus.STOPPED:
|
||||
self.message = "Job was intentionally stopped."
|
||||
elif self.status == JobStatus.SUCCEEDED:
|
||||
self.message = "Job finished successfully."
|
||||
elif self.status == JobStatus.FAILED:
|
||||
self.message = "Job failed."
|
||||
|
||||
def to_json(self) -> Dict[str, Any]:
|
||||
"""Convert this object to a JSON-serializable dictionary.
|
||||
|
||||
Note that the runtime_env field is converted to a JSON-serialized string
|
||||
and the field is renamed to runtime_env_json.
|
||||
|
||||
Returns:
|
||||
A JSON-serializable dictionary representing the JobInfo object.
|
||||
"""
|
||||
|
||||
json_dict = asdict(self)
|
||||
|
||||
# Convert enum values to strings.
|
||||
json_dict["status"] = str(json_dict["status"])
|
||||
json_dict["error_type"] = (
|
||||
json_dict["error_type"].value if json_dict.get("error_type") else None
|
||||
)
|
||||
|
||||
# Convert runtime_env to a JSON-serialized string.
|
||||
if "runtime_env" in json_dict:
|
||||
if json_dict["runtime_env"] is not None:
|
||||
json_dict["runtime_env_json"] = json.dumps(json_dict["runtime_env"])
|
||||
del json_dict["runtime_env"]
|
||||
|
||||
# Assert that the dictionary is JSON-serializable.
|
||||
json.dumps(json_dict)
|
||||
|
||||
return json_dict
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, json_dict: Dict[str, Any]) -> None:
|
||||
"""Initialize this object from a JSON dictionary.
|
||||
|
||||
Note that the runtime_env_json field is converted to a dictionary and
|
||||
the field is renamed to runtime_env.
|
||||
|
||||
Args:
|
||||
json_dict: A JSON dictionary to use to initialize the JobInfo object.
|
||||
"""
|
||||
# Convert enum values to enum objects.
|
||||
json_dict["status"] = JobStatus(json_dict["status"])
|
||||
json_dict["error_type"] = (
|
||||
JobErrorType(json_dict["error_type"])
|
||||
if json_dict.get("error_type")
|
||||
else None
|
||||
)
|
||||
|
||||
# Convert runtime_env from a JSON-serialized string to a dictionary.
|
||||
if "runtime_env_json" in json_dict:
|
||||
if json_dict["runtime_env_json"] is not None:
|
||||
json_dict["runtime_env"] = json.loads(json_dict["runtime_env_json"])
|
||||
del json_dict["runtime_env_json"]
|
||||
|
||||
return cls(**json_dict)
|
||||
|
||||
|
||||
class JobInfoStorageClient:
|
||||
"""
|
||||
Interface to put and get job data from the Internal KV store.
|
||||
"""
|
||||
|
||||
# Please keep this format in sync with JobDataKey()
|
||||
# in src/ray/gcs/gcs_server/gcs_job_manager.h.
|
||||
JOB_DATA_KEY_PREFIX = f"{RAY_INTERNAL_NAMESPACE_PREFIX}job_info_"
|
||||
JOB_DATA_KEY = f"{JOB_DATA_KEY_PREFIX}{{job_id}}"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
gcs_client: GcsClient,
|
||||
export_event_log_dir_root: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the JobInfoStorageClient which manages data in the internal KV store.
|
||||
Export Submission Job events are written when the KV store is updated if
|
||||
the feature flag is on and a export_event_log_dir_root is passed.
|
||||
export_event_log_dir_root doesn't need to be passed if the caller
|
||||
is not modifying data in the KV store.
|
||||
"""
|
||||
self._gcs_client = gcs_client
|
||||
self._export_submission_job_event_logger: logging.Logger = None
|
||||
try:
|
||||
if (
|
||||
check_export_api_enabled(ExportEvent.SourceType.EXPORT_SUBMISSION_JOB)
|
||||
and export_event_log_dir_root is not None
|
||||
):
|
||||
self._export_submission_job_event_logger = get_export_event_logger(
|
||||
EventLogType.SUBMISSION_JOB,
|
||||
export_event_log_dir_root,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Unable to initialize export event logger so no export "
|
||||
"events will be written."
|
||||
)
|
||||
|
||||
async def put_info(
|
||||
self,
|
||||
job_id: str,
|
||||
job_info: JobInfo,
|
||||
overwrite: bool = True,
|
||||
timeout: Optional[int] = 30,
|
||||
) -> bool:
|
||||
"""Put job info to the internal kv store.
|
||||
|
||||
Args:
|
||||
job_id: The job id.
|
||||
job_info: The job info.
|
||||
overwrite: Whether to overwrite the existing job info.
|
||||
timeout: The timeout in seconds for the GCS operation.
|
||||
|
||||
Returns:
|
||||
True if a new key is added.
|
||||
"""
|
||||
added_num = await self._gcs_client.async_internal_kv_put(
|
||||
self.JOB_DATA_KEY.format(job_id=job_id).encode(),
|
||||
json.dumps(job_info.to_json()).encode(),
|
||||
overwrite,
|
||||
namespace=ray_constants.KV_NAMESPACE_JOB,
|
||||
timeout=timeout,
|
||||
)
|
||||
if added_num == 1 or overwrite:
|
||||
# Write export event if data was updated in the KV store
|
||||
try:
|
||||
self._write_submission_job_export_event(job_id, job_info)
|
||||
except Exception:
|
||||
logger.exception("Error while writing job submission export event.")
|
||||
return added_num == 1
|
||||
|
||||
def _write_submission_job_export_event(
|
||||
self, job_id: str, job_info: JobInfo
|
||||
) -> None:
|
||||
"""
|
||||
Write Submission Job export event if _export_submission_job_event_logger
|
||||
exists. The logger will exist if the export API feature flag is enabled
|
||||
and a log directory was passed to JobInfoStorageClient.
|
||||
"""
|
||||
if not self._export_submission_job_event_logger:
|
||||
return
|
||||
|
||||
status_value_descriptor = (
|
||||
ExportSubmissionJobEventData.JobStatus.DESCRIPTOR.values_by_name.get(
|
||||
job_info.status.name
|
||||
)
|
||||
)
|
||||
if status_value_descriptor is None:
|
||||
logger.error(
|
||||
f"{job_info.status.name} is not a valid "
|
||||
"ExportSubmissionJobEventData.JobStatus enum value. This event "
|
||||
"will not be written."
|
||||
)
|
||||
return
|
||||
job_status = status_value_descriptor.number
|
||||
submission_event_data = ExportSubmissionJobEventData(
|
||||
submission_job_id=job_id,
|
||||
status=job_status,
|
||||
entrypoint=job_info.entrypoint,
|
||||
message=job_info.message,
|
||||
metadata=job_info.metadata,
|
||||
error_type=job_info.error_type,
|
||||
start_time=job_info.start_time,
|
||||
end_time=job_info.end_time,
|
||||
runtime_env_json=json.dumps(job_info.runtime_env),
|
||||
driver_agent_http_address=job_info.driver_agent_http_address,
|
||||
driver_node_id=job_info.driver_node_id,
|
||||
driver_exit_code=job_info.driver_exit_code,
|
||||
)
|
||||
self._export_submission_job_event_logger.send_event(submission_event_data)
|
||||
|
||||
async def get_info(self, job_id: str, timeout: int = 30) -> Optional[JobInfo]:
|
||||
serialized_info = await self._gcs_client.async_internal_kv_get(
|
||||
self.JOB_DATA_KEY.format(job_id=job_id).encode(),
|
||||
namespace=ray_constants.KV_NAMESPACE_JOB,
|
||||
timeout=timeout,
|
||||
)
|
||||
if serialized_info is None:
|
||||
return None
|
||||
else:
|
||||
return JobInfo.from_json(json.loads(serialized_info))
|
||||
|
||||
async def delete_info(self, job_id: str, timeout: int = 30):
|
||||
await self._gcs_client.async_internal_kv_del(
|
||||
self.JOB_DATA_KEY.format(job_id=job_id).encode(),
|
||||
False,
|
||||
namespace=ray_constants.KV_NAMESPACE_JOB,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def put_status(
|
||||
self,
|
||||
job_id: str,
|
||||
status: JobStatus,
|
||||
message: Optional[str] = None,
|
||||
driver_exit_code: Optional[int] = None,
|
||||
error_type: Optional[JobErrorType] = None,
|
||||
jobinfo_replace_kwargs: Optional[Dict[str, Any]] = None,
|
||||
timeout: Optional[int] = 30,
|
||||
):
|
||||
"""Puts or updates job status. Sets end_time if status is terminal."""
|
||||
|
||||
old_info = await self.get_info(job_id, timeout=timeout)
|
||||
|
||||
if jobinfo_replace_kwargs is None:
|
||||
jobinfo_replace_kwargs = dict()
|
||||
jobinfo_replace_kwargs.update(
|
||||
status=status,
|
||||
message=message,
|
||||
driver_exit_code=driver_exit_code,
|
||||
error_type=error_type,
|
||||
)
|
||||
if old_info is not None:
|
||||
if status != old_info.status and old_info.status.is_terminal():
|
||||
raise RuntimeError(
|
||||
f"Attempted to change job status from a terminal state: "
|
||||
f"{old_info.status} -> {status}"
|
||||
)
|
||||
new_info = replace(old_info, **jobinfo_replace_kwargs)
|
||||
else:
|
||||
new_info = JobInfo(
|
||||
entrypoint="Entrypoint not found.", **jobinfo_replace_kwargs
|
||||
)
|
||||
|
||||
if status.is_terminal():
|
||||
new_info.end_time = int(time.time() * 1000)
|
||||
|
||||
await self.put_info(job_id, new_info, timeout=timeout)
|
||||
|
||||
async def get_status(self, job_id: str, timeout: int = 30) -> Optional[JobStatus]:
|
||||
job_info = await self.get_info(job_id, timeout)
|
||||
if job_info is None:
|
||||
return None
|
||||
else:
|
||||
return job_info.status
|
||||
|
||||
async def get_all_jobs(self, timeout: int = 30) -> Dict[str, JobInfo]:
|
||||
raw_job_ids_with_prefixes = await self._gcs_client.async_internal_kv_keys(
|
||||
self.JOB_DATA_KEY_PREFIX.encode(),
|
||||
namespace=ray_constants.KV_NAMESPACE_JOB,
|
||||
timeout=timeout,
|
||||
)
|
||||
job_ids_with_prefixes = [
|
||||
job_id.decode() for job_id in raw_job_ids_with_prefixes
|
||||
]
|
||||
job_ids = []
|
||||
for job_id_with_prefix in job_ids_with_prefixes:
|
||||
assert job_id_with_prefix.startswith(
|
||||
self.JOB_DATA_KEY_PREFIX
|
||||
), "Unexpected format for internal_kv key for Job submission"
|
||||
job_ids.append(job_id_with_prefix[len(self.JOB_DATA_KEY_PREFIX) :])
|
||||
|
||||
async def get_job_info(job_id: str):
|
||||
job_info = await self.get_info(job_id, timeout)
|
||||
return job_id, job_info
|
||||
|
||||
results = await asyncio.gather(*[get_job_info(job_id) for job_id in job_ids])
|
||||
return {
|
||||
job_id: job_info for job_id, job_info in results if job_info is not None
|
||||
}
|
||||
|
||||
|
||||
def uri_to_http_components(package_uri: str) -> Tuple[str, str]:
|
||||
suffix = Path(package_uri).suffix
|
||||
if suffix not in {".zip", ".whl"}:
|
||||
raise ValueError(f"package_uri ({package_uri}) does not end in .zip or .whl")
|
||||
# We need to strip the <protocol>:// prefix to make it possible to pass
|
||||
# the package_uri over HTTP.
|
||||
protocol, package_name = parse_uri(package_uri)
|
||||
return protocol.value, package_name
|
||||
|
||||
|
||||
def http_uri_components_to_uri(protocol: str, package_name: str) -> str:
|
||||
return f"{protocol}://{package_name}"
|
||||
|
||||
|
||||
def validate_request_type(json_data: Dict[str, Any], request_type: dataclass) -> Any:
|
||||
return request_type(**json_data)
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobSubmitRequest:
|
||||
# Command to start execution, ex: "python script.py"
|
||||
entrypoint: str
|
||||
# Optional submission_id to specify for the job. If the submission_id
|
||||
# is not specified, one will be generated. If a job with the same
|
||||
# submission_id already exists, it will be rejected.
|
||||
submission_id: Optional[str] = None
|
||||
# DEPRECATED. Use submission_id instead
|
||||
job_id: Optional[str] = None
|
||||
# Dict to setup execution environment.
|
||||
runtime_env: Optional[Dict[str, Any]] = None
|
||||
# Metadata to pass in to the JobConfig.
|
||||
metadata: Optional[Dict[str, str]] = None
|
||||
# The quantity of CPU cores to reserve for the execution
|
||||
# of the entrypoint command, separately from any Ray tasks or actors
|
||||
# that are created by it.
|
||||
entrypoint_num_cpus: Optional[Union[int, float]] = None
|
||||
# The quantity of GPUs to reserve for the execution
|
||||
# of the entrypoint command, separately from any Ray tasks or actors
|
||||
# that are created by it.
|
||||
entrypoint_num_gpus: Optional[Union[int, float]] = None
|
||||
# The amount of total available memory for workers requesting memory
|
||||
# for the execution of the entrypoint command, separately from any Ray
|
||||
# tasks or actors that are created by it.
|
||||
entrypoint_memory: Optional[int] = None
|
||||
# The quantity of various custom resources
|
||||
# to reserve for the entrypoint command, separately from any Ray tasks
|
||||
# or actors that are created by it.
|
||||
entrypoint_resources: Optional[Dict[str, float]] = None
|
||||
# Label selector for the entrypoint command.
|
||||
entrypoint_label_selector: Optional[Dict[str, str]] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if not isinstance(self.entrypoint, str):
|
||||
raise TypeError(f"entrypoint must be a string, got {type(self.entrypoint)}")
|
||||
|
||||
if self.submission_id is not None and not isinstance(self.submission_id, str):
|
||||
raise TypeError(
|
||||
"submission_id must be a string if provided, "
|
||||
f"got {type(self.submission_id)}"
|
||||
)
|
||||
|
||||
if self.job_id is not None and not isinstance(self.job_id, str):
|
||||
raise TypeError(
|
||||
"job_id must be a string if provided, " f"got {type(self.job_id)}"
|
||||
)
|
||||
|
||||
if self.runtime_env is not None:
|
||||
if not isinstance(self.runtime_env, dict):
|
||||
raise TypeError(
|
||||
f"runtime_env must be a dict, got {type(self.runtime_env)}"
|
||||
)
|
||||
else:
|
||||
for k in self.runtime_env.keys():
|
||||
if not isinstance(k, str):
|
||||
raise TypeError(
|
||||
f"runtime_env keys must be strings, got {type(k)}"
|
||||
)
|
||||
|
||||
if self.metadata is not None:
|
||||
if not isinstance(self.metadata, dict):
|
||||
raise TypeError(f"metadata must be a dict, got {type(self.metadata)}")
|
||||
else:
|
||||
for k in self.metadata.keys():
|
||||
if not isinstance(k, str):
|
||||
raise TypeError(f"metadata keys must be strings, got {type(k)}")
|
||||
for v in self.metadata.values():
|
||||
if not isinstance(v, str):
|
||||
raise TypeError(
|
||||
f"metadata values must be strings, got {type(v)}"
|
||||
)
|
||||
|
||||
if self.entrypoint_num_cpus is not None and not isinstance(
|
||||
self.entrypoint_num_cpus, (int, float)
|
||||
):
|
||||
raise TypeError(
|
||||
"entrypoint_num_cpus must be a number, "
|
||||
f"got {type(self.entrypoint_num_cpus)}"
|
||||
)
|
||||
|
||||
if self.entrypoint_num_gpus is not None and not isinstance(
|
||||
self.entrypoint_num_gpus, (int, float)
|
||||
):
|
||||
raise TypeError(
|
||||
"entrypoint_num_gpus must be a number, "
|
||||
f"got {type(self.entrypoint_num_gpus)}"
|
||||
)
|
||||
|
||||
if self.entrypoint_memory is not None and not isinstance(
|
||||
self.entrypoint_memory, int
|
||||
):
|
||||
raise TypeError(
|
||||
"entrypoint_memory must be an integer, "
|
||||
f"got {type(self.entrypoint_memory)}"
|
||||
)
|
||||
|
||||
if self.entrypoint_resources is not None:
|
||||
if not isinstance(self.entrypoint_resources, dict):
|
||||
raise TypeError(
|
||||
"entrypoint_resources must be a dict, "
|
||||
f"got {type(self.entrypoint_resources)}"
|
||||
)
|
||||
else:
|
||||
for k in self.entrypoint_resources.keys():
|
||||
if not isinstance(k, str):
|
||||
raise TypeError(
|
||||
"entrypoint_resources keys must be strings, "
|
||||
f"got {type(k)}"
|
||||
)
|
||||
for v in self.entrypoint_resources.values():
|
||||
if not isinstance(v, (int, float)):
|
||||
raise TypeError(
|
||||
"entrypoint_resources values must be numbers, "
|
||||
f"got {type(v)}"
|
||||
)
|
||||
|
||||
if self.entrypoint_label_selector is not None:
|
||||
if not isinstance(self.entrypoint_label_selector, dict):
|
||||
raise TypeError(
|
||||
"entrypoint_label_selector must be a dict, "
|
||||
f"got {type(self.entrypoint_label_selector)}"
|
||||
)
|
||||
else:
|
||||
for k, v in self.entrypoint_label_selector.items():
|
||||
if not isinstance(k, str):
|
||||
raise TypeError(
|
||||
"entrypoint_label_selector keys must be strings, "
|
||||
f"got {type(k)}"
|
||||
)
|
||||
if not isinstance(v, str):
|
||||
raise TypeError(
|
||||
"entrypoint_label_selector values must be strings, "
|
||||
f"got {type(v)}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobSubmitResponse:
|
||||
# DEPRECATED: Use submission_id instead.
|
||||
job_id: str
|
||||
submission_id: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobStopResponse:
|
||||
stopped: bool
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobDeleteResponse:
|
||||
deleted: bool
|
||||
|
||||
|
||||
# TODO(jiaodong): Support log streaming #19415
|
||||
@dataclass
|
||||
class JobLogsResponse:
|
||||
logs: str
|
||||
@@ -0,0 +1,26 @@
|
||||
{
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"$id": "http://github.com/ray-project/ray/dashboard/modules/job/component_activities_schema.json",
|
||||
"type": "object",
|
||||
"patternProperties": {
|
||||
"[0-9a-f]*": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"is_active": {
|
||||
"type": "string",
|
||||
"enum": ["ACTIVE", "INACTIVE", "ERROR"]
|
||||
},
|
||||
"reason": {
|
||||
"type": ["string", "null"]
|
||||
},
|
||||
"timestamp": {
|
||||
"type": ["number"]
|
||||
},
|
||||
"last_activity_at": {
|
||||
"type": ["number", "null"]
|
||||
}
|
||||
},
|
||||
"required": ["is_active"]
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
import dataclasses
|
||||
import json
|
||||
import logging
|
||||
import traceback
|
||||
|
||||
import aiohttp
|
||||
from aiohttp.web import Request, Response
|
||||
|
||||
import ray
|
||||
import ray.dashboard.optional_utils as optional_utils
|
||||
import ray.dashboard.utils as dashboard_utils
|
||||
from ray.dashboard.modules.job.common import (
|
||||
JobDeleteResponse,
|
||||
JobLogsResponse,
|
||||
JobStopResponse,
|
||||
JobSubmitRequest,
|
||||
JobSubmitResponse,
|
||||
)
|
||||
from ray.dashboard.modules.job.job_manager import JobManager
|
||||
from ray.dashboard.modules.job.pydantic_models import JobType
|
||||
from ray.dashboard.modules.job.utils import find_job_by_ids, parse_and_validate_request
|
||||
|
||||
routes = optional_utils.DashboardAgentRouteTable
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class JobAgent(dashboard_utils.DashboardAgentModule):
|
||||
def __init__(self, dashboard_agent):
|
||||
super().__init__(dashboard_agent)
|
||||
self._job_manager = None
|
||||
|
||||
@routes.post("/api/job_agent/jobs/")
|
||||
@optional_utils.init_ray_and_catch_exceptions()
|
||||
async def submit_job(self, req: Request) -> Response:
|
||||
result = await parse_and_validate_request(req, JobSubmitRequest)
|
||||
# Request parsing failed, returned with Response object.
|
||||
if isinstance(result, Response):
|
||||
return result
|
||||
else:
|
||||
submit_request = result
|
||||
|
||||
request_submission_id = submit_request.submission_id or submit_request.job_id
|
||||
try:
|
||||
ray._common.usage.usage_lib.record_library_usage("job_submission")
|
||||
submission_id = await self.get_job_manager().submit_job(
|
||||
entrypoint=submit_request.entrypoint,
|
||||
submission_id=request_submission_id,
|
||||
runtime_env=submit_request.runtime_env,
|
||||
metadata=submit_request.metadata,
|
||||
entrypoint_num_cpus=submit_request.entrypoint_num_cpus,
|
||||
entrypoint_num_gpus=submit_request.entrypoint_num_gpus,
|
||||
entrypoint_memory=submit_request.entrypoint_memory,
|
||||
entrypoint_resources=submit_request.entrypoint_resources,
|
||||
entrypoint_label_selector=submit_request.entrypoint_label_selector,
|
||||
)
|
||||
|
||||
resp = JobSubmitResponse(job_id=submission_id, submission_id=submission_id)
|
||||
except (TypeError, ValueError):
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(resp)),
|
||||
content_type="application/json",
|
||||
status=aiohttp.web.HTTPOk.status_code,
|
||||
)
|
||||
|
||||
@routes.post("/api/job_agent/jobs/{job_or_submission_id}/stop")
|
||||
@optional_utils.init_ray_and_catch_exceptions()
|
||||
async def stop_job(self, req: Request) -> Response:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self._dashboard_agent.gcs_client,
|
||||
self.get_job_manager().job_info_client(),
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
if job.type is not JobType.SUBMISSION:
|
||||
return Response(
|
||||
text="Can only stop submission type jobs",
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
try:
|
||||
stopped = self.get_job_manager().stop_job(job.submission_id)
|
||||
resp = JobStopResponse(stopped=stopped)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(resp)), content_type="application/json"
|
||||
)
|
||||
|
||||
@routes.delete("/api/job_agent/jobs/{job_or_submission_id}")
|
||||
@optional_utils.init_ray_and_catch_exceptions()
|
||||
async def delete_job(self, req: Request) -> Response:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self._dashboard_agent.gcs_client,
|
||||
self.get_job_manager().job_info_client(),
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
if job.type is not JobType.SUBMISSION:
|
||||
return Response(
|
||||
text="Can only delete submission type jobs",
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
try:
|
||||
deleted = await self.get_job_manager().delete_job(job.submission_id)
|
||||
resp = JobDeleteResponse(deleted=deleted)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(resp)), content_type="application/json"
|
||||
)
|
||||
|
||||
@routes.get("/api/job_agent/jobs/{job_or_submission_id}/logs")
|
||||
@optional_utils.init_ray_and_catch_exceptions()
|
||||
async def get_job_logs(self, req: Request) -> Response:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self._dashboard_agent.gcs_client,
|
||||
self.get_job_manager().job_info_client(),
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
|
||||
if job.type is not JobType.SUBMISSION:
|
||||
return Response(
|
||||
text="Can only get logs of submission type jobs",
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
resp = JobLogsResponse(
|
||||
logs=self.get_job_manager().get_job_logs(job.submission_id)
|
||||
)
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(resp)), content_type="application/json"
|
||||
)
|
||||
|
||||
@routes.get("/api/job_agent/jobs/{job_or_submission_id}/logs/tail")
|
||||
@optional_utils.init_ray_and_catch_exceptions()
|
||||
async def tail_job_logs(self, req: Request) -> Response:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self._dashboard_agent.gcs_client,
|
||||
self.get_job_manager().job_info_client(),
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
|
||||
if job.type is not JobType.SUBMISSION:
|
||||
return Response(
|
||||
text="Can only get logs of submission type jobs",
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
ws = aiohttp.web.WebSocketResponse()
|
||||
await ws.prepare(req)
|
||||
|
||||
async for lines in self._job_manager.tail_job_logs(job.submission_id):
|
||||
await ws.send_str(lines)
|
||||
|
||||
return ws
|
||||
|
||||
def get_job_manager(self):
|
||||
if not self._job_manager:
|
||||
self._job_manager = JobManager(
|
||||
self._dashboard_agent.gcs_client, self._dashboard_agent.log_dir
|
||||
)
|
||||
return self._job_manager
|
||||
|
||||
async def run(self, server):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def is_minimal_module():
|
||||
return False
|
||||
@@ -0,0 +1,774 @@
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import enum
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import AsyncIterator, Dict, Optional, Tuple
|
||||
|
||||
import aiohttp.web
|
||||
from aiohttp.client import ClientResponse
|
||||
from aiohttp.web import Request, Response, StreamResponse
|
||||
|
||||
import ray
|
||||
from ray import NodeID
|
||||
from ray._common.network_utils import build_address
|
||||
from ray._common.pydantic_compat import BaseModel, Extra, Field, validator
|
||||
from ray._common.utils import get_or_create_event_loop, load_class
|
||||
from ray._private.authentication.http_token_authentication import (
|
||||
get_auth_headers_if_auth_enabled,
|
||||
)
|
||||
from ray._private.ray_constants import KV_NAMESPACE_DASHBOARD
|
||||
from ray._private.runtime_env.packaging import (
|
||||
package_exists,
|
||||
pin_runtime_env_uri,
|
||||
upload_package_to_gcs,
|
||||
)
|
||||
from ray.dashboard.consts import (
|
||||
DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX,
|
||||
GCS_RPC_TIMEOUT_SECONDS,
|
||||
RAY_CLUSTER_ACTIVITY_HOOK,
|
||||
TRY_TO_GET_AGENT_INFO_INTERVAL_SECONDS,
|
||||
WAIT_AVAILABLE_AGENT_TIMEOUT,
|
||||
)
|
||||
from ray.dashboard.modules.job.common import (
|
||||
JobDeleteResponse,
|
||||
JobInfoStorageClient,
|
||||
JobLogsResponse,
|
||||
JobStopResponse,
|
||||
JobSubmitRequest,
|
||||
JobSubmitResponse,
|
||||
http_uri_components_to_uri,
|
||||
)
|
||||
from ray.dashboard.modules.job.pydantic_models import JobDetails, JobType
|
||||
from ray.dashboard.modules.job.utils import (
|
||||
find_job_by_ids,
|
||||
get_driver_jobs,
|
||||
parse_and_validate_request,
|
||||
)
|
||||
from ray.dashboard.modules.version import CURRENT_VERSION, VersionResponse
|
||||
from ray.dashboard.subprocesses.module import SubprocessModule
|
||||
from ray.dashboard.subprocesses.routes import SubprocessRouteTable as routes
|
||||
from ray.dashboard.subprocesses.utils import ResponseType
|
||||
from ray.dashboard.utils import get_head_node_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.INFO)
|
||||
|
||||
|
||||
class RayActivityStatus(str, enum.Enum):
|
||||
ACTIVE = "ACTIVE"
|
||||
INACTIVE = "INACTIVE"
|
||||
ERROR = "ERROR"
|
||||
|
||||
|
||||
class RayActivityResponse(BaseModel, extra=Extra.allow):
|
||||
"""
|
||||
Pydantic model used to inform if a particular Ray component can be considered
|
||||
active, and metadata about observation.
|
||||
"""
|
||||
|
||||
is_active: RayActivityStatus = Field(
|
||||
...,
|
||||
description=(
|
||||
"Whether the corresponding Ray component is considered active or inactive, "
|
||||
"or if there was an error while collecting this observation."
|
||||
),
|
||||
)
|
||||
reason: Optional[str] = Field(
|
||||
None, description="Reason if Ray component is considered active or errored."
|
||||
)
|
||||
timestamp: float = Field(
|
||||
...,
|
||||
description=(
|
||||
"Timestamp of when this observation about the Ray component was made. "
|
||||
"This is in the format of seconds since unix epoch."
|
||||
),
|
||||
)
|
||||
last_activity_at: Optional[float] = Field(
|
||||
None,
|
||||
description=(
|
||||
"Timestamp when last actvity of this Ray component finished in format of "
|
||||
"seconds since unix epoch. This field does not need to be populated "
|
||||
"for Ray components where it is not meaningful."
|
||||
),
|
||||
)
|
||||
|
||||
@validator("reason", always=True)
|
||||
def reason_required(cls, v, values, **kwargs):
|
||||
if "is_active" in values and values["is_active"] != RayActivityStatus.INACTIVE:
|
||||
if v is None:
|
||||
raise ValueError(
|
||||
'Reason is required if is_active is "active" or "error"'
|
||||
)
|
||||
return v
|
||||
|
||||
|
||||
class JobAgentSubmissionClient:
|
||||
"""A local client for submitting and interacting with jobs on a specific node
|
||||
in the remote cluster.
|
||||
Submits requests over HTTP to the job agent on the specific node using the REST API.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dashboard_agent_address: str,
|
||||
):
|
||||
self._agent_address = dashboard_agent_address
|
||||
self._session = aiohttp.ClientSession()
|
||||
|
||||
def _get_headers(self):
|
||||
"""Get auth headers if token authentication is enabled."""
|
||||
return get_auth_headers_if_auth_enabled({})
|
||||
|
||||
async def _raise_error(self, resp: ClientResponse):
|
||||
status = resp.status
|
||||
error_text = await resp.text()
|
||||
raise RuntimeError(f"Request failed with status code {status}: {error_text}.")
|
||||
|
||||
async def submit_job_internal(self, req: JobSubmitRequest) -> JobSubmitResponse:
|
||||
logger.debug(f"Submitting job with submission_id={req.submission_id}.")
|
||||
|
||||
async with self._session.post(
|
||||
f"{self._agent_address}/api/job_agent/jobs/",
|
||||
json=dataclasses.asdict(req),
|
||||
headers=self._get_headers(),
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
result_json = await resp.json()
|
||||
return JobSubmitResponse(**result_json)
|
||||
else:
|
||||
await self._raise_error(resp)
|
||||
|
||||
async def stop_job_internal(self, job_id: str) -> JobStopResponse:
|
||||
logger.debug(f"Stopping job with job_id={job_id}.")
|
||||
|
||||
async with self._session.post(
|
||||
f"{self._agent_address}/api/job_agent/jobs/{job_id}/stop",
|
||||
headers=self._get_headers(),
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
result_json = await resp.json()
|
||||
return JobStopResponse(**result_json)
|
||||
else:
|
||||
await self._raise_error(resp)
|
||||
|
||||
async def delete_job_internal(self, job_id: str) -> JobDeleteResponse:
|
||||
logger.debug(f"Deleting job with job_id={job_id}.")
|
||||
|
||||
async with self._session.delete(
|
||||
f"{self._agent_address}/api/job_agent/jobs/{job_id}",
|
||||
headers=self._get_headers(),
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
result_json = await resp.json()
|
||||
return JobDeleteResponse(**result_json)
|
||||
else:
|
||||
await self._raise_error(resp)
|
||||
|
||||
async def get_job_logs_internal(self, job_id: str) -> JobLogsResponse:
|
||||
async with self._session.get(
|
||||
f"{self._agent_address}/api/job_agent/jobs/{job_id}/logs",
|
||||
headers=self._get_headers(),
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
result_json = await resp.json()
|
||||
return JobLogsResponse(**result_json)
|
||||
else:
|
||||
await self._raise_error(resp)
|
||||
|
||||
async def tail_job_logs(self, job_id: str) -> AsyncIterator[str]:
|
||||
"""Get an iterator that follows the logs of a job."""
|
||||
ws = await self._session.ws_connect(
|
||||
f"{self._agent_address}/api/job_agent/jobs/{job_id}/logs/tail",
|
||||
headers=self._get_headers(),
|
||||
)
|
||||
|
||||
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 to job agent closed for job {job_id} "
|
||||
f"with close code {ws.close_code}"
|
||||
)
|
||||
break
|
||||
elif msg.type == aiohttp.WSMsgType.ERROR:
|
||||
logger.warning(
|
||||
f"WebSocket to job agent received an error message "
|
||||
f"while tailing logs for job {job_id}: {ws.exception()!r}. "
|
||||
)
|
||||
pass
|
||||
|
||||
async def close(self, ignore_error=True):
|
||||
try:
|
||||
await self._session.close()
|
||||
except Exception:
|
||||
if not ignore_error:
|
||||
raise
|
||||
|
||||
|
||||
class JobHead(SubprocessModule):
|
||||
"""Runs on the head node of a Ray cluster and handles Ray Jobs APIs.
|
||||
|
||||
NOTE(architkulkarni): Please keep this class in sync with the OpenAPI spec at
|
||||
`doc/source/cluster/running-applications/job-submission/openapi.yml`.
|
||||
We currently do not automatically check that the OpenAPI
|
||||
spec is in sync with the implementation. If any changes are made to the
|
||||
paths in the @route decorators or in the Responses returned by the
|
||||
methods (or any nested fields in the Responses), you will need to find the
|
||||
corresponding field of the OpenAPI yaml file and update it manually. Also,
|
||||
bump the version number in the yaml file and in this class's `get_version`.
|
||||
"""
|
||||
|
||||
# Time that we sleep while tailing logs while waiting for
|
||||
# the supervisor actor to start. We don't know which node
|
||||
# to read the logs from until then.
|
||||
WAIT_FOR_SUPERVISOR_ACTOR_INTERVAL_S = 1
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._job_info_client = None
|
||||
|
||||
# To make sure that the internal KV is initialized by getting the lazy property
|
||||
assert self.gcs_client is not None
|
||||
assert ray.experimental.internal_kv._internal_kv_initialized()
|
||||
|
||||
# It contains all `JobAgentSubmissionClient` that
|
||||
# `JobHead` has ever used, and will not be deleted
|
||||
# from it unless `JobAgentSubmissionClient` is no
|
||||
# longer available (the corresponding agent process is dead)
|
||||
# {node_id: JobAgentSubmissionClient}
|
||||
self._agents: Dict[NodeID, JobAgentSubmissionClient] = dict()
|
||||
|
||||
async def get_target_agent(
|
||||
self, timeout_s: float = WAIT_AVAILABLE_AGENT_TIMEOUT
|
||||
) -> JobAgentSubmissionClient:
|
||||
"""
|
||||
Get a `JobAgentSubmissionClient`, which is a client for interacting with jobs
|
||||
via an agent process.
|
||||
|
||||
Args:
|
||||
timeout_s: The timeout for the operation.
|
||||
|
||||
Returns:
|
||||
A `JobAgentSubmissionClient` for interacting with jobs via an agent process.
|
||||
|
||||
Raises:
|
||||
TimeoutError: If the operation times out.
|
||||
"""
|
||||
return await self._get_head_node_agent(timeout_s)
|
||||
|
||||
async def _get_head_node_agent_once(self) -> JobAgentSubmissionClient:
|
||||
head_node_id_hex = await get_head_node_id(self.gcs_client)
|
||||
|
||||
if not head_node_id_hex:
|
||||
raise Exception("Head node id has not yet been persisted in GCS")
|
||||
|
||||
head_node_id = NodeID.from_hex(head_node_id_hex)
|
||||
|
||||
if head_node_id not in self._agents:
|
||||
ip, http_port, _ = await self._fetch_agent_info(head_node_id)
|
||||
agent_http_address = f"http://{build_address(ip, http_port)}"
|
||||
self._agents[head_node_id] = JobAgentSubmissionClient(agent_http_address)
|
||||
|
||||
return self._agents[head_node_id]
|
||||
|
||||
async def _get_head_node_agent(self, timeout_s: float) -> JobAgentSubmissionClient:
|
||||
"""Retrieves HTTP client for `JobAgent` running on the Head node. If the head
|
||||
node does not have an agent, it will retry every
|
||||
`TRY_TO_GET_AGENT_INFO_INTERVAL_SECONDS` seconds indefinitely.
|
||||
|
||||
Args:
|
||||
timeout_s: The timeout for the operation.
|
||||
|
||||
Returns:
|
||||
A `JobAgentSubmissionClient` for interacting with jobs via the head node's agent process.
|
||||
|
||||
Raises:
|
||||
TimeoutError: If the operation times out.
|
||||
"""
|
||||
timeout_point = time.time() + timeout_s
|
||||
exception = None
|
||||
while time.time() < timeout_point:
|
||||
try:
|
||||
return await self._get_head_node_agent_once()
|
||||
except Exception as e:
|
||||
exception = e
|
||||
logger.exception(
|
||||
f"Failed to get head node agent, retrying in {TRY_TO_GET_AGENT_INFO_INTERVAL_SECONDS} seconds..."
|
||||
)
|
||||
await asyncio.sleep(TRY_TO_GET_AGENT_INFO_INTERVAL_SECONDS)
|
||||
raise TimeoutError(
|
||||
f"Failed to get head node agent within {timeout_s} seconds. The last exception is {exception}"
|
||||
)
|
||||
|
||||
async def _fetch_agent_info(self, target_node_id: NodeID) -> Tuple[str, int, int]:
|
||||
"""
|
||||
Fetches agent info by the Node ID. May raise exception if there's network error or the
|
||||
agent info is not found.
|
||||
|
||||
Returns: (ip, http_port, grpc_port)
|
||||
"""
|
||||
key = f"{DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX}{target_node_id.hex()}"
|
||||
value = await self.gcs_client.async_internal_kv_get(
|
||||
key,
|
||||
namespace=KV_NAMESPACE_DASHBOARD,
|
||||
timeout=GCS_RPC_TIMEOUT_SECONDS,
|
||||
)
|
||||
if not value:
|
||||
raise KeyError(
|
||||
f"Agent info not found in internal KV for node {target_node_id}. "
|
||||
"It's possible that the agent didn't launch successfully due to "
|
||||
"port conflicts or other issues. Please check `dashboard_agent.log` "
|
||||
"for more details."
|
||||
)
|
||||
return json.loads(value.decode())
|
||||
|
||||
@routes.get("/api/version")
|
||||
async def get_version(self, req: Request) -> Response:
|
||||
# NOTE(edoakes): CURRENT_VERSION should be bumped and checked on the
|
||||
# client when we have backwards-incompatible changes.
|
||||
resp = VersionResponse(
|
||||
version=CURRENT_VERSION,
|
||||
ray_version=ray.__version__,
|
||||
ray_commit=ray.__commit__,
|
||||
session_name=self.session_name,
|
||||
)
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(resp)),
|
||||
content_type="application/json",
|
||||
status=aiohttp.web.HTTPOk.status_code,
|
||||
)
|
||||
|
||||
@routes.get("/api/packages/{protocol}/{package_name}")
|
||||
async def get_package(self, req: Request) -> Response:
|
||||
package_uri = http_uri_components_to_uri(
|
||||
protocol=req.match_info["protocol"],
|
||||
package_name=req.match_info["package_name"],
|
||||
)
|
||||
|
||||
logger.debug(f"Adding temporary reference to package {package_uri}.")
|
||||
try:
|
||||
pin_runtime_env_uri(package_uri)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
if not package_exists(package_uri):
|
||||
return Response(
|
||||
text=f"Package {package_uri} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
|
||||
return Response()
|
||||
|
||||
@routes.put("/api/packages/{protocol}/{package_name}")
|
||||
async def upload_package(self, req: Request):
|
||||
package_uri = http_uri_components_to_uri(
|
||||
protocol=req.match_info["protocol"],
|
||||
package_name=req.match_info["package_name"],
|
||||
)
|
||||
logger.info(f"Uploading package {package_uri} to the GCS.")
|
||||
try:
|
||||
data = await req.read()
|
||||
await get_or_create_event_loop().run_in_executor(
|
||||
None,
|
||||
upload_package_to_gcs,
|
||||
package_uri,
|
||||
data,
|
||||
)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
return Response(status=aiohttp.web.HTTPOk.status_code)
|
||||
|
||||
@routes.post("/api/jobs/")
|
||||
async def submit_job(self, req: Request) -> Response:
|
||||
result = await parse_and_validate_request(req, JobSubmitRequest)
|
||||
# Request parsing failed, returned with Response object.
|
||||
if isinstance(result, Response):
|
||||
return result
|
||||
else:
|
||||
submit_request: JobSubmitRequest = result
|
||||
|
||||
try:
|
||||
job_agent_client = await self.get_target_agent()
|
||||
resp = await job_agent_client.submit_job_internal(submit_request)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
"Timed out waiting for an available job agent to submit the job."
|
||||
)
|
||||
return Response(
|
||||
text="No available agent to submit job, please try again later.",
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
except (TypeError, ValueError):
|
||||
logger.warning("Failed to submit job due to an invalid request.")
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to submit job due to unexpected exception.")
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(resp)),
|
||||
content_type="application/json",
|
||||
status=aiohttp.web.HTTPOk.status_code,
|
||||
)
|
||||
|
||||
@routes.post("/api/jobs/{job_or_submission_id}/stop")
|
||||
async def stop_job(self, req: Request) -> Response:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self.gcs_client,
|
||||
self._job_info_client,
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
if job.type is not JobType.SUBMISSION:
|
||||
return Response(
|
||||
text="Can only stop submission type jobs",
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
try:
|
||||
job_agent_client = await self.get_target_agent()
|
||||
resp = await job_agent_client.stop_job_internal(job.submission_id)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(resp)), content_type="application/json"
|
||||
)
|
||||
|
||||
@routes.delete("/api/jobs/{job_or_submission_id}")
|
||||
async def delete_job(self, req: Request) -> Response:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self.gcs_client,
|
||||
self._job_info_client,
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
if job.type is not JobType.SUBMISSION:
|
||||
return Response(
|
||||
text="Can only delete submission type jobs",
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
try:
|
||||
job_agent_client = await self.get_target_agent()
|
||||
resp = await job_agent_client.delete_job_internal(job.submission_id)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(resp)), content_type="application/json"
|
||||
)
|
||||
|
||||
@routes.get("/api/jobs/{job_or_submission_id}")
|
||||
async def get_job_info(self, req: Request) -> Response:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self.gcs_client,
|
||||
self._job_info_client,
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
|
||||
return Response(
|
||||
text=json.dumps(job.dict()),
|
||||
content_type="application/json",
|
||||
)
|
||||
|
||||
# TODO(rickyx): This endpoint's logic is also mirrored in state API's endpoint.
|
||||
# We should eventually unify the backend logic (and keep the logic in sync before
|
||||
# that).
|
||||
@routes.get("/api/jobs/")
|
||||
async def list_jobs(self, req: Request) -> Response:
|
||||
(driver_jobs, submission_job_drivers), submission_jobs = await asyncio.gather(
|
||||
get_driver_jobs(self.gcs_client), self._job_info_client.get_all_jobs()
|
||||
)
|
||||
|
||||
submission_jobs = [
|
||||
JobDetails(
|
||||
**dataclasses.asdict(job),
|
||||
submission_id=submission_id,
|
||||
job_id=submission_job_drivers.get(submission_id).id
|
||||
if submission_id in submission_job_drivers
|
||||
else None,
|
||||
driver_info=submission_job_drivers.get(submission_id),
|
||||
type=JobType.SUBMISSION,
|
||||
)
|
||||
for submission_id, job in submission_jobs.items()
|
||||
]
|
||||
return Response(
|
||||
text=json.dumps(
|
||||
[
|
||||
*[submission_job.dict() for submission_job in submission_jobs],
|
||||
*[job_info.dict() for job_info in driver_jobs.values()],
|
||||
]
|
||||
),
|
||||
content_type="application/json",
|
||||
)
|
||||
|
||||
@routes.get("/api/jobs/{job_or_submission_id}/logs")
|
||||
async def get_job_logs(self, req: Request) -> Response:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self.gcs_client,
|
||||
self._job_info_client,
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
|
||||
if job.type is not JobType.SUBMISSION:
|
||||
return Response(
|
||||
text="Can only get logs of submission type jobs",
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
try:
|
||||
job_agent_client = self.get_job_driver_agent_client(job)
|
||||
payload = (
|
||||
await job_agent_client.get_job_logs_internal(job.submission_id)
|
||||
if job_agent_client
|
||||
else JobLogsResponse("")
|
||||
)
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(payload)),
|
||||
content_type="application/json",
|
||||
)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
@routes.get(
|
||||
"/api/jobs/{job_or_submission_id}/logs/tail", resp_type=ResponseType.WEBSOCKET
|
||||
)
|
||||
async def tail_job_logs(self, req: Request) -> StreamResponse:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self.gcs_client,
|
||||
self._job_info_client,
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
|
||||
if job.type is not JobType.SUBMISSION:
|
||||
return Response(
|
||||
text="Can only get logs of submission type jobs",
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
ws = aiohttp.web.WebSocketResponse()
|
||||
await ws.prepare(req)
|
||||
|
||||
driver_agent_http_address = None
|
||||
while driver_agent_http_address is None:
|
||||
job = await find_job_by_ids(
|
||||
self.gcs_client,
|
||||
self._job_info_client,
|
||||
job_or_submission_id,
|
||||
)
|
||||
driver_agent_http_address = job.driver_agent_http_address
|
||||
status = job.status
|
||||
if status.is_terminal() and driver_agent_http_address is None:
|
||||
# Job exited before supervisor actor started.
|
||||
return ws
|
||||
|
||||
await asyncio.sleep(self.WAIT_FOR_SUPERVISOR_ACTOR_INTERVAL_S)
|
||||
|
||||
job_agent_client = self.get_job_driver_agent_client(job)
|
||||
|
||||
async for lines in job_agent_client.tail_job_logs(job.submission_id):
|
||||
await ws.send_str(lines)
|
||||
|
||||
return ws
|
||||
|
||||
def get_job_driver_agent_client(
|
||||
self, job: JobDetails
|
||||
) -> Optional[JobAgentSubmissionClient]:
|
||||
if job.driver_agent_http_address is None:
|
||||
return None
|
||||
|
||||
driver_node_id = job.driver_node_id
|
||||
if driver_node_id not in self._agents:
|
||||
self._agents[driver_node_id] = JobAgentSubmissionClient(
|
||||
job.driver_agent_http_address
|
||||
)
|
||||
|
||||
return self._agents[driver_node_id]
|
||||
|
||||
@routes.get("/api/component_activities")
|
||||
async def get_component_activities(
|
||||
self, req: aiohttp.web.Request
|
||||
) -> aiohttp.web.Response:
|
||||
timeout = req.query.get("timeout", None)
|
||||
if timeout and timeout.isdigit():
|
||||
timeout = int(timeout)
|
||||
else:
|
||||
timeout = 30
|
||||
|
||||
# Get activity information for driver
|
||||
driver_activity_info = await self._get_job_activity_info(timeout=timeout)
|
||||
resp = {"driver": dict(driver_activity_info)}
|
||||
|
||||
if RAY_CLUSTER_ACTIVITY_HOOK in os.environ:
|
||||
try:
|
||||
cluster_activity_callable = load_class(
|
||||
os.environ[RAY_CLUSTER_ACTIVITY_HOOK]
|
||||
)
|
||||
external_activity_output = cluster_activity_callable()
|
||||
assert isinstance(external_activity_output, dict), (
|
||||
f"Output of hook {os.environ[RAY_CLUSTER_ACTIVITY_HOOK]} "
|
||||
"should be Dict[str, RayActivityResponse]. Got "
|
||||
f"output: {external_activity_output}"
|
||||
)
|
||||
for component_type in external_activity_output:
|
||||
try:
|
||||
component_activity_output = external_activity_output[
|
||||
component_type
|
||||
]
|
||||
# Parse and validate output to type RayActivityResponse
|
||||
component_activity_output = RayActivityResponse(
|
||||
**dict(component_activity_output)
|
||||
)
|
||||
resp[component_type] = dict(component_activity_output)
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Failed to get activity status of {component_type} "
|
||||
f"from user hook {os.environ[RAY_CLUSTER_ACTIVITY_HOOK]}."
|
||||
)
|
||||
resp[component_type] = {
|
||||
"is_active": RayActivityStatus.ERROR,
|
||||
"reason": repr(e),
|
||||
"timestamp": datetime.now().timestamp(),
|
||||
}
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
"Failed to get activity status from user "
|
||||
f"hook {os.environ[RAY_CLUSTER_ACTIVITY_HOOK]}."
|
||||
)
|
||||
resp["external_component"] = {
|
||||
"is_active": RayActivityStatus.ERROR,
|
||||
"reason": repr(e),
|
||||
"timestamp": datetime.now().timestamp(),
|
||||
}
|
||||
|
||||
return aiohttp.web.Response(
|
||||
text=json.dumps(resp),
|
||||
content_type="application/json",
|
||||
status=aiohttp.web.HTTPOk.status_code,
|
||||
)
|
||||
|
||||
async def _get_job_activity_info(self, timeout: int) -> RayActivityResponse:
|
||||
# Returns if there is Ray activity from drivers (job).
|
||||
# Drivers in namespaces that start with _ray_internal_ are not
|
||||
# considered activity.
|
||||
# This includes the _ray_internal_dashboard job that gets automatically
|
||||
# created with every cluster
|
||||
try:
|
||||
reply = await self.gcs_client.async_get_all_job_info(
|
||||
skip_submission_job_info_field=True,
|
||||
skip_is_running_tasks_field=True,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
num_active_drivers = 0
|
||||
latest_job_end_time = 0
|
||||
for job_table_entry in reply.values():
|
||||
is_dead = bool(job_table_entry.is_dead)
|
||||
in_internal_namespace = job_table_entry.config.ray_namespace.startswith(
|
||||
"_ray_internal_"
|
||||
)
|
||||
latest_job_end_time = (
|
||||
max(latest_job_end_time, job_table_entry.end_time)
|
||||
if job_table_entry.end_time
|
||||
else latest_job_end_time
|
||||
)
|
||||
if not is_dead and not in_internal_namespace:
|
||||
num_active_drivers += 1
|
||||
|
||||
current_timestamp = datetime.now().timestamp()
|
||||
# Latest job end time must be before or equal to the current timestamp.
|
||||
# Job end times may be provided in epoch milliseconds. Check if this
|
||||
# is true, and convert to seconds
|
||||
if latest_job_end_time > current_timestamp:
|
||||
latest_job_end_time = latest_job_end_time / 1000
|
||||
assert current_timestamp >= latest_job_end_time, (
|
||||
f"Most recent job end time {latest_job_end_time} must be "
|
||||
f"before or equal to the current timestamp {current_timestamp}"
|
||||
)
|
||||
|
||||
is_active = (
|
||||
RayActivityStatus.ACTIVE
|
||||
if num_active_drivers > 0
|
||||
else RayActivityStatus.INACTIVE
|
||||
)
|
||||
return RayActivityResponse(
|
||||
is_active=is_active,
|
||||
reason=f"Number of active drivers: {num_active_drivers}"
|
||||
if num_active_drivers
|
||||
else None,
|
||||
timestamp=current_timestamp,
|
||||
# If latest_job_end_time == 0, no jobs have finished yet so don't
|
||||
# populate last_activity_at
|
||||
last_activity_at=latest_job_end_time if latest_job_end_time else None,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("Failed to get activity status of Ray drivers.")
|
||||
return RayActivityResponse(
|
||||
is_active=RayActivityStatus.ERROR,
|
||||
reason=repr(e),
|
||||
timestamp=datetime.now().timestamp(),
|
||||
)
|
||||
|
||||
async def run(self):
|
||||
await super().run()
|
||||
if not self._job_info_client:
|
||||
self._job_info_client = JobInfoStorageClient(self.gcs_client)
|
||||
@@ -0,0 +1,57 @@
|
||||
import os
|
||||
from typing import AsyncIterator, List, Tuple
|
||||
|
||||
import ray
|
||||
from ray.dashboard.modules.job.common import JOB_LOGS_PATH_TEMPLATE
|
||||
from ray.dashboard.modules.job.utils import fast_tail_last_n_lines, file_tail_iterator
|
||||
|
||||
|
||||
class JobLogStorageClient:
|
||||
"""
|
||||
Disk storage for stdout / stderr of driver script logs.
|
||||
"""
|
||||
|
||||
# Number of last N lines to put in job message upon failure.
|
||||
NUM_LOG_LINES_ON_ERROR = 10
|
||||
# Maximum number of characters to print out of the logs to avoid
|
||||
# HUGE log outputs that bring down the api server
|
||||
MAX_LOG_SIZE = 20000
|
||||
|
||||
def get_logs(self, job_id: str) -> str:
|
||||
try:
|
||||
with open(self.get_log_file_path(job_id), "r") as f:
|
||||
return f.read()
|
||||
except FileNotFoundError:
|
||||
return ""
|
||||
|
||||
def tail_logs(self, job_id: str) -> AsyncIterator[List[str]]:
|
||||
return file_tail_iterator(self.get_log_file_path(job_id))
|
||||
|
||||
async def get_last_n_log_lines(
|
||||
self, job_id: str, num_log_lines: int = NUM_LOG_LINES_ON_ERROR
|
||||
) -> str:
|
||||
"""Returns the last MAX_LOG_SIZE (20000) characters in the last ``num_log_lines`` lines.
|
||||
|
||||
Args:
|
||||
job_id: The id of the job whose logs we want to return
|
||||
num_log_lines: The number of lines to return.
|
||||
|
||||
Returns:
|
||||
Up to ``MAX_LOG_SIZE`` characters drawn from the last
|
||||
``num_log_lines`` lines of the job's log file.
|
||||
"""
|
||||
return fast_tail_last_n_lines(
|
||||
path=self.get_log_file_path(job_id),
|
||||
num_lines=num_log_lines,
|
||||
max_chars=self.MAX_LOG_SIZE,
|
||||
)
|
||||
|
||||
def get_log_file_path(self, job_id: str) -> Tuple[str, str]:
|
||||
"""
|
||||
Get the file path to the logs of a given job. Example:
|
||||
/tmp/ray/session_date/logs/job-driver-{job_id}.log
|
||||
"""
|
||||
return os.path.join(
|
||||
ray._private.worker._global_node.get_logs_dir_path(),
|
||||
JOB_LOGS_PATH_TEMPLATE.format(submission_id=job_id),
|
||||
)
|
||||
@@ -0,0 +1,707 @@
|
||||
import asyncio
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import string
|
||||
import time
|
||||
import traceback
|
||||
from typing import Any, AsyncIterator, Dict, Optional, Union
|
||||
|
||||
import ray
|
||||
import ray._private.ray_constants as ray_constants
|
||||
from ray._common.utils import Timer, run_background_task
|
||||
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.event.event_logger import get_event_logger
|
||||
from ray._private.label_utils import validate_label_selector
|
||||
from ray._raylet import GcsClient
|
||||
from ray.actor import ActorHandle
|
||||
from ray.core.generated.event_pb2 import Event
|
||||
from ray.dashboard.consts import (
|
||||
DEFAULT_JOB_START_TIMEOUT_SECONDS,
|
||||
RAY_JOB_ALLOW_DRIVER_ON_WORKER_NODES_ENV_VAR,
|
||||
RAY_JOB_START_TIMEOUT_SECONDS_ENV_VAR,
|
||||
RAY_STREAM_RUNTIME_ENV_LOG_TO_JOB_DRIVER_LOG_ENV_VAR,
|
||||
)
|
||||
from ray.dashboard.modules.job.common import (
|
||||
JOB_ACTOR_NAME_TEMPLATE,
|
||||
SUPERVISOR_ACTOR_RAY_NAMESPACE,
|
||||
JobInfo,
|
||||
JobInfoStorageClient,
|
||||
)
|
||||
from ray.dashboard.modules.job.job_log_storage_client import JobLogStorageClient
|
||||
from ray.dashboard.modules.job.job_supervisor import JobSupervisor
|
||||
from ray.dashboard.utils import close_logger_file_descriptor, get_head_node_id
|
||||
from ray.exceptions import ActorDiedError, ActorUnschedulableError, RuntimeEnvSetupError
|
||||
from ray.job_submission import JobErrorType, JobStatus
|
||||
from ray.runtime_env import RuntimeEnvConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def generate_job_id() -> str:
|
||||
"""Returns a job_id of the form 'raysubmit_XYZ'.
|
||||
|
||||
Prefixed with 'raysubmit' to avoid confusion with Ray JobID (driver ID).
|
||||
"""
|
||||
rand = random.SystemRandom()
|
||||
possible_characters = list(
|
||||
set(string.ascii_letters + string.digits)
|
||||
- {"I", "l", "o", "O", "0"} # No confusing characters
|
||||
)
|
||||
id_part = "".join(rand.choices(possible_characters, k=16))
|
||||
return f"raysubmit_{id_part}"
|
||||
|
||||
|
||||
class JobManager:
|
||||
"""Provide python APIs for job submission and management.
|
||||
|
||||
It does not provide persistence, all info will be lost if the cluster
|
||||
goes down.
|
||||
"""
|
||||
|
||||
# Time that we will sleep while tailing logs if no new log line is
|
||||
# available.
|
||||
LOG_TAIL_SLEEP_S = 1
|
||||
JOB_MONITOR_LOOP_PERIOD_S = 1
|
||||
WAIT_FOR_ACTOR_DEATH_TIMEOUT_S = 0.1
|
||||
|
||||
def __init__(
|
||||
self, gcs_client: GcsClient, logs_dir: str, timeout_check_timer: Timer = None
|
||||
):
|
||||
self._gcs_client = gcs_client
|
||||
self._logs_dir = logs_dir
|
||||
self._job_info_client = JobInfoStorageClient(gcs_client, logs_dir)
|
||||
self._gcs_address = gcs_client.address
|
||||
self._cluster_id_hex = gcs_client.cluster_id.hex()
|
||||
self._log_client = JobLogStorageClient()
|
||||
self._supervisor_actor_cls = ray.remote(JobSupervisor)
|
||||
self._timeout_check_timer = timeout_check_timer or Timer()
|
||||
self.monitored_jobs = set()
|
||||
try:
|
||||
self.event_logger = get_event_logger(Event.SourceType.JOBS, logs_dir)
|
||||
except Exception:
|
||||
self.event_logger = None
|
||||
|
||||
self._recover_running_jobs_event = asyncio.Event()
|
||||
run_background_task(self._recover_running_jobs())
|
||||
|
||||
def _get_job_driver_logger(self, job_id: str) -> logging.Logger:
|
||||
"""Return job driver logger to log messages to the job driver log file.
|
||||
|
||||
If this function is called for the first time, configure the logger.
|
||||
"""
|
||||
job_driver_logger = logging.getLogger(f"{__name__}.driver-{job_id}")
|
||||
|
||||
# Configure the logger if it's not already configured.
|
||||
if not job_driver_logger.handlers:
|
||||
job_driver_log_path = self._log_client.get_log_file_path(job_id)
|
||||
job_driver_handler = logging.FileHandler(job_driver_log_path)
|
||||
job_driver_formatter = logging.Formatter(ray_constants.LOGGER_FORMAT)
|
||||
job_driver_handler.setFormatter(job_driver_formatter)
|
||||
job_driver_logger.addHandler(job_driver_handler)
|
||||
|
||||
return job_driver_logger
|
||||
|
||||
async def _recover_running_jobs(self):
|
||||
"""Recovers all running jobs from the status client.
|
||||
|
||||
For each job, we will spawn a coroutine to monitor it.
|
||||
Each will be added to self._running_jobs and reconciled.
|
||||
"""
|
||||
try:
|
||||
all_jobs = await self._job_info_client.get_all_jobs()
|
||||
for job_id, job_info in all_jobs.items():
|
||||
if not job_info.status.is_terminal():
|
||||
run_background_task(self._monitor_job(job_id))
|
||||
finally:
|
||||
# This event is awaited in `submit_job` to avoid race conditions between
|
||||
# recovery and new job submission, so it must always get set even if there
|
||||
# are exceptions.
|
||||
self._recover_running_jobs_event.set()
|
||||
|
||||
def _get_actor_for_job(self, job_id: str) -> Optional[ActorHandle]:
|
||||
try:
|
||||
return ray.get_actor(
|
||||
JOB_ACTOR_NAME_TEMPLATE.format(job_id=job_id),
|
||||
namespace=SUPERVISOR_ACTOR_RAY_NAMESPACE,
|
||||
)
|
||||
except ValueError: # Ray returns ValueError for nonexistent actor.
|
||||
return None
|
||||
|
||||
async def _monitor_job(
|
||||
self, job_id: str, job_supervisor: Optional[ActorHandle] = None
|
||||
):
|
||||
"""Monitors the specified job until it enters a terminal state.
|
||||
|
||||
This is necessary because we need to handle the case where the
|
||||
JobSupervisor dies unexpectedly.
|
||||
"""
|
||||
if job_id in self.monitored_jobs:
|
||||
logger.debug(f"Job {job_id} is already being monitored.")
|
||||
return
|
||||
|
||||
self.monitored_jobs.add(job_id)
|
||||
try:
|
||||
await self._monitor_job_internal(job_id, job_supervisor)
|
||||
except Exception as e:
|
||||
logger.error("Unhandled exception in job monitoring!", exc_info=e)
|
||||
raise e
|
||||
finally:
|
||||
self.monitored_jobs.remove(job_id)
|
||||
|
||||
async def _monitor_job_internal(
|
||||
self, job_id: str, job_supervisor: Optional[ActorHandle] = None
|
||||
):
|
||||
timeout = float(
|
||||
os.environ.get(
|
||||
RAY_JOB_START_TIMEOUT_SECONDS_ENV_VAR,
|
||||
DEFAULT_JOB_START_TIMEOUT_SECONDS,
|
||||
)
|
||||
)
|
||||
|
||||
job_status = None
|
||||
job_info = None
|
||||
ping_obj_ref = None
|
||||
|
||||
while True:
|
||||
try:
|
||||
# NOTE: Job monitoring loop sleeps before proceeding with monitoring
|
||||
# sequence to consolidate the control-flow of the pacing
|
||||
# in a single place, rather than having it spread across
|
||||
# many branches
|
||||
await asyncio.sleep(self.JOB_MONITOR_LOOP_PERIOD_S)
|
||||
|
||||
job_status = await self._job_info_client.get_status(
|
||||
job_id, timeout=None
|
||||
)
|
||||
if job_status == JobStatus.PENDING:
|
||||
# Compare the current time with the job start time.
|
||||
# If the job is still pending, we will set the status
|
||||
# to FAILED.
|
||||
if job_info is None:
|
||||
job_info = await self._job_info_client.get_info(
|
||||
job_id, timeout=None
|
||||
)
|
||||
|
||||
if (
|
||||
self._timeout_check_timer.time() - job_info.start_time / 1000
|
||||
> timeout
|
||||
):
|
||||
err_msg = (
|
||||
"Job supervisor actor failed to start within "
|
||||
f"{timeout} seconds. This timeout can be "
|
||||
f"configured by setting the environment "
|
||||
f"variable {RAY_JOB_START_TIMEOUT_SECONDS_ENV_VAR}."
|
||||
)
|
||||
resources_specified = (
|
||||
(
|
||||
job_info.entrypoint_num_cpus is not None
|
||||
and job_info.entrypoint_num_cpus > 0
|
||||
)
|
||||
or (
|
||||
job_info.entrypoint_num_gpus is not None
|
||||
and job_info.entrypoint_num_gpus > 0
|
||||
)
|
||||
or (
|
||||
job_info.entrypoint_memory is not None
|
||||
and job_info.entrypoint_memory > 0
|
||||
)
|
||||
or (
|
||||
job_info.entrypoint_resources is not None
|
||||
and len(job_info.entrypoint_resources) > 0
|
||||
)
|
||||
)
|
||||
if resources_specified:
|
||||
err_msg += (
|
||||
" This may be because the job entrypoint's specified "
|
||||
"resources (entrypoint_num_cpus, entrypoint_num_gpus, "
|
||||
"entrypoint_resources, entrypoint_memory)"
|
||||
"aren't available on the cluster."
|
||||
" Try checking the cluster's available resources with "
|
||||
"`ray status` and specifying fewer resources for the "
|
||||
"job entrypoint."
|
||||
)
|
||||
await self._job_info_client.put_status(
|
||||
job_id,
|
||||
JobStatus.FAILED,
|
||||
message=err_msg,
|
||||
error_type=JobErrorType.JOB_SUPERVISOR_ACTOR_START_TIMEOUT,
|
||||
timeout=None,
|
||||
)
|
||||
logger.error(err_msg)
|
||||
break
|
||||
|
||||
if job_supervisor is None:
|
||||
job_supervisor = self._get_actor_for_job(job_id)
|
||||
|
||||
if job_supervisor is None:
|
||||
if job_status == JobStatus.PENDING:
|
||||
# Maybe the job supervisor actor is not created yet.
|
||||
# We will wait for the next loop.
|
||||
continue
|
||||
else:
|
||||
# The job supervisor actor is not created, but the job
|
||||
# status is not PENDING. This means the job supervisor
|
||||
# actor is not created due to some unexpected errors.
|
||||
# We will set the job status to FAILED.
|
||||
logger.error(f"Failed to get job supervisor for job {job_id}.")
|
||||
await self._job_info_client.put_status(
|
||||
job_id,
|
||||
JobStatus.FAILED,
|
||||
message=(
|
||||
"Unexpected error occurred: "
|
||||
"failed to get job supervisor."
|
||||
),
|
||||
error_type=JobErrorType.JOB_SUPERVISOR_ACTOR_START_FAILURE,
|
||||
timeout=None,
|
||||
)
|
||||
break
|
||||
|
||||
# Check to see if `JobSupervisor` is alive and reachable
|
||||
if ping_obj_ref is None:
|
||||
ping_obj_ref = job_supervisor.ping.options(
|
||||
max_task_retries=-1
|
||||
).remote()
|
||||
ready, _ = ray.wait([ping_obj_ref], timeout=0)
|
||||
if ready:
|
||||
ray.get(ping_obj_ref)
|
||||
ping_obj_ref = None
|
||||
else:
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
job_status = await self._job_info_client.get_status(
|
||||
job_id, timeout=None
|
||||
)
|
||||
target_job_error_message = ""
|
||||
target_job_error_type: Optional[JobErrorType] = None
|
||||
if job_status is not None and job_status.is_terminal():
|
||||
# If the job is already in a terminal state, then the actor
|
||||
# exiting is expected.
|
||||
pass
|
||||
else:
|
||||
if isinstance(e, RuntimeEnvSetupError):
|
||||
logger.error(f"Failed to set up runtime_env for job {job_id}.")
|
||||
|
||||
target_job_error_message = f"runtime_env setup failed: {e}"
|
||||
target_job_error_type = JobErrorType.RUNTIME_ENV_SETUP_FAILURE
|
||||
|
||||
elif isinstance(e, ActorUnschedulableError):
|
||||
logger.error(
|
||||
f"Failed to schedule job {job_id} because the supervisor "
|
||||
f"actor could not be scheduled: {e}"
|
||||
)
|
||||
|
||||
target_job_error_message = (
|
||||
f"Job supervisor actor could not be scheduled: {e}"
|
||||
)
|
||||
target_job_error_type = (
|
||||
JobErrorType.JOB_SUPERVISOR_ACTOR_UNSCHEDULABLE
|
||||
)
|
||||
|
||||
elif isinstance(e, ActorDiedError):
|
||||
logger.error(f"Job supervisor actor for {job_id} died: {e}")
|
||||
target_job_error_message = f"Job supervisor actor died: {e}"
|
||||
target_job_error_type = JobErrorType.JOB_SUPERVISOR_ACTOR_DIED
|
||||
|
||||
else:
|
||||
logger.error(
|
||||
f"Job monitoring for job {job_id} failed "
|
||||
f"unexpectedly: {e}.",
|
||||
exc_info=e,
|
||||
)
|
||||
|
||||
target_job_error_message = f"Unexpected error occurred: {e}"
|
||||
target_job_error_type = (
|
||||
JobErrorType.JOB_SUPERVISOR_ACTOR_UNKNOWN_FAILURE
|
||||
)
|
||||
|
||||
job_status = JobStatus.FAILED
|
||||
await self._job_info_client.put_status(
|
||||
job_id,
|
||||
job_status,
|
||||
message=target_job_error_message,
|
||||
error_type=target_job_error_type
|
||||
or JobErrorType.JOB_SUPERVISOR_ACTOR_UNKNOWN_FAILURE,
|
||||
timeout=None,
|
||||
)
|
||||
|
||||
# Log error message to the job driver file for easy access.
|
||||
if target_job_error_message:
|
||||
log_path = self._log_client.get_log_file_path(job_id)
|
||||
os.makedirs(os.path.dirname(log_path), exist_ok=True)
|
||||
with open(log_path, "a") as log_file:
|
||||
log_file.write(target_job_error_message)
|
||||
|
||||
# Log events
|
||||
if self.event_logger:
|
||||
event_log = (
|
||||
f"Completed a ray job {job_id} with a status {job_status}."
|
||||
)
|
||||
if target_job_error_message:
|
||||
event_log += f" {target_job_error_message}"
|
||||
self.event_logger.error(event_log, submission_id=job_id)
|
||||
else:
|
||||
self.event_logger.info(event_log, submission_id=job_id)
|
||||
|
||||
break
|
||||
|
||||
# Kill the actor defensively to avoid leaking actors in unexpected error cases.
|
||||
if job_supervisor is None:
|
||||
job_supervisor = self._get_actor_for_job(job_id)
|
||||
if job_supervisor is not None:
|
||||
ray.kill(job_supervisor, no_restart=True)
|
||||
|
||||
def _handle_supervisor_startup(self, job_id: str, result: Optional[Exception]):
|
||||
"""Handle the result of starting a job supervisor actor.
|
||||
|
||||
If started successfully, result should be None. Otherwise it should be
|
||||
an Exception.
|
||||
|
||||
On failure, the job will be marked failed with a relevant error
|
||||
message.
|
||||
"""
|
||||
if result is None:
|
||||
return
|
||||
|
||||
def _get_supervisor_runtime_env(
|
||||
self,
|
||||
user_runtime_env: Dict[str, Any],
|
||||
submission_id: str,
|
||||
resources_specified: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
"""Configure and return the runtime_env for the supervisor actor.
|
||||
|
||||
Args:
|
||||
user_runtime_env: The runtime_env specified by the user.
|
||||
submission_id: The submission id of the job; used to derive the log
|
||||
file path piped into the runtime env config.
|
||||
resources_specified: Whether the user specified resources in the
|
||||
submit_job() call. If so, we will skip the workaround introduced
|
||||
in #24546 for GPU detection and just use the user's resource
|
||||
requests, so that the behavior matches that of the user specifying
|
||||
resources for any other actor.
|
||||
|
||||
Returns:
|
||||
The runtime_env for the supervisor actor.
|
||||
"""
|
||||
# Make a copy to avoid mutating passed runtime_env.
|
||||
runtime_env = (
|
||||
copy.deepcopy(user_runtime_env) if user_runtime_env is not None else {}
|
||||
)
|
||||
|
||||
# NOTE(edoakes): Can't use .get(, {}) here because we need to handle the case
|
||||
# where env_vars is explicitly set to `None`.
|
||||
env_vars = runtime_env.get("env_vars")
|
||||
if env_vars is None:
|
||||
env_vars = {}
|
||||
|
||||
env_vars[ray_constants.RAY_WORKER_NICENESS] = "0"
|
||||
|
||||
if not resources_specified:
|
||||
# Don't set CUDA_VISIBLE_DEVICES for the supervisor actor so the
|
||||
# driver can use GPUs if it wants to. This will be removed from
|
||||
# the driver's runtime_env so it isn't inherited by tasks & actors.
|
||||
env_vars[NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR] = "1"
|
||||
env_vars[NOSET_ASCEND_RT_VISIBLE_DEVICES_ENV_VAR] = "1"
|
||||
|
||||
runtime_env["env_vars"] = env_vars
|
||||
|
||||
if os.getenv(RAY_STREAM_RUNTIME_ENV_LOG_TO_JOB_DRIVER_LOG_ENV_VAR, "0") == "1":
|
||||
config = runtime_env.get("config")
|
||||
# Empty fields may be set to None, so we need to check for None explicitly.
|
||||
if config is None:
|
||||
config = RuntimeEnvConfig()
|
||||
config["log_files"] = [self._log_client.get_log_file_path(submission_id)]
|
||||
runtime_env["config"] = config
|
||||
return runtime_env
|
||||
|
||||
async def _get_label_selector(self, resources_specified: bool) -> Dict:
|
||||
"""Determine the scheduling strategy for the job using a label selector.
|
||||
|
||||
If resources_specified is true, or if the environment variable is set to
|
||||
allow the job to run on worker nodes, we will not use any label constraints.
|
||||
Otherwise, we will force the job to use the head node via a label selector
|
||||
specifying the head node id.
|
||||
|
||||
Args:
|
||||
resources_specified: Whether the job specified any resources
|
||||
(CPUs, GPUs, or custom resources).
|
||||
|
||||
Returns:
|
||||
The label selector to use for the job.
|
||||
"""
|
||||
if resources_specified:
|
||||
return {}
|
||||
|
||||
if os.environ.get(RAY_JOB_ALLOW_DRIVER_ON_WORKER_NODES_ENV_VAR, "0") == "1":
|
||||
logger.info(
|
||||
f"{RAY_JOB_ALLOW_DRIVER_ON_WORKER_NODES_ENV_VAR} was set to 1. "
|
||||
"Using Ray's default actor scheduling strategy for the job "
|
||||
"driver instead of running it on the head node via a label selector."
|
||||
)
|
||||
return {}
|
||||
|
||||
# If the user did not specify any resources or set the driver on worker nodes
|
||||
# env var, we will run the driver on the head node.
|
||||
|
||||
head_node_id = await get_head_node_id(self._gcs_client)
|
||||
if head_node_id is None:
|
||||
logger.info(
|
||||
"Head node ID not found in GCS. Using Ray's default actor "
|
||||
"scheduling strategy for the job driver instead of running "
|
||||
"it on the head node via a label selector."
|
||||
)
|
||||
return {}
|
||||
|
||||
logger.info(
|
||||
"Head node ID found in GCS; scheduling job driver on "
|
||||
f"head node {head_node_id} using a label selector"
|
||||
)
|
||||
return {ray._raylet.RAY_NODE_ID_KEY: head_node_id}
|
||||
|
||||
async def submit_job(
|
||||
self,
|
||||
*,
|
||||
entrypoint: str,
|
||||
submission_id: Optional[str] = None,
|
||||
runtime_env: Optional[Dict[str, Any]] = None,
|
||||
metadata: Optional[Dict[str, 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,
|
||||
_start_signal_actor: Optional[ActorHandle] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Job execution happens asynchronously.
|
||||
|
||||
1) Generate a new unique id for this job submission, each call of this
|
||||
method assumes they're independent submission with its own new
|
||||
ID, job supervisor actor, and child process.
|
||||
2) Create new detached actor with same runtime_env as job spec
|
||||
|
||||
Actual setting up runtime_env, subprocess group, driver command
|
||||
execution, subprocess cleaning up and running status update to GCS
|
||||
is all handled by job supervisor actor.
|
||||
|
||||
Args:
|
||||
entrypoint: Driver command to execute in subprocess shell.
|
||||
Represents the entrypoint to start user application.
|
||||
submission_id: Optional caller-provided submission id. When None, a
|
||||
new id is generated via ``generate_job_id()``.
|
||||
runtime_env: Runtime environment used to execute driver command,
|
||||
which could contain its own ray.init() to configure runtime
|
||||
env at ray cluster, task and actor level.
|
||||
metadata: Support passing arbitrary data to driver command in
|
||||
case needed.
|
||||
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 entrypoint command, separately from any tasks or actors launched
|
||||
by it. Defaults to 0.
|
||||
entrypoint_memory: The amount of total available memory for workers
|
||||
requesting memory the entrypoint command, separately from any tasks
|
||||
or actors launched by it. Defaults to 0.
|
||||
entrypoint_resources: The quantity of various custom resources
|
||||
to reserve for the entrypoint command, separately from any tasks or
|
||||
actors launched by it.
|
||||
entrypoint_label_selector: Label selector for the entrypoint command.
|
||||
_start_signal_actor: Used in testing only to capture state
|
||||
transitions between PENDING -> RUNNING. Regular user shouldn't
|
||||
need this.
|
||||
|
||||
Returns:
|
||||
job_id: Generated uuid for further job management. Only valid
|
||||
within the same ray cluster.
|
||||
"""
|
||||
if entrypoint_num_cpus is None:
|
||||
entrypoint_num_cpus = 0
|
||||
if entrypoint_num_gpus is None:
|
||||
entrypoint_num_gpus = 0
|
||||
if entrypoint_memory is None:
|
||||
entrypoint_memory = 0
|
||||
if submission_id is None:
|
||||
submission_id = generate_job_id()
|
||||
|
||||
# Wait for `_recover_running_jobs` to run before accepting submissions to
|
||||
# avoid duplicate monitoring of the same job.
|
||||
await self._recover_running_jobs_event.wait()
|
||||
|
||||
logger.info(f"Starting job with submission_id: {submission_id}")
|
||||
if entrypoint_label_selector:
|
||||
error_message = validate_label_selector(entrypoint_label_selector)
|
||||
if error_message:
|
||||
raise ValueError(error_message)
|
||||
job_info = JobInfo(
|
||||
entrypoint=entrypoint,
|
||||
status=JobStatus.PENDING,
|
||||
start_time=int(time.time() * 1000),
|
||||
metadata=metadata,
|
||||
runtime_env=runtime_env,
|
||||
entrypoint_num_cpus=entrypoint_num_cpus,
|
||||
entrypoint_num_gpus=entrypoint_num_gpus,
|
||||
entrypoint_memory=entrypoint_memory,
|
||||
entrypoint_resources=entrypoint_resources,
|
||||
)
|
||||
new_key_added = await self._job_info_client.put_info(
|
||||
submission_id, job_info, overwrite=False
|
||||
)
|
||||
if not new_key_added:
|
||||
raise ValueError(
|
||||
f"Job with submission_id {submission_id} already exists. "
|
||||
"Please use a different submission_id."
|
||||
)
|
||||
|
||||
driver_logger = self._get_job_driver_logger(submission_id)
|
||||
# Wait for the actor to start up asynchronously so this call always
|
||||
# returns immediately and we can catch errors with the actor starting
|
||||
# up.
|
||||
try:
|
||||
resources_specified = any(
|
||||
[
|
||||
entrypoint_num_cpus is not None and entrypoint_num_cpus > 0,
|
||||
entrypoint_num_gpus is not None and entrypoint_num_gpus > 0,
|
||||
entrypoint_memory is not None and entrypoint_memory > 0,
|
||||
entrypoint_resources not in [None, {}],
|
||||
entrypoint_label_selector not in [None, {}],
|
||||
]
|
||||
)
|
||||
label_selector = await self._get_label_selector(resources_specified)
|
||||
if entrypoint_label_selector:
|
||||
label_selector = {**label_selector, **entrypoint_label_selector}
|
||||
|
||||
if self.event_logger:
|
||||
self.event_logger.info(
|
||||
f"Started a ray job {submission_id}.", submission_id=submission_id
|
||||
)
|
||||
|
||||
driver_logger.info("Runtime env is setting up.")
|
||||
supervisor_options = dict(
|
||||
lifetime="detached",
|
||||
name=JOB_ACTOR_NAME_TEMPLATE.format(job_id=submission_id),
|
||||
num_cpus=entrypoint_num_cpus,
|
||||
num_gpus=entrypoint_num_gpus,
|
||||
memory=entrypoint_memory,
|
||||
resources=entrypoint_resources,
|
||||
label_selector=label_selector,
|
||||
runtime_env=self._get_supervisor_runtime_env(
|
||||
runtime_env, submission_id, resources_specified
|
||||
),
|
||||
namespace=SUPERVISOR_ACTOR_RAY_NAMESPACE,
|
||||
# Don't pollute task events with system actor tasks that users don't
|
||||
# know about.
|
||||
enable_task_events=False,
|
||||
)
|
||||
supervisor = self._supervisor_actor_cls.options(
|
||||
**supervisor_options
|
||||
).remote(
|
||||
submission_id,
|
||||
entrypoint,
|
||||
metadata or {},
|
||||
self._gcs_address,
|
||||
self._cluster_id_hex,
|
||||
self._logs_dir,
|
||||
)
|
||||
supervisor.run.remote(
|
||||
_start_signal_actor=_start_signal_actor,
|
||||
resources_specified=resources_specified,
|
||||
)
|
||||
|
||||
# Monitor the job in the background so we can detect errors without
|
||||
# requiring a client to poll.
|
||||
run_background_task(
|
||||
self._monitor_job(submission_id, job_supervisor=supervisor)
|
||||
)
|
||||
except Exception as e:
|
||||
tb_str = traceback.format_exc()
|
||||
driver_logger.warning(
|
||||
f"Failed to start supervisor actor for job {submission_id}: '{e}'"
|
||||
f". Full traceback:\n{tb_str}"
|
||||
)
|
||||
await self._job_info_client.put_status(
|
||||
submission_id,
|
||||
JobStatus.FAILED,
|
||||
message=(
|
||||
f"Failed to start supervisor actor {submission_id}: '{e}'"
|
||||
f". Full traceback:\n{tb_str}"
|
||||
),
|
||||
error_type=JobErrorType.JOB_SUPERVISOR_ACTOR_START_FAILURE,
|
||||
)
|
||||
finally:
|
||||
close_logger_file_descriptor(driver_logger)
|
||||
|
||||
return submission_id
|
||||
|
||||
def stop_job(self, job_id) -> bool:
|
||||
"""Request a job to exit, fire and forget.
|
||||
|
||||
Returns whether or not the job was running.
|
||||
"""
|
||||
job_supervisor_actor = self._get_actor_for_job(job_id)
|
||||
if job_supervisor_actor is not None:
|
||||
# Actor is still alive, signal it to stop the driver, fire and
|
||||
# forget
|
||||
job_supervisor_actor.stop.remote()
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
async def delete_job(self, job_id):
|
||||
"""Delete a job's info and metadata from the cluster."""
|
||||
job_status = await self._job_info_client.get_status(job_id)
|
||||
|
||||
if job_status is None or not job_status.is_terminal():
|
||||
raise RuntimeError(
|
||||
f"Attempted to delete job '{job_id}', "
|
||||
f"but it is in a non-terminal state {job_status}."
|
||||
)
|
||||
|
||||
await self._job_info_client.delete_info(job_id)
|
||||
return True
|
||||
|
||||
def job_info_client(self) -> JobInfoStorageClient:
|
||||
return self._job_info_client
|
||||
|
||||
async def get_job_status(self, job_id: str) -> Optional[JobStatus]:
|
||||
"""Get latest status of a job."""
|
||||
return await self._job_info_client.get_status(job_id)
|
||||
|
||||
async def get_job_info(self, job_id: str) -> Optional[JobInfo]:
|
||||
"""Get latest info of a job."""
|
||||
return await self._job_info_client.get_info(job_id)
|
||||
|
||||
async def list_jobs(self) -> Dict[str, JobInfo]:
|
||||
"""Get info for all jobs."""
|
||||
return await self._job_info_client.get_all_jobs()
|
||||
|
||||
def get_job_logs(self, job_id: str) -> str:
|
||||
"""Get all logs produced by a job."""
|
||||
return self._log_client.get_logs(job_id)
|
||||
|
||||
async def tail_job_logs(self, job_id: str) -> AsyncIterator[str]:
|
||||
"""Return an iterator following the logs of a job."""
|
||||
if await self.get_job_status(job_id) is None:
|
||||
raise RuntimeError(f"Job '{job_id}' does not exist.")
|
||||
|
||||
job_finished = False
|
||||
async for lines in self._log_client.tail_logs(job_id):
|
||||
if lines is None:
|
||||
if job_finished:
|
||||
# Job has already finished and we have read EOF afterwards,
|
||||
# it's guaranteed that we won't get any more logs.
|
||||
return
|
||||
else:
|
||||
status = await self.get_job_status(job_id)
|
||||
if status.is_terminal():
|
||||
job_finished = True
|
||||
# Continue tailing logs generated between the
|
||||
# last EOF read and the finish of the job.
|
||||
|
||||
await asyncio.sleep(self.LOG_TAIL_SLEEP_S)
|
||||
else:
|
||||
yield "".join(lines)
|
||||
@@ -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()
|
||||
@@ -0,0 +1,110 @@
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from ray._common.pydantic_compat import PYDANTIC_INSTALLED, BaseModel, Field
|
||||
from ray.dashboard.modules.job.common import JobStatus
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
# Pydantic is not part of the minimal Ray installation.
|
||||
if PYDANTIC_INSTALLED:
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class DriverInfo(BaseModel):
|
||||
"""A class for recording information about the driver related to the job."""
|
||||
|
||||
id: str = Field(..., description="The id of the driver")
|
||||
node_ip_address: str = Field(
|
||||
..., description="The IP address of the node the driver is running on."
|
||||
)
|
||||
pid: str = Field(
|
||||
..., description="The PID of the worker process the driver is using."
|
||||
)
|
||||
# TODO(aguo): Add node_id as a field.
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class JobType(str, Enum):
|
||||
"""An enumeration for describing the different job types.
|
||||
|
||||
NOTE:
|
||||
This field is still experimental and may change in the future.
|
||||
"""
|
||||
|
||||
#: A job that was initiated by the Ray Jobs API.
|
||||
SUBMISSION = "SUBMISSION"
|
||||
#: A job that was initiated by a driver script.
|
||||
DRIVER = "DRIVER"
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class JobDetails(BaseModel):
|
||||
"""
|
||||
Job data with extra details about its driver and its submission.
|
||||
"""
|
||||
|
||||
type: JobType = Field(..., description="The type of job.")
|
||||
job_id: Optional[str] = Field(
|
||||
None,
|
||||
description="The job ID. An ID that is created for every job that is "
|
||||
"launched in Ray. This can be used to fetch data about jobs using Ray "
|
||||
"Core APIs.",
|
||||
)
|
||||
submission_id: Optional[str] = Field(
|
||||
None,
|
||||
description="A submission ID is an ID created for every job submitted via"
|
||||
"the Ray Jobs API. It can "
|
||||
"be used to fetch data about jobs using the Ray Jobs API.",
|
||||
)
|
||||
driver_info: Optional[DriverInfo] = Field(
|
||||
None,
|
||||
description="The driver related to this job. For jobs submitted via "
|
||||
"the Ray Jobs API, "
|
||||
"it is the last driver launched by that job submission, "
|
||||
"or None if there is no driver.",
|
||||
)
|
||||
|
||||
# The following fields are copied from JobInfo.
|
||||
# TODO(aguo): Inherit from JobInfo once it's migrated to pydantic.
|
||||
status: JobStatus = Field(..., description="The status of the job.")
|
||||
entrypoint: str = Field(..., description="The entrypoint command for this job.")
|
||||
message: Optional[str] = Field(
|
||||
None, description="A message describing the status in more detail."
|
||||
)
|
||||
error_type: Optional[str] = Field(
|
||||
None, description="Internal error or user script error."
|
||||
)
|
||||
start_time: Optional[int] = Field(
|
||||
None,
|
||||
description="The time when the job was started. A Unix timestamp in ms.",
|
||||
)
|
||||
end_time: Optional[int] = Field(
|
||||
None,
|
||||
description="The time when the job moved into a terminal state. "
|
||||
"A Unix timestamp in ms.",
|
||||
)
|
||||
metadata: Optional[Dict[str, str]] = Field(
|
||||
None, description="Arbitrary user-provided metadata for the job."
|
||||
)
|
||||
runtime_env: Optional[Dict[str, Any]] = Field(
|
||||
None, description="The runtime environment for the job."
|
||||
)
|
||||
# the node info where the driver running on.
|
||||
# - driver_agent_http_address: this node's agent http address
|
||||
# - driver_node_id: this node's id.
|
||||
driver_agent_http_address: Optional[str] = Field(
|
||||
None,
|
||||
description="The HTTP address of the JobAgent on the node the job "
|
||||
"entrypoint command is running on.",
|
||||
)
|
||||
driver_node_id: Optional[str] = Field(
|
||||
None,
|
||||
description="The ID of the node the job entrypoint command is running on.",
|
||||
)
|
||||
driver_exit_code: Optional[int] = Field(
|
||||
None,
|
||||
description="The driver process exit code after the driver executed. "
|
||||
"Return None if driver doesn't finish executing.",
|
||||
)
|
||||
|
||||
else:
|
||||
DriverInfo = None
|
||||
JobType = None
|
||||
JobDetails = None
|
||||
@@ -0,0 +1,542 @@
|
||||
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
|
||||
@@ -0,0 +1,26 @@
|
||||
import requests
|
||||
|
||||
import ray
|
||||
|
||||
ray.init()
|
||||
|
||||
|
||||
@ray.remote
|
||||
class Counter:
|
||||
def __init__(self):
|
||||
self.counter = 0
|
||||
|
||||
def inc(self):
|
||||
self.counter += 1
|
||||
|
||||
def get_counter(self):
|
||||
return self.counter
|
||||
|
||||
|
||||
counter = Counter.remote()
|
||||
|
||||
for _ in range(5):
|
||||
ray.get(counter.inc.remote())
|
||||
print(ray.get(counter.get_counter.remote()))
|
||||
|
||||
print(requests.__version__)
|
||||
+106
@@ -0,0 +1,106 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
set -ex
|
||||
|
||||
unset RAY_ADDRESS
|
||||
|
||||
if ! [ -x "$(command -v conda)" ]; then
|
||||
echo "conda doesn't exist. Please download conda for this machine"
|
||||
exit 1
|
||||
else
|
||||
echo "conda exists"
|
||||
fi
|
||||
|
||||
# This is required to use conda activate
|
||||
source "$(conda info --base)/etc/profile.d/conda.sh"
|
||||
|
||||
PYTHON_VERSION=$(python -c"from platform import python_version; print(python_version())")
|
||||
|
||||
RAY_VERSIONS=("2.0.1")
|
||||
|
||||
for RAY_VERSION in "${RAY_VERSIONS[@]}"
|
||||
do
|
||||
env_name=${JOB_COMPATIBILITY_TEST_TEMP_ENV}
|
||||
|
||||
# Check if the conda env exists
|
||||
if conda env list | grep -q "${env_name}"; then
|
||||
# Clean up if env name is already taken from previous leaking runs
|
||||
conda env remove --name="${env_name}"
|
||||
fi
|
||||
|
||||
printf "\n\n\n"
|
||||
echo "========================================================================================="
|
||||
printf "Creating new conda environment with python %s for ray %s \n" "${PYTHON_VERSION}" "${RAY_VERSION}"
|
||||
echo "========================================================================================="
|
||||
printf "\n\n\n"
|
||||
|
||||
# Include `pip` explicitly: conda-forge's `python` package stopped
|
||||
# bundling pip as a dep, and without it `conda activate` puts us in
|
||||
# an env with python but no pip, so subsequent `pip install` falls
|
||||
# back to the base miniforge env's pip. That clobbers the editable
|
||||
# ray 3.0.0.dev0 in base with ray 2.0.1, and every subsequent
|
||||
# dashboard test that imports `ray._common` fails because 2.0.1
|
||||
# predates that module.
|
||||
conda create -y -n "${env_name}" python="${PYTHON_VERSION}" pip=25.2
|
||||
conda activate "${env_name}"
|
||||
python -m pip install --upgrade pip
|
||||
|
||||
# Pin pydantic version due to: https://github.com/ray-project/ray/issues/36990.
|
||||
# ray<2.9 is only compatible with pydantic<2 and setuptools < 70.
|
||||
python -m pip install -U "pydantic<2" ray=="${RAY_VERSION}" ray[default]=="${RAY_VERSION}" setuptools==69.5.1
|
||||
|
||||
printf "\n\n\n"
|
||||
echo "========================================================="
|
||||
printf "Installed ray job server version: "
|
||||
SERVER_RAY_VERSION=$(python -c "import ray; print(ray.__version__)")
|
||||
printf "%s \n" "${SERVER_RAY_VERSION}"
|
||||
echo "========================================================="
|
||||
printf "\n\n\n"
|
||||
ray stop --force
|
||||
ray start --head
|
||||
|
||||
conda deactivate
|
||||
|
||||
CLIENT_RAY_VERSION=$(python -c "import ray; print(ray.__version__)")
|
||||
CLIENT_RAY_COMMIT=$(python -c "import ray; print(ray.__commit__)")
|
||||
printf "\n\n\n"
|
||||
echo "========================================================================================="
|
||||
printf "Using Ray %s on %s as job client \n" "${CLIENT_RAY_VERSION}" "${CLIENT_RAY_COMMIT}"
|
||||
echo "========================================================================================="
|
||||
printf "\n\n\n"
|
||||
|
||||
export RAY_ADDRESS="http://127.0.0.1:8265"
|
||||
|
||||
cleanup () {
|
||||
unset RAY_ADDRESS
|
||||
ray stop --force
|
||||
conda remove -y --name "${env_name}" --all
|
||||
}
|
||||
|
||||
JOB_ID=$(python -c "import uuid; print(uuid.uuid4().hex)")
|
||||
|
||||
# Get directory of current file. https://stackoverflow.com/questions/59895/
|
||||
DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
|
||||
|
||||
if ! ray job submit --job-id="${JOB_ID}" --working-dir="${DIR}" --runtime-env-json='{"pip": ["requests==2.26.0", "setuptools==69.5.1"]}' -- python script.py; then
|
||||
cleanup
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ! ray job status "${JOB_ID}"; then
|
||||
cleanup
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ! ray job logs "${JOB_ID}"; then
|
||||
cleanup
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ! pytest -vs "${DIR}"/../test_backwards_compatibility.py::test_error_message; then
|
||||
cleanup
|
||||
exit 1
|
||||
fi
|
||||
|
||||
cleanup
|
||||
done
|
||||
@@ -0,0 +1,30 @@
|
||||
import os
|
||||
|
||||
import ray
|
||||
from ray._raylet import GcsClient
|
||||
from ray.dashboard.modules.job.job_manager import JobManager
|
||||
|
||||
TEST_NAMESPACE = "jobs_test_namespace"
|
||||
|
||||
|
||||
def create_ray_cluster(_tracing_startup_hook=None):
|
||||
return ray.init(
|
||||
num_cpus=16,
|
||||
num_gpus=1,
|
||||
resources={"Custom": 1},
|
||||
namespace=TEST_NAMESPACE,
|
||||
log_to_driver=True,
|
||||
_tracing_startup_hook=_tracing_startup_hook,
|
||||
)
|
||||
|
||||
|
||||
def create_job_manager(ray_cluster, tmp_path):
|
||||
address_info = ray_cluster
|
||||
gcs_client = GcsClient(address=address_info["gcs_address"])
|
||||
return JobManager(gcs_client, tmp_path)
|
||||
|
||||
|
||||
def _driver_script_path(file_name: str) -> str:
|
||||
return os.path.join(
|
||||
os.path.dirname(__file__), "subprocess_driver_scripts", file_name
|
||||
)
|
||||
Binary file not shown.
+31
@@ -0,0 +1,31 @@
|
||||
"""
|
||||
A dummy ray driver script that executes in subprocess.
|
||||
Prints global worker's `load_code_from_local` property that ought to be set
|
||||
whenever `JobConfig.code_search_path` is specified
|
||||
"""
|
||||
|
||||
|
||||
def run():
|
||||
import ray
|
||||
from ray.job_config import JobConfig
|
||||
|
||||
ray.init(job_config=JobConfig(code_search_path=["/home/code/"]))
|
||||
|
||||
@ray.remote
|
||||
def foo() -> bool:
|
||||
return ray._private.worker.global_worker.load_code_from_local
|
||||
|
||||
load_code_from_local = ray.get(foo.remote())
|
||||
|
||||
statement = "propagated" if load_code_from_local else "NOT propagated"
|
||||
|
||||
# Step 1: Print the statement indicating that the code_search_path have been
|
||||
# properly respected
|
||||
print(f"Code search path is {statement}")
|
||||
# Step 2: Print the whole runtime_env to validate that it's been passed
|
||||
# appropriately from submit_job API
|
||||
print(ray.get_runtime_context().runtime_env)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run()
|
||||
+22
@@ -0,0 +1,22 @@
|
||||
import os
|
||||
|
||||
import ray
|
||||
|
||||
cuda_env = ray._private.accelerators.nvidia_gpu.NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR
|
||||
if os.environ.get("RAY_TEST_RESOURCES_SPECIFIED") == "1":
|
||||
assert cuda_env not in os.environ
|
||||
if os.environ.get("RAY_TEST_GPUS_SPECIFIED") == "1":
|
||||
assert "CUDA_VISIBLE_DEVICES" in os.environ
|
||||
else:
|
||||
assert "CUDA_VISIBLE_DEVICES" not in os.environ
|
||||
else:
|
||||
assert os.environ[cuda_env] == "1"
|
||||
|
||||
|
||||
@ray.remote
|
||||
def f():
|
||||
assert cuda_env not in os.environ
|
||||
|
||||
|
||||
# Will raise if task fails.
|
||||
ray.get(f.remote())
|
||||
@@ -0,0 +1,23 @@
|
||||
"""
|
||||
A dummy ray driver script that executes in subprocess.
|
||||
Checks that job manager's environment variable is different.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
def run():
|
||||
ray.init()
|
||||
|
||||
@ray.remote
|
||||
def foo():
|
||||
print("worker", os.nice(0))
|
||||
|
||||
ray.get(foo.remote())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("driver", os.nice(0))
|
||||
run()
|
||||
@@ -0,0 +1,13 @@
|
||||
import ray
|
||||
|
||||
ray.init()
|
||||
|
||||
|
||||
@ray.remote(num_cpus=1)
|
||||
def f():
|
||||
pass
|
||||
|
||||
|
||||
print("Hanging...")
|
||||
ray.get(f.remote())
|
||||
print("Success!")
|
||||
+58
@@ -0,0 +1,58 @@
|
||||
import argparse
|
||||
import sys
|
||||
import time
|
||||
|
||||
import ray
|
||||
|
||||
# This prefix is used to identify the output log line that contains the runtime env.
|
||||
RUNTIME_ENV_LOG_LINE_PREFIX = "ray_job_test_runtime_env_output:"
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Dashboard agent.")
|
||||
parser.add_argument(
|
||||
"--conflict",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--worker-process-setup-hook",
|
||||
type=str,
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.worker_process_setup_hook:
|
||||
ray.init(
|
||||
runtime_env={
|
||||
"worker_process_setup_hook": lambda: print(
|
||||
args.worker_process_setup_hook
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
@ray.remote
|
||||
def f():
|
||||
pass
|
||||
|
||||
ray.get(f.remote())
|
||||
time.sleep(5)
|
||||
sys.exit(0)
|
||||
|
||||
if args.conflict == "pip":
|
||||
ray.init(runtime_env={"pip": ["numpy"]})
|
||||
print(
|
||||
RUNTIME_ENV_LOG_LINE_PREFIX + ray._private.worker.global_worker.runtime_env
|
||||
)
|
||||
elif args.conflict == "env_vars":
|
||||
ray.init(runtime_env={"env_vars": {"A": "1"}})
|
||||
print(
|
||||
RUNTIME_ENV_LOG_LINE_PREFIX + ray._private.worker.global_worker.runtime_env
|
||||
)
|
||||
else:
|
||||
ray.init(
|
||||
runtime_env={
|
||||
"env_vars": {"C": "1"},
|
||||
}
|
||||
)
|
||||
print(
|
||||
RUNTIME_ENV_LOG_LINE_PREFIX + ray._private.worker.global_worker.runtime_env
|
||||
)
|
||||
@@ -0,0 +1,27 @@
|
||||
"""
|
||||
Test script that attempts to set its own runtime_env, but we should ensure
|
||||
we ended up using job submission API call's runtime_env instead of scripts
|
||||
"""
|
||||
|
||||
|
||||
def run():
|
||||
import os
|
||||
|
||||
import ray
|
||||
|
||||
ray.init(
|
||||
runtime_env={
|
||||
"env_vars": {"TEST_SUBPROCESS_JOB_CONFIG_ENV_VAR": "SHOULD_BE_OVERRIDEN"}
|
||||
},
|
||||
)
|
||||
|
||||
@ray.remote
|
||||
def foo():
|
||||
return "bar"
|
||||
|
||||
ray.get(foo.remote())
|
||||
print(os.environ.get("TEST_SUBPROCESS_JOB_CONFIG_ENV_VAR", None))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run()
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
import os
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
def run():
|
||||
ray.init()
|
||||
|
||||
@ray.remote(runtime_env={"env_vars": {"FOO": "bar"}})
|
||||
def get_task_working_dir():
|
||||
# Check behavior of working_dir: The cwd should contain the
|
||||
# current file, which is being used as a job entrypoint script.
|
||||
assert os.path.exists("per_task_runtime_env.py")
|
||||
|
||||
return ray.get_runtime_context().runtime_env.working_dir()
|
||||
|
||||
driver_working_dir = ray.get_runtime_context().runtime_env.working_dir()
|
||||
task_working_dir = ray.get(get_task_working_dir.remote())
|
||||
assert driver_working_dir == task_working_dir, (
|
||||
driver_working_dir,
|
||||
task_working_dir,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run()
|
||||
@@ -0,0 +1,22 @@
|
||||
"""
|
||||
A dummy ray driver script that executes in subprocess. Prints namespace
|
||||
from ray's runtime context for job submission API testing.
|
||||
"""
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
def run():
|
||||
ray.init()
|
||||
|
||||
@ray.remote
|
||||
def foo():
|
||||
return "bar"
|
||||
|
||||
ray.get(foo.remote())
|
||||
|
||||
print(ray.get_runtime_context().namespace)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run()
|
||||
@@ -0,0 +1,22 @@
|
||||
"""
|
||||
A dummy ray driver script that executes in subprocess. Prints runtime_env
|
||||
from ray's runtime context for job submission API testing.
|
||||
"""
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
def run():
|
||||
ray.init()
|
||||
|
||||
@ray.remote
|
||||
def foo():
|
||||
return "bar"
|
||||
|
||||
ray.get(foo.remote())
|
||||
|
||||
print(ray.get_runtime_context().runtime_env)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run()
|
||||
@@ -0,0 +1,15 @@
|
||||
"""Tests that Ray Tune works with the working_dir set in Jobs.
|
||||
|
||||
Ray Tune internally sets environment variables using runtime_env.
|
||||
If the inherited internal runtime environment overwrites the working_dir
|
||||
from jobs with an empty working_dir, this test will fail. See #25484"""
|
||||
from ray_tune_dependency import foo
|
||||
|
||||
from ray import tune
|
||||
|
||||
|
||||
def objective(*args):
|
||||
foo()
|
||||
|
||||
|
||||
tune.run(objective)
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
"""A file dependency for testing working_dir behavior with Ray Tune."""
|
||||
|
||||
|
||||
def foo():
|
||||
pass
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
def run():
|
||||
raise Exception("Script failed with exception !")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run()
|
||||
@@ -0,0 +1,124 @@
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
from ray.job_submission import JobStatus, JobSubmissionClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def conda_env(env_name):
|
||||
# Set env name for shell script
|
||||
os.environ["JOB_COMPATIBILITY_TEST_TEMP_ENV"] = env_name
|
||||
# Delete conda env if it already exists
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# Clean up created conda env upon test exit to prevent leaking
|
||||
del os.environ["JOB_COMPATIBILITY_TEST_TEMP_ENV"]
|
||||
subprocess.run(
|
||||
f"conda env remove -y --name {env_name}", shell=True, stdout=subprocess.PIPE
|
||||
)
|
||||
|
||||
|
||||
def _compatibility_script_path(file_name: str) -> str:
|
||||
return os.path.join(
|
||||
os.path.dirname(__file__), "backwards_compatibility_scripts", file_name
|
||||
)
|
||||
|
||||
|
||||
class TestBackwardsCompatibility:
|
||||
@pytest.mark.skipif(
|
||||
sys.platform == "darwin",
|
||||
reason="ray 2.0.1 runs differently on apple silicon than today's.",
|
||||
)
|
||||
def test_cli(self):
|
||||
"""
|
||||
Test that the current commit's CLI works with old server-side Ray versions.
|
||||
|
||||
1) Create a new conda environment with old ray version X installed;
|
||||
inherits same env as current conda envionment except ray version
|
||||
2) (Server) Start head node and dashboard with old ray version X
|
||||
3) (Client) Use current commit's CLI code to do sample job submission flow
|
||||
4) Deactivate the new conda environment and back to original place
|
||||
"""
|
||||
# Shell script creates and cleans up tmp conda environment regardless
|
||||
# of the outcome
|
||||
env_name = f"jobs-backwards-compatibility-{uuid.uuid4().hex}"
|
||||
with conda_env(env_name):
|
||||
shell_cmd = f"{_compatibility_script_path('test_backwards_compatibility.sh')}" # noqa: E501
|
||||
|
||||
try:
|
||||
subprocess.check_output(shell_cmd, shell=True, stderr=subprocess.STDOUT)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.error(str(e))
|
||||
logger.error(e.stdout.decode())
|
||||
raise e
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("JOB_COMPATIBILITY_TEST_TEMP_ENV") is None,
|
||||
reason="This test is only meant to be run from the "
|
||||
"test_backwards_compatibility.sh shell script.",
|
||||
)
|
||||
def test_error_message():
|
||||
"""
|
||||
Check that we get a good error message when running against an old server version.
|
||||
"""
|
||||
# Import lazily so the module still loads when the compatibility script
|
||||
# installs an older Ray that does not expose `ray._common`.
|
||||
from ray._common.test_utils import wait_for_condition
|
||||
|
||||
client = JobSubmissionClient("http://127.0.0.1:8265")
|
||||
|
||||
# Check that a basic job successfully runs.
|
||||
job_id = client.submit_job(
|
||||
entrypoint="echo 'hello world'",
|
||||
)
|
||||
wait_for_condition(lambda: client.get_job_status(job_id) == JobStatus.SUCCEEDED)
|
||||
|
||||
# `entrypoint_num_cpus`, `entrypoint_num_gpus`, `entrypoint_resources`, and
|
||||
# `entrypoint_label_selector`
|
||||
# are not supported in ray<2.2.0.
|
||||
for unsupported_submit_kwargs in [
|
||||
{"entrypoint_num_cpus": 1},
|
||||
{"entrypoint_num_gpus": 1},
|
||||
{"entrypoint_resources": {"custom": 1}},
|
||||
{"entrypoint_label_selector": {"fragile_node": "!1"}},
|
||||
]:
|
||||
with pytest.raises(
|
||||
Exception,
|
||||
match="Ray version 2.0.1 is running on the cluster. "
|
||||
"`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.",
|
||||
):
|
||||
client.submit_job(
|
||||
entrypoint="echo hello",
|
||||
**unsupported_submit_kwargs,
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
Exception,
|
||||
match="Ray version 2.0.1 is running on the cluster. "
|
||||
"`entrypoint_memory` kwarg"
|
||||
" is not supported on the Ray cluster. Please ensure the cluster is "
|
||||
"running Ray 2.8 or higher.",
|
||||
):
|
||||
client.submit_job(
|
||||
entrypoint="echo hello",
|
||||
entrypoint_memory=4,
|
||||
)
|
||||
|
||||
assert True
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,695 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shlex
|
||||
import sys
|
||||
import tempfile
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from subprocess import list2cmdline
|
||||
from typing import Optional
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from click.testing import CliRunner
|
||||
|
||||
from ray.dashboard.modules.job.cli import job_cli_group
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_sdk_client():
|
||||
class AsyncIterator:
|
||||
def __init__(self, seq):
|
||||
self._seq = seq
|
||||
self.iter = iter(self._seq)
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
try:
|
||||
return next(self.iter)
|
||||
except StopIteration:
|
||||
self.iter = iter(self._seq)
|
||||
raise StopAsyncIteration
|
||||
|
||||
if "RAY_ADDRESS" in os.environ:
|
||||
del os.environ["RAY_ADDRESS"]
|
||||
with mock.patch("ray.dashboard.modules.job.cli.JobSubmissionClient") as mock_client:
|
||||
# In python 3.6 it will fail with error
|
||||
# 'async for' requires an object with __aiter__ method, got MagicMock"
|
||||
mock_client().tail_job_logs.return_value = AsyncIterator(range(10))
|
||||
|
||||
# We need to return a string for the address and not a MagicMock
|
||||
mock_client().get_address.return_value = ""
|
||||
|
||||
yield mock_client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runtime_env_formats():
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
path = Path(tmp_dir)
|
||||
|
||||
test_env = {
|
||||
"working_dir": "s3://bogus.zip",
|
||||
"conda": "conda_env",
|
||||
"pip": ["pip-install-test"],
|
||||
"env_vars": {"hi": "hi2"},
|
||||
}
|
||||
|
||||
yaml_file = path / "env.yaml"
|
||||
with yaml_file.open(mode="w") as f:
|
||||
yaml.dump(test_env, f)
|
||||
|
||||
yield test_env, json.dumps(test_env), yaml_file
|
||||
|
||||
|
||||
@contextmanager
|
||||
def set_env_var(key: str, val: Optional[str] = None):
|
||||
old_val = os.environ.get(key, None)
|
||||
if val is not None:
|
||||
os.environ[key] = val
|
||||
elif key in os.environ:
|
||||
del os.environ[key]
|
||||
|
||||
yield
|
||||
|
||||
if key in os.environ:
|
||||
del os.environ[key]
|
||||
if old_val is not None:
|
||||
os.environ[key] = old_val
|
||||
|
||||
|
||||
def check_exit_code(result, exit_code):
|
||||
assert result.exit_code == exit_code, result.output
|
||||
|
||||
|
||||
def _expected_entrypoint(*args):
|
||||
"""Return the expected entrypoint string for the current platform.
|
||||
|
||||
On Windows, the CLI uses subprocess.list2cmdline (double quotes).
|
||||
On POSIX, it uses shlex.join (single quotes).
|
||||
"""
|
||||
if sys.platform == "win32":
|
||||
return list2cmdline(args)
|
||||
return shlex.join(args)
|
||||
|
||||
|
||||
def _job_cli_group_test_address(mock_sdk_client, cmd, *args):
|
||||
runner = CliRunner()
|
||||
|
||||
create_cluster_if_needed = True if cmd == "submit" else False
|
||||
# Test passing address via command line.
|
||||
result = runner.invoke(job_cli_group, [cmd, "--address=arg_addr", *args])
|
||||
mock_sdk_client.assert_called_with(
|
||||
"arg_addr", create_cluster_if_needed, headers=None, verify=True
|
||||
)
|
||||
with pytest.raises(AssertionError):
|
||||
mock_sdk_client.assert_called_with(
|
||||
"some_other_addr", True, headers=None, verify=True
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
# Test passing address via env var.
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(job_cli_group, [cmd, *args])
|
||||
check_exit_code(result, 0)
|
||||
# RAY_ADDRESS is read inside the SDK client.
|
||||
mock_sdk_client.assert_called_with(
|
||||
None, create_cluster_if_needed, headers=None, verify=True
|
||||
)
|
||||
# Test passing no address.
|
||||
result = runner.invoke(job_cli_group, [cmd, *args])
|
||||
check_exit_code(result, 0)
|
||||
mock_sdk_client.assert_called_with(
|
||||
None, create_cluster_if_needed, headers=None, verify=True
|
||||
)
|
||||
|
||||
|
||||
class TestList:
|
||||
def test_address(self, mock_sdk_client):
|
||||
_job_cli_group_test_address(mock_sdk_client, "list")
|
||||
|
||||
def test_list(self, mock_sdk_client):
|
||||
runner = CliRunner()
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
["list"],
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
|
||||
result = runner.invoke(job_cli_group, ["submit", "--", "echo hello"])
|
||||
check_exit_code(result, 0)
|
||||
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
["list"],
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
|
||||
|
||||
class TestSubmit:
|
||||
def test_address(self, mock_sdk_client):
|
||||
_job_cli_group_test_address(mock_sdk_client, "submit", "--", "echo", "hello")
|
||||
|
||||
def test_working_dir(self, mock_sdk_client):
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(job_cli_group, ["submit", "--", "echo hello"])
|
||||
check_exit_code(result, 0)
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env={},
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
)
|
||||
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
["submit", "--working-dir", "blah", "--", "echo hello"],
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env={"working_dir": "blah"},
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
)
|
||||
|
||||
result = runner.invoke(
|
||||
job_cli_group, ["submit", "--working-dir='.'", "--", "echo hello"]
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env={"working_dir": "'.'"},
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
)
|
||||
|
||||
def test_runtime_env(self, mock_sdk_client, runtime_env_formats):
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
env_dict, env_json, env_yaml = runtime_env_formats
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
# Test passing via file.
|
||||
result = runner.invoke(
|
||||
job_cli_group, ["submit", "--runtime-env", env_yaml, "--", "echo hello"]
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env=env_dict,
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
)
|
||||
|
||||
# Test passing via json.
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
["submit", "--runtime-env-json", env_json, "--", "echo hello"],
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env=env_dict,
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
)
|
||||
|
||||
# Test passing both throws an error.
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
[
|
||||
"submit",
|
||||
"--runtime-env",
|
||||
env_yaml,
|
||||
"--runtime-env-json",
|
||||
env_json,
|
||||
"--",
|
||||
"echo hello",
|
||||
],
|
||||
)
|
||||
check_exit_code(result, 1)
|
||||
assert "Only one of" in str(result.exception)
|
||||
|
||||
# Test overriding working_dir.
|
||||
env_dict.update(working_dir=".")
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
[
|
||||
"submit",
|
||||
"--runtime-env",
|
||||
env_yaml,
|
||||
"--working-dir",
|
||||
".",
|
||||
"--",
|
||||
"echo hello",
|
||||
],
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env=env_dict,
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
)
|
||||
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
[
|
||||
"submit",
|
||||
"--runtime-env-json",
|
||||
env_json,
|
||||
"--working-dir",
|
||||
".",
|
||||
"--",
|
||||
"echo hello",
|
||||
],
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env=env_dict,
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
)
|
||||
|
||||
def test_job_id(self, mock_sdk_client):
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(job_cli_group, ["submit", "--", "echo hello"])
|
||||
check_exit_code(result, 0)
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env={},
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
)
|
||||
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
["submit", "--submission-id=my_job_id", "--", "echo hello"],
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id="my_job_id",
|
||||
runtime_env={},
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
)
|
||||
|
||||
def test_entrypoint_num_cpus(self, mock_sdk_client):
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
["submit", "--entrypoint-num-cpus=2", "--", "echo hello"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env={},
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=2,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
)
|
||||
|
||||
def test_entrypoint_num_gpus(self, mock_sdk_client):
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
["submit", "--entrypoint-num-gpus=2", "--", "echo hello"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env={},
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=2,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
)
|
||||
|
||||
def test_entrypoint_memory(self, mock_sdk_client):
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
["submit", "--entrypoint-memory=4", "--", "echo hello"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env={},
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=4,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"resources",
|
||||
[
|
||||
("--entrypoint-num-cpus=2", {"entrypoint_num_cpus": 2}),
|
||||
("--entrypoint-num-gpus=2", {"entrypoint_num_gpus": 2}),
|
||||
(
|
||||
"""--entrypoint-resources={"Custom":3}""",
|
||||
{"entrypoint_resources": {"Custom": 3}},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_entrypoint_resources(self, mock_sdk_client, resources):
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
["submit", resources[0], "--", "echo hello"],
|
||||
)
|
||||
print(result.output)
|
||||
assert result.exit_code == 0
|
||||
expected_kwargs = {
|
||||
"entrypoint": _expected_entrypoint("echo hello"),
|
||||
"submission_id": None,
|
||||
"runtime_env": {},
|
||||
"metadata": None,
|
||||
"entrypoint_num_cpus": None,
|
||||
"entrypoint_num_gpus": None,
|
||||
"entrypoint_memory": None,
|
||||
"entrypoint_resources": None,
|
||||
"entrypoint_label_selector": None,
|
||||
}
|
||||
expected_kwargs.update(resources[1])
|
||||
mock_client_instance.submit_job.assert_called_with(**expected_kwargs)
|
||||
|
||||
def test_entrypoint_resources_invalid_json(self, mock_sdk_client):
|
||||
runner = CliRunner()
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
[
|
||||
"submit",
|
||||
"""--entrypoint-resources={"Custom":3""",
|
||||
"--",
|
||||
"echo hello world",
|
||||
],
|
||||
)
|
||||
print(result.output)
|
||||
assert result.exit_code == 1
|
||||
assert "not a valid JSON string" in result.output
|
||||
|
||||
def test_entrypoint_label_selector(self, mock_sdk_client):
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
[
|
||||
"submit",
|
||||
"""--entrypoint-label-selector={"fragile_node":"!1"}""",
|
||||
"--",
|
||||
"echo hello",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env={},
|
||||
metadata=None,
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector={"fragile_node": "!1"},
|
||||
)
|
||||
|
||||
def test_metadata(self, mock_sdk_client):
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
[
|
||||
"submit",
|
||||
"--metadata-json",
|
||||
'{"key": "value"}',
|
||||
"--",
|
||||
"echo hello",
|
||||
],
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
mock_client_instance.submit_job.assert_called_with(
|
||||
entrypoint=_expected_entrypoint("echo hello"),
|
||||
submission_id=None,
|
||||
runtime_env={},
|
||||
entrypoint_num_cpus=None,
|
||||
entrypoint_num_gpus=None,
|
||||
entrypoint_memory=None,
|
||||
entrypoint_resources=None,
|
||||
entrypoint_label_selector=None,
|
||||
metadata={"key": "value"},
|
||||
)
|
||||
|
||||
def test_metadata_invalid_json(self, mock_sdk_client):
|
||||
runner = CliRunner()
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
[
|
||||
"submit",
|
||||
"--metadata-json",
|
||||
'{"key": "value"',
|
||||
"--",
|
||||
"echo hello",
|
||||
],
|
||||
)
|
||||
print(result.output)
|
||||
check_exit_code(result, 1)
|
||||
assert "not a valid JSON string" in result.output
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cli_val, verify_param",
|
||||
[
|
||||
("True", True),
|
||||
("true", True),
|
||||
("1", True),
|
||||
("False", False),
|
||||
("false", False),
|
||||
("0", False),
|
||||
("a/rel/path", "a/rel/path"),
|
||||
("/an/abs/path", "/an/abs/path"),
|
||||
],
|
||||
)
|
||||
def test_entrypoint_verify(self, mock_sdk_client, cli_val, verify_param):
|
||||
runner = CliRunner()
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
["submit", f"--verify={cli_val}", "--", "echo hello"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
mock_sdk_client.assert_called_with(
|
||||
None, True, headers=None, verify=verify_param
|
||||
)
|
||||
|
||||
|
||||
class TestDelete:
|
||||
def test_address(self, mock_sdk_client):
|
||||
_job_cli_group_test_address(mock_sdk_client, "delete", "fake_job_id")
|
||||
|
||||
def test_delete(self, mock_sdk_client):
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(job_cli_group, ["delete", "job_id"])
|
||||
check_exit_code(result, 0)
|
||||
mock_client_instance.delete_job.assert_called_with("job_id")
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_address(self, mock_sdk_client):
|
||||
_job_cli_group_test_address(mock_sdk_client, "status", "fake_job_id")
|
||||
|
||||
def test_status(self, mock_sdk_client):
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(job_cli_group, ["status", "job_id"])
|
||||
check_exit_code(result, 0)
|
||||
mock_client_instance.get_job_info.assert_called_with("job_id")
|
||||
|
||||
|
||||
class TestEntrypointShellQuoting:
|
||||
"""Regression test for https://github.com/ray-project/ray/issues/56232.
|
||||
|
||||
`ray job submit` previously used `subprocess.list2cmdline` unconditionally
|
||||
to join entrypoint arguments. That function wraps arguments in double
|
||||
quotes, which causes POSIX shells on the server to expand $VAR references.
|
||||
|
||||
The fix uses `shlex.join` on POSIX platforms (which single-quotes
|
||||
arguments to prevent expansion) and `list2cmdline` on Windows (which
|
||||
double-quotes arguments as expected by cmd.exe).
|
||||
"""
|
||||
|
||||
def test_entrypoint_preserves_shell_variables(self, mock_sdk_client):
|
||||
"""Ensure $VAR in entrypoint is single-quoted on POSIX, not double-quoted."""
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
with mock.patch("ray.dashboard.modules.job.cli.sys") as mock_sys:
|
||||
mock_sys.platform = "linux"
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
[
|
||||
"submit",
|
||||
"--",
|
||||
"python",
|
||||
"-m",
|
||||
"launcher",
|
||||
"--config",
|
||||
"$CONFIG_PATH",
|
||||
],
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
call_kwargs = mock_client_instance.submit_job.call_args
|
||||
entrypoint = call_kwargs.kwargs["entrypoint"]
|
||||
|
||||
# shlex.join must single-quote the $VAR argument so that
|
||||
# the server-side POSIX shell does NOT expand it.
|
||||
assert "'$CONFIG_PATH'" in entrypoint, (
|
||||
f"Expected single-quoted $CONFIG_PATH in entrypoint, "
|
||||
f"got: {entrypoint!r}"
|
||||
)
|
||||
# Double quotes around $CONFIG_PATH would cause expansion.
|
||||
assert '"$CONFIG_PATH"' not in entrypoint, (
|
||||
f"Double-quoted $CONFIG_PATH would be expanded by the "
|
||||
f"server shell, got: {entrypoint!r}"
|
||||
)
|
||||
|
||||
def test_entrypoint_uses_list2cmdline_on_windows(self, mock_sdk_client):
|
||||
"""On Windows, entrypoint should use list2cmdline (double quotes)."""
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
with mock.patch("ray.dashboard.modules.job.cli.sys") as mock_sys:
|
||||
mock_sys.platform = "win32"
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
[
|
||||
"submit",
|
||||
"--",
|
||||
"echo",
|
||||
"hello world",
|
||||
],
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
call_kwargs = mock_client_instance.submit_job.call_args
|
||||
entrypoint = call_kwargs.kwargs["entrypoint"]
|
||||
|
||||
# list2cmdline wraps args with spaces in double quotes
|
||||
assert (
|
||||
entrypoint == 'echo "hello world"'
|
||||
), f"Expected list2cmdline output on Windows, got: {entrypoint!r}"
|
||||
|
||||
def test_entrypoint_simple_args_not_over_quoted(self, mock_sdk_client):
|
||||
"""Simple arguments without special chars should not be quoted."""
|
||||
runner = CliRunner()
|
||||
mock_client_instance = mock_sdk_client.return_value
|
||||
|
||||
with set_env_var("RAY_ADDRESS", "env_addr"):
|
||||
result = runner.invoke(
|
||||
job_cli_group,
|
||||
["submit", "--", "echo", "hello"],
|
||||
)
|
||||
check_exit_code(result, 0)
|
||||
call_kwargs = mock_client_instance.submit_job.call_args
|
||||
entrypoint = call_kwargs.kwargs["entrypoint"]
|
||||
assert entrypoint == "echo hello"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,356 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from contextlib import contextmanager
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def shutdown_only():
|
||||
yield None
|
||||
# The code after the yield will run as teardown code.
|
||||
ray.shutdown()
|
||||
# Delete the cluster address just in case.
|
||||
ray._common.utils.reset_ray_address()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def set_env_var(key: str, val: Optional[str] = None):
|
||||
old_val = os.environ.get(key, None)
|
||||
if val is not None:
|
||||
os.environ[key] = val
|
||||
elif key in os.environ:
|
||||
del os.environ[key]
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if key in os.environ:
|
||||
del os.environ[key]
|
||||
if old_val is not None:
|
||||
os.environ[key] = old_val
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_start_stop():
|
||||
subprocess.check_output(["ray", "start", "--head"])
|
||||
try:
|
||||
with set_env_var("RAY_ADDRESS", "http://127.0.0.1:8265"):
|
||||
yield
|
||||
finally:
|
||||
subprocess.check_output(["ray", "stop", "--force"])
|
||||
|
||||
|
||||
@contextmanager
|
||||
def ray_cluster_manager():
|
||||
"""
|
||||
Used not as fixture in case we want to set RAY_ADDRESS first.
|
||||
"""
|
||||
subprocess.check_output(["ray", "start", "--head"])
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
subprocess.check_output(["ray", "stop", "--force"])
|
||||
|
||||
|
||||
def _run_cmd(cmd: str, should_fail=False) -> Tuple[str, str]:
|
||||
"""Convenience wrapper for subprocess.run.
|
||||
|
||||
We always run with shell=True to simulate the CLI.
|
||||
|
||||
Asserts that the process succeeds/fails depending on should_fail.
|
||||
|
||||
Returns (stdout, stderr).
|
||||
"""
|
||||
print(f"Running command: '{cmd}'")
|
||||
p: subprocess.CompletedProcess = subprocess.run(
|
||||
cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
|
||||
)
|
||||
if p.returncode == 0:
|
||||
print("Command succeeded.")
|
||||
if should_fail:
|
||||
raise RuntimeError(
|
||||
f"Expected command to fail, but got exit code: {p.returncode}."
|
||||
)
|
||||
else:
|
||||
print(f"Command failed with exit code: {p.returncode}.")
|
||||
if not should_fail:
|
||||
raise RuntimeError(
|
||||
f"Expected command to succeed, but got exit code: {p.returncode}."
|
||||
)
|
||||
|
||||
return p.stdout.decode("utf-8"), p.stderr.decode("utf-8")
|
||||
|
||||
|
||||
class TestJobSubmitHook:
|
||||
"""Tests the RAY_JOB_SUBMIT_HOOK env var."""
|
||||
|
||||
def test_hook(self, ray_start_stop):
|
||||
with set_env_var("RAY_JOB_SUBMIT_HOOK", "ray._private.test_utils.job_hook"):
|
||||
stdout, _ = _run_cmd("ray job submit -- echo hello")
|
||||
assert "hook intercepted: echo hello" in stdout
|
||||
|
||||
|
||||
class TestRayJobHeaders:
|
||||
"""
|
||||
Integration version of job CLI test that ensures interaction with the
|
||||
following components are working as expected:
|
||||
1) Ray client: use of RAY_JOB_HEADERS and ray.init() in job_head.py
|
||||
2) Ray dashboard: `ray start --head`
|
||||
"""
|
||||
|
||||
def test_empty_ray_job_headers(self, ray_start_stop):
|
||||
with set_env_var("RAY_JOB_HEADERS", None):
|
||||
stdout, _ = _run_cmd("ray job submit -- echo hello")
|
||||
assert "hello" in stdout
|
||||
assert "succeeded" in stdout
|
||||
|
||||
@pytest.mark.parametrize("ray_job_headers", ['{"key": "value"}'])
|
||||
def test_ray_job_headers(self, ray_start_stop, ray_job_headers: str):
|
||||
with set_env_var("RAY_JOB_HEADERS", ray_job_headers):
|
||||
_run_cmd("ray job submit -- echo hello", should_fail=False)
|
||||
|
||||
@pytest.mark.parametrize("ray_job_headers", ["{key value}"])
|
||||
def test_ray_incorrectly_formatted_job_headers(
|
||||
self, ray_start_stop, ray_job_headers: str
|
||||
):
|
||||
with set_env_var("RAY_JOB_HEADERS", ray_job_headers):
|
||||
_run_cmd("ray job submit -- echo hello", should_fail=True)
|
||||
|
||||
|
||||
class TestRayAddress:
|
||||
"""
|
||||
Integration version of job CLI test that ensures interaction with the
|
||||
following components are working as expected:
|
||||
|
||||
1) Ray client: use of RAY_ADDRESS and ray.init() in job_head.py
|
||||
2) Ray dashboard: `ray start --head`
|
||||
"""
|
||||
|
||||
def test_empty_ray_address(self, ray_start_stop):
|
||||
with set_env_var("RAY_ADDRESS", None):
|
||||
stdout, _ = _run_cmd("ray job submit -- echo hello")
|
||||
assert "hello" in stdout
|
||||
assert "succeeded" in stdout
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_api_server_address,should_fail",
|
||||
[
|
||||
("http://127.0.0.1:8265", False), # correct API server
|
||||
("127.0.0.1:8265", True), # wrong format without http
|
||||
("http://127.0.0.1:9999", True), # wrong port
|
||||
],
|
||||
)
|
||||
def test_ray_api_server_address(
|
||||
self,
|
||||
ray_start_stop,
|
||||
ray_api_server_address: str,
|
||||
should_fail: bool,
|
||||
):
|
||||
# Set a `RAY_ADDRESS` that would not work with the `ray job submit` CLI because it uses the `ray://` prefix.
|
||||
# This verifies that the `RAY_API_SERVER_ADDRESS` env var takes precedence.
|
||||
with set_env_var("RAY_ADDRESS", "ray://127.0.0.1:8265"):
|
||||
with set_env_var("RAY_API_SERVER_ADDRESS", ray_api_server_address):
|
||||
_run_cmd("ray job submit -- echo hello", should_fail=should_fail)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_client_address,should_fail",
|
||||
[
|
||||
("127.0.0.1:8265", True),
|
||||
("ray://127.0.0.1:8265", True),
|
||||
("http://127.0.0.1:8265", False),
|
||||
],
|
||||
)
|
||||
def test_ray_client_address(
|
||||
self, ray_start_stop, ray_client_address: str, should_fail: bool
|
||||
):
|
||||
with set_env_var("RAY_ADDRESS", ray_client_address):
|
||||
_run_cmd("ray job submit -- echo hello", should_fail=should_fail)
|
||||
|
||||
def test_valid_http_ray_address(self, ray_start_stop):
|
||||
stdout, _ = _run_cmd("ray job submit -- echo hello")
|
||||
assert "hello" in stdout
|
||||
assert "succeeded" in stdout
|
||||
|
||||
|
||||
class TestJobSubmit:
|
||||
def test_basic_submit(self, ray_start_stop):
|
||||
"""Should tail logs and wait for process to exit."""
|
||||
cmd = "sleep 1 && echo hello && sleep 1 && echo hello"
|
||||
stdout, _ = _run_cmd(f"ray job submit -- bash -c '{cmd}'")
|
||||
|
||||
# 'hello' should appear four times: twice when we print the entrypoint, then
|
||||
# two more times in the logs from the `echo`.
|
||||
assert stdout.count("hello") == 4
|
||||
assert "succeeded" in stdout
|
||||
|
||||
def test_submit_no_wait(self, ray_start_stop):
|
||||
"""Should exit immediately w/o printing logs."""
|
||||
cmd = "echo hello && sleep 1000"
|
||||
stdout, _ = _run_cmd(f"ray job submit --no-wait -- bash -c '{cmd}'")
|
||||
assert "hello" not in stdout
|
||||
assert "Tailing logs until the job exits" not in stdout
|
||||
|
||||
def test_submit_with_logs_instant_job(self, ray_start_stop):
|
||||
"""Should exit immediately and print logs even if job returns instantly."""
|
||||
cmd = "echo hello"
|
||||
stdout, _ = _run_cmd(f"ray job submit -- bash -c '{cmd}'")
|
||||
|
||||
# 'hello' should appear twice: once when we print the entrypoint, then
|
||||
# again from the `echo`.
|
||||
assert stdout.count("hello") == 2
|
||||
|
||||
def test_multiple_ray_init(self, ray_start_stop):
|
||||
cmd = (
|
||||
"python -c 'import ray; ray.init(); ray.shutdown(); "
|
||||
"ray.init(); ray.shutdown();'"
|
||||
)
|
||||
stdout, _ = _run_cmd(f"ray job submit -- {cmd}")
|
||||
assert "succeeded" in stdout
|
||||
|
||||
def test_metadata(self, ray_start_stop):
|
||||
cmd = "echo hello"
|
||||
stdout, _ = _run_cmd(
|
||||
f'ray job submit --metadata-json=\'{{"key": "value"}}\' -- {cmd}'
|
||||
)
|
||||
assert "hello" in stdout
|
||||
assert "succeeded" in stdout
|
||||
|
||||
def test_job_failed(self, ray_start_stop):
|
||||
cmd = "python -c 'import ray; ray.init(); assert 1 == 2;'"
|
||||
_run_cmd(f"ray job submit -- {cmd}", should_fail=True)
|
||||
|
||||
|
||||
class TestRuntimeEnv:
|
||||
def test_bad_runtime_env(self, ray_start_stop):
|
||||
"""Should fail with helpful error if runtime env setup fails."""
|
||||
stdout, _ = _run_cmd(
|
||||
'ray job submit --runtime-env-json=\'{"pip": '
|
||||
'["does-not-exist"]}\' -- echo hi',
|
||||
should_fail=True,
|
||||
)
|
||||
assert "Tailing logs until the job exits" in stdout
|
||||
assert "runtime_env setup failed" in stdout
|
||||
assert "No matching distribution found for does-not-exist" in stdout
|
||||
|
||||
|
||||
class TestJobStop:
|
||||
def test_basic_stop(self, ray_start_stop):
|
||||
"""Should wait until the job is stopped."""
|
||||
cmd = "sleep 1000"
|
||||
job_id = "test_basic_stop"
|
||||
_run_cmd(f"ray job submit --no-wait --job-id={job_id} -- {cmd}")
|
||||
|
||||
stdout, _ = _run_cmd(f"ray job stop {job_id}")
|
||||
assert "Waiting for job" in stdout
|
||||
assert f"Job '{job_id}' was stopped" in stdout
|
||||
|
||||
def test_stop_no_wait(self, ray_start_stop):
|
||||
"""Should not wait until the job is stopped."""
|
||||
cmd = "echo hello && sleep 1000"
|
||||
job_id = "test_stop_no_wait"
|
||||
_run_cmd(f"ray job submit --no-wait --job-id={job_id} -- bash -c '{cmd}'")
|
||||
|
||||
stdout, _ = _run_cmd(f"ray job stop --no-wait {job_id}")
|
||||
assert "Waiting for job" not in stdout
|
||||
assert f"Job '{job_id}' was stopped" not in stdout
|
||||
|
||||
|
||||
class TestJobList:
|
||||
def test_empty(self, ray_start_stop):
|
||||
stdout, _ = _run_cmd("ray job list")
|
||||
assert "[]" in stdout
|
||||
|
||||
def test_list(self, ray_start_stop):
|
||||
_run_cmd("ray job submit --job-id='hello_id' -- echo hello")
|
||||
|
||||
runtime_env = {"env_vars": {"TEST": "123"}}
|
||||
_run_cmd(
|
||||
"ray job submit --job-id='hi_id' "
|
||||
f"--runtime-env-json='{json.dumps(runtime_env)}' -- echo hi"
|
||||
)
|
||||
stdout, _ = _run_cmd("ray job list")
|
||||
assert "123" in stdout
|
||||
assert "hello_id" in stdout
|
||||
assert "hi_id" in stdout
|
||||
|
||||
|
||||
class TestJobDelete:
|
||||
def test_basic_delete(self, ray_start_stop):
|
||||
cmd = "sleep 1000"
|
||||
job_id = "test_basic_delete"
|
||||
_run_cmd(f"ray job submit --no-wait --submission-id={job_id} -- {cmd}")
|
||||
|
||||
# Job shouldn't be able to be deleted because it is not in a terminal state.
|
||||
stdout, stderr = _run_cmd(f"ray job delete {job_id}", should_fail=True)
|
||||
assert "it is in a non-terminal state" in stderr
|
||||
|
||||
# Submit a job that finishes quickly.
|
||||
cmd = "echo hello"
|
||||
job_id = "test_basic_delete_quick"
|
||||
_run_cmd(f"ray job submit --submission-id={job_id} -- bash -c '{cmd}'")
|
||||
|
||||
# Job should be able to be deleted because it is finished.
|
||||
stdout, _ = _run_cmd(f"ray job delete {job_id}")
|
||||
assert f"Job '{job_id}' deleted successfully" in stdout
|
||||
|
||||
|
||||
class TestJobStatus:
|
||||
# `ray job status` should exit with 0 if the job exists and non-zero if it doesn't.
|
||||
# This is the contract between Ray and KubRay v1.3.0.
|
||||
def test_status_job_exists(self, ray_start_stop):
|
||||
cmd = "echo hello"
|
||||
job_id = "test_job_id"
|
||||
_run_cmd(
|
||||
f"ray job submit --submission-id={job_id} -- bash -c '{cmd}'",
|
||||
should_fail=False,
|
||||
)
|
||||
_run_cmd(f"ray job status {job_id}", should_fail=False)
|
||||
|
||||
def test_status_job_does_not_exist(self, ray_start_stop):
|
||||
job_id = "test_job_id"
|
||||
_run_cmd(f"ray job status {job_id}", should_fail=True)
|
||||
|
||||
|
||||
def test_quote_escaping(ray_start_stop):
|
||||
cmd = "echo \"hello 'world'\""
|
||||
job_id = "test_quote_escaping"
|
||||
stdout, _ = _run_cmd(
|
||||
f"ray job submit --job-id={job_id} -- {cmd}",
|
||||
)
|
||||
assert "hello 'world'" in stdout
|
||||
|
||||
|
||||
def test_resources(shutdown_only):
|
||||
ray.init(num_cpus=1, num_gpus=1, resources={"Custom": 1}, _memory=4)
|
||||
|
||||
# Check the case of too many resources.
|
||||
for id, arg in [
|
||||
("entrypoint_num_cpus", "--entrypoint-num-cpus=2"),
|
||||
("entrypoint_num_gpus", "--entrypoint-num-gpus=2"),
|
||||
("entrypoint_memory", "--entrypoint-memory=5"),
|
||||
("entrypoint_resources", "--entrypoint-resources='{\"Custom\": 2}'"),
|
||||
]:
|
||||
_run_cmd(f"ray job submit --submission-id={id} --no-wait {arg} -- echo hi")
|
||||
stdout, _ = _run_cmd(f"ray job status {id}")
|
||||
assert "waiting for resources" in stdout
|
||||
|
||||
# Check the case of sufficient resources.
|
||||
stdout, _ = _run_cmd(
|
||||
"ray job submit --entrypoint-num-cpus=1 "
|
||||
"--entrypoint-num-gpus=1 --entrypoint-memory=4 --entrypoint-resources='{"
|
||||
'"Custom": 1}\' -- echo hello',
|
||||
)
|
||||
assert "hello" in stdout
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,302 @@
|
||||
import asyncio
|
||||
import json
|
||||
from dataclasses import asdict
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from google.protobuf.json_format import Parse
|
||||
|
||||
from ray.core.generated.gcs_pb2 import JobsAPIInfo
|
||||
from ray.dashboard.modules.job.common import (
|
||||
JobErrorType,
|
||||
JobInfo,
|
||||
JobInfoStorageClient,
|
||||
JobStatus,
|
||||
JobSubmitRequest,
|
||||
http_uri_components_to_uri,
|
||||
uri_to_http_components,
|
||||
validate_request_type,
|
||||
)
|
||||
|
||||
|
||||
class TestJobSubmitRequestValidation:
|
||||
def test_validate_entrypoint(self):
|
||||
r = validate_request_type({"entrypoint": "abc"}, JobSubmitRequest)
|
||||
assert r.entrypoint == "abc"
|
||||
|
||||
with pytest.raises(TypeError, match="required positional argument"):
|
||||
validate_request_type({}, JobSubmitRequest)
|
||||
|
||||
with pytest.raises(TypeError, match="must be a string"):
|
||||
validate_request_type({"entrypoint": 123}, JobSubmitRequest)
|
||||
|
||||
def test_validate_submission_id(self):
|
||||
r = validate_request_type({"entrypoint": "abc"}, JobSubmitRequest)
|
||||
assert r.entrypoint == "abc"
|
||||
assert r.submission_id is None
|
||||
|
||||
r = validate_request_type(
|
||||
{"entrypoint": "abc", "submission_id": "123"}, JobSubmitRequest
|
||||
)
|
||||
assert r.entrypoint == "abc"
|
||||
assert r.submission_id == "123"
|
||||
|
||||
with pytest.raises(TypeError, match="must be a string"):
|
||||
validate_request_type(
|
||||
{"entrypoint": 123, "submission_id": 1}, JobSubmitRequest
|
||||
)
|
||||
|
||||
def test_validate_runtime_env(self):
|
||||
r = validate_request_type({"entrypoint": "abc"}, JobSubmitRequest)
|
||||
assert r.entrypoint == "abc"
|
||||
assert r.runtime_env is None
|
||||
|
||||
r = validate_request_type(
|
||||
{"entrypoint": "abc", "runtime_env": {"hi": "hi2"}}, JobSubmitRequest
|
||||
)
|
||||
assert r.entrypoint == "abc"
|
||||
assert r.runtime_env == {"hi": "hi2"}
|
||||
|
||||
with pytest.raises(TypeError, match="must be a dict"):
|
||||
validate_request_type(
|
||||
{"entrypoint": "abc", "runtime_env": 123}, JobSubmitRequest
|
||||
)
|
||||
|
||||
with pytest.raises(TypeError, match="keys must be strings"):
|
||||
validate_request_type(
|
||||
{"entrypoint": "abc", "runtime_env": {1: "hi"}}, JobSubmitRequest
|
||||
)
|
||||
|
||||
def test_validate_metadata(self):
|
||||
r = validate_request_type({"entrypoint": "abc"}, JobSubmitRequest)
|
||||
assert r.entrypoint == "abc"
|
||||
assert r.metadata is None
|
||||
|
||||
r = validate_request_type(
|
||||
{"entrypoint": "abc", "metadata": {"hi": "hi2"}}, JobSubmitRequest
|
||||
)
|
||||
assert r.entrypoint == "abc"
|
||||
assert r.metadata == {"hi": "hi2"}
|
||||
|
||||
with pytest.raises(TypeError, match="must be a dict"):
|
||||
validate_request_type(
|
||||
{"entrypoint": "abc", "metadata": 123}, JobSubmitRequest
|
||||
)
|
||||
|
||||
with pytest.raises(TypeError, match="keys must be strings"):
|
||||
validate_request_type(
|
||||
{"entrypoint": "abc", "metadata": {1: "hi"}}, JobSubmitRequest
|
||||
)
|
||||
|
||||
with pytest.raises(TypeError, match="values must be strings"):
|
||||
validate_request_type(
|
||||
{"entrypoint": "abc", "metadata": {"hi": 1}}, JobSubmitRequest
|
||||
)
|
||||
|
||||
def test_validate_entrypoint_label_selector(self):
|
||||
r = validate_request_type(
|
||||
{
|
||||
"entrypoint": "abc",
|
||||
"entrypoint_label_selector": {"fragile_node": "!1"},
|
||||
},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
assert r.entrypoint_label_selector == {"fragile_node": "!1"}
|
||||
|
||||
with pytest.raises(TypeError, match="must be a dict"):
|
||||
validate_request_type(
|
||||
{"entrypoint": "abc", "entrypoint_label_selector": "bad"},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
|
||||
with pytest.raises(TypeError, match="keys must be strings"):
|
||||
validate_request_type(
|
||||
{"entrypoint": "abc", "entrypoint_label_selector": {1: "bad"}},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
|
||||
with pytest.raises(TypeError, match="values must be strings"):
|
||||
validate_request_type(
|
||||
{"entrypoint": "abc", "entrypoint_label_selector": {"k": 1}},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
|
||||
def test_entrypoint_resources_disallow_strings(self):
|
||||
with pytest.raises(TypeError, match="values must be numbers"):
|
||||
validate_request_type(
|
||||
{"entrypoint": "abc", "entrypoint_resources": {"Custom": "1"}},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
|
||||
|
||||
def test_uri_to_http_and_back():
|
||||
assert uri_to_http_components("gcs://hello.zip") == ("gcs", "hello.zip")
|
||||
assert uri_to_http_components("gcs://hello.whl") == ("gcs", "hello.whl")
|
||||
|
||||
with pytest.raises(ValueError, match="'blah' is not a valid Protocol"):
|
||||
uri_to_http_components("blah://halb.zip")
|
||||
|
||||
with pytest.raises(ValueError, match="does not end in .zip or .whl"):
|
||||
assert uri_to_http_components("gcs://hello.not_zip")
|
||||
|
||||
with pytest.raises(ValueError, match="does not end in .zip or .whl"):
|
||||
assert uri_to_http_components("gcs://hello")
|
||||
|
||||
assert http_uri_components_to_uri("gcs", "hello.zip") == "gcs://hello.zip"
|
||||
assert http_uri_components_to_uri("blah", "halb.zip") == "blah://halb.zip"
|
||||
assert http_uri_components_to_uri("blah", "halb.whl") == "blah://halb.whl"
|
||||
|
||||
for original_uri in ["gcs://hello.zip", "gcs://fasdf.whl"]:
|
||||
new_uri = http_uri_components_to_uri(*uri_to_http_components(original_uri))
|
||||
assert new_uri == original_uri
|
||||
|
||||
|
||||
def test_dynamic_status_message():
|
||||
info = JobInfo(
|
||||
status=JobStatus.PENDING, entrypoint="echo hi", entrypoint_num_cpus=1
|
||||
)
|
||||
assert "may be waiting for resources" in info.message
|
||||
|
||||
info = JobInfo(
|
||||
status=JobStatus.PENDING, entrypoint="echo hi", entrypoint_num_gpus=1
|
||||
)
|
||||
assert "may be waiting for resources" in info.message
|
||||
|
||||
info = JobInfo(status=JobStatus.PENDING, entrypoint="echo hi", entrypoint_memory=4)
|
||||
assert "may be waiting for resources" in info.message
|
||||
|
||||
info = JobInfo(
|
||||
status=JobStatus.PENDING,
|
||||
entrypoint="echo hi",
|
||||
entrypoint_resources={"Custom": 1},
|
||||
)
|
||||
assert "may be waiting for resources" in info.message
|
||||
|
||||
info = JobInfo(
|
||||
status=JobStatus.PENDING, entrypoint="echo hi", runtime_env={"conda": "env"}
|
||||
)
|
||||
assert "may be waiting for the runtime environment" in info.message
|
||||
|
||||
|
||||
def test_job_info_to_json():
|
||||
info = JobInfo(
|
||||
status=JobStatus.PENDING,
|
||||
entrypoint="echo hi",
|
||||
entrypoint_num_cpus=1,
|
||||
entrypoint_num_gpus=1,
|
||||
entrypoint_memory=4,
|
||||
entrypoint_resources={"Custom": 1},
|
||||
runtime_env={"pip": ["pkg"]},
|
||||
)
|
||||
expected_items = {
|
||||
"status": "PENDING",
|
||||
"message": (
|
||||
"Job has not started yet. It may be waiting for resources "
|
||||
"(CPUs, GPUs, memory, custom resources) to become available. "
|
||||
"It may be waiting for the runtime environment to be set up."
|
||||
),
|
||||
"entrypoint": "echo hi",
|
||||
"entrypoint_num_cpus": 1,
|
||||
"entrypoint_num_gpus": 1,
|
||||
"entrypoint_memory": 4,
|
||||
"entrypoint_resources": {"Custom": 1},
|
||||
"runtime_env_json": '{"pip": ["pkg"]}',
|
||||
}
|
||||
|
||||
# Check that the expected items are in the JSON.
|
||||
assert expected_items.items() <= info.to_json().items()
|
||||
|
||||
new_job_info = JobInfo.from_json(info.to_json())
|
||||
assert new_job_info == info
|
||||
|
||||
# If `status` is just a string, then operations like status.is_terminal()
|
||||
# would fail, so we should make sure that it's a JobStatus.
|
||||
assert isinstance(new_job_info.status, JobStatus)
|
||||
|
||||
|
||||
def test_job_info_json_to_proto():
|
||||
"""Test that JobInfo JSON can be converted to JobsAPIInfo protobuf."""
|
||||
info = JobInfo(
|
||||
status=JobStatus.PENDING,
|
||||
entrypoint="echo hi",
|
||||
error_type=JobErrorType.JOB_SUPERVISOR_ACTOR_UNSCHEDULABLE,
|
||||
start_time=123,
|
||||
end_time=456,
|
||||
metadata={"hi": "hi2"},
|
||||
entrypoint_num_cpus=1,
|
||||
entrypoint_num_gpus=1,
|
||||
entrypoint_memory=4,
|
||||
entrypoint_resources={"Custom": 1},
|
||||
runtime_env={"pip": ["pkg"]},
|
||||
driver_agent_http_address="http://localhost:1234",
|
||||
driver_node_id="node_id",
|
||||
)
|
||||
info_json = json.dumps(info.to_json())
|
||||
info_proto = Parse(info_json, JobsAPIInfo())
|
||||
assert info_proto.status == "PENDING"
|
||||
assert info_proto.entrypoint == "echo hi"
|
||||
assert info_proto.start_time == 123
|
||||
assert info_proto.end_time == 456
|
||||
assert info_proto.metadata == {"hi": "hi2"}
|
||||
assert info_proto.entrypoint_num_cpus == 1
|
||||
assert info_proto.entrypoint_num_gpus == 1
|
||||
assert info_proto.entrypoint_memory == 4
|
||||
assert info_proto.entrypoint_resources == {"Custom": 1}
|
||||
assert info_proto.runtime_env_json == '{"pip": ["pkg"]}'
|
||||
assert info_proto.message == (
|
||||
"Job has not started yet. It may be waiting for resources "
|
||||
"(CPUs, GPUs, memory, custom resources) to become available. "
|
||||
"It may be waiting for the runtime environment to be set up."
|
||||
)
|
||||
assert info_proto.error_type == "JOB_SUPERVISOR_ACTOR_UNSCHEDULABLE"
|
||||
assert info_proto.driver_agent_http_address == "http://localhost:1234"
|
||||
assert info_proto.driver_node_id == "node_id"
|
||||
|
||||
minimal_info = JobInfo(status=JobStatus.PENDING, entrypoint="echo hi")
|
||||
minimal_info_json = json.dumps(minimal_info.to_json())
|
||||
minimal_info_proto = Parse(minimal_info_json, JobsAPIInfo())
|
||||
assert minimal_info_proto.status == "PENDING"
|
||||
assert minimal_info_proto.entrypoint == "echo hi"
|
||||
for unset_optional_field in [
|
||||
"entrypoint_num_cpus",
|
||||
"entrypoint_num_gpus",
|
||||
"entrypoint_memory",
|
||||
"runtime_env_json",
|
||||
"error_type",
|
||||
"driver_agent_http_address",
|
||||
"driver_node_id",
|
||||
]:
|
||||
assert not minimal_info_proto.HasField(unset_optional_field)
|
||||
|
||||
|
||||
def test_get_all_jobs_filters_out_none_job_info():
|
||||
prefix = JobInfoStorageClient.JOB_DATA_KEY_PREFIX
|
||||
mock_gcs = MagicMock()
|
||||
mock_gcs.async_internal_kv_keys = AsyncMock(
|
||||
return_value=[
|
||||
(prefix + "job1").encode(),
|
||||
(prefix + "job2").encode(),
|
||||
]
|
||||
)
|
||||
|
||||
storage = JobInfoStorageClient(mock_gcs)
|
||||
job_info_1 = JobInfo(status=JobStatus.RUNNING, entrypoint="echo 1")
|
||||
|
||||
async def mock_get_info(job_id, timeout=30):
|
||||
if job_id == "job1":
|
||||
return job_info_1
|
||||
return None
|
||||
|
||||
storage.get_info = mock_get_info
|
||||
|
||||
result = asyncio.run(storage.get_all_jobs())
|
||||
|
||||
assert result == {"job1": job_info_1}
|
||||
for job_id, job_info in result.items():
|
||||
asdict(job_info) # This should not raise an exception
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,177 @@
|
||||
import json
|
||||
import os
|
||||
import pprint
|
||||
import sys
|
||||
|
||||
import jsonschema
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from ray._common.test_utils import (
|
||||
run_string_as_driver,
|
||||
wait_for_condition,
|
||||
)
|
||||
from ray._private.test_utils import (
|
||||
format_web_url,
|
||||
run_string_as_driver_nonblocking,
|
||||
)
|
||||
from ray.dashboard import dashboard
|
||||
from ray.dashboard.consts import RAY_CLUSTER_ACTIVITY_HOOK
|
||||
from ray.dashboard.modules.job.job_head import RayActivityResponse
|
||||
from ray.dashboard.tests.conftest import * # noqa
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def set_ray_cluster_activity_hook(request):
|
||||
"""
|
||||
Fixture that sets RAY_CLUSTER_ACTIVITY_HOOK environment variable
|
||||
for test_e2e_component_activities_hook.
|
||||
"""
|
||||
external_hook = request.param
|
||||
assert (
|
||||
external_hook
|
||||
), "Please pass value of RAY_CLUSTER_ACTIVITY_HOOK env var to this fixture"
|
||||
old_hook = os.environ.get(RAY_CLUSTER_ACTIVITY_HOOK)
|
||||
os.environ[RAY_CLUSTER_ACTIVITY_HOOK] = external_hook
|
||||
|
||||
yield external_hook
|
||||
|
||||
if old_hook is not None:
|
||||
os.environ[RAY_CLUSTER_ACTIVITY_HOOK] = old_hook
|
||||
else:
|
||||
del os.environ[RAY_CLUSTER_ACTIVITY_HOOK]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"set_ray_cluster_activity_hook",
|
||||
[
|
||||
"ray._private.test_utils.external_ray_cluster_activity_hook1",
|
||||
"ray._private.test_utils.external_ray_cluster_activity_hook2",
|
||||
"ray._private.test_utils.external_ray_cluster_activity_hook3",
|
||||
"ray._private.test_utils.external_ray_cluster_activity_hook4",
|
||||
"ray._private.test_utils.external_ray_cluster_activity_hook5",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_component_activities_hook(set_ray_cluster_activity_hook, call_ray_start):
|
||||
"""
|
||||
Tests /api/component_activities returns correctly for various
|
||||
responses of RAY_CLUSTER_ACTIVITY_HOOK.
|
||||
|
||||
Verify no active drivers are correctly reflected in response.
|
||||
"""
|
||||
external_hook = set_ray_cluster_activity_hook
|
||||
|
||||
response = requests.get("http://127.0.0.1:8265/api/component_activities")
|
||||
response.raise_for_status()
|
||||
|
||||
# Validate schema of response
|
||||
data = response.json()
|
||||
schema_path = os.path.join(
|
||||
os.path.dirname(dashboard.__file__),
|
||||
"modules/job/component_activities_schema.json",
|
||||
)
|
||||
pprint.pprint(data)
|
||||
jsonschema.validate(instance=data, schema=json.load(open(schema_path)))
|
||||
|
||||
# Validate driver response can be cast to RayActivityResponse object
|
||||
# and that there are no active drivers.
|
||||
driver_ray_activity_response = RayActivityResponse(**data["driver"])
|
||||
assert driver_ray_activity_response.is_active == "INACTIVE"
|
||||
assert driver_ray_activity_response.reason is None
|
||||
|
||||
# Validate external component response can be cast to RayActivityResponse object
|
||||
if external_hook[-1] == "5":
|
||||
external_activity_response = RayActivityResponse(**data["test_component5"])
|
||||
assert external_activity_response.is_active == "ACTIVE"
|
||||
assert external_activity_response.reason == "Counter: 1"
|
||||
elif external_hook[-1] == "4":
|
||||
external_activity_response = RayActivityResponse(**data["external_component"])
|
||||
assert external_activity_response.is_active == "ERROR"
|
||||
assert (
|
||||
"'Error in external cluster activity hook'"
|
||||
in external_activity_response.reason
|
||||
)
|
||||
elif external_hook[-1] == "3":
|
||||
external_activity_response = RayActivityResponse(**data["external_component"])
|
||||
assert external_activity_response.is_active == "ERROR"
|
||||
elif external_hook[-1] == "2":
|
||||
external_activity_response = RayActivityResponse(**data["test_component2"])
|
||||
assert external_activity_response.is_active == "ERROR"
|
||||
elif external_hook[-1] == "1":
|
||||
external_activity_response = RayActivityResponse(**data["test_component1"])
|
||||
assert external_activity_response.is_active == "ACTIVE"
|
||||
assert external_activity_response.reason == "Counter: 1"
|
||||
|
||||
# Call endpoint again to validate different response
|
||||
response = requests.get("http://127.0.0.1:8265/api/component_activities")
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
jsonschema.validate(instance=data, schema=json.load(open(schema_path)))
|
||||
|
||||
external_activity_response = RayActivityResponse(**data["test_component1"])
|
||||
assert external_activity_response.is_active == "ACTIVE"
|
||||
assert external_activity_response.reason == "Counter: 2"
|
||||
|
||||
|
||||
def test_active_component_activities(ray_start_with_dashboard):
|
||||
# Verify drivers which don't have namespace starting with _ray_internal_
|
||||
# are considered active.
|
||||
|
||||
webui_url = ray_start_with_dashboard["webui_url"]
|
||||
webui_url = format_web_url(webui_url)
|
||||
|
||||
driver_template = """
|
||||
import ray
|
||||
|
||||
ray.init(address="auto", namespace="{namespace}")
|
||||
import time
|
||||
time.sleep({sleep_time_s})
|
||||
"""
|
||||
run_string_as_driver(
|
||||
driver_template.format(namespace="my_namespace", sleep_time_s=0)
|
||||
)
|
||||
|
||||
run_string_as_driver_nonblocking(
|
||||
driver_template.format(namespace="my_namespace", sleep_time_s=5)
|
||||
)
|
||||
run_string_as_driver_nonblocking(
|
||||
driver_template.format(namespace="_ray_internal_job_info_id1", sleep_time_s=5)
|
||||
)
|
||||
# Simulate the default driver that gets created by dashboard
|
||||
run_string_as_driver_nonblocking(
|
||||
driver_template.format(namespace="_ray_internal_dashboard", sleep_time_s=5)
|
||||
)
|
||||
|
||||
def verify_driver_response():
|
||||
# Verify drivers are considered active after running script
|
||||
response = requests.get(f"{webui_url}/api/component_activities")
|
||||
response.raise_for_status()
|
||||
|
||||
# Validate schema of response
|
||||
data = response.json()
|
||||
schema_path = os.path.join(
|
||||
os.path.dirname(dashboard.__file__),
|
||||
"modules/job/component_activities_schema.json",
|
||||
)
|
||||
|
||||
jsonschema.validate(instance=data, schema=json.load(open(schema_path)))
|
||||
|
||||
# Validate ray_activity_response field can be cast to RayActivityResponse object
|
||||
driver_ray_activity_response = RayActivityResponse(**data["driver"])
|
||||
print(driver_ray_activity_response)
|
||||
|
||||
assert driver_ray_activity_response.is_active == "ACTIVE"
|
||||
# Drivers with namespace starting with "_ray_internal" are not
|
||||
# considered active drivers. Two active drivers are the second one
|
||||
# run with namespace "my_namespace" and the one started
|
||||
# from ray_start_with_dashboard
|
||||
assert driver_ray_activity_response.reason == "Number of active drivers: 2"
|
||||
|
||||
return True
|
||||
|
||||
wait_for_condition(verify_driver_response)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,777 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
import yaml
|
||||
|
||||
import ray
|
||||
from ray._common.test_utils import wait_for_condition
|
||||
from ray._private.runtime_env.packaging import (
|
||||
create_package,
|
||||
download_and_unpack_package,
|
||||
get_uri_for_file,
|
||||
)
|
||||
from ray._private.test_utils import (
|
||||
chdir,
|
||||
format_web_url,
|
||||
wait_until_server_available,
|
||||
)
|
||||
from ray.dashboard.modules.dashboard_sdk import ClusterInfo, parse_cluster_info
|
||||
from ray.dashboard.modules.job.common import uri_to_http_components
|
||||
from ray.dashboard.modules.job.pydantic_models import JobDetails
|
||||
from ray.dashboard.modules.job.tests.test_cli_integration import set_env_var
|
||||
from ray.dashboard.modules.version import CURRENT_VERSION
|
||||
from ray.dashboard.tests.conftest import * # noqa
|
||||
from ray.job_submission import JobStatus, JobSubmissionClient
|
||||
from ray.runtime_env.runtime_env import RuntimeEnv, RuntimeEnvConfig
|
||||
from ray.tests.conftest import _ray_start
|
||||
|
||||
# This test requires you have AWS credentials set up (any AWS credentials will
|
||||
# do, this test only accesses a public bucket).
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DRIVER_SCRIPT_DIR = os.path.join(os.path.dirname(__file__), "subprocess_driver_scripts")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def headers():
|
||||
return {"Connection": "keep-alive", "Authorization": "TOK:<MY_TOKEN>"}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def ray_start_context():
|
||||
with _ray_start(include_dashboard=True, num_cpus=1) as ctx:
|
||||
yield ctx
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def job_sdk_client(headers, ray_start_context) -> JobSubmissionClient:
|
||||
address = ray_start_context.address_info["webui_url"]
|
||||
assert wait_until_server_available(address)
|
||||
yield JobSubmissionClient(format_web_url(address), headers=headers)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def shutdown_only():
|
||||
yield None
|
||||
# The code after the yield will run as teardown code.
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
def test_submit_job_with_resources(shutdown_only):
|
||||
ctx = ray.init(
|
||||
include_dashboard=True,
|
||||
num_cpus=1,
|
||||
num_gpus=1,
|
||||
resources={"Custom": 1},
|
||||
dashboard_port=8269,
|
||||
_memory=4,
|
||||
)
|
||||
address = ctx.address_info["webui_url"]
|
||||
client = JobSubmissionClient(format_web_url(address))
|
||||
# Check the case of too many resources.
|
||||
for kwargs in [
|
||||
{"entrypoint_num_cpus": 2},
|
||||
{"entrypoint_num_gpus": 2},
|
||||
{"entrypoint_memory": 4},
|
||||
{"entrypoint_resources": {"Custom": 2}},
|
||||
]:
|
||||
job_id = client.submit_job(entrypoint="echo hello", **kwargs)
|
||||
data = client.get_job_info(job_id)
|
||||
assert "waiting for resources" in data.message
|
||||
|
||||
# Check the case of sufficient resources.
|
||||
job_id = client.submit_job(
|
||||
entrypoint="echo hello",
|
||||
entrypoint_num_cpus=1,
|
||||
entrypoint_num_gpus=1,
|
||||
entrypoint_memory=4,
|
||||
entrypoint_resources={"Custom": 1},
|
||||
)
|
||||
wait_for_condition(_check_job_succeeded, client=client, job_id=job_id, timeout=10)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_sdk", [True, False])
|
||||
def test_list_jobs_empty(headers, use_sdk: bool):
|
||||
# Create a cluster using `ray start` instead of `ray.init` to avoid creating a job
|
||||
subprocess.check_output(["ray", "start", "--head"])
|
||||
address = "http://127.0.0.1:8265"
|
||||
try:
|
||||
with set_env_var("RAY_ADDRESS", address):
|
||||
client = JobSubmissionClient(format_web_url(address), headers=headers)
|
||||
|
||||
if use_sdk:
|
||||
assert client.list_jobs() == []
|
||||
else:
|
||||
r = client._do_request(
|
||||
"GET",
|
||||
"/api/jobs/",
|
||||
)
|
||||
|
||||
assert r.status_code == 200
|
||||
assert json.loads(r.text) == []
|
||||
|
||||
finally:
|
||||
subprocess.check_output(["ray", "stop", "--force"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_sdk", [True, False])
|
||||
def test_list_jobs(job_sdk_client: JobSubmissionClient, use_sdk: bool):
|
||||
client = job_sdk_client
|
||||
|
||||
runtime_env = {"env_vars": {"TEST": "123"}}
|
||||
metadata = {"foo": "bar"}
|
||||
entrypoint = "echo hello"
|
||||
submission_id = client.submit_job(
|
||||
entrypoint=entrypoint, runtime_env=runtime_env, metadata=metadata
|
||||
)
|
||||
|
||||
wait_for_condition(_check_job_succeeded, client=client, job_id=submission_id)
|
||||
if use_sdk:
|
||||
info: JobDetails = next(
|
||||
job_info
|
||||
for job_info in client.list_jobs()
|
||||
if job_info.submission_id == submission_id
|
||||
)
|
||||
else:
|
||||
r = client._do_request(
|
||||
"GET",
|
||||
"/api/jobs/",
|
||||
)
|
||||
|
||||
assert r.status_code == 200
|
||||
jobs_info_json = json.loads(r.text)
|
||||
info_json = next(
|
||||
job_info
|
||||
for job_info in jobs_info_json
|
||||
if job_info["submission_id"] == submission_id
|
||||
)
|
||||
info = JobDetails(**info_json)
|
||||
|
||||
assert info.entrypoint == entrypoint
|
||||
assert info.status == JobStatus.SUCCEEDED
|
||||
assert info.message is not None
|
||||
assert info.end_time >= info.start_time
|
||||
assert info.runtime_env == runtime_env
|
||||
assert info.metadata == metadata
|
||||
|
||||
# Test get job status by job / driver id
|
||||
status = client.get_job_status(info.submission_id)
|
||||
assert status == JobStatus.SUCCEEDED
|
||||
|
||||
|
||||
def _check_job_succeeded(client: JobSubmissionClient, job_id: str) -> bool:
|
||||
status = client.get_job_status(job_id)
|
||||
if status == JobStatus.FAILED:
|
||||
logs = client.get_job_logs(job_id)
|
||||
raise RuntimeError(
|
||||
f"Job failed\nlogs:\n{logs}, info: {client.get_job_info(job_id)}"
|
||||
)
|
||||
assert status == JobStatus.SUCCEEDED
|
||||
return True
|
||||
|
||||
|
||||
def _check_job_failed(client: JobSubmissionClient, job_id: str) -> bool:
|
||||
status = client.get_job_status(job_id)
|
||||
return status == JobStatus.FAILED
|
||||
|
||||
|
||||
def _check_job_stopped(client: JobSubmissionClient, job_id: str) -> bool:
|
||||
status = client.get_job_status(job_id)
|
||||
return status == JobStatus.STOPPED
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
scope="module",
|
||||
params=[
|
||||
"no_working_dir",
|
||||
"local_working_dir",
|
||||
"s3_working_dir",
|
||||
"local_py_modules",
|
||||
"working_dir_and_local_py_modules_whl",
|
||||
"local_working_dir_zip",
|
||||
"pip_txt",
|
||||
"conda_yaml",
|
||||
"local_py_modules",
|
||||
],
|
||||
)
|
||||
def runtime_env_option(request):
|
||||
import_in_task_script = """
|
||||
import ray
|
||||
ray.init(address="auto")
|
||||
|
||||
@ray.remote
|
||||
def f():
|
||||
import pip_install_test
|
||||
|
||||
ray.get(f.remote())
|
||||
"""
|
||||
if request.param == "no_working_dir":
|
||||
yield {
|
||||
"runtime_env": {},
|
||||
"entrypoint": "echo hello",
|
||||
"expected_logs": "hello\n",
|
||||
}
|
||||
elif request.param in {
|
||||
"local_working_dir",
|
||||
"local_working_dir_zip",
|
||||
"local_py_modules",
|
||||
"working_dir_and_local_py_modules_whl",
|
||||
}:
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
path = Path(tmp_dir)
|
||||
|
||||
hello_file = path / "test.py"
|
||||
with hello_file.open(mode="w") as f:
|
||||
f.write("from test_module import run_test\n")
|
||||
f.write("print(run_test())")
|
||||
|
||||
module_path = path / "test_module"
|
||||
module_path.mkdir(parents=True)
|
||||
|
||||
test_file = module_path / "test.py"
|
||||
with test_file.open(mode="w") as f:
|
||||
f.write("def run_test():\n")
|
||||
f.write(" return 'Hello from test_module!'\n") # noqa: Q000
|
||||
|
||||
init_file = module_path / "__init__.py"
|
||||
with init_file.open(mode="w") as f:
|
||||
f.write("from test_module.test import run_test\n")
|
||||
|
||||
if request.param == "local_working_dir":
|
||||
yield {
|
||||
"runtime_env": {"working_dir": tmp_dir},
|
||||
"entrypoint": "python test.py",
|
||||
"expected_logs": "Hello from test_module!\n",
|
||||
}
|
||||
elif request.param == "local_working_dir_zip":
|
||||
local_zipped_dir = shutil.make_archive(
|
||||
os.path.join(tmp_dir, "test"), "zip", tmp_dir
|
||||
)
|
||||
yield {
|
||||
"runtime_env": {"working_dir": local_zipped_dir},
|
||||
"entrypoint": "python test.py",
|
||||
"expected_logs": "Hello from test_module!\n",
|
||||
}
|
||||
elif request.param == "local_py_modules":
|
||||
yield {
|
||||
"runtime_env": {"py_modules": [str(Path(tmp_dir) / "test_module")]},
|
||||
"entrypoint": (
|
||||
"python -c 'import test_module;print(test_module.run_test())'"
|
||||
),
|
||||
"expected_logs": "Hello from test_module!\n",
|
||||
}
|
||||
elif request.param == "working_dir_and_local_py_modules_whl":
|
||||
yield {
|
||||
"runtime_env": {
|
||||
"working_dir": "s3://runtime-env-test/script_runtime_env.zip",
|
||||
"py_modules": [
|
||||
Path(os.path.dirname(__file__))
|
||||
/ "pip_install_test-0.5-py3-none-any.whl"
|
||||
],
|
||||
},
|
||||
"entrypoint": (
|
||||
"python script.py && python -c 'import pip_install_test'"
|
||||
),
|
||||
"expected_logs": (
|
||||
"Executing main() from script.py !!\n"
|
||||
"Good job! You installed a pip module."
|
||||
),
|
||||
}
|
||||
else:
|
||||
raise ValueError(f"Unexpected pytest fixture option {request.param}")
|
||||
elif request.param == "s3_working_dir":
|
||||
yield {
|
||||
"runtime_env": {
|
||||
"working_dir": "s3://runtime-env-test/script_runtime_env.zip",
|
||||
},
|
||||
"entrypoint": "python script.py",
|
||||
"expected_logs": "Executing main() from script.py !!\n",
|
||||
}
|
||||
elif request.param == "pip_txt":
|
||||
with tempfile.TemporaryDirectory() as tmpdir, chdir(tmpdir):
|
||||
pip_list = ["pip-install-test==0.5"]
|
||||
relative_filepath = "requirements.txt"
|
||||
pip_file = Path(relative_filepath)
|
||||
pip_file.write_text("\n".join(pip_list))
|
||||
runtime_env = {"pip": {"packages": relative_filepath, "pip_check": False}}
|
||||
yield {
|
||||
"runtime_env": runtime_env,
|
||||
"entrypoint": (
|
||||
f"python -c 'import pip_install_test' && "
|
||||
f"python -c '{import_in_task_script}'"
|
||||
),
|
||||
"expected_logs": "Good job! You installed a pip module.",
|
||||
}
|
||||
elif request.param == "conda_yaml":
|
||||
with tempfile.TemporaryDirectory() as tmpdir, chdir(tmpdir):
|
||||
conda_dict = {"dependencies": ["pip", {"pip": ["pip-install-test==0.5"]}]}
|
||||
relative_filepath = "environment.yml"
|
||||
conda_file = Path(relative_filepath)
|
||||
conda_file.write_text(yaml.dump(conda_dict))
|
||||
runtime_env = {"conda": relative_filepath}
|
||||
|
||||
yield {
|
||||
"runtime_env": runtime_env,
|
||||
"entrypoint": f"python -c '{import_in_task_script}'",
|
||||
# TODO(architkulkarni): Uncomment after #22968 is fixed.
|
||||
# "entrypoint": "python -c 'import pip_install_test'",
|
||||
"expected_logs": "Good job! You installed a pip module.",
|
||||
}
|
||||
else:
|
||||
assert False, f"Unrecognized option: {request.param}."
|
||||
|
||||
|
||||
def test_submit_job(job_sdk_client, runtime_env_option, monkeypatch):
|
||||
# This flag allows for local testing of runtime env conda functionality
|
||||
# without needing a built Ray wheel. Rather than insert the link to the
|
||||
# wheel into the conda spec, it links to the current Python site.
|
||||
monkeypatch.setenv("RAY_RUNTIME_ENV_LOCAL_DEV_MODE", "1")
|
||||
|
||||
client = job_sdk_client
|
||||
|
||||
job_id = client.submit_job(
|
||||
entrypoint=runtime_env_option["entrypoint"],
|
||||
runtime_env=runtime_env_option["runtime_env"],
|
||||
)
|
||||
|
||||
try:
|
||||
job_start_time = time.time()
|
||||
wait_for_condition(
|
||||
_check_job_succeeded, client=client, job_id=job_id, timeout=300
|
||||
)
|
||||
job_duration = time.time() - job_start_time
|
||||
print(f"The job took {job_duration}s to succeed.")
|
||||
except RuntimeError as e:
|
||||
# If the job is still pending, include job logs and info in error.
|
||||
if client.get_job_status(job_id) == JobStatus.PENDING:
|
||||
logs = client.get_job_logs(job_id)
|
||||
info = client.get_job_info(job_id)
|
||||
raise RuntimeError(
|
||||
f"Job was stuck in PENDING.\nLogs: {logs}\nInfo: {info}"
|
||||
) from e
|
||||
|
||||
logs = client.get_job_logs(job_id)
|
||||
assert runtime_env_option["expected_logs"] in logs
|
||||
|
||||
|
||||
def test_timeout(job_sdk_client):
|
||||
client = job_sdk_client
|
||||
|
||||
job_id = client.submit_job(
|
||||
entrypoint="echo hello",
|
||||
# Assume pip packages take > 1s to download, or this test will spuriously fail.
|
||||
runtime_env=RuntimeEnv(
|
||||
pip={
|
||||
"packages": ["tensorflow", "requests", "botocore", "torch"],
|
||||
"pip_check": False,
|
||||
"pip_version": "==23.3.2;python_version=='3.9.16'",
|
||||
},
|
||||
config=RuntimeEnvConfig(setup_timeout_seconds=1),
|
||||
),
|
||||
)
|
||||
|
||||
wait_for_condition(_check_job_failed, client=client, job_id=job_id, timeout=10)
|
||||
data = client.get_job_info(job_id)
|
||||
assert "Failed to set up runtime environment" in data.message
|
||||
assert "Timeout" in data.message
|
||||
assert "setup_timeout_seconds" in data.message
|
||||
|
||||
|
||||
def test_per_task_runtime_env(job_sdk_client: JobSubmissionClient):
|
||||
run_cmd = "python per_task_runtime_env.py"
|
||||
job_id = job_sdk_client.submit_job(
|
||||
entrypoint=run_cmd,
|
||||
runtime_env={"working_dir": DRIVER_SCRIPT_DIR},
|
||||
)
|
||||
|
||||
wait_for_condition(_check_job_succeeded, client=job_sdk_client, job_id=job_id)
|
||||
|
||||
|
||||
def test_ray_tune_basic(job_sdk_client: JobSubmissionClient):
|
||||
run_cmd = "python ray_tune_basic.py"
|
||||
job_id = job_sdk_client.submit_job(
|
||||
entrypoint=run_cmd,
|
||||
runtime_env={"working_dir": DRIVER_SCRIPT_DIR},
|
||||
)
|
||||
wait_for_condition(
|
||||
_check_job_succeeded, timeout=30, client=job_sdk_client, job_id=job_id
|
||||
)
|
||||
|
||||
|
||||
def test_http_bad_request(job_sdk_client):
|
||||
"""
|
||||
Send bad requests to job http server and ensure right return code and
|
||||
error message is returned via http.
|
||||
"""
|
||||
client = job_sdk_client
|
||||
|
||||
# 400 - HTTPBadRequest
|
||||
r = client._do_request(
|
||||
"POST",
|
||||
"/api/jobs/",
|
||||
json_data={"key": "baaaad request"},
|
||||
)
|
||||
|
||||
assert r.status_code == 400
|
||||
assert "__init__() got an unexpected keyword argument" in r.text
|
||||
|
||||
|
||||
def test_invalid_runtime_env(job_sdk_client):
|
||||
client = job_sdk_client
|
||||
with pytest.raises(ValueError, match="Only .zip, .tar.gz, and .tgz files"):
|
||||
client.submit_job(
|
||||
entrypoint="echo hello", runtime_env={"working_dir": "s3://not_a_zip"}
|
||||
)
|
||||
|
||||
|
||||
def test_runtime_env_setup_failure(job_sdk_client):
|
||||
client = job_sdk_client
|
||||
job_id = client.submit_job(
|
||||
entrypoint="echo hello", runtime_env={"working_dir": "s3://does_not_exist.zip"}
|
||||
)
|
||||
|
||||
wait_for_condition(_check_job_failed, client=client, job_id=job_id)
|
||||
data = client.get_job_info(job_id)
|
||||
assert "Failed to set up runtime environment" in data.message
|
||||
|
||||
|
||||
def test_submit_job_with_exception_in_driver(job_sdk_client):
|
||||
"""
|
||||
Submit a job that's expected to throw exception while executing.
|
||||
"""
|
||||
client = job_sdk_client
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
path = Path(tmp_dir)
|
||||
driver_script = """
|
||||
print('Hello !')
|
||||
raise RuntimeError('Intentionally failed.')
|
||||
"""
|
||||
test_script_file = path / "test_script.py"
|
||||
with open(test_script_file, "w+") as file:
|
||||
file.write(driver_script)
|
||||
|
||||
job_id = client.submit_job(
|
||||
entrypoint="python test_script.py", runtime_env={"working_dir": tmp_dir}
|
||||
)
|
||||
|
||||
wait_for_condition(_check_job_failed, client=client, job_id=job_id)
|
||||
logs = client.get_job_logs(job_id)
|
||||
assert "Hello !" in logs
|
||||
assert "RuntimeError: Intentionally failed." in logs
|
||||
|
||||
|
||||
def test_stop_long_running_job(job_sdk_client):
|
||||
"""
|
||||
Submit a job that runs for a while and stop it in the middle.
|
||||
"""
|
||||
client = job_sdk_client
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
path = Path(tmp_dir)
|
||||
driver_script = """
|
||||
print('Hello !')
|
||||
import time
|
||||
time.sleep(300) # This should never finish
|
||||
raise RuntimeError('Intentionally failed.')
|
||||
"""
|
||||
test_script_file = path / "test_script.py"
|
||||
with open(test_script_file, "w+") as file:
|
||||
file.write(driver_script)
|
||||
|
||||
job_id = client.submit_job(
|
||||
entrypoint="python test_script.py", runtime_env={"working_dir": tmp_dir}
|
||||
)
|
||||
assert client.stop_job(job_id) is True
|
||||
wait_for_condition(_check_job_stopped, client=client, job_id=job_id)
|
||||
|
||||
|
||||
def test_delete_job(job_sdk_client, capsys):
|
||||
"""
|
||||
Submit a job and delete it.
|
||||
"""
|
||||
client: JobSubmissionClient = job_sdk_client
|
||||
|
||||
job_id = client.submit_job(entrypoint="sleep 300 && echo hello")
|
||||
with pytest.raises(Exception, match="but it is in a non-terminal state"):
|
||||
# This should fail because the job is not in a terminal state.
|
||||
client.delete_job(job_id)
|
||||
|
||||
# Check that the job appears in list_jobs
|
||||
jobs = client.list_jobs()
|
||||
assert job_id in [job.submission_id for job in jobs]
|
||||
|
||||
finished_job_id = client.submit_job(entrypoint="echo hello")
|
||||
wait_for_condition(_check_job_succeeded, client=client, job_id=finished_job_id)
|
||||
deleted = client.delete_job(finished_job_id)
|
||||
assert deleted is True
|
||||
|
||||
# Check that the job no longer appears in list_jobs
|
||||
jobs = client.list_jobs()
|
||||
assert finished_job_id not in [job.submission_id for job in jobs]
|
||||
|
||||
|
||||
def test_job_metadata(job_sdk_client):
|
||||
client = job_sdk_client
|
||||
|
||||
print_metadata_cmd = (
|
||||
'python -c"'
|
||||
"import ray;"
|
||||
"ray.init();"
|
||||
"job_config=ray._private.worker.global_worker.core_worker.get_job_config();"
|
||||
"print(dict(sorted(job_config.metadata.items())))"
|
||||
'"'
|
||||
)
|
||||
|
||||
job_id = client.submit_job(
|
||||
entrypoint=print_metadata_cmd, metadata={"key1": "val1", "key2": "val2"}
|
||||
)
|
||||
|
||||
wait_for_condition(_check_job_succeeded, client=client, job_id=job_id)
|
||||
|
||||
assert str(
|
||||
{
|
||||
"job_name": job_id,
|
||||
"job_submission_id": job_id,
|
||||
"key1": "val1",
|
||||
"key2": "val2",
|
||||
}
|
||||
) in client.get_job_logs(job_id)
|
||||
|
||||
|
||||
def test_pass_job_id(job_sdk_client):
|
||||
client = job_sdk_client
|
||||
|
||||
job_id = "my_custom_id"
|
||||
returned_id = client.submit_job(entrypoint="echo hello", job_id=job_id)
|
||||
|
||||
assert returned_id == job_id
|
||||
wait_for_condition(_check_job_succeeded, client=client, job_id=returned_id)
|
||||
|
||||
# Test that a duplicate job_id is rejected.
|
||||
with pytest.raises(Exception, match=f"{job_id} already exists"):
|
||||
returned_id = client.submit_job(entrypoint="echo hello", job_id=job_id)
|
||||
|
||||
|
||||
def test_nonexistent_job(job_sdk_client):
|
||||
client = job_sdk_client
|
||||
|
||||
with pytest.raises(RuntimeError, match="nonexistent_job does not exist"):
|
||||
client.get_job_status("nonexistent_job")
|
||||
|
||||
|
||||
def test_submit_optional_args(job_sdk_client):
|
||||
"""Check that job_id, runtime_env, and metadata are optional."""
|
||||
client = job_sdk_client
|
||||
|
||||
r = client._do_request(
|
||||
"POST",
|
||||
"/api/jobs/",
|
||||
json_data={"entrypoint": "ls"},
|
||||
)
|
||||
|
||||
wait_for_condition(
|
||||
_check_job_succeeded, client=client, job_id=r.json()["submission_id"]
|
||||
)
|
||||
|
||||
|
||||
def test_submit_still_accepts_job_id_or_submission_id(job_sdk_client):
|
||||
"""Check that job_id, runtime_env, and metadata are optional."""
|
||||
client = job_sdk_client
|
||||
|
||||
client._do_request(
|
||||
"POST",
|
||||
"/api/jobs/",
|
||||
json_data={"entrypoint": "ls", "job_id": "raysubmit_12345"},
|
||||
)
|
||||
|
||||
wait_for_condition(_check_job_succeeded, client=client, job_id="raysubmit_12345")
|
||||
|
||||
client._do_request(
|
||||
"POST",
|
||||
"/api/jobs/",
|
||||
json_data={"entrypoint": "ls", "submission_id": "raysubmit_23456"},
|
||||
)
|
||||
|
||||
wait_for_condition(_check_job_succeeded, client=client, job_id="raysubmit_23456")
|
||||
|
||||
|
||||
def test_missing_resources(job_sdk_client):
|
||||
"""Check that 404s are raised for resources that don't exist."""
|
||||
client = job_sdk_client
|
||||
|
||||
conditions = [
|
||||
("GET", "/api/jobs/fake_job_id"),
|
||||
("GET", "/api/jobs/fake_job_id/logs"),
|
||||
("POST", "/api/jobs/fake_job_id/stop"),
|
||||
("GET", "/api/packages/fake_package_uri"),
|
||||
]
|
||||
|
||||
for method, route in conditions:
|
||||
assert client._do_request(method, route).status_code == 404
|
||||
|
||||
|
||||
def test_version_endpoint(job_sdk_client):
|
||||
client = job_sdk_client
|
||||
|
||||
r = client._do_request("GET", "/api/version")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body == {
|
||||
"version": CURRENT_VERSION,
|
||||
"ray_version": ray.__version__,
|
||||
"ray_commit": ray.__commit__,
|
||||
"session_name": body["session_name"],
|
||||
}
|
||||
|
||||
|
||||
def test_request_headers(job_sdk_client):
|
||||
client = job_sdk_client
|
||||
with patch("requests.request") as mock_request:
|
||||
_ = client._do_request(
|
||||
"POST",
|
||||
"/api/jobs/",
|
||||
json_data={"entrypoint": "ls"},
|
||||
)
|
||||
mock_request.assert_called_with(
|
||||
"POST",
|
||||
"http://127.0.0.1:8265/api/jobs/",
|
||||
cookies=None,
|
||||
data=None,
|
||||
json={"entrypoint": "ls"},
|
||||
headers={"Connection": "keep-alive", "Authorization": "TOK:<MY_TOKEN>"},
|
||||
verify=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheme", ["http", "https", "fake_module"])
|
||||
@pytest.mark.parametrize("host", ["127.0.0.1", "localhost", "fake.dns.name"])
|
||||
@pytest.mark.parametrize("port", [None, 8265, 10000])
|
||||
def test_parse_cluster_info(scheme: str, host: str, port: Optional[int]):
|
||||
address = f"{scheme}://{host}"
|
||||
if port is not None:
|
||||
address += f":{port}"
|
||||
|
||||
if scheme in {"http", "https"}:
|
||||
assert parse_cluster_info(address, False) == ClusterInfo(
|
||||
address=address,
|
||||
cookies=None,
|
||||
metadata=None,
|
||||
headers=None,
|
||||
)
|
||||
else:
|
||||
with pytest.raises(RuntimeError):
|
||||
parse_cluster_info(address, False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tail_job_logs(job_sdk_client):
|
||||
client = job_sdk_client
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
path = Path(tmp_dir)
|
||||
driver_script = """
|
||||
import time
|
||||
for i in range(100):
|
||||
print("Hello", i)
|
||||
time.sleep(0.1)
|
||||
"""
|
||||
test_script_file = path / "test_script.py"
|
||||
with open(test_script_file, "w+") as f:
|
||||
f.write(driver_script)
|
||||
|
||||
job_id = client.submit_job(
|
||||
entrypoint="python test_script.py", runtime_env={"working_dir": tmp_dir}
|
||||
)
|
||||
|
||||
st = time.time()
|
||||
while time.time() - st <= 10:
|
||||
try:
|
||||
i = 0
|
||||
async for lines in client.tail_job_logs(job_id):
|
||||
print(lines, end="")
|
||||
for line in lines.strip().split("\n"):
|
||||
assert line.split(" ") == ["Hello", str(i)]
|
||||
i += 1
|
||||
except Exception as ex:
|
||||
print("Exception:", ex)
|
||||
|
||||
wait_for_condition(_check_job_succeeded, client=client, job_id=job_id)
|
||||
|
||||
|
||||
def _hook(env):
|
||||
with open(env["env_vars"]["TEMPPATH"], "w+") as f:
|
||||
f.write(env["env_vars"]["TOKEN"])
|
||||
return env
|
||||
|
||||
|
||||
def test_jobs_env_hook(job_sdk_client: JobSubmissionClient):
|
||||
client = job_sdk_client
|
||||
|
||||
_, path = tempfile.mkstemp()
|
||||
runtime_env = {"env_vars": {"TEMPPATH": path, "TOKEN": "Ray rocks!"}}
|
||||
run_job_script = """
|
||||
import os
|
||||
import ray
|
||||
os.environ["RAY_RUNTIME_ENV_HOOK"] =\
|
||||
"ray.dashboard.modules.job.tests.test_http_job_server._hook"
|
||||
ray.init(address="auto")
|
||||
"""
|
||||
entrypoint = f"python -c '{run_job_script}'"
|
||||
job_id = client.submit_job(entrypoint=entrypoint, runtime_env=runtime_env)
|
||||
|
||||
wait_for_condition(_check_job_succeeded, client=client, job_id=job_id)
|
||||
|
||||
with open(path) as f:
|
||||
assert f.read().strip() == "Ray rocks!"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_upload_package(ray_start_context, tmp_path):
|
||||
assert wait_until_server_available(ray_start_context["webui_url"])
|
||||
webui_url = format_web_url(ray_start_context["webui_url"])
|
||||
gcs_client = ray._private.worker.global_worker.gcs_client
|
||||
url = webui_url + "/api/packages/{protocol}/{package_name}"
|
||||
|
||||
pkg_dir = tmp_path / "pkg"
|
||||
pkg_dir.mkdir()
|
||||
filename = "task.py"
|
||||
|
||||
file_content = b"Hello world"
|
||||
with (pkg_dir / filename).open("wb") as f:
|
||||
f.write(file_content)
|
||||
|
||||
package_uri = get_uri_for_file(str(pkg_dir / filename))
|
||||
protocol, package_name = uri_to_http_components(package_uri)
|
||||
package_file = tmp_path / package_name
|
||||
create_package(str(pkg_dir), package_file, include_gitignore=True)
|
||||
|
||||
resp = requests.get(url.format(protocol=protocol, package_name=package_name))
|
||||
assert resp.status_code == 404
|
||||
|
||||
resp = requests.put(
|
||||
url.format(protocol=protocol, package_name=package_name),
|
||||
data=package_file.read_bytes(),
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
resp = requests.get(url.format(protocol=protocol, package_name=package_name))
|
||||
assert resp.status_code == 200
|
||||
|
||||
await download_and_unpack_package(package_uri, str(tmp_path), gcs_client)
|
||||
assert (package_file.with_suffix("") / filename).read_bytes() == file_content
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,47 @@
|
||||
import ssl
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import trustme
|
||||
|
||||
import ray
|
||||
from ray.job_submission import JobSubmissionClient
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def ca():
|
||||
return trustme.CA()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def httpserver_ssl_context(ca):
|
||||
context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
||||
localhost_cert = ca.issue_cert("localhost")
|
||||
localhost_cert.configure_cert(context)
|
||||
return context
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def httpclient_ssl_context(ca):
|
||||
with ca.cert_pem.tempfile() as ca_temp_path:
|
||||
return ssl.create_default_context(cafile=ca_temp_path)
|
||||
|
||||
|
||||
def test_mock_https_connection(httpserver, ca):
|
||||
"""Test connections to a mock HTTPS job submission server."""
|
||||
httpserver.expect_request("/api/version").respond_with_json(
|
||||
{"ray_version": ray.__version__}
|
||||
)
|
||||
mock_url = httpserver.url_for("/")
|
||||
# Connection without SSL certificate should fail
|
||||
with pytest.raises(ConnectionError):
|
||||
JobSubmissionClient(mock_url)
|
||||
# Connecton with SSL verification skipped should succeed
|
||||
JobSubmissionClient(mock_url, verify=False)
|
||||
# Connection with SSL verification should succeed
|
||||
with ca.cert_pem.tempfile() as ca_temp_path:
|
||||
JobSubmissionClient(mock_url, verify=ca_temp_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,692 @@
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import requests
|
||||
import yaml
|
||||
|
||||
import ray
|
||||
from ray._common.network_utils import build_address
|
||||
from ray._common.test_utils import async_wait_for_condition, wait_for_condition
|
||||
from ray._common.utils import get_or_create_event_loop
|
||||
from ray._private.ray_constants import DEFAULT_DASHBOARD_AGENT_LISTEN_PORT
|
||||
from ray._private.runtime_env.py_modules import upload_py_modules_if_needed
|
||||
from ray._private.runtime_env.working_dir import upload_working_dir_if_needed
|
||||
from ray._private.test_utils import (
|
||||
chdir,
|
||||
format_web_url,
|
||||
get_current_unused_port,
|
||||
run_string_as_driver_nonblocking,
|
||||
wait_until_server_available,
|
||||
)
|
||||
from ray.dashboard.modules.job.common import (
|
||||
JOB_ACTOR_NAME_TEMPLATE,
|
||||
SUPERVISOR_ACTOR_RAY_NAMESPACE,
|
||||
JobSubmitRequest,
|
||||
validate_request_type,
|
||||
)
|
||||
from ray.dashboard.modules.job.job_head import JobAgentSubmissionClient
|
||||
from ray.dashboard.tests.conftest import * # noqa
|
||||
from ray.job_submission import JobStatus, JobSubmissionClient
|
||||
from ray.runtime_env.runtime_env import RuntimeEnv, RuntimeEnvConfig
|
||||
from ray.tests.conftest import _ray_start
|
||||
from ray.util.state import get_node, list_actors, list_nodes
|
||||
|
||||
# This test requires you have AWS credentials set up (any AWS credentials will
|
||||
# do, this test only accesses a public bucket).
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DRIVER_SCRIPT_DIR = os.path.join(os.path.dirname(__file__), "subprocess_driver_scripts")
|
||||
EVENT_LOOP = get_or_create_event_loop()
|
||||
|
||||
|
||||
def get_node_id_for_supervisor_actor_for_job(
|
||||
address: str, job_submission_id: str
|
||||
) -> str:
|
||||
actors = list_actors(
|
||||
address=address,
|
||||
filters=[("ray_namespace", "=", SUPERVISOR_ACTOR_RAY_NAMESPACE)],
|
||||
)
|
||||
for actor in actors:
|
||||
if actor.name == JOB_ACTOR_NAME_TEMPLATE.format(job_id=job_submission_id):
|
||||
return actor.node_id
|
||||
raise ValueError(f"actor not found for job_submission_id {job_submission_id}")
|
||||
|
||||
|
||||
def get_node_ip_by_id(node_id: str) -> str:
|
||||
node = get_node(id=node_id)
|
||||
return node.node_ip
|
||||
|
||||
|
||||
class JobAgentSubmissionBrowserClient(JobAgentSubmissionClient):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._session.headers[
|
||||
"User-Agent"
|
||||
] = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36" # noqa: E501
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def job_sdk_client(make_sure_dashboard_http_port_unused):
|
||||
with _ray_start(include_dashboard=True, num_cpus=1) as ctx:
|
||||
node_ip = ctx.address_info["node_ip_address"]
|
||||
agent_address = build_address(node_ip, DEFAULT_DASHBOARD_AGENT_LISTEN_PORT)
|
||||
assert wait_until_server_available(agent_address)
|
||||
head_address = ctx.address_info["webui_url"]
|
||||
assert wait_until_server_available(head_address)
|
||||
yield (
|
||||
JobAgentSubmissionClient(format_web_url(agent_address)),
|
||||
JobSubmissionClient(format_web_url(head_address)),
|
||||
)
|
||||
|
||||
|
||||
def _check_job(
|
||||
client: JobSubmissionClient, job_id: str, status: JobStatus, timeout: int = 10
|
||||
) -> bool:
|
||||
res_status = client.get_job_status(job_id)
|
||||
assert res_status == status
|
||||
return True
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
scope="module",
|
||||
params=[
|
||||
"no_working_dir",
|
||||
"local_working_dir",
|
||||
"s3_working_dir",
|
||||
"local_py_modules",
|
||||
"working_dir_and_local_py_modules_whl",
|
||||
"local_working_dir_zip",
|
||||
"pip_txt",
|
||||
"conda_yaml",
|
||||
"local_py_modules",
|
||||
],
|
||||
)
|
||||
def runtime_env_option(request):
|
||||
import_in_task_script = """
|
||||
import ray
|
||||
ray.init(address="auto")
|
||||
|
||||
@ray.remote
|
||||
def f():
|
||||
import pip_install_test
|
||||
|
||||
ray.get(f.remote())
|
||||
"""
|
||||
if request.param == "no_working_dir":
|
||||
yield {
|
||||
"runtime_env": {},
|
||||
"entrypoint": "echo hello",
|
||||
"expected_logs": "hello\n",
|
||||
}
|
||||
elif request.param in {
|
||||
"local_working_dir",
|
||||
"local_working_dir_zip",
|
||||
"local_py_modules",
|
||||
"working_dir_and_local_py_modules_whl",
|
||||
}:
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
path = Path(tmp_dir)
|
||||
|
||||
hello_file = path / "test.py"
|
||||
with hello_file.open(mode="w") as f:
|
||||
f.write("from test_module import run_test\n")
|
||||
f.write("print(run_test())")
|
||||
|
||||
module_path = path / "test_module"
|
||||
module_path.mkdir(parents=True)
|
||||
|
||||
test_file = module_path / "test.py"
|
||||
with test_file.open(mode="w") as f:
|
||||
f.write("def run_test():\n")
|
||||
f.write(" return 'Hello from test_module!'\n") # noqa: Q000
|
||||
|
||||
init_file = module_path / "__init__.py"
|
||||
with init_file.open(mode="w") as f:
|
||||
f.write("from test_module.test import run_test\n")
|
||||
|
||||
if request.param == "local_working_dir":
|
||||
yield {
|
||||
"runtime_env": {"working_dir": tmp_dir},
|
||||
"entrypoint": "python test.py",
|
||||
"expected_logs": "Hello from test_module!\n",
|
||||
}
|
||||
elif request.param == "local_working_dir_zip":
|
||||
local_zipped_dir = shutil.make_archive(
|
||||
os.path.join(tmp_dir, "test"), "zip", tmp_dir
|
||||
)
|
||||
yield {
|
||||
"runtime_env": {"working_dir": local_zipped_dir},
|
||||
"entrypoint": "python test.py",
|
||||
"expected_logs": "Hello from test_module!\n",
|
||||
}
|
||||
elif request.param == "local_py_modules":
|
||||
yield {
|
||||
"runtime_env": {"py_modules": [str(Path(tmp_dir) / "test_module")]},
|
||||
"entrypoint": (
|
||||
"python -c 'import test_module;print(test_module.run_test())'"
|
||||
),
|
||||
"expected_logs": "Hello from test_module!\n",
|
||||
}
|
||||
elif request.param == "working_dir_and_local_py_modules_whl":
|
||||
yield {
|
||||
"runtime_env": {
|
||||
"working_dir": "s3://runtime-env-test/script_runtime_env.zip",
|
||||
"py_modules": [
|
||||
Path(os.path.dirname(__file__))
|
||||
/ "pip_install_test-0.5-py3-none-any.whl"
|
||||
],
|
||||
},
|
||||
"entrypoint": (
|
||||
"python script.py && python -c 'import pip_install_test'"
|
||||
),
|
||||
"expected_logs": (
|
||||
"Executing main() from script.py !!\n"
|
||||
"Good job! You installed a pip module."
|
||||
),
|
||||
}
|
||||
else:
|
||||
raise ValueError(f"Unexpected pytest fixture option {request.param}")
|
||||
elif request.param == "s3_working_dir":
|
||||
yield {
|
||||
"runtime_env": {
|
||||
"working_dir": "s3://runtime-env-test/script_runtime_env.zip",
|
||||
},
|
||||
"entrypoint": "python script.py",
|
||||
"expected_logs": "Executing main() from script.py !!\n",
|
||||
}
|
||||
elif request.param == "pip_txt":
|
||||
with tempfile.TemporaryDirectory() as tmpdir, chdir(tmpdir):
|
||||
pip_list = ["pip-install-test==0.5"]
|
||||
relative_filepath = "requirements.txt"
|
||||
pip_file = Path(relative_filepath)
|
||||
pip_file.write_text("\n".join(pip_list))
|
||||
runtime_env = {"pip": {"packages": relative_filepath, "pip_check": False}}
|
||||
yield {
|
||||
"runtime_env": runtime_env,
|
||||
"entrypoint": (
|
||||
f"python -c 'import pip_install_test' && "
|
||||
f"python -c '{import_in_task_script}'"
|
||||
),
|
||||
"expected_logs": "Good job! You installed a pip module.",
|
||||
}
|
||||
elif request.param == "conda_yaml":
|
||||
with tempfile.TemporaryDirectory() as tmpdir, chdir(tmpdir):
|
||||
conda_dict = {"dependencies": ["pip", {"pip": ["pip-install-test==0.5"]}]}
|
||||
relative_filepath = "environment.yml"
|
||||
conda_file = Path(relative_filepath)
|
||||
conda_file.write_text(yaml.dump(conda_dict))
|
||||
runtime_env = {"conda": relative_filepath}
|
||||
|
||||
yield {
|
||||
"runtime_env": runtime_env,
|
||||
"entrypoint": f"python -c '{import_in_task_script}'",
|
||||
# TODO(architkulkarni): Uncomment after #22968 is fixed.
|
||||
# "entrypoint": "python -c 'import pip_install_test'",
|
||||
"expected_logs": "Good job! You installed a pip module.",
|
||||
}
|
||||
else:
|
||||
assert False, f"Unrecognized option: {request.param}."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_job(job_sdk_client, runtime_env_option, monkeypatch):
|
||||
# This flag allows for local testing of runtime env conda functionality
|
||||
# without needing a built Ray wheel. Rather than insert the link to the
|
||||
# wheel into the conda spec, it links to the current Python site.
|
||||
monkeypatch.setenv("RAY_RUNTIME_ENV_LOCAL_DEV_MODE", "1")
|
||||
|
||||
agent_client, head_client = job_sdk_client
|
||||
|
||||
runtime_env = runtime_env_option["runtime_env"]
|
||||
runtime_env = upload_working_dir_if_needed(
|
||||
runtime_env, include_gitignore=True, logger=logger
|
||||
)
|
||||
runtime_env = upload_py_modules_if_needed(
|
||||
runtime_env, include_gitignore=True, logger=logger
|
||||
)
|
||||
runtime_env = RuntimeEnv(**runtime_env_option["runtime_env"]).to_dict()
|
||||
request = validate_request_type(
|
||||
{"runtime_env": runtime_env, "entrypoint": runtime_env_option["entrypoint"]},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
submit_result = await agent_client.submit_job_internal(request)
|
||||
job_id = submit_result.submission_id
|
||||
|
||||
try:
|
||||
job_start_time = time.time()
|
||||
wait_for_condition(
|
||||
partial(
|
||||
_check_job,
|
||||
client=head_client,
|
||||
job_id=job_id,
|
||||
status=JobStatus.SUCCEEDED,
|
||||
),
|
||||
timeout=300,
|
||||
)
|
||||
job_duration = time.time() - job_start_time
|
||||
print(f"The job took {job_duration}s to succeed.")
|
||||
except RuntimeError as e:
|
||||
# If the job is still pending, include job logs and info in error.
|
||||
if head_client.get_job_status(job_id) == JobStatus.PENDING:
|
||||
logs = head_client.get_job_logs(job_id)
|
||||
info = head_client.get_job_info(job_id)
|
||||
raise RuntimeError(
|
||||
f"Job was stuck in PENDING.\nLogs: {logs}\nInfo: {info}"
|
||||
) from e
|
||||
|
||||
# There is only one node, so there is no need to replace the client of the JobAgent
|
||||
resp = await agent_client.get_job_logs_internal(job_id)
|
||||
assert runtime_env_option["expected_logs"] in resp.logs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_job_rejects_browsers(
|
||||
job_sdk_client, runtime_env_option, monkeypatch
|
||||
):
|
||||
# This flag allows for local testing of runtime env conda functionality
|
||||
# without needing a built Ray wheel. Rather than insert the link to the
|
||||
# wheel into the conda spec, it links to the current Python site.
|
||||
monkeypatch.setenv("RAY_RUNTIME_ENV_LOCAL_DEV_MODE", "1")
|
||||
|
||||
agent_client, head_client = job_sdk_client
|
||||
agent_address = agent_client._agent_address
|
||||
agent_client = JobAgentSubmissionBrowserClient(agent_address)
|
||||
|
||||
runtime_env = runtime_env_option["runtime_env"]
|
||||
runtime_env = upload_working_dir_if_needed(
|
||||
runtime_env, include_gitignore=True, logger=logger
|
||||
)
|
||||
runtime_env = upload_py_modules_if_needed(
|
||||
runtime_env, include_gitignore=True, logger=logger
|
||||
)
|
||||
runtime_env = RuntimeEnv(**runtime_env_option["runtime_env"]).to_dict()
|
||||
request = validate_request_type(
|
||||
{"runtime_env": runtime_env, "entrypoint": runtime_env_option["entrypoint"]},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError) as exc:
|
||||
_ = await agent_client.submit_job_internal(request)
|
||||
|
||||
assert "status code 403" in str(exc.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_job_rejects_browsers(job_sdk_client, monkeypatch):
|
||||
"""Test that DELETE job requests from browsers are rejected."""
|
||||
monkeypatch.setenv("RAY_RUNTIME_ENV_LOCAL_DEV_MODE", "1")
|
||||
|
||||
agent_client, head_client = job_sdk_client
|
||||
|
||||
# First, submit a job using the normal client
|
||||
runtime_env = RuntimeEnv().to_dict()
|
||||
request = validate_request_type(
|
||||
{"runtime_env": runtime_env, "entrypoint": "echo hello"},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
submit_result = await agent_client.submit_job_internal(request)
|
||||
job_id = submit_result.submission_id
|
||||
|
||||
# Now try to delete the job using browser-like headers
|
||||
agent_address = agent_client._agent_address
|
||||
browser_client = JobAgentSubmissionBrowserClient(agent_address)
|
||||
|
||||
with pytest.raises(RuntimeError) as exc:
|
||||
_ = await browser_client.delete_job_internal(job_id)
|
||||
|
||||
assert "status code 403" in str(exc.value)
|
||||
|
||||
await browser_client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout(job_sdk_client):
|
||||
agent_client, head_client = job_sdk_client
|
||||
|
||||
runtime_env = RuntimeEnv(
|
||||
pip={
|
||||
"packages": ["tensorflow", "requests", "botocore", "torch"],
|
||||
"pip_check": False,
|
||||
"pip_version": "==23.3.2;python_version=='3.9.16'",
|
||||
},
|
||||
config=RuntimeEnvConfig(setup_timeout_seconds=1),
|
||||
).to_dict()
|
||||
request = validate_request_type(
|
||||
{"runtime_env": runtime_env, "entrypoint": "echo hello"},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
|
||||
submit_result = await agent_client.submit_job_internal(request)
|
||||
job_id = submit_result.submission_id
|
||||
|
||||
wait_for_condition(
|
||||
partial(_check_job, client=head_client, job_id=job_id, status=JobStatus.FAILED),
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
data = head_client.get_job_info(job_id)
|
||||
assert "Failed to set up runtime environment" in data.message
|
||||
assert "Timeout" in data.message
|
||||
assert "setup_timeout_seconds" in data.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_env_setup_failure(job_sdk_client):
|
||||
agent_client, head_client = job_sdk_client
|
||||
|
||||
runtime_env = RuntimeEnv(working_dir="s3://does_not_exist.zip").to_dict()
|
||||
request = validate_request_type(
|
||||
{"runtime_env": runtime_env, "entrypoint": "echo hello"},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
|
||||
submit_result = await agent_client.submit_job_internal(request)
|
||||
job_id = submit_result.submission_id
|
||||
|
||||
wait_for_condition(
|
||||
partial(_check_job, client=head_client, job_id=job_id, status=JobStatus.FAILED),
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
data = head_client.get_job_info(job_id)
|
||||
assert "Failed to set up runtime environment" in data.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_long_running_job(job_sdk_client):
|
||||
"""
|
||||
Submit a job that runs for a while and stop it in the middle.
|
||||
"""
|
||||
agent_client, head_client = job_sdk_client
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
path = Path(tmp_dir)
|
||||
driver_script = """
|
||||
print('Hello !')
|
||||
import time
|
||||
time.sleep(300) # This should never finish
|
||||
raise RuntimeError('Intentionally failed.')
|
||||
"""
|
||||
test_script_file = path / "test_script.py"
|
||||
with open(test_script_file, "w+") as file:
|
||||
file.write(driver_script)
|
||||
|
||||
runtime_env = {"working_dir": tmp_dir}
|
||||
runtime_env = upload_working_dir_if_needed(
|
||||
runtime_env, include_gitignore=True, scratch_dir=tmp_dir, logger=logger
|
||||
)
|
||||
runtime_env = RuntimeEnv(**runtime_env).to_dict()
|
||||
|
||||
request = validate_request_type(
|
||||
{"runtime_env": runtime_env, "entrypoint": "python test_script.py"},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
submit_result = await agent_client.submit_job_internal(request)
|
||||
job_id = submit_result.submission_id
|
||||
|
||||
resp = await agent_client.stop_job_internal(job_id)
|
||||
assert resp.stopped is True
|
||||
|
||||
wait_for_condition(
|
||||
partial(
|
||||
_check_job, client=head_client, job_id=job_id, status=JobStatus.STOPPED
|
||||
),
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tail_job_logs_with_echo(job_sdk_client):
|
||||
agent_client, head_client = job_sdk_client
|
||||
|
||||
runtime_env = RuntimeEnv().to_dict()
|
||||
entrypoint = "python -c \"import time; [(print('Hello', i), time.sleep(0.1)) for i in range(100)]\"" # noqa: E501
|
||||
request = validate_request_type(
|
||||
{
|
||||
"runtime_env": runtime_env,
|
||||
"entrypoint": entrypoint,
|
||||
},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
|
||||
submit_result = await agent_client.submit_job_internal(request)
|
||||
job_id = submit_result.submission_id
|
||||
|
||||
i = 0
|
||||
async for lines in agent_client.tail_job_logs(job_id):
|
||||
print(lines, end="")
|
||||
for line in lines.strip().split("\n"):
|
||||
if (
|
||||
"Runtime env is setting up." in line
|
||||
or "Running entrypoint for job" in line
|
||||
):
|
||||
continue
|
||||
assert line.split(" ") == ["Hello", str(i)]
|
||||
i += 1
|
||||
|
||||
wait_for_condition(
|
||||
partial(
|
||||
_check_job, client=head_client, job_id=job_id, status=JobStatus.SUCCEEDED
|
||||
),
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster_head",
|
||||
[
|
||||
{
|
||||
"include_dashboard": True,
|
||||
"dashboard_agent_listen_port": DEFAULT_DASHBOARD_AGENT_LISTEN_PORT,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_job_log_in_multiple_node(
|
||||
make_sure_dashboard_http_port_unused,
|
||||
enable_test_module,
|
||||
disable_aiohttp_cache,
|
||||
ray_start_cluster_head,
|
||||
):
|
||||
cluster = ray_start_cluster_head
|
||||
assert wait_until_server_available(cluster.webui_url) is True
|
||||
webui_url = cluster.webui_url
|
||||
webui_url = format_web_url(webui_url)
|
||||
cluster.add_node(
|
||||
dashboard_agent_listen_port=DEFAULT_DASHBOARD_AGENT_LISTEN_PORT + 1
|
||||
)
|
||||
cluster.add_node(
|
||||
dashboard_agent_listen_port=DEFAULT_DASHBOARD_AGENT_LISTEN_PORT + 2
|
||||
)
|
||||
|
||||
node_ip = cluster.head_node.node_ip_address
|
||||
agent_address = build_address(node_ip, DEFAULT_DASHBOARD_AGENT_LISTEN_PORT)
|
||||
assert wait_until_server_available(agent_address)
|
||||
client = JobAgentSubmissionClient(format_web_url(agent_address))
|
||||
|
||||
def _check_nodes():
|
||||
try:
|
||||
assert len(list_nodes()) == 3
|
||||
return True
|
||||
except Exception as ex:
|
||||
logger.info(ex)
|
||||
return False
|
||||
|
||||
wait_for_condition(_check_nodes, timeout=15)
|
||||
|
||||
job_ids = []
|
||||
job_check_status = []
|
||||
JOB_NUM = 10
|
||||
job_agent_ports = [
|
||||
DEFAULT_DASHBOARD_AGENT_LISTEN_PORT,
|
||||
DEFAULT_DASHBOARD_AGENT_LISTEN_PORT + 1,
|
||||
DEFAULT_DASHBOARD_AGENT_LISTEN_PORT + 2,
|
||||
]
|
||||
for index in range(JOB_NUM):
|
||||
runtime_env = RuntimeEnv().to_dict()
|
||||
request = validate_request_type(
|
||||
{
|
||||
"runtime_env": runtime_env,
|
||||
"entrypoint": f"while true; do echo hello index-{index}"
|
||||
" && sleep 3600; done",
|
||||
},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
|
||||
submit_result = await client.submit_job_internal(request)
|
||||
job_ids.append(submit_result.submission_id)
|
||||
job_check_status.append(False)
|
||||
|
||||
async def _check_all_jobs_log():
|
||||
response = requests.get(webui_url + "/nodes?view=summary")
|
||||
response.raise_for_status()
|
||||
summary = response.json()
|
||||
assert summary["result"] is True, summary["msg"]
|
||||
summary = summary["data"]["summary"]
|
||||
|
||||
for index, job_id in enumerate(job_ids):
|
||||
if job_check_status[index]:
|
||||
continue
|
||||
result_log = f"hello index-{index}"
|
||||
# Try to get the node id which supervisor actor running in.
|
||||
node_id = get_node_id_for_supervisor_actor_for_job(cluster.address, job_id)
|
||||
for node_info in summary:
|
||||
if node_info["raylet"]["nodeId"] == node_id:
|
||||
break
|
||||
assert node_info["raylet"]["nodeId"] == node_id, f"node id: {node_id}"
|
||||
|
||||
# Try to get the agent HTTP port by node id.
|
||||
for agent_port in job_agent_ports:
|
||||
if f"--listen-port={agent_port}" in " ".join(node_info["cmdline"]):
|
||||
break
|
||||
assert f"--listen-port={agent_port}" in " ".join(
|
||||
node_info["cmdline"]
|
||||
), f"port: {agent_port}"
|
||||
|
||||
# Finally, we got the whole agent address, and try to get the job log.
|
||||
ip = get_node_ip_by_id(node_id)
|
||||
agent_address = f"{ip}:{agent_port}"
|
||||
assert wait_until_server_available(agent_address)
|
||||
client = JobAgentSubmissionClient(format_web_url(agent_address))
|
||||
resp = await client.get_job_logs_internal(job_id)
|
||||
assert result_log in resp.logs, f"logs: {resp.logs}"
|
||||
|
||||
job_check_status[index] = True
|
||||
return True
|
||||
|
||||
st = time.time()
|
||||
while time.time() - st <= 30:
|
||||
try:
|
||||
await _check_all_jobs_log()
|
||||
break
|
||||
except Exception as ex:
|
||||
print("error:", ex)
|
||||
time.sleep(1)
|
||||
assert all(job_check_status), job_check_status
|
||||
|
||||
|
||||
def test_agent_logs_not_streamed_to_drivers():
|
||||
"""Ensure when the job submission is used,
|
||||
(ray.init is called from an agent), the agent logs are
|
||||
not streamed to drivers.
|
||||
|
||||
Related: https://github.com/ray-project/ray/issues/29944
|
||||
"""
|
||||
script = """
|
||||
import ray
|
||||
from ray.job_submission import JobSubmissionClient, JobStatus
|
||||
from ray._private.test_utils import format_web_url
|
||||
from ray._common.test_utils import wait_for_condition
|
||||
|
||||
ray.init()
|
||||
address = ray._private.worker._global_node.webui_url
|
||||
address = format_web_url(address)
|
||||
client = JobSubmissionClient(address)
|
||||
submission_id = client.submit_job(entrypoint="ls")
|
||||
wait_for_condition(
|
||||
lambda: client.get_job_status(submission_id) == JobStatus.SUCCEEDED
|
||||
)
|
||||
"""
|
||||
|
||||
proc = run_string_as_driver_nonblocking(script)
|
||||
out_str = proc.stdout.read().decode("ascii")
|
||||
err_str = proc.stderr.read().decode("ascii")
|
||||
|
||||
print(out_str, err_str)
|
||||
|
||||
assert "(raylet)" not in out_str
|
||||
assert "(raylet)" not in err_str
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_default_dashboard_agent_http_port(tmp_path):
|
||||
"""Test that we can connect to the dashboard agent with a non-default
|
||||
http port.
|
||||
"""
|
||||
import subprocess
|
||||
|
||||
dashboard_agent_port = get_current_unused_port()
|
||||
cmd = f"ray start --head --dashboard-agent-listen-port {dashboard_agent_port}"
|
||||
subprocess.check_output(cmd, shell=True)
|
||||
|
||||
try:
|
||||
# We will need to wait for the ray to be started in the subprocess.
|
||||
address_info = ray.init("auto", ignore_reinit_error=True).address_info
|
||||
|
||||
node_ip = address_info["node_ip_address"]
|
||||
|
||||
dashboard_agent_listen_port = address_info["dashboard_agent_listen_port"]
|
||||
agent_address = build_address(node_ip, dashboard_agent_listen_port)
|
||||
print("agent address = ", agent_address)
|
||||
|
||||
agent_client = JobAgentSubmissionClient(format_web_url(agent_address))
|
||||
head_client = JobSubmissionClient(format_web_url(address_info["webui_url"]))
|
||||
|
||||
assert wait_until_server_available(agent_address)
|
||||
|
||||
# Submit a job through the agent.
|
||||
runtime_env = RuntimeEnv().to_dict()
|
||||
request = validate_request_type(
|
||||
{
|
||||
"runtime_env": runtime_env,
|
||||
"entrypoint": "echo hello",
|
||||
},
|
||||
JobSubmitRequest,
|
||||
)
|
||||
submit_result = await agent_client.submit_job_internal(request)
|
||||
job_id = submit_result.submission_id
|
||||
|
||||
async def verify():
|
||||
# Wait for job finished.
|
||||
wait_for_condition(
|
||||
partial(
|
||||
_check_job,
|
||||
client=head_client,
|
||||
job_id=job_id,
|
||||
status=JobStatus.SUCCEEDED,
|
||||
),
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
resp = await agent_client.get_job_logs_internal(job_id)
|
||||
assert "hello" in resp.logs, resp.logs
|
||||
|
||||
return True
|
||||
|
||||
await async_wait_for_condition(verify, retry_interval_ms=2000)
|
||||
finally:
|
||||
subprocess.check_output("ray stop --force", shell=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,182 @@
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from ray.cluster_utils import Cluster
|
||||
from ray.dashboard.consts import RAY_STREAM_RUNTIME_ENV_LOG_TO_JOB_DRIVER_LOG_ENV_VAR
|
||||
from ray.dashboard.modules.job.tests.conftest import _driver_script_path
|
||||
from ray.dashboard.modules.job.tests.subprocess_driver_scripts.driver_runtime_env_inheritance import ( # noqa: E501
|
||||
RUNTIME_ENV_LOG_LINE_PREFIX,
|
||||
)
|
||||
from ray.job_submission import JobStatus, JobSubmissionClient
|
||||
|
||||
|
||||
def wait_until_status(client, job_id, status_to_wait_for, timeout_seconds=20):
|
||||
start = time.time()
|
||||
while time.time() - start <= timeout_seconds:
|
||||
status = client.get_job_status(job_id)
|
||||
print(f"status: {status}")
|
||||
if status in status_to_wait_for:
|
||||
return
|
||||
time.sleep(1)
|
||||
raise Exception
|
||||
|
||||
|
||||
def wait(client, job_id):
|
||||
wait_until_status(
|
||||
client,
|
||||
job_id,
|
||||
{JobStatus.SUCCEEDED, JobStatus.STOPPED, JobStatus.FAILED},
|
||||
timeout_seconds=60,
|
||||
)
|
||||
|
||||
|
||||
def get_runtime_env_from_logs(client, job_id):
|
||||
wait(client, job_id)
|
||||
logs = client.get_job_logs(job_id)
|
||||
print(logs)
|
||||
assert client.get_job_status(job_id) == JobStatus.SUCCEEDED
|
||||
# Split logs by line, find the unique line that starts with
|
||||
# RUNTIME_ENV_LOG_LINE_PREFIX, strip it and parse it as JSON.
|
||||
lines = logs.strip().split("\n")
|
||||
assert len(lines) > 0
|
||||
for line in lines:
|
||||
if line.startswith(RUNTIME_ENV_LOG_LINE_PREFIX):
|
||||
return json.loads(line[len(RUNTIME_ENV_LOG_LINE_PREFIX) :])
|
||||
|
||||
|
||||
def test_job_driver_inheritance():
|
||||
try:
|
||||
c = Cluster()
|
||||
c.add_node(num_cpus=1)
|
||||
# If using a remote cluster, replace 127.0.0.1 with the head node's IP address.
|
||||
client = JobSubmissionClient("http://127.0.0.1:8265")
|
||||
driver_script_path = _driver_script_path("driver_runtime_env_inheritance.py")
|
||||
job_id = client.submit_job(
|
||||
entrypoint=f"python {driver_script_path}",
|
||||
runtime_env={
|
||||
"env_vars": {"A": "1", "B": "2"},
|
||||
"pip": ["requests"],
|
||||
},
|
||||
)
|
||||
|
||||
# Test key is merged
|
||||
print("Test key merged")
|
||||
runtime_env = get_runtime_env_from_logs(client, job_id)
|
||||
assert runtime_env["env_vars"] == {"A": "1", "B": "2", "C": "1"}
|
||||
assert runtime_env["pip"] == {"packages": ["requests"], "pip_check": False}
|
||||
|
||||
# Test worker process setuphook works.
|
||||
print("Test key setup hook")
|
||||
expected_str = "HELLOWORLD"
|
||||
job_id = client.submit_job(
|
||||
entrypoint=(
|
||||
f"python {driver_script_path} "
|
||||
f"--worker-process-setup-hook {expected_str}"
|
||||
),
|
||||
runtime_env={
|
||||
"env_vars": {"A": "1", "B": "2"},
|
||||
},
|
||||
)
|
||||
wait(client, job_id)
|
||||
logs = client.get_job_logs(job_id)
|
||||
assert expected_str in logs
|
||||
|
||||
# Test raise an exception upon key conflict
|
||||
print("Test conflicting pip")
|
||||
job_id = client.submit_job(
|
||||
entrypoint=f"python {driver_script_path} --conflict=pip",
|
||||
runtime_env={"pip": ["numpy"]},
|
||||
)
|
||||
wait(client, job_id)
|
||||
status = client.get_job_status(job_id)
|
||||
logs = client.get_job_logs(job_id)
|
||||
assert status == JobStatus.FAILED
|
||||
assert "Failed to merge the Job's runtime env" in logs
|
||||
|
||||
# Test raise an exception upon env var conflict
|
||||
print("Test conflicting env vars")
|
||||
job_id = client.submit_job(
|
||||
entrypoint=f"python {driver_script_path} --conflict=env_vars",
|
||||
runtime_env={
|
||||
"env_vars": {"A": "1"},
|
||||
},
|
||||
)
|
||||
wait(client, job_id)
|
||||
status = client.get_job_status(job_id)
|
||||
logs = client.get_job_logs(job_id)
|
||||
assert status == JobStatus.FAILED
|
||||
assert "Failed to merge the Job's runtime env" in logs
|
||||
finally:
|
||||
c.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream_runtime_env_log", ["1", "0"])
|
||||
def test_runtime_env_logs_streamed_to_job_driver_log(
|
||||
monkeypatch, stream_runtime_env_log
|
||||
):
|
||||
monkeypatch.setenv(
|
||||
RAY_STREAM_RUNTIME_ENV_LOG_TO_JOB_DRIVER_LOG_ENV_VAR, stream_runtime_env_log
|
||||
)
|
||||
try:
|
||||
c = Cluster()
|
||||
c.add_node(num_cpus=1)
|
||||
client = JobSubmissionClient("http://127.0.0.1:8265")
|
||||
job_id = client.submit_job(
|
||||
entrypoint="echo hello world",
|
||||
runtime_env={"pip": ["requests==2.25.1"]},
|
||||
)
|
||||
wait(client, job_id)
|
||||
logs = client.get_job_logs(job_id)
|
||||
if stream_runtime_env_log == "0":
|
||||
assert "Creating virtualenv at" not in logs
|
||||
else:
|
||||
assert "Creating virtualenv at" in logs
|
||||
finally:
|
||||
c.shutdown()
|
||||
|
||||
|
||||
def test_job_driver_inheritance_override(monkeypatch):
|
||||
monkeypatch.setenv("RAY_OVERRIDE_JOB_RUNTIME_ENV", "1")
|
||||
|
||||
try:
|
||||
c = Cluster()
|
||||
c.add_node(num_cpus=1)
|
||||
# If using a remote cluster, replace 127.0.0.1 with the head node's IP address.
|
||||
client = JobSubmissionClient("http://127.0.0.1:8265")
|
||||
driver_script_path = _driver_script_path("driver_runtime_env_inheritance.py")
|
||||
job_id = client.submit_job(
|
||||
entrypoint=f"python {driver_script_path}",
|
||||
runtime_env={
|
||||
"env_vars": {"A": "1", "B": "2"},
|
||||
"pip": ["requests"],
|
||||
},
|
||||
)
|
||||
|
||||
# Test conflict resolution regular field
|
||||
job_id = client.submit_job(
|
||||
entrypoint=f"python {driver_script_path} --conflict=pip",
|
||||
runtime_env={"pip": ["pip-install-test==0.5"]},
|
||||
)
|
||||
runtime_env = get_runtime_env_from_logs(client, job_id)
|
||||
print(runtime_env)
|
||||
assert runtime_env["pip"] == {"packages": ["numpy"], "pip_check": False}
|
||||
|
||||
# Test raise an exception upon env var conflict
|
||||
job_id = client.submit_job(
|
||||
entrypoint=f"python {driver_script_path} --conflict=env_vars",
|
||||
runtime_env={
|
||||
"env_vars": {"A": "2"},
|
||||
},
|
||||
)
|
||||
runtime_env = get_runtime_env_from_logs(client, job_id)
|
||||
print(runtime_env)
|
||||
assert runtime_env["env_vars"]["A"] == "1"
|
||||
finally:
|
||||
c.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,73 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from ray._common.test_utils import async_wait_for_condition
|
||||
from ray.dashboard.modules.job.tests.conftest import (
|
||||
_driver_script_path,
|
||||
create_job_manager,
|
||||
create_ray_cluster,
|
||||
)
|
||||
from ray.dashboard.modules.job.tests.test_job_manager import check_job_succeeded
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestRuntimeEnvStandalone:
|
||||
"""NOTE: PLEASE READ CAREFULLY BEFORE MODIFYING
|
||||
This test is extracted into a standalone module such that it can bootstrap its own
|
||||
(standalone) Ray cluster while avoiding affecting the shared one used by other
|
||||
JobManager tests
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tracing_enabled",
|
||||
[
|
||||
False,
|
||||
# TODO(issues/38633): local code loading is broken when tracing is enabled
|
||||
# True,
|
||||
],
|
||||
)
|
||||
async def test_user_provided_job_config_honored_by_worker(
|
||||
self, tracing_enabled, tmp_path
|
||||
):
|
||||
"""Ensures that the JobConfig instance injected into ray.init in the driver
|
||||
script is honored even in case when job is submitted via JobManager.submit_job
|
||||
API (involving RAY_JOB_CONFIG_JSON_ENV_VAR being set in child process env)
|
||||
|
||||
"""
|
||||
|
||||
if tracing_enabled:
|
||||
tracing_startup_hook = (
|
||||
"ray.util.tracing.setup_local_tmp_tracing:setup_tracing"
|
||||
)
|
||||
else:
|
||||
tracing_startup_hook = None
|
||||
|
||||
with create_ray_cluster(_tracing_startup_hook=tracing_startup_hook) as cluster:
|
||||
job_manager = create_job_manager(cluster, tmp_path)
|
||||
|
||||
driver_script_path = _driver_script_path(
|
||||
"check_code_search_path_is_propagated.py"
|
||||
)
|
||||
|
||||
job_id = await job_manager.submit_job(
|
||||
entrypoint=f"python {driver_script_path}",
|
||||
# NOTE: We inject runtime_env in here, but also specify the JobConfig in
|
||||
# the driver script: settings to JobConfig (other than the
|
||||
# runtime_env) passed in via ray.init(...) have to be respected
|
||||
# along with the runtime_env passed from submit_job API
|
||||
runtime_env={"env_vars": {"TEST_SUBPROCESS_RANDOM_VAR": "0xDEEDDEED"}},
|
||||
)
|
||||
|
||||
await async_wait_for_condition(
|
||||
check_job_succeeded, job_manager=job_manager, job_id=job_id
|
||||
)
|
||||
|
||||
logs = job_manager.get_job_logs(job_id)
|
||||
|
||||
assert "Code search path is propagated" in logs, logs
|
||||
assert "0xDEEDDEED" in logs, logs
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,435 @@
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import aiohttp
|
||||
import pytest
|
||||
|
||||
import ray.experimental.internal_kv as kv
|
||||
from ray._common.test_utils import wait_for_condition
|
||||
from ray._private import worker
|
||||
from ray._private.ray_constants import (
|
||||
KV_NAMESPACE_DASHBOARD,
|
||||
PROCESS_TYPE_DASHBOARD,
|
||||
)
|
||||
from ray._private.test_utils import (
|
||||
format_web_url,
|
||||
wait_until_server_available,
|
||||
)
|
||||
from ray._raylet import GcsClient
|
||||
from ray.dashboard.consts import (
|
||||
DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX,
|
||||
GCS_RPC_TIMEOUT_SECONDS,
|
||||
RAY_JOB_ALLOW_DRIVER_ON_WORKER_NODES_ENV_VAR,
|
||||
)
|
||||
from ray.dashboard.modules.dashboard_sdk import (
|
||||
DEFAULT_DASHBOARD_ADDRESS,
|
||||
ClusterInfo,
|
||||
parse_cluster_info,
|
||||
)
|
||||
from ray.dashboard.modules.job.pydantic_models import JobType
|
||||
from ray.dashboard.modules.job.sdk import JobStatus, JobSubmissionClient
|
||||
from ray.dashboard.tests.conftest import * # noqa
|
||||
from ray.runtime_env.runtime_env import RuntimeEnv
|
||||
from ray.tests.conftest import _ray_start
|
||||
from ray.util.state import list_nodes
|
||||
|
||||
import psutil
|
||||
|
||||
|
||||
def _check_job_succeeded(client: JobSubmissionClient, job_id: str) -> bool:
|
||||
status = client.get_job_status(job_id)
|
||||
if status == JobStatus.FAILED:
|
||||
logs = client.get_job_logs(job_id)
|
||||
raise RuntimeError(f"Job failed\nlogs:\n{logs}")
|
||||
return status == JobStatus.SUCCEEDED
|
||||
|
||||
|
||||
def check_internal_kv_gced():
|
||||
return len(kv._internal_kv_list("gcs://")) == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"address_param",
|
||||
[
|
||||
("ray://1.2.3.4:10001", "ray", "1.2.3.4:10001"),
|
||||
("other_module://", "other_module", ""),
|
||||
("other_module://address", "other_module", "address"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("create_cluster_if_needed", [True, False])
|
||||
@pytest.mark.parametrize("cookies", [None, {"test_cookie_key": "test_cookie_val"}])
|
||||
@pytest.mark.parametrize("metadata", [None, {"test_metadata_key": "test_metadata_val"}])
|
||||
@pytest.mark.parametrize("headers", [None, {"test_headers_key": "test_headers_val"}])
|
||||
@pytest.mark.parametrize("extra_kwargs", [{}, {"cloud": "my-cloud"}])
|
||||
def test_parse_cluster_info(
|
||||
address_param: Tuple[str, str, str],
|
||||
create_cluster_if_needed: bool,
|
||||
cookies: Optional[Dict[str, str]],
|
||||
metadata: Optional[Dict[str, str]],
|
||||
headers: Optional[Dict[str, str]],
|
||||
extra_kwargs: Dict[str, str],
|
||||
):
|
||||
"""
|
||||
Test ray.dashboard.modules.dashboard_sdk.parse_cluster_info for different
|
||||
format of addresses.
|
||||
"""
|
||||
mock_get_job_submission_client_cluster = Mock(return_value="Ray ClusterInfo")
|
||||
mock_module = Mock()
|
||||
mock_module.get_job_submission_client_cluster_info = Mock(
|
||||
return_value="Other module ClusterInfo"
|
||||
)
|
||||
mock_import_module = Mock(return_value=mock_module)
|
||||
|
||||
address, module_string, inner_address = address_param
|
||||
|
||||
with (
|
||||
patch.multiple(
|
||||
"ray.dashboard.modules.dashboard_sdk",
|
||||
get_job_submission_client_cluster_info=mock_get_job_submission_client_cluster,
|
||||
),
|
||||
patch.multiple("importlib", import_module=mock_import_module),
|
||||
):
|
||||
if module_string == "ray":
|
||||
with pytest.raises(ValueError, match="ray://"):
|
||||
parse_cluster_info(
|
||||
address,
|
||||
create_cluster_if_needed=create_cluster_if_needed,
|
||||
cookies=cookies,
|
||||
metadata=metadata,
|
||||
headers=headers,
|
||||
**extra_kwargs,
|
||||
)
|
||||
elif module_string == "other_module":
|
||||
assert (
|
||||
parse_cluster_info(
|
||||
address,
|
||||
create_cluster_if_needed=create_cluster_if_needed,
|
||||
cookies=cookies,
|
||||
metadata=metadata,
|
||||
headers=headers,
|
||||
**extra_kwargs,
|
||||
)
|
||||
== "Other module ClusterInfo"
|
||||
)
|
||||
mock_import_module.assert_called_once_with(module_string)
|
||||
mock_module.get_job_submission_client_cluster_info.assert_called_once_with(
|
||||
inner_address,
|
||||
create_cluster_if_needed=create_cluster_if_needed,
|
||||
cookies=cookies,
|
||||
metadata=metadata,
|
||||
headers=headers,
|
||||
**extra_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def test_parse_cluster_info_default_address():
|
||||
assert parse_cluster_info(
|
||||
address=None,
|
||||
) == ClusterInfo(address=DEFAULT_DASHBOARD_ADDRESS)
|
||||
|
||||
|
||||
def test_submit_job_does_not_mutate_runtime_env():
|
||||
class TestClient(JobSubmissionClient):
|
||||
def __init__(self):
|
||||
self._default_metadata = {}
|
||||
|
||||
def _upload_working_dir_if_needed(self, runtime_env):
|
||||
runtime_env["working_dir"] = "gcs://test.zip"
|
||||
|
||||
def _upload_py_modules_if_needed(self, runtime_env):
|
||||
runtime_env["py_modules"] = ["gcs://test_module.zip"]
|
||||
|
||||
def _do_request(self, method, endpoint, **kwargs):
|
||||
return MagicMock(
|
||||
status_code=200,
|
||||
json=lambda: {"job_id": "test_job", "submission_id": "test_job"},
|
||||
)
|
||||
|
||||
runtime_env = {"working_dir": "/tmp/test", "py_modules": ["/tmp/test_module"]}
|
||||
original_runtime_env = {
|
||||
"working_dir": runtime_env["working_dir"],
|
||||
"py_modules": list(runtime_env["py_modules"]),
|
||||
}
|
||||
|
||||
assert (
|
||||
TestClient().submit_job(entrypoint="echo hi", runtime_env=runtime_env)
|
||||
== "test_job"
|
||||
)
|
||||
assert runtime_env == original_runtime_env
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expiration_s", [0, 10])
|
||||
def test_temporary_uri_reference(monkeypatch, expiration_s):
|
||||
"""Test that temporary GCS URI references are deleted after expiration_s."""
|
||||
monkeypatch.setenv(
|
||||
"RAY_RUNTIME_ENV_TEMPORARY_REFERENCE_EXPIRATION_S", str(expiration_s)
|
||||
)
|
||||
# We can't use a fixture with a shared Ray runtime because we need to set the
|
||||
# expiration_s env var before Ray starts.
|
||||
with _ray_start(include_dashboard=True, num_cpus=1) as ctx:
|
||||
headers = {"Connection": "keep-alive", "Authorization": "TOK:<MY_TOKEN>"}
|
||||
address = ctx.address_info["webui_url"]
|
||||
assert wait_until_server_available(address)
|
||||
client = JobSubmissionClient(format_web_url(address), headers=headers)
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
path = Path(tmp_dir)
|
||||
|
||||
hello_file = path / "hi.txt"
|
||||
with hello_file.open(mode="w") as f:
|
||||
f.write("hi\n")
|
||||
|
||||
start = time.time()
|
||||
|
||||
runtime_env = {"working_dir": tmp_dir}
|
||||
job_id = client.submit_job(entrypoint="echo hi", runtime_env=runtime_env)
|
||||
assert runtime_env == {"working_dir": tmp_dir}
|
||||
wait_for_condition(
|
||||
_check_job_succeeded, client=client, job_id=job_id, timeout=30
|
||||
)
|
||||
|
||||
# Give time for deletion to occur if expiration_s is 0.
|
||||
time.sleep(2)
|
||||
# Need to connect to Ray to check internal_kv.
|
||||
# ray.init(address="auto")
|
||||
|
||||
print("Starting Internal KV checks at time ", time.time() - start)
|
||||
if expiration_s > 0:
|
||||
assert not check_internal_kv_gced()
|
||||
wait_for_condition(check_internal_kv_gced, timeout=2 * expiration_s)
|
||||
assert expiration_s < time.time() - start < 2 * expiration_s
|
||||
print("Internal KV was GC'ed at time ", time.time() - start)
|
||||
else:
|
||||
wait_for_condition(check_internal_kv_gced)
|
||||
print("Internal KV was GC'ed at time ", time.time() - start)
|
||||
|
||||
# Regression test for #46625: reusing the same runtime_env after
|
||||
# the package has been GC'ed should re-upload the local working_dir.
|
||||
job_id = client.submit_job(
|
||||
entrypoint="echo hi", runtime_env=runtime_env
|
||||
)
|
||||
wait_for_condition(
|
||||
_check_job_succeeded, client=client, job_id=job_id, timeout=30
|
||||
)
|
||||
|
||||
|
||||
def get_register_agents_number(gcs_client):
|
||||
keys = gcs_client.internal_kv_keys(
|
||||
prefix=DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX,
|
||||
namespace=KV_NAMESPACE_DASHBOARD,
|
||||
timeout=GCS_RPC_TIMEOUT_SECONDS,
|
||||
)
|
||||
return len(keys)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster_head_with_env_vars",
|
||||
[
|
||||
{
|
||||
"include_dashboard": True,
|
||||
"env_vars": {RAY_JOB_ALLOW_DRIVER_ON_WORKER_NODES_ENV_VAR: "1"},
|
||||
},
|
||||
{
|
||||
"include_dashboard": True,
|
||||
"env_vars": {RAY_JOB_ALLOW_DRIVER_ON_WORKER_NODES_ENV_VAR: "0"},
|
||||
},
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_jobs_run_on_head_by_default_E2E(ray_start_cluster_head_with_env_vars):
|
||||
allow_driver_on_worker_nodes = (
|
||||
os.environ.get(RAY_JOB_ALLOW_DRIVER_ON_WORKER_NODES_ENV_VAR) == "1"
|
||||
)
|
||||
# Cluster setup
|
||||
cluster = ray_start_cluster_head_with_env_vars
|
||||
cluster.add_node(dashboard_agent_listen_port=52366)
|
||||
cluster.add_node(dashboard_agent_listen_port=52367)
|
||||
assert wait_until_server_available(cluster.webui_url) is True
|
||||
webui_url = cluster.webui_url
|
||||
webui_url = format_web_url(webui_url)
|
||||
client = JobSubmissionClient(webui_url)
|
||||
gcs_client = GcsClient(address=cluster.gcs_address)
|
||||
|
||||
def _check_nodes(num_nodes):
|
||||
try:
|
||||
assert len(list_nodes()) == num_nodes
|
||||
return True
|
||||
except Exception as ex:
|
||||
print(ex)
|
||||
return False
|
||||
|
||||
wait_for_condition(lambda: _check_nodes(num_nodes=3), timeout=15)
|
||||
wait_for_condition(lambda: get_register_agents_number(gcs_client) == 3, timeout=20)
|
||||
|
||||
# Submit 20 simple jobs.
|
||||
for i in range(20):
|
||||
client.submit_job(entrypoint="echo hi", submission_id=f"job_{i}")
|
||||
import pprint
|
||||
|
||||
def check_all_jobs_succeeded():
|
||||
submission_jobs = [
|
||||
job for job in client.list_jobs() if job.type == JobType.SUBMISSION
|
||||
]
|
||||
for job in submission_jobs:
|
||||
pprint.pprint(job)
|
||||
if job.status != JobStatus.SUCCEEDED:
|
||||
return False
|
||||
return True
|
||||
|
||||
# Wait until all jobs have finished.
|
||||
wait_for_condition(check_all_jobs_succeeded, timeout=60, retry_interval_ms=1000)
|
||||
|
||||
# Check driver_node_id of all jobs.
|
||||
submission_jobs = [
|
||||
job for job in client.list_jobs() if job.type == JobType.SUBMISSION
|
||||
]
|
||||
driver_node_ids = [job.driver_node_id for job in submission_jobs]
|
||||
|
||||
# Spuriously fails with probability (1/3)^20.
|
||||
pprint.pprint(driver_node_ids)
|
||||
num_ids = len(set(driver_node_ids))
|
||||
assert (num_ids > 1) if allow_driver_on_worker_nodes else (num_ids == 1), [
|
||||
id[:5] for id in driver_node_ids
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runtime_env_working_dir():
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
path = Path(tmp_dir)
|
||||
working_dir = path / "working_dir"
|
||||
working_dir.mkdir(parents=True)
|
||||
yield working_dir
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def py_module_whl():
|
||||
with tempfile.NamedTemporaryFile(suffix=".whl") as tmp_file:
|
||||
yield tmp_file.name
|
||||
|
||||
|
||||
def test_job_submission_with_runtime_env_as_dict(
|
||||
runtime_env_working_dir, py_module_whl
|
||||
):
|
||||
working_dir_str = str(runtime_env_working_dir)
|
||||
with _ray_start(num_cpus=1):
|
||||
client = JobSubmissionClient()
|
||||
runtime_env = {"working_dir": working_dir_str, "py_modules": [py_module_whl]}
|
||||
job_id = client.submit_job(entrypoint="echo hi", runtime_env=runtime_env)
|
||||
job_details = client.get_job_info(job_id)
|
||||
parsed_runtime_env = job_details.runtime_env
|
||||
assert "gcs://" in parsed_runtime_env["working_dir"]
|
||||
assert len(parsed_runtime_env["py_modules"]) == 1
|
||||
assert "gcs://" in parsed_runtime_env["py_modules"][0]
|
||||
|
||||
|
||||
def test_job_submission_with_runtime_env_as_object(
|
||||
runtime_env_working_dir, py_module_whl
|
||||
):
|
||||
working_dir_str = str(runtime_env_working_dir)
|
||||
with _ray_start(num_cpus=1):
|
||||
client = JobSubmissionClient()
|
||||
runtime_env = RuntimeEnv(
|
||||
working_dir=working_dir_str, py_modules=[py_module_whl]
|
||||
)
|
||||
job_id = client.submit_job(entrypoint="echo hi", runtime_env=runtime_env)
|
||||
job_details = client.get_job_info(job_id)
|
||||
parsed_runtime_env = job_details.runtime_env
|
||||
assert "gcs://" in parsed_runtime_env["working_dir"]
|
||||
assert len(parsed_runtime_env["py_modules"]) == 1
|
||||
assert "gcs://" in parsed_runtime_env["py_modules"][0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tail_job_logs_passes_headers_to_websocket(ray_start_regular):
|
||||
"""
|
||||
Test that authentication headers are passed to WebSocket connections.
|
||||
|
||||
This test verifies that headers provided to JobSubmissionClient are
|
||||
explicitly passed to the ws_connect() method, not just to the ClientSession.
|
||||
This is required because aiohttp's ClientSession does not automatically
|
||||
include session headers in WebSocket upgrade requests.
|
||||
"""
|
||||
dashboard_url = ray_start_regular.dashboard_url
|
||||
test_headers = {"Authorization": "Bearer test-token"}
|
||||
client = JobSubmissionClient(format_web_url(dashboard_url), headers=test_headers)
|
||||
|
||||
# Submit a simple job
|
||||
job_id = client.submit_job(entrypoint="echo hello")
|
||||
|
||||
# Mock the aiohttp ClientSession and WebSocket
|
||||
mock_ws = AsyncMock()
|
||||
mock_ws.receive = AsyncMock()
|
||||
mock_ws.receive.side_effect = [
|
||||
# First call returns a text message
|
||||
MagicMock(type=aiohttp.WSMsgType.TEXT, data="test log line\n"),
|
||||
# Second call indicates WebSocket is closed
|
||||
MagicMock(type=aiohttp.WSMsgType.CLOSED),
|
||||
]
|
||||
mock_ws.close_code = 1000 # Normal closure
|
||||
|
||||
mock_session = AsyncMock()
|
||||
mock_session.ws_connect = AsyncMock(return_value=mock_ws)
|
||||
mock_session.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_session.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
# Patch ClientSession to use our mock
|
||||
with patch("aiohttp.ClientSession", return_value=mock_session):
|
||||
# Tail logs
|
||||
log_lines = []
|
||||
async for lines in client.tail_job_logs(job_id):
|
||||
log_lines.append(lines)
|
||||
|
||||
# Verify ws_connect was called with headers
|
||||
mock_session.ws_connect.assert_called_once()
|
||||
call_args = mock_session.ws_connect.call_args
|
||||
|
||||
assert "headers" in call_args.kwargs
|
||||
assert call_args.kwargs["headers"] == test_headers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tail_job_logs_websocket_abnormal_closure(ray_start_regular):
|
||||
"""
|
||||
Test that ABNORMAL_CLOSURE raises RuntimeError when tailing logs.
|
||||
|
||||
This test uses its own Ray cluster and kills the dashboard while tailing logs
|
||||
to simulate an abnormal WebSocket closure.
|
||||
"""
|
||||
dashboard_url = ray_start_regular.dashboard_url
|
||||
client = JobSubmissionClient(format_web_url(dashboard_url))
|
||||
|
||||
# Submit a long-running job
|
||||
driver_script = """
|
||||
import time
|
||||
for i in range(100):
|
||||
print("Hello", i)
|
||||
time.sleep(0.5)
|
||||
"""
|
||||
entrypoint = f"python -c '{driver_script}'"
|
||||
job_id = client.submit_job(entrypoint=entrypoint)
|
||||
|
||||
# Start tailing logs and stop Ray while tailing
|
||||
# Expect RuntimeError when WebSocket closes abnormally
|
||||
with pytest.raises(
|
||||
RuntimeError,
|
||||
match="WebSocket connection closed unexpectedly with close code",
|
||||
):
|
||||
i = 0
|
||||
async for lines in client.tail_job_logs(job_id):
|
||||
print(lines, end="")
|
||||
i += 1
|
||||
|
||||
# Kill the dashboard after receiving a few log lines
|
||||
if i == 3:
|
||||
print("\nKilling the dashboard to close websocket abnormally...")
|
||||
dash_info = worker._global_node.all_processes[PROCESS_TYPE_DASHBOARD][0]
|
||||
psutil.Process(dash_info.process.pid).kill()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,303 @@
|
||||
import os
|
||||
import sys
|
||||
from tempfile import NamedTemporaryFile
|
||||
|
||||
import pytest
|
||||
|
||||
from ray.dashboard.modules.job.common import JobSubmitRequest
|
||||
from ray.dashboard.modules.job.utils import (
|
||||
fast_tail_last_n_lines,
|
||||
file_tail_iterator,
|
||||
parse_and_validate_request,
|
||||
redact_url_password,
|
||||
strip_keys_with_value_none,
|
||||
)
|
||||
|
||||
|
||||
# Polyfill anext() function for Python 3.9 compatibility
|
||||
# May raise StopAsyncIteration.
|
||||
async def anext_polyfill(iterator):
|
||||
return await iterator.__anext__()
|
||||
|
||||
|
||||
# Use the built-in anext() for Python 3.10+, otherwise use our polyfilled function
|
||||
if sys.version_info < (3, 10):
|
||||
anext = anext_polyfill
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tmp():
|
||||
with NamedTemporaryFile() as f:
|
||||
yield f.name
|
||||
|
||||
|
||||
def test_strip_keys_with_value_none():
|
||||
d = {"a": 1, "b": None, "c": 3}
|
||||
assert strip_keys_with_value_none(d) == {"a": 1, "c": 3}
|
||||
d = {"a": 1, "b": 2, "c": 3}
|
||||
assert strip_keys_with_value_none(d) == d
|
||||
d = {"a": 1, "b": None, "c": None}
|
||||
assert strip_keys_with_value_none(d) == {"a": 1}
|
||||
|
||||
|
||||
def test_redact_url_password():
|
||||
url = "http://user:password@host:port"
|
||||
assert redact_url_password(url) == "http://user:<redacted>@host:port"
|
||||
url = "http://user:password@host:port?query=1"
|
||||
assert redact_url_password(url) == "http://user:<redacted>@host:port?query=1"
|
||||
url = "http://user:password@host:port?query=1&password=2"
|
||||
assert (
|
||||
redact_url_password(url)
|
||||
== "http://user:<redacted>@host:port?query=1&password=2"
|
||||
)
|
||||
url = "https://user:password@127.0.0.1:8080"
|
||||
assert redact_url_password(url) == "https://user:<redacted>@127.0.0.1:8080"
|
||||
url = "https://user:password@host:port?query=1"
|
||||
assert redact_url_password(url) == "https://user:<redacted>@host:port?query=1"
|
||||
url = "https://user:password@host:port?query=1&password=2"
|
||||
assert (
|
||||
redact_url_password(url)
|
||||
== "https://user:<redacted>@host:port?query=1&password=2"
|
||||
)
|
||||
|
||||
|
||||
# Mock for aiohttp.web.Request, which should not be constructed directly.
|
||||
class MockRequest:
|
||||
def __init__(self, **kwargs):
|
||||
self._json = kwargs
|
||||
|
||||
async def json(self):
|
||||
return self._json
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mock_request():
|
||||
request = MockRequest(a=1, b=2)
|
||||
assert await request.json() == {"a": 1, "b": 2}
|
||||
request = MockRequest(a=1, b=None)
|
||||
assert await request.json() == {"a": 1, "b": None}
|
||||
|
||||
|
||||
# async test
|
||||
@pytest.mark.asyncio
|
||||
class TestParseAndValidateRequest:
|
||||
async def test_basic(self):
|
||||
request = MockRequest(entrypoint="echo hi")
|
||||
expected = JobSubmitRequest(entrypoint="echo hi")
|
||||
assert await parse_and_validate_request(request, JobSubmitRequest) == expected
|
||||
|
||||
async def test_forward_compatibility(self):
|
||||
request = MockRequest(entrypoint="echo hi", new_client_field=None)
|
||||
expected = JobSubmitRequest(entrypoint="echo hi")
|
||||
assert await parse_and_validate_request(request, JobSubmitRequest) == expected
|
||||
|
||||
|
||||
class TestIterLine:
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_type(self):
|
||||
with pytest.raises(TypeError, match="path must be a string"):
|
||||
await anext(file_tail_iterator(1))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_not_created(self, tmp):
|
||||
it = file_tail_iterator(tmp)
|
||||
assert await anext(it) is None
|
||||
f = open(tmp, "w")
|
||||
f.write("hi\n")
|
||||
f.flush()
|
||||
assert await anext(it) is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wait_for_newline(self, tmp):
|
||||
it = file_tail_iterator(tmp)
|
||||
assert await anext(it) is None
|
||||
|
||||
f = open(tmp, "w")
|
||||
f.write("no_newline_yet")
|
||||
assert await anext(it) is None
|
||||
f.write("\n")
|
||||
f.flush()
|
||||
assert await anext(it) == ["no_newline_yet\n"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_lines(self, tmp):
|
||||
it = file_tail_iterator(tmp)
|
||||
assert await anext(it) is None
|
||||
|
||||
f = open(tmp, "w")
|
||||
|
||||
num_lines = 10
|
||||
for i in range(num_lines):
|
||||
s = f"{i}\n"
|
||||
f.write(s)
|
||||
f.flush()
|
||||
assert await anext(it) == [s]
|
||||
|
||||
assert await anext(it) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batching(self, tmp):
|
||||
it = file_tail_iterator(tmp)
|
||||
assert await anext(it) is None
|
||||
|
||||
f = open(tmp, "w")
|
||||
|
||||
# Write lines in batches of 10, check that we get them back in batches.
|
||||
for _ in range(100):
|
||||
num_lines = 10
|
||||
for i in range(num_lines):
|
||||
f.write(f"{i}\n")
|
||||
f.flush()
|
||||
|
||||
assert await anext(it) == [f"{i}\n" for i in range(10)]
|
||||
|
||||
assert await anext(it) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_line_batching(self, tmp):
|
||||
it = file_tail_iterator(tmp)
|
||||
assert await anext(it) is None
|
||||
|
||||
f = open(tmp, "w")
|
||||
|
||||
# Write lines in batches of 50, check that we get them back in batches of 10.
|
||||
for _ in range(100):
|
||||
num_lines = 50
|
||||
for i in range(num_lines):
|
||||
f.write(f"{i}\n")
|
||||
f.flush()
|
||||
|
||||
assert await anext(it) == [f"{i}\n" for i in range(10)]
|
||||
assert await anext(it) == [f"{i}\n" for i in range(10, 20)]
|
||||
assert await anext(it) == [f"{i}\n" for i in range(20, 30)]
|
||||
assert await anext(it) == [f"{i}\n" for i in range(30, 40)]
|
||||
assert await anext(it) == [f"{i}\n" for i in range(40, 50)]
|
||||
|
||||
assert await anext(it) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_char_batching(self, tmp):
|
||||
it = file_tail_iterator(tmp)
|
||||
assert await anext(it) is None
|
||||
|
||||
f = open(tmp, "w")
|
||||
|
||||
# Write a single line that is 60k characters
|
||||
f.write(f"{'1234567890' * 6000}\n")
|
||||
# Write a 4 lines that are 10k characters each
|
||||
for _ in range(4):
|
||||
f.write(f"{'1234567890' * 500}\n")
|
||||
f.flush()
|
||||
|
||||
# First line will come in a batch of its own
|
||||
assert await anext(it) == [f"{'1234567890' * 6000}\n"]
|
||||
# Other 4 lines will be batched together
|
||||
assert (
|
||||
await anext(it)
|
||||
== [
|
||||
f"{'1234567890' * 500}\n",
|
||||
]
|
||||
* 4
|
||||
)
|
||||
assert await anext(it) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_file(self):
|
||||
with NamedTemporaryFile() as tmp:
|
||||
it = file_tail_iterator(tmp.name)
|
||||
f = open(tmp.name, "w")
|
||||
|
||||
assert await anext(it) is None
|
||||
|
||||
f.write("hi\n")
|
||||
f.flush()
|
||||
|
||||
assert await anext(it) == ["hi\n"]
|
||||
|
||||
# Calls should continue returning None after file deleted.
|
||||
assert await anext(it) is None
|
||||
|
||||
|
||||
class TestFastTailLastNLines:
|
||||
def test_nonexistent_path(self, tmp):
|
||||
missing = tmp + ".missing"
|
||||
assert not os.path.exists(missing)
|
||||
with pytest.raises(FileNotFoundError):
|
||||
fast_tail_last_n_lines(missing, num_lines=10, max_chars=1000)
|
||||
|
||||
def test_basic_last_n(self, tmp):
|
||||
# Write 100 lines, check that we get the last 10 lines.
|
||||
with open(tmp, "w") as f:
|
||||
for i in range(100):
|
||||
f.write(f"line-{i}\n")
|
||||
out = fast_tail_last_n_lines(tmp, num_lines=10, max_chars=1000)
|
||||
expected = "".join([f"line-{i}\n" for i in range(90, 100)])
|
||||
assert out == expected
|
||||
|
||||
def test_truncate_max_chars(self, tmp):
|
||||
# Construct a log file with two lines, each over max_chars,
|
||||
# check that we truncate to max_chars.
|
||||
with open(tmp, "w") as f:
|
||||
f.write("x" * 5000 + "\n")
|
||||
f.write("y" * 5000 + "\n")
|
||||
out = fast_tail_last_n_lines(tmp, num_lines=2, max_chars=3000)
|
||||
assert len(out) == 3000
|
||||
# Check that we truncate to max_chars, and include the last line.
|
||||
assert out.endswith("\n")
|
||||
|
||||
def test_partial_last_line(self, tmp):
|
||||
# Write a log file with a partial last line, check that we include it.
|
||||
with open(tmp, "w") as f:
|
||||
f.write("a\n")
|
||||
f.write("b\n")
|
||||
f.write("partial_last_line") # No newline at end
|
||||
out = fast_tail_last_n_lines(tmp, num_lines=3, max_chars=1000)
|
||||
assert out == "a\nb\npartial_last_line"
|
||||
|
||||
def test_small_block_size(self, tmp):
|
||||
# Write 30 lines, check that we can read a small block size and get the last N lines.
|
||||
with open(tmp, "w") as f:
|
||||
for i in range(30):
|
||||
f.write(f"{i}\n")
|
||||
out = fast_tail_last_n_lines(tmp, num_lines=5, max_chars=1000, block_size=16)
|
||||
expected = "".join([f"{i}\n" for i in range(25, 30)])
|
||||
assert out == expected
|
||||
|
||||
def test_mixed_long_lines(self, tmp):
|
||||
# Write a log file with a mix of short and long lines, check that we get the last N lines.
|
||||
with open(tmp, "w") as f:
|
||||
f.write("short-1\n")
|
||||
f.write("short-2\n")
|
||||
f.write("long-" + ("Z" * 10000) + "\n")
|
||||
f.write("short-3\n")
|
||||
f.write("short-4\n")
|
||||
out = fast_tail_last_n_lines(tmp, num_lines=3, max_chars=20000)
|
||||
# Check that we get the last 3 lines, including the long line.
|
||||
assert out.splitlines()[-1] == "short-4"
|
||||
assert out.splitlines()[-2] == "short-3"
|
||||
assert out.splitlines()[-3].startswith("long-Z")
|
||||
|
||||
def test_sparse_large_file_tail_max_chars(self, tmp):
|
||||
"""Simulate ~8 GiB sparse file tail and verify max_chars=20000 truncation."""
|
||||
size_8g = 8 * 1024 * 1024 * 1024
|
||||
# Build tail of two extremely long lines
|
||||
tail = "\n" + ("Q" * 25000 + "\n") + ("R" * 25000 + "\n")
|
||||
tail_bytes = tail.encode("utf-8")
|
||||
|
||||
print("Start writing sparse file tail...")
|
||||
# Create a sparse file: seek to near EOF then write only the tail.
|
||||
with open(tmp, "wb") as f:
|
||||
f.seek(size_8g - len(tail_bytes))
|
||||
f.write(tail_bytes)
|
||||
f.flush()
|
||||
|
||||
print("Finish writing sparse file tail.")
|
||||
out = fast_tail_last_n_lines(tmp, num_lines=2, max_chars=20000)
|
||||
print("Finish reading sparse file tail.")
|
||||
assert len(out) == 20000
|
||||
assert out.endswith("\n")
|
||||
assert "R" * 100 in out # sampling check for last line content
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,361 @@
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import traceback
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, AsyncIterator, Dict, List, Optional, Tuple, Union
|
||||
|
||||
from ray._raylet import RAY_INTERNAL_NAMESPACE_PREFIX, GcsClient
|
||||
from ray.dashboard.modules.job.common import (
|
||||
JOB_ID_METADATA_KEY,
|
||||
JobInfoStorageClient,
|
||||
JobStatus,
|
||||
validate_request_type,
|
||||
)
|
||||
from ray.dashboard.modules.job.pydantic_models import DriverInfo, JobDetails, JobType
|
||||
from ray.runtime_env import RuntimeEnv
|
||||
|
||||
try:
|
||||
# package `aiohttp` is not in ray's minimal dependencies
|
||||
import aiohttp
|
||||
from aiohttp.web import Request, Response
|
||||
except Exception:
|
||||
aiohttp = None
|
||||
Request = None
|
||||
Response = None
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MAX_CHUNK_LINE_LENGTH = 10
|
||||
MAX_CHUNK_CHAR_LENGTH = 20000
|
||||
|
||||
|
||||
def strip_keys_with_value_none(d: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Strip keys with value None from a dictionary."""
|
||||
return {k: v for k, v in d.items() if v is not None}
|
||||
|
||||
|
||||
def redact_url_password(url: str) -> str:
|
||||
"""Redact any passwords in a URL."""
|
||||
secret = re.findall(r"https?:\/\/.*:(.*)@.*", url)
|
||||
if len(secret) > 0:
|
||||
url = url.replace(f":{secret[0]}@", ":<redacted>@")
|
||||
|
||||
return url
|
||||
|
||||
|
||||
async def file_tail_iterator(path: str) -> AsyncIterator[Optional[List[str]]]:
|
||||
"""Yield lines from a file as it's written.
|
||||
|
||||
Returns lines in batches of up to 10 lines or 20000 characters,
|
||||
whichever comes first. If it's a chunk of 20000 characters, then
|
||||
the last line that is yielded could be an incomplete line.
|
||||
New line characters are kept in the line string.
|
||||
|
||||
Returns None until the file exists or if no new line has been written.
|
||||
"""
|
||||
if not isinstance(path, str):
|
||||
raise TypeError(f"path must be a string, got {type(path)}.")
|
||||
|
||||
while not os.path.exists(path):
|
||||
logger.debug(f"Path {path} doesn't exist yet.")
|
||||
yield None
|
||||
|
||||
EOF = ""
|
||||
|
||||
with open(path, "r") as f:
|
||||
lines = []
|
||||
|
||||
chunk_char_count = 0
|
||||
curr_line = None
|
||||
|
||||
while True:
|
||||
# We want to flush current chunk in following cases:
|
||||
# - We accumulated 10 lines
|
||||
# - We accumulated at least MAX_CHUNK_CHAR_LENGTH total chars
|
||||
# - We reached EOF
|
||||
if (
|
||||
len(lines) >= 10
|
||||
or chunk_char_count > MAX_CHUNK_CHAR_LENGTH
|
||||
or curr_line == EOF
|
||||
):
|
||||
# Too many lines, return 10 lines in this chunk, and then
|
||||
# continue reading the file.
|
||||
yield lines or None
|
||||
|
||||
lines = []
|
||||
chunk_char_count = 0
|
||||
|
||||
# Read next line
|
||||
curr_line = f.readline()
|
||||
|
||||
# `readline` will return
|
||||
# - '' for EOF
|
||||
# - '\n' for an empty line in the file
|
||||
if curr_line != EOF:
|
||||
# Add line to current chunk
|
||||
lines.append(curr_line)
|
||||
chunk_char_count += len(curr_line)
|
||||
else:
|
||||
# If EOF is reached sleep for 1s before continuing
|
||||
await asyncio.sleep(1)
|
||||
|
||||
|
||||
async def parse_and_validate_request(
|
||||
req: Request, request_type: dataclass
|
||||
) -> Union[dataclass, Response]:
|
||||
"""Parse request and cast to request type.
|
||||
|
||||
Remove keys with value None to allow newer client versions with new optional fields
|
||||
to work with older servers.
|
||||
|
||||
If parsing failed, return a Response object with status 400 and stacktrace instead.
|
||||
|
||||
Args:
|
||||
req: aiohttp request object.
|
||||
request_type: dataclass type to cast request to.
|
||||
|
||||
Returns:
|
||||
Parsed request object or Response object with status 400 and stacktrace.
|
||||
"""
|
||||
import aiohttp
|
||||
|
||||
json_data = strip_keys_with_value_none(await req.json())
|
||||
try:
|
||||
return validate_request_type(json_data, request_type)
|
||||
except Exception as e:
|
||||
logger.info(f"Got invalid request type: {e}")
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
|
||||
async def get_driver_jobs(
|
||||
gcs_client: GcsClient,
|
||||
job_or_submission_id: Optional[str] = None,
|
||||
timeout: Optional[int] = None,
|
||||
) -> Tuple[Dict[str, JobDetails], Dict[str, DriverInfo]]:
|
||||
"""Returns a tuple of dictionaries related to drivers.
|
||||
|
||||
The first dictionary contains all driver jobs and is keyed by the job's id.
|
||||
The second dictionary contains drivers that belong to submission jobs.
|
||||
It's keyed by the submission job's submission id.
|
||||
Only the last driver of a submission job is returned.
|
||||
|
||||
An optional job_or_submission_id filter can be provided to only return
|
||||
jobs with the job id or submission id.
|
||||
"""
|
||||
job_infos = await gcs_client.async_get_all_job_info(
|
||||
job_or_submission_id=job_or_submission_id,
|
||||
skip_submission_job_info_field=True,
|
||||
skip_is_running_tasks_field=True,
|
||||
timeout=timeout,
|
||||
)
|
||||
# Sort jobs from GCS to follow convention of returning only last driver
|
||||
# of submission job.
|
||||
sorted_job_infos = sorted(
|
||||
job_infos.values(), key=lambda job_table_entry: job_table_entry.job_id.hex()
|
||||
)
|
||||
|
||||
jobs = {}
|
||||
submission_job_drivers = {}
|
||||
for job_table_entry in sorted_job_infos:
|
||||
if job_table_entry.config.ray_namespace.startswith(
|
||||
RAY_INTERNAL_NAMESPACE_PREFIX
|
||||
):
|
||||
# Skip jobs in any _ray_internal_ namespace
|
||||
continue
|
||||
job_id = job_table_entry.job_id.hex()
|
||||
metadata = dict(job_table_entry.config.metadata)
|
||||
job_submission_id = metadata.get(JOB_ID_METADATA_KEY)
|
||||
if not job_submission_id:
|
||||
driver = DriverInfo(
|
||||
id=job_id,
|
||||
node_ip_address=job_table_entry.driver_address.ip_address,
|
||||
pid=str(job_table_entry.driver_pid),
|
||||
)
|
||||
job = JobDetails(
|
||||
job_id=job_id,
|
||||
type=JobType.DRIVER,
|
||||
status=JobStatus.SUCCEEDED
|
||||
if job_table_entry.is_dead
|
||||
else JobStatus.RUNNING,
|
||||
entrypoint=job_table_entry.entrypoint,
|
||||
start_time=job_table_entry.start_time,
|
||||
end_time=job_table_entry.end_time,
|
||||
metadata=metadata,
|
||||
runtime_env=RuntimeEnv.deserialize(
|
||||
job_table_entry.config.runtime_env_info.serialized_runtime_env
|
||||
).to_dict(),
|
||||
driver_info=driver,
|
||||
)
|
||||
jobs[job_id] = job
|
||||
else:
|
||||
driver = DriverInfo(
|
||||
id=job_id,
|
||||
node_ip_address=job_table_entry.driver_address.ip_address,
|
||||
pid=str(job_table_entry.driver_pid),
|
||||
)
|
||||
submission_job_drivers[job_submission_id] = driver
|
||||
|
||||
return jobs, submission_job_drivers
|
||||
|
||||
|
||||
async def find_job_by_ids(
|
||||
gcs_client: GcsClient,
|
||||
job_info_client: JobInfoStorageClient,
|
||||
job_or_submission_id: str,
|
||||
) -> Optional[JobDetails]:
|
||||
"""
|
||||
Attempts to find the job with a given submission_id or job id.
|
||||
"""
|
||||
# First try to find by job_id
|
||||
driver_jobs, submission_job_drivers = await get_driver_jobs(
|
||||
gcs_client, job_or_submission_id=job_or_submission_id
|
||||
)
|
||||
job = driver_jobs.get(job_or_submission_id)
|
||||
if job:
|
||||
return job
|
||||
# Try to find a driver with the given id
|
||||
submission_id = next(
|
||||
(
|
||||
id
|
||||
for id, driver in submission_job_drivers.items()
|
||||
if driver.id == job_or_submission_id
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
if not submission_id:
|
||||
# If we didn't find a driver with the given id,
|
||||
# then lets try to search for a submission with given id
|
||||
submission_id = job_or_submission_id
|
||||
|
||||
job_info = await job_info_client.get_info(submission_id)
|
||||
if job_info:
|
||||
driver = submission_job_drivers.get(submission_id)
|
||||
job = JobDetails(
|
||||
**dataclasses.asdict(job_info),
|
||||
submission_id=submission_id,
|
||||
job_id=driver.id if driver else None,
|
||||
driver_info=driver,
|
||||
type=JobType.SUBMISSION,
|
||||
)
|
||||
return job
|
||||
|
||||
return None
|
||||
|
||||
|
||||
async def find_jobs_by_job_ids(
|
||||
gcs_client: GcsClient,
|
||||
job_info_client: JobInfoStorageClient,
|
||||
job_ids: List[str],
|
||||
) -> Dict[str, JobDetails]:
|
||||
"""
|
||||
Returns a dictionary of submission jobs with the given job ids, keyed by the job id.
|
||||
|
||||
This only accepts job ids and not submission ids.
|
||||
"""
|
||||
driver_jobs, submission_job_drivers = await get_driver_jobs(gcs_client)
|
||||
|
||||
# Filter down to the request job_ids
|
||||
driver_jobs = {key: job for key, job in driver_jobs.items() if key in job_ids}
|
||||
submission_job_drivers = {
|
||||
key: job for key, job in submission_job_drivers.items() if job.id in job_ids
|
||||
}
|
||||
|
||||
# Fetch job details for each job
|
||||
job_submission_ids = submission_job_drivers.keys()
|
||||
job_infos = await asyncio.gather(
|
||||
*[
|
||||
job_info_client.get_info(submission_id)
|
||||
for submission_id in job_submission_ids
|
||||
]
|
||||
)
|
||||
|
||||
return {
|
||||
**driver_jobs,
|
||||
**{
|
||||
submission_job_drivers.get(submission_id).id: JobDetails(
|
||||
**dataclasses.asdict(job_info),
|
||||
submission_id=submission_id,
|
||||
job_id=submission_job_drivers.get(submission_id).id,
|
||||
driver_info=submission_job_drivers.get(submission_id),
|
||||
type=JobType.SUBMISSION,
|
||||
)
|
||||
for job_info, submission_id in zip(job_infos, job_submission_ids)
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def fast_tail_last_n_lines(
|
||||
path: str,
|
||||
num_lines: int,
|
||||
max_chars: int,
|
||||
block_size: int = 8192,
|
||||
) -> str:
|
||||
"""Return the last ``num_lines`` lines from a large log file efficiently.
|
||||
|
||||
This function avoids scanning the entire file. It seeks to the end of
|
||||
the file and reads backwards in fixed-size blocks until enough lines are
|
||||
collected. This is much faster for large files compared to using
|
||||
``file_tail_iterator()``, which performs a full sequential scan.
|
||||
|
||||
Args:
|
||||
path: The file path to read.
|
||||
num_lines: Number of lines to return.
|
||||
max_chars: Maximum number of characters in the returned string.
|
||||
block_size: Read size for each backward block.
|
||||
|
||||
Returns:
|
||||
A string containing at most ``num_lines`` of the last lines in the file,
|
||||
truncated to ``max_chars`` characters.
|
||||
"""
|
||||
if num_lines < 0:
|
||||
raise ValueError(f"num_lines must be non-negative, got {num_lines}")
|
||||
if num_lines == 0:
|
||||
return ""
|
||||
if max_chars < 0:
|
||||
raise ValueError(f"max_chars must be non-negative, got {max_chars}")
|
||||
if max_chars == 0:
|
||||
return ""
|
||||
if block_size <= 0:
|
||||
raise ValueError(f"block_size must be positive, got {block_size}")
|
||||
|
||||
logger.debug(
|
||||
f"Start reading log file {path} with num_lines={num_lines} max_chars={max_chars} block_size={block_size}"
|
||||
)
|
||||
with open(path, "rb") as f:
|
||||
f.seek(0, os.SEEK_END)
|
||||
file_size = f.tell()
|
||||
if file_size == 0:
|
||||
return ""
|
||||
|
||||
chunks = []
|
||||
position = file_size
|
||||
newlines_found = 0
|
||||
|
||||
# We read backwards in chunks until we have enough newlines for num_lines.
|
||||
# We may need one more newline to capture the content before the first newline.
|
||||
while position > 0 and newlines_found < num_lines + 1:
|
||||
read_size = min(block_size, position)
|
||||
position -= read_size
|
||||
f.seek(position)
|
||||
|
||||
chunk = f.read(read_size)
|
||||
newlines_found += chunk.count(b"\n")
|
||||
chunks.insert(0, chunk)
|
||||
|
||||
buffer = b"".join(chunks)
|
||||
lines = buffer.decode("utf-8", errors="replace").splitlines(keepends=True)
|
||||
|
||||
if len(lines) <= num_lines:
|
||||
result = "".join(lines)
|
||||
else:
|
||||
result = "".join(lines[-num_lines:])
|
||||
|
||||
return result[-max_chars:]
|
||||
Reference in New Issue
Block a user