chore: import upstream snapshot with attribution

python/ray/dashboard/modules/job/cli.py

@@ -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)