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 1/9] 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 2/9] 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 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 3/9] 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 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 4/9] 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 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 5/9] 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 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 6/9] 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 7/9] 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 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 8/9] 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 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 9/9] 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