From 70664c53b59abc4886202cfee445af23e4365b11 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:51:08 -0700 Subject: [PATCH 1/7] feat: add rp endpoint command group --- runpod/rp_cli/endpoint.py | 1775 +++++++++++++++++++++++ runpod/rp_cli/main.py | 3 +- tests/test_cli/test_endpoint_command.py | 1126 ++++++++++++++ 3 files changed, 2903 insertions(+), 1 deletion(-) create mode 100644 runpod/rp_cli/endpoint.py create mode 100644 tests/test_cli/test_endpoint_command.py diff --git a/runpod/rp_cli/endpoint.py b/runpod/rp_cli/endpoint.py new file mode 100644 index 00000000..b8a69709 --- /dev/null +++ b/runpod/rp_cli/endpoint.py @@ -0,0 +1,1775 @@ +"""`rp endpoint`: create, inspect, scale, observe, and call serverless endpoints. + +configuration goes through the rest v2 control plane (/v2/serverless). +jobs go through the serverless data plane (runpod.endpoint_url_base), +which has no rest v2 equivalent. +""" + +import contextlib +import fnmatch +import json +import queue +import shlex +import threading +import time +from datetime import datetime, timezone +from typing import Any, Iterator, Optional +from urllib.parse import quote + +import click +import requests + +from runpod import error +from runpod.api.rest import ( + _build_headers, + _build_url, + _raise_for_error, + _resolve_api_key, + run_rest_request, +) +from runpod.endpoint.helpers import is_completed +from runpod.rp_cli import compute +from runpod.rp_cli.output import ( + api_errors, + cell, + completer, + confirm, + interactive, + json_option, + pick, + print_fields, + print_json, +) + +TYPES = {"queue": "QUEUE", "load-balancer": "LOAD_BALANCER"} +FLASHBOOT = {"off": "OFF", "standard": "FLASHBOOT", "priority": "PRIORITY_FLASHBOOT"} +DEFAULT_QUEUE_DELAY = 4 +WORKER_POLL_SECONDS = 5.0 + +MUTABLE_FIELDS = { + "name": (str,), + "image": (str,), + "args": (str,), + "env": (dict,), + "disk": (int,), + "ports": (list,), + "registry": (str, type(None)), + "gpu": (dict,), + "cpu": (list,), + "workers": (dict,), + "scaling": (dict,), + "timeout": (int,), + "flashboot": (str,), + "dataCenterIds": (list,), + "mounts": (dict,), +} +NESTED_FIELDS = { + "gpu": {"pools", "excludedTypes", "count", "allowedCudaVersions", "minCudaVersion"}, + "workers": {"min", "max", "idleTimeout"}, + "scaling": {"type", "queueDelay", "requestCount"}, + "mounts": {"network", "global"}, +} + + +def _path(endpoint_id: str, *parts: str) -> str: + segments = [quote(value, safe="") for value in (endpoint_id, *parts)] + return "/v2/serverless/" + "/".join(segments) + + +def _get(endpoint_id: str) -> dict: + return run_rest_request("GET", _path(endpoint_id)) or {} + + +def _list_endpoints() -> list[dict]: + return compute.list_all("/v2/serverless", "endpoints") + + +def _list_workers(endpoint_id: str) -> dict: + return run_rest_request("GET", _path(endpoint_id, "workers")) or {} + + +def _choices(path: str, collection: str): + def fetch(): + return [ + (record["id"], record.get("name") or "") + for record in compute.list_all(path, collection) + ] + + return completer(fetch) + + +complete_endpoint = _choices("/v2/serverless", "endpoints") +complete_network_volume = _choices("/v2/network-volumes", "networkVolumes") +complete_global_volume = _choices("/v2/global-volumes", "globalVolumes") +complete_registry = _choices("/v2/registries", "registries") +complete_template = _choices("/v2/templates", "templates") +complete_datacenter = completer( + lambda: [ + (dc["id"], (dc.get("region") or "").lower()) + for dc in compute.list_datacenters() + ] +) + + +def complete_worker(ctx, param, incomplete): + """worker ids of the endpoint named earlier on the command line.""" + endpoint_id = ctx.params.get("endpoint_id") + if not endpoint_id: + return [] + return completer( + lambda: [ + (worker["id"], (worker.get("status") or "").lower()) + for worker in _list_workers(endpoint_id).get("workers") or [] + ] + )(ctx, param, incomplete) + + +def endpoint_argument(required: bool = True): + return click.argument( + "endpoint_id", required=required, shell_complete=complete_endpoint + ) + + +def _pick_endpoint(endpoint_id: Optional[str]) -> str: + if endpoint_id: + return endpoint_id + if not interactive(): + raise click.UsageError("pass an ENDPOINT_ID") + with api_errors(): + endpoints = _list_endpoints() + return pick( + "endpoint", + [ + (e["id"], f"{e['id']} {e.get('name') or ''} {_type(e) or ''}") + for e in endpoints + ], + "pass an ENDPOINT_ID", + ) + + +# ----------------------------------------------------------------- display + + +def _type(endpoint: dict) -> Optional[str]: + names = {value: key for key, value in TYPES.items()} + return names.get(endpoint.get("type"), endpoint.get("type")) + + +def _flashboot(endpoint: dict) -> Optional[str]: + names = {value: key for key, value in FLASHBOOT.items()} + return names.get(endpoint.get("flashboot"), endpoint.get("flashboot")) + + +def _compute(endpoint: dict, compact: bool = False) -> Optional[str]: + gpu = endpoint.get("gpu") + if gpu: + text = f"{', '.join(gpu.get('pools') or [])} x{gpu.get('count') or 1}" + excluded = gpu.get("excludedTypes") or [] + if compact or not excluded: + return text + return f"{text} (excluding {', '.join(excluded)})" + cpus = endpoint.get("cpu") or [] + return ", ".join(f"{cpu['id']} {cpu.get('vcpuCount')} vcpu" for cpu in cpus) or None + + +def _cuda(endpoint: dict) -> Optional[str]: + gpu = endpoint.get("gpu") + if not gpu: + return None + if gpu.get("minCudaVersion"): + return f">={gpu['minCudaVersion']}" + return ", ".join(gpu.get("allowedCudaVersions") or []) or "any" + + +def _scaling(endpoint: dict, compact: bool = False) -> Optional[str]: + scaling = endpoint.get("scaling") or {} + if scaling.get("type") == "QUEUE_DELAY": + delay = f"{scaling.get('queueDelay'):g}s" + return f"{delay} delay" if compact else f"queue delay {delay}" + if scaling.get("type") == "REQUEST_COUNT": + count = scaling.get("requestCount") + return f"{count}/worker" if compact else f"{count} requests per worker" + return None + + +def _workers(endpoint: dict) -> Optional[str]: + workers = endpoint.get("workers") or {} + if "min" not in workers and "max" not in workers: + return None + return f"{workers.get('min', 0)}-{workers.get('max', 0)}" + + +def _seconds(value: Optional[float]) -> Optional[str]: + return f"{value:g}s" if value else None + + +def _volumes(endpoint: dict) -> list[str]: + mounts = endpoint.get("mounts") or {} + return [ + f"{mount['volumeId']} ({kind})" + for kind in ("network", "global") + for mount in mounts.get(kind) or [] + ] + + +def _url(endpoint: dict) -> Optional[str]: + urls = endpoint.get("requestUrls") or {} + return urls.get("base") or urls.get("run") + + +_LIST_COLUMNS = [ + ("ID", lambda e: e["id"]), + ("NAME", lambda e: e.get("name")), + ("TYPE", _type), + ("COMPUTE", lambda e: _compute(e, compact=True)), + ("WORKERS", _workers), + ("SCALING", lambda e: _scaling(e, compact=True)), + ("FLASHBOOT", _flashboot), + ("CREATED", lambda e: (e.get("createdAt") or "")[:10]), +] + +_FIELDS = [ + ("id", lambda e: e.get("id")), + ("name", lambda e: e.get("name")), + ("type", _type), + ("image", lambda e: e.get("image")), + ("compute", _compute), + ("cuda", _cuda), + ("workers", _workers), + ("idle timeout", lambda e: _seconds((e.get("workers") or {}).get("idleTimeout"))), + ("scaling", _scaling), + ("request timeout", lambda e: _seconds((e.get("timeout") or 0) / 1000)), + ("flashboot", _flashboot), + ("data centers", lambda e: e.get("dataCenterIds")), + ("volumes", _volumes), + ("env", lambda e: sorted(e.get("env") or {})), + ("ports", lambda e: e.get("ports")), + ("disk", lambda e: f"{e['disk']}GB" if e.get("disk") else None), + ("args", lambda e: e.get("args")), + ("registry", lambda e: e.get("registry")), + ("url", _url), + ("created", lambda e: e.get("createdAt")), +] + + +def _escape(text: Any) -> str: + from rich.markup import escape + + return escape(str(text)) + + +def _table(rows: list[dict], columns: list): + """an aligned table whose first column, the id, never wraps or truncates.""" + from rich.markup import escape + from rich.table import Table + + table = Table(box=None, header_style="dim", pad_edge=False) + for index, (header, _) in enumerate(columns): + table.add_column(header, no_wrap=index == 0) + for row in rows: + table.add_row(*(escape(cell(value(row))) for _, value in columns)) + return table + + +def _print_table(rows: list[dict], columns: list, empty: str) -> None: + from runpod.rp_cli import console as ui + + ui.console.print(_table(rows, columns) if rows else f"[dim]{empty}[/dim]") + + +# ------------------------------------------------------------ request body + + +def _env_pairs(values: tuple[str, ...]) -> dict[str, str]: + env = {} + for value in values: + key, sep, val = value.partition("=") + if not sep or not key: + raise click.UsageError(f"invalid --env '{value}'. use KEY=VALUE") + env[key] = val + return env + + +def _gpu_body(specs, count, cuda, updating: bool) -> dict: + gpu: dict[str, Any] = {} + if specs: + catalog = [ + entry + for entry in compute.list_gpus(product="SERVERLESS") + if entry.get("pool") + ] + matches: list[dict] = [] + counts = {count} if count else set() + for spec in specs: + name, spec_count = compute.parse_gpu_spec(spec) + if spec_count: + counts.add(spec_count) + for match in compute.resolve_gpus(name, catalog): + if match not in matches: + matches.append(match) + if len(counts) > 1: + raise click.UsageError("conflicting gpu counts; use one count per worker") + pools, excluded = compute.endpoint_gpu_selection(matches, catalog) + gpu["pools"] = pools + if excluded or updating: + gpu["excludedTypes"] = excluded + count = counts.pop() if counts else None + if count: + gpu["count"] = count + if cuda: + if [value.lower() for value in cuda] == ["any"]: + gpu["allowedCudaVersions"] = [] + if updating: + gpu["minCudaVersion"] = "" + return gpu + minimum, exact = compute.parse_cuda(cuda) + if minimum: + gpu["minCudaVersion"] = minimum + if updating: + gpu["allowedCudaVersions"] = [] + else: + gpu["allowedCudaVersions"] = exact + if updating: + gpu["minCudaVersion"] = "" + return gpu + + +def _cpu_body(specs, vcpu_count, current: Optional[dict]) -> Optional[list]: + if not specs: + if vcpu_count and current and current.get("cpu"): + return [ + {"id": cpu["id"], "vcpuCount": vcpu_count} for cpu in current["cpu"] + ] + if vcpu_count: + raise click.UsageError("--vcpu-count applies to CPU endpoints (pass --cpu)") + return None + flavors = compute.list_cpus(product="SERVERLESS") + entries = [] + for spec in specs: + name, sep, count = spec.rpartition(":") + if not sep: + name, count = spec, "" + if count and not count.isdigit(): + raise click.UsageError( + f"invalid --cpu '{spec}'. use a flavor like cpu3c or cpu3c:8" + ) + vcpus = int(count) if count else vcpu_count + if not vcpus: + raise click.UsageError(f"pass --vcpu-count or --cpu {name}:VCPUS") + entries.append( + {"id": compute.resolve_cpu(name, flavors)["id"], "vcpuCount": vcpus} + ) + return entries + + +def _workers_body(options: dict, current: Optional[dict]) -> Optional[dict]: + changes = { + key: options[name] + for key, name in ( + ("min", "min_workers"), + ("max", "max_workers"), + ("idleTimeout", "idle_timeout_seconds"), + ) + if options.get(name) is not None + } + if not changes: + return None + existing = (current or {}).get("workers") or {} + merged = { + key: existing[key] + for key in ("min", "max", "idleTimeout") + if existing.get(key) is not None + } + merged.update(changes) + return merged + + +def _scaling_body(options: dict, endpoint_type: Optional[str]) -> Optional[dict]: + delay = options.get("queue_delay_seconds") + per_worker = options.get("requests_per_worker") + if delay is not None and per_worker is not None: + raise click.UsageError( + "use either --queue-delay-seconds or --requests-per-worker" + ) + if delay is not None: + if endpoint_type == "LOAD_BALANCER": + raise click.UsageError( + "load-balancer endpoints scale on --requests-per-worker" + ) + return {"type": "QUEUE_DELAY", "queueDelay": delay} + if per_worker is not None: + return {"type": "REQUEST_COUNT", "requestCount": per_worker} + return None + + +def _mounts_body(options: dict, current: Optional[dict]) -> Optional[dict]: + network = options.get("network_volume") or () + global_ = options.get("global_volume") or () + if not network and not global_: + return None + existing = (current or {}).get("mounts") or {} + mounts = {} + for kind, ids in (("network", network), ("global", global_)): + if ids: + mounts[kind] = [{"volumeId": volume} for volume in ids] + elif existing.get(kind): + mounts[kind] = [{"volumeId": m["volumeId"]} for m in existing[kind]] + return mounts + + +def _env_body(options: dict, current: Optional[dict]) -> Optional[dict]: + env = _env_pairs(options.get("env") or ()) + unset = options.get("unset_env") or () + if current is None or not (env or unset): + return env or None + merged = dict(current.get("env") or {}) + merged.update(env) + for key in unset: + if key not in merged: + raise click.UsageError(f"--unset-env {key}: no such variable") + merged.pop(key) + return merged + + +def _check_compute(options: dict, current: Optional[dict]) -> None: + current = current or {} + if options.get("gpu") and options.get("cpu"): + raise click.UsageError("use either --gpu or --cpu") + if options.get("gpu") and current.get("cpu"): + raise click.UsageError("a CPU endpoint cannot switch to GPUs") + if options.get("cpu") and current.get("gpu"): + raise click.UsageError("a GPU endpoint cannot switch to CPUs") + if options.get("cuda") and (options.get("cpu") or current.get("cpu")): + raise click.UsageError("--cuda applies to GPU endpoints") + + +def _config_body(options: dict, current: Optional[dict] = None) -> dict: + """request body for a create (current is None) or an update.""" + updating = current is not None + endpoint_type = current.get("type") if updating else options.get("endpoint_type") + body: dict[str, Any] = {} + if not updating: + body["type"] = endpoint_type + + for key, name in ( + ("name", "name"), + ("image", "image"), + ("templateId", "template"), + ("args", "args"), + ("disk", "disk"), + ("registry", "registry"), + ): + if options.get(name) is not None: + body[key] = options[name] + for key in ("entrypoint", "cmd"): + if options.get(key) is not None: + body[key] = shlex.split(options[key]) + if options.get("port"): + body["ports"] = list(options["port"]) + if options.get("flashboot"): + body["flashboot"] = FLASHBOOT[options["flashboot"].lower()] + if options.get("request_timeout_seconds") is not None: + body["timeout"] = options["request_timeout_seconds"] * 1000 + env = _env_body(options, current) + if env is not None: + body["env"] = env + + _check_compute(options, current) + gpu = _gpu_body( + options.get("gpu"), options.get("gpu_count"), options.get("cuda"), updating + ) + if gpu: + body["gpu"] = gpu + cpu = _cpu_body(options.get("cpu"), options.get("vcpu_count"), current) + if cpu: + body["cpu"] = cpu + + workers = _workers_body(options, current) + if workers: + body["workers"] = workers + scaling = _scaling_body(options, endpoint_type) + if scaling: + body["scaling"] = scaling + elif not updating: + if endpoint_type == "LOAD_BALANCER": + raise click.UsageError("load-balancer endpoints need --requests-per-worker") + body["scaling"] = {"type": "QUEUE_DELAY", "queueDelay": DEFAULT_QUEUE_DELAY} + + places = tuple(options.get("data_center") or ()) + tuple( + options.get("region") or () + ) + if places: + body["dataCenterIds"] = compute.resolve_regions( + places, compute.list_datacenters() + ) + mounts = _mounts_body(options, current) + if mounts: + body["mounts"] = mounts + return body + + +def _scale_options(func): + for option in reversed( + [ + click.option( + "--min-workers", type=click.IntRange(min=0), help="Minimum workers." + ), + click.option( + "--max-workers", type=click.IntRange(min=0), help="Maximum workers." + ), + click.option( + "--idle-timeout-seconds", + type=click.IntRange(1, 3600), + help="Seconds before idle workers scale down.", + ), + click.option( + "--queue-delay-seconds", + type=click.FloatRange(min=0.5), + help="Scale on queue wait time (queue endpoints).", + ), + click.option( + "--requests-per-worker", + type=click.IntRange(min=1), + help="Scale on in-flight requests per worker.", + ), + ] + ): + func = option(func) + return func + + +def _config_options(creating: bool): + """the flags shared by create and update, one per api field.""" + head = [click.option("--name", required=creating, help="Endpoint name.")] + if creating: + head.append( + click.option( + "--type", + "endpoint_type", + required=True, + type=click.Choice(list(TYPES), case_sensitive=False), + callback=lambda ctx, param, value: ( + TYPES[value.lower()] if value else value + ), + help="Request routing: queued jobs or load-balanced HTTP.", + ) + ) + options = head + [ + click.option("--image", help="Container image."), + click.option( + "--template", + help="Serverless template ID to base the container on.", + shell_complete=complete_template, + ), + click.option( + "--gpu", + multiple=True, + help="GPU type (4090, h100, a100-sxm4) or pool id, optionally with a " + "count (h100:2). Repeatable.", + ), + click.option( + "--gpu-count", type=click.IntRange(min=1), help="GPUs per worker." + ), + click.option( + "--cuda", + multiple=True, + help="CUDA constraint: '>=12.4' for a floor, 12.4 or 12.4,12.8 for exact " + "versions, any for no constraint.", + ), + click.option( + "--cpu", + multiple=True, + help="CPU flavor (cpu3c, cpu5g), optionally with vCPUs (cpu3c:8). Repeatable.", + ), + click.option( + "--vcpu-count", type=click.IntRange(min=2), help="vCPUs per worker." + ), + _scale_options, + click.option( + "--request-timeout-seconds", + type=click.IntRange(min=0), + help="Per-request execution timeout.", + ), + click.option( + "--data-center", + multiple=True, + help="Data center id (EU-RO-1). Repeatable.", + shell_complete=complete_datacenter, + ), + click.option( + "--region", + multiple=True, + help="Region (eu, na, asia) or id prefix (us, eu-ro). Repeatable.", + ), + click.option( + "--network-volume", + multiple=True, + help="Network volume ID, at most one per data center. Repeatable.", + shell_complete=complete_network_volume, + ), + click.option( + "--global-volume", + multiple=True, + help="Global volume ID. Repeatable.", + shell_complete=complete_global_volume, + ), + click.option( + "--env", multiple=True, help="Environment variable KEY=VALUE. Repeatable." + ), + click.option( + "--flashboot", + type=click.Choice(list(FLASHBOOT), case_sensitive=False), + help="FlashBoot mode.", + ), + click.option( + "--disk", type=click.IntRange(min=1), help="Container disk in GB." + ), + click.option( + "--port", multiple=True, help="Exposed port as port/protocol. Repeatable." + ), + click.option("--args", help="Container start command as one raw string."), + click.option("--entrypoint", help="Container ENTRYPOINT, shell-quoted."), + click.option("--cmd", help="Container CMD, shell-quoted."), + click.option( + "--registry", + help="Registry credential ID for private images.", + shell_complete=complete_registry, + ), + ] + if not creating: + options.append( + click.option( + "--unset-env", + multiple=True, + help="Remove an environment variable. Repeatable.", + ) + ) + + def decorate(func): + for option in reversed(options): + func = option(func) + return func + + return decorate + + +# ------------------------------------------------------------------ rollout + + +def _rollout_state(endpoint_id: str) -> dict: + return ( + run_rest_request("GET", _path(endpoint_id, "releases"), params={"limit": 1}) + or {} + ) + + +def _rollout_line(state: dict) -> str: + rollout = state.get("rollout") or {} + on_latest = f"{rollout.get('workersOnLatest', 0)}/{rollout.get('workersTotal', 0)}" + phase = "rolling out" if rollout.get("inProgress") else "live" + return f"v{state.get('endpointVersion')} {phase}, {on_latest} workers on latest" + + +def _follow_rollout( + endpoint_id: str, + previous_version: Optional[int], + timeout: float, + interval: float = 2.0, +) -> None: + """poll releases until the new version runs on every worker.""" + from runpod.rp_cli import console as ui + + started = time.monotonic() + last = None + try: + while True: + with api_errors(): + state = _rollout_state(endpoint_id) + version = state.get("endpointVersion") + elapsed = time.monotonic() - started + if previous_version is not None and version == previous_version: + if elapsed >= min(10.0, timeout): + ui.console.print("[dim]no new release to roll out[/dim]") + return + else: + line = _rollout_line(state) + if line != last: + ui.console.print(f"[dim]{line}[/dim]") + last = line + if not (state.get("rollout") or {}).get("inProgress"): + ui.success(f"release v{version} is live") + return + if elapsed >= timeout: + ui.warn( + f"rollout still in progress after {timeout:g}s. " + f"follow with: rp endpoint releases {endpoint_id}" + ) + return + time.sleep(interval) + except KeyboardInterrupt: + ui.console.print( + "\n[dim]detached; the rollout continues. " + f"follow with: rp endpoint releases {endpoint_id}[/dim]" + ) + + +def _apply( + endpoint_id: str, patch: dict, follow: bool, timeout: float, as_json: bool +) -> None: + from runpod.rp_cli import console as ui + + with api_errors(): + before = _rollout_state(endpoint_id).get("endpointVersion") if follow else None + updated = run_rest_request("PATCH", _path(endpoint_id), json=patch) or {} + if as_json: + print_json(updated) + else: + ui.success( + f"updated [white]{_escape(updated.get('name') or endpoint_id)}[/white] " + f"[dim]{endpoint_id}[/dim]" + ) + if follow: + _follow_rollout(endpoint_id, before, timeout) + + +# --------------------------------------------------------------------- edit + + +def _mutable(endpoint: dict) -> dict: + """the endpoint's configuration as UpdateEndpointRequest fields.""" + doc: dict[str, Any] = {} + for key in ("name", "image", "args", "env", "disk", "ports"): + if endpoint.get(key) is not None: + doc[key] = endpoint[key] + doc["registry"] = endpoint.get("registry") + gpu = endpoint.get("gpu") + if gpu: + doc["gpu"] = { + "pools": gpu.get("pools") or [], + "excludedTypes": gpu.get("excludedTypes") or [], + "count": gpu.get("count") or 1, + "allowedCudaVersions": gpu.get("allowedCudaVersions") or [], + "minCudaVersion": gpu.get("minCudaVersion"), + } + if endpoint.get("cpu"): + doc["cpu"] = [ + {"id": cpu["id"], "vcpuCount": cpu["vcpuCount"]} for cpu in endpoint["cpu"] + ] + for key in ("workers", "scaling", "timeout", "flashboot", "dataCenterIds"): + if endpoint.get(key) is not None: + doc[key] = endpoint[key] + mounts = endpoint.get("mounts") or {} + doc["mounts"] = { + kind: [{"volumeId": m["volumeId"]} for m in mounts.get(kind) or []] + for kind in ("network", "global") + } + return doc + + +def _kind(kind: type) -> str: + return { + str: "a string", + int: "an integer", + dict: "an object", + list: "a list", + type(None): "null", + }[kind] + + +def _validate(text: str, before: dict) -> dict: + """parse an edited document, raising ClickException on any problem.""" + try: + doc = json.loads(text) + except json.JSONDecodeError as exc: + raise click.ClickException(f"invalid JSON: {exc}") from exc + if not isinstance(doc, dict): + raise click.ClickException("the document must be a JSON object") + problems = [] + for key, value in doc.items(): + kinds = MUTABLE_FIELDS.get(key) + if kinds is None: + problems.append(f"unknown field '{key}'") + continue + if isinstance(value, bool) or not isinstance(value, kinds): + problems.append(f"'{key}' must be {' or '.join(_kind(k) for k in kinds)}") + continue + allowed = NESTED_FIELDS.get(key) + if allowed: + problems.extend( + f"unknown field '{key}.{nested}'" + for nested in value + if nested not in allowed + ) + env = doc.get("env") + if isinstance(env, dict) and not all(isinstance(v, str) for v in env.values()): + problems.append("env values must be strings") + flashboot = doc.get("flashboot") + if isinstance(flashboot, str) and flashboot not in FLASHBOOT.values(): + problems.append(f"'flashboot' must be one of {', '.join(FLASHBOOT.values())}") + if "gpu" in doc and "gpu" not in before: + problems.append("a CPU endpoint cannot switch to GPUs") + if "cpu" in doc and "cpu" not in before: + problems.append("a GPU endpoint cannot switch to CPUs") + if problems: + raise click.ClickException("; ".join(problems)) + return doc + + +def _patch(before: dict, after: dict) -> dict: + """top-level fields that changed, with gpu narrowed to its changed keys.""" + patch = {} + for key, value in after.items(): + if before.get(key) == value: + continue + if key != "gpu": + patch[key] = value + continue + old = before.get("gpu") or {} + gpu = {k: v for k, v in value.items() if old.get(k) != v} + if "pools" in gpu or "excludedTypes" in gpu: + gpu["pools"] = value.get("pools", old.get("pools")) + gpu["excludedTypes"] = value.get("excludedTypes") or [] + if "minCudaVersion" in gpu and gpu["minCudaVersion"] is None: + gpu["minCudaVersion"] = "" + patch["gpu"] = gpu + return patch + + +def _edit_document(before: dict) -> Optional[dict]: + """open the document in $EDITOR until it validates or the user gives up.""" + from runpod.rp_cli import console as ui + + text = json.dumps(before, indent=2) + "\n" + while True: + edited = click.edit(text, extension=".json") + if edited is None: + return None + try: + return _validate(edited, before) + except click.ClickException as exc: + ui.error(_escape(exc.message)) + if not click.confirm("re-open the editor?", default=True): + raise + text = edited + + +# ------------------------------------------------------------------ workers + + +def _uptime(seconds: Optional[int]) -> Optional[str]: + if not seconds: + return None + hours, rest = divmod(int(seconds), 3600) + minutes, secs = divmod(rest, 60) + if hours: + return f"{hours}h{minutes:02d}m" + return f"{minutes}m{secs:02d}s" if minutes else f"{secs}s" + + +def _worker_version(worker: dict) -> Optional[str]: + stale = " (stale)" if worker.get("isStale") else "" + if worker.get("version") is None: + return stale.strip(" ()") or None + return f"v{worker['version']}{stale}" + + +_WORKER_COLUMNS = [ + ("WORKER", lambda w: w["id"]), + ("STATUS", lambda w: (w.get("status") or "").lower()), + ( + "GPU", + lambda w: ( + f"{w['gpuTypeId']} x{w.get('gpuCount') or 1}" + if w.get("gpuTypeId") + else None + ), + ), + ("DATA CENTER", lambda w: w.get("dataCenterId")), + ("VERSION", _worker_version), + ("UPTIME", lambda w: _uptime(w.get("uptimeSeconds"))), + ("IMAGE", lambda w: w.get("image")), +] + + +def _worker_summary(data: dict) -> str: + summary = data.get("summary") or {} + counts = ", ".join( + f"{summary.get(status, 0)} {status}" + for status in ("running", "idle", "initializing", "throttled", "unhealthy") + ) + version = data.get("endpointVersion") + suffix = f", endpoint v{version}" if version is not None else "" + return f"{summary.get('total', 0)} workers: {counts}{suffix}" + + +def _workers_view(data: dict): + from rich.console import Group + from rich.text import Text + + header = Text(_worker_summary(data), style="dim") + workers = data.get("workers") or [] + if not workers: + return header + return Group(header, _table(workers, _WORKER_COLUMNS)) + + +# --------------------------------------------------------------------- logs + + +def _log_stream( + endpoint_id: str, + worker_id: str, + params: dict, + resume: Optional[str], + events: "queue.Queue", + stop: threading.Event, + read_timeout: float, +) -> None: + """push one worker's sse log events onto the shared queue.""" + path = _path(endpoint_id, "workers", worker_id, "logs") + event_id = None + try: + headers = dict( + _build_headers(_resolve_api_key(None)), Accept="text/event-stream" + ) + if resume: + headers["Last-Event-ID"] = resume + with requests.get( + _build_url(path), + params=params, + headers=headers, + stream=True, + timeout=(10, read_timeout), + ) as response: + if response.status_code >= 400: + _raise_for_error(response.status_code, "GET", path, response.text) + for raw in response.iter_lines(decode_unicode=True): + if stop.is_set(): + return + if raw.startswith("id:"): + event_id = raw[3:].strip() + elif raw.startswith("data:"): + try: + entry = json.loads(raw[5:].strip()) + except json.JSONDecodeError: + continue + events.put(("line", worker_id, entry, event_id or entry.get("ts"))) + except requests.exceptions.RequestException: + pass + except error.RunPodError as exc: + events.put(("error", worker_id, exc, None)) + finally: + events.put(("done", worker_id, None, None)) + + +def _log_line(worker_id: str, entry: dict, width: int, as_json: bool) -> str: + if as_json: + return json.dumps({"worker": worker_id, **entry}) + return ( + f"{worker_id:<{width}} [{entry.get('source') or '?'}] {entry.get('line', '')}" + ) + + +def _is_live(entry: dict, started: datetime) -> bool: + try: + stamp = datetime.fromisoformat(str(entry.get("ts")).replace("Z", "+00:00")) + except ValueError: + return False + if stamp.tzinfo is None: + stamp = stamp.replace(tzinfo=timezone.utc) + return stamp >= started + + +# --------------------------------------------------------------- data plane + + +@contextlib.contextmanager +def _data_plane_errors() -> Iterator[None]: + """surface serverless data-plane failures as clean cli errors.""" + try: + yield + except requests.HTTPError as exc: + response = exc.response + if response is None: + raise click.ClickException(str(exc)) from exc + detail = response.text.strip() or response.reason or "" + raise click.ClickException(f"{response.status_code} {detail}".strip()) from exc + except (requests.RequestException, RuntimeError) as exc: + raise click.ClickException(str(exc)) from exc + + +def _client(): + from runpod.endpoint.runner import RunPodClient + + with _data_plane_errors(): + return RunPodClient() + + +def _job_path(endpoint_id: str, action: str, job_id: Optional[str] = None) -> str: + parts = [quote(endpoint_id, safe=""), action] + if job_id: + parts.append(quote(job_id, safe="")) + return "/".join(parts) + + +def _job_failure(job: dict) -> click.ClickException: + detail = job.get("error") + status = str(job.get("status") or "unknown").lower() + return click.ClickException( + f"job {job.get('id')} {status}" + (f": {detail}" if detail else "") + ) + + +def _wait_job(client, endpoint_id: str, job_id: str, timeout: Optional[float]) -> dict: + from runpod.rp_cli import console as ui + + deadline = time.monotonic() + timeout if timeout else None + try: + with ui.console.status(f"[dim]waiting for job {job_id} ...[/dim]"): + while True: + with _data_plane_errors(): + job = client.get(_job_path(endpoint_id, "status", job_id)) + if is_completed(job.get("status")): + return job + if deadline and time.monotonic() >= deadline: + raise click.ClickException( + f"job {job_id} is still {job.get('status')} after {timeout:g}s. " + f"resume with: rp endpoint job wait {endpoint_id} {job_id}" + ) + time.sleep(1) + except KeyboardInterrupt: + click.echo( + f"\njob {job_id} keeps running. resume with: " + f"rp endpoint job wait {endpoint_id} {job_id}", + err=True, + ) + raise click.Abort() from None + + +def _print_job(job: dict, as_json: bool) -> None: + if as_json: + print_json(job) + elif job.get("status") == "COMPLETED": + output = job.get("output") + click.echo(output if isinstance(output, str) else json.dumps(output, indent=2)) + if job.get("status") != "COMPLETED": + raise _job_failure(job) + + +def _stream_job(client, endpoint_id: str, job_id: str, as_json: bool) -> None: + open_line = False + try: + while True: + with _data_plane_errors(): + partial = client.get(_job_path(endpoint_id, "stream", job_id)) + chunks = partial.get("stream") or [] + for chunk in chunks: + output = chunk.get("output") + if as_json: + click.echo(json.dumps(output)) + elif isinstance(output, str): + click.echo(output, nl=False) + open_line = not output.endswith("\n") + else: + if open_line: + click.echo() + open_line = False + click.echo(json.dumps(output)) + if is_completed(partial.get("status")) and not chunks: + break + if not chunks: + time.sleep(0.5) + except KeyboardInterrupt: + click.echo( + f"\njob {job_id} keeps running. resume with: " + f"rp endpoint job stream {endpoint_id} {job_id}", + err=True, + ) + raise click.Abort() from None + if open_line: + click.echo() + if partial.get("status") != "COMPLETED": + with _data_plane_errors(): + raise _job_failure(client.get(_job_path(endpoint_id, "status", job_id))) + + +def _request_body(input_file, data: Optional[str]) -> Optional[Any]: + if input_file and data: + raise click.UsageError("use either --input or --data") + text = input_file.read() if input_file else data + if text is None: + return None + try: + return json.loads(text) + except json.JSONDecodeError as exc: + raise click.UsageError(f"request body is not valid JSON: {exc}") from exc + + +def _call_load_balancer( + endpoint: dict, path: str, method, body, timeout, as_json +) -> None: + base = (endpoint.get("requestUrls") or {}).get("base") or ( + f"https://{endpoint['id']}.api.runpod.ai" + ) + method = (method or ("POST" if body is not None else "GET")).upper() + url = f"{base.rstrip('/')}/{path.lstrip('/')}" + with _data_plane_errors(): + response = requests.request( + method, + url, + headers=_build_headers(_resolve_api_key(None)), + json=body, + timeout=timeout or 600, + ) + if response.status_code >= 400: + raise click.ClickException( + f"{method} {url} returned {response.status_code}: {response.text.strip()}" + ) + try: + payload = response.json() + except ValueError: + click.echo(response.text) + return + if isinstance(payload, str) and not as_json: + click.echo(payload) + return + print_json(payload) + + +# ----------------------------------------------------------------- commands + + +@click.group() +def endpoint(): + """Create, inspect, scale, observe, and call serverless endpoints. + + \b + rp endpoint list --type queue + rp endpoint create --name llama --type queue --image IMAGE --gpu h100 + rp endpoint edit ENDPOINT_ID + rp endpoint logs ENDPOINT_ID --follow + rp endpoint run ENDPOINT_ID --data '{"prompt":"hello"}' + """ + + +@endpoint.command(name="list") +@click.option( + "--name", help="Only endpoints with this name (globs like 'llama*' work)." +) +@click.option( + "--type", + "endpoint_type", + type=click.Choice(list(TYPES), case_sensitive=False), + help="Only endpoints of this type.", +) +@json_option +def list_endpoints(name, endpoint_type, as_json): + """List endpoints with their IDs.""" + with api_errors(): + endpoints = _list_endpoints() + if name: + endpoints = [ + e + for e in endpoints + if fnmatch.fnmatch((e.get("name") or "").lower(), name.lower()) + ] + if endpoint_type: + endpoints = [ + e for e in endpoints if e.get("type") == TYPES[endpoint_type.lower()] + ] + if as_json: + print_json(endpoints) + return + _print_table(endpoints, _LIST_COLUMNS, "no endpoints match") + + +@endpoint.command(name="get") +@endpoint_argument(required=False) +@json_option +def get_endpoint(endpoint_id, as_json): + """Show an endpoint's configuration (picker when ENDPOINT_ID is omitted).""" + endpoint_id = _pick_endpoint(endpoint_id) + with api_errors(): + record = _get(endpoint_id) + if as_json: + print_json(record) + return + print_fields(record, _FIELDS) + + +@endpoint.command(name="create") +@_config_options(creating=True) +@json_option +def create_endpoint(as_json, **options): + """Create an endpoint. + + GPU types are human-readable and resolve against the serverless + catalog, and regions expand to their data centers. + + \b + rp endpoint create --name llama-prod --type queue \\ + --image ghcr.io/acme/llama:latest --gpu h100 --max-workers 4 \\ + --region eu --env MODEL_ID=meta-llama/Llama-3.1-8B --flashboot standard + rp endpoint create --name api --type load-balancer --image IMAGE \\ + --gpu 4090 --requests-per-worker 4 --env PORT=8000 + rp endpoint create --name embeddings --type queue --image IMAGE \\ + --cpu cpu3c --vcpu-count 8 + """ + from runpod.rp_cli import console as ui + + if not options.get("image") and not options.get("template"): + raise click.UsageError("pass --image or --template") + with api_errors(): + body = _config_body(options) + created = run_rest_request("POST", "/v2/serverless", json=body) or {} + if as_json: + print_json(created) + return + ui.success( + f"created [white]{_escape(created.get('name'))}[/white] [dim]{created.get('id')}[/dim]" + ) + print_fields(created, _FIELDS) + + +@endpoint.command(name="update") +@endpoint_argument() +@_config_options(creating=False) +@click.option( + "--wait", is_flag=True, help="Follow the rollout until the new release is live." +) +@click.option( + "--timeout", default=600.0, show_default=True, help="Seconds to follow the rollout." +) +@json_option +def update_endpoint(endpoint_id, wait, timeout, as_json, **options): + """Change an endpoint's configuration with flags. + + Only the given fields change. --env merges into the existing variables, + while list flags such as --data-center and --network-volume replace + their list. + """ + with api_errors(): + current = _get(endpoint_id) + patch = _config_body(options, current) + if not patch: + raise click.UsageError( + "nothing to update. pass an option (see --help) or use rp endpoint edit" + ) + _apply(endpoint_id, patch, wait, timeout, as_json) + + +@endpoint.command(name="edit") +@endpoint_argument(required=False) +@click.option( + "--file", + "document", + type=click.File("r"), + help="Apply a JSON document (- for stdin) instead of opening $EDITOR.", +) +@click.option("--yes", "-y", is_flag=True, help="Apply without confirming.") +@click.option("--no-follow", is_flag=True, help="Return without following the rollout.") +@click.option( + "--timeout", default=600.0, show_default=True, help="Seconds to follow the rollout." +) +@json_option +def edit_endpoint(endpoint_id, document, yes, no_follow, timeout, as_json): + """Edit an endpoint's configuration in $EDITOR. + + Opens the mutable configuration as JSON, validates it, shows the + patch, and after confirmation applies it and follows the rollout. + Ctrl-C while following detaches without undoing anything. Fields + removed from the document stay unchanged. + """ + from runpod.rp_cli import console as ui + + endpoint_id = _pick_endpoint(endpoint_id) + with api_errors(): + current = _get(endpoint_id) + before = _mutable(current) + if document: + after = _validate(document.read(), before) + elif not interactive(): + raise click.UsageError( + "edit opens $EDITOR and needs an interactive terminal. " + "pass --file or use rp endpoint update" + ) + else: + after = _edit_document(before) + patch = _patch(before, after) if after is not None else {} + if not patch: + if as_json: + print_json({}) + else: + ui.console.print("[dim]no changes[/dim]") + return + if not as_json: + ui.console.print("[dim]patch:[/dim]") + print_json(patch) + confirm(f"apply to endpoint {endpoint_id}?", yes) + _apply(endpoint_id, patch, not (no_follow or as_json), timeout, as_json) + + +@endpoint.command(name="scale") +@endpoint_argument() +@_scale_options +@json_option +def scale_endpoint(endpoint_id, as_json, **options): + """Change worker limits and the autoscaler.""" + from runpod.rp_cli import console as ui + + with api_errors(): + current = _get(endpoint_id) + patch = {} + workers = _workers_body(options, current) + if workers: + patch["workers"] = workers + scaling = _scaling_body(options, current.get("type")) + if scaling: + patch["scaling"] = scaling + if not patch: + raise click.UsageError("pass --min-workers, --max-workers, or a scaler option") + with api_errors(): + updated = run_rest_request("PATCH", _path(endpoint_id), json=patch) or {} + if as_json: + print_json(updated) + return + ui.success(f"scaled [white]{_escape(updated.get('name') or endpoint_id)}[/white]") + print_fields( + updated, + [ + ("id", lambda e: e.get("id") or endpoint_id), + ("workers", _workers), + ( + "idle timeout", + lambda e: _seconds((e.get("workers") or {}).get("idleTimeout")), + ), + ("scaling", _scaling), + ], + ) + + +@endpoint.command(name="delete") +@endpoint_argument() +@click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.") +def delete_endpoint(endpoint_id, yes): + """Delete an endpoint and its workers.""" + from runpod.rp_cli import console as ui + + with api_errors(): + name = _get(endpoint_id).get("name") or endpoint_id + confirm(f"delete endpoint '{name}' ({endpoint_id})?", yes) + with api_errors(): + run_rest_request("DELETE", _path(endpoint_id)) + ui.success(f"deleted [white]{_escape(name)}[/white] [dim]{endpoint_id}[/dim]") + + +@endpoint.command(name="workers") +@endpoint_argument() +@click.option("--watch", "-w", is_flag=True, help="Refresh in place until Ctrl-C.") +@click.option( + "--interval", + type=click.FloatRange(min=1), + default=2.0, + show_default=True, + help="Seconds between refreshes.", +) +@json_option +def list_workers(endpoint_id, watch, interval, as_json): + """Show an endpoint's workers. With --watch, --json prints one line per refresh.""" + from runpod.rp_cli import console as ui + + try: + with api_errors(): + data = _list_workers(endpoint_id) + if not watch: + if as_json: + print_json(data) + else: + ui.console.print(_workers_view(data)) + return + if as_json: + while True: + click.echo(json.dumps(data, default=str)) + time.sleep(interval) + with api_errors(): + data = _list_workers(endpoint_id) + from rich.live import Live + + with Live(_workers_view(data), console=ui.console, auto_refresh=False) as live: + while True: + time.sleep(interval) + with api_errors(): + data = _list_workers(endpoint_id) + live.update(_workers_view(data), refresh=True) + except KeyboardInterrupt: + pass + + +@endpoint.command(name="releases") +@endpoint_argument() +@json_option +def list_releases(endpoint_id, as_json): + """Show release history and rollout progress.""" + from runpod.rp_cli import console as ui + + with api_errors(): + state = _rollout_state(endpoint_id) + releases = compute.list_all(_path(endpoint_id, "releases"), "releases") + if as_json: + print_json( + { + "endpointVersion": state.get("endpointVersion"), + "rollout": state.get("rollout"), + "releases": releases, + } + ) + return + if state.get("endpointVersion") is not None: + ui.console.print(f"[dim]{_rollout_line(state)}[/dim]") + _print_table( + releases, + [ + ("RELEASE", lambda r: r["id"]), + ( + "VERSION", + lambda r: f"v{r['version']}" if r.get("version") is not None else None, + ), + ("SOURCE", lambda r: (r.get("source") or "").lower().replace("_", " ")), + ("BUILD", lambda r: r.get("buildId")), + ("WORKERS", lambda r: r.get("workerCount")), + ("CREATED", lambda r: r.get("createdAt")), + ( + "CHANGES", + lambda r: [entry.get("field") for entry in r.get("diff") or []], + ), + ], + "no releases", + ) + + +@endpoint.command(name="logs") +@endpoint_argument() +@click.option( + "--worker", + "worker_ids", + multiple=True, + shell_complete=complete_worker, + help="Only this worker. Repeatable. Omit for every worker.", +) +@click.option( + "--follow", "-f", is_flag=True, help="Keep streaming and attach to new workers." +) +@click.option( + "--source", + type=click.Choice(["container", "system"]), + help="Log source. Omit for both.", +) +@click.option( + "--tail", + type=click.IntRange(0, 5000), + default=100, + show_default=True, + help="Lines of backfill per worker.", +) +@click.option("--since", help="RFC3339 timestamp to start from (overrides --tail).") +@json_option +def logs(endpoint_id, worker_ids, follow, source, tail, since, as_json): + """Show worker logs, each line prefixed with its worker id. + + Without --worker, every worker's logs are multiplexed, and with + --follow newly created workers are attached automatically. + """ + from runpod.rp_cli import console as ui + + params: dict[str, Any] = {"tail": tail} + if source: + params["source"] = source + if since: + params["since"] = since + wanted = set(worker_ids) + started = datetime.now(timezone.utc) + events: queue.Queue = queue.Queue() + stop = threading.Event() + attached: set[str] = set() + cursors: dict[str, str] = {} + collected: list[tuple[str, dict]] = [] + width = max((len(w) for w in wanted), default=0) + + def attach_new() -> None: + nonlocal width + with api_errors(): + listed = [w["id"] for w in _list_workers(endpoint_id).get("workers") or []] + for worker_id in listed: + if (wanted and worker_id not in wanted) or worker_id in attached: + continue + attached.add(worker_id) + width = max(width, len(worker_id)) + threading.Thread( + target=_log_stream, + args=( + endpoint_id, + worker_id, + params, + cursors.get(worker_id), + events, + stop, + 90.0 if follow else 3.0, + ), + daemon=True, + ).start() + + try: + attach_new() + if not attached: + if not follow: + ui.console.print("[dim]no workers[/dim]") + return + ui.console.print("[dim]waiting for workers ...[/dim]") + next_poll = time.monotonic() + WORKER_POLL_SECONDS + while follow or attached: + try: + kind, worker_id, payload, cursor = events.get(timeout=0.5) + except queue.Empty: + kind = None + if kind == "line": + if cursor: + cursors[worker_id] = cursor + if follow: + click.echo(_log_line(worker_id, payload, width, as_json)) + elif _is_live(payload, started): + attached.discard(worker_id) + else: + collected.append((worker_id, payload)) + elif kind == "error": + if not ( + isinstance(payload, error.QueryError) and payload.status_code == 404 + ): + raise click.ClickException(str(payload)) + elif kind == "done": + attached.discard(worker_id) + if follow and time.monotonic() >= next_poll: + attach_new() + next_poll = time.monotonic() + WORKER_POLL_SECONDS + except KeyboardInterrupt: + pass + finally: + stop.set() + for worker_id, entry in sorted( + collected, key=lambda item: str(item[1].get("ts") or "") + ): + click.echo(_log_line(worker_id, entry, width, as_json)) + + +@endpoint.command(name="run") +@endpoint_argument() +@click.option( + "--input", + "input_file", + type=click.File("r"), + help="JSON request file, or - for stdin. A body without an 'input' key is wrapped in one.", +) +@click.option("--data", help="Inline JSON request.") +@click.option("--async", "run_async", is_flag=True, help="Submit and print the job id.") +@click.option("--stream", is_flag=True, help="Print streamed output as it arrives.") +@click.option("--webhook", help="URL notified when the job finishes.") +@click.option( + "--execution-timeout", + type=click.IntRange(min=1), + help="Job execution limit in seconds.", +) +@click.option("--ttl", type=click.IntRange(min=1), help="Job time to live in seconds.") +@click.option( + "--low-priority", is_flag=True, help="Do not scale up workers for this job." +) +@click.option("--timeout", type=float, help="Seconds to wait for the result.") +@click.option("--path", "lb_path", help="Load-balancer endpoints: request path.") +@click.option( + "--method", + help="Load-balancer endpoints: HTTP method (POST with a body, else GET).", +) +@json_option +def run_endpoint( + endpoint_id, + input_file, + data, + run_async, + stream, + webhook, + execution_timeout, + ttl, + low_priority, + timeout, + lb_path, + method, + as_json, +): # pylint: disable=too-many-arguments + """Send a request to an endpoint. + + Queue endpoints run a job on the serverless data plane and print its + output. Load-balancer endpoints receive the body as an HTTP request + to --path on the endpoint URL. + + \b + rp endpoint run ENDPOINT_ID --input request.json + rp endpoint run ENDPOINT_ID --data '{"prompt":"hello"}' + rp endpoint run ENDPOINT_ID --input - --stream + rp endpoint run ENDPOINT_ID --input request.json --async + rp endpoint run LB_ENDPOINT_ID --path /v1/models + """ + from runpod.rp_cli import console as ui + + if run_async and stream: + raise click.UsageError("use either --async or --stream") + body = _request_body(input_file, data) + with api_errors(): + record = _get(endpoint_id) + + if record.get("type") == "LOAD_BALANCER": + if run_async or stream or webhook or execution_timeout or ttl or low_priority: + raise click.UsageError( + "load-balancer endpoints take --path, --method, and a body; " + "job options apply to queue endpoints" + ) + _call_load_balancer(record, lb_path or "/", method, body, timeout, as_json) + return + + if lb_path or method: + raise click.UsageError("--path and --method apply to load-balancer endpoints") + if body is None: + raise click.UsageError("pass --input FILE or --data JSON") + request = body if isinstance(body, dict) and "input" in body else {"input": body} + if webhook: + request["webhook"] = webhook + policy = dict(request.get("policy") or {}) + if execution_timeout: + policy["executionTimeout"] = execution_timeout * 1000 + if ttl: + policy["ttl"] = ttl * 1000 + if low_priority: + policy["lowPriority"] = True + if policy: + request["policy"] = policy + + client = _client() + with _data_plane_errors(): + submitted = client.post(_job_path(endpoint_id, "run"), request) + job_id = submitted.get("id") + if run_async: + if as_json: + print_json(submitted) + return + ui.success( + f"submitted job [white]{job_id}[/white] [dim]{submitted.get('status')}[/dim]" + ) + ui.console.print(f"[dim]rp endpoint job wait {endpoint_id} {job_id}[/dim]") + return + if stream: + if not as_json: + click.echo(f"job {job_id}", err=True) + _stream_job(client, endpoint_id, job_id, as_json) + return + _print_job(_wait_job(client, endpoint_id, job_id, timeout), as_json) + + +@endpoint.group(name="job") +def job(): + """Inspect and control queue jobs.""" + + +def _job_command(name: str): + def decorate(func): + func = json_option(func) + func = click.argument("job_id")(func) + func = endpoint_argument()(func) + return job.command(name=name)(func) + + return decorate + + +def _job_output(record: dict) -> Optional[str]: + output = record.get("output") + if output is None or isinstance(output, str): + return output + return json.dumps(output) + + +_JOB_FIELDS = [ + ("id", lambda j: j.get("id")), + ("status", lambda j: (j.get("status") or "").lower()), + ("delay", lambda j: _seconds((j.get("delayTime") or 0) / 1000)), + ("execution", lambda j: _seconds((j.get("executionTime") or 0) / 1000)), + ("worker", lambda j: j.get("workerId")), + ("error", lambda j: j.get("error")), + ("output", _job_output), +] + + +@_job_command("get") +def job_get(endpoint_id, job_id, as_json): + """Show a job's status and output.""" + with _data_plane_errors(): + record = _client().get(_job_path(endpoint_id, "status", job_id)) + if as_json: + print_json(record) + return + print_fields(record, _JOB_FIELDS) + + +@_job_command("wait") +@click.option("--timeout", type=float, help="Seconds to wait.") +def job_wait(endpoint_id, job_id, timeout, as_json): + """Wait for a job to finish and print its output.""" + _print_job(_wait_job(_client(), endpoint_id, job_id, timeout), as_json) + + +@_job_command("stream") +def job_stream(endpoint_id, job_id, as_json): + """Print a job's streamed output as it arrives.""" + _stream_job(_client(), endpoint_id, job_id, as_json) + + +@_job_command("cancel") +def job_cancel(endpoint_id, job_id, as_json): + """Cancel a queued or running job.""" + from runpod.rp_cli import console as ui + + with _data_plane_errors(): + record = _client().post(_job_path(endpoint_id, "cancel", job_id), None) + if as_json: + print_json(record) + return + ui.success(f"job [white]{job_id}[/white] [dim]{record.get('status')}[/dim]") + + +@_job_command("retry") +def job_retry(endpoint_id, job_id, as_json): + """Requeue a failed or timed-out job.""" + from runpod.rp_cli import console as ui + + with _data_plane_errors(): + record = _client().post(_job_path(endpoint_id, "retry", job_id), None) + if as_json: + print_json(record) + return + ui.success( + f"job [white]{record.get('id') or job_id}[/white] [dim]{record.get('status')}[/dim]" + ) + + +@endpoint.command(name="purge") +@endpoint_argument() +@click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.") +@json_option +def purge_queue(endpoint_id, yes, as_json): + """Remove every queued job. Running jobs are not affected.""" + from runpod.rp_cli import console as ui + + confirm(f"remove every queued job from {endpoint_id}?", yes) + with _data_plane_errors(): + record = _client().post(_job_path(endpoint_id, "purge-queue"), None) + if as_json: + print_json(record) + return + ui.success( + f"purged {record.get('removed', 0)} queued jobs from [white]{endpoint_id}[/white]" + ) diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index d8be1833..ba67ef92 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -1039,8 +1039,9 @@ def _mount_groups() -> None: from runpod.cli.groups.pod.commands import pod_cli from runpod.cli.groups.ssh.commands import ssh_cli from runpod.rp_cli.catalog import catalog + from runpod.rp_cli.endpoint import endpoint - for command in (ssh_cli, pod_cli, catalog): + for command in (ssh_cli, pod_cli, catalog, endpoint): cli.add_command(command) diff --git a/tests/test_cli/test_endpoint_command.py b/tests/test_cli/test_endpoint_command.py new file mode 100644 index 00000000..12ecd6b8 --- /dev/null +++ b/tests/test_cli/test_endpoint_command.py @@ -0,0 +1,1126 @@ +"""tests for rp endpoint.""" + +import copy +import json +import threading +import queue +from unittest.mock import MagicMock + +import click +import pytest +import requests +from click.testing import CliRunner + +from runpod import error +from runpod.endpoint import runner +from runpod.rp_cli import compute, output +from runpod.rp_cli import endpoint as ep +from runpod.rp_cli.main import cli + +GPUS = [ + {"id": "NVIDIA H100 80GB HBM3", "name": "H100 SXM", "pool": "ADA_80_PRO"}, + {"id": "NVIDIA H100 NVL", "name": "H100 NVL", "pool": "ADA_80_PRO"}, + {"id": "NVIDIA A100 80GB PCIe", "name": "A100 PCIe", "pool": "ADA_80_PRO"}, + {"id": "NVIDIA GeForce RTX 4090", "name": "RTX 4090", "pool": "ADA_24"}, + {"id": "NVIDIA GeForce RTX 3070", "name": "RTX 3070", "pool": None}, +] +CPUS = [ + {"id": "cpu3c", "name": "Compute-Optimized", "group": "CPU3"}, + {"id": "cpu5g", "name": "General Purpose", "group": "CPU5"}, +] +DATACENTERS = [ + {"id": "EU-RO-1", "region": "EUROPE"}, + {"id": "EU-CZ-1", "region": "EUROPE"}, + {"id": "US-TX-3", "region": "NORTH_AMERICA"}, +] + +GPU_ENDPOINT = { + "id": "ep1", + "name": "llama-prod", + "type": "QUEUE", + "image": "ghcr.io/acme/llama:latest", + "args": "", + "env": {"A": "1", "B": "2"}, + "disk": 10, + "ports": [], + "gpu": { + "pools": ["ADA_80_PRO"], + "excludedTypes": ["NVIDIA A100 80GB PCIe"], + "count": 1, + "allowedCudaVersions": [], + "minCudaVersion": None, + }, + "workers": {"min": 0, "max": 3, "idleTimeout": 5}, + "scaling": {"type": "QUEUE_DELAY", "queueDelay": 4}, + "timeout": 300000, + "flashboot": "FLASHBOOT", + "dataCenterIds": ["EU-RO-1"], + "networkVolumes": ["vol1"], + "mounts": {"network": [{"volumeId": "vol1"}], "global": []}, + "requestUrls": {"run": "https://api.runpod.ai/v2/ep1/run"}, + "createdAt": "2026-07-10T22:38:43Z", +} +CPU_ENDPOINT = { + "id": "ep2", + "name": "embeddings", + "type": "QUEUE", + "image": "ghcr.io/acme/embeddings:latest", + "env": {}, + "cpu": [{"id": "cpu3c", "vcpuCount": 4, "memory": 8}], + "workers": {"min": 0, "max": 3}, + "scaling": {"type": "QUEUE_DELAY", "queueDelay": 2}, + "timeout": 0, + "flashboot": "OFF", + "dataCenterIds": [], + "mounts": {"network": [], "global": [{"volumeId": "gv1"}]}, + "createdAt": "2026-07-11T00:00:00Z", +} +LB_ENDPOINT = { + "id": "lb1", + "name": "openai-api", + "type": "LOAD_BALANCER", + "image": "ghcr.io/acme/openai-server:latest", + "gpu": { + "pools": ["ADA_24"], + "excludedTypes": [], + "count": 2, + "allowedCudaVersions": [], + "minCudaVersion": "12.4", + }, + "workers": {"min": 1, "max": 10}, + "scaling": {"type": "REQUEST_COUNT", "requestCount": 4}, + "timeout": 0, + "flashboot": "PRIORITY_FLASHBOOT", + "dataCenterIds": [], + "mounts": {"network": [], "global": []}, + "requestUrls": {"base": "https://lb1.api.runpod.ai", "health": "x"}, + "createdAt": "2026-07-12T00:00:00Z", +} +RELEASES = { + "endpointVersion": 3, + "rollout": { + "inProgress": False, + "workersOnLatest": 2, + "workersTotal": 2, + "percentOnLatest": 100, + }, + "releases": [ + { + "id": "rel3", + "version": 3, + "source": "MANUAL", + "workerCount": 2, + "createdAt": "2026-07-12T00:00:00Z", + "diff": [{"field": "template.imageName", "old": "a", "new": "b"}], + }, + { + "id": "rel2", + "version": None, + "source": "GIT_BUILD", + "buildId": "b1", + "workerCount": 0, + "createdAt": "2026-07-11T00:00:00Z", + "diff": [], + }, + ], + "pagination": {"hasNextPage": False, "nextCursor": None}, +} +WORKERS = { + "endpointVersion": 3, + "summary": { + "running": 1, + "idle": 1, + "initializing": 1, + "throttled": 0, + "unhealthy": 0, + "total": 3, + }, + "workers": [ + { + "id": "w1", + "status": "RUNNING", + "isStale": False, + "version": 3, + "gpuCount": 1, + "gpuTypeId": "NVIDIA H100 NVL", + "dataCenterId": "EU-RO-1", + "uptimeSeconds": 3725, + "image": "img", + }, + { + "id": "w2", + "status": "IDLE", + "isStale": True, + "version": 2, + "gpuCount": 1, + "uptimeSeconds": 125, + }, + {"id": "w3", "status": "INITIALIZING", "isStale": True, "gpuCount": 0}, + ], +} + + +class FakeApi: + """routes rest calls by (method, path) and records them.""" + + def __init__(self, routes): + self.routes = routes + self.calls = [] + + def __call__(self, method, path, params=None, json=None, **_): + self.calls.append((method, path, params, json)) + value = self.routes.get((method, path)) + if callable(value): + value = value(params, json) + if isinstance(value, Exception): + raise value + return copy.deepcopy(value) + + def bodies(self, method): + return [body for m, _, _, body in self.calls if m == method] + + +@pytest.fixture(autouse=True) +def wide_console(monkeypatch): + from runpod.rp_cli import console as ui + + monkeypatch.setattr(ui.console, "_width", 200) + + +@pytest.fixture +def api(monkeypatch): + routes = { + ("GET", "/v2/catalog/gpus"): {"gpus": GPUS}, + ("GET", "/v2/catalog/cpus"): {"cpus": CPUS}, + ("GET", "/v2/catalog/datacenters"): {"dataCenters": DATACENTERS}, + ("GET", "/v2/serverless"): { + "endpoints": [GPU_ENDPOINT, CPU_ENDPOINT, LB_ENDPOINT], + "pagination": {"hasNextPage": False, "nextCursor": None}, + }, + ("GET", "/v2/serverless/ep1"): GPU_ENDPOINT, + ("GET", "/v2/serverless/ep2"): CPU_ENDPOINT, + ("GET", "/v2/serverless/lb1"): LB_ENDPOINT, + ("GET", "/v2/serverless/ep1/workers"): WORKERS, + ("GET", "/v2/serverless/ep1/releases"): RELEASES, + ("POST", "/v2/serverless"): lambda params, body: dict(body, id="new1"), + ("PATCH", "/v2/serverless/ep1"): lambda params, body: dict( + GPU_ENDPOINT, **body + ), + ("PATCH", "/v2/serverless/ep2"): lambda params, body: dict( + CPU_ENDPOINT, **body + ), + ("PATCH", "/v2/serverless/lb1"): lambda params, body: dict(LB_ENDPOINT, **body), + ("DELETE", "/v2/serverless/ep1"): None, + } + fake = FakeApi(routes) + monkeypatch.setattr(ep, "run_rest_request", fake) + monkeypatch.setattr(compute, "run_rest_request", fake) + monkeypatch.setattr(ep.time, "sleep", lambda seconds: None) + return fake + + +@pytest.fixture +def client(monkeypatch): + fake = MagicMock() + monkeypatch.setattr(runner, "RunPodClient", lambda: fake) + return fake + + +def invoke(*args, input=None): + return CliRunner().invoke(cli, ["endpoint", *args], input=input) + + +def ok(*args, input=None): + result = invoke(*args, input=input) + assert result.exit_code == 0, result.output + return result.output + + +def fails(*args, input=None): + result = invoke(*args, input=input) + assert result.exit_code != 0, result.output + return result.output + + +# ---------------------------------------------------------------- list/get + + +def test_list_shows_ids_and_summaries(api): + out = ok("list") + for text in ("ep1", "ep2", "lb1", "llama-prod", "load-balancer", "ADA_80_PRO x1"): + assert text in out + assert "4s delay" in out and "4/worker" in out and "excluding" not in out + + +def test_list_filters_by_name_glob_and_type(api): + rows = json.loads(ok("list", "--name", "LLAMA*", "--json")) + assert [row["id"] for row in rows] == ["ep1"] + rows = json.loads(ok("list", "--type", "load-balancer", "--json")) + assert [row["id"] for row in rows] == ["lb1"] + assert "no endpoints match" in ok("list", "--name", "nope") + + +def test_get_shows_fields(api): + out = ok("get", "ep1") + assert "llama-prod" in out + assert "ADA_80_PRO x1 (excluding NVIDIA A100 80GB PCIe)" in out + assert "queue delay 4s" in out + assert "vol1 (network)" in out + assert "300s" in out + out = ok("get", "lb1") + assert ">=12.4" in out and "4 requests per worker" in out and "priority" in out + out = ok("get", "ep2") + assert "cpu3c 4 vcpu" in out and "gv1 (global)" in out + assert json.loads(ok("get", "ep1", "--json"))["id"] == "ep1" + + +def test_get_without_id_needs_a_terminal(api): + assert "pass an ENDPOINT_ID" in fails("get") + + +def test_get_without_id_picks_interactively(api, monkeypatch): + seen = {} + + def fake_pick(message, choices, missing): + seen["choices"] = choices + return "ep2" + + monkeypatch.setattr(ep, "interactive", lambda: True) + monkeypatch.setattr(ep, "pick", fake_pick) + assert "embeddings" in ok("get") + assert seen["choices"][0][0] == "ep1" + assert "llama-prod" in seen["choices"][0][1] + + +def test_api_errors_are_clean(api): + api.routes[("GET", "/v2/serverless/ep1")] = error.QueryError("endpoint not found") + assert "endpoint not found" in fails("get", "ep1") + + +# ------------------------------------------------------------------ create + + +def test_create_queue_endpoint_from_ticket(api): + ok( + "create", + "--name", "llama-prod", + "--type", "queue", + "--image", "ghcr.io/acme/llama:latest", + "--gpu", "H100", + "--gpu-count", "1", + "--min-workers", "0", + "--max-workers", "4", + "--queue-delay-seconds", "4", + "--idle-timeout-seconds", "10", + "--request-timeout-seconds", "300", + "--data-center", "EU-RO-1", + "--network-volume", "vol1", + "--env", "MODEL_ID=meta-llama/Llama-3.1-8B", + "--flashboot", "standard", + ) # fmt: skip + (body,) = api.bodies("POST") + assert body == { + "type": "QUEUE", + "name": "llama-prod", + "image": "ghcr.io/acme/llama:latest", + "flashboot": "FLASHBOOT", + "timeout": 300000, + "env": {"MODEL_ID": "meta-llama/Llama-3.1-8B"}, + "gpu": { + "pools": ["ADA_80_PRO"], + "excludedTypes": ["NVIDIA A100 80GB PCIe"], + "count": 1, + }, + "workers": {"min": 0, "max": 4, "idleTimeout": 10}, + "scaling": {"type": "QUEUE_DELAY", "queueDelay": 4.0}, + "dataCenterIds": ["EU-RO-1"], + "mounts": {"network": [{"volumeId": "vol1"}]}, + } + + +def test_create_load_balancer_endpoint_from_ticket(api): + out = ok( + "create", + "--name", "openai-api", + "--type", "load-balancer", + "--image", "ghcr.io/acme/openai-server:latest", + "--gpu", "h100", + "--gpu", "4090", + "--min-workers", "1", + "--max-workers", "10", + "--requests-per-worker", "4", + "--env", "PORT=8000", + "--env", "PORT_HEALTH=8000", + "--env", "HEALTH_CHECK_PATH=/health", + ) # fmt: skip + assert "new1" in out + (body,) = api.bodies("POST") + assert body["type"] == "LOAD_BALANCER" + assert body["scaling"] == {"type": "REQUEST_COUNT", "requestCount": 4} + assert body["gpu"]["pools"] == ["ADA_80_PRO", "ADA_24"] + assert body["env"] == { + "PORT": "8000", + "PORT_HEALTH": "8000", + "HEALTH_CHECK_PATH": "/health", + } + + +def test_create_cpu_endpoint_from_ticket(api): + ok( + "create", + "--name", "embeddings", + "--type", "queue", + "--image", "ghcr.io/acme/embeddings:latest", + "--cpu", "cpu3c", + "--vcpu-count", "8", + "--min-workers", "0", + "--max-workers", "3", + "--queue-delay-seconds", "2", + ) # fmt: skip + (body,) = api.bodies("POST") + assert body["cpu"] == [{"id": "cpu3c", "vcpuCount": 8}] + assert "gpu" not in body + + +def test_create_maps_container_and_placement_flags(api): + ok( + "create", + "--name", "x", + "--type", "queue", + "--template", "tpl1", + "--gpu", "h100:2", + "--cuda", ">=12.4", + "--region", "eu", + "--global-volume", "gv1", + "--port", "8000/http", + "--disk", "20", + "--args", "python main.py", + "--entrypoint", "python -m 'my app'", + "--cmd", "--port 1", + "--registry", "reg1", + "--json", + ) # fmt: skip + (body,) = api.bodies("POST") + assert body["templateId"] == "tpl1" + assert body["gpu"] == { + "pools": ["ADA_80_PRO"], + "excludedTypes": ["NVIDIA A100 80GB PCIe"], + "count": 2, + "minCudaVersion": "12.4", + } + assert body["scaling"] == {"type": "QUEUE_DELAY", "queueDelay": 4} + assert body["dataCenterIds"] == ["EU-RO-1", "EU-CZ-1"] + assert body["mounts"] == {"global": [{"volumeId": "gv1"}]} + assert body["ports"] == ["8000/http"] + assert body["entrypoint"] == ["python", "-m", "my app"] + assert body["cmd"] == ["--port", "1"] + assert body["registry"] == "reg1" and body["disk"] == 20 + + +def test_create_exact_cuda_and_cpu_vcpu_suffix(api): + ok("create", "--name", "x", "--type", "queue", "--image", "i", "--gpu", "4090", + "--cuda", "12.4,12.8") # fmt: skip + ok("create", "--name", "y", "--type", "queue", "--image", "i", "--cpu", "cpu5g:4") + gpu_body, cpu_body = api.bodies("POST") + assert gpu_body["gpu"] == { + "pools": ["ADA_24"], + "allowedCudaVersions": ["12.4", "12.8"], + } + assert cpu_body["cpu"] == [{"id": "cpu5g", "vcpuCount": 4}] + + +@pytest.mark.parametrize( + "args, message", + [ + (["--name", "x", "--type", "queue"], "pass --image or --template"), + (["--name", "x", "--type", "queue", "--image", "i", "--gpu", "h100", "--cpu", "cpu3c"], + "either --gpu or --cpu"), + (["--name", "x", "--type", "queue", "--image", "i", "--gpu", "h100:2", "--gpu-count", "1"], + "conflicting gpu counts"), + (["--name", "x", "--type", "load-balancer", "--image", "i"], "--requests-per-worker"), + (["--name", "x", "--type", "load-balancer", "--image", "i", "--queue-delay-seconds", "2"], + "scale on --requests-per-worker"), + (["--name", "x", "--type", "queue", "--image", "i", "--queue-delay-seconds", "2", + "--requests-per-worker", "2"], "either --queue-delay-seconds"), + (["--name", "x", "--type", "queue", "--image", "i", "--env", "NOVALUE"], "KEY=VALUE"), + (["--name", "x", "--type", "queue", "--image", "i", "--cpu", "cpu3c"], "--vcpu-count"), + (["--name", "x", "--type", "queue", "--image", "i", "--cpu", "cpu3c:x"], "invalid --cpu"), + (["--name", "x", "--type", "queue", "--image", "i", "--vcpu-count", "4"], "pass --cpu"), + (["--name", "x", "--type", "queue", "--image", "i", "--cpu", "cpu3c:4", "--cuda", "12.4"], + "--cuda applies to GPU"), + (["--name", "x", "--type", "queue", "--image", "i", "--gpu", "b200"], "unknown gpu"), + (["--name", "x", "--type", "queue", "--image", "i", "--region", "mars"], "unknown region"), + ], +) # fmt: skip +def test_create_rejects_bad_flags(api, args, message): + assert message in fails("create", *args) + assert not api.bodies("POST") + + +# ------------------------------------------------------------------ update + + +def test_update_merges_env_and_workers(api): + ok("update", "ep1", "--env", "C=3", "--unset-env", "A", "--max-workers", "8") + (patch,) = api.bodies("PATCH") + assert patch == { + "env": {"B": "2", "C": "3"}, + "workers": {"min": 0, "max": 8, "idleTimeout": 5}, + } + + +def test_update_gpu_replaces_pools_and_exclusions_together(api): + ok("update", "ep1", "--gpu", "4090", "--cuda", "any") + ok("update", "ep1", "--cuda", ">=12.8") + ok("update", "ep1", "--cuda", "12.8") + first, second, third = api.bodies("PATCH") + assert first["gpu"] == { + "pools": ["ADA_24"], + "excludedTypes": [], + "allowedCudaVersions": [], + "minCudaVersion": "", + } + assert second["gpu"] == {"minCudaVersion": "12.8", "allowedCudaVersions": []} + assert third["gpu"] == {"allowedCudaVersions": ["12.8"], "minCudaVersion": ""} + + +def test_update_cpu_vcpus_and_mounts(api): + ok("update", "ep2", "--vcpu-count", "8", "--network-volume", "vol9") + (patch,) = api.bodies("PATCH") + assert patch["cpu"] == [{"id": "cpu3c", "vcpuCount": 8}] + assert patch["mounts"] == { + "network": [{"volumeId": "vol9"}], + "global": [{"volumeId": "gv1"}], + } + + +@pytest.mark.parametrize( + "args, message", + [ + (["ep1"], "nothing to update"), + (["ep1", "--unset-env", "MISSING"], "no such variable"), + (["ep1", "--cpu", "cpu3c:4"], "cannot switch to CPUs"), + (["ep2", "--gpu", "h100"], "cannot switch to GPUs"), + (["ep2", "--cuda", "12.4"], "--cuda applies to GPU"), + (["lb1", "--queue-delay-seconds", "3"], "scale on --requests-per-worker"), + ], +) +def test_update_rejects_bad_flags(api, args, message): + assert message in fails("update", *args) + assert not api.bodies("PATCH") + + +def test_update_json_and_wait_follows_rollout(api): + states = iter( + [ + {"endpointVersion": 3}, + {"endpointVersion": 4, "rollout": {"inProgress": True, "workersOnLatest": 1, + "workersTotal": 2}}, + {"endpointVersion": 4, "rollout": {"inProgress": False, "workersOnLatest": 2, + "workersTotal": 2}}, + ] + ) # fmt: skip + api.routes[("GET", "/v2/serverless/ep1/releases")] = lambda params, body: next( + states + ) + out = ok("update", "ep1", "--image", "img:2", "--wait") + assert "v4 rolling out, 1/2 workers on latest" in out + assert "release v4 is live" in out + out = ok("update", "ep1", "--name", "renamed", "--json") + assert json.loads(out)["name"] == "renamed" + + +# -------------------------------------------------------------------- edit + + +def _doc(**changes): + doc = ep._mutable(GPU_ENDPOINT) + doc.update(changes) + return json.dumps(doc) + + +def test_mutable_document_holds_update_fields_only(): + doc = ep._mutable(GPU_ENDPOINT) + assert set(doc) <= set(ep.MUTABLE_FIELDS) + assert "networkVolumes" not in doc and "id" not in doc and "requestUrls" not in doc + assert ep._mutable(CPU_ENDPOINT)["cpu"] == [{"id": "cpu3c", "vcpuCount": 4}] + + +def test_edit_file_shows_patch_and_applies(api, tmp_path): + path = tmp_path / "doc.json" + gpu = dict(ep._mutable(GPU_ENDPOINT)["gpu"], excludedTypes=[], minCudaVersion=None) + path.write_text( + _doc(image="img:2", workers={"min": 1, "max": 3, "idleTimeout": 5}, gpu=gpu) + ) + out = ok("edit", "ep1", "--file", str(path), "--yes", "--no-follow") + (patch,) = api.bodies("PATCH") + assert patch == { + "image": "img:2", + "workers": {"min": 1, "max": 3, "idleTimeout": 5}, + "gpu": {"pools": ["ADA_80_PRO"], "excludedTypes": []}, + } + assert '"image": "img:2"' in out + assert "updated" in out + + +def test_edit_clears_cuda_floor_with_empty_string(): + before = ep._mutable(LB_ENDPOINT) + after = copy.deepcopy(before) + after["gpu"]["minCudaVersion"] = None + after["gpu"]["count"] = 1 + assert ep._patch(before, after) == {"gpu": {"count": 1, "minCudaVersion": ""}} + + +def test_edit_without_changes(api, tmp_path): + path = tmp_path / "doc.json" + path.write_text(_doc()) + assert "no changes" in ok("edit", "ep1", "--file", str(path)) + assert json.loads(ok("edit", "ep1", "--file", str(path), "--json")) == {} + + +def test_edit_requires_yes_without_a_terminal(api, tmp_path): + path = tmp_path / "doc.json" + path.write_text(_doc(name="renamed")) + assert "pass --yes" in fails("edit", "ep1", "--file", str(path)) + assert not api.bodies("PATCH") + + +def test_edit_needs_terminal_or_file(api): + assert "needs an interactive terminal" in fails("edit", "ep1") + + +@pytest.mark.parametrize( + "text, message", + [ + ("{not json", "invalid JSON"), + ("[]", "must be a JSON object"), + (json.dumps({"bogus": 1}), "unknown field 'bogus'"), + (json.dumps({"disk": "big"}), "'disk' must be an integer"), + (json.dumps({"disk": True}), "'disk' must be an integer"), + (json.dumps({"registry": 1}), "'registry' must be a string or null"), + (json.dumps({"workers": {"maximum": 1}}), "unknown field 'workers.maximum'"), + (json.dumps({"env": {"A": 1}}), "env values must be strings"), + (json.dumps({"flashboot": "ON"}), "'flashboot' must be one of"), + ( + json.dumps({"cpu": [{"id": "cpu3c", "vcpuCount": 2}]}), + "cannot switch to CPUs", + ), + ], +) +def test_edit_validation(text, message): + with pytest.raises(click.ClickException, match=message): + ep._validate(text, ep._mutable(GPU_ENDPOINT)) + with pytest.raises(click.ClickException, match="cannot switch to GPUs"): + ep._validate(json.dumps({"gpu": {}}), ep._mutable(CPU_ENDPOINT)) + + +def test_edit_in_editor_reopens_on_error_then_follows(api, monkeypatch): + monkeypatch.setattr(ep, "interactive", lambda: True) + monkeypatch.setattr(output, "interactive", lambda: True) + edits = iter(['{"bogus": 1}', _doc(name="renamed")]) + seen = [] + + def fake_edit(text, extension): + seen.append(text) + return next(edits) + + monkeypatch.setattr(ep.click, "edit", fake_edit) + states = iter([{"endpointVersion": 3}, {"endpointVersion": 4, "rollout": {}}]) + api.routes[("GET", "/v2/serverless/ep1/releases")] = lambda params, body: next( + states + ) + out = ok("edit", "ep1", input="y\ny\n") + assert "unknown field 'bogus'" in out + assert seen[1] == '{"bogus": 1}' + assert api.bodies("PATCH") == [{"name": "renamed"}] + assert "release v4 is live" in out + + +def test_edit_in_editor_gives_up_or_saves_nothing(api, monkeypatch): + monkeypatch.setattr(ep, "interactive", lambda: True) + monkeypatch.setattr(ep.click, "edit", lambda text, extension: "{oops") + assert "invalid JSON" in fails("edit", "ep1", input="n\n") + monkeypatch.setattr(ep.click, "edit", lambda text, extension: None) + assert "no changes" in ok("edit", "ep1") + + +def test_follow_rollout_without_new_release(api): + path_state = {"endpointVersion": 3, "rollout": {}} + api.routes[("GET", "/v2/serverless/ep1/releases")] = path_state + ep._follow_rollout("ep1", 3, timeout=0) + + +def test_follow_rollout_times_out_and_detaches(api, monkeypatch, capsys): + state = {"endpointVersion": 4, "rollout": {"inProgress": True}} + api.routes[("GET", "/v2/serverless/ep1/releases")] = state + ep._follow_rollout("ep1", 3, timeout=0) + assert "still in progress" in capsys.readouterr().out + + def interrupt(seconds): + raise KeyboardInterrupt + + monkeypatch.setattr(ep.time, "sleep", interrupt) + ep._follow_rollout("ep1", 3, timeout=60) + assert "detached" in capsys.readouterr().out + + +# ------------------------------------------------------------ scale/delete + + +def test_scale_sends_workers_and_scaler(api): + out = ok("scale", "ep1", "--min-workers", "1", "--requests-per-worker", "3") + (patch,) = api.bodies("PATCH") + assert patch == { + "workers": {"min": 1, "max": 3, "idleTimeout": 5}, + "scaling": {"type": "REQUEST_COUNT", "requestCount": 3}, + } + assert "1-3" in out and "3 requests per worker" in out + assert ( + json.loads(ok("scale", "ep1", "--max-workers", "9", "--json"))["workers"]["max"] + == 9 + ) + + +def test_scale_rejects_nothing_and_lb_queue_delay(api): + assert "pass --min-workers" in fails("scale", "ep1") + assert "requests-per-worker" in fails("scale", "lb1", "--queue-delay-seconds", "2") + + +def test_delete_confirms(api, monkeypatch): + assert "pass --yes" in fails("delete", "ep1") + monkeypatch.setattr(output, "interactive", lambda: True) + fails("delete", "ep1", input="n\n") + assert not [c for c in api.calls if c[0] == "DELETE"] + out = ok("delete", "ep1", input="y\n") + assert "deleted" in out and "llama-prod" in out + ok("delete", "ep1", "--yes") + assert len([c for c in api.calls if c[0] == "DELETE"]) == 2 + + +# ------------------------------------------------------- workers/releases + + +def test_workers_table_and_json(api): + out = ok("workers", "ep1") + assert "3 workers: 1 running, 1 idle, 1 initializing" in out + for text in ("w1", "w2", "w3", "1h02m", "2m05s", "v2 (stale)", "stale"): + assert text in out + assert json.loads(ok("workers", "ep1", "--json"))["summary"]["total"] == 3 + + +def test_workers_empty_and_uptime(): + assert ep._uptime(None) is None + assert ep._uptime(42) == "42s" + assert ep._worker_version({"version": None}) is None + assert "0 workers" in str(ep._workers_view({"summary": {}})) + + +def test_workers_watch(api, monkeypatch): + calls = {"n": 0} + + def tick(seconds): + calls["n"] += 1 + if calls["n"] > 1: + raise KeyboardInterrupt + + monkeypatch.setattr(ep.time, "sleep", tick) + out = ok("workers", "ep1", "--watch", "--json") + lines = [json.loads(line) for line in out.splitlines()] + assert len(lines) == 2 and lines[0]["endpointVersion"] == 3 + calls["n"] = 0 + out = ok("workers", "ep1", "--watch") + assert "w1" in out + + +def test_releases(api): + out = ok("releases", "ep1") + assert "v3 live, 2/2 workers on latest" in out + for text in ("rel3", "rel2", "manual", "git build", "template.imageName"): + assert text in out + data = json.loads(ok("releases", "ep1", "--json")) + assert data["endpointVersion"] == 3 and len(data["releases"]) == 2 + + +# -------------------------------------------------------------------- logs + + +class FakeStream: + def __init__(self, lines, status=200, text=""): + self.lines = lines + self.status_code = status + self.text = text + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def iter_lines(self, decode_unicode=True): + for line in self.lines: + if isinstance(line, Exception): + raise line + yield line + + +def _event(ts, line, source="container"): + return [ + f"id: {ts}", + "data: " + json.dumps({"source": source, "line": line, "ts": ts}), + "", + ] + + +@pytest.fixture +def streams(monkeypatch): + calls = [] + by_worker = {} + + def fake_get(url, params=None, headers=None, stream=None, timeout=None): + worker = url.split("/workers/")[1].split("/")[0] + calls.append({"worker": worker, "params": params, "headers": headers}) + return by_worker[worker] + + monkeypatch.setattr(ep.requests, "get", fake_get) + monkeypatch.setattr(ep, "_resolve_api_key", lambda key: "key") + by_worker["calls"] = calls + return by_worker + + +def test_logs_snapshot_multiplexes_sorted_and_stops_at_live_lines(api, streams): + streams["w1"] = FakeStream( + _event("2026-01-01T00:00:02Z", "w1 second") + + ["data: {broken", ": heartbeat"] + + _event("2999-01-01T00:00:00Z", "live line") + ) + streams["w2"] = FakeStream(_event("2026-01-01T00:00:01Z", "w2 first", "system")) + streams["w3"] = FakeStream([], status=404, text='{"title": "worker gone"}') + out = ok("logs", "ep1", "--source", "system", "--tail", "200") + lines = out.splitlines() + assert lines == ["w2 [system] w2 first", "w1 [container] w1 second"] + params = streams["calls"][0]["params"] + assert params == {"tail": 200, "source": "system"} + + +def test_logs_filters_workers_and_prints_json(api, streams): + streams["w2"] = FakeStream(_event("2026-01-01T00:00:01Z", "hello")) + out = ok( + "logs", "ep1", "--worker", "w2", "--since", "2026-01-01T00:00:00Z", "--json" + ) + assert json.loads(out) == { + "worker": "w2", + "source": "container", + "line": "hello", + "ts": "2026-01-01T00:00:01Z", + } + assert [call["worker"] for call in streams["calls"]] == ["w2"] + assert streams["calls"][0]["params"]["since"] == "2026-01-01T00:00:00Z" + + +def test_logs_without_workers(api): + api.routes[("GET", "/v2/serverless/ep1/workers")] = {"workers": []} + assert "no workers" in ok("logs", "ep1") + + +def test_logs_auth_error_fails(api, streams): + for worker in ("w1", "w2", "w3"): + streams[worker] = FakeStream([], status=401) + assert "Unauthorized" in fails("logs", "ep1") + + +def test_logs_follow_attaches_new_workers_and_resumes(api, streams, monkeypatch): + monkeypatch.setattr(ep, "WORKER_POLL_SECONDS", 0) + listings = iter( + [ + {"workers": [{"id": "w1"}]}, + {"workers": [{"id": "w1"}, {"id": "w2"}]}, + ] + ) + + def list_workers(params, body): + try: + return next(listings) + except StopIteration: + raise KeyboardInterrupt from None + + api.routes[("GET", "/v2/serverless/ep1/workers")] = list_workers + streams["w1"] = FakeStream(_event("2026-01-01T00:00:01Z", "one")) + streams["w2"] = FakeStream( + _event("2026-01-01T00:00:02Z", "two") + [requests.exceptions.ReadTimeout()] + ) + out = invoke("logs", "ep1", "--follow").output + assert "w1 [container] one" in out + + +def test_logs_follow_waits_for_workers(api, monkeypatch): + monkeypatch.setattr(ep, "WORKER_POLL_SECONDS", 0) + listings = iter([{"workers": []}]) + + def list_workers(params, body): + try: + return next(listings) + except StopIteration: + raise KeyboardInterrupt from None + + api.routes[("GET", "/v2/serverless/ep1/workers")] = list_workers + assert "waiting for workers" in ok("logs", "ep1", "-f") + + +def test_log_stream_sends_resume_cursor(monkeypatch): + seen = {} + + def fake_get(url, params=None, headers=None, stream=None, timeout=None): + seen["headers"] = headers + return FakeStream(_event("2026-01-01T00:00:01Z", "x")) + + monkeypatch.setattr(ep.requests, "get", fake_get) + monkeypatch.setattr(ep, "_resolve_api_key", lambda key: "key") + events = queue.Queue() + stop = threading.Event() + ep._log_stream("ep1", "w1", {}, "2026-01-01T00:00:00Z", events, stop, 1) + assert seen["headers"]["Last-Event-ID"] == "2026-01-01T00:00:00Z" + assert events.get()[0] == "line" and events.get()[0] == "done" + stop.set() + ep._log_stream("ep1", "w1", {}, None, events, stop, 1) + assert events.get()[0] == "done" + + +def test_is_live_handles_odd_timestamps(): + from datetime import datetime, timezone + + started = datetime(2026, 1, 1, tzinfo=timezone.utc) + assert not ep._is_live({"ts": "garbage"}, started) + assert ep._is_live({"ts": "2026-02-01T00:00:00"}, started) + + +# --------------------------------------------------------------------- run + + +def test_run_waits_for_output(api, client): + client.post.return_value = {"id": "job1", "status": "IN_QUEUE"} + client.get.side_effect = [ + {"id": "job1", "status": "IN_PROGRESS"}, + {"id": "job1", "status": "COMPLETED", "output": {"text": "hi"}}, + ] + out = ok("run", "ep1", "--data", '{"prompt": "hello"}', "--webhook", "https://hook", + "--execution-timeout", "60", "--ttl", "120", "--low-priority") # fmt: skip + assert json.loads(out) == {"text": "hi"} + path, body = client.post.call_args.args + assert path == "ep1/run" + assert body == { + "input": {"prompt": "hello"}, + "webhook": "https://hook", + "policy": {"executionTimeout": 60000, "ttl": 120000, "lowPriority": True}, + } + client.get.assert_called_with("ep1/status/job1") + + +def test_run_reads_input_file_and_stdin(api, client, tmp_path): + client.post.return_value = {"id": "job1"} + client.get.return_value = {"id": "job1", "status": "COMPLETED", "output": "plain"} + path = tmp_path / "request.json" + path.write_text('{"input": {"a": 1}, "policy": {"ttl": 5}}') + assert ok("run", "ep1", "--input", str(path)).strip() == "plain" + assert client.post.call_args.args[1] == {"input": {"a": 1}, "policy": {"ttl": 5}} + full = json.loads(ok("run", "ep1", "--input", "-", "--json", input="[1, 2]")) + assert full["output"] == "plain" + assert client.post.call_args.args[1] == {"input": [1, 2]} + + +def test_run_async(api, client): + client.post.return_value = {"id": "job1", "status": "IN_QUEUE"} + out = ok("run", "ep1", "--data", "{}", "--async") + assert "job1" in out and "rp endpoint job wait ep1 job1" in out + assert ( + json.loads(ok("run", "ep1", "--data", "{}", "--async", "--json"))["id"] + == "job1" + ) + client.get.assert_not_called() + + +def test_run_failed_job_exits_nonzero(api, client): + client.post.return_value = {"id": "job1"} + client.get.return_value = {"id": "job1", "status": "FAILED", "error": "boom"} + assert "job job1 failed: boom" in fails("run", "ep1", "--data", "{}") + + +def test_run_timeout_and_interrupt(api, client, monkeypatch): + client.post.return_value = {"id": "job1"} + client.get.return_value = {"id": "job1", "status": "IN_QUEUE"} + clock = iter(range(0, 1000, 10)) + monkeypatch.setattr(ep.time, "monotonic", lambda: next(clock)) + assert "resume with: rp endpoint job wait ep1 job1" in fails( + "run", "ep1", "--data", "{}", "--timeout", "5" + ) + client.get.side_effect = KeyboardInterrupt + result = invoke("run", "ep1", "--data", "{}") + assert result.exit_code != 0 + + +def test_run_stream(api, client): + client.post.return_value = {"id": "job1"} + client.get.side_effect = [ + {"status": "IN_PROGRESS", "stream": []}, + {"status": "IN_PROGRESS", "stream": [{"output": "hel"}, {"output": "lo"}]}, + {"status": "COMPLETED", "stream": [{"output": {"done": True}}]}, + {"status": "COMPLETED", "stream": []}, + ] + out = ok("run", "ep1", "--input", "-", "--stream", input='{"p": 1}') + assert 'hello\n{"done": true}' in out + + +def test_run_stream_failure_and_json(api, client): + client.post.return_value = {"id": "job1"} + client.get.side_effect = [ + {"status": "FAILED", "stream": [{"output": "partial"}]}, + {"status": "FAILED", "stream": []}, + {"id": "job1", "status": "FAILED", "error": "oom"}, + ] + assert "failed: oom" in fails("run", "ep1", "--data", "{}", "--stream", "--json") + + +@pytest.mark.parametrize( + "args, message", + [ + (["ep1", "--data", "{}", "--async", "--stream"], "either --async or --stream"), + (["ep1"], "pass --input FILE or --data JSON"), + (["ep1", "--data", "{bad"], "not valid JSON"), + (["ep1", "--data", "{}", "--input", "-"], "either --input or --data"), + (["ep1", "--data", "{}", "--path", "/x"], "apply to load-balancer"), + (["lb1", "--data", "{}", "--async"], "job options apply to queue"), + ], +) +def test_run_rejects_bad_flags(api, client, args, message): + assert message in fails("run", *args) + client.post.assert_not_called() + + +def test_run_load_balancer(api, monkeypatch): + calls = [] + + class Response: + def __init__(self, status, payload=None, text=""): + self.status_code = status + self._payload = payload + self.text = text + + def json(self): + if self._payload is None: + raise ValueError + return self._payload + + responses = iter( + [ + Response(200, {"data": [1]}), + Response(200, text="pong"), + Response(200, "a string"), + Response(500, text="kaput"), + ] + ) + + def fake_request(method, url, headers=None, json=None, timeout=None): + calls.append((method, url, json, timeout)) + return next(responses) + + monkeypatch.setattr(ep.requests, "request", fake_request) + monkeypatch.setattr(ep, "_resolve_api_key", lambda key: "key") + assert json.loads(ok("run", "lb1", "--path", "/v1/chat", "--data", '{"m": 1}')) == { + "data": [1] + } + assert ok("run", "lb1", "--path", "ping").strip() == "pong" + assert ok("run", "lb1", "--method", "put", "--timeout", "5").strip() == "a string" + assert "returned 500: kaput" in fails("run", "lb1") + assert calls[0] == ("POST", "https://lb1.api.runpod.ai/v1/chat", {"m": 1}, 600) + assert calls[1][:2] == ("GET", "https://lb1.api.runpod.ai/ping") + assert calls[2][0] == "PUT" and calls[2][3] == 5.0 + + +# -------------------------------------------------------------------- jobs + + +def test_job_get(api, client): + client.get.return_value = { + "id": "job1", + "status": "COMPLETED", + "delayTime": 1500, + "executionTime": 2000, + "workerId": "w1", + "output": {"a": 1}, + } + out = ok("job", "get", "ep1", "job1") + assert "completed" in out and "1.5s" in out and '{"a": 1}' in out + client.get.assert_called_with("ep1/status/job1") + assert json.loads(ok("job", "get", "ep1", "job1", "--json"))["workerId"] == "w1" + + +def test_job_wait_and_stream(api, client): + client.get.return_value = {"id": "job1", "status": "COMPLETED", "output": "done"} + assert ok("job", "wait", "ep1", "job1", "--timeout", "30").strip() == "done" + client.get.return_value = {"status": "COMPLETED", "stream": []} + ok("job", "stream", "ep1", "job1") + client.get.assert_called_with("ep1/stream/job1") + + +def test_job_stream_interrupt(api, client): + client.get.side_effect = KeyboardInterrupt + result = invoke("job", "stream", "ep1", "job1") + assert result.exit_code != 0 + + +def test_job_cancel_and_retry(api, client): + client.post.return_value = {"id": "job1", "status": "CANCELLED"} + assert "CANCELLED" in ok("job", "cancel", "ep1", "job1") + client.post.assert_called_with("ep1/cancel/job1", None) + assert ( + json.loads(ok("job", "cancel", "ep1", "job1", "--json"))["status"] + == "CANCELLED" + ) + client.post.return_value = {"id": "job1", "status": "IN_QUEUE"} + assert "IN_QUEUE" in ok("job", "retry", "ep1", "job1") + client.post.assert_called_with("ep1/retry/job1", None) + assert json.loads(ok("job", "retry", "ep1", "job1", "--json"))["id"] == "job1" + + +def test_purge(api, client): + client.post.return_value = {"removed": 3, "status": "completed"} + assert "pass --yes" in fails("purge", "ep1") + assert "purged 3 queued jobs" in ok("purge", "ep1", "--yes") + client.post.assert_called_with("ep1/purge-queue", None) + assert json.loads(ok("purge", "ep1", "--yes", "--json"))["removed"] == 3 + + +def test_data_plane_errors(api, client, monkeypatch): + response = requests.Response() + response.status_code = 404 + response._content = b"job not found" + client.get.side_effect = requests.HTTPError(response=response) + assert "404 job not found" in fails("job", "get", "ep1", "job1") + client.get.side_effect = requests.HTTPError("no response") + assert "no response" in fails("job", "get", "ep1", "job1") + + def no_key(): + raise RuntimeError("api key not set") + + monkeypatch.setattr(runner, "RunPodClient", no_key) + assert "api key not set" in fails("job", "get", "ep1", "job1") + + +# -------------------------------------------------------------- completion + + +def test_completion_inserts_ids(api): + ctx = click.Context(ep.list_workers, info_name="workers") + items = ep.complete_endpoint(ctx, None, "ep") + assert [item.value for item in items] == ["ep1", "ep2"] + assert items[0].help == "llama-prod" + ctx.params["endpoint_id"] = "ep1" + assert [item.value for item in ep.complete_worker(ctx, None, "w")] == [ + "w1", + "w2", + "w3", + ] + assert ep.complete_worker(click.Context(ep.list_workers), None, "") == [] + assert [item.value for item in ep.complete_datacenter(ctx, None, "EU")] == [ + "EU-RO-1", + "EU-CZ-1", + ] From e7a119d94db6458a5e756e391c6aec9c3dd0e4fc Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:12:45 -0700 Subject: [PATCH 2/7] fix: print final output for non-streaming jobs and tighten endpoint tables --- runpod/rp_cli/endpoint.py | 42 ++++++++++++++++--------- tests/test_cli/test_endpoint_command.py | 37 +++++++++++++++++++--- 2 files changed, 60 insertions(+), 19 deletions(-) diff --git a/runpod/rp_cli/endpoint.py b/runpod/rp_cli/endpoint.py index b8a69709..3e695572 100644 --- a/runpod/rp_cli/endpoint.py +++ b/runpod/rp_cli/endpoint.py @@ -45,6 +45,7 @@ FLASHBOOT = {"off": "OFF", "standard": "FLASHBOOT", "priority": "PRIORITY_FLASHBOOT"} DEFAULT_QUEUE_DELAY = 4 WORKER_POLL_SECONDS = 5.0 +STREAM_POLL_SECONDS = 60 MUTABLE_FIELDS = { "name": (str,), @@ -224,7 +225,6 @@ def _url(endpoint: dict) -> Optional[str]: ("COMPUTE", lambda e: _compute(e, compact=True)), ("WORKERS", _workers), ("SCALING", lambda e: _scaling(e, compact=True)), - ("FLASHBOOT", _flashboot), ("CREATED", lambda e: (e.get("createdAt") or "")[:10]), ] @@ -879,15 +879,12 @@ def _worker_version(worker: dict) -> Optional[str]: ( "GPU", lambda w: ( - f"{w['gpuTypeId']} x{w.get('gpuCount') or 1}" - if w.get("gpuTypeId") - else None + f"{w.get('gpuTypeId')} x{w['gpuCount']}" if w.get("gpuCount") else None ), ), ("DATA CENTER", lambda w: w.get("dataCenterId")), ("VERSION", _worker_version), ("UPTIME", lambda w: _uptime(w.get("uptimeSeconds"))), - ("IMAGE", lambda w: w.get("image")), ] @@ -1057,12 +1054,21 @@ def _print_job(job: dict, as_json: bool) -> None: def _stream_job(client, endpoint_id: str, job_id: str, as_json: bool) -> None: + """print streamed chunks, or the final output when the handler streams nothing.""" open_line = False + streamed = False try: while True: with _data_plane_errors(): - partial = client.get(_job_path(endpoint_id, "stream", job_id)) + try: + partial = client.get( + _job_path(endpoint_id, "stream", job_id), + timeout=STREAM_POLL_SECONDS, + ) + except requests.exceptions.ReadTimeout: + continue chunks = partial.get("stream") or [] + streamed = streamed or bool(chunks) for chunk in chunks: output = chunk.get("output") if as_json: @@ -1088,9 +1094,11 @@ def _stream_job(client, endpoint_id: str, job_id: str, as_json: bool) -> None: raise click.Abort() from None if open_line: click.echo() - if partial.get("status") != "COMPLETED": - with _data_plane_errors(): - raise _job_failure(client.get(_job_path(endpoint_id, "status", job_id))) + if streamed and partial.get("status") == "COMPLETED": + return + with _data_plane_errors(): + final = client.get(_job_path(endpoint_id, "status", job_id)) + _print_job(final, as_json) def _request_body(input_file, data: Optional[str]) -> Optional[Any]: @@ -1405,6 +1413,15 @@ def list_workers(endpoint_id, watch, interval, as_json): pass +def _release_changes(release: dict) -> list[str]: + """changed field names, skipping entries whose old and new values are both empty.""" + return [ + entry["field"].removeprefix("template.") + for entry in release.get("diff") or [] + if entry.get("field") and (entry.get("old") or entry.get("new")) + ] + + @endpoint.command(name="releases") @endpoint_argument() @json_option @@ -1437,11 +1454,8 @@ def list_releases(endpoint_id, as_json): ("SOURCE", lambda r: (r.get("source") or "").lower().replace("_", " ")), ("BUILD", lambda r: r.get("buildId")), ("WORKERS", lambda r: r.get("workerCount")), - ("CREATED", lambda r: r.get("createdAt")), - ( - "CHANGES", - lambda r: [entry.get("field") for entry in r.get("diff") or []], - ), + ("CREATED", lambda r: (r.get("createdAt") or "")[:16].replace("T", " ")), + ("CHANGES", _release_changes), ], "no releases", ) diff --git a/tests/test_cli/test_endpoint_command.py b/tests/test_cli/test_endpoint_command.py index 12ecd6b8..9068fbea 100644 --- a/tests/test_cli/test_endpoint_command.py +++ b/tests/test_cli/test_endpoint_command.py @@ -111,7 +111,10 @@ "source": "MANUAL", "workerCount": 2, "createdAt": "2026-07-12T00:00:00Z", - "diff": [{"field": "template.imageName", "old": "a", "new": "b"}], + "diff": [ + {"field": "template.imageName", "old": "a", "new": "b"}, + {"field": "template.ports", "old": None, "new": ""}, + ], }, { "id": "rel2", @@ -735,8 +738,16 @@ def tick(seconds): def test_releases(api): out = ok("releases", "ep1") assert "v3 live, 2/2 workers on latest" in out - for text in ("rel3", "rel2", "manual", "git build", "template.imageName"): + for text in ( + "rel3", + "rel2", + "manual", + "git build", + "imageName", + "2026-07-12 00:00", + ): assert text in out + assert "template." not in out and "ports" not in out data = json.loads(ok("releases", "ep1", "--json")) assert data["endpointVersion"] == 3 and len(data["releases"]) == 2 @@ -969,6 +980,18 @@ def test_run_stream(api, client): assert 'hello\n{"done": true}' in out +def test_run_stream_prints_output_of_non_streaming_handler(api, client): + client.post.return_value = {"id": "job1"} + client.get.side_effect = [ + requests.exceptions.ReadTimeout(), + {"status": "IN_QUEUE", "stream": []}, + {"status": "COMPLETED", "stream": []}, + {"id": "job1", "status": "COMPLETED", "output": "streamed"}, + ] + assert ok("run", "ep1", "--data", "{}", "--stream").splitlines()[-1] == "streamed" + client.get.assert_called_with("ep1/status/job1") + + def test_run_stream_failure_and_json(api, client): client.post.return_value = {"id": "job1"} client.get.side_effect = [ @@ -1056,9 +1079,13 @@ def test_job_get(api, client): def test_job_wait_and_stream(api, client): client.get.return_value = {"id": "job1", "status": "COMPLETED", "output": "done"} assert ok("job", "wait", "ep1", "job1", "--timeout", "30").strip() == "done" - client.get.return_value = {"status": "COMPLETED", "stream": []} - ok("job", "stream", "ep1", "job1") - client.get.assert_called_with("ep1/stream/job1") + client.get.return_value = None + client.get.side_effect = [ + {"status": "COMPLETED", "stream": [{"output": "chunk"}]}, + {"status": "COMPLETED", "stream": []}, + ] + assert ok("job", "stream", "ep1", "job1").strip() == "chunk" + client.get.assert_called_with("ep1/stream/job1", timeout=ep.STREAM_POLL_SECONDS) def test_job_stream_interrupt(api, client): From f8bf6309b1a2353f9b8ca6b79578d8dbe89ea038 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:15:36 -0700 Subject: [PATCH 3/7] refactor: use shared print_table for endpoint tables --- runpod/rp_cli/endpoint.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/runpod/rp_cli/endpoint.py b/runpod/rp_cli/endpoint.py index 3e695572..d1729fe0 100644 --- a/runpod/rp_cli/endpoint.py +++ b/runpod/rp_cli/endpoint.py @@ -39,6 +39,7 @@ pick, print_fields, print_json, + print_table, ) TYPES = {"queue": "QUEUE", "load-balancer": "LOAD_BALANCER"} @@ -259,24 +260,18 @@ def _escape(text: Any) -> str: def _table(rows: list[dict], columns: list): - """an aligned table whose first column, the id, never wraps or truncates.""" + """the table print_table renders, for views that embed it.""" from rich.markup import escape from rich.table import Table table = Table(box=None, header_style="dim", pad_edge=False) for index, (header, _) in enumerate(columns): - table.add_column(header, no_wrap=index == 0) + table.add_column(header, no_wrap=index == 0, overflow="fold") for row in rows: table.add_row(*(escape(cell(value(row))) for _, value in columns)) return table -def _print_table(rows: list[dict], columns: list, empty: str) -> None: - from runpod.rp_cli import console as ui - - ui.console.print(_table(rows, columns) if rows else f"[dim]{empty}[/dim]") - - # ------------------------------------------------------------ request body @@ -1188,7 +1183,7 @@ def list_endpoints(name, endpoint_type, as_json): if as_json: print_json(endpoints) return - _print_table(endpoints, _LIST_COLUMNS, "no endpoints match") + print_table(endpoints, _LIST_COLUMNS, "no endpoints match") @endpoint.command(name="get") @@ -1443,7 +1438,7 @@ def list_releases(endpoint_id, as_json): return if state.get("endpointVersion") is not None: ui.console.print(f"[dim]{_rollout_line(state)}[/dim]") - _print_table( + print_table( releases, [ ("RELEASE", lambda r: r["id"]), From 250acc9389895ba6ab96b2b0b6b820f5df2abfeb Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:54:27 -0700 Subject: [PATCH 4/7] refactor: consolidate endpoint configuration and job handling --- runpod/rp_cli/endpoint.py | 912 ++++-------------------- runpod/rp_cli/endpoint_config.py | 314 ++++++++ tests/test_cli/test_endpoint_command.py | 273 ++----- 3 files changed, 545 insertions(+), 954 deletions(-) create mode 100644 runpod/rp_cli/endpoint_config.py diff --git a/runpod/rp_cli/endpoint.py b/runpod/rp_cli/endpoint.py index d1729fe0..ae42274d 100644 --- a/runpod/rp_cli/endpoint.py +++ b/runpod/rp_cli/endpoint.py @@ -8,30 +8,22 @@ import contextlib import fnmatch import json -import queue -import shlex -import threading import time -from datetime import datetime, timezone from typing import Any, Iterator, Optional from urllib.parse import quote import click import requests -from runpod import error -from runpod.api.rest import ( - _build_headers, - _build_url, - _raise_for_error, - _resolve_api_key, - run_rest_request, -) +from rich.markup import escape + +from runpod.api.rest import run_rest_request +from runpod.apps.utils.network import api_key +from runpod.endpoint import runner from runpod.endpoint.helpers import is_completed -from runpod.rp_cli import compute +from runpod.rp_cli import compute, console as ui, endpoint_config as config +from runpod.rp_cli.api import api_errors, list_all from runpod.rp_cli.output import ( - api_errors, - cell, completer, confirm, interactive, @@ -40,37 +32,10 @@ print_fields, print_json, print_table, + table, ) TYPES = {"queue": "QUEUE", "load-balancer": "LOAD_BALANCER"} -FLASHBOOT = {"off": "OFF", "standard": "FLASHBOOT", "priority": "PRIORITY_FLASHBOOT"} -DEFAULT_QUEUE_DELAY = 4 -WORKER_POLL_SECONDS = 5.0 -STREAM_POLL_SECONDS = 60 - -MUTABLE_FIELDS = { - "name": (str,), - "image": (str,), - "args": (str,), - "env": (dict,), - "disk": (int,), - "ports": (list,), - "registry": (str, type(None)), - "gpu": (dict,), - "cpu": (list,), - "workers": (dict,), - "scaling": (dict,), - "timeout": (int,), - "flashboot": (str,), - "dataCenterIds": (list,), - "mounts": (dict,), -} -NESTED_FIELDS = { - "gpu": {"pools", "excludedTypes", "count", "allowedCudaVersions", "minCudaVersion"}, - "workers": {"min", "max", "idleTimeout"}, - "scaling": {"type", "queueDelay", "requestCount"}, - "mounts": {"network", "global"}, -} def _path(endpoint_id: str, *parts: str) -> str: @@ -78,23 +43,11 @@ def _path(endpoint_id: str, *parts: str) -> str: return "/v2/serverless/" + "/".join(segments) -def _get(endpoint_id: str) -> dict: - return run_rest_request("GET", _path(endpoint_id)) or {} - - -def _list_endpoints() -> list[dict]: - return compute.list_all("/v2/serverless", "endpoints") - - -def _list_workers(endpoint_id: str) -> dict: - return run_rest_request("GET", _path(endpoint_id, "workers")) or {} - - def _choices(path: str, collection: str): def fetch(): return [ (record["id"], record.get("name") or "") - for record in compute.list_all(path, collection) + for record in list_all(path, collection) ] return completer(fetch) @@ -121,7 +74,9 @@ def complete_worker(ctx, param, incomplete): return completer( lambda: [ (worker["id"], (worker.get("status") or "").lower()) - for worker in _list_workers(endpoint_id).get("workers") or [] + for worker in run_rest_request("GET", _path(endpoint_id, "workers"))[ + "workers" + ] ] )(ctx, param, incomplete) @@ -138,7 +93,7 @@ def _pick_endpoint(endpoint_id: Optional[str]) -> str: if not interactive(): raise click.UsageError("pass an ENDPOINT_ID") with api_errors(): - endpoints = _list_endpoints() + endpoints = list_all("/v2/serverless", "endpoints") return pick( "endpoint", [ @@ -149,16 +104,13 @@ def _pick_endpoint(endpoint_id: Optional[str]) -> str: ) -# ----------------------------------------------------------------- display - - def _type(endpoint: dict) -> Optional[str]: names = {value: key for key, value in TYPES.items()} return names.get(endpoint.get("type"), endpoint.get("type")) def _flashboot(endpoint: dict) -> Optional[str]: - names = {value: key for key, value in FLASHBOOT.items()} + names = {value: key for key, value in config.FLASHBOOT.items()} return names.get(endpoint.get("flashboot"), endpoint.get("flashboot")) @@ -253,256 +205,6 @@ def _url(endpoint: dict) -> Optional[str]: ] -def _escape(text: Any) -> str: - from rich.markup import escape - - return escape(str(text)) - - -def _table(rows: list[dict], columns: list): - """the table print_table renders, for views that embed it.""" - from rich.markup import escape - from rich.table import Table - - table = Table(box=None, header_style="dim", pad_edge=False) - for index, (header, _) in enumerate(columns): - table.add_column(header, no_wrap=index == 0, overflow="fold") - for row in rows: - table.add_row(*(escape(cell(value(row))) for _, value in columns)) - return table - - -# ------------------------------------------------------------ request body - - -def _env_pairs(values: tuple[str, ...]) -> dict[str, str]: - env = {} - for value in values: - key, sep, val = value.partition("=") - if not sep or not key: - raise click.UsageError(f"invalid --env '{value}'. use KEY=VALUE") - env[key] = val - return env - - -def _gpu_body(specs, count, cuda, updating: bool) -> dict: - gpu: dict[str, Any] = {} - if specs: - catalog = [ - entry - for entry in compute.list_gpus(product="SERVERLESS") - if entry.get("pool") - ] - matches: list[dict] = [] - counts = {count} if count else set() - for spec in specs: - name, spec_count = compute.parse_gpu_spec(spec) - if spec_count: - counts.add(spec_count) - for match in compute.resolve_gpus(name, catalog): - if match not in matches: - matches.append(match) - if len(counts) > 1: - raise click.UsageError("conflicting gpu counts; use one count per worker") - pools, excluded = compute.endpoint_gpu_selection(matches, catalog) - gpu["pools"] = pools - if excluded or updating: - gpu["excludedTypes"] = excluded - count = counts.pop() if counts else None - if count: - gpu["count"] = count - if cuda: - if [value.lower() for value in cuda] == ["any"]: - gpu["allowedCudaVersions"] = [] - if updating: - gpu["minCudaVersion"] = "" - return gpu - minimum, exact = compute.parse_cuda(cuda) - if minimum: - gpu["minCudaVersion"] = minimum - if updating: - gpu["allowedCudaVersions"] = [] - else: - gpu["allowedCudaVersions"] = exact - if updating: - gpu["minCudaVersion"] = "" - return gpu - - -def _cpu_body(specs, vcpu_count, current: Optional[dict]) -> Optional[list]: - if not specs: - if vcpu_count and current and current.get("cpu"): - return [ - {"id": cpu["id"], "vcpuCount": vcpu_count} for cpu in current["cpu"] - ] - if vcpu_count: - raise click.UsageError("--vcpu-count applies to CPU endpoints (pass --cpu)") - return None - flavors = compute.list_cpus(product="SERVERLESS") - entries = [] - for spec in specs: - name, sep, count = spec.rpartition(":") - if not sep: - name, count = spec, "" - if count and not count.isdigit(): - raise click.UsageError( - f"invalid --cpu '{spec}'. use a flavor like cpu3c or cpu3c:8" - ) - vcpus = int(count) if count else vcpu_count - if not vcpus: - raise click.UsageError(f"pass --vcpu-count or --cpu {name}:VCPUS") - entries.append( - {"id": compute.resolve_cpu(name, flavors)["id"], "vcpuCount": vcpus} - ) - return entries - - -def _workers_body(options: dict, current: Optional[dict]) -> Optional[dict]: - changes = { - key: options[name] - for key, name in ( - ("min", "min_workers"), - ("max", "max_workers"), - ("idleTimeout", "idle_timeout_seconds"), - ) - if options.get(name) is not None - } - if not changes: - return None - existing = (current or {}).get("workers") or {} - merged = { - key: existing[key] - for key in ("min", "max", "idleTimeout") - if existing.get(key) is not None - } - merged.update(changes) - return merged - - -def _scaling_body(options: dict, endpoint_type: Optional[str]) -> Optional[dict]: - delay = options.get("queue_delay_seconds") - per_worker = options.get("requests_per_worker") - if delay is not None and per_worker is not None: - raise click.UsageError( - "use either --queue-delay-seconds or --requests-per-worker" - ) - if delay is not None: - if endpoint_type == "LOAD_BALANCER": - raise click.UsageError( - "load-balancer endpoints scale on --requests-per-worker" - ) - return {"type": "QUEUE_DELAY", "queueDelay": delay} - if per_worker is not None: - return {"type": "REQUEST_COUNT", "requestCount": per_worker} - return None - - -def _mounts_body(options: dict, current: Optional[dict]) -> Optional[dict]: - network = options.get("network_volume") or () - global_ = options.get("global_volume") or () - if not network and not global_: - return None - existing = (current or {}).get("mounts") or {} - mounts = {} - for kind, ids in (("network", network), ("global", global_)): - if ids: - mounts[kind] = [{"volumeId": volume} for volume in ids] - elif existing.get(kind): - mounts[kind] = [{"volumeId": m["volumeId"]} for m in existing[kind]] - return mounts - - -def _env_body(options: dict, current: Optional[dict]) -> Optional[dict]: - env = _env_pairs(options.get("env") or ()) - unset = options.get("unset_env") or () - if current is None or not (env or unset): - return env or None - merged = dict(current.get("env") or {}) - merged.update(env) - for key in unset: - if key not in merged: - raise click.UsageError(f"--unset-env {key}: no such variable") - merged.pop(key) - return merged - - -def _check_compute(options: dict, current: Optional[dict]) -> None: - current = current or {} - if options.get("gpu") and options.get("cpu"): - raise click.UsageError("use either --gpu or --cpu") - if options.get("gpu") and current.get("cpu"): - raise click.UsageError("a CPU endpoint cannot switch to GPUs") - if options.get("cpu") and current.get("gpu"): - raise click.UsageError("a GPU endpoint cannot switch to CPUs") - if options.get("cuda") and (options.get("cpu") or current.get("cpu")): - raise click.UsageError("--cuda applies to GPU endpoints") - - -def _config_body(options: dict, current: Optional[dict] = None) -> dict: - """request body for a create (current is None) or an update.""" - updating = current is not None - endpoint_type = current.get("type") if updating else options.get("endpoint_type") - body: dict[str, Any] = {} - if not updating: - body["type"] = endpoint_type - - for key, name in ( - ("name", "name"), - ("image", "image"), - ("templateId", "template"), - ("args", "args"), - ("disk", "disk"), - ("registry", "registry"), - ): - if options.get(name) is not None: - body[key] = options[name] - for key in ("entrypoint", "cmd"): - if options.get(key) is not None: - body[key] = shlex.split(options[key]) - if options.get("port"): - body["ports"] = list(options["port"]) - if options.get("flashboot"): - body["flashboot"] = FLASHBOOT[options["flashboot"].lower()] - if options.get("request_timeout_seconds") is not None: - body["timeout"] = options["request_timeout_seconds"] * 1000 - env = _env_body(options, current) - if env is not None: - body["env"] = env - - _check_compute(options, current) - gpu = _gpu_body( - options.get("gpu"), options.get("gpu_count"), options.get("cuda"), updating - ) - if gpu: - body["gpu"] = gpu - cpu = _cpu_body(options.get("cpu"), options.get("vcpu_count"), current) - if cpu: - body["cpu"] = cpu - - workers = _workers_body(options, current) - if workers: - body["workers"] = workers - scaling = _scaling_body(options, endpoint_type) - if scaling: - body["scaling"] = scaling - elif not updating: - if endpoint_type == "LOAD_BALANCER": - raise click.UsageError("load-balancer endpoints need --requests-per-worker") - body["scaling"] = {"type": "QUEUE_DELAY", "queueDelay": DEFAULT_QUEUE_DELAY} - - places = tuple(options.get("data_center") or ()) + tuple( - options.get("region") or () - ) - if places: - body["dataCenterIds"] = compute.resolve_regions( - places, compute.list_datacenters() - ) - mounts = _mounts_body(options, current) - if mounts: - body["mounts"] = mounts - return body - - def _scale_options(func): for option in reversed( [ @@ -613,7 +315,7 @@ def _config_options(creating: bool): ), click.option( "--flashboot", - type=click.Choice(list(FLASHBOOT), case_sensitive=False), + type=click.Choice(list(config.FLASHBOOT), case_sensitive=False), help="FlashBoot mode.", ), click.option( @@ -648,16 +350,6 @@ def decorate(func): return decorate -# ------------------------------------------------------------------ rollout - - -def _rollout_state(endpoint_id: str) -> dict: - return ( - run_rest_request("GET", _path(endpoint_id, "releases"), params={"limit": 1}) - or {} - ) - - def _rollout_line(state: dict) -> str: rollout = state.get("rollout") or {} on_latest = f"{rollout.get('workersOnLatest', 0)}/{rollout.get('workersTotal', 0)}" @@ -672,14 +364,14 @@ def _follow_rollout( interval: float = 2.0, ) -> None: """poll releases until the new version runs on every worker.""" - from runpod.rp_cli import console as ui - started = time.monotonic() last = None try: while True: with api_errors(): - state = _rollout_state(endpoint_id) + state = run_rest_request( + "GET", _path(endpoint_id, "releases"), params={"limit": 1} + ) version = state.get("endpointVersion") elapsed = time.monotonic() - started if previous_version is not None and version == previous_version: @@ -711,146 +403,42 @@ def _follow_rollout( def _apply( endpoint_id: str, patch: dict, follow: bool, timeout: float, as_json: bool ) -> None: - from runpod.rp_cli import console as ui - with api_errors(): - before = _rollout_state(endpoint_id).get("endpointVersion") if follow else None - updated = run_rest_request("PATCH", _path(endpoint_id), json=patch) or {} + before = ( + run_rest_request( + "GET", _path(endpoint_id, "releases"), params={"limit": 1} + ).get("endpointVersion") + if follow + else None + ) + updated = run_rest_request("PATCH", _path(endpoint_id), json=patch) if as_json: print_json(updated) else: ui.success( - f"updated [white]{_escape(updated.get('name') or endpoint_id)}[/white] " + f"updated [white]{escape(updated['name'])}[/white] " f"[dim]{endpoint_id}[/dim]" ) if follow: _follow_rollout(endpoint_id, before, timeout) -# --------------------------------------------------------------------- edit - - -def _mutable(endpoint: dict) -> dict: - """the endpoint's configuration as UpdateEndpointRequest fields.""" - doc: dict[str, Any] = {} - for key in ("name", "image", "args", "env", "disk", "ports"): - if endpoint.get(key) is not None: - doc[key] = endpoint[key] - doc["registry"] = endpoint.get("registry") - gpu = endpoint.get("gpu") - if gpu: - doc["gpu"] = { - "pools": gpu.get("pools") or [], - "excludedTypes": gpu.get("excludedTypes") or [], - "count": gpu.get("count") or 1, - "allowedCudaVersions": gpu.get("allowedCudaVersions") or [], - "minCudaVersion": gpu.get("minCudaVersion"), - } - if endpoint.get("cpu"): - doc["cpu"] = [ - {"id": cpu["id"], "vcpuCount": cpu["vcpuCount"]} for cpu in endpoint["cpu"] - ] - for key in ("workers", "scaling", "timeout", "flashboot", "dataCenterIds"): - if endpoint.get(key) is not None: - doc[key] = endpoint[key] - mounts = endpoint.get("mounts") or {} - doc["mounts"] = { - kind: [{"volumeId": m["volumeId"]} for m in mounts.get(kind) or []] - for kind in ("network", "global") - } - return doc - - -def _kind(kind: type) -> str: - return { - str: "a string", - int: "an integer", - dict: "an object", - list: "a list", - type(None): "null", - }[kind] - - -def _validate(text: str, before: dict) -> dict: - """parse an edited document, raising ClickException on any problem.""" - try: - doc = json.loads(text) - except json.JSONDecodeError as exc: - raise click.ClickException(f"invalid JSON: {exc}") from exc - if not isinstance(doc, dict): - raise click.ClickException("the document must be a JSON object") - problems = [] - for key, value in doc.items(): - kinds = MUTABLE_FIELDS.get(key) - if kinds is None: - problems.append(f"unknown field '{key}'") - continue - if isinstance(value, bool) or not isinstance(value, kinds): - problems.append(f"'{key}' must be {' or '.join(_kind(k) for k in kinds)}") - continue - allowed = NESTED_FIELDS.get(key) - if allowed: - problems.extend( - f"unknown field '{key}.{nested}'" - for nested in value - if nested not in allowed - ) - env = doc.get("env") - if isinstance(env, dict) and not all(isinstance(v, str) for v in env.values()): - problems.append("env values must be strings") - flashboot = doc.get("flashboot") - if isinstance(flashboot, str) and flashboot not in FLASHBOOT.values(): - problems.append(f"'flashboot' must be one of {', '.join(FLASHBOOT.values())}") - if "gpu" in doc and "gpu" not in before: - problems.append("a CPU endpoint cannot switch to GPUs") - if "cpu" in doc and "cpu" not in before: - problems.append("a GPU endpoint cannot switch to CPUs") - if problems: - raise click.ClickException("; ".join(problems)) - return doc - - -def _patch(before: dict, after: dict) -> dict: - """top-level fields that changed, with gpu narrowed to its changed keys.""" - patch = {} - for key, value in after.items(): - if before.get(key) == value: - continue - if key != "gpu": - patch[key] = value - continue - old = before.get("gpu") or {} - gpu = {k: v for k, v in value.items() if old.get(k) != v} - if "pools" in gpu or "excludedTypes" in gpu: - gpu["pools"] = value.get("pools", old.get("pools")) - gpu["excludedTypes"] = value.get("excludedTypes") or [] - if "minCudaVersion" in gpu and gpu["minCudaVersion"] is None: - gpu["minCudaVersion"] = "" - patch["gpu"] = gpu - return patch - - def _edit_document(before: dict) -> Optional[dict]: """open the document in $EDITOR until it validates or the user gives up.""" - from runpod.rp_cli import console as ui - text = json.dumps(before, indent=2) + "\n" while True: edited = click.edit(text, extension=".json") if edited is None: return None try: - return _validate(edited, before) + return config.validate(edited, before) except click.ClickException as exc: - ui.error(_escape(exc.message)) + ui.error(escape(exc.message)) if not click.confirm("re-open the editor?", default=True): raise text = edited -# ------------------------------------------------------------------ workers - - def _uptime(seconds: Optional[int]) -> Optional[str]: if not seconds: return None @@ -883,98 +471,25 @@ def _worker_version(worker: dict) -> Optional[str]: ] -def _worker_summary(data: dict) -> str: - summary = data.get("summary") or {} +def _workers_view(data: dict): + from rich.console import Group + from rich.text import Text + + summary = data["summary"] counts = ", ".join( - f"{summary.get(status, 0)} {status}" + f"{summary[status]} {status}" for status in ("running", "idle", "initializing", "throttled", "unhealthy") ) version = data.get("endpointVersion") suffix = f", endpoint v{version}" if version is not None else "" - return f"{summary.get('total', 0)} workers: {counts}{suffix}" - - -def _workers_view(data: dict): - from rich.console import Group - from rich.text import Text - - header = Text(_worker_summary(data), style="dim") - workers = data.get("workers") or [] - if not workers: - return header - return Group(header, _table(workers, _WORKER_COLUMNS)) - - -# --------------------------------------------------------------------- logs - - -def _log_stream( - endpoint_id: str, - worker_id: str, - params: dict, - resume: Optional[str], - events: "queue.Queue", - stop: threading.Event, - read_timeout: float, -) -> None: - """push one worker's sse log events onto the shared queue.""" - path = _path(endpoint_id, "workers", worker_id, "logs") - event_id = None - try: - headers = dict( - _build_headers(_resolve_api_key(None)), Accept="text/event-stream" - ) - if resume: - headers["Last-Event-ID"] = resume - with requests.get( - _build_url(path), - params=params, - headers=headers, - stream=True, - timeout=(10, read_timeout), - ) as response: - if response.status_code >= 400: - _raise_for_error(response.status_code, "GET", path, response.text) - for raw in response.iter_lines(decode_unicode=True): - if stop.is_set(): - return - if raw.startswith("id:"): - event_id = raw[3:].strip() - elif raw.startswith("data:"): - try: - entry = json.loads(raw[5:].strip()) - except json.JSONDecodeError: - continue - events.put(("line", worker_id, entry, event_id or entry.get("ts"))) - except requests.exceptions.RequestException: - pass - except error.RunPodError as exc: - events.put(("error", worker_id, exc, None)) - finally: - events.put(("done", worker_id, None, None)) - - -def _log_line(worker_id: str, entry: dict, width: int, as_json: bool) -> str: - if as_json: - return json.dumps({"worker": worker_id, **entry}) + header = Text(f"{summary['total']} workers: {counts}{suffix}", style="dim") return ( - f"{worker_id:<{width}} [{entry.get('source') or '?'}] {entry.get('line', '')}" + Group(header, table(data["workers"], _WORKER_COLUMNS)) + if data["workers"] + else header ) -def _is_live(entry: dict, started: datetime) -> bool: - try: - stamp = datetime.fromisoformat(str(entry.get("ts")).replace("Z", "+00:00")) - except ValueError: - return False - if stamp.tzinfo is None: - stamp = stamp.replace(tzinfo=timezone.utc) - return stamp >= started - - -# --------------------------------------------------------------- data plane - - @contextlib.contextmanager def _data_plane_errors() -> Iterator[None]: """surface serverless data-plane failures as clean cli errors.""" @@ -990,13 +505,6 @@ def _data_plane_errors() -> Iterator[None]: raise click.ClickException(str(exc)) from exc -def _client(): - from runpod.endpoint.runner import RunPodClient - - with _data_plane_errors(): - return RunPodClient() - - def _job_path(endpoint_id: str, action: str, job_id: Optional[str] = None) -> str: parts = [quote(endpoint_id, safe=""), action] if job_id: @@ -1004,96 +512,77 @@ def _job_path(endpoint_id: str, action: str, job_id: Optional[str] = None) -> st return "/".join(parts) -def _job_failure(job: dict) -> click.ClickException: - detail = job.get("error") - status = str(job.get("status") or "unknown").lower() - return click.ClickException( - f"job {job.get('id')} {status}" + (f": {detail}" if detail else "") - ) - - -def _wait_job(client, endpoint_id: str, job_id: str, timeout: Optional[float]) -> dict: - from runpod.rp_cli import console as ui - - deadline = time.monotonic() + timeout if timeout else None - try: - with ui.console.status(f"[dim]waiting for job {job_id} ...[/dim]"): - while True: - with _data_plane_errors(): - job = client.get(_job_path(endpoint_id, "status", job_id)) - if is_completed(job.get("status")): - return job - if deadline and time.monotonic() >= deadline: - raise click.ClickException( - f"job {job_id} is still {job.get('status')} after {timeout:g}s. " - f"resume with: rp endpoint job wait {endpoint_id} {job_id}" - ) - time.sleep(1) - except KeyboardInterrupt: - click.echo( - f"\njob {job_id} keeps running. resume with: " - f"rp endpoint job wait {endpoint_id} {job_id}", - err=True, - ) - raise click.Abort() from None - - def _print_job(job: dict, as_json: bool) -> None: if as_json: print_json(job) - elif job.get("status") == "COMPLETED": + elif job["status"] == "COMPLETED": output = job.get("output") click.echo(output if isinstance(output, str) else json.dumps(output, indent=2)) - if job.get("status") != "COMPLETED": - raise _job_failure(job) + if job["status"] != "COMPLETED": + detail = f": {job['error']}" if job.get("error") else "" + raise click.ClickException(f"job {job['id']} {job['status'].lower()}{detail}") -def _stream_job(client, endpoint_id: str, job_id: str, as_json: bool) -> None: - """print streamed chunks, or the final output when the handler streams nothing.""" - open_line = False - streamed = False +def _follow_job(client, endpoint_id, job_id, as_json, *, stream=False, timeout=None): + """wait for completion, draining streamed chunks before reading the final result.""" + action = "stream" if stream else "status" + resume = "stream" if stream else "wait" + deadline = time.monotonic() + timeout if timeout else None + streamed = open_line = False + progress = ( + contextlib.nullcontext() + if stream + else ui.console.status(f"[dim]waiting for job {job_id} ...[/dim]") + ) try: - while True: - with _data_plane_errors(): + with progress: + while True: try: - partial = client.get( - _job_path(endpoint_id, "stream", job_id), - timeout=STREAM_POLL_SECONDS, + if stream: + job = client.get( + _job_path(endpoint_id, action, job_id), timeout=60 + ) + else: + job = client.get(_job_path(endpoint_id, action, job_id)) + except requests.ReadTimeout: + if stream: + continue + raise + chunks = job["stream"] if stream else [] + streamed = streamed or bool(chunks) + for chunk in chunks: + output = chunk["output"] + if isinstance(output, str) and not as_json: + click.echo(output, nl=False) + open_line = not output.endswith("\n") + else: + if open_line: + click.echo() + open_line = False + click.echo(json.dumps(output)) + if is_completed(job["status"]) and not chunks: + break + if deadline and time.monotonic() >= deadline: + raise click.ClickException( + f"job {job_id} is still {job['status']} after {timeout:g}s. " + f"resume with: rp endpoint job {resume} {endpoint_id} {job_id}" ) - except requests.exceptions.ReadTimeout: - continue - chunks = partial.get("stream") or [] - streamed = streamed or bool(chunks) - for chunk in chunks: - output = chunk.get("output") - if as_json: - click.echo(json.dumps(output)) - elif isinstance(output, str): - click.echo(output, nl=False) - open_line = not output.endswith("\n") - else: - if open_line: - click.echo() - open_line = False - click.echo(json.dumps(output)) - if is_completed(partial.get("status")) and not chunks: - break - if not chunks: - time.sleep(0.5) + if not chunks: + time.sleep(0.5 if stream else 1) except KeyboardInterrupt: click.echo( f"\njob {job_id} keeps running. resume with: " - f"rp endpoint job stream {endpoint_id} {job_id}", + f"rp endpoint job {resume} {endpoint_id} {job_id}", err=True, ) raise click.Abort() from None if open_line: click.echo() - if streamed and partial.get("status") == "COMPLETED": + if streamed and job["status"] == "COMPLETED": return - with _data_plane_errors(): - final = client.get(_job_path(endpoint_id, "status", job_id)) - _print_job(final, as_json) + if stream: + job = client.get(_job_path(endpoint_id, "status", job_id)) + _print_job(job, as_json) def _request_body(input_file, data: Optional[str]) -> Optional[Any]: @@ -1116,14 +605,16 @@ def _call_load_balancer( ) method = (method or ("POST" if body is not None else "GET")).upper() url = f"{base.rstrip('/')}/{path.lstrip('/')}" - with _data_plane_errors(): - response = requests.request( - method, - url, - headers=_build_headers(_resolve_api_key(None)), - json=body, - timeout=timeout or 600, - ) + response = requests.request( + method, + url, + headers={ + "Authorization": f"Bearer {api_key()}", + "Accept": "application/json", + }, + json=body, + timeout=timeout or 600, + ) if response.status_code >= 400: raise click.ClickException( f"{method} {url} returned {response.status_code}: {response.text.strip()}" @@ -1139,9 +630,6 @@ def _call_load_balancer( print_json(payload) -# ----------------------------------------------------------------- commands - - @click.group() def endpoint(): """Create, inspect, scale, observe, and call serverless endpoints. @@ -1169,7 +657,7 @@ def endpoint(): def list_endpoints(name, endpoint_type, as_json): """List endpoints with their IDs.""" with api_errors(): - endpoints = _list_endpoints() + endpoints = list_all("/v2/serverless", "endpoints") if name: endpoints = [ e @@ -1193,7 +681,7 @@ def get_endpoint(endpoint_id, as_json): """Show an endpoint's configuration (picker when ENDPOINT_ID is omitted).""" endpoint_id = _pick_endpoint(endpoint_id) with api_errors(): - record = _get(endpoint_id) + record = run_rest_request("GET", _path(endpoint_id)) if as_json: print_json(record) return @@ -1218,18 +706,16 @@ def create_endpoint(as_json, **options): rp endpoint create --name embeddings --type queue --image IMAGE \\ --cpu cpu3c --vcpu-count 8 """ - from runpod.rp_cli import console as ui - if not options.get("image") and not options.get("template"): raise click.UsageError("pass --image or --template") with api_errors(): - body = _config_body(options) - created = run_rest_request("POST", "/v2/serverless", json=body) or {} + body = config.config_body(options) + created = run_rest_request("POST", "/v2/serverless", json=body) if as_json: print_json(created) return ui.success( - f"created [white]{_escape(created.get('name'))}[/white] [dim]{created.get('id')}[/dim]" + f"created [white]{escape(created['name'])}[/white] [dim]{created['id']}[/dim]" ) print_fields(created, _FIELDS) @@ -1252,8 +738,8 @@ def update_endpoint(endpoint_id, wait, timeout, as_json, **options): their list. """ with api_errors(): - current = _get(endpoint_id) - patch = _config_body(options, current) + current = run_rest_request("GET", _path(endpoint_id)) + patch = config.config_body(options, current) if not patch: raise click.UsageError( "nothing to update. pass an option (see --help) or use rp endpoint edit" @@ -1283,14 +769,12 @@ def edit_endpoint(endpoint_id, document, yes, no_follow, timeout, as_json): Ctrl-C while following detaches without undoing anything. Fields removed from the document stay unchanged. """ - from runpod.rp_cli import console as ui - endpoint_id = _pick_endpoint(endpoint_id) with api_errors(): - current = _get(endpoint_id) - before = _mutable(current) + current = run_rest_request("GET", _path(endpoint_id)) + before = config.mutable(current) if document: - after = _validate(document.read(), before) + after = config.validate(document.read(), before) elif not interactive(): raise click.UsageError( "edit opens $EDITOR and needs an interactive terminal. " @@ -1298,7 +782,7 @@ def edit_endpoint(endpoint_id, document, yes, no_follow, timeout, as_json): ) else: after = _edit_document(before) - patch = _patch(before, after) if after is not None else {} + patch = config.patch(before, after) if after is not None else {} if not patch: if as_json: print_json({}) @@ -1318,25 +802,17 @@ def edit_endpoint(endpoint_id, document, yes, no_follow, timeout, as_json): @json_option def scale_endpoint(endpoint_id, as_json, **options): """Change worker limits and the autoscaler.""" - from runpod.rp_cli import console as ui - with api_errors(): - current = _get(endpoint_id) - patch = {} - workers = _workers_body(options, current) - if workers: - patch["workers"] = workers - scaling = _scaling_body(options, current.get("type")) - if scaling: - patch["scaling"] = scaling + current = run_rest_request("GET", _path(endpoint_id)) + patch = config.scale_patch(options, current) if not patch: raise click.UsageError("pass --min-workers, --max-workers, or a scaler option") with api_errors(): - updated = run_rest_request("PATCH", _path(endpoint_id), json=patch) or {} + updated = run_rest_request("PATCH", _path(endpoint_id), json=patch) if as_json: print_json(updated) return - ui.success(f"scaled [white]{_escape(updated.get('name') or endpoint_id)}[/white]") + ui.success(f"scaled [white]{escape(updated['name'])}[/white]") print_fields( updated, [ @@ -1356,14 +832,12 @@ def scale_endpoint(endpoint_id, as_json, **options): @click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.") def delete_endpoint(endpoint_id, yes): """Delete an endpoint and its workers.""" - from runpod.rp_cli import console as ui - with api_errors(): - name = _get(endpoint_id).get("name") or endpoint_id + name = run_rest_request("GET", _path(endpoint_id))["name"] confirm(f"delete endpoint '{name}' ({endpoint_id})?", yes) with api_errors(): run_rest_request("DELETE", _path(endpoint_id)) - ui.success(f"deleted [white]{_escape(name)}[/white] [dim]{endpoint_id}[/dim]") + ui.success(f"deleted [white]{escape(name)}[/white] [dim]{endpoint_id}[/dim]") @endpoint.command(name="workers") @@ -1379,11 +853,9 @@ def delete_endpoint(endpoint_id, yes): @json_option def list_workers(endpoint_id, watch, interval, as_json): """Show an endpoint's workers. With --watch, --json prints one line per refresh.""" - from runpod.rp_cli import console as ui - try: with api_errors(): - data = _list_workers(endpoint_id) + data = run_rest_request("GET", _path(endpoint_id, "workers")) if not watch: if as_json: print_json(data) @@ -1395,14 +867,14 @@ def list_workers(endpoint_id, watch, interval, as_json): click.echo(json.dumps(data, default=str)) time.sleep(interval) with api_errors(): - data = _list_workers(endpoint_id) + data = run_rest_request("GET", _path(endpoint_id, "workers")) from rich.live import Live with Live(_workers_view(data), console=ui.console, auto_refresh=False) as live: while True: time.sleep(interval) with api_errors(): - data = _list_workers(endpoint_id) + data = run_rest_request("GET", _path(endpoint_id, "workers")) live.update(_workers_view(data), refresh=True) except KeyboardInterrupt: pass @@ -1422,11 +894,11 @@ def _release_changes(release: dict) -> list[str]: @json_option def list_releases(endpoint_id, as_json): """Show release history and rollout progress.""" - from runpod.rp_cli import console as ui - with api_errors(): - state = _rollout_state(endpoint_id) - releases = compute.list_all(_path(endpoint_id, "releases"), "releases") + state = run_rest_request( + "GET", _path(endpoint_id, "releases"), params={"limit": 1} + ) + releases = list_all(_path(endpoint_id, "releases"), "releases") if as_json: print_json( { @@ -1488,85 +960,17 @@ def logs(endpoint_id, worker_ids, follow, source, tail, since, as_json): Without --worker, every worker's logs are multiplexed, and with --follow newly created workers are attached automatically. """ - from runpod.rp_cli import console as ui - - params: dict[str, Any] = {"tail": tail} - if source: - params["source"] = source - if since: - params["since"] = since - wanted = set(worker_ids) - started = datetime.now(timezone.utc) - events: queue.Queue = queue.Queue() - stop = threading.Event() - attached: set[str] = set() - cursors: dict[str, str] = {} - collected: list[tuple[str, dict]] = [] - width = max((len(w) for w in wanted), default=0) - - def attach_new() -> None: - nonlocal width - with api_errors(): - listed = [w["id"] for w in _list_workers(endpoint_id).get("workers") or []] - for worker_id in listed: - if (wanted and worker_id not in wanted) or worker_id in attached: - continue - attached.add(worker_id) - width = max(width, len(worker_id)) - threading.Thread( - target=_log_stream, - args=( - endpoint_id, - worker_id, - params, - cursors.get(worker_id), - events, - stop, - 90.0 if follow else 3.0, - ), - daemon=True, - ).start() - - try: - attach_new() - if not attached: - if not follow: - ui.console.print("[dim]no workers[/dim]") - return - ui.console.print("[dim]waiting for workers ...[/dim]") - next_poll = time.monotonic() + WORKER_POLL_SECONDS - while follow or attached: - try: - kind, worker_id, payload, cursor = events.get(timeout=0.5) - except queue.Empty: - kind = None - if kind == "line": - if cursor: - cursors[worker_id] = cursor - if follow: - click.echo(_log_line(worker_id, payload, width, as_json)) - elif _is_live(payload, started): - attached.discard(worker_id) - else: - collected.append((worker_id, payload)) - elif kind == "error": - if not ( - isinstance(payload, error.QueryError) and payload.status_code == 404 - ): - raise click.ClickException(str(payload)) - elif kind == "done": - attached.discard(worker_id) - if follow and time.monotonic() >= next_poll: - attach_new() - next_poll = time.monotonic() + WORKER_POLL_SECONDS - except KeyboardInterrupt: - pass - finally: - stop.set() - for worker_id, entry in sorted( - collected, key=lambda item: str(item[1].get("ts") or "") - ): - click.echo(_log_line(worker_id, entry, width, as_json)) + from runpod.rp_cli.logs import show_logs + + show_logs( + endpoint_id=endpoint_id, + worker_ids=worker_ids, + follow=follow, + source=source, + tail=tail, + since=since, + as_json=as_json, + ) @endpoint.command(name="run") @@ -1597,6 +1001,7 @@ def attach_new() -> None: help="Load-balancer endpoints: HTTP method (POST with a body, else GET).", ) @json_option +@_data_plane_errors() def run_endpoint( endpoint_id, input_file, @@ -1625,13 +1030,11 @@ def run_endpoint( rp endpoint run ENDPOINT_ID --input request.json --async rp endpoint run LB_ENDPOINT_ID --path /v1/models """ - from runpod.rp_cli import console as ui - if run_async and stream: raise click.UsageError("use either --async or --stream") body = _request_body(input_file, data) with api_errors(): - record = _get(endpoint_id) + record = run_rest_request("GET", _path(endpoint_id)) if record.get("type") == "LOAD_BALANCER": if run_async or stream or webhook or execution_timeout or ttl or low_priority: @@ -1659,10 +1062,9 @@ def run_endpoint( if policy: request["policy"] = policy - client = _client() - with _data_plane_errors(): - submitted = client.post(_job_path(endpoint_id, "run"), request) - job_id = submitted.get("id") + client = runner.RunPodClient() + submitted = client.post(_job_path(endpoint_id, "run"), request) + job_id = submitted["id"] if run_async: if as_json: print_json(submitted) @@ -1675,9 +1077,9 @@ def run_endpoint( if stream: if not as_json: click.echo(f"job {job_id}", err=True) - _stream_job(client, endpoint_id, job_id, as_json) + _follow_job(client, endpoint_id, job_id, as_json, stream=True) return - _print_job(_wait_job(client, endpoint_id, job_id, timeout), as_json) + _follow_job(client, endpoint_id, job_id, as_json, timeout=timeout) @endpoint.group(name="job") @@ -1714,10 +1116,10 @@ def _job_output(record: dict) -> Optional[str]: @_job_command("get") +@_data_plane_errors() def job_get(endpoint_id, job_id, as_json): """Show a job's status and output.""" - with _data_plane_errors(): - record = _client().get(_job_path(endpoint_id, "status", job_id)) + record = runner.RunPodClient().get(_job_path(endpoint_id, "status", job_id)) if as_json: print_json(record) return @@ -1726,24 +1128,24 @@ def job_get(endpoint_id, job_id, as_json): @_job_command("wait") @click.option("--timeout", type=float, help="Seconds to wait.") +@_data_plane_errors() def job_wait(endpoint_id, job_id, timeout, as_json): """Wait for a job to finish and print its output.""" - _print_job(_wait_job(_client(), endpoint_id, job_id, timeout), as_json) + _follow_job(runner.RunPodClient(), endpoint_id, job_id, as_json, timeout=timeout) @_job_command("stream") +@_data_plane_errors() def job_stream(endpoint_id, job_id, as_json): """Print a job's streamed output as it arrives.""" - _stream_job(_client(), endpoint_id, job_id, as_json) + _follow_job(runner.RunPodClient(), endpoint_id, job_id, as_json, stream=True) @_job_command("cancel") +@_data_plane_errors() def job_cancel(endpoint_id, job_id, as_json): """Cancel a queued or running job.""" - from runpod.rp_cli import console as ui - - with _data_plane_errors(): - record = _client().post(_job_path(endpoint_id, "cancel", job_id), None) + record = runner.RunPodClient().post(_job_path(endpoint_id, "cancel", job_id), None) if as_json: print_json(record) return @@ -1751,12 +1153,10 @@ def job_cancel(endpoint_id, job_id, as_json): @_job_command("retry") +@_data_plane_errors() def job_retry(endpoint_id, job_id, as_json): """Requeue a failed or timed-out job.""" - from runpod.rp_cli import console as ui - - with _data_plane_errors(): - record = _client().post(_job_path(endpoint_id, "retry", job_id), None) + record = runner.RunPodClient().post(_job_path(endpoint_id, "retry", job_id), None) if as_json: print_json(record) return @@ -1769,13 +1169,11 @@ def job_retry(endpoint_id, job_id, as_json): @endpoint_argument() @click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.") @json_option +@_data_plane_errors() def purge_queue(endpoint_id, yes, as_json): """Remove every queued job. Running jobs are not affected.""" - from runpod.rp_cli import console as ui - confirm(f"remove every queued job from {endpoint_id}?", yes) - with _data_plane_errors(): - record = _client().post(_job_path(endpoint_id, "purge-queue"), None) + record = runner.RunPodClient().post(_job_path(endpoint_id, "purge-queue"), None) if as_json: print_json(record) return diff --git a/runpod/rp_cli/endpoint_config.py b/runpod/rp_cli/endpoint_config.py new file mode 100644 index 00000000..8f37d45b --- /dev/null +++ b/runpod/rp_cli/endpoint_config.py @@ -0,0 +1,314 @@ +"""endpoint request configuration and editable JSON documents.""" + +import json +import shlex +from typing import Any, Optional + +import click + +from runpod.rp_cli import compute +from runpod.rp_cli.options import parse_env + +FLASHBOOT = {"off": "OFF", "standard": "FLASHBOOT", "priority": "PRIORITY_FLASHBOOT"} + + +MUTABLE_FIELDS = { + "name": (str,), + "image": (str,), + "args": (str,), + "env": (dict,), + "disk": (int,), + "ports": (list,), + "registry": (str, type(None)), + "gpu": (dict,), + "cpu": (list,), + "workers": (dict,), + "scaling": (dict,), + "timeout": (int,), + "flashboot": (str,), + "dataCenterIds": (list,), + "mounts": (dict,), +} + + +NESTED_FIELDS = { + "gpu": {"pools", "excludedTypes", "count", "allowedCudaVersions", "minCudaVersion"}, + "workers": {"min", "max", "idleTimeout"}, + "scaling": {"type", "queueDelay", "requestCount"}, + "mounts": {"network", "global"}, +} + + +def _gpu(specs, count, cuda, updating: bool) -> dict: + gpu: dict[str, Any] = {} + if specs: + catalog = [ + entry + for entry in compute.list_gpus(product="SERVERLESS") + if entry.get("pool") + ] + matches: list[dict] = [] + counts = {count} if count else set() + for spec in specs: + name, spec_count = compute.parse_gpu_spec(spec) + if spec_count: + counts.add(spec_count) + for match in compute.resolve_gpus(name, catalog): + if match not in matches: + matches.append(match) + if len(counts) > 1: + raise click.UsageError("conflicting gpu counts; use one count per worker") + pools, excluded = compute.endpoint_gpu_selection(matches, catalog) + gpu["pools"] = pools + if excluded or updating: + gpu["excludedTypes"] = excluded + count = counts.pop() if counts else None + if count: + gpu["count"] = count + if cuda: + if [value.lower() for value in cuda] == ["any"]: + gpu["allowedCudaVersions"] = [] + if updating: + gpu["minCudaVersion"] = "" + return gpu + minimum, exact = compute.parse_cuda(cuda) + if minimum: + gpu["minCudaVersion"] = minimum + if updating: + gpu["allowedCudaVersions"] = [] + else: + gpu["allowedCudaVersions"] = exact + if updating: + gpu["minCudaVersion"] = "" + return gpu + + +def _cpu(specs, vcpu_count, current: Optional[dict]) -> Optional[list]: + if not specs: + if vcpu_count and current and current.get("cpu"): + return [ + {"id": cpu["id"], "vcpuCount": vcpu_count} for cpu in current["cpu"] + ] + if vcpu_count: + raise click.UsageError("--vcpu-count applies to CPU endpoints (pass --cpu)") + return None + flavors = compute.list_cpus(product="SERVERLESS") + entries = [] + for spec in specs: + name, sep, count = spec.rpartition(":") + if not sep: + name, count = spec, "" + if count and not count.isdigit(): + raise click.UsageError( + f"invalid --cpu '{spec}'. use a flavor like cpu3c or cpu3c:8" + ) + vcpus = int(count) if count else vcpu_count + if not vcpus: + raise click.UsageError(f"pass --vcpu-count or --cpu {name}:VCPUS") + entries.append( + {"id": compute.resolve_cpu(name, flavors)["id"], "vcpuCount": vcpus} + ) + return entries + + +def scale_patch(options: dict, current: dict) -> dict: + patch = {} + changes = { + key: options[flag] + for key, flag in ( + ("min", "min_workers"), + ("max", "max_workers"), + ("idleTimeout", "idle_timeout_seconds"), + ) + if options.get(flag) is not None + } + if changes: + patch["workers"] = {**current.get("workers", {}), **changes} + delay = options.get("queue_delay_seconds") + count = options.get("requests_per_worker") + if delay is not None and count is not None: + raise click.UsageError( + "use either --queue-delay-seconds or --requests-per-worker" + ) + if delay is not None: + if current.get("type") == "LOAD_BALANCER": + raise click.UsageError( + "load-balancer endpoints scale on --requests-per-worker" + ) + patch["scaling"] = {"type": "QUEUE_DELAY", "queueDelay": delay} + elif count is not None: + patch["scaling"] = {"type": "REQUEST_COUNT", "requestCount": count} + return patch + + +def config_body(options: dict, current: Optional[dict] = None) -> dict: + """map flags to a create or partial update, preserving untouched configuration.""" + updating = current is not None + current = current if updating else {"type": options["endpoint_type"]} + body = {} + if not updating: + body["type"] = current["type"] + for key, flag in ( + ("name", "name"), + ("image", "image"), + ("templateId", "template"), + ("args", "args"), + ("disk", "disk"), + ("registry", "registry"), + ): + if options.get(flag) is not None: + body[key] = options[flag] + for key in ("entrypoint", "cmd"): + if options.get(key) is not None: + body[key] = shlex.split(options[key]) + if options.get("port"): + body["ports"] = list(options["port"]) + if options.get("flashboot"): + body["flashboot"] = FLASHBOOT[options["flashboot"].lower()] + if options.get("request_timeout_seconds") is not None: + body["timeout"] = options["request_timeout_seconds"] * 1000 + + if options.get("env") or options.get("unset_env"): + env = {**current.get("env", {}), **parse_env(options.get("env", ()))} + for key in options.get("unset_env", ()): + if key not in env: + raise click.UsageError(f"--unset-env {key}: no such variable") + del env[key] + body["env"] = env + + if options.get("gpu") and options.get("cpu"): + raise click.UsageError("use either --gpu or --cpu") + if current.get("cpu") and options.get("gpu"): + raise click.UsageError("a CPU endpoint cannot switch to GPUs") + if current.get("gpu") and options.get("cpu"): + raise click.UsageError("a GPU endpoint cannot switch to CPUs") + if (options.get("cpu") or current.get("cpu")) and ( + options.get("cuda") or options.get("gpu_count") + ): + raise click.UsageError("--cuda and --gpu-count apply to GPU endpoints") + gpu = _gpu( + options.get("gpu"), options.get("gpu_count"), options.get("cuda"), updating + ) + if gpu: + body["gpu"] = gpu + cpu = _cpu(options.get("cpu"), options.get("vcpu_count"), current) + if cpu: + body["cpu"] = cpu + + body.update(scale_patch(options, current)) + if not updating and "scaling" not in body: + if current["type"] == "LOAD_BALANCER": + raise click.UsageError("load-balancer endpoints need --requests-per-worker") + body["scaling"] = {"type": "QUEUE_DELAY", "queueDelay": 4} + places = (*options.get("data_center", ()), *options.get("region", ())) + if places: + body["dataCenterIds"] = compute.resolve_regions( + places, compute.list_datacenters() + ) + if options.get("network_volume") or options.get("global_volume"): + mounts = dict(current.get("mounts", {})) + for kind, flag in (("network", "network_volume"), ("global", "global_volume")): + if options.get(flag): + mounts[kind] = [{"volumeId": volume} for volume in options[flag]] + body["mounts"] = mounts + return body + + +def mutable(endpoint: dict) -> dict: + """the endpoint's configuration as UpdateEndpointRequest fields.""" + doc: dict[str, Any] = {} + for key in ("name", "image", "args", "env", "disk", "ports"): + if endpoint.get(key) is not None: + doc[key] = endpoint[key] + doc["registry"] = endpoint.get("registry") + gpu = endpoint.get("gpu") + if gpu: + doc["gpu"] = { + "pools": gpu.get("pools") or [], + "excludedTypes": gpu.get("excludedTypes") or [], + "count": gpu.get("count") or 1, + "allowedCudaVersions": gpu.get("allowedCudaVersions") or [], + "minCudaVersion": gpu.get("minCudaVersion"), + } + if endpoint.get("cpu"): + doc["cpu"] = [ + {"id": cpu["id"], "vcpuCount": cpu["vcpuCount"]} for cpu in endpoint["cpu"] + ] + for key in ("workers", "scaling", "timeout", "flashboot", "dataCenterIds"): + if endpoint.get(key) is not None: + doc[key] = endpoint[key] + mounts = endpoint.get("mounts") or {} + doc["mounts"] = { + kind: [{"volumeId": m["volumeId"]} for m in mounts.get(kind) or []] + for kind in ("network", "global") + } + return doc + + +def _kind(kind: type) -> str: + return { + str: "a string", + int: "an integer", + dict: "an object", + list: "a list", + type(None): "null", + }[kind] + + +def validate(text: str, before: dict) -> dict: + """parse an edited document, raising ClickException on any problem.""" + try: + doc = json.loads(text) + except json.JSONDecodeError as exc: + raise click.ClickException(f"invalid JSON: {exc}") from exc + if not isinstance(doc, dict): + raise click.ClickException("the document must be a JSON object") + problems = [] + for key, value in doc.items(): + kinds = MUTABLE_FIELDS.get(key) + if kinds is None: + problems.append(f"unknown field '{key}'") + continue + if isinstance(value, bool) or not isinstance(value, kinds): + problems.append(f"'{key}' must be {' or '.join(_kind(k) for k in kinds)}") + continue + allowed = NESTED_FIELDS.get(key) + if allowed: + problems.extend( + f"unknown field '{key}.{nested}'" + for nested in value + if nested not in allowed + ) + env = doc.get("env") + if isinstance(env, dict) and not all(isinstance(v, str) for v in env.values()): + problems.append("env values must be strings") + flashboot = doc.get("flashboot") + if isinstance(flashboot, str) and flashboot not in FLASHBOOT.values(): + problems.append(f"'flashboot' must be one of {', '.join(FLASHBOOT.values())}") + if "gpu" in doc and "gpu" not in before: + problems.append("a CPU endpoint cannot switch to GPUs") + if "cpu" in doc and "cpu" not in before: + problems.append("a GPU endpoint cannot switch to CPUs") + if problems: + raise click.ClickException("; ".join(problems)) + return doc + + +def patch(before: dict, after: dict) -> dict: + """top-level fields that changed, with gpu narrowed to its changed keys.""" + patch = {} + for key, value in after.items(): + if before.get(key) == value: + continue + if key != "gpu": + patch[key] = value + continue + old = before.get("gpu") or {} + gpu = {k: v for k, v in value.items() if old.get(k) != v} + if "pools" in gpu or "excludedTypes" in gpu: + gpu["pools"] = value.get("pools", old.get("pools")) + gpu["excludedTypes"] = value.get("excludedTypes") or [] + if "minCudaVersion" in gpu and gpu["minCudaVersion"] is None: + gpu["minCudaVersion"] = "" + patch["gpu"] = gpu + return patch diff --git a/tests/test_cli/test_endpoint_command.py b/tests/test_cli/test_endpoint_command.py index 9068fbea..e966ac86 100644 --- a/tests/test_cli/test_endpoint_command.py +++ b/tests/test_cli/test_endpoint_command.py @@ -2,8 +2,6 @@ import copy import json -import threading -import queue from unittest.mock import MagicMock import click @@ -13,10 +11,12 @@ from runpod import error from runpod.endpoint import runner -from runpod.rp_cli import compute, output +from runpod.rp_cli import api as cli_api, compute, endpoint_config as config, output from runpod.rp_cli import endpoint as ep from runpod.rp_cli.main import cli +PAGE = {"hasNextPage": False, "nextCursor": None} + GPUS = [ {"id": "NVIDIA H100 80GB HBM3", "name": "H100 SXM", "pool": "ADA_80_PRO"}, {"id": "NVIDIA H100 NVL", "name": "H100 NVL", "pool": "ADA_80_PRO"}, @@ -158,7 +158,13 @@ "gpuCount": 1, "uptimeSeconds": 125, }, - {"id": "w3", "status": "INITIALIZING", "isStale": True, "gpuCount": 0}, + { + "id": "w3", + "status": "INITIALIZING", + "isStale": True, + "gpuCount": 0, + "gpuTypeId": "unknown", + }, ], } @@ -172,7 +178,7 @@ def __init__(self, routes): def __call__(self, method, path, params=None, json=None, **_): self.calls.append((method, path, params, json)) - value = self.routes.get((method, path)) + value = self.routes[(method, path)] if callable(value): value = value(params, json) if isinstance(value, Exception): @@ -184,7 +190,11 @@ def bodies(self, method): @pytest.fixture(autouse=True) -def wide_console(monkeypatch): +def cli_environment(monkeypatch): + def no_network(*args, **kwargs): + pytest.fail("endpoint tests must not make real network requests") + + monkeypatch.setattr(requests.sessions.Session, "request", no_network) from runpod.rp_cli import console as ui monkeypatch.setattr(ui.console, "_width", 200) @@ -193,9 +203,12 @@ def wide_console(monkeypatch): @pytest.fixture def api(monkeypatch): routes = { - ("GET", "/v2/catalog/gpus"): {"gpus": GPUS}, - ("GET", "/v2/catalog/cpus"): {"cpus": CPUS}, - ("GET", "/v2/catalog/datacenters"): {"dataCenters": DATACENTERS}, + ("GET", "/v2/catalog/gpus"): {"gpus": GPUS, "pagination": PAGE}, + ("GET", "/v2/catalog/cpus"): {"cpus": CPUS, "pagination": PAGE}, + ("GET", "/v2/catalog/datacenters"): { + "dataCenters": DATACENTERS, + "pagination": PAGE, + }, ("GET", "/v2/serverless"): { "endpoints": [GPU_ENDPOINT, CPU_ENDPOINT, LB_ENDPOINT], "pagination": {"hasNextPage": False, "nextCursor": None}, @@ -218,6 +231,7 @@ def api(monkeypatch): fake = FakeApi(routes) monkeypatch.setattr(ep, "run_rest_request", fake) monkeypatch.setattr(compute, "run_rest_request", fake) + monkeypatch.setattr(cli_api, "run_rest_request", fake) monkeypatch.setattr(ep.time, "sleep", lambda seconds: None) return fake @@ -245,9 +259,6 @@ def fails(*args, input=None): return result.output -# ---------------------------------------------------------------- list/get - - def test_list_shows_ids_and_summaries(api): out = ok("list") for text in ("ep1", "ep2", "lb1", "llama-prod", "load-balancer", "ADA_80_PRO x1"): @@ -300,9 +311,6 @@ def test_api_errors_are_clean(api): assert "endpoint not found" in fails("get", "ep1") -# ------------------------------------------------------------------ create - - def test_create_queue_endpoint_from_ticket(api): ok( "create", @@ -450,7 +458,7 @@ def test_create_exact_cuda_and_cpu_vcpu_suffix(api): (["--name", "x", "--type", "queue", "--image", "i", "--cpu", "cpu3c:x"], "invalid --cpu"), (["--name", "x", "--type", "queue", "--image", "i", "--vcpu-count", "4"], "pass --cpu"), (["--name", "x", "--type", "queue", "--image", "i", "--cpu", "cpu3c:4", "--cuda", "12.4"], - "--cuda applies to GPU"), + "apply to GPU"), (["--name", "x", "--type", "queue", "--image", "i", "--gpu", "b200"], "unknown gpu"), (["--name", "x", "--type", "queue", "--image", "i", "--region", "mars"], "unknown region"), ], @@ -460,9 +468,6 @@ def test_create_rejects_bad_flags(api, args, message): assert not api.bodies("POST") -# ------------------------------------------------------------------ update - - def test_update_merges_env_and_workers(api): ok("update", "ep1", "--env", "C=3", "--unset-env", "A", "--max-workers", "8") (patch,) = api.bodies("PATCH") @@ -504,7 +509,8 @@ def test_update_cpu_vcpus_and_mounts(api): (["ep1", "--unset-env", "MISSING"], "no such variable"), (["ep1", "--cpu", "cpu3c:4"], "cannot switch to CPUs"), (["ep2", "--gpu", "h100"], "cannot switch to GPUs"), - (["ep2", "--cuda", "12.4"], "--cuda applies to GPU"), + (["ep2", "--cuda", "12.4"], "apply to GPU"), + (["ep2", "--gpu-count", "1"], "apply to GPU"), (["lb1", "--queue-delay-seconds", "3"], "scale on --requests-per-worker"), ], ) @@ -533,25 +539,24 @@ def test_update_json_and_wait_follows_rollout(api): assert json.loads(out)["name"] == "renamed" -# -------------------------------------------------------------------- edit - - def _doc(**changes): - doc = ep._mutable(GPU_ENDPOINT) + doc = config.mutable(GPU_ENDPOINT) doc.update(changes) return json.dumps(doc) def test_mutable_document_holds_update_fields_only(): - doc = ep._mutable(GPU_ENDPOINT) - assert set(doc) <= set(ep.MUTABLE_FIELDS) + doc = config.mutable(GPU_ENDPOINT) + assert set(doc) <= set(config.MUTABLE_FIELDS) assert "networkVolumes" not in doc and "id" not in doc and "requestUrls" not in doc - assert ep._mutable(CPU_ENDPOINT)["cpu"] == [{"id": "cpu3c", "vcpuCount": 4}] + assert config.mutable(CPU_ENDPOINT)["cpu"] == [{"id": "cpu3c", "vcpuCount": 4}] def test_edit_file_shows_patch_and_applies(api, tmp_path): path = tmp_path / "doc.json" - gpu = dict(ep._mutable(GPU_ENDPOINT)["gpu"], excludedTypes=[], minCudaVersion=None) + gpu = dict( + config.mutable(GPU_ENDPOINT)["gpu"], excludedTypes=[], minCudaVersion=None + ) path.write_text( _doc(image="img:2", workers={"min": 1, "max": 3, "idleTimeout": 5}, gpu=gpu) ) @@ -567,11 +572,11 @@ def test_edit_file_shows_patch_and_applies(api, tmp_path): def test_edit_clears_cuda_floor_with_empty_string(): - before = ep._mutable(LB_ENDPOINT) + before = config.mutable(LB_ENDPOINT) after = copy.deepcopy(before) after["gpu"]["minCudaVersion"] = None after["gpu"]["count"] = 1 - assert ep._patch(before, after) == {"gpu": {"count": 1, "minCudaVersion": ""}} + assert config.patch(before, after) == {"gpu": {"count": 1, "minCudaVersion": ""}} def test_edit_without_changes(api, tmp_path): @@ -612,9 +617,12 @@ def test_edit_needs_terminal_or_file(api): ) def test_edit_validation(text, message): with pytest.raises(click.ClickException, match=message): - ep._validate(text, ep._mutable(GPU_ENDPOINT)) + config.validate(text, config.mutable(GPU_ENDPOINT)) + + +def test_edit_rejects_switching_cpu_to_gpu(): with pytest.raises(click.ClickException, match="cannot switch to GPUs"): - ep._validate(json.dumps({"gpu": {}}), ep._mutable(CPU_ENDPOINT)) + config.validate('{"gpu": {}}', config.mutable(CPU_ENDPOINT)) def test_edit_in_editor_reopens_on_error_then_follows(api, monkeypatch): @@ -667,9 +675,6 @@ def interrupt(seconds): assert "detached" in capsys.readouterr().out -# ------------------------------------------------------------ scale/delete - - def test_scale_sends_workers_and_scaler(api): out = ok("scale", "ep1", "--min-workers", "1", "--requests-per-worker", "3") (patch,) = api.bodies("PATCH") @@ -700,11 +705,9 @@ def test_delete_confirms(api, monkeypatch): assert len([c for c in api.calls if c[0] == "DELETE"]) == 2 -# ------------------------------------------------------- workers/releases - - def test_workers_table_and_json(api): out = ok("workers", "ep1") + assert "unknown" not in out assert "3 workers: 1 running, 1 idle, 1 initializing" in out for text in ("w1", "w2", "w3", "1h02m", "2m05s", "v2 (stale)", "stale"): assert text in out @@ -715,7 +718,10 @@ def test_workers_empty_and_uptime(): assert ep._uptime(None) is None assert ep._uptime(42) == "42s" assert ep._worker_version({"version": None}) is None - assert "0 workers" in str(ep._workers_view({"summary": {}})) + empty = copy.deepcopy(WORKERS) + empty["summary"] = dict.fromkeys(WORKERS["summary"], 0) + empty["workers"] = [] + assert "0 workers" in str(ep._workers_view(empty)) def test_workers_watch(api, monkeypatch): @@ -752,159 +758,36 @@ def test_releases(api): assert data["endpointVersion"] == 3 and len(data["releases"]) == 2 -# -------------------------------------------------------------------- logs - - -class FakeStream: - def __init__(self, lines, status=200, text=""): - self.lines = lines - self.status_code = status - self.text = text - - def __enter__(self): - return self - - def __exit__(self, *exc): - return False - - def iter_lines(self, decode_unicode=True): - for line in self.lines: - if isinstance(line, Exception): - raise line - yield line - - -def _event(ts, line, source="container"): - return [ - f"id: {ts}", - "data: " + json.dumps({"source": source, "line": line, "ts": ts}), - "", - ] - - -@pytest.fixture -def streams(monkeypatch): - calls = [] - by_worker = {} - - def fake_get(url, params=None, headers=None, stream=None, timeout=None): - worker = url.split("/workers/")[1].split("/")[0] - calls.append({"worker": worker, "params": params, "headers": headers}) - return by_worker[worker] - - monkeypatch.setattr(ep.requests, "get", fake_get) - monkeypatch.setattr(ep, "_resolve_api_key", lambda key: "key") - by_worker["calls"] = calls - return by_worker - - -def test_logs_snapshot_multiplexes_sorted_and_stops_at_live_lines(api, streams): - streams["w1"] = FakeStream( - _event("2026-01-01T00:00:02Z", "w1 second") - + ["data: {broken", ": heartbeat"] - + _event("2999-01-01T00:00:00Z", "live line") - ) - streams["w2"] = FakeStream(_event("2026-01-01T00:00:01Z", "w2 first", "system")) - streams["w3"] = FakeStream([], status=404, text='{"title": "worker gone"}') - out = ok("logs", "ep1", "--source", "system", "--tail", "200") - lines = out.splitlines() - assert lines == ["w2 [system] w2 first", "w1 [container] w1 second"] - params = streams["calls"][0]["params"] - assert params == {"tail": 200, "source": "system"} - - -def test_logs_filters_workers_and_prints_json(api, streams): - streams["w2"] = FakeStream(_event("2026-01-01T00:00:01Z", "hello")) - out = ok( - "logs", "ep1", "--worker", "w2", "--since", "2026-01-01T00:00:00Z", "--json" - ) - assert json.loads(out) == { - "worker": "w2", - "source": "container", - "line": "hello", - "ts": "2026-01-01T00:00:01Z", - } - assert [call["worker"] for call in streams["calls"]] == ["w2"] - assert streams["calls"][0]["params"]["since"] == "2026-01-01T00:00:00Z" - - -def test_logs_without_workers(api): - api.routes[("GET", "/v2/serverless/ep1/workers")] = {"workers": []} - assert "no workers" in ok("logs", "ep1") - - -def test_logs_auth_error_fails(api, streams): - for worker in ("w1", "w2", "w3"): - streams[worker] = FakeStream([], status=401) - assert "Unauthorized" in fails("logs", "ep1") +def test_logs_routes_options_to_shared_streamer(api, monkeypatch): + from runpod.rp_cli import logs - -def test_logs_follow_attaches_new_workers_and_resumes(api, streams, monkeypatch): - monkeypatch.setattr(ep, "WORKER_POLL_SECONDS", 0) - listings = iter( - [ - {"workers": [{"id": "w1"}]}, - {"workers": [{"id": "w1"}, {"id": "w2"}]}, - ] + show = MagicMock() + monkeypatch.setattr(logs, "show_logs", show) + ok( + "logs", + "ep1", + "--worker", + "w1", + "--worker", + "w2", + "--follow", + "--source", + "system", + "--tail", + "20", + "--since", + "2026-01-01T00:00:00Z", + "--json", ) - - def list_workers(params, body): - try: - return next(listings) - except StopIteration: - raise KeyboardInterrupt from None - - api.routes[("GET", "/v2/serverless/ep1/workers")] = list_workers - streams["w1"] = FakeStream(_event("2026-01-01T00:00:01Z", "one")) - streams["w2"] = FakeStream( - _event("2026-01-01T00:00:02Z", "two") + [requests.exceptions.ReadTimeout()] + show.assert_called_once_with( + endpoint_id="ep1", + worker_ids=("w1", "w2"), + follow=True, + source="system", + tail=20, + since="2026-01-01T00:00:00Z", + as_json=True, ) - out = invoke("logs", "ep1", "--follow").output - assert "w1 [container] one" in out - - -def test_logs_follow_waits_for_workers(api, monkeypatch): - monkeypatch.setattr(ep, "WORKER_POLL_SECONDS", 0) - listings = iter([{"workers": []}]) - - def list_workers(params, body): - try: - return next(listings) - except StopIteration: - raise KeyboardInterrupt from None - - api.routes[("GET", "/v2/serverless/ep1/workers")] = list_workers - assert "waiting for workers" in ok("logs", "ep1", "-f") - - -def test_log_stream_sends_resume_cursor(monkeypatch): - seen = {} - - def fake_get(url, params=None, headers=None, stream=None, timeout=None): - seen["headers"] = headers - return FakeStream(_event("2026-01-01T00:00:01Z", "x")) - - monkeypatch.setattr(ep.requests, "get", fake_get) - monkeypatch.setattr(ep, "_resolve_api_key", lambda key: "key") - events = queue.Queue() - stop = threading.Event() - ep._log_stream("ep1", "w1", {}, "2026-01-01T00:00:00Z", events, stop, 1) - assert seen["headers"]["Last-Event-ID"] == "2026-01-01T00:00:00Z" - assert events.get()[0] == "line" and events.get()[0] == "done" - stop.set() - ep._log_stream("ep1", "w1", {}, None, events, stop, 1) - assert events.get()[0] == "done" - - -def test_is_live_handles_odd_timestamps(): - from datetime import datetime, timezone - - started = datetime(2026, 1, 1, tzinfo=timezone.utc) - assert not ep._is_live({"ts": "garbage"}, started) - assert ep._is_live({"ts": "2026-02-01T00:00:00"}, started) - - -# --------------------------------------------------------------------- run def test_run_waits_for_output(api, client): @@ -1046,7 +929,7 @@ def fake_request(method, url, headers=None, json=None, timeout=None): return next(responses) monkeypatch.setattr(ep.requests, "request", fake_request) - monkeypatch.setattr(ep, "_resolve_api_key", lambda key: "key") + monkeypatch.setattr(ep, "api_key", lambda: "test-key") assert json.loads(ok("run", "lb1", "--path", "/v1/chat", "--data", '{"m": 1}')) == { "data": [1] } @@ -1058,9 +941,6 @@ def fake_request(method, url, headers=None, json=None, timeout=None): assert calls[2][0] == "PUT" and calls[2][3] == 5.0 -# -------------------------------------------------------------------- jobs - - def test_job_get(api, client): client.get.return_value = { "id": "job1", @@ -1085,7 +965,7 @@ def test_job_wait_and_stream(api, client): {"status": "COMPLETED", "stream": []}, ] assert ok("job", "stream", "ep1", "job1").strip() == "chunk" - client.get.assert_called_with("ep1/stream/job1", timeout=ep.STREAM_POLL_SECONDS) + client.get.assert_called_with("ep1/stream/job1", timeout=60) def test_job_stream_interrupt(api, client): @@ -1124,6 +1004,8 @@ def test_data_plane_errors(api, client, monkeypatch): assert "404 job not found" in fails("job", "get", "ep1", "job1") client.get.side_effect = requests.HTTPError("no response") assert "no response" in fails("job", "get", "ep1", "job1") + client.get.side_effect = requests.ReadTimeout("status request timed out") + assert "status request timed out" in fails("job", "wait", "ep1", "job1") def no_key(): raise RuntimeError("api key not set") @@ -1132,9 +1014,6 @@ def no_key(): assert "api key not set" in fails("job", "get", "ep1", "job1") -# -------------------------------------------------------------- completion - - def test_completion_inserts_ids(api): ctx = click.Context(ep.list_workers, info_name="workers") items = ep.complete_endpoint(ctx, None, "ep") From 542813ee99eae127fbe71ae1dc850ce20061c929 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:03:12 -0700 Subject: [PATCH 5/7] style: remove redundant endpoint docstrings --- runpod/rp_cli/endpoint.py | 13 ------------- runpod/rp_cli/endpoint_config.py | 6 ------ tests/test_cli/test_endpoint_command.py | 3 --- 3 files changed, 22 deletions(-) diff --git a/runpod/rp_cli/endpoint.py b/runpod/rp_cli/endpoint.py index ae42274d..3ff6f4e1 100644 --- a/runpod/rp_cli/endpoint.py +++ b/runpod/rp_cli/endpoint.py @@ -1,10 +1,3 @@ -"""`rp endpoint`: create, inspect, scale, observe, and call serverless endpoints. - -configuration goes through the rest v2 control plane (/v2/serverless). -jobs go through the serverless data plane (runpod.endpoint_url_base), -which has no rest v2 equivalent. -""" - import contextlib import fnmatch import json @@ -67,7 +60,6 @@ def fetch(): def complete_worker(ctx, param, incomplete): - """worker ids of the endpoint named earlier on the command line.""" endpoint_id = ctx.params.get("endpoint_id") if not endpoint_id: return [] @@ -236,7 +228,6 @@ def _scale_options(func): def _config_options(creating: bool): - """the flags shared by create and update, one per api field.""" head = [click.option("--name", required=creating, help="Endpoint name.")] if creating: head.append( @@ -363,7 +354,6 @@ def _follow_rollout( timeout: float, interval: float = 2.0, ) -> None: - """poll releases until the new version runs on every worker.""" started = time.monotonic() last = None try: @@ -424,7 +414,6 @@ def _apply( def _edit_document(before: dict) -> Optional[dict]: - """open the document in $EDITOR until it validates or the user gives up.""" text = json.dumps(before, indent=2) + "\n" while True: edited = click.edit(text, extension=".json") @@ -492,7 +481,6 @@ def _workers_view(data: dict): @contextlib.contextmanager def _data_plane_errors() -> Iterator[None]: - """surface serverless data-plane failures as clean cli errors.""" try: yield except requests.HTTPError as exc: @@ -881,7 +869,6 @@ def list_workers(endpoint_id, watch, interval, as_json): def _release_changes(release: dict) -> list[str]: - """changed field names, skipping entries whose old and new values are both empty.""" return [ entry["field"].removeprefix("template.") for entry in release.get("diff") or [] diff --git a/runpod/rp_cli/endpoint_config.py b/runpod/rp_cli/endpoint_config.py index 8f37d45b..5fa1f40e 100644 --- a/runpod/rp_cli/endpoint_config.py +++ b/runpod/rp_cli/endpoint_config.py @@ -1,5 +1,3 @@ -"""endpoint request configuration and editable JSON documents.""" - import json import shlex from typing import Any, Optional @@ -142,7 +140,6 @@ def scale_patch(options: dict, current: dict) -> dict: def config_body(options: dict, current: Optional[dict] = None) -> dict: - """map flags to a create or partial update, preserving untouched configuration.""" updating = current is not None current = current if updating else {"type": options["endpoint_type"]} body = {} @@ -215,7 +212,6 @@ def config_body(options: dict, current: Optional[dict] = None) -> dict: def mutable(endpoint: dict) -> dict: - """the endpoint's configuration as UpdateEndpointRequest fields.""" doc: dict[str, Any] = {} for key in ("name", "image", "args", "env", "disk", "ports"): if endpoint.get(key) is not None: @@ -256,7 +252,6 @@ def _kind(kind: type) -> str: def validate(text: str, before: dict) -> dict: - """parse an edited document, raising ClickException on any problem.""" try: doc = json.loads(text) except json.JSONDecodeError as exc: @@ -295,7 +290,6 @@ def validate(text: str, before: dict) -> dict: def patch(before: dict, after: dict) -> dict: - """top-level fields that changed, with gpu narrowed to its changed keys.""" patch = {} for key, value in after.items(): if before.get(key) == value: diff --git a/tests/test_cli/test_endpoint_command.py b/tests/test_cli/test_endpoint_command.py index e966ac86..91e13312 100644 --- a/tests/test_cli/test_endpoint_command.py +++ b/tests/test_cli/test_endpoint_command.py @@ -1,5 +1,3 @@ -"""tests for rp endpoint.""" - import copy import json from unittest.mock import MagicMock @@ -170,7 +168,6 @@ class FakeApi: - """routes rest calls by (method, path) and records them.""" def __init__(self, routes): self.routes = routes From b6b908d6ccbfc74f418d25517fcfc0e4f0b3f47d Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 20 Oct 2026 10:54:32 -0700 Subject: [PATCH 6/7] fix: harden endpoint configuration and job following --- runpod/rp_cli/endpoint.py | 70 +++++++++---- runpod/rp_cli/endpoint_config.py | 42 +++----- tests/test_cli/test_endpoint_command.py | 134 +++++++++++++++++++++++- 3 files changed, 195 insertions(+), 51 deletions(-) diff --git a/runpod/rp_cli/endpoint.py b/runpod/rp_cli/endpoint.py index 3ff6f4e1..62766aff 100644 --- a/runpod/rp_cli/endpoint.py +++ b/runpod/rp_cli/endpoint.py @@ -1,6 +1,7 @@ import contextlib import fnmatch import json +import sys import time from typing import Any, Iterator, Optional from urllib.parse import quote @@ -320,7 +321,7 @@ def _config_options(creating: bool): click.option("--cmd", help="Container CMD, shell-quoted."), click.option( "--registry", - help="Registry credential ID for private images.", + help="Registry credential name or ID for private images; empty clears it.", shell_complete=complete_registry, ), ] @@ -393,14 +394,17 @@ def _follow_rollout( def _apply( endpoint_id: str, patch: dict, follow: bool, timeout: float, as_json: bool ) -> None: + before = None + if follow: + try: + with api_errors(): + before = run_rest_request( + "GET", _path(endpoint_id, "releases"), params={"limit": 1} + ).get("endpointVersion") + except click.ClickException as exc: + click.echo(f"warning: cannot follow rollout: {exc.message}", err=True) + follow = False with api_errors(): - before = ( - run_rest_request( - "GET", _path(endpoint_id, "releases"), params={"limit": 1} - ).get("endpointVersion") - if follow - else None - ) updated = run_rest_request("PATCH", _path(endpoint_id), json=patch) if as_json: print_json(updated) @@ -410,7 +414,12 @@ def _apply( f"[dim]{endpoint_id}[/dim]" ) if follow: - _follow_rollout(endpoint_id, before, timeout) + with ( + contextlib.redirect_stdout(sys.stderr) + if as_json + else contextlib.nullcontext() + ): + _follow_rollout(endpoint_id, before, timeout) def _edit_document(before: dict) -> Optional[dict]: @@ -515,7 +524,20 @@ def _follow_job(client, endpoint_id, job_id, as_json, *, stream=False, timeout=N """wait for completion, draining streamed chunks before reading the final result.""" action = "stream" if stream else "status" resume = "stream" if stream else "wait" - deadline = time.monotonic() + timeout if timeout else None + deadline = time.monotonic() + timeout if timeout is not None else None + status = "pending" + + def remaining_time(limit): + if deadline is None: + return limit + remaining = deadline - time.monotonic() + if remaining <= 0: + raise click.ClickException( + f"job {job_id} is still {status} after {timeout:g}s. " + f"resume with: rp endpoint job {resume} {endpoint_id} {job_id}" + ) + return min(limit, remaining) + streamed = open_line = False progress = ( contextlib.nullcontext() @@ -528,7 +550,13 @@ def _follow_job(client, endpoint_id, job_id, as_json, *, stream=False, timeout=N try: if stream: job = client.get( - _job_path(endpoint_id, action, job_id), timeout=60 + _job_path(endpoint_id, action, job_id), + timeout=remaining_time(60), + ) + elif deadline is not None: + job = client.get( + _job_path(endpoint_id, action, job_id), + timeout=remaining_time(10), ) else: job = client.get(_job_path(endpoint_id, action, job_id)) @@ -536,7 +564,8 @@ def _follow_job(client, endpoint_id, job_id, as_json, *, stream=False, timeout=N if stream: continue raise - chunks = job["stream"] if stream else [] + status = job["status"] + chunks = (job.get("stream") or []) if stream else [] streamed = streamed or bool(chunks) for chunk in chunks: output = chunk["output"] @@ -550,13 +579,8 @@ def _follow_job(client, endpoint_id, job_id, as_json, *, stream=False, timeout=N click.echo(json.dumps(output)) if is_completed(job["status"]) and not chunks: break - if deadline and time.monotonic() >= deadline: - raise click.ClickException( - f"job {job_id} is still {job['status']} after {timeout:g}s. " - f"resume with: rp endpoint job {resume} {endpoint_id} {job_id}" - ) if not chunks: - time.sleep(0.5 if stream else 1) + time.sleep(remaining_time(0.5 if stream else 1)) except KeyboardInterrupt: click.echo( f"\njob {job_id} keeps running. resume with: " @@ -569,7 +593,8 @@ def _follow_job(client, endpoint_id, job_id, as_json, *, stream=False, timeout=N if streamed and job["status"] == "COMPLETED": return if stream: - job = client.get(_job_path(endpoint_id, "status", job_id)) + kwargs = {"timeout": remaining_time(10)} if deadline is not None else {} + job = client.get(_job_path(endpoint_id, "status", job_id), **kwargs) _print_job(job, as_json) @@ -610,7 +635,10 @@ def _call_load_balancer( try: payload = response.json() except ValueError: - click.echo(response.text) + if as_json: + print_json(response.text) + else: + click.echo(response.text) return if isinstance(payload, str) and not as_json: click.echo(payload) @@ -1064,7 +1092,7 @@ def run_endpoint( if stream: if not as_json: click.echo(f"job {job_id}", err=True) - _follow_job(client, endpoint_id, job_id, as_json, stream=True) + _follow_job(client, endpoint_id, job_id, as_json, stream=True, timeout=timeout) return _follow_job(client, endpoint_id, job_id, as_json, timeout=timeout) diff --git a/runpod/rp_cli/endpoint_config.py b/runpod/rp_cli/endpoint_config.py index 5fa1f40e..5a775747 100644 --- a/runpod/rp_cli/endpoint_config.py +++ b/runpod/rp_cli/endpoint_config.py @@ -1,11 +1,10 @@ import json -import shlex from typing import Any, Optional import click from runpod.rp_cli import compute -from runpod.rp_cli.options import parse_env +from runpod.rp_cli.options import container_fields FLASHBOOT = {"off": "OFF", "standard": "FLASHBOOT", "priority": "PRIORITY_FLASHBOOT"} @@ -121,7 +120,7 @@ def scale_patch(options: dict, current: dict) -> dict: if options.get(flag) is not None } if changes: - patch["workers"] = {**current.get("workers", {}), **changes} + patch["workers"] = {**(current.get("workers") or {}), **changes} delay = options.get("queue_delay_seconds") count = options.get("requests_per_worker") if delay is not None and count is not None: @@ -142,37 +141,20 @@ def scale_patch(options: dict, current: dict) -> dict: def config_body(options: dict, current: Optional[dict] = None) -> dict: updating = current is not None current = current if updating else {"type": options["endpoint_type"]} - body = {} + body = container_fields( + {**options, "ports": options.get("port")}, + current, + options.get("unset_env", ()), + ) if not updating: body["type"] = current["type"] - for key, flag in ( - ("name", "name"), - ("image", "image"), - ("templateId", "template"), - ("args", "args"), - ("disk", "disk"), - ("registry", "registry"), - ): - if options.get(flag) is not None: - body[key] = options[flag] - for key in ("entrypoint", "cmd"): - if options.get(key) is not None: - body[key] = shlex.split(options[key]) - if options.get("port"): - body["ports"] = list(options["port"]) + if options.get("name") is not None: + body["name"] = options["name"] if options.get("flashboot"): body["flashboot"] = FLASHBOOT[options["flashboot"].lower()] if options.get("request_timeout_seconds") is not None: body["timeout"] = options["request_timeout_seconds"] * 1000 - if options.get("env") or options.get("unset_env"): - env = {**current.get("env", {}), **parse_env(options.get("env", ()))} - for key in options.get("unset_env", ()): - if key not in env: - raise click.UsageError(f"--unset-env {key}: no such variable") - del env[key] - body["env"] = env - if options.get("gpu") and options.get("cpu"): raise click.UsageError("use either --gpu or --cpu") if current.get("cpu") and options.get("gpu"): @@ -203,7 +185,7 @@ def config_body(options: dict, current: Optional[dict] = None) -> dict: places, compute.list_datacenters() ) if options.get("network_volume") or options.get("global_volume"): - mounts = dict(current.get("mounts", {})) + mounts = dict(current.get("mounts") or {}) for kind, flag in (("network", "network_volume"), ("global", "global_volume")): if options.get(flag): mounts[kind] = [{"volumeId": volume} for volume in options[flag]] @@ -301,7 +283,9 @@ def patch(before: dict, after: dict) -> dict: gpu = {k: v for k, v in value.items() if old.get(k) != v} if "pools" in gpu or "excludedTypes" in gpu: gpu["pools"] = value.get("pools", old.get("pools")) - gpu["excludedTypes"] = value.get("excludedTypes") or [] + gpu["excludedTypes"] = ( + value.get("excludedTypes", old.get("excludedTypes")) or [] + ) if "minCudaVersion" in gpu and gpu["minCudaVersion"] is None: gpu["minCudaVersion"] = "" patch["gpu"] = gpu diff --git a/tests/test_cli/test_endpoint_command.py b/tests/test_cli/test_endpoint_command.py index 91e13312..c9d8109f 100644 --- a/tests/test_cli/test_endpoint_command.py +++ b/tests/test_cli/test_endpoint_command.py @@ -215,6 +215,10 @@ def api(monkeypatch): ("GET", "/v2/serverless/lb1"): LB_ENDPOINT, ("GET", "/v2/serverless/ep1/workers"): WORKERS, ("GET", "/v2/serverless/ep1/releases"): RELEASES, + ("GET", "/v2/registries"): { + "registries": [{"id": "reg1", "name": "ghcr"}], + "pagination": PAGE, + }, ("POST", "/v2/serverless"): lambda params, body: dict(body, id="new1"), ("PATCH", "/v2/serverless/ep1"): lambda params, body: dict( GPU_ENDPOINT, **body @@ -503,7 +507,6 @@ def test_update_cpu_vcpus_and_mounts(api): "args, message", [ (["ep1"], "nothing to update"), - (["ep1", "--unset-env", "MISSING"], "no such variable"), (["ep1", "--cpu", "cpu3c:4"], "cannot switch to CPUs"), (["ep2", "--gpu", "h100"], "cannot switch to GPUs"), (["ep2", "--cuda", "12.4"], "apply to GPU"), @@ -536,6 +539,91 @@ def test_update_json_and_wait_follows_rollout(api): assert json.loads(out)["name"] == "renamed" +def test_wait_json_keeps_rollout_messages_on_stderr(api): + states = iter( + [ + {"endpointVersion": 3}, + {"endpointVersion": 4, "rollout": {"inProgress": True}}, + {"endpointVersion": 4, "rollout": {"inProgress": False}}, + ] + ) + api.routes[("GET", "/v2/serverless/ep1/releases")] = lambda *_: next(states) + result = invoke("update", "ep1", "--image", "img:2", "--wait", "--json") + assert result.exit_code == 0, result.output + assert json.loads(result.stdout)["image"] == "img:2" + assert "rolling out" in result.stderr + assert "release v4 is live" in result.stderr + + +@pytest.mark.parametrize( + "failure", + [ + error.QueryError("releases unavailable", "GET"), + requests.ConnectionError("offline"), + ], +) +def test_release_baseline_failure_does_not_block_update(api, failure): + api.routes[("GET", "/v2/serverless/ep1/releases")] = failure + result = invoke("update", "ep1", "--image", "img:2", "--wait", "--json") + assert result.exit_code == 0, result.output + assert api.bodies("PATCH") == [{"image": "img:2"}] + assert json.loads(result.stdout)["image"] == "img:2" + assert "cannot follow rollout" in result.stderr + + +@pytest.mark.parametrize( + "value, expected", [("ghcr", "reg1"), ("reg1", "reg1"), ("", None)] +) +def test_registry_selection_and_clearing(api, value, expected): + ok("update", "ep1", "--registry", value) + assert api.bodies("PATCH") == [{"registry": expected}] + + +def test_unset_missing_environment_variable_is_noop(api): + ok("update", "ep1", "--unset-env", "MISSING") + assert api.bodies("PATCH") == [{"env": GPU_ENDPOINT["env"]}] + + +@pytest.mark.parametrize("option", ["--cmd", "--entrypoint"]) +def test_invalid_container_quoting_is_usage_error(api, option): + result = invoke("update", "ep1", option, "'unterminated") + assert result.exit_code == 2 + assert option in result.output + assert not api.bodies("PATCH") + + +@pytest.mark.parametrize( + "key, args, expected", + [ + ("env", ["update", "ep1", "--env", "A=1"], {"env": {"A": "1"}}), + ("workers", ["scale", "ep1", "--max-workers", "2"], {"workers": {"max": 2}}), + ( + "mounts", + ["update", "ep1", "--network-volume", "vol1"], + {"mounts": {"network": [{"volumeId": "vol1"}]}}, + ), + ], +) +def test_nullable_configuration_maps(api, key, args, expected): + api.routes[("GET", "/v2/serverless/ep1")] = {**GPU_ENDPOINT, key: None} + ok(*args) + assert api.bodies("PATCH") == [expected] + + +@pytest.mark.parametrize("excluded", [None, []]) +def test_partial_gpu_edit_preserves_omitted_exclusions(excluded): + before = {"gpu": {"pools": ["ADA_80_PRO"], "excludedTypes": ["excluded"]}} + edited = {"pools": ["ADA_24"]} + if excluded is not None: + edited["excludedTypes"] = excluded + assert config.patch(before, {"gpu": edited}) == { + "gpu": { + "pools": ["ADA_24"], + "excludedTypes": ["excluded"] if excluded is None else [], + } + } + + def _doc(**changes): doc = config.mutable(GPU_ENDPOINT) doc.update(changes) @@ -848,6 +936,42 @@ def test_run_timeout_and_interrupt(api, client, monkeypatch): assert result.exit_code != 0 +@pytest.mark.parametrize("retry", [False, True]) +def test_run_stream_timeout_bounds_every_request(api, client, monkeypatch, retry): + clock = [0.0] + monkeypatch.setattr(ep.time, "monotonic", lambda: clock[0]) + client.post.return_value = {"id": "job1", "status": "IN_QUEUE"} + + def poll(path, timeout): + assert path == "ep1/stream/job1" + assert 0 < timeout <= 5 - clock[0] + clock[0] += min(2.0, timeout) + if retry: + raise requests.ReadTimeout("still waiting") + return {"status": "IN_QUEUE"} + + client.get.side_effect = poll + result = invoke("run", "ep1", "--data", "{}", "--stream", "--timeout", "5") + assert result.exit_code == 1 + assert "after 5s" in result.stderr + assert "resume with: rp endpoint job stream ep1 job1" in result.stderr + assert client.get.call_count == 3 + assert clock[0] == 5 + + +@pytest.mark.parametrize("empty_stream", [{}, {"stream": None}]) +def test_stream_responses_can_omit_chunks(api, client, empty_stream): + client.post.return_value = {"id": "job1", "status": "IN_QUEUE"} + client.get.side_effect = [ + {"status": "IN_QUEUE", **empty_stream}, + {"status": "COMPLETED", **empty_stream}, + {"id": "job1", "status": "COMPLETED", "output": "done"}, + ] + result = invoke("run", "ep1", "--data", "{}", "--stream", "--json") + assert result.exit_code == 0, result.output + assert json.loads(result.stdout)["output"] == "done" + + def test_run_stream(api, client): client.post.return_value = {"id": "job1"} client.get.side_effect = [ @@ -898,6 +1022,14 @@ def test_run_rejects_bad_flags(api, client, args, message): client.post.assert_not_called() +def test_load_balancer_plain_text_json(api, monkeypatch): + response = MagicMock(status_code=200, text="pong") + response.json.side_effect = ValueError("not json") + monkeypatch.setattr(ep.requests, "request", lambda *a, **k: response) + monkeypatch.setattr(ep, "api_key", lambda: "test-key") + assert json.loads(ok("run", "lb1", "--json")) == "pong" + + def test_run_load_balancer(api, monkeypatch): calls = [] From 44409c6108b6c3381b2c3e9d06d37adcc638a830 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Wed, 21 Oct 2026 12:35:46 -0700 Subject: [PATCH 7/7] fix: handle endpoint rollout and job wait deadlines --- runpod/rp_cli/endpoint.py | 37 ++++--- tests/test_cli/test_endpoint_command.py | 126 ++++++++++++++++++++++++ 2 files changed, 152 insertions(+), 11 deletions(-) diff --git a/runpod/rp_cli/endpoint.py b/runpod/rp_cli/endpoint.py index 62766aff..ca6fdbdf 100644 --- a/runpod/rp_cli/endpoint.py +++ b/runpod/rp_cli/endpoint.py @@ -365,11 +365,10 @@ def _follow_rollout( ) version = state.get("endpointVersion") elapsed = time.monotonic() - started - if previous_version is not None and version == previous_version: - if elapsed >= min(10.0, timeout): - ui.console.print("[dim]no new release to roll out[/dim]") - return - else: + waiting_for_release = ( + previous_version is not None and version == previous_version + ) + if not waiting_for_release: line = _rollout_line(state) if line != last: ui.console.print(f"[dim]{line}[/dim]") @@ -378,8 +377,13 @@ def _follow_rollout( ui.success(f"release v{version} is live") return if elapsed >= timeout: + phase = ( + "no new release observed" + if waiting_for_release + else "rollout still in progress" + ) ui.warn( - f"rollout still in progress after {timeout:g}s. " + f"{phase} after {timeout:g}s. " f"follow with: rp endpoint releases {endpoint_id}" ) return @@ -528,7 +532,7 @@ def _follow_job(client, endpoint_id, job_id, as_json, *, stream=False, timeout=N status = "pending" def remaining_time(limit): - if deadline is None: + if deadline is None or is_completed(status): return limit remaining = deadline - time.monotonic() if remaining <= 0: @@ -560,8 +564,13 @@ def remaining_time(limit): ) else: job = client.get(_job_path(endpoint_id, action, job_id)) - except requests.ReadTimeout: - if stream: + except requests.Timeout as exc: + remaining_time(0) + if ( + stream + and isinstance(exc, requests.ReadTimeout) + and not is_completed(status) + ): continue raise status = job["status"] @@ -1009,7 +1018,11 @@ def logs(endpoint_id, worker_ids, follow, source, tail, since, as_json): @click.option( "--low-priority", is_flag=True, help="Do not scale up workers for this job." ) -@click.option("--timeout", type=float, help="Seconds to wait for the result.") +@click.option( + "--timeout", + type=click.FloatRange(min=0, min_open=True), + help="Seconds to wait for the result.", +) @click.option("--path", "lb_path", help="Load-balancer endpoints: request path.") @click.option( "--method", @@ -1142,7 +1155,9 @@ def job_get(endpoint_id, job_id, as_json): @_job_command("wait") -@click.option("--timeout", type=float, help="Seconds to wait.") +@click.option( + "--timeout", type=click.FloatRange(min=0, min_open=True), help="Seconds to wait." +) @_data_plane_errors() def job_wait(endpoint_id, job_id, timeout, as_json): """Wait for a job to finish and print its output.""" diff --git a/tests/test_cli/test_endpoint_command.py b/tests/test_cli/test_endpoint_command.py index c9d8109f..a05f2baa 100644 --- a/tests/test_cli/test_endpoint_command.py +++ b/tests/test_cli/test_endpoint_command.py @@ -746,6 +746,45 @@ def test_follow_rollout_without_new_release(api): ep._follow_rollout("ep1", 3, timeout=0) +def test_follow_rollout_waits_for_delayed_release(api, monkeypatch, capsys): + clock = [0.0] + monkeypatch.setattr(ep.time, "monotonic", lambda: clock[0]) + monkeypatch.setattr( + ep.time, "sleep", lambda seconds: clock.__setitem__(0, clock[0] + seconds) + ) + api.routes[("GET", "/v2/serverless/ep1/releases")] = lambda *_: { + "endpointVersion": 3 if clock[0] < 12 else 4, + "rollout": {"inProgress": False}, + } + + ep._follow_rollout("ep1", 3, timeout=30) + + assert clock[0] == 12 + assert "release v4 is live" in capsys.readouterr().out + + +def test_follow_rollout_reports_unobserved_release_at_deadline( + api, monkeypatch, capsys +): + clock = [0.0] + monkeypatch.setattr(ep.time, "monotonic", lambda: clock[0]) + monkeypatch.setattr( + ep.time, "sleep", lambda seconds: clock.__setitem__(0, clock[0] + seconds) + ) + api.routes[("GET", "/v2/serverless/ep1/releases")] = { + "endpointVersion": 3, + "rollout": {"inProgress": False}, + } + + ep._follow_rollout("ep1", 3, timeout=20) + + assert clock[0] == 20 + output = capsys.readouterr().out + assert "no new release observed after 20s" in output + assert "follow with: rp endpoint releases ep1" in output + assert "no new release to roll out" not in output + + def test_follow_rollout_times_out_and_detaches(api, monkeypatch, capsys): state = {"endpointVersion": 4, "rollout": {"inProgress": True}} api.routes[("GET", "/v2/serverless/ep1/releases")] = state @@ -972,6 +1011,93 @@ def test_stream_responses_can_omit_chunks(api, client, empty_stream): assert json.loads(result.stdout)["output"] == "done" +@pytest.mark.parametrize( + "status, chunks", + [("COMPLETED", [{"output": "done"}]), ("COMPLETED", []), ("FAILED", [])], +) +def test_stream_terminal_status_can_drain_after_deadline( + api, client, monkeypatch, status, chunks +): + clock = [0.0] + monkeypatch.setattr(ep.time, "monotonic", lambda: clock[0]) + client.post.return_value = {"id": "job1"} + responses = iter( + [ + {"status": status, "stream": chunks}, + {"status": status, "stream": []} + if chunks + else { + "id": "job1", + "status": status, + "output": "done", + "error": "worker failed", + }, + ] + ) + + def poll(path, timeout): + assert timeout > 0 + clock[0] = 5.1 + return next(responses) + + client.get.side_effect = poll + result = invoke("run", "ep1", "--data", "{}", "--stream", "--timeout", "5") + + assert result.exit_code == (1 if status == "FAILED" else 0), result.output + assert "still COMPLETED" not in result.output + assert "still FAILED" not in result.output + expected = "failed: worker failed" if status == "FAILED" else "done" + assert expected in result.output + assert client.get.call_count == 2 + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + "timeout_error", [requests.ConnectTimeout, requests.ReadTimeout] +) +def test_request_timeout_at_deadline_includes_resume_hint( + api, client, monkeypatch, stream, timeout_error +): + clock = [0.0] + monkeypatch.setattr(ep.time, "monotonic", lambda: clock[0]) + client.post.return_value = {"id": "job1"} + + def poll(path, timeout): + clock[0] += timeout + raise timeout_error("connection timed out") + + client.get.side_effect = poll + result = invoke( + "run", "ep1", "--data", "{}", "--timeout", "5", + *(["--stream"] if stream else []), + ) + + assert result.exit_code == 1 + assert "after 5s" in result.stderr + action = "stream" if stream else "wait" + assert f"resume with: rp endpoint job {action} ep1 job1" in result.stderr + assert "connection timed out" not in result.stderr + + +@pytest.mark.parametrize("timeout", ["0", "-1"]) +@pytest.mark.parametrize( + "args", + [ + ["run", "ep1", "--data", "{}"], + ["run", "lb1", "--data", "{}"], + ["job", "wait", "ep1", "job1"], + ], +) +def test_job_timeouts_must_be_positive_before_requests(api, client, timeout, args): + result = invoke(*args, "--timeout", timeout) + + assert result.exit_code == 2 + assert "Invalid value for '--timeout'" in result.stderr + assert not api.calls + client.post.assert_not_called() + client.get.assert_not_called() + + def test_run_stream(api, client): client.post.return_value = {"id": "job1"} client.get.side_effect = [