diff --git a/runpod/apps/logs.py b/runpod/apps/logs.py index c6ec89991..6b0aed8a7 100644 --- a/runpod/apps/logs.py +++ b/runpod/apps/logs.py @@ -4,9 +4,10 @@ streams and finite snapshots use the host api, which supports both. """ +import asyncio 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 +47,21 @@ 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, + idle_timeout: Optional[float] = 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. + idle_timeout ends a stream after that many seconds without content, + measured only after the response headers have been received. """ params: Dict[str, Any] = {"tail": str(tail)} if log_type == "all": @@ -69,28 +76,51 @@ 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() - async for raw in resp.content: - line = raw.decode("utf-8", "replace").strip() - # sse frames: "data: {...}"; comments (heartbeats) start - # with ":" - if not line.startswith("data:"): + event_id = None + data = [] + content = resp.content.__aiter__() + while True: + try: + raw = await asyncio.wait_for(anext(content), timeout=idle_timeout) + except StopAsyncIteration: + return + except asyncio.TimeoutError: + if idle_timeout is None: + raise + return + line = raw.decode("utf-8", "replace").rstrip("\r\n") + if line: + field, _, value = line.partition(":") + if value.startswith(" "): + value = value[1:] + if field == "id" and "\x00" not in value: + event_id = value + elif field == "data": + data.append(value) + continue + if not data: continue try: - yield json.loads(line[len("data:") :].strip()) + event = json.loads("\n".join(data)) except json.JSONDecodeError: continue + finally: + data = [] + if event_id is not None and on_event_id is not None: + on_event_id(event_id) + yield event def tail_summary(logs: Dict[str, List[str]], lines: int = 20) -> str: diff --git a/runpod/cli/groups/ssh/functions.py b/runpod/cli/groups/ssh/functions.py index a6dbe0136..43e4e7205 100644 --- a/runpod/cli/groups/ssh/functions.py +++ b/runpod/cli/groups/ssh/functions.py @@ -9,6 +9,7 @@ import paramiko from runpod.api.ctl_commands import get_user, update_user_settings +from runpod.api.rest import run_rest_request SSH_FILES = os.path.expanduser("~/.runpod/ssh") @@ -30,18 +31,12 @@ def get_user_pub_keys(): """ Get the current user's SSH keys """ - user = get_user() - keys = "" if user["pubKey"] is None else user["pubKey"] - - # Parse the keys - keys = keys.split("\n") - keys = [key for key in keys if key != ""] + keys = run_rest_request("GET", "/v2/account/ssh-keys")["keys"] key_list = [] for key in keys: key_parts = key.split(" ") - # Basic validation if len(key_parts) < 2: continue diff --git a/runpod/rp_cli/logs.py b/runpod/rp_cli/logs.py new file mode 100644 index 000000000..d433ade82 --- /dev/null +++ b/runpod/rp_cli/logs.py @@ -0,0 +1,156 @@ +"""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, + idle_timeout=None if follow else SNAPSHOT_IDLE_SECONDS, + ) + try: + async for event in stream: + if follow: + emit(worker, event) + elif ( + datetime.fromisoformat(event["ts"].replace("Z", "+00:00")) + >= started + ): + return + else: + records.append((worker, event)) + except aiohttp.ClientResponseError as exc: + if not endpoint_id or exc.status != 404: + raise + return + except (aiohttp.ClientConnectionError, asyncio.TimeoutError): + 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/main.py b/runpod/rp_cli/main.py index c27a9d2cc..d068a9ff7 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -1104,12 +1104,12 @@ def undeploy(app_name, env_name, yes): def _mount_groups() -> None: - from runpod.cli.groups.pod.commands import pod_cli from runpod.cli.groups.ssh.commands import ssh_cli from runpod.rp_cli.catalog import catalog + from runpod.rp_cli.pod import pod from runpod.rp_cli.volume import volume - for command in (ssh_cli, pod_cli, catalog, volume): + for command in (ssh_cli, pod, catalog, volume): cli.add_command(command) diff --git a/runpod/rp_cli/options.py b/runpod/rp_cli/options.py new file mode 100644 index 000000000..f9ffe61c8 --- /dev/null +++ b/runpod/rp_cli/options.py @@ -0,0 +1,64 @@ +"""value parsing shared by resource command options.""" + +import shlex +from collections.abc import Iterable + +import click + +from runpod.rp_cli.api import list_all +from runpod.rp_cli.compute import ResolutionError + + +def parse_env(values: Iterable[str]) -> dict[str, str]: + env = {} + 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 + + +def container_fields( + options: dict, current: dict | None = None, unset_env: Iterable[str] = () +) -> dict: + current = current or {} + body = { + field: options[option] + for option, field in ( + ("image", "image"), + ("template", "templateId"), + ("args", "args"), + ("disk", "disk"), + ("global_networking", "globalNetworking"), + ) + if options.get(option) is not None + } + for option in ("cmd", "entrypoint"): + if options.get(option) is not None: + try: + body[option] = shlex.split(options[option]) + except ValueError as exc: + raise click.BadParameter(str(exc), param_hint=f"--{option}") from exc + if options.get("ports"): + body["ports"] = list(options["ports"]) + if options.get("registry") is not None: + value = options["registry"] + body["registry"] = None + if value: + matches = [ + registry["id"] + for registry in list_all("/v2/registries", "registries") + if value in (registry["id"], registry["name"]) + ] + if not matches: + raise ResolutionError( + f"unknown registry credential '{value}'. list them with rp registry list" + ) + body["registry"] = matches[0] + if options.get("env") or unset_env: + env = {**(current.get("env") or {}), **parse_env(options.get("env") or ())} + for key in unset_env: + env.pop(key, None) + body["env"] = env + return body diff --git a/runpod/rp_cli/pod.py b/runpod/rp_cli/pod.py new file mode 100644 index 000000000..cb2de32a1 --- /dev/null +++ b/runpod/rp_cli/pod.py @@ -0,0 +1,777 @@ +import time +from operator import itemgetter +from typing import Optional + +import click + +from runpod.api.rest import run_rest_request +from runpod.cli.groups.pod.commands import sync_pods +from runpod.rp_cli import compute +from runpod.rp_cli.api import api_errors, list_all +from runpod.rp_cli.options import container_fields +from runpod.rp_cli.output import ( + completer, + confirm, + interactive, + json_option, + pick, + print_fields, + print_json, + print_table, +) + +STATUSES = ["PROVISIONING", "STARTING", "RUNNING", "EXITED", "ERROR", "TERMINATED"] +DEFAULT_MOUNT_PATH = "/workspace" +WAIT_TIMEOUT = 600 +WAIT_INTERVAL = 5 + + +def _pod_choices(action: Optional[str] = None) -> list[tuple[str, str]]: + return [ + (pod["id"], f"{pod['name']} {pod['status'].lower()}") + for pod in list_all("/v2/pods", "pods") + if action is None or action in pod["actions"] + ] + + +def _resolve_pod_id(pod_id: Optional[str], action: Optional[str] = None) -> str: + if pod_id: + return pod_id + choices = [] + if interactive(): + with api_errors(): + choices = _pod_choices(action) + return pick("pod", choices, "pass a POD_ID") + + +_complete_pods = completer(_pod_choices) +_complete_datacenters = completer( + lambda: [ + (dc["id"], (dc.get("region") or "").lower()) + for dc in compute.list_datacenters() + ] +) +_complete_network_volumes = completer( + lambda: [ + (volume["id"], f"{volume.get('name')} {volume.get('dataCenter')}") + for volume in list_all("/v2/network-volumes", "networkVolumes") + ] +) +_complete_registries = completer( + lambda: [ + (registry["name"], registry["id"]) + for registry in list_all("/v2/registries", "registries") + ] +) + + +def _say(message: str, as_json: bool, warning: bool = False) -> None: + if as_json: + click.echo(message, err=True) + return + from rich.markup import escape + + from runpod.rp_cli import console as ui + + (ui.warn if warning else ui.info)(escape(message)) + + +def _status(pod: dict) -> str: + status = pod["status"].lower() + if status == "running" and not (pod.get("runtime") or {}).get("uptime"): + return "starting" + return status + + +def _summary(pod: dict) -> dict: + hardware = pod.get("gpu") or pod.get("cpu") or {} + description = None + if pod.get("gpu"): + description = f"{hardware['count']}x {hardware['id']}" + elif pod.get("cpu"): + description = f"{hardware['id']} {hardware['vcpuCount']} vcpu" + return { + "id": pod["id"], + "name": pod["name"], + "status": _status(pod), + "hardware": description, + "data center": pod.get("dataCenterId"), + "cost": f"${pod['cost']:.2f}/hr" if pod.get("cost") else None, + "image": pod["image"], + } + + +def _show(pod: dict, as_json: bool) -> None: + if as_json: + print_json(pod) + return + runtime = pod.get("runtime") or {} + mounts = pod.get("mounts") or {} + networking = pod.get("globalNetworking") or {} + ssh = pod.get("ssh") or {} + hardware = pod.get("gpu") or pod.get("cpu") or {} + volumes = [] + if mounts.get("persistent"): + volume = mounts["persistent"] + volumes.append(f"persistent {volume['size']}GB at {volume['path']}") + for kind in ("network", "global"): + for volume in mounts.get(kind, []): + volumes.append( + f"{kind} {volume['volumeId']} at {volume.get('path') or '-'}" + ) + uptime = None + if runtime.get("uptime"): + minutes, seconds = divmod(runtime["uptime"], 60) + hours, minutes = divmod(minutes, 60) + uptime = f"{hours}h {minutes}m" if hours else f"{minutes}m {seconds}s" + if not hours and not minutes: + uptime = f"{seconds}s" + fields = { + **_summary(pod), + "resources": [ + f"{hardware[key]:g}{unit}" + for key, unit in (("vcpuCount", " vcpu"), ("memory", "GB ram")) + if hardware.get(key) + ], + "cloud": (pod.get("cloud") or "").lower(), + "cuda": pod.get("cudaVersion"), + "disk": f"{pod['disk']}GB" if pod.get("disk") else None, + "volumes": volumes, + "ports": pod.get("ports"), + "env": sorted(pod.get("env") or {}), + "command": pod.get("args"), + "locked": pod.get("locked"), + "global network": bool(networking.get("enabled")) + and (networking.get("internalDns") or networking.get("ip") or True), + "cluster": (pod.get("cluster") or {}).get("id"), + "created": pod.get("createdAt"), + "started": pod.get("startedAt"), + "uptime": uptime, + "actions": pod["actions"], + "ssh": (ssh.get("proxy") or {}).get("command"), + "ssh direct": (ssh.get("direct") or {}).get("command"), + "urls": [ + f"https://{pod['id']}-{port.split('/')[0]}.proxy.runpod.net" + for port in pod.get("ports", []) + if port.endswith("/http") + ], + "tcp": [ + f"{port['ip']}:{port['public']} -> {port['private']}" + for port in runtime.get("ports", []) + if port["type"] == "tcp" and port.get("ip") and port.get("public") + ], + } + print_fields(fields, [(label, itemgetter(label)) for label in fields]) + + +def _wait_running(pod_id: str, since: float, previous: Optional[dict] = None) -> dict: + from runpod.rp_cli import console as ui + + deadline = time.monotonic() + WAIT_TIMEOUT + previous_uptime = ((previous or {}).get("runtime") or {}).get("uptime") or 0 + runtime_restarted = previous is None + with ui.console.status(f"[dim]waiting for {pod_id} to start ...[/dim]"): + while True: + pod = run_rest_request("GET", f"/v2/pods/{pod_id}") + status = _status(pod) + uptime = pod["runtime"]["uptime"] if status == "running" else 0 + if status != "running" or uptime < previous_uptime: + runtime_restarted = True + new_start = previous is None or ( + pod.get("startedAt") and pod["startedAt"] != previous.get("startedAt") + ) + if ( + status == "running" + and new_start + and runtime_restarted + and uptime <= time.monotonic() - since + 1 + ): + return pod + if status in ("error", "terminated"): + raise click.ClickException(f"pod {pod_id} is {status}") + if time.monotonic() >= deadline: + raise click.ClickException( + f"pod {pod_id} is still {status or 'unknown'} after {WAIT_TIMEOUT}s" + ) + time.sleep(WAIT_INTERVAL) + + +def _container_body( + options: dict, current: Optional[dict] = None, unset_env=() +) -> dict: + """build container settings, preserving existing environment and mount paths.""" + current = current or {} + body = container_fields(options, current, unset_env) + mounts = {} + for kind in ("persistent", "network", "global"): + value = options[f"{kind}_volume"] + if value is None: + continue + head, sep, path = value.partition(":") + flag = f"--{kind}-volume" + if not head or (sep and not path): + raise click.BadParameter( + f"'{value}' should look like VALUE[:PATH]", param_hint=flag + ) + existing = current.get("mounts", {}).get(kind) + if kind != "persistent": + existing = existing[0] if existing else None + path = path or (existing or {}).get("path") + if kind != "global": + path = path or DEFAULT_MOUNT_PATH + if kind == "persistent": + if not head.isdigit(): + raise click.BadParameter( + "size is in GB, like 50 or 50:/workspace", param_hint=flag + ) + mounts[kind] = {"size": int(head), "path": path} + else: + mounts[kind] = [{"volumeId": head, **({"path": path} if path else {})}] + if mounts: + body["mounts"] = mounts + return body + + +_CONTAINER_OPTIONS = [ + click.option( + "--image", default=None, help="Container image, like runpod/pytorch:latest." + ), + click.option( + "--template", + default=None, + help="Template id whose container settings to apply.", + ), + click.option( + "--env", + "-e", + multiple=True, + metavar="KEY=VALUE", + help="Environment variable. Repeatable.", + ), + click.option( + "--port", + "-p", + "ports", + multiple=True, + metavar="PORT/PROTO", + help="Exposed port like 8888/http or 22/tcp. Repeatable.", + ), + click.option( + "--disk", type=click.IntRange(min=1), default=None, help="Container disk in GB." + ), + click.option( + "--persistent-volume", + metavar="GB[:PATH]", + default=None, + help=f"Host-local persistent volume (path defaults to {DEFAULT_MOUNT_PATH}).", + ), + click.option( + "--network-volume", + metavar="ID[:PATH]", + default=None, + shell_complete=_complete_network_volumes, + help=f"Network volume to mount (path defaults to {DEFAULT_MOUNT_PATH}).", + ), + click.option( + "--global-volume", + metavar="ID[:PATH]", + default=None, + help="Global volume to mount.", + ), + click.option( + "--args", + default=None, + help="Container start command as one raw string ('' clears).", + ), + click.option( + "--cmd", + default=None, + help="Container CMD as a shell string, split into exec form.", + ), + click.option( + "--entrypoint", + default=None, + help="Container ENTRYPOINT as a shell string, split into exec form.", + ), + click.option( + "--registry", + default=None, + shell_complete=_complete_registries, + help="Registry credential name or id for private images ('' clears).", + ), + click.option( + "--global-networking/--no-global-networking", + default=None, + help="Private networking across data centers (NVIDIA GPUs only).", + ), +] + + +def _container_options(command): + for option in reversed(_CONTAINER_OPTIONS): + command = option(command) + return command + + +def _gpu_body( + spec: str, + count: Optional[int], + cloud: str, + cuda: tuple[str, ...], + datacenter_ids: list[str], + min_ram: Optional[int], + min_vcpu: Optional[int], + as_json: bool, +) -> dict: + """resolve a GPU specification by stock and price, then build its request.""" + query, spec_count = compute.parse_gpu_spec(spec) + if spec_count and count and spec_count != count: + raise click.UsageError(f"--gpu {spec} conflicts with --gpu-count {count}") + count = spec_count or count or 1 + min_cuda, exact_cuda = compute.parse_cuda(cuda) + gpus = compute.list_gpus( + product="POD", + count=count, + cloud=cloud.upper(), + cuda_versions=exact_cuda, + min_cuda_version=min_cuda, + ) + matches = compute.resolve_gpus(query, gpus) + allowed = set(datacenter_ids) + + def stock(gpu: dict) -> str: + if not allowed: + return gpu.get("availability") or "NONE" + return compute.best_availability( + dc.get("availability") + for dc in gpu.get("dataCenters") or [] + if dc["id"] in allowed + ) + + offered = [ + gpu + for gpu in matches + if len(matches) == 1 + or ( + gpu.get(cloud, True) + and (gpu.get("maxCount") or {}).get(cloud, count) >= count + ) + ] + if not offered: + names = ", ".join(gpu.get("name") or gpu["id"] for gpu in matches) + raise compute.ResolutionError( + f"no '{query}' gpu ({names}) offers {count} per pod on {cloud} cloud" + ) + best = min( + offered, + key=lambda gpu: ( + -compute.availability_rank(stock(gpu)), + (gpu.get("price") or {}).get(cloud) or float("inf"), + gpu.get("name") or gpu["id"], + ), + ) + availability = stock(best) + label = best.get("name") or best["id"] + price = (best.get("price") or {}).get(cloud) + if len(matches) > 1: + cost = f", ${price * count:.2f}/hr" if price else "" + _say( + f"'{query}' matches {len(matches)} gpu types; using {label} " + f"({best['id']}): {'no' if availability == 'NONE' else availability.lower()} stock{cost}", + as_json, + ) + if availability == "NONE": + _say( + f"the catalog reports no {cloud} stock for {count}x {label} " + "in the requested data centers; creation may fail", + as_json, + warning=True, + ) + body = {"id": best["id"], "count": count} + for field, value in ( + ("minCudaVersion", min_cuda), + ("allowedCudaVersions", exact_cuda), + ("minRamPerGpu", min_ram), + ("minVcpuCountPerGpu", min_vcpu), + ): + if value: + body[field] = value + return body + + +def _datacenter_ids(regions, data_centers, network_volume) -> list[str]: + ids = ( + compute.resolve_regions(regions, compute.list_datacenters()) if regions else [] + ) + ids += [dc for dc in data_centers if dc not in ids] + if not ids and network_volume: + volume_id = network_volume.partition(":")[0] + volume = run_rest_request("GET", f"/v2/network-volumes/{volume_id}") + ids = [volume["dataCenter"]] + return ids + + +def _create_body( + name, + gpu, + gpu_count, + cpu, + vcpu, + cloud, + regions, + data_centers, + cuda, + min_ram_per_gpu, + min_vcpu_per_gpu, + ssh, + jupyter, + as_json, + **container, +): + if gpu and cpu: + raise click.UsageError("use either --gpu or --cpu") + if not gpu and not cpu: + raise click.UsageError("pass --gpu (like 4090 or h100:2) or --cpu (like cpu3c)") + if cpu and (gpu_count or cuda or min_ram_per_gpu or min_vcpu_per_gpu): + raise click.UsageError( + "--gpu-count, --cuda, --min-ram-per-gpu, and --min-vcpu-per-gpu apply to gpu pods" + ) + if vcpu and not cpu: + raise click.UsageError("--vcpu applies to cpu pods") + if not container["image"] and not container["template"]: + raise click.UsageError("pass --image or --template") + if not name: + if not interactive(): + raise click.UsageError("pass a NAME") + name = click.prompt("pod name") + + body = {"name": name, **_container_body(container)} + datacenter_ids = _datacenter_ids(regions, data_centers, container["network_volume"]) + if datacenter_ids: + body["dataCenterIds"] = datacenter_ids + if cloud: + body["cloud"] = cloud.upper() + + if gpu: + body["gpu"] = _gpu_body( + gpu, + gpu_count, + (cloud or "secure").lower(), + cuda, + datacenter_ids, + min_ram_per_gpu, + min_vcpu_per_gpu, + as_json, + ) + else: + flavor = compute.resolve_cpu( + cpu, compute.list_cpus(product="POD", vcpu_count=vcpu) + ) + body["cpu"] = { + "id": flavor["id"], + "vcpuCount": vcpu or (flavor.get("vcpu") or {}).get("min") or 2, + } + + if ssh: + body["startSsh"] = True + keys = run_rest_request("GET", "/v2/account/ssh-keys")["keys"] + if not keys: + _say( + "your account has no ssh keys, so --ssh has no effect. " + "add one with rp ssh add", + as_json, + warning=True, + ) + if jupyter: + body["startJupyter"] = True + + return body + + +@click.group("pod") +def pod(): + """Create and manage pods.""" + + +@pod.command(name="list") +@click.option("--name", default=None, help="Only pods whose name contains this text.") +@click.option( + "--status", + "statuses", + multiple=True, + type=click.Choice(STATUSES, case_sensitive=False), + help="Only pods in this status. Repeatable.", +) +@click.option( + "--include-cluster-pods", is_flag=True, help="Include cluster member pods." +) +@json_option +def list_pods(name, statuses, include_cluster_pods, as_json): + """List your pods.""" + with api_errors(): + params = {"includeClusterPods": "true"} if include_cluster_pods else None + pods = list_all("/v2/pods", "pods", params) + if name: + pods = [p for p in pods if name.lower() in (p.get("name") or "").lower()] + if statuses: + wanted = {status.lower() for status in statuses} + pods = [p for p in pods if _status(p) in wanted] + if as_json: + print_json(pods) + return + columns = ["id", "name", "status", "hardware", "data center", "cost", "image"] + print_table( + [_summary(p) for p in pods], + [(label.upper(), itemgetter(label)) for label in columns], + "no pods. create one with rp pod create", + ) + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@json_option +def get(pod_id, as_json): + """Show a pod with its connection details.""" + pod_id = _resolve_pod_id(pod_id) + with api_errors(): + _show(run_rest_request("GET", f"/v2/pods/{pod_id}"), as_json) + + +@pod.command() +@click.argument("name", required=False) +@click.option( + "--gpu", default=None, metavar="TYPE[:COUNT]", help="GPU type like 4090 or h100:2." +) +@click.option( + "--gpu-count", type=click.IntRange(min=1), default=None, help="GPUs per pod." +) +@click.option( + "--cpu", default=None, help="CPU flavor like cpu3c or cpu5g, for a CPU pod." +) +@click.option( + "--vcpu", + type=click.IntRange(min=2), + default=None, + help="vCPUs for a CPU pod (defaults to the flavor's minimum).", +) +@click.option( + "--cloud", + type=click.Choice(["secure", "community"], case_sensitive=False), + default=None, + help="Cloud tier (the api defaults to secure).", +) +@click.option( + "--region", + "regions", + multiple=True, + help="Region (eu, na, asia), id prefix (us, eu-ro), or data center id. Repeatable.", +) +@click.option( + "--data-center", + "data_centers", + multiple=True, + shell_complete=_complete_datacenters, + help="Exact data center id like US-TX-3. Repeatable.", +) +@click.option( + "--cuda", + multiple=True, + help="Host CUDA: '>=12.4' for a floor, 12.4 or 12.4,12.8 for exact versions.", +) +@click.option( + "--min-ram-per-gpu", + type=click.IntRange(min=1), + default=None, + help="Minimum host RAM in GB per GPU.", +) +@click.option( + "--min-vcpu-per-gpu", + type=click.IntRange(min=1), + default=None, + help="Minimum host vCPUs per GPU.", +) +@_container_options +@click.option("--ssh", is_flag=True, help="Start sshd with your account's SSH keys.") +@click.option( + "--jupyter", is_flag=True, help="Start JupyterLab (expose 8888/http to reach it)." +) +@click.option("--wait", is_flag=True, help="Wait until the pod is running.") +@json_option +def create(name, wait, as_json, **options): + """Create a pod named NAME. + + GPU types are human-readable (4090, h100, a100-sxm). When a name + matches several types, the one with the most stock in the requested + regions wins, then the cheapest. + + \b + rp pod create trainer --gpu h100:2 --image runpod/pytorch:latest --region eu --ssh + rp pod create box --cpu cpu3c --vcpu 4 --image ubuntu:24.04 + rp pod create nb --gpu 4090 --image runpod/pytorch:latest --jupyter -p 8888/http + """ + with api_errors(): + body = _create_body(name, as_json=as_json, **options) + since = time.monotonic() + created = run_rest_request("POST", "/v2/pods", json=body) + _say(f"created pod {created['name']} ({created['id']})", as_json) + if wait: + created = _wait_running(created["id"], since) + _show(created, as_json) + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@click.option("--name", default=None, help="New pod name.") +@_container_options +@click.option( + "--unset-env", + multiple=True, + metavar="KEY", + help="Remove an environment variable. Repeatable.", +) +@click.option( + "--lock/--unlock", + "locked", + default=None, + help="Lock the pod against stop and reset.", +) +@json_option +def update(pod_id, name, unset_env, locked, as_json, **container): + """Change a pod's settings. + + --env and --unset-env edit the existing environment, --port replaces + the exposed ports, and volume flags restate the full mount (the + volume id cannot change; the path defaults to the current one). + Most changes apply on the pod's next start. + """ + requested = [name, locked, *unset_env, *container.values()] + if all(value in (None, ()) for value in requested): + raise click.UsageError("nothing to update. see rp pod update --help") + pod_id = _resolve_pod_id(pod_id) + + with api_errors(): + needs_current = ( + container["env"] + or unset_env + or any( + container[key] + for key in ("persistent_volume", "network_volume", "global_volume") + ) + ) + current = ( + run_rest_request("GET", f"/v2/pods/{pod_id}") if needs_current else None + ) + body = _container_body(container, current, unset_env) + if name: + body["name"] = name + if locked is not None: + body["locked"] = locked + updated = run_rest_request("PATCH", f"/v2/pods/{pod_id}", json=body) + _say(f"updated pod {pod_id}", as_json) + _show(updated, as_json) + + +def _act( + pod_id: Optional[str], action: str, done: str, as_json: bool, wait: bool = False +): + pod_id = _resolve_pod_id(pod_id, action) + with api_errors(): + previous = ( + run_rest_request("GET", f"/v2/pods/{pod_id}") + if wait and action == "restart" + else None + ) + since = time.monotonic() + result = run_rest_request( + "POST", f"/v2/pods/{pod_id}/action", json={"action": action} + ) + _say(f"{done} pod {pod_id}", as_json) + if wait: + result = _wait_running(pod_id, since, previous) + _show(result, as_json) + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@click.option("--wait", is_flag=True, help="Wait until the pod is running.") +@json_option +def start(pod_id, wait, as_json): + """Start a stopped pod.""" + _act(pod_id, "start", "started", as_json, wait) + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@json_option +def stop(pod_id, as_json): + """Stop a running pod, keeping its volume.""" + _act(pod_id, "stop", "stopped", as_json) + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@click.option("--wait", is_flag=True, help="Wait until the pod is running.") +@json_option +def restart(pod_id, wait, as_json): + """Restart a running pod.""" + _act(pod_id, "restart", "restarted", as_json, wait) + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.") +def terminate(pod_id, yes): + """Terminate a pod permanently.""" + from runpod.rp_cli import console as ui + + pod_id = _resolve_pod_id(pod_id, "terminate") + confirm(f"terminate pod {pod_id}? its disk and persistent volume are deleted", yes) + with api_errors(): + run_rest_request("DELETE", f"/v2/pods/{pod_id}") + ui.success(f"terminated pod {pod_id}") + + +pod.add_command(terminate, "delete") + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +@click.option("--follow", "-f", is_flag=True, help="Follow new log lines.") +@click.option("--tail", type=click.IntRange(0, 5000), default=100, show_default=True) +@click.option("--source", type=click.Choice(["container", "system"])) +@click.option("--since", help="Only logs after this timestamp.") +@json_option +def logs(pod_id, follow, tail, source, since, as_json): + """Read a pod's container or system logs.""" + from runpod.rp_cli.logs import show_logs + + pod_id = _resolve_pod_id(pod_id) + show_logs( + pod_id, + follow=follow, + tail=tail, + source=source, + since=since, + as_json=as_json, + ) + + +@pod.command() +@click.argument("pod_id", required=False, shell_complete=_complete_pods) +def connect(pod_id): + """Open an interactive terminal on a pod over SSH. + + Needs an SSH key on your account (rp ssh add) and 22/tcp exposed. + """ + from runpod.cli.utils import ssh_cmd + + pod_id = _resolve_pod_id(pod_id) + click.echo(f"connecting to pod {pod_id}...") + try: + ssh = ssh_cmd.SSHConnection(pod_id) + except (ValueError, TimeoutError) as exc: + raise click.ClickException(str(exc)) from exc + ssh.launch_terminal() + + +pod.add_command(sync_pods, "sync") diff --git a/tests/test_apps/test_logs.py b/tests/test_apps/test_logs.py index a49510f4a..13a5f3a5c 100644 --- a/tests/test_apps/test_logs.py +++ b/tests/test_apps/test_logs.py @@ -1,5 +1,6 @@ """unit tests for pod log access.""" +import asyncio import json from unittest.mock import AsyncMock, MagicMock, patch @@ -77,6 +78,44 @@ 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', + b"\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 = [ @@ -84,11 +123,14 @@ async def test_parses_sse_frames(self, monkeypatch): b"data: " + json.dumps({"source": "container", "line": "hi", "ts": "t1"}).encode() + b"\n", + b"\n", b"garbage\n", b"data: not-json\n", + b"\n", b"data: " + json.dumps({"source": "system", "line": "boot", "ts": "t2"}).encode() + b"\n", + b"\n", ] response = MagicMock() response.raise_for_status = MagicMock() @@ -106,6 +148,107 @@ async def test_parses_sse_frames(self, monkeypatch): assert params["stream"] == "true" assert params["tail"] == "100" + @pytest.mark.parametrize("id_first", [False, True]) + async def test_cursor_matches_each_complete_frame(self, monkeypatch, id_first): + monkeypatch.setenv("RUNPOD_API_KEY", "test-key") + frames = [] + for index in range(2): + fields = [ + f"id: cursor{index}\n".encode(), + f'data: {{"line": {index}}}\n'.encode(), + ] + frames.extend(fields if id_first else reversed(fields)) + frames.append(b"\n") + response = MagicMock() + response.content = _FakeContent(frames) + session_cm, _ = _mock_session(response) + cursors = [] + with patch("aiohttp.ClientSession", return_value=session_cm): + stream = stream_pod_logs("pod1", on_event_id=cursors.append) + assert await anext(stream) == {"line": 0} + assert cursors == ["cursor0"] + assert await anext(stream) == {"line": 1} + assert cursors == ["cursor0", "cursor1"] + await stream.aclose() + + async def test_multiline_data_and_incomplete_frames(self, monkeypatch): + monkeypatch.setenv("RUNPOD_API_KEY", "test-key") + response = MagicMock() + response.content = _FakeContent( + [ + b'data: {"source": "system",\r\n', + b'data: "line": "ready"}\r\n', + b"id: delivered\r\n", + b"\r\n", + b'data: {"line": "incomplete"}\n', + b"id: incomplete\n", + ] + ) + session_cm, _ = _mock_session(response) + cursors = [] + with patch("aiohttp.ClientSession", return_value=session_cm): + events = [ + e async for e in stream_pod_logs("pod1", on_event_id=cursors.append) + ] + assert events == [{"source": "system", "line": "ready"}] + assert cursors == ["delivered"] + + @pytest.mark.parametrize("status", [200, 403]) + async def test_idle_timeout_starts_after_response_headers( + self, monkeypatch, status + ): + monkeypatch.setenv("RUNPOD_API_KEY", "test-key") + response = MagicMock() + response.content = _FakeContent([b'data: {"line": "ready"}\n', b"\n"]) + if status == 403: + response.raise_for_status.side_effect = web.HTTPForbidden() + session_cm, session = _mock_session(response) + + async def connect(): + await asyncio.sleep(0.03) + return response + + session.get.return_value.__aenter__.side_effect = connect + with patch("aiohttp.ClientSession", return_value=session_cm): + stream = stream_pod_logs("pod1", idle_timeout=0.001) + if status == 403: + with pytest.raises(web.HTTPForbidden): + await anext(stream) + else: + assert [e async for e in stream] == [{"line": "ready"}] + session.get.return_value.__aexit__.assert_awaited_once() + + async def test_idle_timeout_closes_established_stream(self, monkeypatch): + monkeypatch.setenv("RUNPOD_API_KEY", "test-key") + closed = [] + + async def content(): + try: + yield b'data: {"line": "ready"}\n' + yield b"\n" + await asyncio.sleep(10) + finally: + closed.append(True) + + response = MagicMock() + response.content = content() + session_cm, session = _mock_session(response) + with patch("aiohttp.ClientSession", return_value=session_cm): + events = [e async for e in stream_pod_logs("pod1", idle_timeout=0.001)] + assert events == [{"line": "ready"}] + assert closed == [True] + session.get.return_value.__aexit__.assert_awaited_once() + + async def test_connection_timeout_propagates(self, monkeypatch): + monkeypatch.setenv("RUNPOD_API_KEY", "test-key") + session_cm, session = _mock_session(MagicMock()) + session.get.return_value.__aenter__.side_effect = asyncio.TimeoutError( + "connection timed out" + ) + with patch("aiohttp.ClientSession", return_value=session_cm): + with pytest.raises(asyncio.TimeoutError, match="connection timed out"): + await anext(stream_pod_logs("pod1", idle_timeout=0.001)) + @pytest.mark.parametrize( "endpoint_id,log_type,owner", [(None, "system", "pod"), ("ep1", "container", "worker")], diff --git a/tests/test_cli/test_cli_groups/test_ssh_functions.py b/tests/test_cli/test_cli_groups/test_ssh_functions.py index bac88130f..33d37d7eb 100644 --- a/tests/test_cli/test_cli_groups/test_ssh_functions.py +++ b/tests/test_cli/test_cli_groups/test_ssh_functions.py @@ -27,22 +27,49 @@ def test_get_ssh_key_fingerprint_invalid(self): with self.assertRaises(ValueError): get_ssh_key_fingerprint("invalidkey") - @patch("runpod.cli.groups.ssh.functions.get_user") - def test_get_user_pub_keys(self, mock_get_user): - """Test the get_user_pub_keys function""" - - # Create dummy base64 data for our mock SSH keys + @patch( + "runpod.api.ctl_commands.run_graphql_query", + side_effect=AssertionError("GraphQL forbidden"), + ) + @patch("runpod.cli.groups.ssh.functions.run_rest_request") + def test_get_user_pub_keys(self, rest, graphql): dummy_data1 = base64.b64encode("test data 1".encode("utf-8")).decode("utf-8") dummy_data2 = base64.b64encode("test data 2".encode("utf-8")).decode("utf-8") - - mock_get_user.return_value = { - "pubKey": f"ssh-rsa {dummy_data1} key1\nssh-rsa {dummy_data2} key2\n1" - } - + named = f"ssh-rsa {dummy_data1} key1" + unnamed = f"ssh-ed25519 {dummy_data2}" + rest.return_value = {"keys": [named, unnamed, "", "invalid"]} keys = get_user_pub_keys() - - self.assertEqual(len(keys), 2) - self.assertEqual(keys[0]["fingerprint"].startswith("SHA256:"), True) + self.assertEqual( + keys, + [ + { + "type": "ssh-rsa", + "key": dummy_data1, + "fingerprint": get_ssh_key_fingerprint(named), + "name": "key1", + }, + { + "type": "ssh-ed25519", + "key": dummy_data2, + "fingerprint": get_ssh_key_fingerprint(unnamed), + "name": "N/A", + }, + ], + ) + rest.assert_called_once_with("GET", "/v2/account/ssh-keys") + graphql.assert_not_called() + + @patch( + "runpod.api.ctl_commands.run_graphql_query", + side_effect=AssertionError("GraphQL forbidden"), + ) + @patch( + "runpod.cli.groups.ssh.functions.run_rest_request", return_value={"keys": []} + ) + def test_get_user_pub_keys_empty(self, rest, graphql): + self.assertEqual(get_user_pub_keys(), []) + rest.assert_called_once_with("GET", "/v2/account/ssh-keys") + graphql.assert_not_called() @patch("runpod.cli.groups.ssh.functions.get_user") def test_add_ssh_key_already_exists(self, mock_get_user): diff --git a/tests/test_cli/test_cli_logs.py b/tests/test_cli/test_cli_logs.py new file mode 100644 index 000000000..1028c9aea --- /dev/null +++ b/tests/test_cli/test_cli_logs.py @@ -0,0 +1,196 @@ +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" + assert calls[0][1]["idle_timeout"] == logs.SNAPSHOT_IDLE_SECONDS + + +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(kwargs["idle_timeout"]) + 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 + + +def test_snapshot_propagates_connection_timeout(monkeypatch): + async def stream(*args, **kwargs): + raise asyncio.TimeoutError("connection timed out") + yield + + monkeypatch.setattr(logs, "stream_pod_logs", stream) + with pytest.raises(click.ClickException, match="connection timed out"): + logs.show_logs("pod1") + + +def test_snapshot_surfaces_delayed_authentication_failure(monkeypatch): + monkeypatch.setattr(logs, "SNAPSHOT_IDLE_SECONDS", 0.001) + + async def stream(*args, **kwargs): + await asyncio.sleep(0.03) + raise aiohttp.ClientResponseError( + SimpleNamespace(real_url="https://api.runpod.io/v2/pods/pod1/logs"), + (), + status=403, + message="denied", + ) + yield + + monkeypatch.setattr(logs, "stream_pod_logs", stream) + with pytest.raises(click.ClickException, match="denied"): + logs.show_logs("pod1") + + +@pytest.mark.parametrize( + "endpoint_id,status,raises", + [(None, 404, True), ("ep1", 404, False), ("ep1", 403, True)], +) +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): + assert kwargs["idle_timeout"] is None + 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 new file mode 100644 index 000000000..6587bc953 --- /dev/null +++ b/tests/test_cli/test_cli_options.py @@ -0,0 +1,98 @@ +"""shared command option parsing.""" + +import click +import pytest + +from runpod.rp_cli.compute import ResolutionError +from runpod.rp_cli.options import container_fields, parse_env + + +def test_env_values_preserve_equals_and_allow_empty_values(): + 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) + + +def test_container_fields_preserve_omitted_values(): + assert container_fields({}) == {} + assert container_fields({"env": (), "ports": (), "registry": None}) == {} + + +def test_container_fields_map_options_and_parse_commands(): + assert container_fields( + { + "image": "image:tag", + "template": "tpl1", + "args": "bash start.sh", + "disk": 0, + "global_networking": False, + "cmd": "python -c 'print(1)'", + "entrypoint": "", + "ports": ("8000/http", "22/tcp"), + } + ) == { + "image": "image:tag", + "templateId": "tpl1", + "args": "bash start.sh", + "disk": 0, + "globalNetworking": False, + "cmd": ["python", "-c", "print(1)"], + "entrypoint": [], + "ports": ["8000/http", "22/tcp"], + } + + +@pytest.mark.parametrize("option", ["cmd", "entrypoint"]) +def test_container_fields_reject_malformed_shell_quoting(option): + with pytest.raises(click.BadParameter) as exc: + container_fields({option: "python 'unterminated"}) + assert exc.value.param_hint == f"--{option}" + + +@pytest.mark.parametrize("registry", ["reg1", "private"]) +def test_container_fields_resolve_registry_id_or_name(monkeypatch, registry): + def list_registries(path, collection): + assert (path, collection) == ("/v2/registries", "registries") + return [{"id": "reg1", "name": "private"}] + + monkeypatch.setattr("runpod.rp_cli.options.list_all", list_registries) + assert container_fields({"registry": registry}) == {"registry": "reg1"} + + +def test_container_fields_clear_registry_without_lookup(monkeypatch): + def unexpected_lookup(*args): + pytest.fail("clearing a registry must not fetch credentials") + + monkeypatch.setattr("runpod.rp_cli.options.list_all", unexpected_lookup) + assert container_fields({"registry": ""}) == {"registry": None} + + +def test_container_fields_reject_unknown_registry(monkeypatch): + monkeypatch.setattr("runpod.rp_cli.options.list_all", lambda *args: []) + with pytest.raises(ResolutionError, match="unknown registry credential"): + container_fields({"registry": "missing"}) + + +def test_container_fields_merge_then_unset_environment(): + current = {"env": {"KEEP": "old", "EDIT": "old", "DROP": "old"}} + assert container_fields( + {"env": ("EDIT=new", "DROP=new", "EMPTY=")}, current, ("DROP", "MISSING") + ) == {"env": {"KEEP": "old", "EDIT": "new", "EMPTY": ""}} + assert current["env"] == {"KEEP": "old", "EDIT": "old", "DROP": "old"} + assert container_fields({}, {"env": {"DROP": "old"}}, ("DROP",)) == {"env": {}} + + +@pytest.mark.parametrize("current", [None, {}, {"env": None}]) +def test_container_fields_allow_empty_current_environment(current): + assert container_fields({"env": ("KEY=value",)}, current) == { + "env": {"KEY": "value"} + } diff --git a/tests/test_cli/test_pod_command.py b/tests/test_cli/test_pod_command.py new file mode 100644 index 000000000..057f4e6d1 --- /dev/null +++ b/tests/test_cli/test_pod_command.py @@ -0,0 +1,902 @@ +import json +import shlex +from io import BytesIO +from unittest.mock import MagicMock, patch + +import pytest +from click.testing import CliRunner + +from runpod import error +from runpod.rp_cli import pod as pod_module +from runpod.rp_cli.main import cli + +GPUS = [ + { + "id": "NVIDIA H100 80GB HBM3", + "name": "H100 SXM", + "price": {"secure": 3.99, "community": 2.69}, + "maxCount": {"secure": 8, "community": 1}, + "secure": True, + "community": True, + "availability": "LOW", + "dataCenters": [ + {"id": "US-TX-3", "availability": "LOW"}, + {"id": "EU-RO-1", "availability": "LOW"}, + ], + }, + { + "id": "NVIDIA H100 NVL", + "name": "H100 NVL", + "price": {"secure": 3.19, "community": 2.59}, + "maxCount": {"secure": 1, "community": 10}, + "secure": True, + "community": True, + "availability": "HIGH", + "dataCenters": [{"id": "US-TX-3", "availability": "HIGH"}], + }, + { + "id": "NVIDIA H100 PCIe", + "name": "H100 PCIe", + "price": {"secure": 2.39, "community": 1.99}, + "maxCount": {"secure": 8, "community": 8}, + "secure": True, + "community": True, + "availability": "LOW", + "dataCenters": [{"id": "US-TX-3", "availability": "LOW"}], + }, + { + "id": "NVIDIA GeForce RTX 4090", + "name": "RTX 4090", + "price": {"secure": 0.69, "community": 0.34}, + "maxCount": {"secure": 8, "community": 8}, + "secure": True, + "community": True, + "availability": "NONE", + }, +] + +CPUS = [ + { + "id": "cpu3c", + "name": "Compute-Optimized", + "group": "CPU3", + "vcpu": {"min": 2, "max": 32}, + }, + { + "id": "cpu5g", + "name": "General Purpose", + "group": "CPU5", + "vcpu": {"min": 4, "max": 32}, + }, +] + +DATACENTERS = [ + {"id": "EU-RO-1", "region": "EUROPE"}, + {"id": "EU-CZ-1", "region": "EUROPE"}, + {"id": "US-TX-3", "region": "NORTH_AMERICA"}, +] + +POD = { + "id": "pod123", + "name": "trainer", + "status": "RUNNING", + "actions": ["stop", "restart", "terminate"], + "image": "runpod/pytorch:latest", + "args": "", + "disk": 50, + "ports": ["8888/http", "22/tcp"], + "env": {"PRIVATE_SETTING": "redacted-test-value", "MODEL": "llama"}, + "mounts": {"network": [{"volumeId": "vol1", "path": "/data"}]}, + "gpu": {"id": "NVIDIA H100 80GB HBM3", "count": 2, "vcpuCount": 32, "memory": 250}, + "cloud": "SECURE", + "dataCenterId": "EU-RO-1", + "cudaVersion": "12.8", + "cost": 7.98, + "locked": False, + "globalNetworking": {"enabled": True, "internalDns": "pod123.runpod.internal"}, + "ssh": { + "proxy": {"command": "ssh pod123-abc@ssh.runpod.io"}, + "direct": {"command": "ssh root@1.2.3.4 -p 40022"}, + }, + "runtime": { + "ports": [ + {"private": 22, "public": 40022, "type": "tcp", "ip": "1.2.3.4"}, + {"private": 8888, "public": None, "type": "http", "ip": None}, + ], + "uptime": 3725, + }, + "createdAt": "2026-10-19T20:00:00Z", + "startedAt": "2026-10-19T20:01:00Z", +} + +STOPPED = { + "id": "pod456", + "name": "box", + "status": "EXITED", + "actions": ["start", "terminate"], + "image": "ubuntu:24.04", + "cpu": {"id": "cpu3c", "vcpuCount": 4, "memory": 8}, + "cost": 0, +} + + +class FakeApi: + def __init__(self, routes=None): + self.calls = [] + self.routes = { + ("GET", "/v2/pods"): { + "pods": [POD, STOPPED], + "pagination": {"hasNextPage": False, "nextCursor": None}, + }, + ("GET", "/v2/pods/pod123"): POD, + ("GET", "/v2/account/ssh-keys"): {"keys": ["redacted-test-key"]}, + ("GET", "/v2/registries"): { + "registries": [{"id": "reg1", "name": "ghcr"}], + "pagination": {"hasNextPage": False, "nextCursor": None}, + }, + ("GET", "/v2/network-volumes"): { + "networkVolumes": [ + {"id": "vol1", "name": "data", "dataCenter": "EU-RO-1"} + ], + "pagination": {"hasNextPage": False, "nextCursor": None}, + }, + ("GET", "/v2/network-volumes/vol1"): { + "id": "vol1", + "dataCenter": "EU-RO-1", + }, + ("POST", "/v2/pods"): {**POD, "status": "PROVISIONING"}, + ("PATCH", "/v2/pods/pod123"): POD, + ("POST", "/v2/pods/pod123/action"): POD, + ("POST", "/v2/pods/pod456/action"): {**STOPPED, "status": "STARTING"}, + ("DELETE", "/v2/pods/pod123"): None, + **(routes or {}), + } + + def __call__(self, method, path, params=None, json=None): + self.calls.append((method, path, params, json)) + response = self.routes[(method, path)] + if isinstance(response, list): + return response.pop(0) + if isinstance(response, Exception): + raise response + return response + + def body(self, method, path): + return next(call[3] for call in self.calls if call[:2] == (method, path)) + + def made(self, method, path): + return any(call[:2] == (method, path) for call in self.calls) + + +@pytest.fixture(autouse=True) +def cli_environment(monkeypatch): + from runpod.rp_cli import console as ui + + monkeypatch.setattr(ui.console, "_width", 200) + monkeypatch.setattr(pod_module.time, "sleep", lambda seconds: None) + monkeypatch.setattr( + "requests.request", + MagicMock(side_effect=AssertionError("unexpected HTTP request")), + ) + + +def _invoke(args, api=None, gpus=GPUS, interactive=False, input=None): + api = api or FakeApi() + with ( + patch("runpod.rp_cli.pod.run_rest_request", api), + patch("runpod.rp_cli.api.run_rest_request", api), + patch("runpod.rp_cli.compute.list_gpus", return_value=gpus) as list_gpus, + patch("runpod.rp_cli.compute.list_cpus", return_value=CPUS) as list_cpus, + patch("runpod.rp_cli.compute.list_datacenters", return_value=DATACENTERS), + patch("runpod.rp_cli.pod.interactive", return_value=interactive), + patch("runpod.rp_cli.output.interactive", return_value=interactive), + ): + result = CliRunner().invoke(cli, ["pod", *args], input=input) + result.api, result.list_gpus, result.list_cpus = api, list_gpus, list_cpus + return result + + +def _create(args, api=None, gpus=GPUS): + result = _invoke(["create", *args], api=api, gpus=gpus) + assert result.exit_code == 0, result.output + return result, result.api.body("POST", "/v2/pods") + + +class TestList: + def test_table(self): + result = _invoke(["list"]) + assert result.exit_code == 0, result.output + assert "pod123" in result.output and "pod456" in result.output + assert "2x NVIDIA H100 80GB HBM3" in result.output + assert "cpu3c 4 vcpu" in result.output + assert "$7.98/hr" in result.output + assert result.api.calls[0][2] == {} + + def test_filters(self): + result = _invoke(["list", "--name", "BO", "--status", "exited", "--json"]) + assert [p["id"] for p in json.loads(result.stdout)] == ["pod456"] + result = _invoke(["list", "--status", "running", "--json"]) + assert [p["id"] for p in json.loads(result.stdout)] == ["pod123"] + + def test_cluster_pods_and_empty(self): + api = FakeApi( + { + ("GET", "/v2/pods"): { + "pods": [], + "pagination": {"hasNextPage": False, "nextCursor": None}, + } + } + ) + result = _invoke(["list", "--include-cluster-pods"], api=api) + assert "no pods" in result.output + assert api.calls[0][2] == {"includeClusterPods": "true"} + + @pytest.mark.parametrize("runtime", [None, {}, {"uptime": 0}]) + @pytest.mark.parametrize("as_json", [False, True]) + def test_filters_use_displayed_status(self, runtime, as_json): + starting = {**POD, "id": "starting-pod", "runtime": runtime} + api = FakeApi( + { + ("GET", "/v2/pods"): { + "pods": [POD, starting], + "pagination": {"hasNextPage": False, "nextCursor": None}, + } + } + ) + for status, expected, excluded in ( + ("starting", "starting-pod", "pod123"), + ("running", "pod123", "starting-pod"), + ): + result = _invoke( + ["list", "--status", status, *(["--json"] if as_json else [])], api=api + ) + assert result.exit_code == 0, result.output + if as_json: + assert [p["id"] for p in json.loads(result.stdout)] == [expected] + else: + assert expected in result.stdout + assert excluded not in result.stdout + + def test_api_error_is_clean(self): + api = FakeApi({("GET", "/v2/pods"): error.AuthenticationError("Unauthorized")}) + result = _invoke(["list"], api=api) + assert result.exit_code == 1 + assert "Error: Unauthorized" in result.output + + +class TestGet: + def test_fields_and_connection_info(self): + result = _invoke(["get", "pod123"]) + assert result.exit_code == 0, result.output + output = result.output + assert "ssh pod123-abc@ssh.runpod.io" in output + assert "ssh root@1.2.3.4 -p 40022" in output + assert "https://pod123-8888.proxy.runpod.net" in output + assert "1.2.3.4:40022 -> 22" in output + assert "network vol1 at /data" in output + assert "pod123.runpod.internal" in output + assert "32 vcpu, 250GB ram" in output + assert "MODEL, PRIVATE_SETTING" in output + assert "redacted-test-value" not in output + assert "1h 2m" in output + + @pytest.mark.parametrize( + "runtime, status, uptime", + [(None, "starting", None), ({"uptime": 0}, "starting", None)] + + [({"uptime": 65}, "running", "1m 5s"), ({"uptime": 9}, "running", "9s")], + ) + def test_status_follows_container(self, runtime, status, uptime): + api = FakeApi({("GET", "/v2/pods/pod123"): {**POD, "runtime": runtime}}) + result = _invoke(["get", "pod123"], api=api) + assert f"status {status}" in result.output + if uptime: + assert f"uptime {uptime}" in result.output + + def test_other_shapes(self): + pod = { + **STOPPED, + "mounts": { + "persistent": {"size": 20, "path": "/workspace"}, + "global": [{"volumeId": "gv1"}], + }, + "globalNetworking": {"enabled": False}, + } + api = FakeApi({("GET", "/v2/pods/pod456"): pod}) + result = _invoke(["get", "pod456"], api=api) + assert "persistent 20GB at /workspace" in result.output + assert "global gv1 at -" in result.output + assert "cpu3c 4 vcpu" in result.output + + def test_json(self): + result = _invoke(["get", "pod123", "--json"]) + assert json.loads(result.stdout) == POD + + def test_missing_id_without_tty(self): + result = _invoke(["get"]) + assert result.exit_code == 2 + assert "pass a POD_ID" in result.output + + def test_picker(self): + with patch("InquirerPy.inquirer.fuzzy") as fuzzy: + fuzzy.return_value.execute.return_value = "pod123" + result = _invoke(["get"], interactive=True) + assert result.exit_code == 0, result.output + choices = fuzzy.call_args.kwargs["choices"] + assert [choice.value for choice in choices] == ["pod123", "pod456"] + assert choices[0].name == "trainer running" + + +class TestCreateGpu: + def test_ticket_example(self): + result, body = _create( + [ + "trainer", + "--gpu", + "h100:2", + "--image", + "runpod/pytorch:latest", + "--region", + "eu", + "--ssh", + ] + ) + assert body == { + "name": "trainer", + "image": "runpod/pytorch:latest", + "dataCenterIds": ["EU-RO-1", "EU-CZ-1"], + "gpu": {"id": "NVIDIA H100 80GB HBM3", "count": 2}, + "startSsh": True, + } + assert "'h100' matches 3 gpu types; using H100 SXM" in result.output + assert "low stock, $7.98/hr" in result.output + assert "no ssh keys" not in result.output + assert "created pod trainer (pod123)" in result.output + assert result.list_gpus.call_args.kwargs == { + "product": "POD", + "count": 2, + "cloud": "SECURE", + "cuda_versions": [], + "min_cuda_version": None, + } + + @pytest.mark.parametrize( + "options, expected", + [ + ([], "NVIDIA H100 NVL"), + (["--gpu-count", "2"], "NVIDIA H100 PCIe"), + (["--data-center", "EU-RO-1"], "NVIDIA H100 80GB HBM3"), + ], + ) + def test_selects_stock_then_price_within_placement(self, options, expected): + _, body = _create(["a", "--gpu", "h100", "--image", "x", *options]) + assert body["gpu"]["id"] == expected + + def test_unique_match_used_directly(self): + result, body = _create(["a", "--gpu", "4090", "--image", "x", "--json"]) + assert body["gpu"] == {"id": "NVIDIA GeForce RTX 4090", "count": 1} + assert "matches" not in result.output + assert "no secure stock for 1x RTX 4090" in result.stderr + assert json.loads(result.stdout)["id"] == "pod123" + + def test_gpu_settings(self): + _, body = _create( + [ + "a", + "--gpu", + "h100 nvl", + "--cloud", + "COMMUNITY", + "--gpu-count", + "4", + "--cuda", + ">=12.4", + "--min-ram-per-gpu", + "32", + "--min-vcpu-per-gpu", + "8", + "--image", + "x", + ] + ) + assert body["cloud"] == "COMMUNITY" + assert body["gpu"] == { + "id": "NVIDIA H100 NVL", + "count": 4, + "minCudaVersion": "12.4", + "minRamPerGpu": 32, + "minVcpuCountPerGpu": 8, + } + + def test_exact_cuda(self): + result, body = _create( + ["a", "--gpu", "4090", "--cuda", "12.4,12.8", "--image", "x"] + ) + assert body["gpu"]["allowedCudaVersions"] == ["12.4", "12.8"] + assert result.list_gpus.call_args.kwargs["cuda_versions"] == ["12.4", "12.8"] + + def test_ssh_without_keys_warns(self): + api = FakeApi({("GET", "/v2/account/ssh-keys"): {"keys": []}}) + result, body = _create(["a", "--gpu", "4090", "--image", "x", "--ssh"], api=api) + assert body["startSsh"] is True + assert "no ssh keys" in result.output + + +class TestCreateCpu: + @pytest.mark.parametrize( + "flavor, options, requested, expected", + [("cpu5g", [], None, 4), ("cpu3c", ["--vcpu", "8"], 8, 8)], + ) + def test_vcpu_resolution(self, flavor, options, requested, expected): + result, body = _create(["box", "--cpu", flavor, "--image", "ubuntu", *options]) + assert body["cpu"] == {"id": flavor, "vcpuCount": expected} + assert result.list_cpus.call_args.kwargs == { + "product": "POD", + "vcpu_count": requested, + } + + +class TestCreateContainer: + def test_every_container_field(self): + _, body = _create( + shlex.split( + """a --gpu 4090 --template tpl1 -e A=1 --env B=x=y + -p 8888/http --port 22/tcp --disk 40 --args 'bash start.sh' + --cmd "python -m train --lr '1e-4'" --entrypoint '/bin/sh -c' + --registry ghcr --global-networking --jupyter --persistent-volume 100""" + ) + ) + assert body == { + "name": "a", + "templateId": "tpl1", + "args": "bash start.sh", + "cmd": ["python", "-m", "train", "--lr", "1e-4"], + "entrypoint": ["/bin/sh", "-c"], + "disk": 40, + "ports": ["8888/http", "22/tcp"], + "registry": "reg1", + "globalNetworking": True, + "env": {"A": "1", "B": "x=y"}, + "mounts": {"persistent": {"size": 100, "path": "/workspace"}}, + "gpu": {"id": "NVIDIA GeForce RTX 4090", "count": 1}, + "startJupyter": True, + } + + def test_network_volume_sets_its_data_center(self): + _, body = _create( + [ + "a", + "--gpu", + "h100", + "--image", + "x", + "--network-volume", + "vol1", + "--global-volume", + "gv1:/g", + ] + ) + assert body["mounts"] == { + "network": [{"volumeId": "vol1", "path": "/workspace"}], + "global": [{"volumeId": "gv1", "path": "/g"}], + } + assert body["dataCenterIds"] == ["EU-RO-1"] + assert body["gpu"]["id"] == "NVIDIA H100 80GB HBM3" + + def test_regions_and_data_centers_merge(self): + _, body = _create( + [ + "a", + "--gpu", + "4090", + "--image", + "x", + "--region", + "eu-ro", + "--data-center", + "EU-RO-1", + "--data-center", + "US-TX-3", + "--network-volume", + "vol1:/data", + ] + ) + assert body["dataCenterIds"] == ["EU-RO-1", "US-TX-3"] + assert body["mounts"]["network"] == [{"volumeId": "vol1", "path": "/data"}] + + +class TestCreateValidation: + @pytest.mark.parametrize("option", ["--cmd", "--entrypoint"]) + @pytest.mark.parametrize("value", ["python 'unclosed", 'python "unclosed']) + def test_malformed_command_quoting(self, option, value): + result = _invoke( + ["create", "a", "--gpu", "4090", "--image", "x", option, value] + ) + assert result.exit_code == 2 + assert f"Invalid value for {option}" in result.output + assert not result.api.made("POST", "/v2/pods") + + @pytest.mark.parametrize( + "args, message", + [ + (["a", "--gpu", "h100:2", "--gpu-count", "3", "--image", "x"], "conflicts"), + (["a", "--gpu", "4090", "--image", "x", "--env", "NOPE"], "use KEY=VALUE"), + ( + ["a", "--gpu", "4090", "--image", "x", "--persistent-volume", "big"], + "size is in GB", + ), + ( + ["a", "--gpu", "4090", "--image", "x", "--network-volume", "vol1:"], + "VALUE[:PATH]", + ), + ( + ["a", "--gpu", "4090", "--cpu", "cpu3c", "--image", "x"], + "either --gpu or --cpu", + ), + (["a", "--image", "x"], "pass --gpu"), + ( + ["a", "--cpu", "cpu3c", "--cuda", "12.4", "--image", "x"], + "apply to gpu pods", + ), + (["a", "--gpu", "4090", "--vcpu", "4", "--image", "x"], "--vcpu applies"), + (["a", "--gpu", "4090"], "pass --image or --template"), + (["--gpu", "4090", "--image", "x"], "pass a NAME"), + ], + ) + def test_usage_errors(self, args, message): + result = _invoke(["create", *args]) + assert result.exit_code == 2 + assert message in result.output + assert not result.api.made("POST", "/v2/pods") + + @pytest.mark.parametrize( + "options, message", + [ + (["--gpu", "h100:9"], "offers 9 per pod on secure cloud"), + (["--gpu", "b9000"], "unknown gpu 'b9000'"), + ( + ["--gpu", "4090", "--registry", "nope"], + "unknown registry credential 'nope'", + ), + ], + ) + def test_resolution_errors(self, options, message): + result = _invoke(["create", "a", "--image", "x", *options]) + assert result.exit_code == 1 + assert message in result.output + assert not result.api.made("POST", "/v2/pods") + + def test_prompts_for_name(self): + result = _invoke( + ["create", "--gpu", "4090", "--image", "x"], + interactive=True, + input="mine\n", + ) + assert result.exit_code == 0, result.output + assert result.api.body("POST", "/v2/pods")["name"] == "mine" + + +class TestWait: + def test_create_waits_until_running(self): + api = FakeApi( + { + ("GET", "/v2/pods/pod123"): [ + {**POD, "status": "PROVISIONING"}, + {**POD, "status": "STARTING"}, + {**POD, "runtime": None}, + {**POD, "runtime": {"uptime": 0}}, + {**POD, "runtime": {"uptime": 1}}, + ] + } + ) + result = _invoke( + ["create", "a", "--gpu", "4090", "--image", "x", "--wait", "--json"], + api=api, + ) + assert result.exit_code == 0, result.output + assert json.loads(result.stdout)["status"] == "RUNNING" + assert sum(call[:2] == ("GET", "/v2/pods/pod123") for call in api.calls) == 5 + + def test_error_status_fails(self): + api = FakeApi({("GET", "/v2/pods/pod456"): {**STOPPED, "status": "ERROR"}}) + result = _invoke(["start", "pod456", "--wait"], api=api) + assert result.exit_code == 1 + assert "pod pod456 is error" in result.output + + def test_timeout(self, monkeypatch): + monkeypatch.setattr(pod_module, "WAIT_TIMEOUT", -1) + api = FakeApi({("GET", "/v2/pods/pod456"): STOPPED}) + result = _invoke(["start", "pod456", "--wait"], api=api) + assert result.exit_code == 1 + assert "still exited" in result.output + + +class TestUpdate: + @pytest.mark.parametrize("option", ["--cmd", "--entrypoint"]) + def test_malformed_command_quoting(self, option): + result = _invoke(["update", "pod123", option, "python 'unclosed"]) + assert result.exit_code == 2 + assert f"Invalid value for {option}" in result.output + assert not result.api.made("PATCH", "/v2/pods/pod123") + + def test_nothing_to_update(self): + result = _invoke(["update", "pod123"]) + assert result.exit_code == 2 + assert "nothing to update" in result.output + + def test_simple_fields_skip_the_read(self): + result = _invoke( + shlex.split("""update pod123 --name renamed --lock --no-global-networking + --image img:2 --registry '' --port 8000/http --disk 80 --cmd ''""") + ) + assert result.exit_code == 0, result.output + assert result.api.body("PATCH", "/v2/pods/pod123") == { + "image": "img:2", + "cmd": [], + "disk": 80, + "ports": ["8000/http"], + "registry": None, + "globalNetworking": False, + "name": "renamed", + "locked": True, + } + assert not result.api.made("GET", "/v2/pods/pod123") + assert "updated pod pod123" in result.output + + def test_env_edits_existing(self): + result = _invoke( + [ + "update", + "pod123", + "-e", + "MODEL=qwen", + "-e", + "NEW=1", + "--unset-env", + "PRIVATE_SETTING", + "--unlock", + ] + ) + assert result.exit_code == 0, result.output + assert result.api.body("PATCH", "/v2/pods/pod123") == { + "locked": False, + "env": {"MODEL": "qwen", "NEW": "1"}, + } + + def test_mount_path_defaults_to_current(self): + result = _invoke(["update", "pod123", "--network-volume", "vol1", "--json"]) + assert result.exit_code == 0, result.output + assert result.api.body("PATCH", "/v2/pods/pod123") == { + "mounts": {"network": [{"volumeId": "vol1", "path": "/data"}]} + } + assert json.loads(result.stdout)["id"] == "pod123" + + +class TestActions: + @pytest.mark.parametrize("action", ["stop", "restart"]) + def test_action(self, action): + result = _invoke([action, "pod123"]) + assert result.exit_code == 0, result.output + assert result.api.body("POST", "/v2/pods/pod123/action") == {"action": action} + + def test_start_wait(self): + api = FakeApi( + { + ("GET", "/v2/pods/pod456"): [ + {**STOPPED, "status": "STARTING"}, + {**STOPPED, "status": "RUNNING", "runtime": {"uptime": 46}}, + {**STOPPED, "status": "RUNNING", "runtime": {"uptime": 1}}, + ] + } + ) + result = _invoke(["start", "pod456", "--wait", "--json"], api=api) + assert result.exit_code == 0, result.output + assert json.loads(result.stdout)["runtime"] == {"uptime": 1} + assert "started pod pod456" in result.stderr + + def test_restart_wait_rejects_recent_pre_action_runtime(self): + previous = {**POD, "runtime": {"uptime": 1}} + restarted = {**previous, "startedAt": "2026-10-19T20:02:00Z"} + api = FakeApi( + { + ("GET", "/v2/pods/pod123"): [ + previous, + previous, + {**previous, "runtime": None}, + restarted, + ] + } + ) + result = _invoke(["restart", "pod123", "--wait", "--json"], api=api) + assert result.exit_code == 0, result.output + assert json.loads(result.stdout) == restarted + assert [call[:2] for call in api.calls] == [ + ("GET", "/v2/pods/pod123"), + ("POST", "/v2/pods/pod123/action"), + ("GET", "/v2/pods/pod123"), + ("GET", "/v2/pods/pod123"), + ("GET", "/v2/pods/pod123"), + ] + + @pytest.mark.parametrize("started_at", [None, POD["startedAt"]]) + def test_restart_wait_requires_new_start_timestamp(self, monkeypatch, started_at): + monkeypatch.setattr(pod_module, "WAIT_TIMEOUT", -1) + stale = {**POD, "runtime": {"uptime": 1}, "startedAt": started_at} + api = FakeApi({("GET", "/v2/pods/pod123"): stale}) + result = _invoke(["restart", "pod123", "--wait"], api=api) + assert result.exit_code == 1 + assert "still running" in result.output + + def test_restart_wait_rejects_stale_runtime_with_new_timestamp(self): + restarted = {**POD, "startedAt": "2026-10-19T20:02:00Z"} + ready = {**restarted, "runtime": {"uptime": 1}} + api = FakeApi({("GET", "/v2/pods/pod123"): [POD, restarted, ready]}) + result = _invoke(["restart", "pod123", "--wait", "--json"], api=api) + assert result.exit_code == 0, result.output + assert json.loads(result.stdout) == ready + assert sum(call[:2] == ("GET", "/v2/pods/pod123") for call in api.calls) == 3 + + def test_restart_wait_rejects_recent_runtime_with_new_timestamp(self, monkeypatch): + monkeypatch.setattr(pod_module, "WAIT_TIMEOUT", -1) + previous = {**POD, "runtime": {"uptime": 1}} + restarted = {**previous, "startedAt": "2026-10-19T20:02:00Z"} + api = FakeApi({("GET", "/v2/pods/pod123"): [previous, restarted]}) + result = _invoke(["restart", "pod123", "--wait"], api=api) + assert result.exit_code == 1 + assert "still running" in result.output + + def test_picker_offers_pods_that_allow_the_action(self): + with patch("InquirerPy.inquirer.fuzzy") as fuzzy: + fuzzy.return_value.execute.return_value = "pod456" + result = _invoke(["start"], interactive=True) + assert result.exit_code == 0, result.output + assert [c.value for c in fuzzy.call_args.kwargs["choices"]] == ["pod456"] + + def test_api_error(self): + api = FakeApi( + { + ("POST", "/v2/pods/pod123/action"): error.QueryError( + "pod is locked", "POST" + ) + } + ) + result = _invoke(["stop", "pod123"], api=api) + assert result.exit_code == 1 + assert "pod is locked" in result.output + + +class TestTerminate: + @pytest.mark.parametrize("command", ["terminate", "delete"]) + def test_with_yes(self, command): + result = _invoke([command, "pod123", "--yes"]) + assert result.exit_code == 0, result.output + assert result.api.made("DELETE", "/v2/pods/pod123") + assert "terminated pod pod123" in result.output + + def test_needs_confirmation(self): + result = _invoke(["terminate", "pod123"]) + assert result.exit_code == 2 + assert "--yes" in result.output + assert not result.api.made("DELETE", "/v2/pods/pod123") + + def test_interactive_confirm(self): + result = _invoke(["terminate", "pod123"], interactive=True, input="n\n") + assert result.exit_code == 1 + assert not result.api.made("DELETE", "/v2/pods/pod123") + + +class TestConnect: + def test_launches_terminal(self): + ssh = MagicMock() + with patch( + "runpod.cli.utils.ssh_cmd.SSHConnection", return_value=ssh + ) as connection: + result = _invoke(["connect", "pod123"]) + assert result.exit_code == 0, result.output + connection.assert_called_once_with("pod123") + ssh.launch_terminal.assert_called_once() + + def test_failure_is_clean(self): + with patch( + "runpod.cli.utils.ssh_cmd.SSHConnection", side_effect=TimeoutError("no ssh") + ): + result = _invoke(["connect", "pod123"]) + assert result.exit_code == 1 + assert "Error: no ssh" in result.output + + def test_sync_is_mounted(self): + result = _invoke(["sync", "--help"]) + assert result.exit_code == 0 + + +@pytest.mark.parametrize( + "args, expected", + [ + ([], dict(follow=False, tail=100, source=None, since=None, as_json=False)), + ( + [ + "--follow", + "--tail", + "5", + "--source", + "container", + "--since", + "2026-10-19T00:00:00Z", + "--json", + ], + dict( + follow=True, + tail=5, + source="container", + since="2026-10-19T00:00:00Z", + as_json=True, + ), + ), + ], +) +def test_logs_forward_options(args, expected): + with patch("runpod.rp_cli.logs.show_logs") as show_logs: + result = _invoke(["logs", "pod123", *args]) + assert result.exit_code == 0, result.output + show_logs.assert_called_once_with("pod123", **expected) + + +def test_logs_reject_tail_above_api_limit(): + with patch("runpod.rp_cli.logs.show_logs") as show_logs: + result = _invoke(["logs", "pod123", "--tail", "5001"]) + assert result.exit_code == 2 + assert "5000" in result.output + show_logs.assert_not_called() + + +class TestCompletion: + def test_completers(self): + api = FakeApi() + with ( + patch("runpod.rp_cli.pod.run_rest_request", api), + patch("runpod.rp_cli.api.run_rest_request", api), + patch("runpod.rp_cli.compute.list_datacenters", return_value=DATACENTERS), + ): + pods = pod_module._complete_pods(None, None, "pod1") + datacenters = pod_module._complete_datacenters(None, None, "EU") + volumes = pod_module._complete_network_volumes(None, None, "") + registries = pod_module._complete_registries(None, None, "g") + assert [(item.value, item.help) for item in pods] == [ + ("pod123", "trainer running") + ] + assert [item.value for item in datacenters] == ["EU-RO-1", "EU-CZ-1"] + assert [(item.value, item.help) for item in volumes] == [ + ("vol1", "data EU-RO-1") + ] + assert [item.value for item in registries] == ["ghcr"] + + +def test_sync_checks_account_keys_over_rest_only(): + source, destination = MagicMock(), MagicMock() + source.__enter__.return_value = source + destination.__enter__.return_value = destination + source.ssh.exec_command.side_effect = [ + (None, BytesIO(value), None) for value in (b"1", b"exists", b"created", b"1K") + ] + destination.ssh.exec_command.return_value = (None, BytesIO(b"1"), None) + with ( + patch( + "runpod.cli.groups.ssh.functions.run_rest_request", + return_value={"keys": ["ssh-ed25519 dGVzdC1rZXk= test-fixture"]}, + ) as rest, + patch( + "runpod.api.ctl_commands.run_graphql_query", + side_effect=AssertionError("GraphQL forbidden"), + ) as graphql, + patch( + "runpod.cli.utils.ssh_cmd.SSHConnection", + side_effect=[source, destination], + ), + ): + result = _invoke(["sync", "source-pod", "destination-pod"]) + assert result.exit_code == 0, result.output + assert "Found 1 SSH key(s)" in result.output + assert "Sync completed successfully" in result.output + rest.assert_called_once_with("GET", "/v2/account/ssh-keys") + graphql.assert_not_called() + source.get_file.assert_called_once() + destination.put_file.assert_called_once()