From 2759ad21068cd3bdd3980cd2ebbc14b0147e8319 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:28:56 -0700 Subject: [PATCH 01/25] feat: add rp catalog for gpu and cpu pricing and stock --- runpod/rp_cli/catalog.py | 277 +++++++++++++++++++++++++++++++++++++ runpod/rp_cli/compute.py | 287 +++++++++++++++++++++++++++++++++++++++ runpod/rp_cli/main.py | 22 ++- runpod/rp_cli/output.py | 120 ++++++++++++++++ 4 files changed, 704 insertions(+), 2 deletions(-) create mode 100644 runpod/rp_cli/catalog.py create mode 100644 runpod/rp_cli/compute.py create mode 100644 runpod/rp_cli/output.py diff --git a/runpod/rp_cli/catalog.py b/runpod/rp_cli/catalog.py new file mode 100644 index 00000000..14e02e62 --- /dev/null +++ b/runpod/rp_cli/catalog.py @@ -0,0 +1,277 @@ +"""`rp catalog`: gpu and cpu pricing and stock search.""" + +from typing import Optional + +import click + +from runpod.rp_cli import compute +from runpod.rp_cli.output import api_errors, json_option, print_json, print_table + + +def _price(value: Optional[float], count: int) -> Optional[float]: + return round(value * count, 4) if value else None + + +def _money(value: Optional[float]) -> str: + return f"${value:.2f}" if value else "-" + + +def _stock(entries: list[dict]) -> str: + ranked = sorted( + entries, key=lambda dc: -compute.availability_rank(dc.get("availability")) + ) + return ", ".join( + f"{dc['id']} {(dc.get('availability') or 'NONE').lower()}" for dc in ranked + ) + + +def _in_regions(record: dict, allowed: Optional[set[str]]) -> dict: + """the record with stock narrowed to the allowed data centers.""" + if allowed is None: + return dict(record) + datacenters = [ + dc for dc in record.get("dataCenters") or [] if dc["id"] in allowed + ] + return dict( + record, + dataCenters=datacenters, + availability=compute.best_availability( + dc.get("availability") for dc in datacenters + ), + ) + + +def _allowed(regions: tuple[str, ...]) -> Optional[set[str]]: + if not regions: + return None + return set(compute.resolve_regions(regions, compute.list_datacenters())) + + +def _unit_price(gpu: dict, product: str, cloud: Optional[str]) -> Optional[float]: + price = gpu.get("price") or {} + if product == "SERVERLESS": + return price.get("serverless") or None + if cloud: + return price.get(cloud) or None + offered = [value for value in (price.get("secure"), price.get("community")) if value] + return min(offered, default=None) + + +def _gpu_rows( + query, product, count, regions, cloud, cuda, min_vram, max_price, available +) -> list[dict]: + min_cuda, exact_cuda = compute.parse_cuda(cuda) + gpus = compute.list_gpus( + product=product, + count=count, + cloud=cloud.upper() if cloud else None, + cuda_versions=exact_cuda, + min_cuda_version=min_cuda, + ) + if query: + gpus = compute.resolve_gpus(query, gpus) + if product == "SERVERLESS": + gpus = [gpu for gpu in gpus if gpu.get("pool")] + allowed = _allowed(regions) + rows = [] + for gpu in gpus: + row = _in_regions(gpu, allowed) + unit = _unit_price(row, product, cloud) + if min_vram and (row.get("memory") or 0) < min_vram: + continue + if max_price is not None and (unit is None or unit > max_price): + continue + if available and row.get("availability") in (None, "NONE"): + continue + price = row.get("price") or {} + row["count"] = count + row["totalPrice"] = { + tier: _price(value, count) for tier, value in price.items() + } + rows.append(row) + rows.sort( + key=lambda row: ( + -compute.availability_rank(row.get("availability")), + _unit_price(row, product, cloud) or float("inf"), + row.get("name") or row["id"], + ) + ) + return rows + + +def _gpu_columns(product: str, count: int, cloud: Optional[str]): + suffix = f" x{count}" if count > 1 else "" + columns = [ + ("GPU", lambda row: row.get("name") or row["id"]), + ("VRAM", lambda row: f"{row['memory']}GB" if row.get("memory") else None), + ] + if product == "SERVERLESS": + columns += [ + ("POOL", lambda row: row.get("pool")), + (f"$/HR{suffix}", lambda row: _money(row["totalPrice"].get("serverless"))), + ] + else: + tiers = [cloud] if cloud else ["secure", "community"] + for tier in tiers: + columns.append( + ( + f"{tier.upper()} $/HR{suffix}", + lambda row, tier=tier: _money(row["totalPrice"].get(tier)), + ) + ) + columns.append( + ("MAX", lambda row: max((row.get("maxCount") or {}).values(), default=None)) + ) + columns += [ + ("STOCK", lambda row: (row.get("availability") or "NONE").lower()), + ("DATA CENTERS", lambda row: _stock(row.get("dataCenters") or [])), + ] + return columns + + +def _cpu_rows(query, product, vcpu, regions, available) -> list[dict]: + cpus = compute.list_cpus(product=product, vcpu_count=vcpu) + if query: + cpus = compute.match_cpus(query, cpus) + if not cpus: + raise compute.ResolutionError(f"unknown cpu '{query}'") + allowed = _allowed(regions) + rows = [_in_regions(cpu, allowed) for cpu in cpus] + if available: + rows = [row for row in rows if row.get("availability") not in (None, "NONE")] + rows.sort(key=lambda row: -compute.availability_rank(row.get("availability"))) + return rows + + +def _cpu_columns(product: str): + tier = "serverlessPerVcpu" if product == "SERVERLESS" else "securePerVcpu" + return [ + ("CPU", lambda row: row["id"]), + ("NAME", lambda row: row.get("name")), + ("GROUP", lambda row: row.get("group")), + ( + "VCPU", + lambda row: "{min}-{max}".format(**row["vcpu"]) if row.get("vcpu") else None, + ), + ("RAM/VCPU", lambda row: f"{row['ramGbPerVcpu']:g}GB" if row.get("ramGbPerVcpu") else None), + ("$/VCPU-HR", lambda row: _money((row.get("price") or {}).get(tier))), + ("STOCK", lambda row: (row.get("availability") or "NONE").lower()), + ("DATA CENTERS", lambda row: _stock(row.get("dataCenters") or [])), + ] + + +_DATACENTER_COLUMNS = [ + ("DATA CENTER", lambda row: row["id"]), + ("REGION", lambda row: (row.get("region") or "").lower().replace("_", " ")), + ( + "VOLUME TYPES", + lambda row: [kind.lower() for kind in row.get("networkVolumeTypes") or []], + ), + ("GLOBAL NETWORK", lambda row: row.get("globalNetwork")), + ("COMPLIANCE", lambda row: row.get("compliance")), +] + + +@click.command() +@click.argument("query", required=False) +@click.option( + "--pod", + "product", + flag_value="POD", + default=True, + help="Pod pricing and stock (default).", +) +@click.option( + "--serverless", + "product", + flag_value="SERVERLESS", + help="Serverless pricing, pools, and stock.", +) +@click.option("--cpu", "cpu_mode", is_flag=True, help="Search CPU flavors instead of GPUs.") +@click.option("--datacenters", is_flag=True, help="List data centers and their regions.") +@click.option( + "--count", + type=click.IntRange(min=1), + default=1, + show_default=True, + help="GPUs per pod or worker; prices and stock account for it.", +) +@click.option( + "--vcpu", type=click.IntRange(min=2), default=None, help="vCPUs per pod or worker (with --cpu)." +) +@click.option( + "--region", + "regions", + multiple=True, + help="Region (eu, na, asia), id prefix (us, eu-ro), or data center id. Repeatable.", +) +@click.option( + "--cloud", + type=click.Choice(["secure", "community"], case_sensitive=False), + default=None, + help="Pod cloud tier.", +) +@click.option( + "--cuda", + multiple=True, + help="CUDA filter: '>=12.4' for a floor, 12.4 or 12.4,12.8 for exact versions.", +) +@click.option("--min-vram", type=int, default=None, help="Minimum VRAM per GPU in GB.") +@click.option( + "--max-price", type=float, default=None, help="Maximum USD per GPU-hour." +) +@click.option("--available", is_flag=True, help="Only show hardware that is in stock.") +@json_option +def catalog( + query, + product, + cpu_mode, + datacenters, + count, + vcpu, + regions, + cloud, + cuda, + min_vram, + max_price, + available, + as_json, +): # pylint: disable=too-many-arguments + """Search GPU and CPU pricing and stock by data center. + + QUERY is a human-readable GPU (4090, h100, a100-sxm4), a serverless + pool id, or with --cpu a CPU flavor. + + \b + rp catalog h100 --count 2 --region eu --cuda '>=12.4' + rp catalog 4090 --serverless + rp catalog --cpu --vcpu 8 --available + rp catalog --datacenters --region us + """ + if cpu_mode and datacenters: + raise click.UsageError("use either --cpu or --datacenters") + if cloud and product == "SERVERLESS": + raise click.UsageError("--cloud applies to pods only") + cloud = cloud.lower() if cloud else None + + with api_errors(): + if datacenters: + values = regions + ((query,) if query else ()) + rows = compute.list_datacenters() + if values: + allowed = set(compute.resolve_regions(values, rows)) + rows = [row for row in rows if row["id"] in allowed] + columns, empty = _DATACENTER_COLUMNS, "no data centers match" + elif cpu_mode: + rows = _cpu_rows(query, product, vcpu, regions, available) + columns, empty = _cpu_columns(product), "no cpu flavors match" + else: + rows = _gpu_rows( + query, product, count, regions, cloud, cuda, min_vram, max_price, available + ) + columns, empty = _gpu_columns(product, count, cloud), "no gpus match" + + if as_json: + print_json(rows) + return + print_table(rows, columns, empty) diff --git a/runpod/rp_cli/compute.py b/runpod/rp_cli/compute.py new file mode 100644 index 00000000..c072a07b --- /dev/null +++ b/runpod/rp_cli/compute.py @@ -0,0 +1,287 @@ +"""gpu, cpu, and region resolution against the rest v2 catalog.""" + +import re +from typing import Any, Iterable, Optional + +from runpod.api.rest import run_rest_request + +AVAILABILITY_RANK = {"NONE": 0, "LOW": 1, "MEDIUM": 2, "HIGH": 3} + +REGION_ALIASES = { + "na": "NORTH_AMERICA", + "northamerica": "NORTH_AMERICA", + "sa": "SOUTH_AMERICA", + "southamerica": "SOUTH_AMERICA", + "eu": "EUROPE", + "europe": "EUROPE", + "asia": "ASIA", + "apac": "ASIA", + "me": "MIDDLE_EAST", + "middleeast": "MIDDLE_EAST", + "africa": "AFRICA", + "oceania": "OCEANIA", +} + +_NOISE_TOKENS = {"nvidia", "geforce", "amd", "instinct"} + + +class ResolutionError(ValueError): + """a user-supplied gpu, cpu, or region that does not resolve.""" + + +def list_all( + path: str, collection: str, params: Optional[dict[str, Any]] = None +) -> list[dict]: + """every record of a cursor-paginated list endpoint.""" + query = dict(params or {}) + records: list[dict] = [] + while True: + page = run_rest_request("GET", path, params=query) or {} + records.extend(page.get(collection) or []) + pagination = page.get("pagination") or {} + cursor = pagination.get("nextCursor") + if not pagination.get("hasNextPage") or not cursor: + return records + query["cursor"] = cursor + + +def _query_params(**values: Any) -> dict[str, Any]: + params = {} + for key, value in values.items(): + if value is None or value == []: + continue + if isinstance(value, (list, tuple)): + value = ",".join(str(item) for item in value) + params[key] = value + return params + + +def list_gpus( + product: Optional[str] = None, + count: Optional[int] = None, + cloud: Optional[str] = None, + cuda_versions: Optional[list[str]] = None, + min_cuda_version: Optional[str] = None, +) -> list[dict]: + """gpu types, with availability when a product context is given.""" + params = _query_params( + include="AVAILABILITY" if product else None, + product=product, + count=count if product else None, + cloud=cloud if product else None, + cudaVersions=cuda_versions if product else None, + minCudaVersion=min_cuda_version if product else None, + ) + response = run_rest_request("GET", "/v2/catalog/gpus", params=params) or {} + return response.get("gpus") or [] + + +def list_cpus( + product: Optional[str] = None, vcpu_count: Optional[int] = None +) -> list[dict]: + """cpu flavors, with availability when a product context is given.""" + params = _query_params( + include="AVAILABILITY" if product else None, + product=product, + vcpuCount=vcpu_count if product else None, + ) + response = run_rest_request("GET", "/v2/catalog/cpus", params=params) or {} + return response.get("cpus") or [] + + +def list_datacenters(include: Optional[list[str]] = None) -> list[dict]: + """every data center in the catalog.""" + response = ( + run_rest_request( + "GET", "/v2/catalog/datacenters", params=_query_params(include=include) + ) + or {} + ) + return response.get("dataCenters") or [] + + +def _tokens(text: str) -> list[str]: + return [token for token in re.split(r"[\s_\-]+", text.lower()) if token] + + +def _compact(text: str) -> str: + return re.sub(r"[\s_\-]+", "", text.lower()) + + +def match_gpus(query: str, gpus: Iterable[dict]) -> list[dict]: + """gpu types a human-readable query names. + + an exact id, short name, or serverless pool id wins. otherwise every + query token must appear in the type's id or name, so "4090" matches + the rtx 4090, "h100" matches every h100 variant, and "l4" does not + match the l40. + """ + gpus = list(gpus) + text = query.strip().lower() + exact = [ + gpu + for gpu in gpus + if text in (gpu["id"].lower(), (gpu.get("name") or "").lower()) + ] + if exact: + return exact + pool = [gpu for gpu in gpus if (gpu.get("pool") or "").lower() == text] + if pool: + return pool + wanted = [token for token in _tokens(text) if token not in _NOISE_TOKENS] + if not wanted: + return [] + hits = [ + gpu + for gpu in gpus + if set(wanted) + <= set(_tokens(gpu["id"])) | set(_tokens(gpu.get("name") or "")) + ] + if hits: + return hits + compact = "".join(wanted) + return [gpu for gpu in gpus if compact in _compact(gpu["id"])] + + +def resolve_gpus(query: str, gpus: Iterable[dict]) -> list[dict]: + """like match_gpus, but raise when nothing matches.""" + gpus = list(gpus) + matches = match_gpus(query, gpus) + if not matches: + known = ", ".join(sorted(gpu.get("name") or gpu["id"] for gpu in gpus)) + raise ResolutionError(f"unknown gpu '{query}'. known gpus: {known}") + return matches + + +def parse_gpu_spec(value: str) -> tuple[str, Optional[int]]: + """split "h100:2" into ("h100", 2); a bare name has no count.""" + name, sep, count = value.strip().rpartition(":") + if not sep: + return value.strip(), None + if not name or not count.isdigit() or int(count) < 1: + raise ResolutionError( + f"invalid gpu '{value}'. use a name with an optional count, like h100:2" + ) + return name, int(count) + + +def parse_cuda(values: Iterable[str]) -> tuple[Optional[str], list[str]]: + """split cuda filters into (minimum, exact versions). + + ">=12.4" sets a floor; "12.4" or "12.4,12.8" lists exact versions. + """ + minimum = None + exact: list[str] = [] + for value in values: + for part in value.split(","): + part = part.strip() + if not part: + continue + floor = part.startswith(">=") + version = part[2:].strip() if floor else part + if not re.fullmatch(r"\d+\.\d+", version): + raise ResolutionError( + f"invalid cuda version '{part}'. use 12.4, 12.4,12.8, or '>=12.4'" + ) + if floor: + minimum = version + else: + exact.append(version) + if minimum and exact: + raise ResolutionError("use either a cuda floor (>=) or exact versions, not both") + return minimum, exact + + +def endpoint_gpu_selection( + matches: Iterable[dict], gpus: Iterable[dict] +) -> tuple[list[str], list[str]]: + """serverless (pools, excludedTypes) that select exactly the matched types.""" + matches = list(matches) + pools: list[str] = [] + for gpu in matches: + if gpu.get("pool") and gpu["pool"] not in pools: + pools.append(gpu["pool"]) + if not pools: + names = ", ".join(gpu.get("name") or gpu["id"] for gpu in matches) + raise ResolutionError(f"{names} is not offered on serverless") + chosen = {gpu["id"] for gpu in matches} + excluded = [ + gpu["id"] + for gpu in gpus + if gpu.get("pool") in pools and gpu["id"] not in chosen + ] + return pools, excluded + + +def resolve_regions(values: Iterable[str], datacenters: Iterable[dict]) -> list[str]: + """data center ids named by ids, id prefixes, or region names. + + "EU-RO-1" is one data center, "us" and "eu-ro" match every id with + that prefix, and "eu" or "europe" match every data center in the + region. + """ + datacenters = list(datacenters) + resolved: list[str] = [] + for value in values: + text = value.strip().upper().replace("_", "-") + region = REGION_ALIASES.get(_compact(value)) + matches = [ + dc["id"] + for dc in datacenters + if dc["id"].upper() == text + or dc["id"].upper().startswith(text + "-") + or (region and dc.get("region") == region) + ] + if not matches: + ids = ", ".join(sorted(dc["id"] for dc in datacenters)) + names = ", ".join( + sorted( + { + dc["region"].lower().replace("_", "-") + for dc in datacenters + if dc.get("region") + } + ) + ) + raise ResolutionError( + f"unknown region '{value}'. use a region (eu, na, {names}), " + f"an id prefix like us or eu-ro, or a data center ({ids})" + ) + resolved.extend(match for match in matches if match not in resolved) + return resolved + + +def match_cpus(query: str, cpus: Iterable[dict]) -> list[dict]: + """cpu flavors named by id ("cpu3c"), group ("cpu3"), or name.""" + text = query.strip().lower() + return [ + cpu + for cpu in cpus + if text + in ( + cpu["id"].lower(), + (cpu.get("group") or "").lower(), + (cpu.get("name") or "").lower(), + ) + or _compact(text) == _compact(cpu.get("name") or "") + ] + + +def resolve_cpu(query: str, cpus: Iterable[dict]) -> dict: + """the single cpu flavor a query names.""" + cpus = list(cpus) + matches = match_cpus(query, cpus) + if len(matches) != 1: + known = ", ".join(sorted(cpu["id"] for cpu in cpus)) + problem = "ambiguous" if matches else "unknown" + raise ResolutionError(f"{problem} cpu '{query}'. cpu flavors: {known}") + return matches[0] + + +def availability_rank(level: Optional[str]) -> int: + return AVAILABILITY_RANK.get(level or "NONE", 0) + + +def best_availability(levels: Iterable[Optional[str]]) -> str: + """the highest stock level among levels, NONE when empty.""" + return max(levels, key=availability_rank, default="NONE") or "NONE" diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index 5c4b4bfd..d8be1833 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -16,7 +16,7 @@ # commands whose output would be polluted by (or that manage) the # update notice -_UPDATE_CHECK_EXCLUDED = frozenset({"flash", "update"}) +_UPDATE_CHECK_EXCLUDED = frozenset({"completion", "flash", "update"}) _FLASH_UPDATE_CHECK_EXCLUDED = frozenset({"dev"}) @@ -486,6 +486,23 @@ async def _session() -> int: raise click.ClickException(str(exc)) from exc +# ------------------------------------------------------------ completion + + +@cli.command() +@click.argument("shell", type=click.Choice(["bash", "zsh", "fish"])) +def completion(shell): + """Print the shell completion script for SHELL. + + Completes commands, flags, and resource IDs. Add this to your shell + profile: eval "$(rp completion zsh)" + """ + from click.shell_completion import get_completion_class + + complete = get_completion_class(shell)(cli, {}, "rp", "_RP_COMPLETE") + click.echo(complete.source()) + + # ------------------------------------------------------------------ logs @@ -1021,8 +1038,9 @@ def undeploy(app_name, env_name, yes): 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 - for command in (ssh_cli, pod_cli): + for command in (ssh_cli, pod_cli, catalog): cli.add_command(command) diff --git a/runpod/rp_cli/output.py b/runpod/rp_cli/output.py new file mode 100644 index 00000000..ce0a0dad --- /dev/null +++ b/runpod/rp_cli/output.py @@ -0,0 +1,120 @@ +"""output, picker, prompt, and completion helpers shared by resource commands.""" + +import contextlib +import json +import sys +from typing import Any, Callable, Iterable, Iterator, Sequence + +import click +from click.shell_completion import CompletionItem + +from runpod import error +from runpod.rp_cli.compute import ResolutionError + +Column = tuple[str, Callable[[dict], Any]] + +json_option = click.option( + "--json", "as_json", is_flag=True, help="Print JSON for scripts and agents." +) + + +def print_json(data: Any) -> None: + click.echo(json.dumps(data, indent=2, default=str)) + + +def cell(value: Any) -> str: + if value is None or value == "" or value == [] or value == {}: + return "-" + if isinstance(value, bool): + return "yes" if value else "no" + if isinstance(value, (list, tuple)): + return ", ".join(str(item) for item in value) + return str(value) + + +def print_table(rows: Sequence[dict], columns: Sequence[Column], empty: str) -> None: + """render rows as an aligned table, or a dim notice when empty.""" + from rich.markup import escape + from rich.table import Table + + from runpod.rp_cli import console as ui + + if not rows: + ui.console.print(f"[dim]{empty}[/dim]") + return + table = Table(box=None, header_style="dim", pad_edge=False) + for header, _ in columns: + table.add_column(header) + for row in rows: + table.add_row(*(escape(cell(value(row))) for _, value in columns)) + ui.console.print(table) + + +def print_fields(record: dict, fields: Sequence[Column]) -> None: + """render one record as aligned label/value lines.""" + from rich.markup import escape + + from runpod.rp_cli import console as ui + + width = max(len(label) for label, _ in fields) + for label, value in fields: + ui.console.print(f"[dim]{label:<{width}}[/dim] {escape(cell(value(record)))}") + + +def interactive() -> bool: + return sys.stdin.isatty() and sys.stdout.isatty() + + +def pick(message: str, choices: Sequence[tuple[str, str]], missing: str) -> str: + """let a person choose a value from (value, label) pairs. + + scripts and agents get a usage error naming what to pass instead. + """ + if not interactive(): + raise click.UsageError(missing) + if not choices: + raise click.UsageError(f"{missing} (none found)") + from InquirerPy import inquirer + from InquirerPy.base.control import Choice + + return inquirer.fuzzy( + message=message, + choices=[Choice(value, name=label) for value, label in choices], + ).execute() + + +def confirm(message: str, yes: bool) -> None: + """ask before a destructive action unless --yes was passed.""" + if yes: + return + if not interactive(): + raise click.UsageError("pass --yes to confirm in a non-interactive shell") + click.confirm(message, abort=True) + + +def completer( + fetch: Callable[[], Iterable[tuple[str, str]]], +) -> Callable[[click.Context, click.Parameter, str], list[CompletionItem]]: + """shell completion offering (value, help) pairs from fetch.""" + + def complete(ctx, param, incomplete): + try: + entries = list(fetch()) + except Exception: # noqa: BLE001 - completion must never break the shell + return [] + return [ + CompletionItem(value, help=label) + for value, label in entries + if value.startswith(incomplete) + ] + + return complete + + +@contextlib.contextmanager +def api_errors() -> Iterator[None]: + """surface api and resolution failures as clean cli errors.""" + try: + yield + except (error.RunPodError, ResolutionError) as exc: + raise click.ClickException(str(exc)) from exc From 15a69b4a1d49a9dd37c8a92a52068ba2bce6343a Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:37:27 -0700 Subject: [PATCH 02/25] test: cover rp catalog and shared compute resolution --- runpod/rp_cli/catalog.py | 46 ++-- runpod/rp_cli/compute.py | 7 +- tests/test_cli/test_catalog_command.py | 287 +++++++++++++++++++++++++ tests/test_cli/test_compute.py | 222 +++++++++++++++++++ 4 files changed, 546 insertions(+), 16 deletions(-) create mode 100644 tests/test_cli/test_catalog_command.py create mode 100644 tests/test_cli/test_compute.py diff --git a/runpod/rp_cli/catalog.py b/runpod/rp_cli/catalog.py index 14e02e62..16d864d0 100644 --- a/runpod/rp_cli/catalog.py +++ b/runpod/rp_cli/catalog.py @@ -29,9 +29,7 @@ def _in_regions(record: dict, allowed: Optional[set[str]]) -> dict: """the record with stock narrowed to the allowed data centers.""" if allowed is None: return dict(record) - datacenters = [ - dc for dc in record.get("dataCenters") or [] if dc["id"] in allowed - ] + datacenters = [dc for dc in record.get("dataCenters") or [] if dc["id"] in allowed] return dict( record, dataCenters=datacenters, @@ -53,7 +51,9 @@ def _unit_price(gpu: dict, product: str, cloud: Optional[str]) -> Optional[float return price.get("serverless") or None if cloud: return price.get(cloud) or None - offered = [value for value in (price.get("secure"), price.get("community")) if value] + offered = [ + value for value in (price.get("secure"), price.get("community")) if value + ] return min(offered, default=None) @@ -151,9 +151,16 @@ def _cpu_columns(product: str): ("GROUP", lambda row: row.get("group")), ( "VCPU", - lambda row: "{min}-{max}".format(**row["vcpu"]) if row.get("vcpu") else None, + lambda row: ( + "{min}-{max}".format(**row["vcpu"]) if row.get("vcpu") else None + ), + ), + ( + "RAM/VCPU", + lambda row: ( + f"{row['ramGbPerVcpu']:g}GB" if row.get("ramGbPerVcpu") else None + ), ), - ("RAM/VCPU", lambda row: f"{row['ramGbPerVcpu']:g}GB" if row.get("ramGbPerVcpu") else None), ("$/VCPU-HR", lambda row: _money((row.get("price") or {}).get(tier))), ("STOCK", lambda row: (row.get("availability") or "NONE").lower()), ("DATA CENTERS", lambda row: _stock(row.get("dataCenters") or [])), @@ -187,8 +194,12 @@ def _cpu_columns(product: str): flag_value="SERVERLESS", help="Serverless pricing, pools, and stock.", ) -@click.option("--cpu", "cpu_mode", is_flag=True, help="Search CPU flavors instead of GPUs.") -@click.option("--datacenters", is_flag=True, help="List data centers and their regions.") +@click.option( + "--cpu", "cpu_mode", is_flag=True, help="Search CPU flavors instead of GPUs." +) +@click.option( + "--datacenters", is_flag=True, help="List data centers and their regions." +) @click.option( "--count", type=click.IntRange(min=1), @@ -197,7 +208,10 @@ def _cpu_columns(product: str): help="GPUs per pod or worker; prices and stock account for it.", ) @click.option( - "--vcpu", type=click.IntRange(min=2), default=None, help="vCPUs per pod or worker (with --cpu)." + "--vcpu", + type=click.IntRange(min=2), + default=None, + help="vCPUs per pod or worker (with --cpu).", ) @click.option( "--region", @@ -217,9 +231,7 @@ def _cpu_columns(product: str): help="CUDA filter: '>=12.4' for a floor, 12.4 or 12.4,12.8 for exact versions.", ) @click.option("--min-vram", type=int, default=None, help="Minimum VRAM per GPU in GB.") -@click.option( - "--max-price", type=float, default=None, help="Maximum USD per GPU-hour." -) +@click.option("--max-price", type=float, default=None, help="Maximum USD per GPU-hour.") @click.option("--available", is_flag=True, help="Only show hardware that is in stock.") @json_option def catalog( @@ -267,7 +279,15 @@ def catalog( columns, empty = _cpu_columns(product), "no cpu flavors match" else: rows = _gpu_rows( - query, product, count, regions, cloud, cuda, min_vram, max_price, available + query, + product, + count, + regions, + cloud, + cuda, + min_vram, + max_price, + available, ) columns, empty = _gpu_columns(product, count, cloud), "no gpus match" diff --git a/runpod/rp_cli/compute.py b/runpod/rp_cli/compute.py index c072a07b..af0a0540 100644 --- a/runpod/rp_cli/compute.py +++ b/runpod/rp_cli/compute.py @@ -134,8 +134,7 @@ def match_gpus(query: str, gpus: Iterable[dict]) -> list[dict]: hits = [ gpu for gpu in gpus - if set(wanted) - <= set(_tokens(gpu["id"])) | set(_tokens(gpu.get("name") or "")) + if set(wanted) <= set(_tokens(gpu["id"])) | set(_tokens(gpu.get("name") or "")) ] if hits: return hits @@ -188,7 +187,9 @@ def parse_cuda(values: Iterable[str]) -> tuple[Optional[str], list[str]]: else: exact.append(version) if minimum and exact: - raise ResolutionError("use either a cuda floor (>=) or exact versions, not both") + raise ResolutionError( + "use either a cuda floor (>=) or exact versions, not both" + ) return minimum, exact diff --git a/tests/test_cli/test_catalog_command.py b/tests/test_cli/test_catalog_command.py new file mode 100644 index 00000000..8efb5fd5 --- /dev/null +++ b/tests/test_cli/test_catalog_command.py @@ -0,0 +1,287 @@ +"""tests for rp catalog and rp completion.""" + +import json +from unittest.mock import patch + +import click +import pytest +from click.testing import CliRunner + +from runpod import error +from runpod.rp_cli import output +from runpod.rp_cli.main import cli + +GPUS = [ + { + "id": "NVIDIA H100 80GB HBM3", + "name": "H100 SXM", + "pool": "ADA_80_PRO", + "memory": 80, + "price": {"secure": 3.99, "community": 2.69, "serverless": 4.79}, + "maxCount": {"secure": 8, "community": 1}, + "availability": "LOW", + "dataCenters": [ + {"id": "US-TX-3", "availability": "LOW"}, + {"id": "EU-RO-1", "availability": "LOW"}, + ], + }, + { + "id": "NVIDIA H100 NVL", + "name": "H100 NVL", + "pool": "ADA_80_PRO", + "memory": 94, + "price": {"secure": 3.19, "community": 0}, + "maxCount": {"secure": 1, "community": 10}, + "availability": "HIGH", + "dataCenters": [{"id": "US-TX-3", "availability": "HIGH"}], + }, + { + "id": "NVIDIA GeForce RTX 3070", + "name": "RTX 3070", + "pool": None, + "memory": 8, + "price": {"secure": 0, "community": 0.13}, + "maxCount": {"secure": 0, "community": 7}, + "availability": "NONE", + "dataCenters": [], + }, +] + +CPUS = [ + { + "id": "cpu3c", + "name": "Compute-Optimized", + "group": "CPU3", + "vcpu": {"min": 2, "max": 32}, + "ramGbPerVcpu": 2, + "price": {"securePerVcpu": 0.03, "serverlessPerVcpu": 0.036}, + "availability": "HIGH", + "dataCenters": [{"id": "EU-RO-1", "availability": "HIGH"}], + }, + { + "id": "cpu5c", + "name": "Compute-Optimized", + "group": "CPU5", + "vcpu": {"min": 2, "max": 32}, + "ramGbPerVcpu": 2, + "price": {"securePerVcpu": 0.04}, + "availability": "NONE", + "dataCenters": [], + }, +] + +DATACENTERS = [ + { + "id": "EU-RO-1", + "region": "EUROPE", + "networkVolumeTypes": ["STANDARD"], + "globalNetwork": True, + "compliance": ["GDPR"], + }, + { + "id": "US-TX-3", + "region": "NORTH_AMERICA", + "networkVolumeTypes": [], + "globalNetwork": False, + "compliance": [], + }, +] + + +@pytest.fixture(autouse=True) +def wide_console(monkeypatch): + from runpod.rp_cli import console as ui + + monkeypatch.setattr(ui.console, "_width", 200) + + +def _invoke(args, gpus=GPUS, cpus=CPUS, datacenters=DATACENTERS): + with ( + patch("runpod.rp_cli.compute.list_gpus", return_value=gpus) as list_gpus, + patch("runpod.rp_cli.compute.list_cpus", return_value=cpus) as list_cpus, + patch("runpod.rp_cli.compute.list_datacenters", return_value=datacenters), + ): + result = CliRunner().invoke(cli, ["catalog", *args]) + return result, list_gpus, list_cpus + + +class TestGpuCatalog: + def test_pod_table_sorted_by_stock_then_price(self): + result, list_gpus, _ = _invoke(["h100", "--count", "2"]) + assert result.exit_code == 0, result.output + assert "SECURE $/HR x2" in result.output + assert result.output.index("H100 NVL") < result.output.index("H100 SXM") + assert "$7.98" in result.output + assert "US-TX-3 low, EU-RO-1 low" in result.output + assert list_gpus.call_args.kwargs == { + "product": "POD", + "count": 2, + "cloud": None, + "cuda_versions": [], + "min_cuda_version": None, + } + + def test_region_narrows_stock(self): + result, _, _ = _invoke(["h100", "--region", "eu", "--json"]) + rows = json.loads(result.output) + assert [row["availability"] for row in rows] == ["LOW", "NONE"] + assert rows[0]["dataCenters"] == [{"id": "EU-RO-1", "availability": "LOW"}] + assert rows[1]["dataCenters"] == [] + + def test_cuda_and_cloud_filters(self): + result, list_gpus, _ = _invoke( + ["--cloud", "COMMUNITY", "--cuda", ">=12.4", "--max-price", "1"] + ) + assert result.exit_code == 0, result.output + assert list_gpus.call_args.kwargs["cloud"] == "COMMUNITY" + assert list_gpus.call_args.kwargs["min_cuda_version"] == "12.4" + assert "COMMUNITY $/HR" in result.output + assert "SECURE" not in result.output + assert "RTX 3070" in result.output + assert "H100" not in result.output + + def test_available_and_vram_filters(self): + result, _, _ = _invoke(["--available", "--min-vram", "90", "--json"]) + assert [row["id"] for row in json.loads(result.output)] == ["NVIDIA H100 NVL"] + + def test_json_totals(self): + result, _, _ = _invoke(["h100 sxm", "--count", "2", "--json"]) + (row,) = json.loads(result.output) + assert row["count"] == 2 + assert row["totalPrice"] == { + "secure": 7.98, + "community": 5.38, + "serverless": 9.58, + } + + def test_serverless_shows_pools_only(self): + result, list_gpus, _ = _invoke(["--serverless"]) + assert result.exit_code == 0, result.output + assert list_gpus.call_args.kwargs["product"] == "SERVERLESS" + assert "ADA_80_PRO" in result.output + assert "$4.79" in result.output + assert "RTX 3070" not in result.output + + def test_unknown_gpu(self): + result, _, _ = _invoke(["b9000"]) + assert result.exit_code == 1 + assert "unknown gpu 'b9000'" in result.output + + def test_empty(self): + result, _, _ = _invoke(["--max-price", "0.01"]) + assert "no gpus match" in result.output + + def test_conflicting_flags(self): + assert _invoke(["--cpu", "--datacenters"])[0].exit_code == 2 + assert _invoke(["--serverless", "--cloud", "secure"])[0].exit_code == 2 + + def test_api_error_is_clean(self): + with patch( + "runpod.rp_cli.compute.list_gpus", + side_effect=error.AuthenticationError("Unauthorized request"), + ): + result = CliRunner().invoke(cli, ["catalog"]) + assert result.exit_code == 1 + assert "Error: Unauthorized request" in result.output + + +class TestCpuCatalog: + def test_table(self): + result, _, list_cpus = _invoke(["--cpu", "--vcpu", "8", "--serverless"]) + assert result.exit_code == 0, result.output + assert list_cpus.call_args.kwargs == {"product": "SERVERLESS", "vcpu_count": 8} + assert result.output.index("cpu3c") < result.output.index("cpu5c") + assert "$0.04" in result.output + assert "2-32" in result.output + + def test_query_region_and_available(self): + result, _, _ = _invoke( + ["cpu3", "--cpu", "--region", "us", "--available", "--json"] + ) + assert json.loads(result.output) == [] + + def test_unknown(self): + result, _, _ = _invoke(["cpu9", "--cpu"]) + assert "unknown cpu 'cpu9'" in result.output + + +class TestDatacenterCatalog: + def test_all(self): + result, _, _ = _invoke(["--datacenters"]) + assert result.exit_code == 0, result.output + assert "north america" in result.output + assert "standard" in result.output + + def test_query_is_a_region(self): + result, _, _ = _invoke(["--datacenters", "eu", "--json"]) + assert [row["id"] for row in json.loads(result.output)] == ["EU-RO-1"] + + +class TestCompletion: + def test_prints_script(self): + result = CliRunner().invoke(cli, ["completion", "zsh"]) + assert result.exit_code == 0 + assert "_RP_COMPLETE" in result.output + + +class TestOutputHelpers: + def test_cell(self): + assert output.cell(None) == "-" + assert output.cell(True) == "yes" + assert output.cell(["a", "b"]) == "a, b" + assert output.cell(3) == "3" + + def test_print_fields(self, capsys): + output.print_fields( + {"id": "abc", "tags": []}, + [("id", lambda r: r["id"]), ("tags", lambda r: r["tags"])], + ) + out = capsys.readouterr().out + assert "id abc" in out + assert "tags -" in out + + def test_pick_requires_tty(self): + with patch.object(output, "interactive", return_value=False): + with pytest.raises(click.UsageError, match="pass an ID"): + output.pick("pick", [("a", "a")], "pass an ID") + + def test_pick_with_no_choices(self): + with patch.object(output, "interactive", return_value=True): + with pytest.raises(click.UsageError, match="none found"): + output.pick("pick", [], "pass an ID") + + def test_pick_uses_fuzzy_prompt(self): + with ( + patch.object(output, "interactive", return_value=True), + patch("InquirerPy.inquirer.fuzzy") as fuzzy, + ): + fuzzy.return_value.execute.return_value = "b" + assert output.pick("pick", [("a", "first"), ("b", "second")], "x") == "b" + choices = fuzzy.call_args.kwargs["choices"] + assert [(choice.value, choice.name) for choice in choices] == [ + ("a", "first"), + ("b", "second"), + ] + + def test_confirm(self): + output.confirm("sure?", yes=True) + with patch.object(output, "interactive", return_value=False): + with pytest.raises(click.UsageError, match="--yes"): + output.confirm("sure?", yes=False) + with ( + patch.object(output, "interactive", return_value=True), + patch("click.confirm") as ask, + ): + output.confirm("sure?", yes=False) + ask.assert_called_once_with("sure?", abort=True) + + def test_completer(self): + complete = output.completer(lambda: [("abc", "first"), ("xyz", "second")]) + items = complete(None, None, "a") + assert [(item.value, item.help) for item in items] == [("abc", "first")] + + def test_completer_swallows_errors(self): + def broken(): + raise error.AuthenticationError("no key") + + assert output.completer(broken)(None, None, "") == [] diff --git a/tests/test_cli/test_compute.py b/tests/test_cli/test_compute.py new file mode 100644 index 00000000..4debe7cb --- /dev/null +++ b/tests/test_cli/test_compute.py @@ -0,0 +1,222 @@ +"""tests for gpu, cpu, and region resolution.""" + +from unittest.mock import patch + +import pytest + +from runpod.rp_cli import compute + +GPUS = [ + {"id": "NVIDIA GeForce RTX 4090", "name": "RTX 4090", "pool": "ADA_24"}, + {"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 H100 PCIe", "name": "H100 PCIe", "pool": "ADA_80_PRO"}, + {"id": "NVIDIA L4", "name": "L4", "pool": "AMPERE_24"}, + {"id": "NVIDIA L40", "name": "L40", "pool": "ADA_48_PRO"}, + {"id": "NVIDIA A100-SXM4-40GB", "name": "A100 SXM 40GB", "pool": None}, +] + +DATACENTERS = [ + {"id": "EU-RO-1", "region": "EUROPE"}, + {"id": "EUR-IS-1", "region": "EUROPE"}, + {"id": "US-TX-3", "region": "NORTH_AMERICA"}, + {"id": "US-TX-4", "region": "NORTH_AMERICA"}, + {"id": "CA-MTL-1", "region": "NORTH_AMERICA"}, + {"id": "AP-JP-1", "region": "ASIA"}, +] + + +def _ids(gpus): + return [gpu["id"] for gpu in gpus] + + +class TestMatchGpus: + @pytest.mark.parametrize( + "query, expected", + [ + ("4090", ["NVIDIA GeForce RTX 4090"]), + ("rtx 4090", ["NVIDIA GeForce RTX 4090"]), + ("rtx4090", ["NVIDIA GeForce RTX 4090"]), + ("NVIDIA L4", ["NVIDIA L4"]), + ("l4", ["NVIDIA L4"]), + ("h100 sxm", ["NVIDIA H100 80GB HBM3"]), + ( + "ADA_80_PRO", + ["NVIDIA H100 80GB HBM3", "NVIDIA H100 NVL", "NVIDIA H100 PCIe"], + ), + ("a100-sxm4", ["NVIDIA A100-SXM4-40GB"]), + ("h100", ["NVIDIA H100 80GB HBM3", "NVIDIA H100 NVL", "NVIDIA H100 PCIe"]), + ("nvidia", []), + ("mi300x", []), + ], + ) + def test_queries(self, query, expected): + assert _ids(compute.match_gpus(query, GPUS)) == expected + + def test_resolve_raises_with_known_names(self): + with pytest.raises( + compute.ResolutionError, match="unknown gpu 'b9000'.*RTX 4090" + ): + compute.resolve_gpus("b9000", GPUS) + + +class TestParsing: + def test_gpu_spec(self): + assert compute.parse_gpu_spec("h100:2") == ("h100", 2) + assert compute.parse_gpu_spec(" 4090 ") == ("4090", None) + + @pytest.mark.parametrize("value", ["h100:0", "h100:x", ":2"]) + def test_invalid_gpu_spec(self, value): + with pytest.raises(compute.ResolutionError, match="invalid gpu"): + compute.parse_gpu_spec(value) + + def test_cuda(self): + assert compute.parse_cuda([">=12.4"]) == ("12.4", []) + assert compute.parse_cuda(["12.4,12.8", "13.0", ""]) == ( + None, + ["12.4", "12.8", "13.0"], + ) + + def test_invalid_cuda(self): + with pytest.raises(compute.ResolutionError, match="invalid cuda"): + compute.parse_cuda(["twelve"]) + with pytest.raises(compute.ResolutionError, match="not both"): + compute.parse_cuda([">=12.4", "12.8"]) + + +class TestEndpointSelection: + def test_excludes_unselected_pool_members(self): + matches = compute.match_gpus("h100 nvl", GPUS) + assert compute.endpoint_gpu_selection(matches, GPUS) == ( + ["ADA_80_PRO"], + ["NVIDIA H100 80GB HBM3", "NVIDIA H100 PCIe"], + ) + + def test_whole_pool(self): + matches = compute.match_gpus("h100", GPUS) + assert compute.endpoint_gpu_selection(matches, GPUS) == (["ADA_80_PRO"], []) + + def test_not_on_serverless(self): + matches = compute.match_gpus("a100-sxm4", GPUS) + with pytest.raises(compute.ResolutionError, match="not offered on serverless"): + compute.endpoint_gpu_selection(matches, GPUS) + + +class TestRegions: + @pytest.mark.parametrize( + "values, expected", + [ + (["eu"], ["EU-RO-1", "EUR-IS-1"]), + (["Europe"], ["EU-RO-1", "EUR-IS-1"]), + (["north-america"], ["US-TX-3", "US-TX-4", "CA-MTL-1"]), + (["us"], ["US-TX-3", "US-TX-4"]), + (["us-tx"], ["US-TX-3", "US-TX-4"]), + (["eu_ro_1"], ["EU-RO-1"]), + (["ap", "EU-RO-1", "eu-ro"], ["AP-JP-1", "EU-RO-1"]), + ], + ) + def test_resolution(self, values, expected): + assert compute.resolve_regions(values, DATACENTERS) == expected + + def test_unknown_region_lists_options(self): + with pytest.raises(compute.ResolutionError) as exc: + compute.resolve_regions(["mars"], DATACENTERS) + assert "north-america" in str(exc.value) + assert "AP-JP-1" in str(exc.value) + + +class TestCpus: + CPUS = [ + {"id": "cpu3c", "name": "Compute-Optimized", "group": "CPU3"}, + {"id": "cpu3g", "name": "General Purpose", "group": "CPU3"}, + {"id": "cpu5c", "name": "Compute-Optimized", "group": "CPU5"}, + ] + + def test_match(self): + assert [cpu["id"] for cpu in compute.match_cpus("CPU3", self.CPUS)] == [ + "cpu3c", + "cpu3g", + ] + assert [ + cpu["id"] for cpu in compute.match_cpus("general purpose", self.CPUS) + ] == ["cpu3g"] + + def test_resolve(self): + assert compute.resolve_cpu("cpu5c", self.CPUS)["id"] == "cpu5c" + with pytest.raises(compute.ResolutionError, match="ambiguous cpu"): + compute.resolve_cpu("cpu3", self.CPUS) + with pytest.raises(compute.ResolutionError, match="unknown cpu"): + compute.resolve_cpu("cpu9", self.CPUS) + + +class TestAvailability: + def test_best(self): + assert compute.best_availability(["LOW", None, "HIGH", "MEDIUM"]) == "HIGH" + assert compute.best_availability([]) == "NONE" + assert compute.availability_rank(None) == 0 + + +class TestCatalogRequests: + def test_list_all_follows_cursors(self): + pages = [ + { + "pods": [{"id": "a"}], + "pagination": {"hasNextPage": True, "nextCursor": "c1"}, + }, + { + "pods": [{"id": "b"}], + "pagination": {"hasNextPage": False, "nextCursor": None}, + }, + ] + with patch.object(compute, "run_rest_request", side_effect=pages) as request: + assert compute.list_all("/v2/pods", "pods", {"limit": 1}) == [ + {"id": "a"}, + {"id": "b"}, + ] + assert request.call_args_list[1].kwargs["params"] == { + "limit": 1, + "cursor": "c1", + } + + def test_list_gpus_with_availability(self): + with patch.object( + compute, "run_rest_request", return_value={"gpus": GPUS} + ) as request: + assert ( + compute.list_gpus( + product="POD", + count=2, + cloud="SECURE", + cuda_versions=["12.4", "12.8"], + ) + == GPUS + ) + assert request.call_args.kwargs["params"] == { + "include": "AVAILABILITY", + "product": "POD", + "count": 2, + "cloud": "SECURE", + "cudaVersions": "12.4,12.8", + } + + def test_list_gpus_without_availability(self): + with patch.object(compute, "run_rest_request", return_value=None) as request: + assert compute.list_gpus(count=2, min_cuda_version="12.4") == [] + assert request.call_args.kwargs["params"] == {} + + def test_list_cpus_and_datacenters(self): + with patch.object( + compute, + "run_rest_request", + side_effect=[{"cpus": [{"id": "cpu3c"}]}, {"dataCenters": DATACENTERS}], + ) as request: + assert compute.list_cpus(product="SERVERLESS", vcpu_count=4) == [ + {"id": "cpu3c"} + ] + assert compute.list_datacenters() == DATACENTERS + assert request.call_args_list[0].kwargs["params"] == { + "include": "AVAILABILITY", + "product": "SERVERLESS", + "vcpuCount": 4, + } + assert request.call_args_list[1].args[1] == "/v2/catalog/datacenters" From 17f26d71693eada3220718e2543d58b804c3d60e Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:38:44 -0700 Subject: [PATCH 03/25] feat: add rp volume commands for network and global volumes --- runpod/rp_cli/main.py | 3 +- runpod/rp_cli/volume.py | 372 +++++++++++++++++++++ tests/test_cli/test_volume_command.py | 449 ++++++++++++++++++++++++++ 3 files changed, 823 insertions(+), 1 deletion(-) create mode 100644 runpod/rp_cli/volume.py create mode 100644 tests/test_cli/test_volume_command.py diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index d8be1833..9c6f1546 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.volume import volume - for command in (ssh_cli, pod_cli, catalog): + for command in (ssh_cli, pod_cli, catalog, volume): cli.add_command(command) diff --git a/runpod/rp_cli/volume.py b/runpod/rp_cli/volume.py new file mode 100644 index 00000000..43be7cc5 --- /dev/null +++ b/runpod/rp_cli/volume.py @@ -0,0 +1,372 @@ +"""`rp volume`: network and global volume management.""" + +import re +from typing import NamedTuple, Optional + +import click + +from runpod.api.rest import run_rest_request +from runpod.rp_cli import compute +from runpod.rp_cli.output import ( + api_errors, + completer, + confirm, + json_option, + pick, + print_fields, + print_json, + print_table, +) + +MIN_SIZE = 10 +MAX_SIZE = 4096 + + +class Kind(NamedTuple): + path: str + collection: str + noun: str + + +NETWORK = Kind("/v2/network-volumes", "networkVolumes", "network volume") +GLOBAL = Kind("/v2/global-volumes", "globalVolumes", "global volume") + + +def _kind(global_volume: bool) -> Kind: + return GLOBAL if global_volume else NETWORK + + +def _type_label(record: dict) -> Optional[str]: + return (record.get("type") or "").lower() or None + + +def _size_label(record: dict) -> Optional[str]: + return f"{record['size']} GB" if record.get("size") else None + + +def _columns(kind: Kind): + if kind is GLOBAL: + return [ + ("ID", lambda row: row["id"]), + ("NAME", lambda row: row.get("name")), + ("CREATED", lambda row: row.get("createdAt")), + ("CREATED BY", lambda row: row.get("createdByUserId")), + ] + return [ + ("ID", lambda row: row["id"]), + ("NAME", lambda row: row.get("name")), + ("SIZE", _size_label), + ("DATA CENTER", lambda row: row.get("dataCenter")), + ("TYPE", _type_label), + ] + + +def _list(kind: Kind) -> list[dict]: + return compute.list_all(kind.path, kind.collection) + + +def _get(kind: Kind, volume_id: str) -> dict: + return run_rest_request("GET", f"{kind.path}/{volume_id}") or {} + + +def _label(kind: Kind, record: dict) -> str: + if kind is GLOBAL: + return record.get("name") or "" + details = [_size_label(record), record.get("dataCenter"), _type_label(record)] + return " ".join([record.get("name") or ""] + [d for d in details if d]) + + +def _choices(kind: Kind) -> list[tuple[str, str]]: + return [(record["id"], _label(kind, record)) for record in _list(kind)] + + +def _complete_ids(ctx, param, incomplete): + kind = _kind(bool(ctx.params.get("global_volume"))) + return completer(lambda: _choices(kind))(ctx, param, incomplete) + + +def _show(kind: Kind, record: dict, as_json: bool) -> None: + if as_json: + print_json(record) + return + print_fields(record, [(header.lower(), fn) for header, fn in _columns(kind)]) + + +def _data_center( + regions: tuple[str, ...], data_center: Optional[str], volume_type: Optional[str] +) -> str: + """the single data center a network volume is created in. + + regions narrow to data centers offering the requested volume type + (or any type), then a person picks one; scripts get the candidates. + """ + if data_center: + return data_center + datacenters = compute.list_datacenters() + if regions: + allowed = set(compute.resolve_regions(regions, datacenters)) + datacenters = [dc for dc in datacenters if dc["id"] in allowed] + candidates = [ + dc + for dc in datacenters + if (volume_type in (dc.get("networkVolumeTypes") or [])) + or (not volume_type and dc.get("networkVolumeTypes")) + ] + scope = f"region {', '.join(regions)}" if regions else "the catalog" + offered = ( + f"{volume_type.lower()} network volumes" if volume_type else "network volumes" + ) + if not candidates: + raise compute.ResolutionError(f"no data center in {scope} offers {offered}") + if len(candidates) == 1: + return candidates[0]["id"] + ids = ", ".join(dc["id"] for dc in candidates) + return pick( + "data center", + [ + ( + dc["id"], + f"{dc['id']} {(dc.get('region') or '').lower().replace('_', ' ')} " + + ", ".join(t.lower() for t in dc["networkVolumeTypes"]), + ) + for dc in candidates + ], + f"several data centers in {scope} offer {offered} ({ids}). " + "pass one with --data-center", + ) + + +def _parse_growth(value: str) -> tuple[bool, int]: + """split an expand size into (relative, gigabytes): N or +N.""" + text = value.strip() + if not re.fullmatch(r"\+?\d+", text): + raise click.BadParameter( + f"'{value}' is not a size. use N for a total or +N to add N GB", + param_hint="--size", + ) + return text.startswith("+"), int(text.lstrip("+")) + + +def _check_growth(volume_id: str, current: int, size: int, strict: bool) -> None: + if size < current or (strict and size == current): + raise click.UsageError( + f"network volumes can only grow: {volume_id} is {current} GB, " + f"requested {size} GB" + ) + if size > MAX_SIZE: + raise click.UsageError( + f"network volumes hold at most {MAX_SIZE} GB, requested {size} GB" + ) + + +global_option = click.option( + "--global", + "global_volume", + is_flag=True, + help="Manage global volumes, which mount in any data center.", +) + + +@click.group() +def volume(): + """Create, list, resize, and delete network and global volumes.""" + + +@volume.command(name="list") +@global_option +@json_option +def volume_list(global_volume, as_json): + """List volumes with their size, data center, and type.""" + kind = _kind(global_volume) + with api_errors(): + rows = _list(kind) + if as_json: + print_json(rows) + return + print_table( + rows, _columns(kind), f"no {kind.noun}s. create one with rp volume create" + ) + + +@volume.command(name="get") +@click.argument("volume_id", required=False, shell_complete=_complete_ids) +@global_option +@json_option +def volume_get(volume_id, global_volume, as_json): + """Show one volume (pick one when VOLUME_ID is omitted).""" + kind = _kind(global_volume) + with api_errors(): + if volume_id is None: + volume_id = pick(kind.noun, _choices(kind), "pass VOLUME_ID") + record = _get(kind, volume_id) + _show(kind, record, as_json) + + +@volume.command(name="create") +@click.option("--name", required=True, help="Human-readable name.") +@click.option( + "--size", + type=click.IntRange(MIN_SIZE, MAX_SIZE), + default=None, + help=f"Storage in GB ({MIN_SIZE}-{MAX_SIZE}). Required for network volumes.", +) +@click.option( + "--region", + "regions", + multiple=True, + help="Region (eu, na), id prefix (us, eu-ro), or data center id. Repeatable.", +) +@click.option("--data-center", default=None, help="Exact data center id, like EU-RO-1.") +@click.option( + "--type", + "volume_type", + type=click.Choice(["standard", "high_performance"], case_sensitive=False), + default=None, + help="Storage tier, immutable after creation. Defaults to the data center's " + "primary tier.", +) +@global_option +@json_option +def volume_create( + name, size, regions, data_center, volume_type, global_volume, as_json +): # pylint: disable=too-many-arguments + """Create a network volume in one data center, or a global volume. + + A region matching several data centers narrows to those offering + the requested --type, then prompts for one. + + \b + rp volume create --name models --size 100 --region eu + rp volume create --name scratch --size 50 --data-center EU-RO-1 + rp volume create --name shared --global + """ + from runpod.rp_cli import console as ui + + kind = _kind(global_volume) + if global_volume: + given = [ + flag + for flag, value in ( + ("--size", size), + ("--region", regions), + ("--data-center", data_center), + ("--type", volume_type), + ) + if value + ] + if given: + raise click.UsageError(f"{', '.join(given)} do not apply to global volumes") + body = {"name": name} + else: + if size is None: + raise click.UsageError("pass --size in GB for a network volume") + if regions and data_center: + raise click.UsageError("use either --region or --data-center") + volume_type = volume_type.upper() if volume_type else None + with api_errors(): + body = { + "name": name, + "size": size, + "dataCenter": _data_center(regions, data_center, volume_type), + } + if volume_type: + body["type"] = volume_type + + with api_errors(): + record = run_rest_request("POST", kind.path, json=body) or {} + if as_json: + print_json(record) + return + ui.success(f"created {kind.noun} [white]{record.get('id')}[/white]") + _show(kind, record, as_json) + + +@volume.command(name="update") +@click.argument("volume_id", shell_complete=_complete_ids) +@click.option("--name", default=None, help="New name.") +@click.option( + "--size", + type=click.IntRange(MIN_SIZE, MAX_SIZE), + default=None, + help="New size in GB; network volumes can only grow.", +) +@global_option +@json_option +def volume_update(volume_id, name, size, global_volume, as_json): + """Rename and/or resize a volume.""" + from runpod.rp_cli import console as ui + + kind = _kind(global_volume) + if global_volume and size is not None: + raise click.UsageError("global volumes have no size to change") + if name is None and size is None: + raise click.UsageError("pass --name and/or --size") + body = { + key: value + for key, value in (("name", name), ("size", size)) + if value is not None + } + with api_errors(): + if size is not None: + current = _get(kind, volume_id).get("size") or 0 + _check_growth(volume_id, current, size, strict=False) + record = run_rest_request("PATCH", f"{kind.path}/{volume_id}", json=body) or {} + if as_json: + print_json(record) + return + ui.success(f"updated {kind.noun} [white]{volume_id}[/white]") + _show(kind, record, as_json) + + +@volume.command(name="expand") +@click.argument("volume_id", shell_complete=_complete_ids) +@click.option( + "--size", + required=True, + metavar="N|+N", + help="New total size in GB, or +N to add N GB.", +) +@json_option +def volume_expand(volume_id, size, as_json): + """Grow a network volume; storage can never shrink.""" + from runpod.rp_cli import console as ui + + relative, amount = _parse_growth(size) + with api_errors(): + current = _get(NETWORK, volume_id).get("size") or 0 + target = current + amount if relative else amount + _check_growth(volume_id, current, target, strict=True) + record = ( + run_rest_request( + "PATCH", f"{NETWORK.path}/{volume_id}", json={"size": target} + ) + or {} + ) + if as_json: + print_json(record) + return + ui.success( + f"expanded [white]{volume_id}[/white] from {current} GB to " + f"{record.get('size', target)} GB" + ) + + +@volume.command(name="delete") +@click.argument("volume_id", shell_complete=_complete_ids) +@click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.") +@global_option +def volume_delete(volume_id, yes, global_volume): + """Delete a volume and everything stored on it.""" + from runpod.rp_cli import console as ui + + kind = _kind(global_volume) + with api_errors(): + if not yes: + record = _get(kind, volume_id) + confirm( + f"delete {kind.noun} {volume_id} ({_label(kind, record)}) " + "and all its data?", + yes, + ) + run_rest_request("DELETE", f"{kind.path}/{volume_id}") + ui.success(f"deleted {kind.noun} [white]{volume_id}[/white]") diff --git a/tests/test_cli/test_volume_command.py b/tests/test_cli/test_volume_command.py new file mode 100644 index 00000000..0aa0c8ad --- /dev/null +++ b/tests/test_cli/test_volume_command.py @@ -0,0 +1,449 @@ +"""tests for the rp volume command group.""" + +import json +from unittest.mock import patch + +import pytest +from click.testing import CliRunner + +from runpod import error +from runpod.rp_cli.main import cli + +VOLUME = { + "id": "vol1", + "name": "models", + "size": 100, + "dataCenter": "EU-RO-1", + "type": "STANDARD", +} +GLOBAL_VOLUME = { + "id": "gsb_abc", + "name": "shared", + "createdAt": "2026-09-24T00:00:00Z", + "createdByUserId": "user_123", +} +DATACENTERS = [ + {"id": "EU-RO-1", "region": "EUROPE", "networkVolumeTypes": ["STANDARD"]}, + {"id": "EU-FR-1", "region": "EUROPE", "networkVolumeTypes": ["HIGH_PERFORMANCE"]}, + { + "id": "EU-CZ-1", + "region": "EUROPE", + "networkVolumeTypes": ["STANDARD", "HIGH_PERFORMANCE"], + }, + {"id": "EU-SE-1", "region": "EUROPE", "networkVolumeTypes": []}, + {"id": "US-TX-3", "region": "NORTH_AMERICA", "networkVolumeTypes": ["STANDARD"]}, +] + + +def invoke(*args): + return CliRunner().invoke(cli, ["volume", *args]) + + +@pytest.fixture +def rest(): + with patch("runpod.rp_cli.volume.run_rest_request") as mock: + yield mock + + +@pytest.fixture +def listing(): + with patch("runpod.rp_cli.compute.list_all") as mock: + yield mock + + +@pytest.fixture +def datacenters(): + with patch( + "runpod.rp_cli.compute.list_datacenters", return_value=DATACENTERS + ) as mock: + yield mock + + +def test_volume_is_mounted(): + assert "volume" in cli.commands + + +class TestList: + def test_network_table(self, listing): + listing.return_value = [VOLUME] + result = invoke("list") + assert result.exit_code == 0, result.output + listing.assert_called_once_with("/v2/network-volumes", "networkVolumes") + for text in ( + "ID", + "DATA CENTER", + "TYPE", + "vol1", + "models", + "100 GB", + "EU-RO-1", + "standard", + ): + assert text in result.output + + def test_global_json(self, listing): + listing.return_value = [GLOBAL_VOLUME] + result = invoke("list", "--global", "--json") + assert result.exit_code == 0, result.output + listing.assert_called_once_with("/v2/global-volumes", "globalVolumes") + assert json.loads(result.output) == [GLOBAL_VOLUME] + + def test_global_table(self, listing): + listing.return_value = [GLOBAL_VOLUME] + result = invoke("list", "--global") + assert "gsb_abc" in result.output and "user_123" in result.output + + def test_empty(self, listing): + listing.return_value = [] + result = invoke("list") + assert result.exit_code == 0 + assert "no network volumes" in result.output + + def test_api_error(self, listing): + listing.side_effect = error.AuthenticationError("bad key") + result = invoke("list") + assert result.exit_code == 1 + assert "bad key" in result.output + + +class TestGet: + def test_fields(self, rest): + rest.return_value = VOLUME + result = invoke("get", "vol1") + assert result.exit_code == 0, result.output + rest.assert_called_once_with("GET", "/v2/network-volumes/vol1") + assert "data center" in result.output and "EU-RO-1" in result.output + assert "standard" in result.output + + def test_global_json(self, rest): + rest.return_value = GLOBAL_VOLUME + result = invoke("get", "gsb_abc", "--global", "--json") + rest.assert_called_once_with("GET", "/v2/global-volumes/gsb_abc") + assert json.loads(result.output) == GLOBAL_VOLUME + + def test_picker(self, rest, listing): + listing.return_value = [VOLUME] + rest.return_value = VOLUME + with patch("runpod.rp_cli.volume.pick", return_value="vol1") as pick: + result = invoke("get") + assert result.exit_code == 0, result.output + choices = pick.call_args.args[1] + assert choices == [("vol1", "models 100 GB EU-RO-1 standard")] + rest.assert_called_once_with("GET", "/v2/network-volumes/vol1") + + def test_missing_id_without_tty(self, rest, listing): + listing.return_value = [VOLUME] + result = invoke("get") + assert result.exit_code == 2 + assert "pass VOLUME_ID" in result.output + rest.assert_not_called() + + def test_not_found(self, rest): + rest.side_effect = error.QueryError("network volume not found", status_code=404) + result = invoke("get", "nope") + assert result.exit_code == 1 + assert "not found" in result.output + + +class TestCreate: + def test_exact_data_center(self, rest, datacenters): + rest.return_value = VOLUME + result = invoke( + "create", + "--name", + "models", + "--size", + "100", + "--data-center", + "EU-RO-1", + "--type", + "standard", + ) + assert result.exit_code == 0, result.output + datacenters.assert_not_called() + rest.assert_called_once_with( + "POST", + "/v2/network-volumes", + json={ + "name": "models", + "size": 100, + "dataCenter": "EU-RO-1", + "type": "STANDARD", + }, + ) + assert "created network volume" in result.output and "vol1" in result.output + + def test_region_narrows_by_type(self, rest, datacenters): + rest.return_value = VOLUME + with patch("runpod.rp_cli.volume.pick", return_value="EU-CZ-1") as pick: + result = invoke( + "create", + "--name", + "m", + "--size", + "10", + "--region", + "eu", + "--type", + "high_performance", + "--json", + ) + assert result.exit_code == 0, result.output + assert [value for value, _ in pick.call_args.args[1]] == ["EU-FR-1", "EU-CZ-1"] + assert rest.call_args.kwargs["json"] == { + "name": "m", + "size": 10, + "dataCenter": "EU-CZ-1", + "type": "HIGH_PERFORMANCE", + } + assert json.loads(result.output) == VOLUME + + def test_region_with_one_candidate(self, rest, datacenters): + rest.return_value = VOLUME + result = invoke("create", "--name", "m", "--size", "10", "--region", "na") + assert result.exit_code == 0, result.output + assert rest.call_args.kwargs["json"] == { + "name": "m", + "size": 10, + "dataCenter": "US-TX-3", + } + + def test_ambiguous_region_without_tty_lists_candidates(self, rest, datacenters): + result = invoke("create", "--name", "m", "--size", "10", "--region", "eu") + assert result.exit_code == 2 + assert "EU-RO-1, EU-FR-1, EU-CZ-1" in result.output + assert "EU-SE-1" not in result.output + assert "--data-center" in result.output + rest.assert_not_called() + + def test_no_region_picks_from_every_data_center(self, rest, datacenters): + rest.return_value = VOLUME + with patch("runpod.rp_cli.volume.pick", return_value="US-TX-3") as pick: + result = invoke( + "create", "--name", "m", "--size", "10", "--type", "standard" + ) + assert result.exit_code == 0, result.output + labels = dict(pick.call_args.args[1]) + assert set(labels) == {"EU-RO-1", "EU-CZ-1", "US-TX-3"} + assert labels["EU-CZ-1"] == "EU-CZ-1 europe standard, high_performance" + + def test_region_without_type_support(self, rest, datacenters): + result = invoke( + "create", + "--name", + "m", + "--size", + "10", + "--region", + "na", + "--type", + "high_performance", + ) + assert result.exit_code == 1 + assert "no data center in region na offers high_performance" in result.output + rest.assert_not_called() + + def test_unknown_region(self, rest, datacenters): + result = invoke("create", "--name", "m", "--size", "10", "--region", "mars") + assert result.exit_code == 1 + assert "unknown region 'mars'" in result.output + + def test_requires_size(self, rest): + result = invoke("create", "--name", "m", "--data-center", "EU-RO-1") + assert result.exit_code == 2 + assert "--size" in result.output + rest.assert_not_called() + + def test_size_bounds(self, rest): + result = invoke( + "create", "--name", "m", "--size", "5000", "--data-center", "EU-RO-1" + ) + assert result.exit_code == 2 + rest.assert_not_called() + + def test_region_and_data_center_conflict(self, rest): + result = invoke( + "create", + "--name", + "m", + "--size", + "10", + "--region", + "eu", + "--data-center", + "EU-RO-1", + ) + assert result.exit_code == 2 + assert "either --region or --data-center" in result.output + + def test_global(self, rest): + rest.return_value = GLOBAL_VOLUME + result = invoke("create", "--name", "shared", "--global") + assert result.exit_code == 0, result.output + rest.assert_called_once_with( + "POST", "/v2/global-volumes", json={"name": "shared"} + ) + assert "created global volume" in result.output and "gsb_abc" in result.output + + def test_global_rejects_network_flags(self, rest): + result = invoke( + "create", "--name", "s", "--global", "--size", "10", "--region", "eu" + ) + assert result.exit_code == 2 + assert "--size, --region do not apply to global volumes" in result.output + rest.assert_not_called() + + +class TestUpdate: + def test_rename(self, rest): + rest.return_value = dict(VOLUME, name="renamed") + result = invoke("update", "vol1", "--name", "renamed") + assert result.exit_code == 0, result.output + rest.assert_called_once_with( + "PATCH", "/v2/network-volumes/vol1", json={"name": "renamed"} + ) + assert "renamed" in result.output + + def test_rename_and_resize(self, rest): + rest.side_effect = [VOLUME, dict(VOLUME, name="r", size=200)] + result = invoke("update", "vol1", "--name", "r", "--size", "200", "--json") + assert result.exit_code == 0, result.output + assert rest.call_args.args == ("PATCH", "/v2/network-volumes/vol1") + assert rest.call_args.kwargs["json"] == {"name": "r", "size": 200} + assert json.loads(result.output)["size"] == 200 + + def test_same_size_is_allowed(self, rest): + rest.side_effect = [VOLUME, VOLUME] + result = invoke("update", "vol1", "--size", "100") + assert result.exit_code == 0, result.output + + def test_rejects_shrinking(self, rest): + rest.return_value = VOLUME + result = invoke("update", "vol1", "--size", "50") + assert result.exit_code == 2 + assert "can only grow: vol1 is 100 GB, requested 50 GB" in result.output + rest.assert_called_once_with("GET", "/v2/network-volumes/vol1") + + def test_requires_a_change(self, rest): + result = invoke("update", "vol1") + assert result.exit_code == 2 + assert "--name and/or --size" in result.output + + def test_global_rename(self, rest): + rest.return_value = dict(GLOBAL_VOLUME, name="renamed") + result = invoke("update", "gsb_abc", "--global", "--name", "renamed") + assert result.exit_code == 0, result.output + rest.assert_called_once_with( + "PATCH", "/v2/global-volumes/gsb_abc", json={"name": "renamed"} + ) + + def test_global_rejects_size(self, rest): + result = invoke("update", "gsb_abc", "--global", "--size", "20") + assert result.exit_code == 2 + assert "global volumes have no size" in result.output + rest.assert_not_called() + + +class TestExpand: + def test_absolute(self, rest): + rest.side_effect = [VOLUME, dict(VOLUME, size=250)] + result = invoke("expand", "vol1", "--size", "250") + assert result.exit_code == 0, result.output + rest.assert_called_with("PATCH", "/v2/network-volumes/vol1", json={"size": 250}) + assert "from 100 GB to 250 GB" in result.output + + def test_relative(self, rest): + rest.side_effect = [VOLUME, dict(VOLUME, size=150)] + result = invoke("expand", "vol1", "--size", "+50", "--json") + assert result.exit_code == 0, result.output + rest.assert_called_with("PATCH", "/v2/network-volumes/vol1", json={"size": 150}) + assert json.loads(result.output)["size"] == 150 + + @pytest.mark.parametrize("size", ["50", "100", "+0"]) + def test_rejects_shrinking_or_no_growth(self, rest, size): + rest.return_value = VOLUME + result = invoke("expand", "vol1", "--size", size) + assert result.exit_code == 2 + assert "network volumes can only grow: vol1 is 100 GB" in result.output + assert rest.call_count == 1 + + def test_rejects_past_maximum(self, rest): + rest.return_value = VOLUME + result = invoke("expand", "vol1", "--size", "+4000") + assert result.exit_code == 2 + assert "at most 4096 GB, requested 4100 GB" in result.output + assert rest.call_count == 1 + + @pytest.mark.parametrize("size", ["abc", "-5", "1.5"]) + def test_rejects_malformed_size(self, rest, size): + result = invoke("expand", "vol1", "--size", size) + assert result.exit_code == 2 + assert "is not a size" in result.output + rest.assert_not_called() + + +class TestDelete: + def test_yes_skips_confirmation(self, rest): + result = invoke("delete", "vol1", "--yes") + assert result.exit_code == 0, result.output + rest.assert_called_once_with("DELETE", "/v2/network-volumes/vol1") + assert "deleted network volume" in result.output + + def test_confirms_interactively(self, rest): + rest.return_value = VOLUME + with patch("runpod.rp_cli.output.interactive", return_value=True): + result = CliRunner().invoke(cli, ["volume", "delete", "vol1"], input="y\n") + assert result.exit_code == 0, result.output + assert "models 100 GB EU-RO-1 standard" in result.output + rest.assert_called_with("DELETE", "/v2/network-volumes/vol1") + + def test_declined(self, rest): + rest.return_value = VOLUME + with patch("runpod.rp_cli.output.interactive", return_value=True): + result = CliRunner().invoke(cli, ["volume", "delete", "vol1"], input="n\n") + assert result.exit_code == 1 + assert all(call.args[0] == "GET" for call in rest.call_args_list) + + def test_requires_yes_without_tty(self, rest): + rest.return_value = VOLUME + result = invoke("delete", "vol1") + assert result.exit_code == 2 + assert "--yes" in result.output + assert all(call.args[0] == "GET" for call in rest.call_args_list) + + def test_global(self, rest): + result = invoke("delete", "gsb_abc", "--global", "--yes") + assert result.exit_code == 0, result.output + rest.assert_called_once_with("DELETE", "/v2/global-volumes/gsb_abc") + + def test_mounted_global_volume(self, rest): + rest.side_effect = error.QueryError("volume is still mounted", status_code=409) + result = invoke("delete", "gsb_abc", "--global", "--yes") + assert result.exit_code == 1 + assert "still mounted" in result.output + + +class TestCompletion: + def _complete(self, args, incomplete): + from click.shell_completion import ShellComplete + + return [ + item.value + for item in ShellComplete(cli, {}, "rp", "_RP_COMPLETE").get_completions( + args, incomplete + ) + ] + + def test_network_ids(self, listing): + listing.return_value = [VOLUME, dict(VOLUME, id="other")] + assert self._complete(["volume", "get"], "v") == ["vol1"] + listing.assert_called_with("/v2/network-volumes", "networkVolumes") + + def test_global_ids(self, listing): + listing.return_value = [GLOBAL_VOLUME] + assert self._complete(["volume", "delete", "--global"], "") == ["gsb_abc"] + listing.assert_called_with("/v2/global-volumes", "globalVolumes") + + def test_errors_complete_nothing(self, listing): + listing.side_effect = error.AuthenticationError("no key") + assert self._complete(["volume", "update"], "") == [] From 232e7cb4cc40341da43b9c7db955b07347243d3b Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:44:40 -0700 Subject: [PATCH 04/25] feat: add rp pod lifecycle commands --- runpod/rp_cli/main.py | 4 +- runpod/rp_cli/pod.py | 825 +++++++++++++++++++++++++++++ tests/test_cli/test_pod_command.py | 765 ++++++++++++++++++++++++++ 3 files changed, 1592 insertions(+), 2 deletions(-) create mode 100644 runpod/rp_cli/pod.py create mode 100644 tests/test_cli/test_pod_command.py diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index d8be1833..7387e564 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -1036,11 +1036,11 @@ def undeploy(app_name, env_name, yes): 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.pod import pod - for command in (ssh_cli, pod_cli, catalog): + for command in (ssh_cli, pod, catalog): cli.add_command(command) diff --git a/runpod/rp_cli/pod.py b/runpod/rp_cli/pod.py new file mode 100644 index 00000000..35f62d5b --- /dev/null +++ b/runpod/rp_cli/pod.py @@ -0,0 +1,825 @@ +"""`rp pod`: create, inspect, and manage pods.""" + +import shlex +import time +from typing import Optional + +import click + +from runpod.api.rest import run_rest_request +from runpod.cli.groups.pod.commands import sync_pods +from runpod.rp_cli import compute +from runpod.rp_cli.output import ( + api_errors, + completer, + confirm, + interactive, + json_option, + pick, + print_fields, + print_json, + print_table, +) + +STATUSES = ["PROVISIONING", "STARTING", "RUNNING", "EXITED", "ERROR", "TERMINATED"] +DEFAULT_MOUNT_PATH = "/workspace" +WAIT_TIMEOUT = 600 +WAIT_INTERVAL = 5 + + +def _list_pods(include_cluster_pods: bool = False) -> list[dict]: + params = {"includeClusterPods": "true"} if include_cluster_pods else None + return compute.list_all("/v2/pods", "pods", params) + + +def _get_pod(pod_id: str) -> dict: + return run_rest_request("GET", f"/v2/pods/{pod_id}") or {} + + +def _pod_choices(action: Optional[str] = None) -> list[tuple[str, str]]: + return [ + (pod["id"], f"{pod.get('name')} {(pod.get('status') or '').lower()}") + for pod in _list_pods() + if action is None or action in (pod.get("actions") or []) + ] + + +def _resolve_pod_id(pod_id: Optional[str], action: Optional[str] = None) -> str: + """the given pod id, or one picked interactively.""" + if pod_id: + return pod_id + choices = [] + if interactive(): + with api_errors(): + choices = _pod_choices(action) + return pick("pod", choices, "pass a POD_ID") + + +_complete_pods = completer(_pod_choices) +_complete_datacenters = completer( + lambda: [ + (dc["id"], (dc.get("region") or "").lower()) + for dc in compute.list_datacenters() + ] +) +_complete_network_volumes = completer( + lambda: [ + (volume["id"], f"{volume.get('name')} {volume.get('dataCenter')}") + for volume in compute.list_all("/v2/network-volumes", "networkVolumes") + ] +) +_complete_registries = completer( + lambda: [ + (registry["name"], registry["id"]) + for registry in compute.list_all("/v2/registries", "registries") + ] +) + + +def _say(message: str, as_json: bool, warning: bool = False) -> None: + """a status line on stdout, or stderr when stdout carries json.""" + if as_json: + click.echo(message, err=True) + return + from rich.markup import escape + + from runpod.rp_cli import console as ui + + (ui.warn if warning else ui.info)(escape(message)) + + +def _hardware(pod: dict) -> Optional[str]: + gpu = pod.get("gpu") + if gpu: + return f"{gpu.get('count') or 1}x {gpu['id']}" + cpu = pod.get("cpu") + if cpu: + return f"{cpu['id']} {cpu.get('vcpuCount')} vcpu" + return None + + +def _cost(pod: dict) -> Optional[str]: + return f"${pod['cost']:.2f}/hr" if pod.get("cost") else None + + +def _mounts(pod: dict) -> list[str]: + mounts = pod.get("mounts") or {} + described = [] + persistent = mounts.get("persistent") + if persistent: + described.append(f"persistent {persistent['size']}GB at {persistent['path']}") + for mount in mounts.get("network") or []: + described.append(f"network {mount['volumeId']} at {mount.get('path')}") + for mount in mounts.get("global") or []: + described.append(f"global {mount['volumeId']} at {mount.get('path') or '-'}") + return described + + +def _urls(pod: dict) -> list[str]: + urls = [] + for port in pod.get("ports") or []: + number, _, protocol = port.partition("/") + if protocol == "http": + urls.append(f"https://{pod['id']}-{number}.proxy.runpod.net") + return urls + + +def _tcp(pod: dict) -> list[str]: + ports = (pod.get("runtime") or {}).get("ports") or [] + return [ + f"{port['ip']}:{port['public']} -> {port['private']}" + for port in ports + if port.get("type") == "tcp" and port.get("ip") and port.get("public") + ] + + +def _global_network(pod: dict) -> object: + networking = pod.get("globalNetworking") or {} + if not networking.get("enabled"): + return False + return networking.get("internalDns") or networking.get("ip") or True + + +def _resources(pod: dict) -> Optional[str]: + hardware = pod.get("gpu") or pod.get("cpu") or {} + parts = [] + if hardware.get("vcpuCount"): + parts.append(f"{hardware['vcpuCount']:g} vcpu") + if hardware.get("memory"): + parts.append(f"{hardware['memory']:g}GB ram") + return ", ".join(parts) or None + + +_LIST_COLUMNS = [ + ("ID", lambda pod: pod["id"]), + ("NAME", lambda pod: pod.get("name")), + ("STATUS", lambda pod: (pod.get("status") or "").lower()), + ("HARDWARE", _hardware), + ("DATA CENTER", lambda pod: pod.get("dataCenterId")), + ("COST", _cost), + ("IMAGE", lambda pod: pod.get("image")), +] + +_FIELDS = [ + ("id", lambda pod: pod["id"]), + ("name", lambda pod: pod.get("name")), + ("status", lambda pod: (pod.get("status") or "").lower()), + ("image", lambda pod: pod.get("image")), + ("hardware", _hardware), + ("resources", _resources), + ("cloud", lambda pod: (pod.get("cloud") or "").lower()), + ("data center", lambda pod: pod.get("dataCenterId")), + ("cuda", lambda pod: pod.get("cudaVersion")), + ("cost", _cost), + ("disk", lambda pod: f"{pod['disk']}GB" if pod.get("disk") else None), + ("volumes", _mounts), + ("ports", lambda pod: pod.get("ports")), + ("env", lambda pod: sorted(pod.get("env") or {})), + ("command", lambda pod: pod.get("args")), + ("locked", lambda pod: pod.get("locked")), + ("global network", _global_network), + ("cluster", lambda pod: (pod.get("cluster") or {}).get("id")), + ("created", lambda pod: pod.get("createdAt")), + ("started", lambda pod: pod.get("startedAt")), + ("actions", lambda pod: pod.get("actions")), + ("ssh", lambda pod: ((pod.get("ssh") or {}).get("proxy") or {}).get("command")), + ( + "ssh direct", + lambda pod: ((pod.get("ssh") or {}).get("direct") or {}).get("command"), + ), + ("urls", _urls), + ("tcp", _tcp), +] + + +def _show(pod: dict, as_json: bool) -> None: + if as_json: + print_json(pod) + return + print_fields(pod, _FIELDS) + + +def _wait_running(pod_id: str) -> dict: + """poll a pod until it is running.""" + from runpod.rp_cli import console as ui + + deadline = time.monotonic() + WAIT_TIMEOUT + with ui.console.status(f"[dim]waiting for {pod_id} to start ...[/dim]"): + while True: + pod = _get_pod(pod_id) + status = pod.get("status") + if status == "RUNNING": + return pod + if status in ("ERROR", "TERMINATED"): + raise click.ClickException(f"pod {pod_id} is {status.lower()}") + if time.monotonic() >= deadline: + raise click.ClickException( + f"pod {pod_id} is still {(status or 'unknown').lower()} " + f"after {WAIT_TIMEOUT}s" + ) + time.sleep(WAIT_INTERVAL) + + +def _parse_env(values: tuple[str, ...]) -> dict[str, str]: + env = {} + for value in values: + key, sep, setting = value.partition("=") + if not sep or not key: + raise click.BadParameter(f"'{value}' is not KEY=VALUE", param_hint="--env") + env[key] = setting + return env + + +def _split_mount(value: str, flag: str) -> tuple[str, Optional[str]]: + head, sep, path = value.partition(":") + if not head or (sep and not path): + raise click.BadParameter( + f"'{value}' should look like VALUE[:PATH]", param_hint=flag + ) + return head, path or None + + +def _mount_path(current: dict, kind: str) -> Optional[str]: + entries = current.get(kind) or [] + return entries[0].get("path") if entries else None + + +def _mounts_body(options: dict, current: dict) -> dict: + """the mounts object for the volume flags, defaulting paths from current.""" + mounts: dict = {} + if options["persistent_volume"]: + size, path = _split_mount(options["persistent_volume"], "--persistent-volume") + if not size.isdigit(): + raise click.BadParameter( + "size is in GB, like 50 or 50:/workspace", + param_hint="--persistent-volume", + ) + mounts["persistent"] = { + "size": int(size), + "path": path + or (current.get("persistent") or {}).get("path") + or DEFAULT_MOUNT_PATH, + } + if options["network_volume"]: + volume_id, path = _split_mount(options["network_volume"], "--network-volume") + mounts["network"] = [ + { + "volumeId": volume_id, + "path": path or _mount_path(current, "network") or DEFAULT_MOUNT_PATH, + } + ] + if options["global_volume"]: + volume_id, path = _split_mount(options["global_volume"], "--global-volume") + path = path or _mount_path(current, "global") + mounts["global"] = [{"volumeId": volume_id, **({"path": path} if path else {})}] + return mounts + + +def _resolve_registry(value: str) -> Optional[str]: + if not value: + return None + for registry in compute.list_all("/v2/registries", "registries"): + if value in (registry["id"], registry.get("name")): + return registry["id"] + raise compute.ResolutionError( + f"unknown registry credential '{value}'. list them with rp registry list" + ) + + +def _container_body(options: dict) -> dict: + """request fields shared by create and update, except env and mounts.""" + body: dict = {} + if options["image"]: + body["image"] = options["image"] + if options["template"]: + body["templateId"] = options["template"] + if options["args"] is not None: + body["args"] = options["args"] + if options["cmd"] is not None: + body["cmd"] = shlex.split(options["cmd"]) + if options["entrypoint"] is not None: + body["entrypoint"] = shlex.split(options["entrypoint"]) + if options["disk"]: + body["disk"] = options["disk"] + if options["ports"]: + body["ports"] = list(options["ports"]) + if options["registry"] is not None: + body["registry"] = _resolve_registry(options["registry"]) + if options["global_networking"] is not None: + body["globalNetworking"] = options["global_networking"] + return body + + +_CONTAINER_OPTIONS = [ + click.option( + "--image", default=None, help="Container image, like runpod/pytorch:latest." + ), + click.option( + "--template", + default=None, + help="Template id whose container settings to apply.", + ), + click.option( + "--env", + "-e", + multiple=True, + metavar="KEY=VALUE", + help="Environment variable. Repeatable.", + ), + click.option( + "--port", + "-p", + "ports", + multiple=True, + metavar="PORT/PROTO", + help="Exposed port like 8888/http or 22/tcp. Repeatable.", + ), + click.option( + "--disk", type=click.IntRange(min=1), default=None, help="Container disk in GB." + ), + click.option( + "--persistent-volume", + metavar="GB[:PATH]", + default=None, + help=f"Host-local persistent volume (path defaults to {DEFAULT_MOUNT_PATH}).", + ), + click.option( + "--network-volume", + metavar="ID[:PATH]", + default=None, + shell_complete=_complete_network_volumes, + help=f"Network volume to mount (path defaults to {DEFAULT_MOUNT_PATH}).", + ), + click.option( + "--global-volume", + metavar="ID[:PATH]", + default=None, + help="Global volume to mount.", + ), + click.option( + "--args", + default=None, + help="Container start command as one raw string ('' clears).", + ), + click.option( + "--cmd", + default=None, + help="Container CMD as a shell string, split into exec form.", + ), + click.option( + "--entrypoint", + default=None, + help="Container ENTRYPOINT as a shell string, split into exec form.", + ), + click.option( + "--registry", + default=None, + shell_complete=_complete_registries, + help="Registry credential name or id for private images ('' clears).", + ), + click.option( + "--global-networking/--no-global-networking", + default=None, + help="Private networking across data centers (NVIDIA GPUs only).", + ), +] + + +def _container_options(command): + for option in reversed(_CONTAINER_OPTIONS): + command = option(command) + return command + + +def _choose_gpu( + query: str, + count: int, + cloud: str, + cuda: tuple[Optional[str], list[str]], + datacenter_ids: list[str], +) -> tuple[dict, str, int]: + """the gpu type to request, its stock level, and how many types matched. + + a query naming several types picks the one with the most stock in the + requested data centers, then the lowest price. + """ + min_cuda, exact_cuda = cuda + gpus = compute.list_gpus( + product="POD", + count=count, + cloud=cloud.upper(), + cuda_versions=exact_cuda, + min_cuda_version=min_cuda, + ) + matches = compute.resolve_gpus(query, gpus) + allowed = set(datacenter_ids) + + def stock(gpu: dict) -> str: + if not allowed: + return gpu.get("availability") or "NONE" + return compute.best_availability( + dc.get("availability") + for dc in gpu.get("dataCenters") or [] + if dc["id"] in allowed + ) + + offered = [ + gpu + for gpu in matches + if len(matches) == 1 + or ( + gpu.get(cloud, True) + and (gpu.get("maxCount") or {}).get(cloud, count) >= count + ) + ] + if not offered: + names = ", ".join(gpu.get("name") or gpu["id"] for gpu in matches) + raise compute.ResolutionError( + f"no '{query}' gpu ({names}) offers {count} per pod on {cloud} cloud" + ) + best = min( + offered, + key=lambda gpu: ( + -compute.availability_rank(stock(gpu)), + (gpu.get("price") or {}).get(cloud) or float("inf"), + gpu.get("name") or gpu["id"], + ), + ) + return best, stock(best), len(matches) + + +def _datacenter_ids(regions, data_centers, network_volume) -> list[str]: + ids = ( + compute.resolve_regions(regions, compute.list_datacenters()) if regions else [] + ) + ids += [dc for dc in data_centers if dc not in ids] + if not ids and network_volume: + volume_id = network_volume.partition(":")[0] + volume = run_rest_request("GET", f"/v2/network-volumes/{volume_id}") or {} + if volume.get("dataCenter"): + ids = [volume["dataCenter"]] + return ids + + +@click.group("pod") +def pod(): + """Create and manage pods.""" + + +@pod.command(name="list") +@click.option("--name", default=None, help="Only pods whose name contains this text.") +@click.option( + "--status", + "statuses", + multiple=True, + type=click.Choice(STATUSES, case_sensitive=False), + help="Only pods in this status. Repeatable.", +) +@click.option( + "--include-cluster-pods", is_flag=True, help="Include cluster member pods." +) +@json_option +def list_pods(name, statuses, include_cluster_pods, as_json): + """List your pods.""" + with api_errors(): + pods = _list_pods(include_cluster_pods) + if name: + pods = [p for p in pods if name.lower() in (p.get("name") or "").lower()] + if statuses: + wanted = {status.upper() for status in statuses} + pods = [p for p in pods if p.get("status") in wanted] + if as_json: + print_json(pods) + return + print_table(pods, _LIST_COLUMNS, "no pods. create one with rp pod create") + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@json_option +def get(pod_id, as_json): + """Show a pod with its connection details.""" + pod_id = _resolve_pod_id(pod_id) + with api_errors(): + _show(_get_pod(pod_id), as_json) + + +@pod.command() +@click.argument("name", required=False) +@click.option( + "--gpu", default=None, metavar="TYPE[:COUNT]", help="GPU type like 4090 or h100:2." +) +@click.option( + "--gpu-count", type=click.IntRange(min=1), default=None, help="GPUs per pod." +) +@click.option( + "--cpu", default=None, help="CPU flavor like cpu3c or cpu5g, for a CPU pod." +) +@click.option( + "--vcpu", + type=click.IntRange(min=2), + default=None, + help="vCPUs for a CPU pod (defaults to the flavor's minimum).", +) +@click.option( + "--cloud", + type=click.Choice(["secure", "community"], case_sensitive=False), + default=None, + help="Cloud tier (the api defaults to secure).", +) +@click.option( + "--region", + "regions", + multiple=True, + help="Region (eu, na, asia), id prefix (us, eu-ro), or data center id. Repeatable.", +) +@click.option( + "--data-center", + "data_centers", + multiple=True, + shell_complete=_complete_datacenters, + help="Exact data center id like US-TX-3. Repeatable.", +) +@click.option( + "--cuda", + multiple=True, + help="Host CUDA: '>=12.4' for a floor, 12.4 or 12.4,12.8 for exact versions.", +) +@click.option( + "--min-ram-per-gpu", + type=click.IntRange(min=1), + default=None, + help="Minimum host RAM in GB per GPU.", +) +@click.option( + "--min-vcpu-per-gpu", + type=click.IntRange(min=1), + default=None, + help="Minimum host vCPUs per GPU.", +) +@_container_options +@click.option("--ssh", is_flag=True, help="Start sshd with your account's SSH keys.") +@click.option( + "--jupyter", is_flag=True, help="Start JupyterLab (expose 8888/http to reach it)." +) +@click.option("--wait", is_flag=True, help="Wait until the pod is running.") +@json_option +def create( + name, + gpu, + gpu_count, + cpu, + vcpu, + cloud, + regions, + data_centers, + cuda, + min_ram_per_gpu, + min_vcpu_per_gpu, + ssh, + jupyter, + wait, + as_json, + **container, +): # pylint: disable=too-many-arguments,too-many-locals + """Create a pod named NAME. + + GPU types are human-readable (4090, h100, a100-sxm). When a name + matches several types, the one with the most stock in the requested + regions wins, then the cheapest. + + \b + rp pod create trainer --gpu h100:2 --image runpod/pytorch:latest --region eu --ssh + rp pod create box --cpu cpu3c --vcpu 4 --image ubuntu:24.04 + rp pod create nb --gpu 4090 --image runpod/pytorch:latest --jupyter -p 8888/http + """ + if gpu and cpu: + raise click.UsageError("use either --gpu or --cpu") + if not gpu and not cpu: + raise click.UsageError("pass --gpu (like 4090 or h100:2) or --cpu (like cpu3c)") + if cpu and (gpu_count or cuda or min_ram_per_gpu or min_vcpu_per_gpu): + raise click.UsageError( + "--gpu-count, --cuda, --min-ram-per-gpu, and --min-vcpu-per-gpu apply to gpu pods" + ) + if vcpu and not cpu: + raise click.UsageError("--vcpu applies to cpu pods") + if not container["image"] and not container["template"]: + raise click.UsageError("pass --image or --template") + if not name: + if not interactive(): + raise click.UsageError("pass a NAME") + name = click.prompt("pod name") + + with api_errors(): + body = {"name": name, **_container_body(container)} + env = _parse_env(container["env"]) + if env: + body["env"] = env + mounts = _mounts_body(container, {}) + if mounts: + body["mounts"] = mounts + datacenter_ids = _datacenter_ids( + regions, data_centers, container["network_volume"] + ) + if datacenter_ids: + body["dataCenterIds"] = datacenter_ids + if cloud: + body["cloud"] = cloud.upper() + + if gpu: + query, spec_count = compute.parse_gpu_spec(gpu) + if spec_count and gpu_count and spec_count != gpu_count: + raise click.UsageError( + f"--gpu {gpu} conflicts with --gpu-count {gpu_count}" + ) + count = spec_count or gpu_count or 1 + min_cuda, exact_cuda = compute.parse_cuda(cuda) + tier = (cloud or "secure").lower() + chosen, stock, matched = _choose_gpu( + query, count, tier, (min_cuda, exact_cuda), datacenter_ids + ) + label = chosen.get("name") or chosen["id"] + price = (chosen.get("price") or {}).get(tier) + if matched > 1: + cost = f", ${price * count:.2f}/hr" if price else "" + _say( + f"'{query}' matches {matched} gpu types; using {label} " + f"({chosen['id']}): {'no' if stock == 'NONE' else stock.lower()} stock{cost}", + as_json, + ) + if stock == "NONE": + _say( + f"the catalog reports no {tier} stock for {count}x {label} " + "in the requested data centers; creation may fail", + as_json, + warning=True, + ) + gpu_body = {"id": chosen["id"], "count": count} + if min_cuda: + gpu_body["minCudaVersion"] = min_cuda + if exact_cuda: + gpu_body["allowedCudaVersions"] = exact_cuda + if min_ram_per_gpu: + gpu_body["minRamPerGpu"] = min_ram_per_gpu + if min_vcpu_per_gpu: + gpu_body["minVcpuCountPerGpu"] = min_vcpu_per_gpu + body["gpu"] = gpu_body + else: + flavor = compute.resolve_cpu( + cpu, compute.list_cpus(product="POD", vcpu_count=vcpu) + ) + body["cpu"] = { + "id": flavor["id"], + "vcpuCount": vcpu or (flavor.get("vcpu") or {}).get("min") or 2, + } + + if ssh: + body["startSsh"] = True + keys = (run_rest_request("GET", "/v2/account/ssh-keys") or {}).get("keys") + if not keys: + _say( + "your account has no ssh keys, so --ssh has no effect. " + "add one with rp ssh add", + as_json, + warning=True, + ) + if jupyter: + body["startJupyter"] = True + + created = run_rest_request("POST", "/v2/pods", json=body) or {} + _say(f"created pod {created.get('name')} ({created.get('id')})", as_json) + if wait: + created = _wait_running(created["id"]) + _show(created, as_json) + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@click.option("--name", default=None, help="New pod name.") +@_container_options +@click.option( + "--unset-env", + multiple=True, + metavar="KEY", + help="Remove an environment variable. Repeatable.", +) +@click.option( + "--lock/--unlock", + "locked", + default=None, + help="Lock the pod against stop and reset.", +) +@json_option +def update(pod_id, name, unset_env, locked, as_json, **container): + """Change a pod's settings. + + --env and --unset-env edit the existing environment, --port replaces + the exposed ports, and volume flags restate the full mount (the + volume id cannot change; the path defaults to the current one). + Most changes apply on the pod's next start. + """ + requested = [name, locked, *unset_env, *container.values()] + if all(value in (None, ()) for value in requested): + raise click.UsageError("nothing to update. see rp pod update --help") + pod_id = _resolve_pod_id(pod_id) + + with api_errors(): + body = _container_body(container) + if name: + body["name"] = name + if locked is not None: + body["locked"] = locked + mounts_changed = any( + container[key] + for key in ("persistent_volume", "network_volume", "global_volume") + ) + if container["env"] or unset_env or mounts_changed: + current = _get_pod(pod_id) + if container["env"] or unset_env: + env = {**(current.get("env") or {}), **_parse_env(container["env"])} + for key in unset_env: + env.pop(key, None) + body["env"] = env + if mounts_changed: + body["mounts"] = _mounts_body(container, current.get("mounts") or {}) + updated = run_rest_request("PATCH", f"/v2/pods/{pod_id}", json=body) or {} + _say(f"updated pod {pod_id}", as_json) + _show(updated, as_json) + + +def _act( + pod_id: Optional[str], action: str, done: str, as_json: bool, wait: bool = False +): + pod_id = _resolve_pod_id(pod_id, action) + with api_errors(): + result = run_rest_request( + "POST", f"/v2/pods/{pod_id}/action", json={"action": action} + ) + _say(f"{done} pod {pod_id}", as_json) + if wait: + result = _wait_running(pod_id) + _show(result or {}, as_json) + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@click.option("--wait", is_flag=True, help="Wait until the pod is running.") +@json_option +def start(pod_id, wait, as_json): + """Start a stopped pod.""" + _act(pod_id, "start", "started", as_json, wait) + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@json_option +def stop(pod_id, as_json): + """Stop a running pod, keeping its volume.""" + _act(pod_id, "stop", "stopped", as_json) + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@click.option("--wait", is_flag=True, help="Wait until the pod is running.") +@json_option +def restart(pod_id, wait, as_json): + """Restart a running pod.""" + _act(pod_id, "restart", "restarted", as_json, wait) + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.") +def terminate(pod_id, yes): + """Terminate a pod permanently.""" + from runpod.rp_cli import console as ui + + pod_id = _resolve_pod_id(pod_id, "terminate") + confirm(f"terminate pod {pod_id}? its disk and persistent volume are deleted", yes) + with api_errors(): + run_rest_request("DELETE", f"/v2/pods/{pod_id}") + ui.success(f"terminated pod {pod_id}") + + +pod.add_command(terminate, "delete") + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +def connect(pod_id): + """Open an interactive terminal on a pod over SSH. + + Needs an SSH key on your account (rp ssh add) and 22/tcp exposed. + """ + from runpod.cli.utils import ssh_cmd + + pod_id = _resolve_pod_id(pod_id) + click.echo(f"connecting to pod {pod_id}...") + try: + ssh = ssh_cmd.SSHConnection(pod_id) + except (ValueError, TimeoutError) as exc: + raise click.ClickException(str(exc)) from exc + ssh.launch_terminal() + + +pod.add_command(sync_pods, "sync") diff --git a/tests/test_cli/test_pod_command.py b/tests/test_cli/test_pod_command.py new file mode 100644 index 00000000..b65fafcd --- /dev/null +++ b/tests/test_cli/test_pod_command.py @@ -0,0 +1,765 @@ +"""tests for rp pod.""" + +import json +from unittest.mock import MagicMock, patch + +import pytest +from click.testing import CliRunner + +from runpod import error +from runpod.rp_cli import pod as pod_module +from runpod.rp_cli.main import cli + +GPUS = [ + { + "id": "NVIDIA H100 80GB HBM3", + "name": "H100 SXM", + "price": {"secure": 3.99, "community": 2.69}, + "maxCount": {"secure": 8, "community": 1}, + "secure": True, + "community": True, + "availability": "LOW", + "dataCenters": [ + {"id": "US-TX-3", "availability": "LOW"}, + {"id": "EU-RO-1", "availability": "LOW"}, + ], + }, + { + "id": "NVIDIA H100 NVL", + "name": "H100 NVL", + "price": {"secure": 3.19, "community": 2.59}, + "maxCount": {"secure": 1, "community": 10}, + "secure": True, + "community": True, + "availability": "HIGH", + "dataCenters": [{"id": "US-TX-3", "availability": "HIGH"}], + }, + { + "id": "NVIDIA H100 PCIe", + "name": "H100 PCIe", + "price": {"secure": 2.39, "community": 1.99}, + "maxCount": {"secure": 8, "community": 8}, + "secure": True, + "community": True, + "availability": "LOW", + "dataCenters": [{"id": "US-TX-3", "availability": "LOW"}], + }, + { + "id": "NVIDIA GeForce RTX 4090", + "name": "RTX 4090", + "price": {"secure": 0.69, "community": 0.34}, + "maxCount": {"secure": 8, "community": 8}, + "secure": True, + "community": True, + "availability": "NONE", + }, +] + +CPUS = [ + { + "id": "cpu3c", + "name": "Compute-Optimized", + "group": "CPU3", + "vcpu": {"min": 2, "max": 32}, + }, + { + "id": "cpu5g", + "name": "General Purpose", + "group": "CPU5", + "vcpu": {"min": 4, "max": 32}, + }, +] + +DATACENTERS = [ + {"id": "EU-RO-1", "region": "EUROPE"}, + {"id": "EU-CZ-1", "region": "EUROPE"}, + {"id": "US-TX-3", "region": "NORTH_AMERICA"}, +] + +POD = { + "id": "pod123", + "name": "trainer", + "status": "RUNNING", + "actions": ["stop", "restart", "terminate"], + "image": "runpod/pytorch:latest", + "args": "", + "disk": 50, + "ports": ["8888/http", "22/tcp"], + "env": {"HF_TOKEN": "secret", "MODEL": "llama"}, + "mounts": {"network": [{"volumeId": "vol1", "path": "/data"}]}, + "gpu": {"id": "NVIDIA H100 80GB HBM3", "count": 2, "vcpuCount": 32, "memory": 250}, + "cloud": "SECURE", + "dataCenterId": "EU-RO-1", + "cudaVersion": "12.8", + "cost": 7.98, + "locked": False, + "globalNetworking": {"enabled": True, "internalDns": "pod123.runpod.internal"}, + "ssh": { + "proxy": {"command": "ssh pod123-abc@ssh.runpod.io"}, + "direct": {"command": "ssh root@1.2.3.4 -p 40022"}, + }, + "runtime": { + "ports": [ + {"private": 22, "public": 40022, "type": "tcp", "ip": "1.2.3.4"}, + {"private": 8888, "public": None, "type": "http", "ip": None}, + ] + }, + "createdAt": "2026-10-19T20:00:00Z", + "startedAt": "2026-10-19T20:01:00Z", +} + +STOPPED = { + "id": "pod456", + "name": "box", + "status": "EXITED", + "actions": ["start", "terminate"], + "image": "ubuntu:24.04", + "cpu": {"id": "cpu3c", "vcpuCount": 4, "memory": 8}, + "cost": 0, +} + + +class FakeApi: + """routes rest calls by method and path, recording each one.""" + + def __init__(self, routes=None): + self.calls = [] + self.routes = { + ("GET", "/v2/pods"): {"pods": [POD, STOPPED]}, + ("GET", "/v2/pods/pod123"): POD, + ("GET", "/v2/account/ssh-keys"): {"keys": ["ssh-ed25519 AAAA me"]}, + ("GET", "/v2/registries"): {"registries": [{"id": "reg1", "name": "ghcr"}]}, + ("GET", "/v2/network-volumes"): { + "networkVolumes": [ + {"id": "vol1", "name": "data", "dataCenter": "EU-RO-1"} + ] + }, + ("GET", "/v2/network-volumes/vol1"): { + "id": "vol1", + "dataCenter": "EU-RO-1", + }, + ("POST", "/v2/pods"): {**POD, "status": "PROVISIONING"}, + ("PATCH", "/v2/pods/pod123"): POD, + ("POST", "/v2/pods/pod123/action"): POD, + ("POST", "/v2/pods/pod456/action"): {**STOPPED, "status": "STARTING"}, + ("DELETE", "/v2/pods/pod123"): None, + **(routes or {}), + } + + def __call__(self, method, path, params=None, json=None): + self.calls.append((method, path, params, json)) + response = self.routes[(method, path)] + if isinstance(response, list): + return response.pop(0) + if isinstance(response, Exception): + raise response + return response + + def body(self, method, path): + return next(call[3] for call in self.calls if call[:2] == (method, path)) + + def made(self, method, path): + return any(call[:2] == (method, path) for call in self.calls) + + +@pytest.fixture(autouse=True) +def wide_console(monkeypatch): + from runpod.rp_cli import console as ui + + monkeypatch.setattr(ui.console, "_width", 200) + + +@pytest.fixture(autouse=True) +def no_sleep(monkeypatch): + monkeypatch.setattr(pod_module.time, "sleep", lambda seconds: None) + + +def _invoke(args, api=None, gpus=GPUS, interactive=False, input=None): + api = api or FakeApi() + with ( + patch("runpod.rp_cli.pod.run_rest_request", api), + patch("runpod.rp_cli.compute.run_rest_request", api), + patch("runpod.rp_cli.compute.list_gpus", return_value=gpus) as list_gpus, + patch("runpod.rp_cli.compute.list_cpus", return_value=CPUS) as list_cpus, + patch("runpod.rp_cli.compute.list_datacenters", return_value=DATACENTERS), + patch("runpod.rp_cli.pod.interactive", return_value=interactive), + patch("runpod.rp_cli.output.interactive", return_value=interactive), + ): + result = CliRunner().invoke(cli, ["pod", *args], input=input) + result.api, result.list_gpus, result.list_cpus = api, list_gpus, list_cpus + return result + + +def _create(args, api=None, gpus=GPUS): + result = _invoke(["create", *args], api=api, gpus=gpus) + assert result.exit_code == 0, result.output + return result, result.api.body("POST", "/v2/pods") + + +class TestList: + def test_table(self): + result = _invoke(["list"]) + assert result.exit_code == 0, result.output + assert "pod123" in result.output and "pod456" in result.output + assert "2x NVIDIA H100 80GB HBM3" in result.output + assert "cpu3c 4 vcpu" in result.output + assert "$7.98/hr" in result.output + assert result.api.calls[0][2] == {} + + def test_filters(self): + result = _invoke(["list", "--name", "BO", "--status", "exited", "--json"]) + assert [p["id"] for p in json.loads(result.stdout)] == ["pod456"] + result = _invoke(["list", "--status", "running", "--json"]) + assert [p["id"] for p in json.loads(result.stdout)] == ["pod123"] + + def test_cluster_pods_and_empty(self): + api = FakeApi({("GET", "/v2/pods"): {"pods": []}}) + result = _invoke(["list", "--include-cluster-pods"], api=api) + assert "no pods" in result.output + assert api.calls[0][2] == {"includeClusterPods": "true"} + + def test_api_error_is_clean(self): + api = FakeApi({("GET", "/v2/pods"): error.AuthenticationError("Unauthorized")}) + result = _invoke(["list"], api=api) + assert result.exit_code == 1 + assert "Error: Unauthorized" in result.output + + +class TestGet: + def test_fields_and_connection_info(self): + result = _invoke(["get", "pod123"]) + assert result.exit_code == 0, result.output + output = result.output + assert "ssh pod123-abc@ssh.runpod.io" in output + assert "ssh root@1.2.3.4 -p 40022" in output + assert "https://pod123-8888.proxy.runpod.net" in output + assert "1.2.3.4:40022 -> 22" in output + assert "network vol1 at /data" in output + assert "pod123.runpod.internal" in output + assert "32 vcpu, 250GB ram" in output + assert "HF_TOKEN, MODEL" in output + assert "secret" not in output + + def test_other_shapes(self): + pod = { + **STOPPED, + "mounts": { + "persistent": {"size": 20, "path": "/workspace"}, + "global": [{"volumeId": "gv1"}], + }, + "globalNetworking": {"enabled": False}, + } + api = FakeApi({("GET", "/v2/pods/pod456"): pod}) + result = _invoke(["get", "pod456"], api=api) + assert "persistent 20GB at /workspace" in result.output + assert "global gv1 at -" in result.output + assert "cpu3c 4 vcpu" in result.output + + def test_json(self): + result = _invoke(["get", "pod123", "--json"]) + assert json.loads(result.stdout)["env"]["HF_TOKEN"] == "secret" + + def test_missing_id_without_tty(self): + result = _invoke(["get"]) + assert result.exit_code == 2 + assert "pass a POD_ID" in result.output + + def test_picker(self): + with patch("InquirerPy.inquirer.fuzzy") as fuzzy: + fuzzy.return_value.execute.return_value = "pod123" + result = _invoke(["get"], interactive=True) + assert result.exit_code == 0, result.output + choices = fuzzy.call_args.kwargs["choices"] + assert [choice.value for choice in choices] == ["pod123", "pod456"] + assert choices[0].name == "trainer running" + + +class TestCreateGpu: + def test_ticket_example(self): + result, body = _create( + [ + "trainer", + "--gpu", + "h100:2", + "--image", + "runpod/pytorch:latest", + "--region", + "eu", + "--ssh", + ] + ) + assert body == { + "name": "trainer", + "image": "runpod/pytorch:latest", + "dataCenterIds": ["EU-RO-1", "EU-CZ-1"], + "gpu": {"id": "NVIDIA H100 80GB HBM3", "count": 2}, + "startSsh": True, + } + assert "'h100' matches 3 gpu types; using H100 SXM" in result.output + assert "low stock, $7.98/hr" in result.output + assert "no ssh keys" not in result.output + assert "created pod trainer (pod123)" in result.output + assert result.list_gpus.call_args.kwargs == { + "product": "POD", + "count": 2, + "cloud": "SECURE", + "cuda_versions": [], + "min_cuda_version": None, + } + + def test_prefers_stock_then_price(self): + _, body = _create(["a", "--gpu", "h100", "--image", "x"]) + assert body["gpu"]["id"] == "NVIDIA H100 NVL" + _, body = _create(["a", "--gpu", "h100", "--gpu-count", "2", "--image", "x"]) + assert body["gpu"]["id"] == "NVIDIA H100 PCIe" + + def test_stock_in_requested_data_centers(self): + _, body = _create( + ["a", "--gpu", "h100", "--data-center", "EU-RO-1", "--image", "x"] + ) + assert body["gpu"]["id"] == "NVIDIA H100 80GB HBM3" + assert body["dataCenterIds"] == ["EU-RO-1"] + + def test_unique_match_used_directly(self): + result, body = _create(["a", "--gpu", "4090", "--image", "x", "--json"]) + assert body["gpu"] == {"id": "NVIDIA GeForce RTX 4090", "count": 1} + assert "matches" not in result.output + assert "no secure stock for 1x RTX 4090" in result.stderr + assert json.loads(result.stdout)["id"] == "pod123" + + def test_gpu_settings(self): + _, body = _create( + [ + "a", + "--gpu", + "h100 nvl", + "--cloud", + "COMMUNITY", + "--gpu-count", + "4", + "--cuda", + ">=12.4", + "--min-ram-per-gpu", + "32", + "--min-vcpu-per-gpu", + "8", + "--image", + "x", + ] + ) + assert body["cloud"] == "COMMUNITY" + assert body["gpu"] == { + "id": "NVIDIA H100 NVL", + "count": 4, + "minCudaVersion": "12.4", + "minRamPerGpu": 32, + "minVcpuCountPerGpu": 8, + } + + def test_exact_cuda(self): + result, body = _create( + ["a", "--gpu", "4090", "--cuda", "12.4,12.8", "--image", "x"] + ) + assert body["gpu"]["allowedCudaVersions"] == ["12.4", "12.8"] + assert result.list_gpus.call_args.kwargs["cuda_versions"] == ["12.4", "12.8"] + + def test_no_type_offers_count(self): + result = _invoke(["create", "a", "--gpu", "h100:9", "--image", "x"]) + assert result.exit_code == 1 + assert "offers 9 per pod on secure cloud" in result.output + + def test_count_conflict(self): + result = _invoke( + ["create", "a", "--gpu", "h100:2", "--gpu-count", "3", "--image", "x"] + ) + assert result.exit_code == 2 + assert "conflicts" in result.output + + def test_unknown_gpu(self): + result = _invoke(["create", "a", "--gpu", "b9000", "--image", "x"]) + assert result.exit_code == 1 + assert "unknown gpu 'b9000'" in result.output + + def test_ssh_without_keys_warns(self): + api = FakeApi({("GET", "/v2/account/ssh-keys"): {"keys": []}}) + result, body = _create(["a", "--gpu", "4090", "--image", "x", "--ssh"], api=api) + assert body["startSsh"] is True + assert "no ssh keys" in result.output + + +class TestCreateCpu: + def test_default_vcpu(self): + result, body = _create(["box", "--cpu", "cpu5g", "--image", "ubuntu"]) + assert body["cpu"] == {"id": "cpu5g", "vcpuCount": 4} + assert result.list_cpus.call_args.kwargs == { + "product": "POD", + "vcpu_count": None, + } + + def test_explicit_vcpu(self): + _, body = _create(["box", "--cpu", "cpu3c", "--vcpu", "8", "--image", "ubuntu"]) + assert body["cpu"] == {"id": "cpu3c", "vcpuCount": 8} + + +class TestCreateContainer: + def test_every_container_field(self): + _, body = _create( + [ + "a", + "--gpu", + "4090", + "--template", + "tpl1", + "--image", + "img", + "-e", + "A=1", + "--env", + "B=x=y", + "-p", + "8888/http", + "--port", + "22/tcp", + "--disk", + "40", + "--args", + "bash start.sh", + "--cmd", + "python -m train --lr '1e-4'", + "--entrypoint", + "/bin/sh -c", + "--registry", + "ghcr", + "--global-networking", + "--jupyter", + "--persistent-volume", + "100", + ] + ) + assert body == { + "name": "a", + "image": "img", + "templateId": "tpl1", + "args": "bash start.sh", + "cmd": ["python", "-m", "train", "--lr", "1e-4"], + "entrypoint": ["/bin/sh", "-c"], + "disk": 40, + "ports": ["8888/http", "22/tcp"], + "registry": "reg1", + "globalNetworking": True, + "env": {"A": "1", "B": "x=y"}, + "mounts": {"persistent": {"size": 100, "path": "/workspace"}}, + "gpu": {"id": "NVIDIA GeForce RTX 4090", "count": 1}, + "startJupyter": True, + } + + def test_network_volume_sets_its_data_center(self): + _, body = _create( + [ + "a", + "--gpu", + "h100", + "--image", + "x", + "--network-volume", + "vol1", + "--global-volume", + "gv1:/g", + ] + ) + assert body["mounts"] == { + "network": [{"volumeId": "vol1", "path": "/workspace"}], + "global": [{"volumeId": "gv1", "path": "/g"}], + } + assert body["dataCenterIds"] == ["EU-RO-1"] + assert body["gpu"]["id"] == "NVIDIA H100 80GB HBM3" + + def test_regions_and_data_centers_merge(self): + _, body = _create( + [ + "a", + "--gpu", + "4090", + "--image", + "x", + "--region", + "eu-ro", + "--data-center", + "EU-RO-1", + "--data-center", + "US-TX-3", + "--network-volume", + "vol1:/data", + ] + ) + assert body["dataCenterIds"] == ["EU-RO-1", "US-TX-3"] + assert body["mounts"]["network"] == [{"volumeId": "vol1", "path": "/data"}] + + def test_template_without_image(self): + _, body = _create(["a", "--gpu", "4090", "--template", "tpl1"]) + assert body["templateId"] == "tpl1" + assert "image" not in body + + @pytest.mark.parametrize( + "args, message", + [ + (["--env", "NOPE"], "is not KEY=VALUE"), + (["--persistent-volume", "big"], "size is in GB"), + (["--network-volume", "vol1:"], "VALUE[:PATH]"), + ], + ) + def test_bad_values(self, args, message): + result = _invoke(["create", "a", "--gpu", "4090", "--image", "x", *args]) + assert result.exit_code == 2 + assert message in result.output + + def test_unknown_registry(self): + result = _invoke( + ["create", "a", "--gpu", "4090", "--image", "x", "--registry", "nope"] + ) + assert result.exit_code == 1 + assert "unknown registry credential 'nope'" in result.output + + +class TestCreateValidation: + @pytest.mark.parametrize( + "args, message", + [ + ( + ["a", "--gpu", "4090", "--cpu", "cpu3c", "--image", "x"], + "either --gpu or --cpu", + ), + (["a", "--image", "x"], "pass --gpu"), + ( + ["a", "--cpu", "cpu3c", "--cuda", "12.4", "--image", "x"], + "apply to gpu pods", + ), + (["a", "--gpu", "4090", "--vcpu", "4", "--image", "x"], "--vcpu applies"), + (["a", "--gpu", "4090"], "pass --image or --template"), + (["--gpu", "4090", "--image", "x"], "pass a NAME"), + ], + ) + def test_usage_errors(self, args, message): + result = _invoke(["create", *args]) + assert result.exit_code == 2 + assert message in result.output + assert not result.api.made("POST", "/v2/pods") + + def test_prompts_for_name(self): + result = _invoke( + ["create", "--gpu", "4090", "--image", "x"], + interactive=True, + input="mine\n", + ) + assert result.exit_code == 0, result.output + assert result.api.body("POST", "/v2/pods")["name"] == "mine" + + +class TestWait: + def test_create_waits_until_running(self): + api = FakeApi( + { + ("GET", "/v2/pods/pod123"): [ + {**POD, "status": "PROVISIONING"}, + {**POD, "status": "STARTING"}, + POD, + ] + } + ) + result = _invoke( + ["create", "a", "--gpu", "4090", "--image", "x", "--wait", "--json"], + api=api, + ) + assert result.exit_code == 0, result.output + assert json.loads(result.stdout)["status"] == "RUNNING" + assert sum(call[:2] == ("GET", "/v2/pods/pod123") for call in api.calls) == 3 + + def test_error_status_fails(self): + api = FakeApi({("GET", "/v2/pods/pod456"): {**STOPPED, "status": "ERROR"}}) + result = _invoke(["start", "pod456", "--wait"], api=api) + assert result.exit_code == 1 + assert "pod pod456 is error" in result.output + + def test_timeout(self, monkeypatch): + monkeypatch.setattr(pod_module, "WAIT_TIMEOUT", -1) + api = FakeApi({("GET", "/v2/pods/pod456"): STOPPED}) + result = _invoke(["start", "pod456", "--wait"], api=api) + assert result.exit_code == 1 + assert "still exited" in result.output + + +class TestUpdate: + def test_nothing_to_update(self): + result = _invoke(["update", "pod123"]) + assert result.exit_code == 2 + assert "nothing to update" in result.output + + def test_simple_fields_skip_the_read(self): + result = _invoke( + [ + "update", + "pod123", + "--name", + "renamed", + "--lock", + "--no-global-networking", + "--image", + "img:2", + "--registry", + "", + "--port", + "8000/http", + "--disk", + "80", + "--cmd", + "", + ] + ) + assert result.exit_code == 0, result.output + assert result.api.body("PATCH", "/v2/pods/pod123") == { + "image": "img:2", + "cmd": [], + "disk": 80, + "ports": ["8000/http"], + "registry": None, + "globalNetworking": False, + "name": "renamed", + "locked": True, + } + assert not result.api.made("GET", "/v2/pods/pod123") + assert "updated pod pod123" in result.output + + def test_env_edits_existing(self): + result = _invoke( + [ + "update", + "pod123", + "-e", + "MODEL=qwen", + "-e", + "NEW=1", + "--unset-env", + "HF_TOKEN", + "--unlock", + ] + ) + assert result.exit_code == 0, result.output + assert result.api.body("PATCH", "/v2/pods/pod123") == { + "locked": False, + "env": {"MODEL": "qwen", "NEW": "1"}, + } + + def test_mount_path_defaults_to_current(self): + result = _invoke(["update", "pod123", "--network-volume", "vol1", "--json"]) + assert result.exit_code == 0, result.output + assert result.api.body("PATCH", "/v2/pods/pod123") == { + "mounts": {"network": [{"volumeId": "vol1", "path": "/data"}]} + } + assert json.loads(result.stdout)["id"] == "pod123" + + +class TestActions: + @pytest.mark.parametrize("action", ["stop", "restart"]) + def test_action(self, action): + result = _invoke([action, "pod123"]) + assert result.exit_code == 0, result.output + assert result.api.body("POST", "/v2/pods/pod123/action") == {"action": action} + + def test_start_wait(self): + api = FakeApi( + { + ("GET", "/v2/pods/pod456"): [ + {**STOPPED, "status": "STARTING"}, + {**STOPPED, "status": "RUNNING"}, + ] + } + ) + result = _invoke(["start", "pod456", "--wait", "--json"], api=api) + assert result.exit_code == 0, result.output + assert json.loads(result.stdout)["status"] == "RUNNING" + assert "started pod pod456" in result.stderr + + def test_picker_offers_pods_that_allow_the_action(self): + with patch("InquirerPy.inquirer.fuzzy") as fuzzy: + fuzzy.return_value.execute.return_value = "pod456" + result = _invoke(["start"], interactive=True) + assert result.exit_code == 0, result.output + assert [c.value for c in fuzzy.call_args.kwargs["choices"]] == ["pod456"] + + def test_api_error(self): + api = FakeApi( + { + ("POST", "/v2/pods/pod123/action"): error.QueryError( + "pod is locked", "POST" + ) + } + ) + result = _invoke(["stop", "pod123"], api=api) + assert result.exit_code == 1 + assert "pod is locked" in result.output + + +class TestTerminate: + @pytest.mark.parametrize("command", ["terminate", "delete"]) + def test_with_yes(self, command): + result = _invoke([command, "pod123", "--yes"]) + assert result.exit_code == 0, result.output + assert result.api.made("DELETE", "/v2/pods/pod123") + assert "terminated pod pod123" in result.output + + def test_needs_confirmation(self): + result = _invoke(["terminate", "pod123"]) + assert result.exit_code == 2 + assert "--yes" in result.output + assert not result.api.made("DELETE", "/v2/pods/pod123") + + def test_interactive_confirm(self): + result = _invoke(["terminate", "pod123"], interactive=True, input="n\n") + assert result.exit_code == 1 + assert not result.api.made("DELETE", "/v2/pods/pod123") + + +class TestConnect: + def test_launches_terminal(self): + ssh = MagicMock() + with patch( + "runpod.cli.utils.ssh_cmd.SSHConnection", return_value=ssh + ) as connection: + result = _invoke(["connect", "pod123"]) + assert result.exit_code == 0, result.output + connection.assert_called_once_with("pod123") + ssh.launch_terminal.assert_called_once() + + def test_failure_is_clean(self): + with patch( + "runpod.cli.utils.ssh_cmd.SSHConnection", side_effect=TimeoutError("no ssh") + ): + result = _invoke(["connect", "pod123"]) + assert result.exit_code == 1 + assert "Error: no ssh" in result.output + + def test_sync_is_mounted(self): + result = _invoke(["sync", "--help"]) + assert result.exit_code == 0 + + +class TestCompletion: + def test_completers(self): + api = FakeApi() + with ( + patch("runpod.rp_cli.pod.run_rest_request", api), + patch("runpod.rp_cli.compute.run_rest_request", api), + patch("runpod.rp_cli.compute.list_datacenters", return_value=DATACENTERS), + ): + pods = pod_module._complete_pods(None, None, "pod1") + datacenters = pod_module._complete_datacenters(None, None, "EU") + volumes = pod_module._complete_network_volumes(None, None, "") + registries = pod_module._complete_registries(None, None, "g") + assert [(item.value, item.help) for item in pods] == [ + ("pod123", "trainer running") + ] + assert [item.value for item in datacenters] == ["EU-RO-1", "EU-CZ-1"] + assert [(item.value, item.help) for item in volumes] == [ + ("vol1", "data EU-RO-1") + ] + assert [item.value for item in registries] == ["ghcr"] 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 05/25] 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 0e98c3bb4592d06c29db205928ee6dfea7d4c089 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:01:01 -0700 Subject: [PATCH 06/25] fix: keep global volume ids whole in rp volume list --- runpod/rp_cli/volume.py | 11 ++++++++--- tests/test_cli/test_volume_command.py | 21 ++++++++++++++++++--- 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/runpod/rp_cli/volume.py b/runpod/rp_cli/volume.py index 43be7cc5..f28212e4 100644 --- a/runpod/rp_cli/volume.py +++ b/runpod/rp_cli/volume.py @@ -49,8 +49,7 @@ def _columns(kind: Kind): return [ ("ID", lambda row: row["id"]), ("NAME", lambda row: row.get("name")), - ("CREATED", lambda row: row.get("createdAt")), - ("CREATED BY", lambda row: row.get("createdByUserId")), + ("CREATED", lambda row: (row.get("createdAt") or "")[:10] or None), ] return [ ("ID", lambda row: row["id"]), @@ -89,7 +88,13 @@ def _show(kind: Kind, record: dict, as_json: bool) -> None: if as_json: print_json(record) return - print_fields(record, [(header.lower(), fn) for header, fn in _columns(kind)]) + fields = [(header.lower(), fn) for header, fn in _columns(kind)] + if kind is GLOBAL: + fields[2:] = [ + ("created", lambda row: row.get("createdAt")), + ("created by", lambda row: row.get("createdByUserId")), + ] + print_fields(record, fields) def _data_center( diff --git a/tests/test_cli/test_volume_command.py b/tests/test_cli/test_volume_command.py index 0aa0c8ad..f66441a4 100644 --- a/tests/test_cli/test_volume_command.py +++ b/tests/test_cli/test_volume_command.py @@ -88,10 +88,18 @@ def test_global_json(self, listing): listing.assert_called_once_with("/v2/global-volumes", "globalVolumes") assert json.loads(result.output) == [GLOBAL_VOLUME] - def test_global_table(self, listing): - listing.return_value = [GLOBAL_VOLUME] + def test_global_table_keeps_long_ids_whole(self, listing): + listing.return_value = [ + dict( + GLOBAL_VOLUME, + id="cmvfsn771000007jp15tn8esm", + name="capy-test-global-renamed", + ) + ] result = invoke("list", "--global") - assert "gsb_abc" in result.output and "user_123" in result.output + assert "cmvfsn771000007jp15tn8esm" in result.output + assert "2026-09-24" in result.output + assert "user_123" not in result.output def test_empty(self, listing): listing.return_value = [] @@ -115,6 +123,13 @@ def test_fields(self, rest): assert "data center" in result.output and "EU-RO-1" in result.output assert "standard" in result.output + def test_global_fields(self, rest): + rest.return_value = GLOBAL_VOLUME + result = invoke("get", "gsb_abc", "--global") + assert result.exit_code == 0, result.output + assert "2026-09-24T00:00:00Z" in result.output + assert "created by" in result.output and "user_123" in result.output + def test_global_json(self, rest): rest.return_value = GLOBAL_VOLUME result = invoke("get", "gsb_abc", "--global", "--json") From ed085997d28038a68d3c92b4659e0bd40cd5cd2f Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:06:31 -0700 Subject: [PATCH 07/25] fix: keep table ids and values whole on narrow terminals --- runpod/rp_cli/output.py | 4 ++-- tests/test_cli/test_catalog_command.py | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/runpod/rp_cli/output.py b/runpod/rp_cli/output.py index ce0a0dad..69f7a0f7 100644 --- a/runpod/rp_cli/output.py +++ b/runpod/rp_cli/output.py @@ -43,8 +43,8 @@ def print_table(rows: Sequence[dict], columns: Sequence[Column], empty: str) -> ui.console.print(f"[dim]{empty}[/dim]") return table = Table(box=None, header_style="dim", pad_edge=False) - for header, _ in columns: - table.add_column(header) + 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)) ui.console.print(table) diff --git a/tests/test_cli/test_catalog_command.py b/tests/test_cli/test_catalog_command.py index 8efb5fd5..f30abcc6 100644 --- a/tests/test_cli/test_catalog_command.py +++ b/tests/test_cli/test_catalog_command.py @@ -231,6 +231,24 @@ def test_cell(self): assert output.cell(["a", "b"]) == "a, b" assert output.cell(3) == "3" + def test_print_table_keeps_ids_and_values_whole(self, monkeypatch, capsys): + from runpod.rp_cli import console as ui + + monkeypatch.setattr(ui.console, "_width", 40) + rows = [{"id": "cmvfsn771000007jp15tn8esm", "type": "high_performance"}] + output.print_table( + rows, + [ + ("ID", lambda r: r["id"]), + ("TYPE", lambda r: r["type"]), + ("PAD", lambda r: "x" * 30), + ], + "none", + ) + out = capsys.readouterr().out + assert "cmvfsn771000007jp15tn8esm" in out + assert "…" not in out + def test_print_fields(self, capsys): output.print_fields( {"id": "abc", "tags": []}, 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 08/25] 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 09/25] 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 1173258eb4da17ca4e99eb113871a77229c5520d Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:22:27 -0700 Subject: [PATCH 10/25] fix: wait for the pod container to start, not just its desired status --- runpod/rp_cli/pod.py | 48 ++++++++++++++++++++++-------- tests/test_cli/test_pod_command.py | 27 +++++++++++++---- 2 files changed, 58 insertions(+), 17 deletions(-) diff --git a/runpod/rp_cli/pod.py b/runpod/rp_cli/pod.py index 35f62d5b..3cfd74ec 100644 --- a/runpod/rp_cli/pod.py +++ b/runpod/rp_cli/pod.py @@ -140,6 +140,26 @@ def _global_network(pod: dict) -> object: return networking.get("internalDns") or networking.get("ip") or True +def _status(pod: dict) -> str: + status = (pod.get("status") or "").lower() + if status == "running" and not (pod.get("runtime") or {}).get("uptime"): + return "starting" + return status + + +def _uptime(pod: dict) -> Optional[str]: + seconds = (pod.get("runtime") or {}).get("uptime") + if not seconds: + return None + minutes, seconds = divmod(seconds, 60) + hours, minutes = divmod(minutes, 60) + if hours: + return f"{hours}h {minutes}m" + if minutes: + return f"{minutes}m {seconds}s" + return f"{seconds}s" + + def _resources(pod: dict) -> Optional[str]: hardware = pod.get("gpu") or pod.get("cpu") or {} parts = [] @@ -153,7 +173,7 @@ def _resources(pod: dict) -> Optional[str]: _LIST_COLUMNS = [ ("ID", lambda pod: pod["id"]), ("NAME", lambda pod: pod.get("name")), - ("STATUS", lambda pod: (pod.get("status") or "").lower()), + ("STATUS", _status), ("HARDWARE", _hardware), ("DATA CENTER", lambda pod: pod.get("dataCenterId")), ("COST", _cost), @@ -163,7 +183,7 @@ def _resources(pod: dict) -> Optional[str]: _FIELDS = [ ("id", lambda pod: pod["id"]), ("name", lambda pod: pod.get("name")), - ("status", lambda pod: (pod.get("status") or "").lower()), + ("status", _status), ("image", lambda pod: pod.get("image")), ("hardware", _hardware), ("resources", _resources), @@ -181,6 +201,7 @@ def _resources(pod: dict) -> Optional[str]: ("cluster", lambda pod: (pod.get("cluster") or {}).get("id")), ("created", lambda pod: pod.get("createdAt")), ("started", lambda pod: pod.get("startedAt")), + ("uptime", _uptime), ("actions", lambda pod: pod.get("actions")), ("ssh", lambda pod: ((pod.get("ssh") or {}).get("proxy") or {}).get("command")), ( @@ -199,23 +220,24 @@ def _show(pod: dict, as_json: bool) -> None: print_fields(pod, _FIELDS) -def _wait_running(pod_id: str) -> dict: - """poll a pod until it is running.""" +def _wait_running(pod_id: str, since: float) -> dict: + """poll a pod until a container started after `since` is up.""" from runpod.rp_cli import console as ui deadline = time.monotonic() + WAIT_TIMEOUT with ui.console.status(f"[dim]waiting for {pod_id} to start ...[/dim]"): while True: pod = _get_pod(pod_id) - status = pod.get("status") - if status == "RUNNING": + status = _status(pod) + uptime = pod["runtime"]["uptime"] if status == "running" else 0 + # a longer uptime is a stale report from the container before the action + if status == "running" and uptime <= time.monotonic() - since + 1: return pod - if status in ("ERROR", "TERMINATED"): - raise click.ClickException(f"pod {pod_id} is {status.lower()}") + if status in ("error", "terminated"): + raise click.ClickException(f"pod {pod_id} is {status}") if time.monotonic() >= deadline: raise click.ClickException( - f"pod {pod_id} is still {(status or 'unknown').lower()} " - f"after {WAIT_TIMEOUT}s" + f"pod {pod_id} is still {status or 'unknown'} after {WAIT_TIMEOUT}s" ) time.sleep(WAIT_INTERVAL) @@ -686,10 +708,11 @@ def create( if jupyter: body["startJupyter"] = True + since = time.monotonic() created = run_rest_request("POST", "/v2/pods", json=body) or {} _say(f"created pod {created.get('name')} ({created.get('id')})", as_json) if wait: - created = _wait_running(created["id"]) + created = _wait_running(created["id"], since) _show(created, as_json) @@ -752,12 +775,13 @@ def _act( ): pod_id = _resolve_pod_id(pod_id, action) with api_errors(): + since = time.monotonic() result = run_rest_request( "POST", f"/v2/pods/{pod_id}/action", json={"action": action} ) _say(f"{done} pod {pod_id}", as_json) if wait: - result = _wait_running(pod_id) + result = _wait_running(pod_id, since) _show(result or {}, as_json) diff --git a/tests/test_cli/test_pod_command.py b/tests/test_cli/test_pod_command.py index b65fafcd..25678127 100644 --- a/tests/test_cli/test_pod_command.py +++ b/tests/test_cli/test_pod_command.py @@ -102,7 +102,8 @@ "ports": [ {"private": 22, "public": 40022, "type": "tcp", "ip": "1.2.3.4"}, {"private": 8888, "public": None, "type": "http", "ip": None}, - ] + ], + "uptime": 3725, }, "createdAt": "2026-10-19T20:00:00Z", "startedAt": "2026-10-19T20:01:00Z", @@ -239,6 +240,19 @@ def test_fields_and_connection_info(self): assert "32 vcpu, 250GB ram" in output assert "HF_TOKEN, MODEL" in output assert "secret" not in output + assert "1h 2m" in output + + @pytest.mark.parametrize( + "runtime, status, uptime", + [(None, "starting", None), ({"uptime": 0}, "starting", None)] + + [({"uptime": 65}, "running", "1m 5s"), ({"uptime": 9}, "running", "9s")], + ) + def test_status_follows_container(self, runtime, status, uptime): + api = FakeApi({("GET", "/v2/pods/pod123"): {**POD, "runtime": runtime}}) + result = _invoke(["get", "pod123"], api=api) + assert f"status {status}" in result.output + if uptime: + assert f"uptime {uptime}" in result.output def test_other_shapes(self): pod = { @@ -562,7 +576,9 @@ def test_create_waits_until_running(self): ("GET", "/v2/pods/pod123"): [ {**POD, "status": "PROVISIONING"}, {**POD, "status": "STARTING"}, - POD, + {**POD, "runtime": None}, + {**POD, "runtime": {"uptime": 0}}, + {**POD, "runtime": {"uptime": 1}}, ] } ) @@ -572,7 +588,7 @@ def test_create_waits_until_running(self): ) assert result.exit_code == 0, result.output assert json.loads(result.stdout)["status"] == "RUNNING" - assert sum(call[:2] == ("GET", "/v2/pods/pod123") for call in api.calls) == 3 + assert sum(call[:2] == ("GET", "/v2/pods/pod123") for call in api.calls) == 5 def test_error_status_fails(self): api = FakeApi({("GET", "/v2/pods/pod456"): {**STOPPED, "status": "ERROR"}}) @@ -670,13 +686,14 @@ def test_start_wait(self): { ("GET", "/v2/pods/pod456"): [ {**STOPPED, "status": "STARTING"}, - {**STOPPED, "status": "RUNNING"}, + {**STOPPED, "status": "RUNNING", "runtime": {"uptime": 46}}, + {**STOPPED, "status": "RUNNING", "runtime": {"uptime": 1}}, ] } ) result = _invoke(["start", "pod456", "--wait", "--json"], api=api) assert result.exit_code == 0, result.output - assert json.loads(result.stdout)["status"] == "RUNNING" + assert json.loads(result.stdout)["runtime"] == {"uptime": 1} assert "started pod pod456" in result.stderr def test_picker_offers_pods_that_allow_the_action(self): From 619e23bac26df572ffa67823fc6418611146841d Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:47:57 -0700 Subject: [PATCH 11/25] refactor: consolidate CLI resolution and rendering helpers --- runpod/rp_cli/api.py | 30 +++++++++ runpod/rp_cli/catalog.py | 62 ++++++++---------- runpod/rp_cli/compute.py | 61 ++++------------- runpod/rp_cli/options.py | 15 +++++ runpod/rp_cli/output.py | 54 ++++++---------- tests/test_cli/test_catalog_command.py | 83 ------------------------ tests/test_cli/test_cli_api.py | 42 ++++++++++++ tests/test_cli/test_cli_options.py | 21 ++++++ tests/test_cli/test_cli_output.py | 90 ++++++++++++++++++++++++++ tests/test_cli/test_compute.py | 31 ++------- 10 files changed, 262 insertions(+), 227 deletions(-) create mode 100644 runpod/rp_cli/api.py create mode 100644 runpod/rp_cli/options.py create mode 100644 tests/test_cli/test_cli_api.py create mode 100644 tests/test_cli/test_cli_options.py create mode 100644 tests/test_cli/test_cli_output.py diff --git a/runpod/rp_cli/api.py b/runpod/rp_cli/api.py new file mode 100644 index 00000000..05d493f8 --- /dev/null +++ b/runpod/rp_cli/api.py @@ -0,0 +1,30 @@ +"""pagination and error reporting for rest-backed cli commands.""" + +from contextlib import contextmanager + +import click +from requests import RequestException + +from runpod import error +from runpod.api.rest import run_rest_request +from runpod.rp_cli.compute import ResolutionError + + +def list_all(path: str, collection: str, params: dict | None = None) -> list[dict]: + query = dict(params or {}) + records = [] + while True: + page = run_rest_request("GET", path, params=query) + records.extend(page[collection]) + pagination = page["pagination"] + if not pagination["hasNextPage"]: + return records + query["cursor"] = pagination["nextCursor"] + + +@contextmanager +def api_errors(): + try: + yield + except (error.RunPodError, ResolutionError, RequestException) as exc: + raise click.ClickException(str(exc)) from exc diff --git a/runpod/rp_cli/catalog.py b/runpod/rp_cli/catalog.py index 16d864d0..ffe437fd 100644 --- a/runpod/rp_cli/catalog.py +++ b/runpod/rp_cli/catalog.py @@ -5,11 +5,8 @@ import click from runpod.rp_cli import compute -from runpod.rp_cli.output import api_errors, json_option, print_json, print_table - - -def _price(value: Optional[float], count: int) -> Optional[float]: - return round(value * count, 4) if value else None +from runpod.rp_cli.api import api_errors +from runpod.rp_cli.output import json_option, print_json, print_table def _money(value: Optional[float]) -> str: @@ -25,28 +22,27 @@ def _stock(entries: list[dict]) -> str: ) -def _in_regions(record: dict, allowed: Optional[set[str]]) -> dict: - """the record with stock narrowed to the allowed data centers.""" - if allowed is None: - return dict(record) - datacenters = [dc for dc in record.get("dataCenters") or [] if dc["id"] in allowed] - return dict( - record, - dataCenters=datacenters, - availability=compute.best_availability( - dc.get("availability") for dc in datacenters - ), +def _filter_stock(records: list[dict], regions: tuple[str, ...], available: bool): + allowed = ( + set(compute.resolve_regions(regions, compute.list_datacenters())) + if regions + else None ) - - -def _allowed(regions: tuple[str, ...]) -> Optional[set[str]]: - if not regions: - return None - return set(compute.resolve_regions(regions, compute.list_datacenters())) + for record in records: + row = dict(record) + if allowed is not None: + row["dataCenters"] = [ + dc for dc in record["dataCenters"] if dc["id"] in allowed + ] + row["availability"] = compute.best_availability( + dc["availability"] for dc in row["dataCenters"] + ) + if not available or compute.availability_rank(row["availability"]) > 0: + yield row def _unit_price(gpu: dict, product: str, cloud: Optional[str]) -> Optional[float]: - price = gpu.get("price") or {} + price = gpu["price"] if product == "SERVERLESS": return price.get("serverless") or None if cloud: @@ -72,21 +68,17 @@ def _gpu_rows( gpus = compute.resolve_gpus(query, gpus) if product == "SERVERLESS": gpus = [gpu for gpu in gpus if gpu.get("pool")] - allowed = _allowed(regions) rows = [] - for gpu in gpus: - row = _in_regions(gpu, allowed) + for row in _filter_stock(gpus, regions, available): unit = _unit_price(row, product, cloud) if min_vram and (row.get("memory") or 0) < min_vram: continue if max_price is not None and (unit is None or unit > max_price): continue - if available and row.get("availability") in (None, "NONE"): - continue - price = row.get("price") or {} row["count"] = count row["totalPrice"] = { - tier: _price(value, count) for tier, value in price.items() + tier: round(value * count, 4) if value else None + for tier, value in row["price"].items() } rows.append(row) rows.sort( @@ -135,12 +127,10 @@ def _cpu_rows(query, product, vcpu, regions, available) -> list[dict]: cpus = compute.match_cpus(query, cpus) if not cpus: raise compute.ResolutionError(f"unknown cpu '{query}'") - allowed = _allowed(regions) - rows = [_in_regions(cpu, allowed) for cpu in cpus] - if available: - rows = [row for row in rows if row.get("availability") not in (None, "NONE")] - rows.sort(key=lambda row: -compute.availability_rank(row.get("availability"))) - return rows + return sorted( + _filter_stock(cpus, regions, available), + key=lambda row: -compute.availability_rank(row["availability"]), + ) def _cpu_columns(product: str): diff --git a/runpod/rp_cli/compute.py b/runpod/rp_cli/compute.py index af0a0540..8339323f 100644 --- a/runpod/rp_cli/compute.py +++ b/runpod/rp_cli/compute.py @@ -1,7 +1,8 @@ """gpu, cpu, and region resolution against the rest v2 catalog.""" import re -from typing import Any, Iterable, Optional +from collections.abc import Iterable +from typing import Any, Optional from runpod.api.rest import run_rest_request @@ -29,31 +30,12 @@ class ResolutionError(ValueError): """a user-supplied gpu, cpu, or region that does not resolve.""" -def list_all( - path: str, collection: str, params: Optional[dict[str, Any]] = None -) -> list[dict]: - """every record of a cursor-paginated list endpoint.""" - query = dict(params or {}) - records: list[dict] = [] - while True: - page = run_rest_request("GET", path, params=query) or {} - records.extend(page.get(collection) or []) - pagination = page.get("pagination") or {} - cursor = pagination.get("nextCursor") - if not pagination.get("hasNextPage") or not cursor: - return records - query["cursor"] = cursor - - def _query_params(**values: Any) -> dict[str, Any]: - params = {} - for key, value in values.items(): - if value is None or value == []: - continue - if isinstance(value, (list, tuple)): - value = ",".join(str(item) for item in value) - params[key] = value - return params + return { + key: ",".join(value) if isinstance(value, (list, tuple)) else value + for key, value in values.items() + if value is not None and value != [] + } def list_gpus( @@ -72,8 +54,7 @@ def list_gpus( cudaVersions=cuda_versions if product else None, minCudaVersion=min_cuda_version if product else None, ) - response = run_rest_request("GET", "/v2/catalog/gpus", params=params) or {} - return response.get("gpus") or [] + return run_rest_request("GET", "/v2/catalog/gpus", params=params)["gpus"] def list_cpus( @@ -85,19 +66,14 @@ def list_cpus( product=product, vcpuCount=vcpu_count if product else None, ) - response = run_rest_request("GET", "/v2/catalog/cpus", params=params) or {} - return response.get("cpus") or [] + return run_rest_request("GET", "/v2/catalog/cpus", params=params)["cpus"] def list_datacenters(include: Optional[list[str]] = None) -> list[dict]: """every data center in the catalog.""" - response = ( - run_rest_request( - "GET", "/v2/catalog/datacenters", params=_query_params(include=include) - ) - or {} - ) - return response.get("dataCenters") or [] + return run_rest_request( + "GET", "/v2/catalog/datacenters", params=_query_params(include=include) + )["dataCenters"] def _tokens(text: str) -> list[str]: @@ -109,13 +85,7 @@ def _compact(text: str) -> str: def match_gpus(query: str, gpus: Iterable[dict]) -> list[dict]: - """gpu types a human-readable query names. - - an exact id, short name, or serverless pool id wins. otherwise every - query token must appear in the type's id or name, so "4090" matches - the rtx 4090, "h100" matches every h100 variant, and "l4" does not - match the l40. - """ + """prefer exact ids, names, and pools, then match hardware-name tokens.""" gpus = list(gpus) text = query.strip().lower() exact = [ @@ -198,10 +168,7 @@ def endpoint_gpu_selection( ) -> tuple[list[str], list[str]]: """serverless (pools, excludedTypes) that select exactly the matched types.""" matches = list(matches) - pools: list[str] = [] - for gpu in matches: - if gpu.get("pool") and gpu["pool"] not in pools: - pools.append(gpu["pool"]) + pools = list(dict.fromkeys(gpu["pool"] for gpu in matches if gpu.get("pool"))) if not pools: names = ", ".join(gpu.get("name") or gpu["id"] for gpu in matches) raise ResolutionError(f"{names} is not offered on serverless") diff --git a/runpod/rp_cli/options.py b/runpod/rp_cli/options.py new file mode 100644 index 00000000..ee396289 --- /dev/null +++ b/runpod/rp_cli/options.py @@ -0,0 +1,15 @@ +"""value parsing shared by resource command options.""" + +from collections.abc import Iterable + +import click + + +def parse_env(values: Iterable[str]) -> dict[str, str]: + env = {} + for value in values: + key, separator, content = value.partition("=") + if not key or not separator: + raise click.UsageError("invalid --env value; use KEY=VALUE") + env[key] = content + return env diff --git a/runpod/rp_cli/output.py b/runpod/rp_cli/output.py index 69f7a0f7..31fd9baf 100644 --- a/runpod/rp_cli/output.py +++ b/runpod/rp_cli/output.py @@ -1,16 +1,13 @@ """output, picker, prompt, and completion helpers shared by resource commands.""" -import contextlib import json import sys -from typing import Any, Callable, Iterable, Iterator, Sequence +from collections.abc import Callable, Iterable, Sequence +from typing import Any import click from click.shell_completion import CompletionItem -from runpod import error -from runpod.rp_cli.compute import ResolutionError - Column = tuple[str, Callable[[dict], Any]] json_option = click.option( @@ -32,33 +29,36 @@ def cell(value: Any) -> str: return str(value) -def print_table(rows: Sequence[dict], columns: Sequence[Column], empty: str) -> None: - """render rows as an aligned table, or a dim notice when empty.""" +def table(rows: Sequence[dict], columns: Sequence[Column]): from rich.markup import escape from rich.table import Table - from runpod.rp_cli import console as ui - - if not rows: - ui.console.print(f"[dim]{empty}[/dim]") - return - table = Table(box=None, header_style="dim", pad_edge=False) + result = 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") + result.add_column(header, no_wrap=index == 0, overflow="fold") for row in rows: - table.add_row(*(escape(cell(value(row))) for _, value in columns)) - ui.console.print(table) + result.add_row(*(escape(cell(value(row))) for _, value in columns)) + return result + + +def print_table(rows: Sequence[dict], columns: Sequence[Column], empty: str) -> None: + from runpod.rp_cli import console as ui + + ui.console.print(table(rows, columns) if rows else f"[dim]{empty}[/dim]") def print_fields(record: dict, fields: Sequence[Column]) -> None: - """render one record as aligned label/value lines.""" from rich.markup import escape + from rich.table import Table from runpod.rp_cli import console as ui - width = max(len(label) for label, _ in fields) + result = Table.grid(padding=(0, 2)) + result.add_column(style="dim", no_wrap=True) + result.add_column(overflow="fold") for label, value in fields: - ui.console.print(f"[dim]{label:<{width}}[/dim] {escape(cell(value(record)))}") + result.add_row(label, escape(cell(value(record)))) + ui.console.print(result) def interactive() -> bool: @@ -66,10 +66,6 @@ def interactive() -> bool: def pick(message: str, choices: Sequence[tuple[str, str]], missing: str) -> str: - """let a person choose a value from (value, label) pairs. - - scripts and agents get a usage error naming what to pass instead. - """ if not interactive(): raise click.UsageError(missing) if not choices: @@ -84,7 +80,6 @@ def pick(message: str, choices: Sequence[tuple[str, str]], missing: str) -> str: def confirm(message: str, yes: bool) -> None: - """ask before a destructive action unless --yes was passed.""" if yes: return if not interactive(): @@ -100,7 +95,7 @@ def completer( def complete(ctx, param, incomplete): try: entries = list(fetch()) - except Exception: # noqa: BLE001 - completion must never break the shell + except Exception: return [] return [ CompletionItem(value, help=label) @@ -109,12 +104,3 @@ def complete(ctx, param, incomplete): ] return complete - - -@contextlib.contextmanager -def api_errors() -> Iterator[None]: - """surface api and resolution failures as clean cli errors.""" - try: - yield - except (error.RunPodError, ResolutionError) as exc: - raise click.ClickException(str(exc)) from exc diff --git a/tests/test_cli/test_catalog_command.py b/tests/test_cli/test_catalog_command.py index f30abcc6..9c27ca3f 100644 --- a/tests/test_cli/test_catalog_command.py +++ b/tests/test_cli/test_catalog_command.py @@ -3,12 +3,10 @@ import json from unittest.mock import patch -import click import pytest from click.testing import CliRunner from runpod import error -from runpod.rp_cli import output from runpod.rp_cli.main import cli GPUS = [ @@ -222,84 +220,3 @@ def test_prints_script(self): result = CliRunner().invoke(cli, ["completion", "zsh"]) assert result.exit_code == 0 assert "_RP_COMPLETE" in result.output - - -class TestOutputHelpers: - def test_cell(self): - assert output.cell(None) == "-" - assert output.cell(True) == "yes" - assert output.cell(["a", "b"]) == "a, b" - assert output.cell(3) == "3" - - def test_print_table_keeps_ids_and_values_whole(self, monkeypatch, capsys): - from runpod.rp_cli import console as ui - - monkeypatch.setattr(ui.console, "_width", 40) - rows = [{"id": "cmvfsn771000007jp15tn8esm", "type": "high_performance"}] - output.print_table( - rows, - [ - ("ID", lambda r: r["id"]), - ("TYPE", lambda r: r["type"]), - ("PAD", lambda r: "x" * 30), - ], - "none", - ) - out = capsys.readouterr().out - assert "cmvfsn771000007jp15tn8esm" in out - assert "…" not in out - - def test_print_fields(self, capsys): - output.print_fields( - {"id": "abc", "tags": []}, - [("id", lambda r: r["id"]), ("tags", lambda r: r["tags"])], - ) - out = capsys.readouterr().out - assert "id abc" in out - assert "tags -" in out - - def test_pick_requires_tty(self): - with patch.object(output, "interactive", return_value=False): - with pytest.raises(click.UsageError, match="pass an ID"): - output.pick("pick", [("a", "a")], "pass an ID") - - def test_pick_with_no_choices(self): - with patch.object(output, "interactive", return_value=True): - with pytest.raises(click.UsageError, match="none found"): - output.pick("pick", [], "pass an ID") - - def test_pick_uses_fuzzy_prompt(self): - with ( - patch.object(output, "interactive", return_value=True), - patch("InquirerPy.inquirer.fuzzy") as fuzzy, - ): - fuzzy.return_value.execute.return_value = "b" - assert output.pick("pick", [("a", "first"), ("b", "second")], "x") == "b" - choices = fuzzy.call_args.kwargs["choices"] - assert [(choice.value, choice.name) for choice in choices] == [ - ("a", "first"), - ("b", "second"), - ] - - def test_confirm(self): - output.confirm("sure?", yes=True) - with patch.object(output, "interactive", return_value=False): - with pytest.raises(click.UsageError, match="--yes"): - output.confirm("sure?", yes=False) - with ( - patch.object(output, "interactive", return_value=True), - patch("click.confirm") as ask, - ): - output.confirm("sure?", yes=False) - ask.assert_called_once_with("sure?", abort=True) - - def test_completer(self): - complete = output.completer(lambda: [("abc", "first"), ("xyz", "second")]) - items = complete(None, None, "a") - assert [(item.value, item.help) for item in items] == [("abc", "first")] - - def test_completer_swallows_errors(self): - def broken(): - raise error.AuthenticationError("no key") - - assert output.completer(broken)(None, None, "") == [] diff --git a/tests/test_cli/test_cli_api.py b/tests/test_cli/test_cli_api.py new file mode 100644 index 00000000..30bdf537 --- /dev/null +++ b/tests/test_cli/test_cli_api.py @@ -0,0 +1,42 @@ +"""cursor pagination and cli errors.""" + +from unittest.mock import patch + +import click +import pytest +import requests + +from runpod import error +from runpod.rp_cli import api + + +class TestApi: + @pytest.mark.parametrize( + "failure", + [error.AuthenticationError("unauthorized"), requests.Timeout("timed out")], + ) + def test_api_errors_are_click_errors(self, failure): + with pytest.raises(click.ClickException, match=str(failure)): + with api.api_errors(): + raise failure + + def test_list_all_follows_cursors(self): + pages = [ + { + "pods": [{"id": "a"}], + "pagination": {"hasNextPage": True, "nextCursor": "c1"}, + }, + { + "pods": [{"id": "b"}], + "pagination": {"hasNextPage": False, "nextCursor": None}, + }, + ] + with patch.object(api, "run_rest_request", side_effect=pages) as request: + assert api.list_all("/v2/pods", "pods", {"limit": 1}) == [ + {"id": "a"}, + {"id": "b"}, + ] + assert request.call_args_list[1].kwargs["params"] == { + "limit": 1, + "cursor": "c1", + } diff --git a/tests/test_cli/test_cli_options.py b/tests/test_cli/test_cli_options.py new file mode 100644 index 00000000..aaae9516 --- /dev/null +++ b/tests/test_cli/test_cli_options.py @@ -0,0 +1,21 @@ +"""shared command option parsing.""" + +import click +import pytest + +from runpod.rp_cli.options import parse_env + + +def test_env_values_preserve_equals_and_allow_empty_values(): + assert parse_env(["TOKEN=a=b", "EMPTY=", "TOKEN=c=d"]) == { + "TOKEN": "c=d", + "EMPTY": "", + } + + +@pytest.mark.parametrize("value", ["secret-without-equals", "=secret"]) +def test_env_error_does_not_echo_the_value(value): + with pytest.raises(click.UsageError) as exc: + parse_env([value]) + assert "KEY=VALUE" in str(exc.value) + assert "secret" not in str(exc.value) diff --git a/tests/test_cli/test_cli_output.py b/tests/test_cli/test_cli_output.py new file mode 100644 index 00000000..b8d59249 --- /dev/null +++ b/tests/test_cli/test_cli_output.py @@ -0,0 +1,90 @@ +"""shared cli rendering and terminal interaction.""" + +from unittest.mock import patch + +import click +import pytest + +from runpod import error +from runpod.rp_cli import output + + +class TestOutputHelpers: + def test_cell(self): + assert output.cell(None) == "-" + assert output.cell(True) == "yes" + assert output.cell(["a", "b"]) == "a, b" + assert output.cell(3) == "3" + + def test_print_table_keeps_ids_and_values_whole(self, monkeypatch, capsys): + from runpod.rp_cli import console as ui + + monkeypatch.setattr(ui.console, "_width", 40) + rows = [{"id": "cmvfsn771000007jp15tn8esm", "type": "high_performance"}] + output.print_table( + rows, + [ + ("ID", lambda r: r["id"]), + ("TYPE", lambda r: r["type"]), + ("PAD", lambda r: "x" * 30), + ], + "none", + ) + out = capsys.readouterr().out + assert "cmvfsn771000007jp15tn8esm" in out + assert "…" not in out + + def test_print_fields(self, capsys): + output.print_fields( + {"id": "abc", "tags": []}, + [("id", lambda r: r["id"]), ("tags", lambda r: r["tags"])], + ) + out = capsys.readouterr().out + assert "id abc" in out + assert "tags -" in out + + def test_pick_requires_tty(self): + with patch.object(output, "interactive", return_value=False): + with pytest.raises(click.UsageError, match="pass an ID"): + output.pick("pick", [("a", "a")], "pass an ID") + + def test_pick_with_no_choices(self): + with patch.object(output, "interactive", return_value=True): + with pytest.raises(click.UsageError, match="none found"): + output.pick("pick", [], "pass an ID") + + def test_pick_uses_fuzzy_prompt(self): + with ( + patch.object(output, "interactive", return_value=True), + patch("InquirerPy.inquirer.fuzzy") as fuzzy, + ): + fuzzy.return_value.execute.return_value = "b" + assert output.pick("pick", [("a", "first"), ("b", "second")], "x") == "b" + choices = fuzzy.call_args.kwargs["choices"] + assert [(choice.value, choice.name) for choice in choices] == [ + ("a", "first"), + ("b", "second"), + ] + + def test_confirm(self): + output.confirm("sure?", yes=True) + with patch.object(output, "interactive", return_value=False): + with pytest.raises(click.UsageError, match="--yes"): + output.confirm("sure?", yes=False) + with ( + patch.object(output, "interactive", return_value=True), + patch("click.confirm") as ask, + ): + output.confirm("sure?", yes=False) + ask.assert_called_once_with("sure?", abort=True) + + def test_completer(self): + complete = output.completer(lambda: [("abc", "first"), ("xyz", "second")]) + items = complete(None, None, "a") + assert [(item.value, item.help) for item in items] == [("abc", "first")] + + def test_completer_swallows_errors(self): + def broken(): + raise error.AuthenticationError("no key") + + assert output.completer(broken)(None, None, "") == [] diff --git a/tests/test_cli/test_compute.py b/tests/test_cli/test_compute.py index 4debe7cb..d7124b4c 100644 --- a/tests/test_cli/test_compute.py +++ b/tests/test_cli/test_compute.py @@ -26,10 +26,6 @@ ] -def _ids(gpus): - return [gpu["id"] for gpu in gpus] - - class TestMatchGpus: @pytest.mark.parametrize( "query, expected", @@ -51,7 +47,7 @@ class TestMatchGpus: ], ) def test_queries(self, query, expected): - assert _ids(compute.match_gpus(query, GPUS)) == expected + assert [gpu["id"] for gpu in compute.match_gpus(query, GPUS)] == expected def test_resolve_raises_with_known_names(self): with pytest.raises( @@ -157,27 +153,6 @@ def test_best(self): class TestCatalogRequests: - def test_list_all_follows_cursors(self): - pages = [ - { - "pods": [{"id": "a"}], - "pagination": {"hasNextPage": True, "nextCursor": "c1"}, - }, - { - "pods": [{"id": "b"}], - "pagination": {"hasNextPage": False, "nextCursor": None}, - }, - ] - with patch.object(compute, "run_rest_request", side_effect=pages) as request: - assert compute.list_all("/v2/pods", "pods", {"limit": 1}) == [ - {"id": "a"}, - {"id": "b"}, - ] - assert request.call_args_list[1].kwargs["params"] == { - "limit": 1, - "cursor": "c1", - } - def test_list_gpus_with_availability(self): with patch.object( compute, "run_rest_request", return_value={"gpus": GPUS} @@ -200,7 +175,9 @@ def test_list_gpus_with_availability(self): } def test_list_gpus_without_availability(self): - with patch.object(compute, "run_rest_request", return_value=None) as request: + with patch.object( + compute, "run_rest_request", return_value={"gpus": []} + ) as request: assert compute.list_gpus(count=2, min_cuda_version="12.4") == [] assert request.call_args.kwargs["params"] == {} From 139713ab3cd3a38c66233e6be0d3f1aa2824521c Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:48:34 -0700 Subject: [PATCH 12/25] refactor: simplify volume commands and shared API integration --- runpod/rp_cli/volume.py | 330 +++++-------- tests/test_cli/test_volume_command.py | 687 +++++++++++--------------- 2 files changed, 409 insertions(+), 608 deletions(-) diff --git a/runpod/rp_cli/volume.py b/runpod/rp_cli/volume.py index f28212e4..61898d79 100644 --- a/runpod/rp_cli/volume.py +++ b/runpod/rp_cli/volume.py @@ -1,14 +1,13 @@ """`rp volume`: network and global volume management.""" import re -from typing import NamedTuple, Optional import click from runpod.api.rest import run_rest_request -from runpod.rp_cli import compute +from runpod.rp_cli import compute, console as ui +from runpod.rp_cli.api import api_errors, list_all from runpod.rp_cli.output import ( - api_errors, completer, confirm, json_option, @@ -18,104 +17,69 @@ print_table, ) -MIN_SIZE = 10 -MAX_SIZE = 4096 - - -class Kind(NamedTuple): - path: str - collection: str - noun: str - - -NETWORK = Kind("/v2/network-volumes", "networkVolumes", "network volume") -GLOBAL = Kind("/v2/global-volumes", "globalVolumes", "global volume") - - -def _kind(global_volume: bool) -> Kind: - return GLOBAL if global_volume else NETWORK - - -def _type_label(record: dict) -> Optional[str]: - return (record.get("type") or "").lower() or None +NETWORK_FIELDS = [ + ("id", lambda row: row["id"]), + ("name", lambda row: row["name"]), + ("size", lambda row: f"{row['size']} GB"), + ("data center", lambda row: row["dataCenter"]), + ("type", lambda row: row["type"].lower()), +] +GLOBAL_FIELDS = [ + ("id", lambda row: row["id"]), + ("name", lambda row: row["name"]), + ("created", lambda row: row["createdAt"]), + ("created by", lambda row: row["createdByUserId"]), +] +global_option = click.option( + "--global", + "kind", + flag_value="global", + default="network", + help="Manage global volumes, which mount in any data center.", +) -def _size_label(record: dict) -> Optional[str]: - return f"{record['size']} GB" if record.get("size") else None +def _label(record): + if "size" not in record: + return record["name"] + return ( + f"{record['name']} {record['size']} GB " + f"{record['dataCenter']} {record['type'].lower()}" + ) -def _columns(kind: Kind): - if kind is GLOBAL: - return [ - ("ID", lambda row: row["id"]), - ("NAME", lambda row: row.get("name")), - ("CREATED", lambda row: (row.get("createdAt") or "")[:10] or None), - ] +def _choices(kind): return [ - ("ID", lambda row: row["id"]), - ("NAME", lambda row: row.get("name")), - ("SIZE", _size_label), - ("DATA CENTER", lambda row: row.get("dataCenter")), - ("TYPE", _type_label), + (row["id"], _label(row)) + for row in list_all(f"/v2/{kind}-volumes", f"{kind}Volumes") ] -def _list(kind: Kind) -> list[dict]: - return compute.list_all(kind.path, kind.collection) - - -def _get(kind: Kind, volume_id: str) -> dict: - return run_rest_request("GET", f"{kind.path}/{volume_id}") or {} - - -def _label(kind: Kind, record: dict) -> str: - if kind is GLOBAL: - return record.get("name") or "" - details = [_size_label(record), record.get("dataCenter"), _type_label(record)] - return " ".join([record.get("name") or ""] + [d for d in details if d]) - - -def _choices(kind: Kind) -> list[tuple[str, str]]: - return [(record["id"], _label(kind, record)) for record in _list(kind)] - - def _complete_ids(ctx, param, incomplete): - kind = _kind(bool(ctx.params.get("global_volume"))) - return completer(lambda: _choices(kind))(ctx, param, incomplete) + return completer(lambda: _choices(ctx.params.get("kind", "network")))( + ctx, param, incomplete + ) -def _show(kind: Kind, record: dict, as_json: bool) -> None: +def _show(record, kind, as_json, action=None): if as_json: print_json(record) return - fields = [(header.lower(), fn) for header, fn in _columns(kind)] - if kind is GLOBAL: - fields[2:] = [ - ("created", lambda row: row.get("createdAt")), - ("created by", lambda row: row.get("createdByUserId")), - ] - print_fields(record, fields) - - -def _data_center( - regions: tuple[str, ...], data_center: Optional[str], volume_type: Optional[str] -) -> str: - """the single data center a network volume is created in. - - regions narrow to data centers offering the requested volume type - (or any type), then a person picks one; scripts get the candidates. - """ - if data_center: - return data_center + if action: + ui.success(f"{action} {kind} volume [white]{record['id']}[/white]") + print_fields(record, GLOBAL_FIELDS if kind == "global" else NETWORK_FIELDS) + + +def _data_center(regions, volume_type): datacenters = compute.list_datacenters() if regions: - allowed = set(compute.resolve_regions(regions, datacenters)) + allowed = compute.resolve_regions(regions, datacenters) datacenters = [dc for dc in datacenters if dc["id"] in allowed] candidates = [ dc for dc in datacenters - if (volume_type in (dc.get("networkVolumeTypes") or [])) - or (not volume_type and dc.get("networkVolumeTypes")) + if dc["networkVolumeTypes"] + and (not volume_type or volume_type in dc["networkVolumeTypes"]) ] scope = f"region {', '.join(regions)}" if regions else "the catalog" offered = ( @@ -125,51 +89,41 @@ def _data_center( raise compute.ResolutionError(f"no data center in {scope} offers {offered}") if len(candidates) == 1: return candidates[0]["id"] - ids = ", ".join(dc["id"] for dc in candidates) return pick( "data center", [ ( dc["id"], - f"{dc['id']} {(dc.get('region') or '').lower().replace('_', ' ')} " + f"{dc['id']} {dc['region'].lower().replace('_', ' ')} " + ", ".join(t.lower() for t in dc["networkVolumeTypes"]), ) for dc in candidates ], - f"several data centers in {scope} offer {offered} ({ids}). " - "pass one with --data-center", + f"several data centers in {scope} offer {offered} " + f"({', '.join(dc['id'] for dc in candidates)}). pass one with --data-center", ) -def _parse_growth(value: str) -> tuple[bool, int]: - """split an expand size into (relative, gigabytes): N or +N.""" - text = value.strip() +def _growth(volume_id, size, expand=False): + """resolve and validate a size against the current network volume.""" + text = str(size).strip() if not re.fullmatch(r"\+?\d+", text): raise click.BadParameter( - f"'{value}' is not a size. use N for a total or +N to add N GB", + f"'{size}' is not a size. use N for a total or +N to add N GB", param_hint="--size", ) - return text.startswith("+"), int(text.lstrip("+")) - - -def _check_growth(volume_id: str, current: int, size: int, strict: bool) -> None: - if size < current or (strict and size == current): + current = run_rest_request("GET", f"/v2/network-volumes/{volume_id}")["size"] + target = current + int(text) if text.startswith("+") else int(text) + if target < current or (expand and target == current): raise click.UsageError( f"network volumes can only grow: {volume_id} is {current} GB, " - f"requested {size} GB" + f"requested {target} GB" ) - if size > MAX_SIZE: + if target > 4096: raise click.UsageError( - f"network volumes hold at most {MAX_SIZE} GB, requested {size} GB" + f"network volumes hold at most 4096 GB, requested {target} GB" ) - - -global_option = click.option( - "--global", - "global_volume", - is_flag=True, - help="Manage global volumes, which mount in any data center.", -) + return current, target @click.group() @@ -180,16 +134,20 @@ def volume(): @volume.command(name="list") @global_option @json_option -def volume_list(global_volume, as_json): +def volume_list(kind, as_json): """List volumes with their size, data center, and type.""" - kind = _kind(global_volume) with api_errors(): - rows = _list(kind) + rows = list_all(f"/v2/{kind}-volumes", f"{kind}Volumes") if as_json: print_json(rows) return + fields = NETWORK_FIELDS + if kind == "global": + fields = GLOBAL_FIELDS[:2] + [("created", lambda row: row["createdAt"][:10])] print_table( - rows, _columns(kind), f"no {kind.noun}s. create one with rp volume create" + rows, + [(label.upper(), value) for label, value in fields], + f"no {kind} volumes. create one with rp volume create", ) @@ -197,23 +155,21 @@ def volume_list(global_volume, as_json): @click.argument("volume_id", required=False, shell_complete=_complete_ids) @global_option @json_option -def volume_get(volume_id, global_volume, as_json): +def volume_get(volume_id, kind, as_json): """Show one volume (pick one when VOLUME_ID is omitted).""" - kind = _kind(global_volume) with api_errors(): if volume_id is None: - volume_id = pick(kind.noun, _choices(kind), "pass VOLUME_ID") - record = _get(kind, volume_id) - _show(kind, record, as_json) + volume_id = pick(f"{kind} volume", _choices(kind), "pass VOLUME_ID") + record = run_rest_request("GET", f"/v2/{kind}-volumes/{volume_id}") + _show(record, kind, as_json) @volume.command(name="create") @click.option("--name", required=True, help="Human-readable name.") @click.option( "--size", - type=click.IntRange(MIN_SIZE, MAX_SIZE), - default=None, - help=f"Storage in GB ({MIN_SIZE}-{MAX_SIZE}). Required for network volumes.", + type=click.IntRange(10, 4096), + help="Storage in GB (10-4096). Required for network volumes.", ) @click.option( "--region", @@ -221,20 +177,17 @@ def volume_get(volume_id, global_volume, as_json): multiple=True, help="Region (eu, na), id prefix (us, eu-ro), or data center id. Repeatable.", ) -@click.option("--data-center", default=None, help="Exact data center id, like EU-RO-1.") +@click.option("--data-center", help="Exact data center id, like EU-RO-1.") @click.option( "--type", "volume_type", type=click.Choice(["standard", "high_performance"], case_sensitive=False), - default=None, help="Storage tier, immutable after creation. Defaults to the data center's " "primary tier.", ) @global_option @json_option -def volume_create( - name, size, regions, data_center, volume_type, global_volume, as_json -): # pylint: disable=too-many-arguments +def volume_create(name, size, regions, data_center, volume_type, kind, as_json): """Create a network volume in one data center, or a global volume. A region matching several data centers narrows to those offering @@ -245,82 +198,65 @@ def volume_create( rp volume create --name scratch --size 50 --data-center EU-RO-1 rp volume create --name shared --global """ - from runpod.rp_cli import console as ui - - kind = _kind(global_volume) - if global_volume: - given = [ - flag - for flag, value in ( - ("--size", size), - ("--region", regions), - ("--data-center", data_center), - ("--type", volume_type), - ) - if value - ] - if given: - raise click.UsageError(f"{', '.join(given)} do not apply to global volumes") - body = {"name": name} - else: - if size is None: - raise click.UsageError("pass --size in GB for a network volume") - if regions and data_center: - raise click.UsageError("use either --region or --data-center") - volume_type = volume_type.upper() if volume_type else None - with api_errors(): - body = { - "name": name, - "size": size, - "dataCenter": _data_center(regions, data_center, volume_type), - } - if volume_type: - body["type"] = volume_type - + body = {"name": name} with api_errors(): - record = run_rest_request("POST", kind.path, json=body) or {} - if as_json: - print_json(record) - return - ui.success(f"created {kind.noun} [white]{record.get('id')}[/white]") - _show(kind, record, as_json) + if kind == "global": + given = [ + flag + for flag, value in ( + ("--size", size), + ("--region", regions), + ("--data-center", data_center), + ("--type", volume_type), + ) + if value + ] + if given: + raise click.UsageError( + f"{', '.join(given)} do not apply to global volumes" + ) + else: + if size is None: + raise click.UsageError("pass --size in GB for a network volume") + if regions and data_center: + raise click.UsageError("use either --region or --data-center") + volume_type = volume_type.upper() if volume_type else None + body.update( + size=size, + dataCenter=data_center or _data_center(regions, volume_type), + ) + if volume_type: + body["type"] = volume_type + record = run_rest_request("POST", f"/v2/{kind}-volumes", json=body) + _show(record, kind, as_json, "created") @volume.command(name="update") @click.argument("volume_id", shell_complete=_complete_ids) -@click.option("--name", default=None, help="New name.") +@click.option("--name", help="New name.") @click.option( "--size", - type=click.IntRange(MIN_SIZE, MAX_SIZE), - default=None, + type=click.IntRange(10, 4096), help="New size in GB; network volumes can only grow.", ) @global_option @json_option -def volume_update(volume_id, name, size, global_volume, as_json): +def volume_update(volume_id, name, size, kind, as_json): """Rename and/or resize a volume.""" - from runpod.rp_cli import console as ui - - kind = _kind(global_volume) - if global_volume and size is not None: + if kind == "global" and size is not None: raise click.UsageError("global volumes have no size to change") - if name is None and size is None: - raise click.UsageError("pass --name and/or --size") body = { key: value for key, value in (("name", name), ("size", size)) if value is not None } + if not body: + raise click.UsageError("pass --name and/or --size") with api_errors(): if size is not None: - current = _get(kind, volume_id).get("size") or 0 - _check_growth(volume_id, current, size, strict=False) - record = run_rest_request("PATCH", f"{kind.path}/{volume_id}", json=body) or {} - if as_json: - print_json(record) - return - ui.success(f"updated {kind.noun} [white]{volume_id}[/white]") - _show(kind, record, as_json) + _growth(volume_id, size) + record = run_rest_request("PATCH", f"/v2/{kind}-volumes/{volume_id}", json=body) + _show(record, kind, as_json, "updated") @volume.command(name="expand") @@ -334,44 +270,32 @@ def volume_update(volume_id, name, size, global_volume, as_json): @json_option def volume_expand(volume_id, size, as_json): """Grow a network volume; storage can never shrink.""" - from runpod.rp_cli import console as ui - - relative, amount = _parse_growth(size) with api_errors(): - current = _get(NETWORK, volume_id).get("size") or 0 - target = current + amount if relative else amount - _check_growth(volume_id, current, target, strict=True) - record = ( - run_rest_request( - "PATCH", f"{NETWORK.path}/{volume_id}", json={"size": target} - ) - or {} + current, target = _growth(volume_id, size, expand=True) + record = run_rest_request( + "PATCH", f"/v2/network-volumes/{volume_id}", json={"size": target} ) if as_json: print_json(record) - return - ui.success( - f"expanded [white]{volume_id}[/white] from {current} GB to " - f"{record.get('size', target)} GB" - ) + else: + ui.success( + f"expanded [white]{volume_id}[/white] from {current} GB to {record['size']} GB" + ) @volume.command(name="delete") @click.argument("volume_id", shell_complete=_complete_ids) @click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.") @global_option -def volume_delete(volume_id, yes, global_volume): +def volume_delete(volume_id, yes, kind): """Delete a volume and everything stored on it.""" - from runpod.rp_cli import console as ui - - kind = _kind(global_volume) + path = f"/v2/{kind}-volumes/{volume_id}" with api_errors(): if not yes: - record = _get(kind, volume_id) + record = run_rest_request("GET", path) confirm( - f"delete {kind.noun} {volume_id} ({_label(kind, record)}) " - "and all its data?", + f"delete {kind} volume {volume_id} ({_label(record)}) and all its data?", yes, ) - run_rest_request("DELETE", f"{kind.path}/{volume_id}") - ui.success(f"deleted {kind.noun} [white]{volume_id}[/white]") + run_rest_request("DELETE", path) + ui.success(f"deleted {kind} volume [white]{volume_id}[/white]") diff --git a/tests/test_cli/test_volume_command.py b/tests/test_cli/test_volume_command.py index f66441a4..5faa4a29 100644 --- a/tests/test_cli/test_volume_command.py +++ b/tests/test_cli/test_volume_command.py @@ -1,9 +1,10 @@ """tests for the rp volume command group.""" import json -from unittest.mock import patch +from unittest.mock import call, patch import pytest +from click.shell_completion import ShellComplete from click.testing import CliRunner from runpod import error @@ -17,10 +18,10 @@ "type": "STANDARD", } GLOBAL_VOLUME = { - "id": "gsb_abc", - "name": "shared", + "id": "cmvfsn771000007jp15tn8esm", + "name": "capy-test-global-renamed", "createdAt": "2026-09-24T00:00:00Z", - "createdByUserId": "user_123", + "createdByUserId": "user_example", } DATACENTERS = [ {"id": "EU-RO-1", "region": "EUROPE", "networkVolumeTypes": ["STANDARD"]}, @@ -35,19 +36,19 @@ ] -def invoke(*args): - return CliRunner().invoke(cli, ["volume", *args]) +def invoke(command, **kwargs): + return CliRunner().invoke(cli, f"volume {command}", **kwargs) @pytest.fixture def rest(): - with patch("runpod.rp_cli.volume.run_rest_request") as mock: + with patch("runpod.rp_cli.volume.run_rest_request", return_value=VOLUME) as mock: yield mock @pytest.fixture def listing(): - with patch("runpod.rp_cli.compute.list_all") as mock: + with patch("runpod.rp_cli.volume.list_all", return_value=[VOLUME]) as mock: yield mock @@ -59,406 +60,282 @@ def datacenters(): yield mock -def test_volume_is_mounted(): - assert "volume" in cli.commands - - -class TestList: - def test_network_table(self, listing): - listing.return_value = [VOLUME] - result = invoke("list") - assert result.exit_code == 0, result.output - listing.assert_called_once_with("/v2/network-volumes", "networkVolumes") - for text in ( - "ID", - "DATA CENTER", - "TYPE", - "vol1", - "models", - "100 GB", - "EU-RO-1", - "standard", - ): - assert text in result.output +@pytest.fixture(params=["network", "global"]) +def resource(request, rest): + kind = request.param + record = GLOBAL_VOLUME if kind == "global" else VOLUME + rest.return_value = record + return kind, record, "--global" if kind == "global" else "" + - def test_global_json(self, listing): - listing.return_value = [GLOBAL_VOLUME] - result = invoke("list", "--global", "--json") - assert result.exit_code == 0, result.output - listing.assert_called_once_with("/v2/global-volumes", "globalVolumes") - assert json.loads(result.output) == [GLOBAL_VOLUME] - - def test_global_table_keeps_long_ids_whole(self, listing): - listing.return_value = [ - dict( - GLOBAL_VOLUME, - id="cmvfsn771000007jp15tn8esm", - name="capy-test-global-renamed", - ) - ] - result = invoke("list", "--global") - assert "cmvfsn771000007jp15tn8esm" in result.output +def test_list_table_and_json(resource, listing): + kind, record, flag = resource + listing.return_value = [record] + result = invoke(f"list {flag}") + assert result.exit_code == 0, result.output + assert record["id"] in result.output + if kind == "global": assert "2026-09-24" in result.output - assert "user_123" not in result.output - - def test_empty(self, listing): - listing.return_value = [] - result = invoke("list") - assert result.exit_code == 0 - assert "no network volumes" in result.output - - def test_api_error(self, listing): - listing.side_effect = error.AuthenticationError("bad key") - result = invoke("list") - assert result.exit_code == 1 - assert "bad key" in result.output - - -class TestGet: - def test_fields(self, rest): - rest.return_value = VOLUME - result = invoke("get", "vol1") - assert result.exit_code == 0, result.output - rest.assert_called_once_with("GET", "/v2/network-volumes/vol1") - assert "data center" in result.output and "EU-RO-1" in result.output - assert "standard" in result.output - - def test_global_fields(self, rest): - rest.return_value = GLOBAL_VOLUME - result = invoke("get", "gsb_abc", "--global") - assert result.exit_code == 0, result.output - assert "2026-09-24T00:00:00Z" in result.output - assert "created by" in result.output and "user_123" in result.output - - def test_global_json(self, rest): - rest.return_value = GLOBAL_VOLUME - result = invoke("get", "gsb_abc", "--global", "--json") - rest.assert_called_once_with("GET", "/v2/global-volumes/gsb_abc") - assert json.loads(result.output) == GLOBAL_VOLUME - - def test_picker(self, rest, listing): - listing.return_value = [VOLUME] - rest.return_value = VOLUME - with patch("runpod.rp_cli.volume.pick", return_value="vol1") as pick: - result = invoke("get") - assert result.exit_code == 0, result.output - choices = pick.call_args.args[1] - assert choices == [("vol1", "models 100 GB EU-RO-1 standard")] - rest.assert_called_once_with("GET", "/v2/network-volumes/vol1") - - def test_missing_id_without_tty(self, rest, listing): - listing.return_value = [VOLUME] + assert record["createdByUserId"] not in result.output + else: + for text in ("ID", "DATA CENTER", "TYPE", "100 GB", "EU-RO-1", "standard"): + assert text in result.output + result = invoke(f"list {flag} --json") + assert json.loads(result.output) == [record] + listing.assert_called_with(f"/v2/{kind}-volumes", f"{kind}Volumes") + + +def test_list_empty(listing): + listing.return_value = [] + assert "no network volumes" in invoke("list").output + + +def test_get_fields_and_json(resource, rest): + kind, record, flag = resource + result = invoke(f"get {record['id']} {flag}") + assert result.exit_code == 0, result.output + for value in record.values(): + assert str(value).lower() in result.output.lower() + result = invoke(f"get {record['id']} {flag} --json") + assert json.loads(result.output) == record + rest.assert_called_with("GET", f"/v2/{kind}-volumes/{record['id']}") + + +def test_get_picker(rest, listing): + with patch("runpod.rp_cli.volume.pick", return_value="vol1") as pick: result = invoke("get") - assert result.exit_code == 2 - assert "pass VOLUME_ID" in result.output - rest.assert_not_called() - - def test_not_found(self, rest): - rest.side_effect = error.QueryError("network volume not found", status_code=404) - result = invoke("get", "nope") - assert result.exit_code == 1 - assert "not found" in result.output - - -class TestCreate: - def test_exact_data_center(self, rest, datacenters): - rest.return_value = VOLUME - result = invoke( - "create", - "--name", - "models", - "--size", - "100", - "--data-center", - "EU-RO-1", - "--type", - "standard", - ) - assert result.exit_code == 0, result.output - datacenters.assert_not_called() - rest.assert_called_once_with( - "POST", - "/v2/network-volumes", - json={ - "name": "models", - "size": 100, - "dataCenter": "EU-RO-1", - "type": "STANDARD", - }, - ) - assert "created network volume" in result.output and "vol1" in result.output - - def test_region_narrows_by_type(self, rest, datacenters): - rest.return_value = VOLUME - with patch("runpod.rp_cli.volume.pick", return_value="EU-CZ-1") as pick: - result = invoke( - "create", - "--name", - "m", - "--size", - "10", - "--region", - "eu", - "--type", - "high_performance", - "--json", - ) - assert result.exit_code == 0, result.output - assert [value for value, _ in pick.call_args.args[1]] == ["EU-FR-1", "EU-CZ-1"] - assert rest.call_args.kwargs["json"] == { - "name": "m", - "size": 10, - "dataCenter": "EU-CZ-1", - "type": "HIGH_PERFORMANCE", - } - assert json.loads(result.output) == VOLUME - - def test_region_with_one_candidate(self, rest, datacenters): - rest.return_value = VOLUME - result = invoke("create", "--name", "m", "--size", "10", "--region", "na") - assert result.exit_code == 0, result.output - assert rest.call_args.kwargs["json"] == { + assert result.exit_code == 0, result.output + assert pick.call_args.args[1] == [("vol1", "models 100 GB EU-RO-1 standard")] + rest.assert_called_once_with("GET", "/v2/network-volumes/vol1") + + +def test_missing_id_requires_terminal(rest, listing): + result = invoke("get") + assert result.exit_code == 2 + assert "pass VOLUME_ID" in result.output + rest.assert_not_called() + + +@pytest.mark.parametrize( + "command,exception", + [ + ("get vol1", error.QueryError("not found", status_code=404)), + ("get vol1", error.AuthenticationError("bad key")), + ( + "delete vol1 --global --yes", + error.QueryError("still mounted", status_code=409), + ), + ], +) +def test_api_errors_are_clean(rest, command, exception): + rest.side_effect = exception + result = invoke(command) + assert result.exit_code == 1 + assert str(exception) in result.output + + +def test_create_exact_data_center(rest, datacenters): + result = invoke( + "create --name models --size 100 --data-center EU-RO-1 --type standard" + ) + assert result.exit_code == 0, result.output + datacenters.assert_not_called() + rest.assert_called_once_with( + "POST", + "/v2/network-volumes", + json={ + "name": "models", + "size": 100, + "dataCenter": "EU-RO-1", + "type": "STANDARD", + }, + ) + assert "created network volume" in result.output + + +@pytest.mark.parametrize( + "options,candidates", + [ + ("--region eu --type high_performance", ["EU-FR-1", "EU-CZ-1"]), + ("--type standard", ["EU-RO-1", "EU-CZ-1", "US-TX-3"]), + ("--region eu --region na", ["EU-RO-1", "EU-FR-1", "EU-CZ-1", "US-TX-3"]), + ], +) +def test_create_filters_picker_candidates(rest, datacenters, options, candidates): + with patch("runpod.rp_cli.volume.pick", return_value=candidates[0]) as pick: + result = invoke(f"create --name m --size 10 {options} --json") + assert result.exit_code == 0, result.output + assert [value for value, _ in pick.call_args.args[1]] == candidates + assert rest.call_args.kwargs["json"]["dataCenter"] == candidates[0] + assert json.loads(result.output) == VOLUME + + +def test_create_single_region_match_uses_default_tier(rest, datacenters): + result = invoke("create --name m --size 10 --region na") + assert result.exit_code == 0, result.output + rest.assert_called_once_with( + "POST", + "/v2/network-volumes", + json={ "name": "m", "size": 10, "dataCenter": "US-TX-3", - } - - def test_ambiguous_region_without_tty_lists_candidates(self, rest, datacenters): - result = invoke("create", "--name", "m", "--size", "10", "--region", "eu") - assert result.exit_code == 2 - assert "EU-RO-1, EU-FR-1, EU-CZ-1" in result.output - assert "EU-SE-1" not in result.output - assert "--data-center" in result.output - rest.assert_not_called() - - def test_no_region_picks_from_every_data_center(self, rest, datacenters): - rest.return_value = VOLUME - with patch("runpod.rp_cli.volume.pick", return_value="US-TX-3") as pick: - result = invoke( - "create", "--name", "m", "--size", "10", "--type", "standard" - ) - assert result.exit_code == 0, result.output - labels = dict(pick.call_args.args[1]) - assert set(labels) == {"EU-RO-1", "EU-CZ-1", "US-TX-3"} - assert labels["EU-CZ-1"] == "EU-CZ-1 europe standard, high_performance" - - def test_region_without_type_support(self, rest, datacenters): - result = invoke( - "create", - "--name", - "m", - "--size", - "10", - "--region", - "na", - "--type", - "high_performance", - ) - assert result.exit_code == 1 - assert "no data center in region na offers high_performance" in result.output - rest.assert_not_called() - - def test_unknown_region(self, rest, datacenters): - result = invoke("create", "--name", "m", "--size", "10", "--region", "mars") - assert result.exit_code == 1 - assert "unknown region 'mars'" in result.output - - def test_requires_size(self, rest): - result = invoke("create", "--name", "m", "--data-center", "EU-RO-1") - assert result.exit_code == 2 - assert "--size" in result.output - rest.assert_not_called() - - def test_size_bounds(self, rest): - result = invoke( - "create", "--name", "m", "--size", "5000", "--data-center", "EU-RO-1" - ) - assert result.exit_code == 2 - rest.assert_not_called() - - def test_region_and_data_center_conflict(self, rest): - result = invoke( - "create", - "--name", - "m", - "--size", - "10", - "--region", - "eu", - "--data-center", - "EU-RO-1", - ) - assert result.exit_code == 2 - assert "either --region or --data-center" in result.output - - def test_global(self, rest): - rest.return_value = GLOBAL_VOLUME - result = invoke("create", "--name", "shared", "--global") - assert result.exit_code == 0, result.output - rest.assert_called_once_with( - "POST", "/v2/global-volumes", json={"name": "shared"} - ) - assert "created global volume" in result.output and "gsb_abc" in result.output - - def test_global_rejects_network_flags(self, rest): - result = invoke( - "create", "--name", "s", "--global", "--size", "10", "--region", "eu" - ) - assert result.exit_code == 2 - assert "--size, --region do not apply to global volumes" in result.output - rest.assert_not_called() - - -class TestUpdate: - def test_rename(self, rest): - rest.return_value = dict(VOLUME, name="renamed") - result = invoke("update", "vol1", "--name", "renamed") - assert result.exit_code == 0, result.output - rest.assert_called_once_with( - "PATCH", "/v2/network-volumes/vol1", json={"name": "renamed"} - ) - assert "renamed" in result.output - - def test_rename_and_resize(self, rest): - rest.side_effect = [VOLUME, dict(VOLUME, name="r", size=200)] - result = invoke("update", "vol1", "--name", "r", "--size", "200", "--json") - assert result.exit_code == 0, result.output - assert rest.call_args.args == ("PATCH", "/v2/network-volumes/vol1") - assert rest.call_args.kwargs["json"] == {"name": "r", "size": 200} - assert json.loads(result.output)["size"] == 200 - - def test_same_size_is_allowed(self, rest): - rest.side_effect = [VOLUME, VOLUME] - result = invoke("update", "vol1", "--size", "100") - assert result.exit_code == 0, result.output - - def test_rejects_shrinking(self, rest): - rest.return_value = VOLUME - result = invoke("update", "vol1", "--size", "50") - assert result.exit_code == 2 - assert "can only grow: vol1 is 100 GB, requested 50 GB" in result.output - rest.assert_called_once_with("GET", "/v2/network-volumes/vol1") - - def test_requires_a_change(self, rest): - result = invoke("update", "vol1") - assert result.exit_code == 2 - assert "--name and/or --size" in result.output - - def test_global_rename(self, rest): - rest.return_value = dict(GLOBAL_VOLUME, name="renamed") - result = invoke("update", "gsb_abc", "--global", "--name", "renamed") - assert result.exit_code == 0, result.output - rest.assert_called_once_with( - "PATCH", "/v2/global-volumes/gsb_abc", json={"name": "renamed"} - ) - - def test_global_rejects_size(self, rest): - result = invoke("update", "gsb_abc", "--global", "--size", "20") - assert result.exit_code == 2 - assert "global volumes have no size" in result.output - rest.assert_not_called() - - -class TestExpand: - def test_absolute(self, rest): - rest.side_effect = [VOLUME, dict(VOLUME, size=250)] - result = invoke("expand", "vol1", "--size", "250") - assert result.exit_code == 0, result.output - rest.assert_called_with("PATCH", "/v2/network-volumes/vol1", json={"size": 250}) - assert "from 100 GB to 250 GB" in result.output - - def test_relative(self, rest): - rest.side_effect = [VOLUME, dict(VOLUME, size=150)] - result = invoke("expand", "vol1", "--size", "+50", "--json") - assert result.exit_code == 0, result.output - rest.assert_called_with("PATCH", "/v2/network-volumes/vol1", json={"size": 150}) - assert json.loads(result.output)["size"] == 150 - - @pytest.mark.parametrize("size", ["50", "100", "+0"]) - def test_rejects_shrinking_or_no_growth(self, rest, size): - rest.return_value = VOLUME - result = invoke("expand", "vol1", "--size", size) - assert result.exit_code == 2 - assert "network volumes can only grow: vol1 is 100 GB" in result.output - assert rest.call_count == 1 - - def test_rejects_past_maximum(self, rest): - rest.return_value = VOLUME - result = invoke("expand", "vol1", "--size", "+4000") - assert result.exit_code == 2 - assert "at most 4096 GB, requested 4100 GB" in result.output - assert rest.call_count == 1 - - @pytest.mark.parametrize("size", ["abc", "-5", "1.5"]) - def test_rejects_malformed_size(self, rest, size): - result = invoke("expand", "vol1", "--size", size) - assert result.exit_code == 2 - assert "is not a size" in result.output - rest.assert_not_called() - - -class TestDelete: - def test_yes_skips_confirmation(self, rest): - result = invoke("delete", "vol1", "--yes") - assert result.exit_code == 0, result.output - rest.assert_called_once_with("DELETE", "/v2/network-volumes/vol1") - assert "deleted network volume" in result.output - - def test_confirms_interactively(self, rest): - rest.return_value = VOLUME - with patch("runpod.rp_cli.output.interactive", return_value=True): - result = CliRunner().invoke(cli, ["volume", "delete", "vol1"], input="y\n") - assert result.exit_code == 0, result.output - assert "models 100 GB EU-RO-1 standard" in result.output - rest.assert_called_with("DELETE", "/v2/network-volumes/vol1") - - def test_declined(self, rest): - rest.return_value = VOLUME - with patch("runpod.rp_cli.output.interactive", return_value=True): - result = CliRunner().invoke(cli, ["volume", "delete", "vol1"], input="n\n") - assert result.exit_code == 1 - assert all(call.args[0] == "GET" for call in rest.call_args_list) - - def test_requires_yes_without_tty(self, rest): - rest.return_value = VOLUME - result = invoke("delete", "vol1") - assert result.exit_code == 2 - assert "--yes" in result.output - assert all(call.args[0] == "GET" for call in rest.call_args_list) - - def test_global(self, rest): - result = invoke("delete", "gsb_abc", "--global", "--yes") - assert result.exit_code == 0, result.output - rest.assert_called_once_with("DELETE", "/v2/global-volumes/gsb_abc") - - def test_mounted_global_volume(self, rest): - rest.side_effect = error.QueryError("volume is still mounted", status_code=409) - result = invoke("delete", "gsb_abc", "--global", "--yes") - assert result.exit_code == 1 - assert "still mounted" in result.output - - -class TestCompletion: - def _complete(self, args, incomplete): - from click.shell_completion import ShellComplete - - return [ - item.value - for item in ShellComplete(cli, {}, "rp", "_RP_COMPLETE").get_completions( - args, incomplete - ) - ] - - def test_network_ids(self, listing): - listing.return_value = [VOLUME, dict(VOLUME, id="other")] - assert self._complete(["volume", "get"], "v") == ["vol1"] - listing.assert_called_with("/v2/network-volumes", "networkVolumes") - - def test_global_ids(self, listing): - listing.return_value = [GLOBAL_VOLUME] - assert self._complete(["volume", "delete", "--global"], "") == ["gsb_abc"] - listing.assert_called_with("/v2/global-volumes", "globalVolumes") - - def test_errors_complete_nothing(self, listing): - listing.side_effect = error.AuthenticationError("no key") - assert self._complete(["volume", "update"], "") == [] + }, + ) + + +@pytest.mark.parametrize( + "options,code,message", + [ + ("--region eu", 2, "EU-RO-1, EU-FR-1, EU-CZ-1"), + ("--region na --type high_performance", 1, "no data center"), + ("--region mars", 1, "unknown region"), + ("--region eu --data-center EU-RO-1", 2, "either --region or --data-center"), + ("--size 5000", 2, "4096"), + ], +) +def test_create_rejects_invalid_placement_and_size( + rest, datacenters, options, code, message +): + result = invoke(f"create --name m --size 10 {options}") + assert result.exit_code == code + assert message in result.output + rest.assert_not_called() + + +def test_create_requires_network_size(rest): + result = invoke("create --name m --data-center EU-RO-1") + assert result.exit_code == 2 + assert "pass --size" in result.output + rest.assert_not_called() + + +def test_create_global(rest): + rest.return_value = GLOBAL_VOLUME + result = invoke("create --name shared --global") + assert result.exit_code == 0, result.output + rest.assert_called_once_with("POST", "/v2/global-volumes", json={"name": "shared"}) + + +@pytest.mark.parametrize( + "command,message", + [ + ("create --name s --global --size 10 --region eu", "do not apply"), + ("update vol1 --global --size 20", "global volumes have no size"), + ("update vol1", "pass --name and/or --size"), + ], +) +def test_invalid_updates_and_global_options(rest, command, message): + result = invoke(command) + assert result.exit_code == 2 + assert message in result.output + rest.assert_not_called() + + +def test_rename_without_fetching(resource, rest): + kind, record, flag = resource + result = invoke(f"update {record['id']} {flag} --name renamed") + assert result.exit_code == 0, result.output + rest.assert_called_once_with( + "PATCH", f"/v2/{kind}-volumes/{record['id']}", json={"name": "renamed"} + ) + + +@pytest.mark.parametrize( + "command,target", + [ + ("update vol1 --size 200 --name renamed", 200), + ("update vol1 --size 100", 100), + ("expand vol1 --size 250", 250), + ("expand vol1 --size +50", 150), + ], +) +def test_resize_requests(rest, command, target): + updated = dict(VOLUME, size=target) + rest.side_effect = [VOLUME, updated] + result = invoke(f"{command} --json") + assert result.exit_code == 0, result.output + assert json.loads(result.output) == updated + body = {"size": target} + if "--name" in command: + body["name"] = "renamed" + assert rest.call_args_list == [ + call("GET", "/v2/network-volumes/vol1"), + call("PATCH", "/v2/network-volumes/vol1", json=body), + ] + + +def test_expand_human_output(rest): + rest.side_effect = [VOLUME, dict(VOLUME, size=101)] + result = invoke("expand vol1 --size +1") + assert result.exit_code == 0, result.output + assert "from 100 GB to 101 GB" in result.output + + +@pytest.mark.parametrize( + "command,message", + [ + ("update vol1 --size 50", "can only grow"), + ("expand vol1 --size 50", "can only grow"), + ("expand vol1 --size +0", "can only grow"), + ("expand vol1 --size +4000", "at most 4096 GB"), + ], +) +def test_resize_rejects_shrink_no_growth_and_overflow(rest, command, message): + result = invoke(command) + assert result.exit_code == 2 + assert message in result.output + rest.assert_called_once_with("GET", "/v2/network-volumes/vol1") + + +@pytest.mark.parametrize("size", ["abc", "-5", "1.5"]) +def test_expand_rejects_malformed_size_without_fetching(rest, size): + result = invoke(f"expand vol1 --size {size}") + assert result.exit_code == 2 + assert "is not a size" in result.output + rest.assert_not_called() + + +def test_delete_yes(resource, rest): + kind, record, flag = resource + result = invoke(f"delete {record['id']} {flag} --yes") + assert result.exit_code == 0, result.output + rest.assert_called_once_with("DELETE", f"/v2/{kind}-volumes/{record['id']}") + + +@pytest.mark.parametrize("answer,code", [("y", 0), ("n", 1)]) +def test_delete_confirmation(rest, answer, code): + with patch("runpod.rp_cli.output.interactive", return_value=True): + result = invoke("delete vol1", input=f"{answer}\n") + assert result.exit_code == code + assert "models 100 GB EU-RO-1 standard" in result.output + expected = [call("GET", "/v2/network-volumes/vol1")] + if answer == "y": + expected.append(call("DELETE", "/v2/network-volumes/vol1")) + assert rest.call_args_list == expected + + +def test_delete_requires_yes_without_terminal(rest): + result = invoke("delete vol1") + assert result.exit_code == 2 + assert "pass --yes" in result.output + rest.assert_called_once_with("GET", "/v2/network-volumes/vol1") + + +def test_completion_by_kind(resource, listing): + kind, record, flag = resource + listing.return_value = [record, dict(record, id="other")] + args = ["volume", "get"] + ([flag] if flag else []) + items = ShellComplete(cli, {}, "rp", "_RP_COMPLETE").get_completions( + args, record["id"][:1] + ) + assert [item.value for item in items] == [record["id"]] + listing.assert_called_once_with(f"/v2/{kind}-volumes", f"{kind}Volumes") + + +def test_completion_suppresses_errors(listing): + listing.side_effect = error.AuthenticationError("no key") + items = ShellComplete(cli, {}, "rp", "_RP_COMPLETE").get_completions( + ["volume", "expand"], "" + ) + assert items == [] 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 13/25] 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 14/25] 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 bc5e6595ba42179e91ef20843f46e3d851746f2d Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:07:32 -0700 Subject: [PATCH 15/25] style: remove redundant volume docstrings --- runpod/rp_cli/volume.py | 3 --- tests/test_cli/test_volume_command.py | 2 -- 2 files changed, 5 deletions(-) diff --git a/runpod/rp_cli/volume.py b/runpod/rp_cli/volume.py index 61898d79..95d2e550 100644 --- a/runpod/rp_cli/volume.py +++ b/runpod/rp_cli/volume.py @@ -1,5 +1,3 @@ -"""`rp volume`: network and global volume management.""" - import re import click @@ -105,7 +103,6 @@ def _data_center(regions, volume_type): def _growth(volume_id, size, expand=False): - """resolve and validate a size against the current network volume.""" text = str(size).strip() if not re.fullmatch(r"\+?\d+", text): raise click.BadParameter( diff --git a/tests/test_cli/test_volume_command.py b/tests/test_cli/test_volume_command.py index 5faa4a29..40a4d1dc 100644 --- a/tests/test_cli/test_volume_command.py +++ b/tests/test_cli/test_volume_command.py @@ -1,5 +1,3 @@ -"""tests for the rp volume command group.""" - import json from unittest.mock import call, patch From b33a2717875370a260d7192c501e76fc97124972 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:07:50 -0700 Subject: [PATCH 16/25] refactor: share SDK log streaming across resource commands --- runpod/apps/logs.py | 33 +++-- runpod/rp_cli/catalog.py | 54 +++----- runpod/rp_cli/logs.py | 163 ++++++++++++++++++++++++ runpod/rp_cli/main.py | 3 - tests/test_apps/test_logs.py | 37 ++++++ tests/test_cli/test_catalog_command.py | 10 ++ tests/test_cli/test_cli_logs.py | 166 +++++++++++++++++++++++++ 7 files changed, 416 insertions(+), 50 deletions(-) create mode 100644 runpod/rp_cli/logs.py create mode 100644 tests/test_cli/test_cli_logs.py diff --git a/runpod/apps/logs.py b/runpod/apps/logs.py index c6ec8999..4a68d6d3 100644 --- a/runpod/apps/logs.py +++ b/runpod/apps/logs.py @@ -6,7 +6,7 @@ import json import os -from typing import Any, AsyncIterator, Dict, List, Optional +from typing import Any, AsyncIterator, Callable, Dict, List, Optional import aiohttp @@ -46,15 +46,18 @@ async def stream_pod_logs( pod_id: str, *, endpoint_id: Optional[str] = None, - log_type: str = "all", + log_type: Optional[str] = "all", tail: int = 100, since: Optional[str] = None, + last_event_id: Optional[str] = None, + on_event_id: Optional[Callable[[str], None]] = None, ) -> AsyncIterator[Dict[str, Any]]: """follow a pod's logs as they arrive. yields {"source": "container"|"system", "line": str, "ts": str} parsed from the sse stream. pass endpoint_id for serverless workers. - ends when the server closes the stream or the caller breaks out. + log_type=None selects both sources through rest. last_event_id resumes + a stream; on_event_id receives the cursor of each delivered event. """ params: Dict[str, Any] = {"tail": str(tail)} if log_type == "all": @@ -69,28 +72,32 @@ async def stream_pod_logs( f"/workers/{_path_segment(pod_id)}/logs" ) url = _build_url(path) - params["source"] = log_type + if log_type: + params["source"] = log_type if since: params["since"] = since + headers = _headers() + if last_event_id: + headers["Last-Event-ID"] = last_event_id - client_timeout = aiohttp.ClientTimeout( - total=STREAM_TIMEOUT_SECONDS, sock_read=60 - ) + client_timeout = aiohttp.ClientTimeout(total=STREAM_TIMEOUT_SECONDS, sock_read=60) async with aiohttp.ClientSession(timeout=client_timeout) as session: - async with session.get( - url, params=params, headers=_headers() - ) as resp: + async with session.get(url, params=params, headers=headers) as resp: resp.raise_for_status() + event_id = None async for raw in resp.content: line = raw.decode("utf-8", "replace").strip() - # sse frames: "data: {...}"; comments (heartbeats) start - # with ":" + if line.startswith("id:"): + event_id = line[3:].strip() if not line.startswith("data:"): continue try: - yield json.loads(line[len("data:") :].strip()) + event = json.loads(line[len("data:") :].strip()) except json.JSONDecodeError: continue + if event_id is not None and on_event_id is not None: + on_event_id(event_id) + yield event def tail_summary(logs: Dict[str, List[str]], lines: int = 20) -> str: diff --git a/runpod/rp_cli/catalog.py b/runpod/rp_cli/catalog.py index ffe437fd..20db7791 100644 --- a/runpod/rp_cli/catalog.py +++ b/runpod/rp_cli/catalog.py @@ -15,11 +15,9 @@ def _money(value: Optional[float]) -> str: def _stock(entries: list[dict]) -> str: ranked = sorted( - entries, key=lambda dc: -compute.availability_rank(dc.get("availability")) - ) - return ", ".join( - f"{dc['id']} {(dc.get('availability') or 'NONE').lower()}" for dc in ranked + entries, key=lambda dc: -compute.availability_rank(dc["availability"]) ) + return ", ".join(f"{dc['id']} {dc['availability'].lower()}" for dc in ranked) def _filter_stock(records: list[dict], regions: tuple[str, ...], available: bool): @@ -32,12 +30,12 @@ def _filter_stock(records: list[dict], regions: tuple[str, ...], available: bool row = dict(record) if allowed is not None: row["dataCenters"] = [ - dc for dc in record["dataCenters"] if dc["id"] in allowed + dc for dc in record.get("dataCenters", []) if dc["id"] in allowed ] row["availability"] = compute.best_availability( dc["availability"] for dc in row["dataCenters"] ) - if not available or compute.availability_rank(row["availability"]) > 0: + if not available or compute.availability_rank(row.get("availability")) > 0: yield row @@ -71,7 +69,7 @@ def _gpu_rows( rows = [] for row in _filter_stock(gpus, regions, available): unit = _unit_price(row, product, cloud) - if min_vram and (row.get("memory") or 0) < min_vram: + if min_vram and row["memory"] < min_vram: continue if max_price is not None and (unit is None or unit > max_price): continue @@ -85,7 +83,7 @@ def _gpu_rows( key=lambda row: ( -compute.availability_rank(row.get("availability")), _unit_price(row, product, cloud) or float("inf"), - row.get("name") or row["id"], + row["name"], ) ) return rows @@ -94,12 +92,12 @@ def _gpu_rows( def _gpu_columns(product: str, count: int, cloud: Optional[str]): suffix = f" x{count}" if count > 1 else "" columns = [ - ("GPU", lambda row: row.get("name") or row["id"]), - ("VRAM", lambda row: f"{row['memory']}GB" if row.get("memory") else None), + ("GPU", lambda row: row["name"]), + ("VRAM", lambda row: f"{row['memory']}GB"), ] if product == "SERVERLESS": columns += [ - ("POOL", lambda row: row.get("pool")), + ("POOL", lambda row: row["pool"]), (f"$/HR{suffix}", lambda row: _money(row["totalPrice"].get("serverless"))), ] else: @@ -111,9 +109,7 @@ def _gpu_columns(product: str, count: int, cloud: Optional[str]): lambda row, tier=tier: _money(row["totalPrice"].get(tier)), ) ) - columns.append( - ("MAX", lambda row: max((row.get("maxCount") or {}).values(), default=None)) - ) + columns.append(("MAX", lambda row: max(row["maxCount"].values(), default=None))) columns += [ ("STOCK", lambda row: (row.get("availability") or "NONE").lower()), ("DATA CENTERS", lambda row: _stock(row.get("dataCenters") or [])), @@ -129,7 +125,7 @@ def _cpu_rows(query, product, vcpu, regions, available) -> list[dict]: raise compute.ResolutionError(f"unknown cpu '{query}'") return sorted( _filter_stock(cpus, regions, available), - key=lambda row: -compute.availability_rank(row["availability"]), + key=lambda row: -compute.availability_rank(row.get("availability")), ) @@ -137,21 +133,11 @@ def _cpu_columns(product: str): tier = "serverlessPerVcpu" if product == "SERVERLESS" else "securePerVcpu" return [ ("CPU", lambda row: row["id"]), - ("NAME", lambda row: row.get("name")), - ("GROUP", lambda row: row.get("group")), - ( - "VCPU", - lambda row: ( - "{min}-{max}".format(**row["vcpu"]) if row.get("vcpu") else None - ), - ), - ( - "RAM/VCPU", - lambda row: ( - f"{row['ramGbPerVcpu']:g}GB" if row.get("ramGbPerVcpu") else None - ), - ), - ("$/VCPU-HR", lambda row: _money((row.get("price") or {}).get(tier))), + ("NAME", lambda row: row["name"]), + ("GROUP", lambda row: row["group"]), + ("VCPU", lambda row: "{min}-{max}".format(**row["vcpu"])), + ("RAM/VCPU", lambda row: f"{row['ramGbPerVcpu']:g}GB"), + ("$/VCPU-HR", lambda row: _money(row["price"].get(tier))), ("STOCK", lambda row: (row.get("availability") or "NONE").lower()), ("DATA CENTERS", lambda row: _stock(row.get("dataCenters") or [])), ] @@ -159,13 +145,13 @@ def _cpu_columns(product: str): _DATACENTER_COLUMNS = [ ("DATA CENTER", lambda row: row["id"]), - ("REGION", lambda row: (row.get("region") or "").lower().replace("_", " ")), + ("REGION", lambda row: row["region"].lower().replace("_", " ")), ( "VOLUME TYPES", - lambda row: [kind.lower() for kind in row.get("networkVolumeTypes") or []], + lambda row: [kind.lower() for kind in row["networkVolumeTypes"]], ), - ("GLOBAL NETWORK", lambda row: row.get("globalNetwork")), - ("COMPLIANCE", lambda row: row.get("compliance")), + ("GLOBAL NETWORK", lambda row: row["globalNetwork"]), + ("COMPLIANCE", lambda row: row["compliance"]), ] diff --git a/runpod/rp_cli/logs.py b/runpod/rp_cli/logs.py new file mode 100644 index 00000000..3ab78ebb --- /dev/null +++ b/runpod/rp_cli/logs.py @@ -0,0 +1,163 @@ +"""shared pod and serverless log views backed by the sdk stream reader.""" + +import asyncio +import json +from datetime import datetime, timezone + +import aiohttp +import click + +from runpod.api.ctl_commands import _path_segment +from runpod.api.rest import run_rest_request +from runpod.apps.logs import stream_pod_logs +from runpod.rp_cli import console as ui +from runpod.rp_cli.api import api_errors + +SNAPSHOT_IDLE_SECONDS = 3 +WORKER_POLL_SECONDS = 5 +RECONNECT_SECONDS = 1 + + +async def _show_logs( + pod_id, endpoint_id, worker_ids, follow, source, tail, since, as_json +): + started = datetime.now(timezone.utc) + records = [] + tasks = {} + width = 0 + + def emit(worker, event): + if as_json: + click.echo( + json.dumps({"worker": worker, **event} if endpoint_id else event) + ) + else: + prefix = f"{worker:<{width}} " if endpoint_id else "" + click.echo(f"{prefix}[{event['source']}] {event['line']}") + + async def consume(worker): + cursor = None + + def remember(event_id): + nonlocal cursor + cursor = event_id + + while True: + stream = stream_pod_logs( + worker, + endpoint_id=endpoint_id, + log_type=source, + tail=tail, + since=since, + last_event_id=cursor, + on_event_id=remember, + ) + try: + while True: + event = await asyncio.wait_for( + anext(stream), + timeout=None if follow else SNAPSHOT_IDLE_SECONDS, + ) + if follow: + emit(worker, event) + elif ( + datetime.fromisoformat(event["ts"].replace("Z", "+00:00")) + >= started + ): + return + else: + records.append((worker, event)) + except StopAsyncIteration: + pass + except asyncio.TimeoutError: + pass + except aiohttp.ClientResponseError as exc: + if not endpoint_id or exc.status != 404: + raise + return + except aiohttp.ClientConnectionError: + if not follow: + raise + finally: + await stream.aclose() + if not follow: + return + await asyncio.sleep(RECONNECT_SECONDS) + + async def attach(): + nonlocal width + if pod_id: + ids = [pod_id] + else: + response = await asyncio.to_thread( + run_rest_request, + "GET", + f"/v2/serverless/{_path_segment(endpoint_id)}/workers", + ) + ids = [ + worker["id"] + for worker in response["workers"] + if not worker_ids or worker["id"] in worker_ids + ] + width = max([width, *(len(worker) for worker in ids)]) + for worker in ids: + if worker not in tasks: + tasks[worker] = asyncio.create_task(consume(worker)) + + try: + await attach() + if not tasks and not as_json: + ui.console.print( + "[dim]waiting for workers ...[/dim]" + if follow + else "[dim]no workers[/dim]" + ) + if not follow: + await asyncio.gather(*tasks.values()) + else: + while True: + await asyncio.sleep(WORKER_POLL_SECONDS) + for task in tasks.values(): + if task.done(): + task.result() + tasks = { + worker: task for worker, task in tasks.items() if not task.done() + } + await attach() + finally: + for task in tasks.values(): + task.cancel() + await asyncio.gather(*tasks.values(), return_exceptions=True) + for worker, event in sorted(records, key=lambda row: row[1]["ts"]): + emit(worker, event) + + +def show_logs( + pod_id=None, + *, + endpoint_id=None, + worker_ids=(), + follow=False, + source=None, + tail=100, + since=None, + as_json=False, +): + try: + with api_errors(): + asyncio.run( + _show_logs( + pod_id, + endpoint_id, + worker_ids, + follow, + source, + tail, + since, + as_json, + ) + ) + except KeyboardInterrupt: + pass + except (aiohttp.ClientError, asyncio.TimeoutError) as exc: + raise click.ClickException(str(exc)) from exc diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index d8be1833..9001ff04 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -486,9 +486,6 @@ async def _session() -> int: raise click.ClickException(str(exc)) from exc -# ------------------------------------------------------------ completion - - @cli.command() @click.argument("shell", type=click.Choice(["bash", "zsh", "fish"])) def completion(shell): diff --git a/tests/test_apps/test_logs.py b/tests/test_apps/test_logs.py index a49510f4..02ac437d 100644 --- a/tests/test_apps/test_logs.py +++ b/tests/test_apps/test_logs.py @@ -77,6 +77,43 @@ async def __anext__(self): class TestStreamPodLogs: + @pytest.mark.parametrize("endpoint_id", [None, "ep1"]) + async def test_combined_rest_stream_resumes_delivered_events( + self, monkeypatch, endpoint_id + ): + monkeypatch.setenv("RUNPOD_API_KEY", "test-key") + response = MagicMock() + response.content = _FakeContent( + [ + b"id: cursor2\n", + b'data: {"source":"system","line":"ready","ts":"t2"}\n', + ] + ) + session_cm, session = _mock_session(response) + cursors = [] + with patch("aiohttp.ClientSession", return_value=session_cm): + entries = [ + entry + async for entry in stream_pod_logs( + "pod1", + endpoint_id=endpoint_id, + log_type=None, + last_event_id="cursor1", + on_event_id=cursors.append, + ) + ] + url = session.get.call_args.args[0] + suffix = ( + "/v2/serverless/ep1/workers/pod1/logs" + if endpoint_id + else "/v2/pods/pod1/logs" + ) + assert url.endswith(suffix) + assert session.get.call_args.kwargs["params"] == {"tail": "100"} + assert session.get.call_args.kwargs["headers"]["Last-Event-ID"] == "cursor1" + assert cursors == ["cursor2"] + assert entries == [{"source": "system", "line": "ready", "ts": "t2"}] + async def test_parses_sse_frames(self, monkeypatch): monkeypatch.setenv("RUNPOD_API_KEY", "test-key") frames = [ diff --git a/tests/test_cli/test_catalog_command.py b/tests/test_cli/test_catalog_command.py index 9c27ca3f..daccfa0d 100644 --- a/tests/test_cli/test_catalog_command.py +++ b/tests/test_cli/test_catalog_command.py @@ -104,6 +104,16 @@ def _invoke(args, gpus=GPUS, cpus=CPUS, datacenters=DATACENTERS): class TestGpuCatalog: + def test_unavailable_gpu_omits_stock_fields(self): + gpu = { + key: value + for key, value in GPUS[0].items() + if key not in ("dataCenters", "availability") + } + result, _, _ = _invoke(["--region", "eu", "--available", "--json"], gpus=[gpu]) + assert result.exit_code == 0, result.output + assert json.loads(result.output) == [] + def test_pod_table_sorted_by_stock_then_price(self): result, list_gpus, _ = _invoke(["h100", "--count", "2"]) assert result.exit_code == 0, result.output diff --git a/tests/test_cli/test_cli_logs.py b/tests/test_cli/test_cli_logs.py new file mode 100644 index 00000000..c747f7a8 --- /dev/null +++ b/tests/test_cli/test_cli_logs.py @@ -0,0 +1,166 @@ +import asyncio +import json +from types import SimpleNamespace + +import aiohttp +import click +import pytest + +from runpod.rp_cli import logs + + +def event(line="hello", ts="2026-01-01T00:00:00Z"): + return {"source": "container", "line": line, "ts": ts} + + +@pytest.mark.parametrize("as_json", [False, True]) +def test_pod_snapshot_uses_sdk_reader(monkeypatch, capsys, as_json): + calls = [] + + async def stream(pod_id, **kwargs): + calls.append((pod_id, kwargs)) + yield event() + + monkeypatch.setattr(logs, "stream_pod_logs", stream) + logs.show_logs("pod1", tail=7, since="2025-01-01T00:00:00Z", as_json=as_json) + text = capsys.readouterr().out + assert json.loads(text) == event() if as_json else text == "[container] hello\n" + assert calls[0][0] == "pod1" + assert calls[0][1]["log_type"] is None + assert calls[0][1]["tail"] == 7 + assert calls[0][1]["since"] == "2025-01-01T00:00:00Z" + + +def test_worker_snapshot_is_sorted_and_filterable(monkeypatch, capsys): + monkeypatch.setattr( + logs, + "run_rest_request", + lambda *a, **k: {"workers": [{"id": "a"}, {"id": "b"}, {"id": "ignored"}]}, + ) + + async def stream(worker, **kwargs): + assert kwargs["endpoint_id"] == "ep1" + assert kwargs["log_type"] == "system" + yield event( + worker, "2026-01-01T00:00:02Z" if worker == "a" else "2026-01-01T00:00:01Z" + ) + + monkeypatch.setattr(logs, "stream_pod_logs", stream) + logs.show_logs( + endpoint_id="ep1", worker_ids=("a", "b"), source="system", as_json=True + ) + rows = [json.loads(line) for line in capsys.readouterr().out.splitlines()] + assert [row["worker"] for row in rows] == ["b", "a"] + + +@pytest.mark.parametrize("empty", [False, True]) +def test_snapshot_closes_idle_or_live_stream(monkeypatch, capsys, empty): + closed = [] + monkeypatch.setattr(logs, "SNAPSHOT_IDLE_SECONDS", 0.01) + + async def stream(*args, **kwargs): + try: + if not empty: + yield event("backfill") + yield event("live", "2999-01-01T00:00:00Z") + await asyncio.sleep(10) + finally: + closed.append(True) + + monkeypatch.setattr(logs, "stream_pod_logs", stream) + logs.show_logs("pod1") + assert closed == [True] + assert capsys.readouterr().out == ("" if empty else "[container] backfill\n") + + +def test_empty_workers(monkeypatch, capsys): + monkeypatch.setattr(logs, "run_rest_request", lambda *a, **k: {"workers": []}) + logs.show_logs(endpoint_id="ep1") + assert "no workers" in capsys.readouterr().out + + +@pytest.mark.parametrize( + "endpoint_id,status,raises", + [(None, 404, True), ("ep1", 404, False), ("ep1", 403, True)], +) +def test_stream_errors_do_not_hide_auth_failures( + monkeypatch, endpoint_id, status, raises +): + monkeypatch.setattr( + logs, "run_rest_request", lambda *a, **k: {"workers": [{"id": "pod1"}]} + ) + + async def stream(*args, **kwargs): + raise aiohttp.ClientResponseError( + SimpleNamespace(real_url="https://api.runpod.io/v2/pods/pod1/logs"), + (), + status=status, + message="denied", + ) + yield + + monkeypatch.setattr(logs, "stream_pod_logs", stream) + if raises: + with pytest.raises(click.ClickException): + logs.show_logs(None if endpoint_id else "pod1", endpoint_id=endpoint_id) + else: + logs.show_logs(endpoint_id=endpoint_id) + + +async def test_follow_resumes_and_cancels_streams(monkeypatch, capsys): + resumes, closed = [], [] + received = asyncio.Event() + monkeypatch.setattr(logs, "RECONNECT_SECONDS", 0) + monkeypatch.setattr(logs, "WORKER_POLL_SECONDS", 0.01) + + async def stream(worker, **kwargs): + resumes.append(kwargs["last_event_id"]) + try: + if len(resumes) == 1: + kwargs["on_event_id"]("cursor1") + yield event("first") + else: + yield event("second") + received.set() + await asyncio.sleep(10) + finally: + closed.append(True) + + monkeypatch.setattr(logs, "stream_pod_logs", stream) + task = asyncio.create_task( + logs._show_logs("pod1", None, (), True, None, 10, None, False) + ) + await asyncio.wait_for(received.wait(), 1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert resumes == [None, "cursor1"] + assert len(closed) == 2 + assert "[container] first\n[container] second\n" == capsys.readouterr().out + + +async def test_follow_attaches_new_workers(monkeypatch, capsys): + polls = 0 + received = asyncio.Event() + monkeypatch.setattr(logs, "WORKER_POLL_SECONDS", 0.01) + + def workers(*args): + nonlocal polls + polls += 1 + return {"workers": [] if polls == 1 else [{"id": "new"}]} + + async def stream(worker, **kwargs): + yield event(worker) + received.set() + await asyncio.sleep(10) + + monkeypatch.setattr(logs, "run_rest_request", workers) + monkeypatch.setattr(logs, "stream_pod_logs", stream) + task = asyncio.create_task( + logs._show_logs(None, "ep1", (), True, None, 10, None, False) + ) + await asyncio.wait_for(received.wait(), 1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert "new [container] new" in capsys.readouterr().out From c4489ab2bfbcac79d41ea3c0535723bdf6896d13 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:32:20 -0700 Subject: [PATCH 17/25] refactor: simplify pod commands and share log streaming --- runpod/rp_cli/pod.py | 681 +++++++++++++---------------- tests/test_cli/test_pod_command.py | 259 ++++++----- 2 files changed, 440 insertions(+), 500 deletions(-) diff --git a/runpod/rp_cli/pod.py b/runpod/rp_cli/pod.py index 3cfd74ec..c4e2455b 100644 --- a/runpod/rp_cli/pod.py +++ b/runpod/rp_cli/pod.py @@ -1,7 +1,6 @@ -"""`rp pod`: create, inspect, and manage pods.""" - import shlex import time +from operator import itemgetter from typing import Optional import click @@ -9,8 +8,9 @@ from runpod.api.rest import run_rest_request from runpod.cli.groups.pod.commands import sync_pods from runpod.rp_cli import compute +from runpod.rp_cli.api import api_errors, list_all +from runpod.rp_cli.options import parse_env from runpod.rp_cli.output import ( - api_errors, completer, confirm, interactive, @@ -27,25 +27,15 @@ WAIT_INTERVAL = 5 -def _list_pods(include_cluster_pods: bool = False) -> list[dict]: - params = {"includeClusterPods": "true"} if include_cluster_pods else None - return compute.list_all("/v2/pods", "pods", params) - - -def _get_pod(pod_id: str) -> dict: - return run_rest_request("GET", f"/v2/pods/{pod_id}") or {} - - def _pod_choices(action: Optional[str] = None) -> list[tuple[str, str]]: return [ - (pod["id"], f"{pod.get('name')} {(pod.get('status') or '').lower()}") - for pod in _list_pods() - if action is None or action in (pod.get("actions") or []) + (pod["id"], f"{pod['name']} {pod['status'].lower()}") + for pod in list_all("/v2/pods", "pods") + if action is None or action in pod["actions"] ] def _resolve_pod_id(pod_id: Optional[str], action: Optional[str] = None) -> str: - """the given pod id, or one picked interactively.""" if pod_id: return pod_id choices = [] @@ -65,19 +55,18 @@ def _resolve_pod_id(pod_id: Optional[str], action: Optional[str] = None) -> str: _complete_network_volumes = completer( lambda: [ (volume["id"], f"{volume.get('name')} {volume.get('dataCenter')}") - for volume in compute.list_all("/v2/network-volumes", "networkVolumes") + for volume in list_all("/v2/network-volumes", "networkVolumes") ] ) _complete_registries = completer( lambda: [ (registry["name"], registry["id"]) - for registry in compute.list_all("/v2/registries", "registries") + for registry in list_all("/v2/registries", "registries") ] ) def _say(message: str, as_json: bool, warning: bool = False) -> None: - """a status line on stdout, or stderr when stdout carries json.""" if as_json: click.echo(message, err=True) return @@ -88,146 +77,101 @@ def _say(message: str, as_json: bool, warning: bool = False) -> None: (ui.warn if warning else ui.info)(escape(message)) -def _hardware(pod: dict) -> Optional[str]: - gpu = pod.get("gpu") - if gpu: - return f"{gpu.get('count') or 1}x {gpu['id']}" - cpu = pod.get("cpu") - if cpu: - return f"{cpu['id']} {cpu.get('vcpuCount')} vcpu" - return None - - -def _cost(pod: dict) -> Optional[str]: - return f"${pod['cost']:.2f}/hr" if pod.get("cost") else None - - -def _mounts(pod: dict) -> list[str]: - mounts = pod.get("mounts") or {} - described = [] - persistent = mounts.get("persistent") - if persistent: - described.append(f"persistent {persistent['size']}GB at {persistent['path']}") - for mount in mounts.get("network") or []: - described.append(f"network {mount['volumeId']} at {mount.get('path')}") - for mount in mounts.get("global") or []: - described.append(f"global {mount['volumeId']} at {mount.get('path') or '-'}") - return described - - -def _urls(pod: dict) -> list[str]: - urls = [] - for port in pod.get("ports") or []: - number, _, protocol = port.partition("/") - if protocol == "http": - urls.append(f"https://{pod['id']}-{number}.proxy.runpod.net") - return urls - - -def _tcp(pod: dict) -> list[str]: - ports = (pod.get("runtime") or {}).get("ports") or [] - return [ - f"{port['ip']}:{port['public']} -> {port['private']}" - for port in ports - if port.get("type") == "tcp" and port.get("ip") and port.get("public") - ] - - -def _global_network(pod: dict) -> object: - networking = pod.get("globalNetworking") or {} - if not networking.get("enabled"): - return False - return networking.get("internalDns") or networking.get("ip") or True - - def _status(pod: dict) -> str: - status = (pod.get("status") or "").lower() + status = pod["status"].lower() if status == "running" and not (pod.get("runtime") or {}).get("uptime"): return "starting" return status -def _uptime(pod: dict) -> Optional[str]: - seconds = (pod.get("runtime") or {}).get("uptime") - if not seconds: - return None - minutes, seconds = divmod(seconds, 60) - hours, minutes = divmod(minutes, 60) - if hours: - return f"{hours}h {minutes}m" - if minutes: - return f"{minutes}m {seconds}s" - return f"{seconds}s" - - -def _resources(pod: dict) -> Optional[str]: +def _summary(pod: dict) -> dict: hardware = pod.get("gpu") or pod.get("cpu") or {} - parts = [] - if hardware.get("vcpuCount"): - parts.append(f"{hardware['vcpuCount']:g} vcpu") - if hardware.get("memory"): - parts.append(f"{hardware['memory']:g}GB ram") - return ", ".join(parts) or None - - -_LIST_COLUMNS = [ - ("ID", lambda pod: pod["id"]), - ("NAME", lambda pod: pod.get("name")), - ("STATUS", _status), - ("HARDWARE", _hardware), - ("DATA CENTER", lambda pod: pod.get("dataCenterId")), - ("COST", _cost), - ("IMAGE", lambda pod: pod.get("image")), -] - -_FIELDS = [ - ("id", lambda pod: pod["id"]), - ("name", lambda pod: pod.get("name")), - ("status", _status), - ("image", lambda pod: pod.get("image")), - ("hardware", _hardware), - ("resources", _resources), - ("cloud", lambda pod: (pod.get("cloud") or "").lower()), - ("data center", lambda pod: pod.get("dataCenterId")), - ("cuda", lambda pod: pod.get("cudaVersion")), - ("cost", _cost), - ("disk", lambda pod: f"{pod['disk']}GB" if pod.get("disk") else None), - ("volumes", _mounts), - ("ports", lambda pod: pod.get("ports")), - ("env", lambda pod: sorted(pod.get("env") or {})), - ("command", lambda pod: pod.get("args")), - ("locked", lambda pod: pod.get("locked")), - ("global network", _global_network), - ("cluster", lambda pod: (pod.get("cluster") or {}).get("id")), - ("created", lambda pod: pod.get("createdAt")), - ("started", lambda pod: pod.get("startedAt")), - ("uptime", _uptime), - ("actions", lambda pod: pod.get("actions")), - ("ssh", lambda pod: ((pod.get("ssh") or {}).get("proxy") or {}).get("command")), - ( - "ssh direct", - lambda pod: ((pod.get("ssh") or {}).get("direct") or {}).get("command"), - ), - ("urls", _urls), - ("tcp", _tcp), -] + description = None + if pod.get("gpu"): + description = f"{hardware['count']}x {hardware['id']}" + elif pod.get("cpu"): + description = f"{hardware['id']} {hardware['vcpuCount']} vcpu" + return { + "id": pod["id"], + "name": pod["name"], + "status": _status(pod), + "hardware": description, + "data center": pod.get("dataCenterId"), + "cost": f"${pod['cost']:.2f}/hr" if pod.get("cost") else None, + "image": pod["image"], + } def _show(pod: dict, as_json: bool) -> None: if as_json: print_json(pod) return - print_fields(pod, _FIELDS) + runtime = pod.get("runtime") or {} + mounts = pod.get("mounts") or {} + networking = pod.get("globalNetworking") or {} + ssh = pod.get("ssh") or {} + hardware = pod.get("gpu") or pod.get("cpu") or {} + volumes = [] + if mounts.get("persistent"): + volume = mounts["persistent"] + volumes.append(f"persistent {volume['size']}GB at {volume['path']}") + for kind in ("network", "global"): + for volume in mounts.get(kind, []): + volumes.append( + f"{kind} {volume['volumeId']} at {volume.get('path') or '-'}" + ) + uptime = None + if runtime.get("uptime"): + minutes, seconds = divmod(runtime["uptime"], 60) + hours, minutes = divmod(minutes, 60) + uptime = f"{hours}h {minutes}m" if hours else f"{minutes}m {seconds}s" + if not hours and not minutes: + uptime = f"{seconds}s" + fields = { + **_summary(pod), + "resources": [ + f"{hardware[key]:g}{unit}" + for key, unit in (("vcpuCount", " vcpu"), ("memory", "GB ram")) + if hardware.get(key) + ], + "cloud": (pod.get("cloud") or "").lower(), + "cuda": pod.get("cudaVersion"), + "disk": f"{pod['disk']}GB" if pod.get("disk") else None, + "volumes": volumes, + "ports": pod.get("ports"), + "env": sorted(pod.get("env") or {}), + "command": pod.get("args"), + "locked": pod.get("locked"), + "global network": bool(networking.get("enabled")) + and (networking.get("internalDns") or networking.get("ip") or True), + "cluster": (pod.get("cluster") or {}).get("id"), + "created": pod.get("createdAt"), + "started": pod.get("startedAt"), + "uptime": uptime, + "actions": pod["actions"], + "ssh": (ssh.get("proxy") or {}).get("command"), + "ssh direct": (ssh.get("direct") or {}).get("command"), + "urls": [ + f"https://{pod['id']}-{port.split('/')[0]}.proxy.runpod.net" + for port in pod.get("ports", []) + if port.endswith("/http") + ], + "tcp": [ + f"{port['ip']}:{port['public']} -> {port['private']}" + for port in runtime.get("ports", []) + if port["type"] == "tcp" and port.get("ip") and port.get("public") + ], + } + print_fields(fields, [(label, itemgetter(label)) for label in fields]) def _wait_running(pod_id: str, since: float) -> dict: - """poll a pod until a container started after `since` is up.""" from runpod.rp_cli import console as ui deadline = time.monotonic() + WAIT_TIMEOUT with ui.console.status(f"[dim]waiting for {pod_id} to start ...[/dim]"): while True: - pod = _get_pod(pod_id) + pod = run_rest_request("GET", f"/v2/pods/{pod_id}") status = _status(pod) uptime = pod["runtime"]["uptime"] if status == "running" else 0 # a longer uptime is a stale report from the container before the action @@ -242,93 +186,73 @@ def _wait_running(pod_id: str, since: float) -> dict: time.sleep(WAIT_INTERVAL) -def _parse_env(values: tuple[str, ...]) -> dict[str, str]: - env = {} - for value in values: - key, sep, setting = value.partition("=") - if not sep or not key: - raise click.BadParameter(f"'{value}' is not KEY=VALUE", param_hint="--env") - env[key] = setting - return env - - -def _split_mount(value: str, flag: str) -> tuple[str, Optional[str]]: - head, sep, path = value.partition(":") - if not head or (sep and not path): - raise click.BadParameter( - f"'{value}' should look like VALUE[:PATH]", param_hint=flag +def _container_body( + options: dict, current: Optional[dict] = None, unset_env=() +) -> dict: + """build container settings, preserving existing environment and mount paths.""" + current = current or {} + body = { + field: options[option] + for option, field in ( + ("image", "image"), + ("template", "templateId"), + ("args", "args"), + ("disk", "disk"), + ("global_networking", "globalNetworking"), ) - return head, path or None - - -def _mount_path(current: dict, kind: str) -> Optional[str]: - entries = current.get(kind) or [] - return entries[0].get("path") if entries else None - - -def _mounts_body(options: dict, current: dict) -> dict: - """the mounts object for the volume flags, defaulting paths from current.""" - mounts: dict = {} - if options["persistent_volume"]: - size, path = _split_mount(options["persistent_volume"], "--persistent-volume") - if not size.isdigit(): - raise click.BadParameter( - "size is in GB, like 50 or 50:/workspace", - param_hint="--persistent-volume", - ) - mounts["persistent"] = { - "size": int(size), - "path": path - or (current.get("persistent") or {}).get("path") - or DEFAULT_MOUNT_PATH, - } - if options["network_volume"]: - volume_id, path = _split_mount(options["network_volume"], "--network-volume") - mounts["network"] = [ - { - "volumeId": volume_id, - "path": path or _mount_path(current, "network") or DEFAULT_MOUNT_PATH, - } - ] - if options["global_volume"]: - volume_id, path = _split_mount(options["global_volume"], "--global-volume") - path = path or _mount_path(current, "global") - mounts["global"] = [{"volumeId": volume_id, **({"path": path} if path else {})}] - return mounts - - -def _resolve_registry(value: str) -> Optional[str]: - if not value: - return None - for registry in compute.list_all("/v2/registries", "registries"): - if value in (registry["id"], registry.get("name")): - return registry["id"] - raise compute.ResolutionError( - f"unknown registry credential '{value}'. list them with rp registry list" - ) - - -def _container_body(options: dict) -> dict: - """request fields shared by create and update, except env and mounts.""" - body: dict = {} - if options["image"]: - body["image"] = options["image"] - if options["template"]: - body["templateId"] = options["template"] - if options["args"] is not None: - body["args"] = options["args"] - if options["cmd"] is not None: - body["cmd"] = shlex.split(options["cmd"]) - if options["entrypoint"] is not None: - body["entrypoint"] = shlex.split(options["entrypoint"]) - if options["disk"]: - body["disk"] = options["disk"] + if options[option] is not None + } + for option in ("cmd", "entrypoint"): + if options[option] is not None: + body[option] = shlex.split(options[option]) if options["ports"]: body["ports"] = list(options["ports"]) if options["registry"] is not None: - body["registry"] = _resolve_registry(options["registry"]) - if options["global_networking"] is not None: - body["globalNetworking"] = options["global_networking"] + value = options["registry"] + body["registry"] = None + if value: + matches = [ + registry["id"] + for registry in list_all("/v2/registries", "registries") + if value in (registry["id"], registry["name"]) + ] + if not matches: + raise compute.ResolutionError( + f"unknown registry credential '{value}'. list them with rp registry list" + ) + body["registry"] = matches[0] + if options["env"] or unset_env: + env = {**current.get("env", {}), **parse_env(options["env"])} + for key in unset_env: + env.pop(key, None) + body["env"] = env + mounts = {} + for kind in ("persistent", "network", "global"): + value = options[f"{kind}_volume"] + if value is None: + continue + head, sep, path = value.partition(":") + flag = f"--{kind}-volume" + if not head or (sep and not path): + raise click.BadParameter( + f"'{value}' should look like VALUE[:PATH]", param_hint=flag + ) + existing = current.get("mounts", {}).get(kind) + if kind != "persistent": + existing = existing[0] if existing else None + path = path or (existing or {}).get("path") + if kind != "global": + path = path or DEFAULT_MOUNT_PATH + if kind == "persistent": + if not head.isdigit(): + raise click.BadParameter( + "size is in GB, like 50 or 50:/workspace", param_hint=flag + ) + mounts[kind] = {"size": int(head), "path": path} + else: + mounts[kind] = [{"volumeId": head, **({"path": path} if path else {})}] + if mounts: + body["mounts"] = mounts return body @@ -413,19 +337,22 @@ def _container_options(command): return command -def _choose_gpu( - query: str, - count: int, +def _gpu_body( + spec: str, + count: Optional[int], cloud: str, - cuda: tuple[Optional[str], list[str]], + cuda: tuple[str, ...], datacenter_ids: list[str], -) -> tuple[dict, str, int]: - """the gpu type to request, its stock level, and how many types matched. - - a query naming several types picks the one with the most stock in the - requested data centers, then the lowest price. - """ - min_cuda, exact_cuda = cuda + min_ram: Optional[int], + min_vcpu: Optional[int], + as_json: bool, +) -> dict: + """resolve a GPU specification by stock and price, then build its request.""" + query, spec_count = compute.parse_gpu_spec(spec) + if spec_count and count and spec_count != count: + raise click.UsageError(f"--gpu {spec} conflicts with --gpu-count {count}") + count = spec_count or count or 1 + min_cuda, exact_cuda = compute.parse_cuda(cuda) gpus = compute.list_gpus( product="POD", count=count, @@ -467,7 +394,33 @@ def stock(gpu: dict) -> str: gpu.get("name") or gpu["id"], ), ) - return best, stock(best), len(matches) + availability = stock(best) + label = best.get("name") or best["id"] + price = (best.get("price") or {}).get(cloud) + if len(matches) > 1: + cost = f", ${price * count:.2f}/hr" if price else "" + _say( + f"'{query}' matches {len(matches)} gpu types; using {label} " + f"({best['id']}): {'no' if availability == 'NONE' else availability.lower()} stock{cost}", + as_json, + ) + if availability == "NONE": + _say( + f"the catalog reports no {cloud} stock for {count}x {label} " + "in the requested data centers; creation may fail", + as_json, + warning=True, + ) + body = {"id": best["id"], "count": count} + for field, value in ( + ("minCudaVersion", min_cuda), + ("allowedCudaVersions", exact_cuda), + ("minRamPerGpu", min_ram), + ("minVcpuCountPerGpu", min_vcpu), + ): + if value: + body[field] = value + return body def _datacenter_ids(regions, data_centers, network_volume) -> list[str]: @@ -477,12 +430,88 @@ def _datacenter_ids(regions, data_centers, network_volume) -> list[str]: ids += [dc for dc in data_centers if dc not in ids] if not ids and network_volume: volume_id = network_volume.partition(":")[0] - volume = run_rest_request("GET", f"/v2/network-volumes/{volume_id}") or {} - if volume.get("dataCenter"): - ids = [volume["dataCenter"]] + volume = run_rest_request("GET", f"/v2/network-volumes/{volume_id}") + ids = [volume["dataCenter"]] return ids +def _create_body( + name, + gpu, + gpu_count, + cpu, + vcpu, + cloud, + regions, + data_centers, + cuda, + min_ram_per_gpu, + min_vcpu_per_gpu, + ssh, + jupyter, + as_json, + **container, +): + if gpu and cpu: + raise click.UsageError("use either --gpu or --cpu") + if not gpu and not cpu: + raise click.UsageError("pass --gpu (like 4090 or h100:2) or --cpu (like cpu3c)") + if cpu and (gpu_count or cuda or min_ram_per_gpu or min_vcpu_per_gpu): + raise click.UsageError( + "--gpu-count, --cuda, --min-ram-per-gpu, and --min-vcpu-per-gpu apply to gpu pods" + ) + if vcpu and not cpu: + raise click.UsageError("--vcpu applies to cpu pods") + if not container["image"] and not container["template"]: + raise click.UsageError("pass --image or --template") + if not name: + if not interactive(): + raise click.UsageError("pass a NAME") + name = click.prompt("pod name") + + body = {"name": name, **_container_body(container)} + datacenter_ids = _datacenter_ids(regions, data_centers, container["network_volume"]) + if datacenter_ids: + body["dataCenterIds"] = datacenter_ids + if cloud: + body["cloud"] = cloud.upper() + + if gpu: + body["gpu"] = _gpu_body( + gpu, + gpu_count, + (cloud or "secure").lower(), + cuda, + datacenter_ids, + min_ram_per_gpu, + min_vcpu_per_gpu, + as_json, + ) + else: + flavor = compute.resolve_cpu( + cpu, compute.list_cpus(product="POD", vcpu_count=vcpu) + ) + body["cpu"] = { + "id": flavor["id"], + "vcpuCount": vcpu or (flavor.get("vcpu") or {}).get("min") or 2, + } + + if ssh: + body["startSsh"] = True + keys = run_rest_request("GET", "/v2/account/ssh-keys")["keys"] + if not keys: + _say( + "your account has no ssh keys, so --ssh has no effect. " + "add one with rp ssh add", + as_json, + warning=True, + ) + if jupyter: + body["startJupyter"] = True + + return body + + @click.group("pod") def pod(): """Create and manage pods.""" @@ -504,7 +533,8 @@ def pod(): def list_pods(name, statuses, include_cluster_pods, as_json): """List your pods.""" with api_errors(): - pods = _list_pods(include_cluster_pods) + params = {"includeClusterPods": "true"} if include_cluster_pods else None + pods = list_all("/v2/pods", "pods", params) if name: pods = [p for p in pods if name.lower() in (p.get("name") or "").lower()] if statuses: @@ -513,7 +543,12 @@ def list_pods(name, statuses, include_cluster_pods, as_json): if as_json: print_json(pods) return - print_table(pods, _LIST_COLUMNS, "no pods. create one with rp pod create") + columns = ["id", "name", "status", "hardware", "data center", "cost", "image"] + print_table( + [_summary(p) for p in pods], + [(label.upper(), itemgetter(label)) for label in columns], + "no pods. create one with rp pod create", + ) @pod.command() @@ -523,7 +558,7 @@ def get(pod_id, as_json): """Show a pod with its connection details.""" pod_id = _resolve_pod_id(pod_id) with api_errors(): - _show(_get_pod(pod_id), as_json) + _show(run_rest_request("GET", f"/v2/pods/{pod_id}"), as_json) @pod.command() @@ -586,24 +621,7 @@ def get(pod_id, as_json): ) @click.option("--wait", is_flag=True, help="Wait until the pod is running.") @json_option -def create( - name, - gpu, - gpu_count, - cpu, - vcpu, - cloud, - regions, - data_centers, - cuda, - min_ram_per_gpu, - min_vcpu_per_gpu, - ssh, - jupyter, - wait, - as_json, - **container, -): # pylint: disable=too-many-arguments,too-many-locals +def create(name, wait, as_json, **options): """Create a pod named NAME. GPU types are human-readable (4090, h100, a100-sxm). When a name @@ -615,102 +633,11 @@ def create( rp pod create box --cpu cpu3c --vcpu 4 --image ubuntu:24.04 rp pod create nb --gpu 4090 --image runpod/pytorch:latest --jupyter -p 8888/http """ - if gpu and cpu: - raise click.UsageError("use either --gpu or --cpu") - if not gpu and not cpu: - raise click.UsageError("pass --gpu (like 4090 or h100:2) or --cpu (like cpu3c)") - if cpu and (gpu_count or cuda or min_ram_per_gpu or min_vcpu_per_gpu): - raise click.UsageError( - "--gpu-count, --cuda, --min-ram-per-gpu, and --min-vcpu-per-gpu apply to gpu pods" - ) - if vcpu and not cpu: - raise click.UsageError("--vcpu applies to cpu pods") - if not container["image"] and not container["template"]: - raise click.UsageError("pass --image or --template") - if not name: - if not interactive(): - raise click.UsageError("pass a NAME") - name = click.prompt("pod name") - with api_errors(): - body = {"name": name, **_container_body(container)} - env = _parse_env(container["env"]) - if env: - body["env"] = env - mounts = _mounts_body(container, {}) - if mounts: - body["mounts"] = mounts - datacenter_ids = _datacenter_ids( - regions, data_centers, container["network_volume"] - ) - if datacenter_ids: - body["dataCenterIds"] = datacenter_ids - if cloud: - body["cloud"] = cloud.upper() - - if gpu: - query, spec_count = compute.parse_gpu_spec(gpu) - if spec_count and gpu_count and spec_count != gpu_count: - raise click.UsageError( - f"--gpu {gpu} conflicts with --gpu-count {gpu_count}" - ) - count = spec_count or gpu_count or 1 - min_cuda, exact_cuda = compute.parse_cuda(cuda) - tier = (cloud or "secure").lower() - chosen, stock, matched = _choose_gpu( - query, count, tier, (min_cuda, exact_cuda), datacenter_ids - ) - label = chosen.get("name") or chosen["id"] - price = (chosen.get("price") or {}).get(tier) - if matched > 1: - cost = f", ${price * count:.2f}/hr" if price else "" - _say( - f"'{query}' matches {matched} gpu types; using {label} " - f"({chosen['id']}): {'no' if stock == 'NONE' else stock.lower()} stock{cost}", - as_json, - ) - if stock == "NONE": - _say( - f"the catalog reports no {tier} stock for {count}x {label} " - "in the requested data centers; creation may fail", - as_json, - warning=True, - ) - gpu_body = {"id": chosen["id"], "count": count} - if min_cuda: - gpu_body["minCudaVersion"] = min_cuda - if exact_cuda: - gpu_body["allowedCudaVersions"] = exact_cuda - if min_ram_per_gpu: - gpu_body["minRamPerGpu"] = min_ram_per_gpu - if min_vcpu_per_gpu: - gpu_body["minVcpuCountPerGpu"] = min_vcpu_per_gpu - body["gpu"] = gpu_body - else: - flavor = compute.resolve_cpu( - cpu, compute.list_cpus(product="POD", vcpu_count=vcpu) - ) - body["cpu"] = { - "id": flavor["id"], - "vcpuCount": vcpu or (flavor.get("vcpu") or {}).get("min") or 2, - } - - if ssh: - body["startSsh"] = True - keys = (run_rest_request("GET", "/v2/account/ssh-keys") or {}).get("keys") - if not keys: - _say( - "your account has no ssh keys, so --ssh has no effect. " - "add one with rp ssh add", - as_json, - warning=True, - ) - if jupyter: - body["startJupyter"] = True - + body = _create_body(name, as_json=as_json, **options) since = time.monotonic() - created = run_rest_request("POST", "/v2/pods", json=body) or {} - _say(f"created pod {created.get('name')} ({created.get('id')})", as_json) + created = run_rest_request("POST", "/v2/pods", json=body) + _say(f"created pod {created['name']} ({created['id']})", as_json) if wait: created = _wait_running(created["id"], since) _show(created, as_json) @@ -747,25 +674,23 @@ def update(pod_id, name, unset_env, locked, as_json, **container): pod_id = _resolve_pod_id(pod_id) with api_errors(): - body = _container_body(container) + needs_current = ( + container["env"] + or unset_env + or any( + container[key] + for key in ("persistent_volume", "network_volume", "global_volume") + ) + ) + current = ( + run_rest_request("GET", f"/v2/pods/{pod_id}") if needs_current else None + ) + body = _container_body(container, current, unset_env) if name: body["name"] = name if locked is not None: body["locked"] = locked - mounts_changed = any( - container[key] - for key in ("persistent_volume", "network_volume", "global_volume") - ) - if container["env"] or unset_env or mounts_changed: - current = _get_pod(pod_id) - if container["env"] or unset_env: - env = {**(current.get("env") or {}), **_parse_env(container["env"])} - for key in unset_env: - env.pop(key, None) - body["env"] = env - if mounts_changed: - body["mounts"] = _mounts_body(container, current.get("mounts") or {}) - updated = run_rest_request("PATCH", f"/v2/pods/{pod_id}", json=body) or {} + updated = run_rest_request("PATCH", f"/v2/pods/{pod_id}", json=body) _say(f"updated pod {pod_id}", as_json) _show(updated, as_json) @@ -782,7 +707,7 @@ def _act( _say(f"{done} pod {pod_id}", as_json) if wait: result = _wait_running(pod_id, since) - _show(result or {}, as_json) + _show(result, as_json) @pod.command() @@ -828,6 +753,28 @@ def terminate(pod_id, yes): pod.add_command(terminate, "delete") +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@click.option("--follow", "-f", is_flag=True, help="Follow new log lines.") +@click.option("--tail", type=click.IntRange(0, 5000), default=100, show_default=True) +@click.option("--source", type=click.Choice(["container", "system"])) +@click.option("--since", help="Only logs after this timestamp.") +@json_option +def logs(pod_id, follow, tail, source, since, as_json): + """Read a pod's container or system logs.""" + from runpod.rp_cli.logs import show_logs + + pod_id = _resolve_pod_id(pod_id) + show_logs( + pod_id, + follow=follow, + tail=tail, + source=source, + since=since, + as_json=as_json, + ) + + @pod.command() @click.argument("pod_id", required=False, shell_complete=_complete_pods) def connect(pod_id): diff --git a/tests/test_cli/test_pod_command.py b/tests/test_cli/test_pod_command.py index 25678127..81eed4bb 100644 --- a/tests/test_cli/test_pod_command.py +++ b/tests/test_cli/test_pod_command.py @@ -1,6 +1,5 @@ -"""tests for rp pod.""" - import json +import shlex from unittest.mock import MagicMock, patch import pytest @@ -85,7 +84,7 @@ "args": "", "disk": 50, "ports": ["8888/http", "22/tcp"], - "env": {"HF_TOKEN": "secret", "MODEL": "llama"}, + "env": {"PRIVATE_SETTING": "redacted-test-value", "MODEL": "llama"}, "mounts": {"network": [{"volumeId": "vol1", "path": "/data"}]}, "gpu": {"id": "NVIDIA H100 80GB HBM3", "count": 2, "vcpuCount": 32, "memory": 250}, "cloud": "SECURE", @@ -121,19 +120,24 @@ class FakeApi: - """routes rest calls by method and path, recording each one.""" - def __init__(self, routes=None): self.calls = [] self.routes = { - ("GET", "/v2/pods"): {"pods": [POD, STOPPED]}, + ("GET", "/v2/pods"): { + "pods": [POD, STOPPED], + "pagination": {"hasNextPage": False, "nextCursor": None}, + }, ("GET", "/v2/pods/pod123"): POD, - ("GET", "/v2/account/ssh-keys"): {"keys": ["ssh-ed25519 AAAA me"]}, - ("GET", "/v2/registries"): {"registries": [{"id": "reg1", "name": "ghcr"}]}, + ("GET", "/v2/account/ssh-keys"): {"keys": ["redacted-test-key"]}, + ("GET", "/v2/registries"): { + "registries": [{"id": "reg1", "name": "ghcr"}], + "pagination": {"hasNextPage": False, "nextCursor": None}, + }, ("GET", "/v2/network-volumes"): { "networkVolumes": [ {"id": "vol1", "name": "data", "dataCenter": "EU-RO-1"} - ] + ], + "pagination": {"hasNextPage": False, "nextCursor": None}, }, ("GET", "/v2/network-volumes/vol1"): { "id": "vol1", @@ -164,22 +168,22 @@ def made(self, method, path): @pytest.fixture(autouse=True) -def wide_console(monkeypatch): +def cli_environment(monkeypatch): from runpod.rp_cli import console as ui monkeypatch.setattr(ui.console, "_width", 200) - - -@pytest.fixture(autouse=True) -def no_sleep(monkeypatch): monkeypatch.setattr(pod_module.time, "sleep", lambda seconds: None) + monkeypatch.setattr( + "requests.request", + MagicMock(side_effect=AssertionError("unexpected HTTP request")), + ) def _invoke(args, api=None, gpus=GPUS, interactive=False, input=None): api = api or FakeApi() with ( patch("runpod.rp_cli.pod.run_rest_request", api), - patch("runpod.rp_cli.compute.run_rest_request", api), + patch("runpod.rp_cli.api.run_rest_request", api), patch("runpod.rp_cli.compute.list_gpus", return_value=gpus) as list_gpus, patch("runpod.rp_cli.compute.list_cpus", return_value=CPUS) as list_cpus, patch("runpod.rp_cli.compute.list_datacenters", return_value=DATACENTERS), @@ -214,7 +218,14 @@ def test_filters(self): assert [p["id"] for p in json.loads(result.stdout)] == ["pod123"] def test_cluster_pods_and_empty(self): - api = FakeApi({("GET", "/v2/pods"): {"pods": []}}) + api = FakeApi( + { + ("GET", "/v2/pods"): { + "pods": [], + "pagination": {"hasNextPage": False, "nextCursor": None}, + } + } + ) result = _invoke(["list", "--include-cluster-pods"], api=api) assert "no pods" in result.output assert api.calls[0][2] == {"includeClusterPods": "true"} @@ -238,8 +249,8 @@ def test_fields_and_connection_info(self): assert "network vol1 at /data" in output assert "pod123.runpod.internal" in output assert "32 vcpu, 250GB ram" in output - assert "HF_TOKEN, MODEL" in output - assert "secret" not in output + assert "MODEL, PRIVATE_SETTING" in output + assert "redacted-test-value" not in output assert "1h 2m" in output @pytest.mark.parametrize( @@ -271,7 +282,7 @@ def test_other_shapes(self): def test_json(self): result = _invoke(["get", "pod123", "--json"]) - assert json.loads(result.stdout)["env"]["HF_TOKEN"] == "secret" + assert json.loads(result.stdout) == POD def test_missing_id_without_tty(self): result = _invoke(["get"]) @@ -321,18 +332,17 @@ def test_ticket_example(self): "min_cuda_version": None, } - def test_prefers_stock_then_price(self): - _, body = _create(["a", "--gpu", "h100", "--image", "x"]) - assert body["gpu"]["id"] == "NVIDIA H100 NVL" - _, body = _create(["a", "--gpu", "h100", "--gpu-count", "2", "--image", "x"]) - assert body["gpu"]["id"] == "NVIDIA H100 PCIe" - - def test_stock_in_requested_data_centers(self): - _, body = _create( - ["a", "--gpu", "h100", "--data-center", "EU-RO-1", "--image", "x"] - ) - assert body["gpu"]["id"] == "NVIDIA H100 80GB HBM3" - assert body["dataCenterIds"] == ["EU-RO-1"] + @pytest.mark.parametrize( + "options, expected", + [ + ([], "NVIDIA H100 NVL"), + (["--gpu-count", "2"], "NVIDIA H100 PCIe"), + (["--data-center", "EU-RO-1"], "NVIDIA H100 80GB HBM3"), + ], + ) + def test_selects_stock_then_price_within_placement(self, options, expected): + _, body = _create(["a", "--gpu", "h100", "--image", "x", *options]) + assert body["gpu"]["id"] == expected def test_unique_match_used_directly(self): result, body = _create(["a", "--gpu", "4090", "--image", "x", "--json"]) @@ -377,23 +387,6 @@ def test_exact_cuda(self): assert body["gpu"]["allowedCudaVersions"] == ["12.4", "12.8"] assert result.list_gpus.call_args.kwargs["cuda_versions"] == ["12.4", "12.8"] - def test_no_type_offers_count(self): - result = _invoke(["create", "a", "--gpu", "h100:9", "--image", "x"]) - assert result.exit_code == 1 - assert "offers 9 per pod on secure cloud" in result.output - - def test_count_conflict(self): - result = _invoke( - ["create", "a", "--gpu", "h100:2", "--gpu-count", "3", "--image", "x"] - ) - assert result.exit_code == 2 - assert "conflicts" in result.output - - def test_unknown_gpu(self): - result = _invoke(["create", "a", "--gpu", "b9000", "--image", "x"]) - assert result.exit_code == 1 - assert "unknown gpu 'b9000'" in result.output - def test_ssh_without_keys_warns(self): api = FakeApi({("GET", "/v2/account/ssh-keys"): {"keys": []}}) result, body = _create(["a", "--gpu", "4090", "--image", "x", "--ssh"], api=api) @@ -402,57 +395,31 @@ def test_ssh_without_keys_warns(self): class TestCreateCpu: - def test_default_vcpu(self): - result, body = _create(["box", "--cpu", "cpu5g", "--image", "ubuntu"]) - assert body["cpu"] == {"id": "cpu5g", "vcpuCount": 4} + @pytest.mark.parametrize( + "flavor, options, requested, expected", + [("cpu5g", [], None, 4), ("cpu3c", ["--vcpu", "8"], 8, 8)], + ) + def test_vcpu_resolution(self, flavor, options, requested, expected): + result, body = _create(["box", "--cpu", flavor, "--image", "ubuntu", *options]) + assert body["cpu"] == {"id": flavor, "vcpuCount": expected} assert result.list_cpus.call_args.kwargs == { "product": "POD", - "vcpu_count": None, + "vcpu_count": requested, } - def test_explicit_vcpu(self): - _, body = _create(["box", "--cpu", "cpu3c", "--vcpu", "8", "--image", "ubuntu"]) - assert body["cpu"] == {"id": "cpu3c", "vcpuCount": 8} - class TestCreateContainer: def test_every_container_field(self): _, body = _create( - [ - "a", - "--gpu", - "4090", - "--template", - "tpl1", - "--image", - "img", - "-e", - "A=1", - "--env", - "B=x=y", - "-p", - "8888/http", - "--port", - "22/tcp", - "--disk", - "40", - "--args", - "bash start.sh", - "--cmd", - "python -m train --lr '1e-4'", - "--entrypoint", - "/bin/sh -c", - "--registry", - "ghcr", - "--global-networking", - "--jupyter", - "--persistent-volume", - "100", - ] + shlex.split( + """a --gpu 4090 --template tpl1 -e A=1 --env B=x=y + -p 8888/http --port 22/tcp --disk 40 --args 'bash start.sh' + --cmd "python -m train --lr '1e-4'" --entrypoint '/bin/sh -c' + --registry ghcr --global-networking --jupyter --persistent-volume 100""" + ) ) assert body == { "name": "a", - "image": "img", "templateId": "tpl1", "args": "bash start.sh", "cmd": ["python", "-m", "train", "--lr", "1e-4"], @@ -509,36 +476,21 @@ def test_regions_and_data_centers_merge(self): assert body["dataCenterIds"] == ["EU-RO-1", "US-TX-3"] assert body["mounts"]["network"] == [{"volumeId": "vol1", "path": "/data"}] - def test_template_without_image(self): - _, body = _create(["a", "--gpu", "4090", "--template", "tpl1"]) - assert body["templateId"] == "tpl1" - assert "image" not in body - - @pytest.mark.parametrize( - "args, message", - [ - (["--env", "NOPE"], "is not KEY=VALUE"), - (["--persistent-volume", "big"], "size is in GB"), - (["--network-volume", "vol1:"], "VALUE[:PATH]"), - ], - ) - def test_bad_values(self, args, message): - result = _invoke(["create", "a", "--gpu", "4090", "--image", "x", *args]) - assert result.exit_code == 2 - assert message in result.output - - def test_unknown_registry(self): - result = _invoke( - ["create", "a", "--gpu", "4090", "--image", "x", "--registry", "nope"] - ) - assert result.exit_code == 1 - assert "unknown registry credential 'nope'" in result.output - class TestCreateValidation: @pytest.mark.parametrize( "args, message", [ + (["a", "--gpu", "h100:2", "--gpu-count", "3", "--image", "x"], "conflicts"), + (["a", "--gpu", "4090", "--image", "x", "--env", "NOPE"], "use KEY=VALUE"), + ( + ["a", "--gpu", "4090", "--image", "x", "--persistent-volume", "big"], + "size is in GB", + ), + ( + ["a", "--gpu", "4090", "--image", "x", "--network-volume", "vol1:"], + "VALUE[:PATH]", + ), ( ["a", "--gpu", "4090", "--cpu", "cpu3c", "--image", "x"], "either --gpu or --cpu", @@ -559,6 +511,23 @@ def test_usage_errors(self, args, message): assert message in result.output assert not result.api.made("POST", "/v2/pods") + @pytest.mark.parametrize( + "options, message", + [ + (["--gpu", "h100:9"], "offers 9 per pod on secure cloud"), + (["--gpu", "b9000"], "unknown gpu 'b9000'"), + ( + ["--gpu", "4090", "--registry", "nope"], + "unknown registry credential 'nope'", + ), + ], + ) + def test_resolution_errors(self, options, message): + result = _invoke(["create", "a", "--image", "x", *options]) + assert result.exit_code == 1 + assert message in result.output + assert not result.api.made("POST", "/v2/pods") + def test_prompts_for_name(self): result = _invoke( ["create", "--gpu", "4090", "--image", "x"], @@ -612,24 +581,8 @@ def test_nothing_to_update(self): def test_simple_fields_skip_the_read(self): result = _invoke( - [ - "update", - "pod123", - "--name", - "renamed", - "--lock", - "--no-global-networking", - "--image", - "img:2", - "--registry", - "", - "--port", - "8000/http", - "--disk", - "80", - "--cmd", - "", - ] + shlex.split("""update pod123 --name renamed --lock --no-global-networking + --image img:2 --registry '' --port 8000/http --disk 80 --cmd ''""") ) assert result.exit_code == 0, result.output assert result.api.body("PATCH", "/v2/pods/pod123") == { @@ -655,7 +608,7 @@ def test_env_edits_existing(self): "-e", "NEW=1", "--unset-env", - "HF_TOKEN", + "PRIVATE_SETTING", "--unlock", ] ) @@ -760,12 +713,52 @@ def test_sync_is_mounted(self): assert result.exit_code == 0 +@pytest.mark.parametrize( + "args, expected", + [ + ([], dict(follow=False, tail=100, source=None, since=None, as_json=False)), + ( + [ + "--follow", + "--tail", + "5", + "--source", + "container", + "--since", + "2026-10-19T00:00:00Z", + "--json", + ], + dict( + follow=True, + tail=5, + source="container", + since="2026-10-19T00:00:00Z", + as_json=True, + ), + ), + ], +) +def test_logs_forward_options(args, expected): + with patch("runpod.rp_cli.logs.show_logs") as show_logs: + result = _invoke(["logs", "pod123", *args]) + assert result.exit_code == 0, result.output + show_logs.assert_called_once_with("pod123", **expected) + + +def test_logs_reject_tail_above_api_limit(): + with patch("runpod.rp_cli.logs.show_logs") as show_logs: + result = _invoke(["logs", "pod123", "--tail", "5001"]) + assert result.exit_code == 2 + assert "5000" in result.output + show_logs.assert_not_called() + + class TestCompletion: def test_completers(self): api = FakeApi() with ( patch("runpod.rp_cli.pod.run_rest_request", api), - patch("runpod.rp_cli.compute.run_rest_request", api), + patch("runpod.rp_cli.api.run_rest_request", api), patch("runpod.rp_cli.compute.list_datacenters", return_value=DATACENTERS), ): pods = pod_module._complete_pods(None, None, "pod1") From 329f7d5473fbf083aaa2b9949b9aab28271b6da6 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:35:00 -0700 Subject: [PATCH 18/25] fix: clarify log stream completion handling --- runpod/rp_cli/logs.py | 36 +++++++++++++++++------------------- 1 file changed, 17 insertions(+), 19 deletions(-) diff --git a/runpod/rp_cli/logs.py b/runpod/rp_cli/logs.py index 3ab78ebb..2a7d2f63 100644 --- a/runpod/rp_cli/logs.py +++ b/runpod/rp_cli/logs.py @@ -2,6 +2,7 @@ import asyncio import json +from contextlib import suppress from datetime import datetime, timezone import aiohttp @@ -53,24 +54,21 @@ def remember(event_id): on_event_id=remember, ) try: - while True: - event = await asyncio.wait_for( - anext(stream), - timeout=None if follow else SNAPSHOT_IDLE_SECONDS, - ) - if follow: - emit(worker, event) - elif ( - datetime.fromisoformat(event["ts"].replace("Z", "+00:00")) - >= started - ): - return - else: - records.append((worker, event)) - except StopAsyncIteration: - pass - except asyncio.TimeoutError: - pass + with suppress(StopAsyncIteration, asyncio.TimeoutError): + while True: + event = await asyncio.wait_for( + anext(stream), + timeout=None if follow else SNAPSHOT_IDLE_SECONDS, + ) + if follow: + emit(worker, event) + elif ( + datetime.fromisoformat(event["ts"].replace("Z", "+00:00")) + >= started + ): + return + else: + records.append((worker, event)) except aiohttp.ClientResponseError as exc: if not endpoint_id or exc.status != 404: raise @@ -158,6 +156,6 @@ def show_logs( ) ) except KeyboardInterrupt: - pass + return except (aiohttp.ClientError, asyncio.TimeoutError) as exc: raise click.ClickException(str(exc)) from exc From 9880c1dbee525c003ef415f23a9c698d89761249 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:26:26 -0700 Subject: [PATCH 19/25] fix: fit resource picker labels to terminal width --- runpod/rp_cli/output.py | 13 +++++++++++-- tests/test_cli/test_cli_output.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/runpod/rp_cli/output.py b/runpod/rp_cli/output.py index 31fd9baf..c715d990 100644 --- a/runpod/rp_cli/output.py +++ b/runpod/rp_cli/output.py @@ -1,6 +1,7 @@ """output, picker, prompt, and completion helpers shared by resource commands.""" import json +import shutil import sys from collections.abc import Callable, Iterable, Sequence from typing import Any @@ -72,10 +73,18 @@ def pick(message: str, choices: Sequence[tuple[str, str]], missing: str) -> str: raise click.UsageError(f"{missing} (none found)") from InquirerPy import inquirer from InquirerPy.base.control import Choice - + from rich.text import Text + + width = max(1, shutil.get_terminal_size().columns - Text(message).cell_len - 3) + items = [] + for value, label in choices: + name = Text(label) + name.truncate(width, overflow="ellipsis") + items.append(Choice(value, name=name.plain)) return inquirer.fuzzy( message=message, - choices=[Choice(value, name=label) for value, label in choices], + choices=items, + wrap_lines=False, ).execute() diff --git a/tests/test_cli/test_cli_output.py b/tests/test_cli/test_cli_output.py index b8d59249..22260e9f 100644 --- a/tests/test_cli/test_cli_output.py +++ b/tests/test_cli/test_cli_output.py @@ -78,6 +78,36 @@ def test_confirm(self): output.confirm("sure?", yes=False) ask.assert_called_once_with("sure?", abort=True) + @pytest.mark.parametrize("message", ["volume", "pod", "endpoint", "区域"]) + @pytest.mark.parametrize("width", [20, 40]) + def test_pick_truncates_labels_to_terminal_cells(self, message, width): + import os + + from rich.text import Text + + labels = ["capy-test-example-network EU-RO-1 10 GB", "训练资源" * 12] + with ( + patch.object(output, "interactive", return_value=True), + patch.object( + output.shutil, + "get_terminal_size", + return_value=os.terminal_size((width, 24)), + ), + patch("InquirerPy.inquirer.fuzzy") as fuzzy, + ): + fuzzy.return_value.execute.return_value = "second-id" + assert ( + output.pick(message, list(zip(["first-id", "second-id"], labels)), "x") + == "second-id" + ) + choices = fuzzy.call_args.kwargs["choices"] + assert [choice.value for choice in choices] == ["first-id", "second-id"] + assert all(choice.name.endswith("…") for choice in choices) + assert all( + Text(f"? {message} {choice.name}").cell_len <= width for choice in choices + ) + assert fuzzy.call_args.kwargs["wrap_lines"] is False + def test_completer(self): complete = output.completer(lambda: [("abc", "first"), ("xyz", "second")]) items = complete(None, None, "a") From 6d0518db48517796a11229e7c2814c72985034f1 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:47:39 -0700 Subject: [PATCH 20/25] fix: read pod sync SSH keys through REST --- runpod/cli/groups/ssh/functions.py | 9 +--- .../test_cli_groups/test_ssh_functions.py | 53 ++++++++++++++----- tests/test_cli/test_pod_command.py | 33 ++++++++++++ 3 files changed, 75 insertions(+), 20 deletions(-) diff --git a/runpod/cli/groups/ssh/functions.py b/runpod/cli/groups/ssh/functions.py index a6dbe013..43e4e720 100644 --- a/runpod/cli/groups/ssh/functions.py +++ b/runpod/cli/groups/ssh/functions.py @@ -9,6 +9,7 @@ import paramiko from runpod.api.ctl_commands import get_user, update_user_settings +from runpod.api.rest import run_rest_request SSH_FILES = os.path.expanduser("~/.runpod/ssh") @@ -30,18 +31,12 @@ def get_user_pub_keys(): """ Get the current user's SSH keys """ - user = get_user() - keys = "" if user["pubKey"] is None else user["pubKey"] - - # Parse the keys - keys = keys.split("\n") - keys = [key for key in keys if key != ""] + keys = run_rest_request("GET", "/v2/account/ssh-keys")["keys"] key_list = [] for key in keys: key_parts = key.split(" ") - # Basic validation if len(key_parts) < 2: continue diff --git a/tests/test_cli/test_cli_groups/test_ssh_functions.py b/tests/test_cli/test_cli_groups/test_ssh_functions.py index bac88130..33d37d7e 100644 --- a/tests/test_cli/test_cli_groups/test_ssh_functions.py +++ b/tests/test_cli/test_cli_groups/test_ssh_functions.py @@ -27,22 +27,49 @@ def test_get_ssh_key_fingerprint_invalid(self): with self.assertRaises(ValueError): get_ssh_key_fingerprint("invalidkey") - @patch("runpod.cli.groups.ssh.functions.get_user") - def test_get_user_pub_keys(self, mock_get_user): - """Test the get_user_pub_keys function""" - - # Create dummy base64 data for our mock SSH keys + @patch( + "runpod.api.ctl_commands.run_graphql_query", + side_effect=AssertionError("GraphQL forbidden"), + ) + @patch("runpod.cli.groups.ssh.functions.run_rest_request") + def test_get_user_pub_keys(self, rest, graphql): dummy_data1 = base64.b64encode("test data 1".encode("utf-8")).decode("utf-8") dummy_data2 = base64.b64encode("test data 2".encode("utf-8")).decode("utf-8") - - mock_get_user.return_value = { - "pubKey": f"ssh-rsa {dummy_data1} key1\nssh-rsa {dummy_data2} key2\n1" - } - + named = f"ssh-rsa {dummy_data1} key1" + unnamed = f"ssh-ed25519 {dummy_data2}" + rest.return_value = {"keys": [named, unnamed, "", "invalid"]} keys = get_user_pub_keys() - - self.assertEqual(len(keys), 2) - self.assertEqual(keys[0]["fingerprint"].startswith("SHA256:"), True) + self.assertEqual( + keys, + [ + { + "type": "ssh-rsa", + "key": dummy_data1, + "fingerprint": get_ssh_key_fingerprint(named), + "name": "key1", + }, + { + "type": "ssh-ed25519", + "key": dummy_data2, + "fingerprint": get_ssh_key_fingerprint(unnamed), + "name": "N/A", + }, + ], + ) + rest.assert_called_once_with("GET", "/v2/account/ssh-keys") + graphql.assert_not_called() + + @patch( + "runpod.api.ctl_commands.run_graphql_query", + side_effect=AssertionError("GraphQL forbidden"), + ) + @patch( + "runpod.cli.groups.ssh.functions.run_rest_request", return_value={"keys": []} + ) + def test_get_user_pub_keys_empty(self, rest, graphql): + self.assertEqual(get_user_pub_keys(), []) + rest.assert_called_once_with("GET", "/v2/account/ssh-keys") + graphql.assert_not_called() @patch("runpod.cli.groups.ssh.functions.get_user") def test_add_ssh_key_already_exists(self, mock_get_user): diff --git a/tests/test_cli/test_pod_command.py b/tests/test_cli/test_pod_command.py index 81eed4bb..81177b75 100644 --- a/tests/test_cli/test_pod_command.py +++ b/tests/test_cli/test_pod_command.py @@ -1,5 +1,6 @@ import json import shlex +from io import BytesIO from unittest.mock import MagicMock, patch import pytest @@ -773,3 +774,35 @@ def test_completers(self): ("vol1", "data EU-RO-1") ] assert [item.value for item in registries] == ["ghcr"] + + +def test_sync_checks_account_keys_over_rest_only(): + source, destination = MagicMock(), MagicMock() + source.__enter__.return_value = source + destination.__enter__.return_value = destination + source.ssh.exec_command.side_effect = [ + (None, BytesIO(value), None) for value in (b"1", b"exists", b"created", b"1K") + ] + destination.ssh.exec_command.return_value = (None, BytesIO(b"1"), None) + with ( + patch( + "runpod.cli.groups.ssh.functions.run_rest_request", + return_value={"keys": ["ssh-ed25519 dGVzdC1rZXk= test-fixture"]}, + ) as rest, + patch( + "runpod.api.ctl_commands.run_graphql_query", + side_effect=AssertionError("GraphQL forbidden"), + ) as graphql, + patch( + "runpod.cli.utils.ssh_cmd.SSHConnection", + side_effect=[source, destination], + ), + ): + result = _invoke(["sync", "source-pod", "destination-pod"]) + assert result.exit_code == 0, result.output + assert "Found 1 SSH key(s)" in result.output + assert "Sync completed successfully" in result.output + rest.assert_called_once_with("GET", "/v2/account/ssh-keys") + graphql.assert_not_called() + source.get_file.assert_called_once() + destination.put_file.assert_called_once() From eb6a82d05917d0d259e7618e709e955ef4e130e3 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:53:01 -0700 Subject: [PATCH 21/25] refactor: scope catalog shared helpers to hardware and output --- runpod/apps/logs.py | 33 +++--- runpod/rp_cli/logs.py | 161 ---------------------------- runpod/rp_cli/options.py | 15 --- tests/test_apps/test_logs.py | 37 ------- tests/test_cli/test_cli_logs.py | 166 ----------------------------- tests/test_cli/test_cli_options.py | 21 ---- 6 files changed, 13 insertions(+), 420 deletions(-) delete mode 100644 runpod/rp_cli/logs.py delete mode 100644 runpod/rp_cli/options.py delete mode 100644 tests/test_cli/test_cli_logs.py delete mode 100644 tests/test_cli/test_cli_options.py diff --git a/runpod/apps/logs.py b/runpod/apps/logs.py index 4a68d6d3..c6ec8999 100644 --- a/runpod/apps/logs.py +++ b/runpod/apps/logs.py @@ -6,7 +6,7 @@ import json import os -from typing import Any, AsyncIterator, Callable, Dict, List, Optional +from typing import Any, AsyncIterator, Dict, List, Optional import aiohttp @@ -46,18 +46,15 @@ async def stream_pod_logs( pod_id: str, *, endpoint_id: Optional[str] = None, - log_type: Optional[str] = "all", + log_type: str = "all", tail: int = 100, since: Optional[str] = None, - last_event_id: Optional[str] = None, - on_event_id: Optional[Callable[[str], None]] = None, ) -> AsyncIterator[Dict[str, Any]]: """follow a pod's logs as they arrive. yields {"source": "container"|"system", "line": str, "ts": str} parsed from the sse stream. pass endpoint_id for serverless workers. - log_type=None selects both sources through rest. last_event_id resumes - a stream; on_event_id receives the cursor of each delivered event. + ends when the server closes the stream or the caller breaks out. """ params: Dict[str, Any] = {"tail": str(tail)} if log_type == "all": @@ -72,32 +69,28 @@ async def stream_pod_logs( f"/workers/{_path_segment(pod_id)}/logs" ) url = _build_url(path) - if log_type: - params["source"] = log_type + params["source"] = log_type if since: params["since"] = since - headers = _headers() - if last_event_id: - headers["Last-Event-ID"] = last_event_id - client_timeout = aiohttp.ClientTimeout(total=STREAM_TIMEOUT_SECONDS, sock_read=60) + client_timeout = aiohttp.ClientTimeout( + total=STREAM_TIMEOUT_SECONDS, sock_read=60 + ) async with aiohttp.ClientSession(timeout=client_timeout) as session: - async with session.get(url, params=params, headers=headers) as resp: + async with session.get( + url, params=params, headers=_headers() + ) as resp: resp.raise_for_status() - event_id = None async for raw in resp.content: line = raw.decode("utf-8", "replace").strip() - if line.startswith("id:"): - event_id = line[3:].strip() + # sse frames: "data: {...}"; comments (heartbeats) start + # with ":" if not line.startswith("data:"): continue try: - event = json.loads(line[len("data:") :].strip()) + yield json.loads(line[len("data:") :].strip()) except json.JSONDecodeError: continue - if event_id is not None and on_event_id is not None: - on_event_id(event_id) - yield event def tail_summary(logs: Dict[str, List[str]], lines: int = 20) -> str: diff --git a/runpod/rp_cli/logs.py b/runpod/rp_cli/logs.py deleted file mode 100644 index 2a7d2f63..00000000 --- a/runpod/rp_cli/logs.py +++ /dev/null @@ -1,161 +0,0 @@ -"""shared pod and serverless log views backed by the sdk stream reader.""" - -import asyncio -import json -from contextlib import suppress -from datetime import datetime, timezone - -import aiohttp -import click - -from runpod.api.ctl_commands import _path_segment -from runpod.api.rest import run_rest_request -from runpod.apps.logs import stream_pod_logs -from runpod.rp_cli import console as ui -from runpod.rp_cli.api import api_errors - -SNAPSHOT_IDLE_SECONDS = 3 -WORKER_POLL_SECONDS = 5 -RECONNECT_SECONDS = 1 - - -async def _show_logs( - pod_id, endpoint_id, worker_ids, follow, source, tail, since, as_json -): - started = datetime.now(timezone.utc) - records = [] - tasks = {} - width = 0 - - def emit(worker, event): - if as_json: - click.echo( - json.dumps({"worker": worker, **event} if endpoint_id else event) - ) - else: - prefix = f"{worker:<{width}} " if endpoint_id else "" - click.echo(f"{prefix}[{event['source']}] {event['line']}") - - async def consume(worker): - cursor = None - - def remember(event_id): - nonlocal cursor - cursor = event_id - - while True: - stream = stream_pod_logs( - worker, - endpoint_id=endpoint_id, - log_type=source, - tail=tail, - since=since, - last_event_id=cursor, - on_event_id=remember, - ) - try: - with suppress(StopAsyncIteration, asyncio.TimeoutError): - while True: - event = await asyncio.wait_for( - anext(stream), - timeout=None if follow else SNAPSHOT_IDLE_SECONDS, - ) - if follow: - emit(worker, event) - elif ( - datetime.fromisoformat(event["ts"].replace("Z", "+00:00")) - >= started - ): - return - else: - records.append((worker, event)) - except aiohttp.ClientResponseError as exc: - if not endpoint_id or exc.status != 404: - raise - return - except aiohttp.ClientConnectionError: - if not follow: - raise - finally: - await stream.aclose() - if not follow: - return - await asyncio.sleep(RECONNECT_SECONDS) - - async def attach(): - nonlocal width - if pod_id: - ids = [pod_id] - else: - response = await asyncio.to_thread( - run_rest_request, - "GET", - f"/v2/serverless/{_path_segment(endpoint_id)}/workers", - ) - ids = [ - worker["id"] - for worker in response["workers"] - if not worker_ids or worker["id"] in worker_ids - ] - width = max([width, *(len(worker) for worker in ids)]) - for worker in ids: - if worker not in tasks: - tasks[worker] = asyncio.create_task(consume(worker)) - - try: - await attach() - if not tasks and not as_json: - ui.console.print( - "[dim]waiting for workers ...[/dim]" - if follow - else "[dim]no workers[/dim]" - ) - if not follow: - await asyncio.gather(*tasks.values()) - else: - while True: - await asyncio.sleep(WORKER_POLL_SECONDS) - for task in tasks.values(): - if task.done(): - task.result() - tasks = { - worker: task for worker, task in tasks.items() if not task.done() - } - await attach() - finally: - for task in tasks.values(): - task.cancel() - await asyncio.gather(*tasks.values(), return_exceptions=True) - for worker, event in sorted(records, key=lambda row: row[1]["ts"]): - emit(worker, event) - - -def show_logs( - pod_id=None, - *, - endpoint_id=None, - worker_ids=(), - follow=False, - source=None, - tail=100, - since=None, - as_json=False, -): - try: - with api_errors(): - asyncio.run( - _show_logs( - pod_id, - endpoint_id, - worker_ids, - follow, - source, - tail, - since, - as_json, - ) - ) - except KeyboardInterrupt: - return - except (aiohttp.ClientError, asyncio.TimeoutError) as exc: - raise click.ClickException(str(exc)) from exc diff --git a/runpod/rp_cli/options.py b/runpod/rp_cli/options.py deleted file mode 100644 index ee396289..00000000 --- a/runpod/rp_cli/options.py +++ /dev/null @@ -1,15 +0,0 @@ -"""value parsing shared by resource command options.""" - -from collections.abc import Iterable - -import click - - -def parse_env(values: Iterable[str]) -> dict[str, str]: - env = {} - for value in values: - key, separator, content = value.partition("=") - if not key or not separator: - raise click.UsageError("invalid --env value; use KEY=VALUE") - env[key] = content - return env diff --git a/tests/test_apps/test_logs.py b/tests/test_apps/test_logs.py index 02ac437d..a49510f4 100644 --- a/tests/test_apps/test_logs.py +++ b/tests/test_apps/test_logs.py @@ -77,43 +77,6 @@ async def __anext__(self): class TestStreamPodLogs: - @pytest.mark.parametrize("endpoint_id", [None, "ep1"]) - async def test_combined_rest_stream_resumes_delivered_events( - self, monkeypatch, endpoint_id - ): - monkeypatch.setenv("RUNPOD_API_KEY", "test-key") - response = MagicMock() - response.content = _FakeContent( - [ - b"id: cursor2\n", - b'data: {"source":"system","line":"ready","ts":"t2"}\n', - ] - ) - session_cm, session = _mock_session(response) - cursors = [] - with patch("aiohttp.ClientSession", return_value=session_cm): - entries = [ - entry - async for entry in stream_pod_logs( - "pod1", - endpoint_id=endpoint_id, - log_type=None, - last_event_id="cursor1", - on_event_id=cursors.append, - ) - ] - url = session.get.call_args.args[0] - suffix = ( - "/v2/serverless/ep1/workers/pod1/logs" - if endpoint_id - else "/v2/pods/pod1/logs" - ) - assert url.endswith(suffix) - assert session.get.call_args.kwargs["params"] == {"tail": "100"} - assert session.get.call_args.kwargs["headers"]["Last-Event-ID"] == "cursor1" - assert cursors == ["cursor2"] - assert entries == [{"source": "system", "line": "ready", "ts": "t2"}] - async def test_parses_sse_frames(self, monkeypatch): monkeypatch.setenv("RUNPOD_API_KEY", "test-key") frames = [ diff --git a/tests/test_cli/test_cli_logs.py b/tests/test_cli/test_cli_logs.py deleted file mode 100644 index c747f7a8..00000000 --- a/tests/test_cli/test_cli_logs.py +++ /dev/null @@ -1,166 +0,0 @@ -import asyncio -import json -from types import SimpleNamespace - -import aiohttp -import click -import pytest - -from runpod.rp_cli import logs - - -def event(line="hello", ts="2026-01-01T00:00:00Z"): - return {"source": "container", "line": line, "ts": ts} - - -@pytest.mark.parametrize("as_json", [False, True]) -def test_pod_snapshot_uses_sdk_reader(monkeypatch, capsys, as_json): - calls = [] - - async def stream(pod_id, **kwargs): - calls.append((pod_id, kwargs)) - yield event() - - monkeypatch.setattr(logs, "stream_pod_logs", stream) - logs.show_logs("pod1", tail=7, since="2025-01-01T00:00:00Z", as_json=as_json) - text = capsys.readouterr().out - assert json.loads(text) == event() if as_json else text == "[container] hello\n" - assert calls[0][0] == "pod1" - assert calls[0][1]["log_type"] is None - assert calls[0][1]["tail"] == 7 - assert calls[0][1]["since"] == "2025-01-01T00:00:00Z" - - -def test_worker_snapshot_is_sorted_and_filterable(monkeypatch, capsys): - monkeypatch.setattr( - logs, - "run_rest_request", - lambda *a, **k: {"workers": [{"id": "a"}, {"id": "b"}, {"id": "ignored"}]}, - ) - - async def stream(worker, **kwargs): - assert kwargs["endpoint_id"] == "ep1" - assert kwargs["log_type"] == "system" - yield event( - worker, "2026-01-01T00:00:02Z" if worker == "a" else "2026-01-01T00:00:01Z" - ) - - monkeypatch.setattr(logs, "stream_pod_logs", stream) - logs.show_logs( - endpoint_id="ep1", worker_ids=("a", "b"), source="system", as_json=True - ) - rows = [json.loads(line) for line in capsys.readouterr().out.splitlines()] - assert [row["worker"] for row in rows] == ["b", "a"] - - -@pytest.mark.parametrize("empty", [False, True]) -def test_snapshot_closes_idle_or_live_stream(monkeypatch, capsys, empty): - closed = [] - monkeypatch.setattr(logs, "SNAPSHOT_IDLE_SECONDS", 0.01) - - async def stream(*args, **kwargs): - try: - if not empty: - yield event("backfill") - yield event("live", "2999-01-01T00:00:00Z") - await asyncio.sleep(10) - finally: - closed.append(True) - - monkeypatch.setattr(logs, "stream_pod_logs", stream) - logs.show_logs("pod1") - assert closed == [True] - assert capsys.readouterr().out == ("" if empty else "[container] backfill\n") - - -def test_empty_workers(monkeypatch, capsys): - monkeypatch.setattr(logs, "run_rest_request", lambda *a, **k: {"workers": []}) - logs.show_logs(endpoint_id="ep1") - assert "no workers" in capsys.readouterr().out - - -@pytest.mark.parametrize( - "endpoint_id,status,raises", - [(None, 404, True), ("ep1", 404, False), ("ep1", 403, True)], -) -def test_stream_errors_do_not_hide_auth_failures( - monkeypatch, endpoint_id, status, raises -): - monkeypatch.setattr( - logs, "run_rest_request", lambda *a, **k: {"workers": [{"id": "pod1"}]} - ) - - async def stream(*args, **kwargs): - raise aiohttp.ClientResponseError( - SimpleNamespace(real_url="https://api.runpod.io/v2/pods/pod1/logs"), - (), - status=status, - message="denied", - ) - yield - - monkeypatch.setattr(logs, "stream_pod_logs", stream) - if raises: - with pytest.raises(click.ClickException): - logs.show_logs(None if endpoint_id else "pod1", endpoint_id=endpoint_id) - else: - logs.show_logs(endpoint_id=endpoint_id) - - -async def test_follow_resumes_and_cancels_streams(monkeypatch, capsys): - resumes, closed = [], [] - received = asyncio.Event() - monkeypatch.setattr(logs, "RECONNECT_SECONDS", 0) - monkeypatch.setattr(logs, "WORKER_POLL_SECONDS", 0.01) - - async def stream(worker, **kwargs): - resumes.append(kwargs["last_event_id"]) - try: - if len(resumes) == 1: - kwargs["on_event_id"]("cursor1") - yield event("first") - else: - yield event("second") - received.set() - await asyncio.sleep(10) - finally: - closed.append(True) - - monkeypatch.setattr(logs, "stream_pod_logs", stream) - task = asyncio.create_task( - logs._show_logs("pod1", None, (), True, None, 10, None, False) - ) - await asyncio.wait_for(received.wait(), 1) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert resumes == [None, "cursor1"] - assert len(closed) == 2 - assert "[container] first\n[container] second\n" == capsys.readouterr().out - - -async def test_follow_attaches_new_workers(monkeypatch, capsys): - polls = 0 - received = asyncio.Event() - monkeypatch.setattr(logs, "WORKER_POLL_SECONDS", 0.01) - - def workers(*args): - nonlocal polls - polls += 1 - return {"workers": [] if polls == 1 else [{"id": "new"}]} - - async def stream(worker, **kwargs): - yield event(worker) - received.set() - await asyncio.sleep(10) - - monkeypatch.setattr(logs, "run_rest_request", workers) - monkeypatch.setattr(logs, "stream_pod_logs", stream) - task = asyncio.create_task( - logs._show_logs(None, "ep1", (), True, None, 10, None, False) - ) - await asyncio.wait_for(received.wait(), 1) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert "new [container] new" in capsys.readouterr().out diff --git a/tests/test_cli/test_cli_options.py b/tests/test_cli/test_cli_options.py deleted file mode 100644 index aaae9516..00000000 --- a/tests/test_cli/test_cli_options.py +++ /dev/null @@ -1,21 +0,0 @@ -"""shared command option parsing.""" - -import click -import pytest - -from runpod.rp_cli.options import parse_env - - -def test_env_values_preserve_equals_and_allow_empty_values(): - assert parse_env(["TOKEN=a=b", "EMPTY=", "TOKEN=c=d"]) == { - "TOKEN": "c=d", - "EMPTY": "", - } - - -@pytest.mark.parametrize("value", ["secret-without-equals", "=secret"]) -def test_env_error_does_not_echo_the_value(value): - with pytest.raises(click.UsageError) as exc: - parse_env([value]) - assert "KEY=VALUE" in str(exc.value) - assert "secret" not in str(exc.value) 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 22/25] 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 01097e7adb7f9675349115c2f6eba62a084ce952 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 23/25] fix: harden pod readiness and log streaming --- runpod/apps/logs.py | 35 ++++++++-- runpod/rp_cli/logs.py | 29 ++++---- runpod/rp_cli/options.py | 49 +++++++++++++ runpod/rp_cli/pod.py | 67 +++++++----------- tests/test_apps/test_logs.py | 106 +++++++++++++++++++++++++++++ tests/test_cli/test_cli_logs.py | 32 ++++++++- tests/test_cli/test_cli_options.py | 79 ++++++++++++++++++++- tests/test_cli/test_pod_command.py | 94 +++++++++++++++++++++++++ 8 files changed, 423 insertions(+), 68 deletions(-) diff --git a/runpod/apps/logs.py b/runpod/apps/logs.py index 4a68d6d3..6b0aed8a 100644 --- a/runpod/apps/logs.py +++ b/runpod/apps/logs.py @@ -4,6 +4,7 @@ streams and finite snapshots use the host api, which supports both. """ +import asyncio import json import os from typing import Any, AsyncIterator, Callable, Dict, List, Optional @@ -51,6 +52,7 @@ async def stream_pod_logs( since: Optional[str] = None, last_event_id: Optional[str] = None, on_event_id: Optional[Callable[[str], None]] = None, + idle_timeout: Optional[float] = None, ) -> AsyncIterator[Dict[str, Any]]: """follow a pod's logs as they arrive. @@ -58,6 +60,8 @@ async def stream_pod_logs( parsed from the sse stream. pass endpoint_id for serverless workers. log_type=None selects both sources through rest. last_event_id resumes a stream; on_event_id receives the cursor of each delivered event. + idle_timeout ends a stream after that many seconds without content, + measured only after the response headers have been received. """ params: Dict[str, Any] = {"tail": str(tail)} if log_type == "all": @@ -85,16 +89,35 @@ async def stream_pod_logs( async with session.get(url, params=params, headers=headers) as resp: resp.raise_for_status() event_id = None - async for raw in resp.content: - line = raw.decode("utf-8", "replace").strip() - if line.startswith("id:"): - event_id = line[3:].strip() - if not line.startswith("data:"): + data = [] + content = resp.content.__aiter__() + while True: + try: + raw = await asyncio.wait_for(anext(content), timeout=idle_timeout) + except StopAsyncIteration: + return + except asyncio.TimeoutError: + if idle_timeout is None: + raise + return + line = raw.decode("utf-8", "replace").rstrip("\r\n") + if line: + field, _, value = line.partition(":") + if value.startswith(" "): + value = value[1:] + if field == "id" and "\x00" not in value: + event_id = value + elif field == "data": + data.append(value) + continue + if not data: continue try: - event = json.loads(line[len("data:") :].strip()) + event = json.loads("\n".join(data)) except json.JSONDecodeError: continue + finally: + data = [] if event_id is not None and on_event_id is not None: on_event_id(event_id) yield event diff --git a/runpod/rp_cli/logs.py b/runpod/rp_cli/logs.py index 2a7d2f63..d433ade8 100644 --- a/runpod/rp_cli/logs.py +++ b/runpod/rp_cli/logs.py @@ -2,7 +2,6 @@ import asyncio import json -from contextlib import suppress from datetime import datetime, timezone import aiohttp @@ -52,28 +51,24 @@ def remember(event_id): since=since, last_event_id=cursor, on_event_id=remember, + idle_timeout=None if follow else SNAPSHOT_IDLE_SECONDS, ) try: - with suppress(StopAsyncIteration, asyncio.TimeoutError): - while True: - event = await asyncio.wait_for( - anext(stream), - timeout=None if follow else SNAPSHOT_IDLE_SECONDS, - ) - if follow: - emit(worker, event) - elif ( - datetime.fromisoformat(event["ts"].replace("Z", "+00:00")) - >= started - ): - return - else: - records.append((worker, event)) + async for event in stream: + if follow: + emit(worker, event) + elif ( + datetime.fromisoformat(event["ts"].replace("Z", "+00:00")) + >= started + ): + return + else: + records.append((worker, event)) except aiohttp.ClientResponseError as exc: if not endpoint_id or exc.status != 404: raise return - except aiohttp.ClientConnectionError: + except (aiohttp.ClientConnectionError, asyncio.TimeoutError): if not follow: raise finally: diff --git a/runpod/rp_cli/options.py b/runpod/rp_cli/options.py index ee396289..f9ffe61c 100644 --- a/runpod/rp_cli/options.py +++ b/runpod/rp_cli/options.py @@ -1,9 +1,13 @@ """value parsing shared by resource command options.""" +import shlex from collections.abc import Iterable import click +from runpod.rp_cli.api import list_all +from runpod.rp_cli.compute import ResolutionError + def parse_env(values: Iterable[str]) -> dict[str, str]: env = {} @@ -13,3 +17,48 @@ def parse_env(values: Iterable[str]) -> dict[str, str]: raise click.UsageError("invalid --env value; use KEY=VALUE") env[key] = content return env + + +def container_fields( + options: dict, current: dict | None = None, unset_env: Iterable[str] = () +) -> dict: + current = current or {} + body = { + field: options[option] + for option, field in ( + ("image", "image"), + ("template", "templateId"), + ("args", "args"), + ("disk", "disk"), + ("global_networking", "globalNetworking"), + ) + if options.get(option) is not None + } + for option in ("cmd", "entrypoint"): + if options.get(option) is not None: + try: + body[option] = shlex.split(options[option]) + except ValueError as exc: + raise click.BadParameter(str(exc), param_hint=f"--{option}") from exc + if options.get("ports"): + body["ports"] = list(options["ports"]) + if options.get("registry") is not None: + value = options["registry"] + body["registry"] = None + if value: + matches = [ + registry["id"] + for registry in list_all("/v2/registries", "registries") + if value in (registry["id"], registry["name"]) + ] + if not matches: + raise ResolutionError( + f"unknown registry credential '{value}'. list them with rp registry list" + ) + body["registry"] = matches[0] + if options.get("env") or unset_env: + env = {**(current.get("env") or {}), **parse_env(options.get("env") or ())} + for key in unset_env: + env.pop(key, None) + body["env"] = env + return body diff --git a/runpod/rp_cli/pod.py b/runpod/rp_cli/pod.py index c4e2455b..cb2de32a 100644 --- a/runpod/rp_cli/pod.py +++ b/runpod/rp_cli/pod.py @@ -1,4 +1,3 @@ -import shlex import time from operator import itemgetter from typing import Optional @@ -9,7 +8,7 @@ from runpod.cli.groups.pod.commands import sync_pods from runpod.rp_cli import compute from runpod.rp_cli.api import api_errors, list_all -from runpod.rp_cli.options import parse_env +from runpod.rp_cli.options import container_fields from runpod.rp_cli.output import ( completer, confirm, @@ -165,17 +164,28 @@ def _show(pod: dict, as_json: bool) -> None: print_fields(fields, [(label, itemgetter(label)) for label in fields]) -def _wait_running(pod_id: str, since: float) -> dict: +def _wait_running(pod_id: str, since: float, previous: Optional[dict] = None) -> dict: from runpod.rp_cli import console as ui deadline = time.monotonic() + WAIT_TIMEOUT + previous_uptime = ((previous or {}).get("runtime") or {}).get("uptime") or 0 + runtime_restarted = previous is None with ui.console.status(f"[dim]waiting for {pod_id} to start ...[/dim]"): while True: pod = run_rest_request("GET", f"/v2/pods/{pod_id}") status = _status(pod) uptime = pod["runtime"]["uptime"] if status == "running" else 0 - # a longer uptime is a stale report from the container before the action - if status == "running" and uptime <= time.monotonic() - since + 1: + if status != "running" or uptime < previous_uptime: + runtime_restarted = True + new_start = previous is None or ( + pod.get("startedAt") and pod["startedAt"] != previous.get("startedAt") + ) + if ( + status == "running" + and new_start + and runtime_restarted + and uptime <= time.monotonic() - since + 1 + ): return pod if status in ("error", "terminated"): raise click.ClickException(f"pod {pod_id} is {status}") @@ -191,41 +201,7 @@ def _container_body( ) -> dict: """build container settings, preserving existing environment and mount paths.""" current = current or {} - body = { - field: options[option] - for option, field in ( - ("image", "image"), - ("template", "templateId"), - ("args", "args"), - ("disk", "disk"), - ("global_networking", "globalNetworking"), - ) - if options[option] is not None - } - for option in ("cmd", "entrypoint"): - if options[option] is not None: - body[option] = shlex.split(options[option]) - if options["ports"]: - body["ports"] = list(options["ports"]) - if options["registry"] is not None: - value = options["registry"] - body["registry"] = None - if value: - matches = [ - registry["id"] - for registry in list_all("/v2/registries", "registries") - if value in (registry["id"], registry["name"]) - ] - if not matches: - raise compute.ResolutionError( - f"unknown registry credential '{value}'. list them with rp registry list" - ) - body["registry"] = matches[0] - if options["env"] or unset_env: - env = {**current.get("env", {}), **parse_env(options["env"])} - for key in unset_env: - env.pop(key, None) - body["env"] = env + body = container_fields(options, current, unset_env) mounts = {} for kind in ("persistent", "network", "global"): value = options[f"{kind}_volume"] @@ -538,8 +514,8 @@ def list_pods(name, statuses, include_cluster_pods, as_json): if name: pods = [p for p in pods if name.lower() in (p.get("name") or "").lower()] if statuses: - wanted = {status.upper() for status in statuses} - pods = [p for p in pods if p.get("status") in wanted] + wanted = {status.lower() for status in statuses} + pods = [p for p in pods if _status(p) in wanted] if as_json: print_json(pods) return @@ -700,13 +676,18 @@ def _act( ): pod_id = _resolve_pod_id(pod_id, action) with api_errors(): + previous = ( + run_rest_request("GET", f"/v2/pods/{pod_id}") + if wait and action == "restart" + else None + ) since = time.monotonic() result = run_rest_request( "POST", f"/v2/pods/{pod_id}/action", json={"action": action} ) _say(f"{done} pod {pod_id}", as_json) if wait: - result = _wait_running(pod_id, since) + result = _wait_running(pod_id, since, previous) _show(result, as_json) diff --git a/tests/test_apps/test_logs.py b/tests/test_apps/test_logs.py index 02ac437d..13a5f3a5 100644 --- a/tests/test_apps/test_logs.py +++ b/tests/test_apps/test_logs.py @@ -1,5 +1,6 @@ """unit tests for pod log access.""" +import asyncio import json from unittest.mock import AsyncMock, MagicMock, patch @@ -87,6 +88,7 @@ async def test_combined_rest_stream_resumes_delivered_events( [ b"id: cursor2\n", b'data: {"source":"system","line":"ready","ts":"t2"}\n', + b"\n", ] ) session_cm, session = _mock_session(response) @@ -121,11 +123,14 @@ async def test_parses_sse_frames(self, monkeypatch): b"data: " + json.dumps({"source": "container", "line": "hi", "ts": "t1"}).encode() + b"\n", + b"\n", b"garbage\n", b"data: not-json\n", + b"\n", b"data: " + json.dumps({"source": "system", "line": "boot", "ts": "t2"}).encode() + b"\n", + b"\n", ] response = MagicMock() response.raise_for_status = MagicMock() @@ -143,6 +148,107 @@ async def test_parses_sse_frames(self, monkeypatch): assert params["stream"] == "true" assert params["tail"] == "100" + @pytest.mark.parametrize("id_first", [False, True]) + async def test_cursor_matches_each_complete_frame(self, monkeypatch, id_first): + monkeypatch.setenv("RUNPOD_API_KEY", "test-key") + frames = [] + for index in range(2): + fields = [ + f"id: cursor{index}\n".encode(), + f'data: {{"line": {index}}}\n'.encode(), + ] + frames.extend(fields if id_first else reversed(fields)) + frames.append(b"\n") + response = MagicMock() + response.content = _FakeContent(frames) + session_cm, _ = _mock_session(response) + cursors = [] + with patch("aiohttp.ClientSession", return_value=session_cm): + stream = stream_pod_logs("pod1", on_event_id=cursors.append) + assert await anext(stream) == {"line": 0} + assert cursors == ["cursor0"] + assert await anext(stream) == {"line": 1} + assert cursors == ["cursor0", "cursor1"] + await stream.aclose() + + async def test_multiline_data_and_incomplete_frames(self, monkeypatch): + monkeypatch.setenv("RUNPOD_API_KEY", "test-key") + response = MagicMock() + response.content = _FakeContent( + [ + b'data: {"source": "system",\r\n', + b'data: "line": "ready"}\r\n', + b"id: delivered\r\n", + b"\r\n", + b'data: {"line": "incomplete"}\n', + b"id: incomplete\n", + ] + ) + session_cm, _ = _mock_session(response) + cursors = [] + with patch("aiohttp.ClientSession", return_value=session_cm): + events = [ + e async for e in stream_pod_logs("pod1", on_event_id=cursors.append) + ] + assert events == [{"source": "system", "line": "ready"}] + assert cursors == ["delivered"] + + @pytest.mark.parametrize("status", [200, 403]) + async def test_idle_timeout_starts_after_response_headers( + self, monkeypatch, status + ): + monkeypatch.setenv("RUNPOD_API_KEY", "test-key") + response = MagicMock() + response.content = _FakeContent([b'data: {"line": "ready"}\n', b"\n"]) + if status == 403: + response.raise_for_status.side_effect = web.HTTPForbidden() + session_cm, session = _mock_session(response) + + async def connect(): + await asyncio.sleep(0.03) + return response + + session.get.return_value.__aenter__.side_effect = connect + with patch("aiohttp.ClientSession", return_value=session_cm): + stream = stream_pod_logs("pod1", idle_timeout=0.001) + if status == 403: + with pytest.raises(web.HTTPForbidden): + await anext(stream) + else: + assert [e async for e in stream] == [{"line": "ready"}] + session.get.return_value.__aexit__.assert_awaited_once() + + async def test_idle_timeout_closes_established_stream(self, monkeypatch): + monkeypatch.setenv("RUNPOD_API_KEY", "test-key") + closed = [] + + async def content(): + try: + yield b'data: {"line": "ready"}\n' + yield b"\n" + await asyncio.sleep(10) + finally: + closed.append(True) + + response = MagicMock() + response.content = content() + session_cm, session = _mock_session(response) + with patch("aiohttp.ClientSession", return_value=session_cm): + events = [e async for e in stream_pod_logs("pod1", idle_timeout=0.001)] + assert events == [{"line": "ready"}] + assert closed == [True] + session.get.return_value.__aexit__.assert_awaited_once() + + async def test_connection_timeout_propagates(self, monkeypatch): + monkeypatch.setenv("RUNPOD_API_KEY", "test-key") + session_cm, session = _mock_session(MagicMock()) + session.get.return_value.__aenter__.side_effect = asyncio.TimeoutError( + "connection timed out" + ) + with patch("aiohttp.ClientSession", return_value=session_cm): + with pytest.raises(asyncio.TimeoutError, match="connection timed out"): + await anext(stream_pod_logs("pod1", idle_timeout=0.001)) + @pytest.mark.parametrize( "endpoint_id,log_type,owner", [(None, "system", "pod"), ("ep1", "container", "worker")], diff --git a/tests/test_cli/test_cli_logs.py b/tests/test_cli/test_cli_logs.py index c747f7a8..1028c9ae 100644 --- a/tests/test_cli/test_cli_logs.py +++ b/tests/test_cli/test_cli_logs.py @@ -29,6 +29,7 @@ async def stream(pod_id, **kwargs): assert calls[0][1]["log_type"] is None assert calls[0][1]["tail"] == 7 assert calls[0][1]["since"] == "2025-01-01T00:00:00Z" + assert calls[0][1]["idle_timeout"] == logs.SNAPSHOT_IDLE_SECONDS def test_worker_snapshot_is_sorted_and_filterable(monkeypatch, capsys): @@ -63,7 +64,7 @@ async def stream(*args, **kwargs): if not empty: yield event("backfill") yield event("live", "2999-01-01T00:00:00Z") - await asyncio.sleep(10) + await asyncio.sleep(kwargs["idle_timeout"]) finally: closed.append(True) @@ -79,6 +80,34 @@ def test_empty_workers(monkeypatch, capsys): assert "no workers" in capsys.readouterr().out +def test_snapshot_propagates_connection_timeout(monkeypatch): + async def stream(*args, **kwargs): + raise asyncio.TimeoutError("connection timed out") + yield + + monkeypatch.setattr(logs, "stream_pod_logs", stream) + with pytest.raises(click.ClickException, match="connection timed out"): + logs.show_logs("pod1") + + +def test_snapshot_surfaces_delayed_authentication_failure(monkeypatch): + monkeypatch.setattr(logs, "SNAPSHOT_IDLE_SECONDS", 0.001) + + async def stream(*args, **kwargs): + await asyncio.sleep(0.03) + raise aiohttp.ClientResponseError( + SimpleNamespace(real_url="https://api.runpod.io/v2/pods/pod1/logs"), + (), + status=403, + message="denied", + ) + yield + + monkeypatch.setattr(logs, "stream_pod_logs", stream) + with pytest.raises(click.ClickException, match="denied"): + logs.show_logs("pod1") + + @pytest.mark.parametrize( "endpoint_id,status,raises", [(None, 404, True), ("ep1", 404, False), ("ep1", 403, True)], @@ -114,6 +143,7 @@ async def test_follow_resumes_and_cancels_streams(monkeypatch, capsys): monkeypatch.setattr(logs, "WORKER_POLL_SECONDS", 0.01) async def stream(worker, **kwargs): + assert kwargs["idle_timeout"] is None resumes.append(kwargs["last_event_id"]) try: if len(resumes) == 1: diff --git a/tests/test_cli/test_cli_options.py b/tests/test_cli/test_cli_options.py index aaae9516..6587bc95 100644 --- a/tests/test_cli/test_cli_options.py +++ b/tests/test_cli/test_cli_options.py @@ -3,7 +3,8 @@ import click import pytest -from runpod.rp_cli.options import parse_env +from runpod.rp_cli.compute import ResolutionError +from runpod.rp_cli.options import container_fields, parse_env def test_env_values_preserve_equals_and_allow_empty_values(): @@ -19,3 +20,79 @@ def test_env_error_does_not_echo_the_value(value): parse_env([value]) assert "KEY=VALUE" in str(exc.value) assert "secret" not in str(exc.value) + + +def test_container_fields_preserve_omitted_values(): + assert container_fields({}) == {} + assert container_fields({"env": (), "ports": (), "registry": None}) == {} + + +def test_container_fields_map_options_and_parse_commands(): + assert container_fields( + { + "image": "image:tag", + "template": "tpl1", + "args": "bash start.sh", + "disk": 0, + "global_networking": False, + "cmd": "python -c 'print(1)'", + "entrypoint": "", + "ports": ("8000/http", "22/tcp"), + } + ) == { + "image": "image:tag", + "templateId": "tpl1", + "args": "bash start.sh", + "disk": 0, + "globalNetworking": False, + "cmd": ["python", "-c", "print(1)"], + "entrypoint": [], + "ports": ["8000/http", "22/tcp"], + } + + +@pytest.mark.parametrize("option", ["cmd", "entrypoint"]) +def test_container_fields_reject_malformed_shell_quoting(option): + with pytest.raises(click.BadParameter) as exc: + container_fields({option: "python 'unterminated"}) + assert exc.value.param_hint == f"--{option}" + + +@pytest.mark.parametrize("registry", ["reg1", "private"]) +def test_container_fields_resolve_registry_id_or_name(monkeypatch, registry): + def list_registries(path, collection): + assert (path, collection) == ("/v2/registries", "registries") + return [{"id": "reg1", "name": "private"}] + + monkeypatch.setattr("runpod.rp_cli.options.list_all", list_registries) + assert container_fields({"registry": registry}) == {"registry": "reg1"} + + +def test_container_fields_clear_registry_without_lookup(monkeypatch): + def unexpected_lookup(*args): + pytest.fail("clearing a registry must not fetch credentials") + + monkeypatch.setattr("runpod.rp_cli.options.list_all", unexpected_lookup) + assert container_fields({"registry": ""}) == {"registry": None} + + +def test_container_fields_reject_unknown_registry(monkeypatch): + monkeypatch.setattr("runpod.rp_cli.options.list_all", lambda *args: []) + with pytest.raises(ResolutionError, match="unknown registry credential"): + container_fields({"registry": "missing"}) + + +def test_container_fields_merge_then_unset_environment(): + current = {"env": {"KEEP": "old", "EDIT": "old", "DROP": "old"}} + assert container_fields( + {"env": ("EDIT=new", "DROP=new", "EMPTY=")}, current, ("DROP", "MISSING") + ) == {"env": {"KEEP": "old", "EDIT": "new", "EMPTY": ""}} + assert current["env"] == {"KEEP": "old", "EDIT": "old", "DROP": "old"} + assert container_fields({}, {"env": {"DROP": "old"}}, ("DROP",)) == {"env": {}} + + +@pytest.mark.parametrize("current", [None, {}, {"env": None}]) +def test_container_fields_allow_empty_current_environment(current): + assert container_fields({"env": ("KEY=value",)}, current) == { + "env": {"KEY": "value"} + } diff --git a/tests/test_cli/test_pod_command.py b/tests/test_cli/test_pod_command.py index 81177b75..057f4e6d 100644 --- a/tests/test_cli/test_pod_command.py +++ b/tests/test_cli/test_pod_command.py @@ -231,6 +231,32 @@ def test_cluster_pods_and_empty(self): assert "no pods" in result.output assert api.calls[0][2] == {"includeClusterPods": "true"} + @pytest.mark.parametrize("runtime", [None, {}, {"uptime": 0}]) + @pytest.mark.parametrize("as_json", [False, True]) + def test_filters_use_displayed_status(self, runtime, as_json): + starting = {**POD, "id": "starting-pod", "runtime": runtime} + api = FakeApi( + { + ("GET", "/v2/pods"): { + "pods": [POD, starting], + "pagination": {"hasNextPage": False, "nextCursor": None}, + } + } + ) + for status, expected, excluded in ( + ("starting", "starting-pod", "pod123"), + ("running", "pod123", "starting-pod"), + ): + result = _invoke( + ["list", "--status", status, *(["--json"] if as_json else [])], api=api + ) + assert result.exit_code == 0, result.output + if as_json: + assert [p["id"] for p in json.loads(result.stdout)] == [expected] + else: + assert expected in result.stdout + assert excluded not in result.stdout + def test_api_error_is_clean(self): api = FakeApi({("GET", "/v2/pods"): error.AuthenticationError("Unauthorized")}) result = _invoke(["list"], api=api) @@ -479,6 +505,16 @@ def test_regions_and_data_centers_merge(self): class TestCreateValidation: + @pytest.mark.parametrize("option", ["--cmd", "--entrypoint"]) + @pytest.mark.parametrize("value", ["python 'unclosed", 'python "unclosed']) + def test_malformed_command_quoting(self, option, value): + result = _invoke( + ["create", "a", "--gpu", "4090", "--image", "x", option, value] + ) + assert result.exit_code == 2 + assert f"Invalid value for {option}" in result.output + assert not result.api.made("POST", "/v2/pods") + @pytest.mark.parametrize( "args, message", [ @@ -575,6 +611,13 @@ def test_timeout(self, monkeypatch): class TestUpdate: + @pytest.mark.parametrize("option", ["--cmd", "--entrypoint"]) + def test_malformed_command_quoting(self, option): + result = _invoke(["update", "pod123", option, "python 'unclosed"]) + assert result.exit_code == 2 + assert f"Invalid value for {option}" in result.output + assert not result.api.made("PATCH", "/v2/pods/pod123") + def test_nothing_to_update(self): result = _invoke(["update", "pod123"]) assert result.exit_code == 2 @@ -650,6 +693,57 @@ def test_start_wait(self): assert json.loads(result.stdout)["runtime"] == {"uptime": 1} assert "started pod pod456" in result.stderr + def test_restart_wait_rejects_recent_pre_action_runtime(self): + previous = {**POD, "runtime": {"uptime": 1}} + restarted = {**previous, "startedAt": "2026-10-19T20:02:00Z"} + api = FakeApi( + { + ("GET", "/v2/pods/pod123"): [ + previous, + previous, + {**previous, "runtime": None}, + restarted, + ] + } + ) + result = _invoke(["restart", "pod123", "--wait", "--json"], api=api) + assert result.exit_code == 0, result.output + assert json.loads(result.stdout) == restarted + assert [call[:2] for call in api.calls] == [ + ("GET", "/v2/pods/pod123"), + ("POST", "/v2/pods/pod123/action"), + ("GET", "/v2/pods/pod123"), + ("GET", "/v2/pods/pod123"), + ("GET", "/v2/pods/pod123"), + ] + + @pytest.mark.parametrize("started_at", [None, POD["startedAt"]]) + def test_restart_wait_requires_new_start_timestamp(self, monkeypatch, started_at): + monkeypatch.setattr(pod_module, "WAIT_TIMEOUT", -1) + stale = {**POD, "runtime": {"uptime": 1}, "startedAt": started_at} + api = FakeApi({("GET", "/v2/pods/pod123"): stale}) + result = _invoke(["restart", "pod123", "--wait"], api=api) + assert result.exit_code == 1 + assert "still running" in result.output + + def test_restart_wait_rejects_stale_runtime_with_new_timestamp(self): + restarted = {**POD, "startedAt": "2026-10-19T20:02:00Z"} + ready = {**restarted, "runtime": {"uptime": 1}} + api = FakeApi({("GET", "/v2/pods/pod123"): [POD, restarted, ready]}) + result = _invoke(["restart", "pod123", "--wait", "--json"], api=api) + assert result.exit_code == 0, result.output + assert json.loads(result.stdout) == ready + assert sum(call[:2] == ("GET", "/v2/pods/pod123") for call in api.calls) == 3 + + def test_restart_wait_rejects_recent_runtime_with_new_timestamp(self, monkeypatch): + monkeypatch.setattr(pod_module, "WAIT_TIMEOUT", -1) + previous = {**POD, "runtime": {"uptime": 1}} + restarted = {**previous, "startedAt": "2026-10-19T20:02:00Z"} + api = FakeApi({("GET", "/v2/pods/pod123"): [previous, restarted]}) + result = _invoke(["restart", "pod123", "--wait"], api=api) + assert result.exit_code == 1 + assert "still running" in result.output + def test_picker_offers_pods_that_allow_the_action(self): with patch("InquirerPy.inquirer.fuzzy") as fuzzy: fuzzy.return_value.execute.return_value = "pod456" From f03757bbe2e605b74bf3e01199f6fb4a1cc35c7a 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 24/25] fix: honor catalog cloud and cuda constraints --- runpod/rp_cli/catalog.py | 13 ++++++++++++- runpod/rp_cli/compute.py | 5 ++++- tests/test_cli/test_catalog_command.py | 24 ++++++++++++++++++++++++ 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/runpod/rp_cli/catalog.py b/runpod/rp_cli/catalog.py index 20db7791..4a08a251 100644 --- a/runpod/rp_cli/catalog.py +++ b/runpod/rp_cli/catalog.py @@ -109,7 +109,16 @@ def _gpu_columns(product: str, count: int, cloud: Optional[str]): lambda row, tier=tier: _money(row["totalPrice"].get(tier)), ) ) - columns.append(("MAX", lambda row: max(row["maxCount"].values(), default=None))) + columns.append( + ( + "MAX", + lambda row: ( + row["maxCount"].get(cloud) + if cloud + else max(row["maxCount"].values(), default=None) + ), + ) + ) columns += [ ("STOCK", lambda row: (row.get("availability") or "NONE").lower()), ("DATA CENTERS", lambda row: _stock(row.get("dataCenters") or [])), @@ -241,6 +250,8 @@ def catalog( if cloud and product == "SERVERLESS": raise click.UsageError("--cloud applies to pods only") cloud = cloud.lower() if cloud else None + if cpu_mode and cloud == "community": + raise click.UsageError("--cpu is only available in the secure cloud") with api_errors(): if datacenters: diff --git a/runpod/rp_cli/compute.py b/runpod/rp_cli/compute.py index 8339323f..50ca02ec 100644 --- a/runpod/rp_cli/compute.py +++ b/runpod/rp_cli/compute.py @@ -153,7 +153,10 @@ def parse_cuda(values: Iterable[str]) -> tuple[Optional[str], list[str]]: f"invalid cuda version '{part}'. use 12.4, 12.4,12.8, or '>=12.4'" ) if floor: - minimum = version + if minimum is None or tuple(map(int, version.split("."))) > tuple( + map(int, minimum.split(".")) + ): + minimum = version else: exact.append(version) if minimum and exact: diff --git a/tests/test_cli/test_catalog_command.py b/tests/test_cli/test_catalog_command.py index daccfa0d..f2137729 100644 --- a/tests/test_cli/test_catalog_command.py +++ b/tests/test_cli/test_catalog_command.py @@ -104,6 +104,30 @@ def _invoke(args, gpus=GPUS, cpus=CPUS, datacenters=DATACENTERS): class TestGpuCatalog: + @pytest.mark.parametrize( + "cloud, expected", [(None, 10), ("secure", 1), ("community", 10)] + ) + def test_max_count_matches_cloud(self, cloud, expected): + from runpod.rp_cli.catalog import _gpu_columns + + columns = dict(_gpu_columns("POD", 1, cloud)) + assert columns["MAX"](GPUS[1]) == expected + + def test_cpu_community_rejected_before_fetch(self): + result, list_gpus, list_cpus = _invoke(["--cpu", "--cloud", "community"]) + assert result.exit_code == 2 + assert "secure cloud" in result.output + list_gpus.assert_not_called() + list_cpus.assert_not_called() + + @pytest.mark.parametrize("floors", [("12.4", "12.10"), ("12.10", "12.4")]) + def test_strongest_cuda_floor_is_sent(self, floors): + result, list_gpus, _ = _invoke( + ["--cuda", ">=" + floors[0], "--cuda", ">=" + floors[1]] + ) + assert result.exit_code == 0, result.output + assert list_gpus.call_args.kwargs["min_cuda_version"] == "12.10" + def test_unavailable_gpu_omits_stock_fields(self): gpu = { key: value 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 25/25] 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 = [