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 new file mode 100644 index 00000000..4a08a251 --- /dev/null +++ b/runpod/rp_cli/catalog.py @@ -0,0 +1,284 @@ +"""`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.api import api_errors +from runpod.rp_cli.output import json_option, print_json, print_table + + +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["availability"]) + ) + return ", ".join(f"{dc['id']} {dc['availability'].lower()}" for dc in ranked) + + +def _filter_stock(records: list[dict], regions: tuple[str, ...], available: bool): + allowed = ( + set(compute.resolve_regions(regions, compute.list_datacenters())) + if regions + else None + ) + for record in records: + row = dict(record) + if allowed is not None: + row["dataCenters"] = [ + 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.get("availability")) > 0: + yield row + + +def _unit_price(gpu: dict, product: str, cloud: Optional[str]) -> Optional[float]: + price = gpu["price"] + 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")] + rows = [] + for row in _filter_stock(gpus, regions, available): + unit = _unit_price(row, product, cloud) + if min_vram and row["memory"] < min_vram: + continue + if max_price is not None and (unit is None or unit > max_price): + continue + row["count"] = count + row["totalPrice"] = { + tier: round(value * count, 4) if value else None + for tier, value in row["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["name"], + ) + ) + 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["name"]), + ("VRAM", lambda row: f"{row['memory']}GB"), + ] + if product == "SERVERLESS": + columns += [ + ("POOL", lambda row: row["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: ( + 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 [])), + ] + 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}'") + return sorted( + _filter_stock(cpus, regions, available), + key=lambda row: -compute.availability_rank(row.get("availability")), + ) + + +def _cpu_columns(product: str): + tier = "serverlessPerVcpu" if product == "SERVERLESS" else "securePerVcpu" + return [ + ("CPU", lambda row: row["id"]), + ("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 [])), + ] + + +_DATACENTER_COLUMNS = [ + ("DATA CENTER", lambda row: row["id"]), + ("REGION", lambda row: row["region"].lower().replace("_", " ")), + ( + "VOLUME TYPES", + lambda row: [kind.lower() for kind in row["networkVolumeTypes"]], + ), + ("GLOBAL NETWORK", lambda row: row["globalNetwork"]), + ("COMPLIANCE", lambda row: row["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 + if cpu_mode and cloud == "community": + raise click.UsageError("--cpu is only available in the secure cloud") + + 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..50ca02ec --- /dev/null +++ b/runpod/rp_cli/compute.py @@ -0,0 +1,258 @@ +"""gpu, cpu, and region resolution against the rest v2 catalog.""" + +import re +from collections.abc import Iterable +from typing import Any, 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 _query_params(**values: Any) -> dict[str, Any]: + 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( + 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, + ) + return run_rest_request("GET", "/v2/catalog/gpus", params=params)["gpus"] + + +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, + ) + 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.""" + return run_rest_request( + "GET", "/v2/catalog/datacenters", params=_query_params(include=include) + )["dataCenters"] + + +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]: + """prefer exact ids, names, and pools, then match hardware-name tokens.""" + 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: + 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: + 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(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") + 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 5be6aaef..73f306f2 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -17,7 +17,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"}) @@ -550,6 +550,20 @@ async def _session() -> int: raise click.ClickException(str(exc)) from exc +@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()) + + def _discard_result(future: "asyncio.Future[None]") -> None: """mark an abandoned entrypoint's error as retrieved, so asyncio does not log it as never retrieved.""" @@ -1092,8 +1106,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..c715d990 --- /dev/null +++ b/runpod/rp_cli/output.py @@ -0,0 +1,115 @@ +"""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 + +import click +from click.shell_completion import CompletionItem + +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 table(rows: Sequence[dict], columns: Sequence[Column]): + from rich.markup import escape + from rich.table import Table + + result = Table(box=None, header_style="dim", pad_edge=False) + for index, (header, _) in enumerate(columns): + result.add_column(header, no_wrap=index == 0, overflow="fold") + for row in rows: + 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: + from rich.markup import escape + from rich.table import Table + + from runpod.rp_cli import console as ui + + result = Table.grid(padding=(0, 2)) + result.add_column(style="dim", no_wrap=True) + result.add_column(overflow="fold") + for label, value in fields: + result.add_row(label, escape(cell(value(record)))) + ui.console.print(result) + + +def interactive() -> bool: + return sys.stdin.isatty() and sys.stdout.isatty() + + +def pick(message: str, choices: Sequence[tuple[str, str]], missing: str) -> str: + 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 + 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=items, + wrap_lines=False, + ).execute() + + +def confirm(message: str, yes: bool) -> None: + 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: + return [] + return [ + CompletionItem(value, help=label) + for value, label in entries + if value.startswith(incomplete) + ] + + return complete diff --git a/tests/test_cli/test_catalog_command.py b/tests/test_cli/test_catalog_command.py new file mode 100644 index 00000000..f2137729 --- /dev/null +++ b/tests/test_cli/test_catalog_command.py @@ -0,0 +1,256 @@ +"""tests for rp catalog and rp completion.""" + +import json +from unittest.mock import patch + +import pytest +from click.testing import CliRunner + +from runpod import error +from runpod.rp_cli.main import cli + +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: + @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 + 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 + 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 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_output.py b/tests/test_cli/test_cli_output.py new file mode 100644 index 00000000..22260e9f --- /dev/null +++ b/tests/test_cli/test_cli_output.py @@ -0,0 +1,120 @@ +"""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) + + @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") + 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..d7124b4c --- /dev/null +++ b/tests/test_cli/test_compute.py @@ -0,0 +1,199 @@ +"""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"}, +] + + +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 [gpu["id"] for gpu in 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_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={"gpus": []} + ) 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"