feat: add rp pod lifecycle commands

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:

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
 

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

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)
 
 

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