From 8ceb0bd4762467a9512ffb4a1bea3e3b858e1c88 Mon Sep 17 00:00:00 2001 From: Muhib Waqar Date: Fri, 25 Sep 2026 16:56:10 -0400 Subject: [PATCH] feat: add pod and serverless worker log readers Read the REST v2 SSE log streams (GET /v2/pods/{id}/logs and GET /v2/serverless/{id}/workers/{workerId}/logs): - get_pod_logs / get_endpoint_worker_logs return a bounded snapshot (tail backfill plus up to max_wait seconds of live output, capped at max_bytes with the oldest lines dropped). - iter_pod_logs / iter_endpoint_worker_logs follow the stream, resuming from Last-Event-ID on close or idle, and waiting out 429s on reconnect. - get_endpoint_workers lists worker IDs for the worker log readers. - QueryError.retry_after exposes the Retry-After header. Refs #400 Co-Authored-By: Claude Opus 5.5 (1M context) --- README.md | 32 ++++ docs/api/handling_errors.md | 15 ++ runpod/__init__.py | 10 ++ runpod/api/ctl_commands.py | 210 ++++++++++++++++++++++- runpod/api/rest.py | 91 +++++++++- runpod/error.py | 3 + tests/test_api/test_ctl_commands.py | 252 ++++++++++++++++++++++++++++ tests/test_api/test_logs_server.py | 228 +++++++++++++++++++++++++ tests/test_api/test_rest.py | 122 +++++++++++++- tests/test_init.py | 6 +- 10 files changed, 962 insertions(+), 7 deletions(-) create mode 100644 tests/test_api/test_logs_server.py diff --git a/README.md b/README.md index a372aa8b..ad007b82 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,7 @@ Welcome to the official Python library for Runpod API & SDK. - [📚 | REST API v2 Wrapper](#--rest-api-v2-wrapper) - [Endpoints](#endpoints) - [GPU Cloud (Pods)](#gpu-cloud-pods) + - [Logs](#logs) - [📁 | Directory](#--directory) - [🤝 | Community and Contributing](#--community-and-contributing) @@ -305,6 +306,37 @@ runpod.resume_pod(pod["id"]) runpod.terminate_pod(pod["id"]) ``` +### Logs + +Pod and Serverless worker logs are read from the REST v2 log streams. Each entry +is a dict with `id`, `ts`, `source` (`container` or `system`) and `line`. + +```python +import runpod + +# snapshot: backfill the last 200 container lines, then read live output for 5s +logs = runpod.get_pod_logs(pod["id"], tail=200, source="container", max_wait=5) +print("\n".join(entry["line"] for entry in logs)) + +# follow: yield lines as they arrive until you stop iterating +for entry in runpod.iter_pod_logs(pod["id"], tail=0): + print(entry["ts"], entry["line"]) + +# Serverless workers +workers = runpod.get_endpoint_workers("ENDPOINT_ID") +logs = runpod.get_endpoint_worker_logs("ENDPOINT_ID", workers[0]["id"]) +``` + +- `tail` backfills 0–5000 historical lines (API default 100) and is ignored when + `since` is set. `since` takes an RFC3339 string or a timezone-aware `datetime`. +- `get_*_logs` returns once `max_wait` seconds pass, or once the stream has been + idle that long. Past `max_bytes` of log text (default 4 MiB) the oldest lines + are dropped, so the newest output is always kept. +- `iter_*_logs` reconnects from the last event ID when the stream closes or goes + idle, so lines are neither repeated nor skipped. It waits out `429` responses + on reconnect. Errors on the first connection raise. Pass `max_wait` to stop + after that many seconds. + ### Template and placement options - With `create_pod(template_id=...)`, omitting `docker_args` inherits the template's diff --git a/docs/api/handling_errors.md b/docs/api/handling_errors.md index e8b1a28b..5fca1af5 100644 --- a/docs/api/handling_errors.md +++ b/docs/api/handling_errors.md @@ -26,3 +26,18 @@ except runpod.error.QueryError as err: print(err.query) print(err.errors) ``` + +Rate-limited responses (`429`) also set `retry_after`, the seconds to wait from +the `Retry-After` header, or `None` when the header is missing: + +```python +import time + +try: + pods = runpod.get_pods() +except runpod.error.QueryError as err: + if err.status_code != 429: + raise + time.sleep(err.retry_after or 1) + pods = runpod.get_pods() +``` diff --git a/runpod/__init__.py b/runpod/__init__.py index d1ffb33a..2be0b019 100644 --- a/runpod/__init__.py +++ b/runpod/__init__.py @@ -14,12 +14,17 @@ create_pod, create_template, delete_container_registry_auth, + get_endpoint_worker_logs, + get_endpoint_workers, get_endpoints, get_gpu, get_gpus, get_pod, + get_pod_logs, get_pods, get_user, + iter_endpoint_worker_logs, + iter_pod_logs, resume_pod, stop_pod, terminate_pod, @@ -43,12 +48,17 @@ "create_pod", "create_template", "delete_container_registry_auth", + "get_endpoint_worker_logs", + "get_endpoint_workers", "get_endpoints", "get_gpu", "get_gpus", "get_pod", + "get_pod_logs", "get_pods", "get_user", + "iter_endpoint_worker_logs", + "iter_pod_logs", "resume_pod", "stop_pod", "terminate_pod", diff --git a/runpod/api/ctl_commands.py b/runpod/api/ctl_commands.py index 1291df49..664b6032 100644 --- a/runpod/api/ctl_commands.py +++ b/runpod/api/ctl_commands.py @@ -2,16 +2,33 @@ # pylint: disable=too-many-arguments,too-many-locals +import json +import math import re -from typing import Any, Iterable, Optional +import time +from collections import deque +from datetime import datetime +from typing import Any, Iterable, Iterator, Optional, Union from urllib.parse import quote +import requests + from runpod import error from .graphql import run_graphql_query from .mutations import container_register_auth as container_register_auth_mutations from .queries import user as user_queries -from .rest import HTTP_STATUS_NOT_FOUND, run_rest_request +from .rest import ( + HTTP_STATUS_NOT_FOUND, + HTTP_STATUS_TOO_MANY_REQUESTS, + read_event_stream, + run_rest_request, +) + +LOG_SOURCES = ("container", "system") +LOG_MAX_TAIL = 5000 +LOG_MAX_BYTES = 4 * 1024 * 1024 +LOG_RECONNECT_DELAY = 1 def _path_segment(value: str) -> str: @@ -109,6 +126,185 @@ def get_pod(pod_id: str, api_key: Optional[str] = None) -> Optional[dict]: raise +def _log_params( + tail: Optional[int], + since: Optional[Union[str, datetime]], + source: Optional[str], +) -> dict[str, Any]: + if tail is not None and not 0 <= tail <= LOG_MAX_TAIL: + raise ValueError(f"tail must be between 0 and {LOG_MAX_TAIL}") + if source is not None and source not in LOG_SOURCES: + raise ValueError(f"source must be one of {LOG_SOURCES} or None") + if isinstance(since, datetime): + if since.utcoffset() is None: + raise ValueError("since must be timezone-aware") + since = since.isoformat() + params = {"tail": tail, "since": since, "source": source} + return {key: value for key, value in params.items() if value is not None} + + +def _check_max_wait(max_wait: Optional[float], required: bool) -> None: + if max_wait is None and not required: + return + if max_wait is None or max_wait <= 0: + raise ValueError("max_wait must be positive") + + +def _iter_logs( + path: str, + params: dict[str, Any], + max_wait: Optional[float], + follow: bool, + api_key: Optional[str], +) -> Iterator[dict]: + """Yield log entries from an SSE log endpoint. + + With `follow`, a closed or idle stream is reopened from the last event ID + so no line is repeated or skipped. Errors on the first connection raise; + on a reconnect, network errors and 429s are retried. + """ + deadline = math.inf if max_wait is None else time.monotonic() + max_wait + last_event_id = None + connected = False + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + return + delay = LOG_RECONNECT_DELAY + try: + for event in read_event_stream( + path, + api_key=api_key, + params=params, + max_wait=None if max_wait is None else remaining, + last_event_id=last_event_id, + ): + connected = True + entry = json.loads(event["data"]) + if "id" in event: + entry["id"] = last_event_id = event["id"] + yield entry + connected = True + except (requests.exceptions.RequestException, error.QueryError) as exc: + retryable = not isinstance(exc, error.QueryError) or ( + exc.status_code == HTTP_STATUS_TOO_MANY_REQUESTS + ) + if not (follow and connected and retryable): + raise + delay = getattr(exc, "retry_after", None) or delay + if not follow: + return + time.sleep(max(0, min(delay, deadline - time.monotonic()))) + + +def _snapshot_logs( + path: str, + params: dict[str, Any], + max_wait: float, + max_bytes: int, + api_key: Optional[str], +) -> list[dict]: + logs: deque[dict] = deque() + size = 0 + for entry in _iter_logs(path, params, max_wait, follow=False, api_key=api_key): + logs.append(entry) + size += len(entry.get("line", "")) + while size > max_bytes and len(logs) > 1: + size -= len(logs.popleft().get("line", "")) + return list(logs) + + +def _pod_logs_path(pod_id: str) -> str: + return f"/v2/pods/{_path_segment(pod_id)}/logs" + + +def _worker_logs_path(endpoint_id: str, worker_id: str) -> str: + return ( + f"/v2/serverless/{_path_segment(endpoint_id)}" + f"/workers/{_path_segment(worker_id)}/logs" + ) + + +def get_pod_logs( + pod_id: str, + tail: Optional[int] = None, + since: Optional[Union[str, datetime]] = None, + source: Optional[str] = None, + max_wait: float = 5, + max_bytes: int = LOG_MAX_BYTES, + api_key: Optional[str] = None, +) -> list[dict]: + """Get a snapshot of a pod's logs. + + Reads the live log stream for up to `max_wait` seconds and returns the + lines received, oldest first, as dicts with `id`, `ts`, `source` and `line`. + `tail` backfills that many historical lines (API default 100, max 5000) and + is ignored when `since` is set. `source` is "container", "system", or None + for both. Past `max_bytes` of log text, the oldest lines are dropped. + """ + params = _log_params(tail, since, source) + _check_max_wait(max_wait, required=True) + return _snapshot_logs(_pod_logs_path(pod_id), params, max_wait, max_bytes, api_key) + + +def iter_pod_logs( + pod_id: str, + tail: Optional[int] = None, + since: Optional[Union[str, datetime]] = None, + source: Optional[str] = None, + max_wait: Optional[float] = None, + api_key: Optional[str] = None, +) -> Iterator[dict]: + """Follow a pod's logs, yielding entries as they arrive. + + Takes the same `tail`, `since` and `source` as `get_pod_logs`. Dropped or + idle connections are resumed from the last event ID. Runs until the caller + stops iterating, or for `max_wait` seconds when it is set. + """ + params = _log_params(tail, since, source) + _check_max_wait(max_wait, required=False) + return _iter_logs(_pod_logs_path(pod_id), params, max_wait, True, api_key) + + +def get_endpoint_worker_logs( + endpoint_id: str, + worker_id: str, + tail: Optional[int] = None, + since: Optional[Union[str, datetime]] = None, + source: Optional[str] = None, + max_wait: float = 5, + max_bytes: int = LOG_MAX_BYTES, + api_key: Optional[str] = None, +) -> list[dict]: + """Get a snapshot of a Serverless worker's logs. + + Behaves like `get_pod_logs`. A crash-looping worker can still report as + running, so its logs are the reliable signal when jobs stay in queue. + """ + params = _log_params(tail, since, source) + _check_max_wait(max_wait, required=True) + return _snapshot_logs( + _worker_logs_path(endpoint_id, worker_id), params, max_wait, max_bytes, api_key + ) + + +def iter_endpoint_worker_logs( + endpoint_id: str, + worker_id: str, + tail: Optional[int] = None, + since: Optional[Union[str, datetime]] = None, + source: Optional[str] = None, + max_wait: Optional[float] = None, + api_key: Optional[str] = None, +) -> Iterator[dict]: + """Follow a Serverless worker's logs. Behaves like `iter_pod_logs`.""" + params = _log_params(tail, since, source) + _check_max_wait(max_wait, required=False) + return _iter_logs( + _worker_logs_path(endpoint_id, worker_id), params, max_wait, True, api_key + ) + + def create_pod( name: str, image_name: Optional[str] = "", @@ -283,6 +479,16 @@ def get_endpoints() -> list[dict]: return response["endpoints"] +def get_endpoint_workers(endpoint_id: str, api_key: Optional[str] = None) -> list[dict]: + """Get the active workers of a Serverless endpoint.""" + response = run_rest_request( + "GET", + f"/v2/serverless/{_path_segment(endpoint_id)}/workers", + api_key=api_key, + ) + return response["workers"] + + def create_endpoint( name: str, template_id: str, diff --git a/runpod/api/rest.py b/runpod/api/rest.py index 987eb143..31636e2a 100644 --- a/runpod/api/rest.py +++ b/runpod/api/rest.py @@ -1,7 +1,9 @@ """Runpod REST API transport.""" +import math import os -from typing import Any, Mapping, Optional +import time +from typing import Any, Iterator, Mapping, Optional import requests @@ -12,6 +14,8 @@ HTTP_STATUS_BAD_REQUEST = 400 HTTP_STATUS_UNAUTHORIZED = 401 HTTP_STATUS_NOT_FOUND = 404 +HTTP_STATUS_TOO_MANY_REQUESTS = 429 +STREAM_IDLE_TIMEOUT = 60 def _resolve_api_key(api_key: Optional[str]) -> str: @@ -45,6 +49,13 @@ def _response_json(response: requests.Response) -> dict[str, Any]: return payload if isinstance(payload, dict) else {} +def _retry_after(response: requests.Response) -> Optional[float]: + try: + return float(response.headers["Retry-After"]) + except (KeyError, TypeError, ValueError): + return None + + def _raise_for_error( response: requests.Response, method: str, path: str ) -> None: @@ -66,6 +77,7 @@ def _raise_for_error( f"{method.upper()} {path}", status_code=response.status_code, errors=payload.get("errors"), + retry_after=_retry_after(response), ) @@ -92,3 +104,80 @@ def run_rest_request( if response.status_code == HTTP_STATUS_NO_CONTENT or not response.content: return None return response.json() + + +def _parse_event_stream(chunks: Iterator[bytes]) -> Iterator[dict[str, str]]: + """Parse text/event-stream bytes into events with `id`, `event` and `data`. + + Only events terminated by a blank line are yielded, so an event cut off + mid-frame when the stream is closed early is discarded. + """ + buffer = b"" + event: dict[str, str] = {} + data: list[str] = [] + for chunk in chunks: + buffer += chunk + *lines, buffer = buffer.replace(b"\r\n", b"\n").split(b"\n") + for raw_line in lines: + line = raw_line.decode("utf-8") + if not line: + if data: + event["data"] = "\n".join(data) + yield event + event, data = {}, [] + continue + if line.startswith(":"): + continue + field, _, value = line.partition(":") + value = value[1:] if value.startswith(" ") else value + if field == "data": + data.append(value) + elif field in ("id", "event"): + event[field] = value + + +def read_event_stream( + path: str, + *, + api_key: Optional[str] = None, + params: Optional[Mapping[str, Any]] = None, + max_wait: Optional[float] = 5, + last_event_id: Optional[str] = None, +) -> Iterator[dict[str, str]]: + """Read events from a Runpod REST SSE endpoint for up to `max_wait` seconds. + + The endpoints behind this hold the connection open to tail live output, so + the read ends at the deadline, or once the stream has been idle for + `max_wait` seconds (`STREAM_IDLE_TIMEOUT` when `max_wait` is None, which + sets no deadline). A timeout before the response headers arrive raises. + `last_event_id` resumes the stream after that event. + """ + deadline = math.inf if max_wait is None else time.monotonic() + max_wait + headers = _build_headers(_resolve_api_key(api_key)) + headers["Accept"] = "text/event-stream" + if last_event_id is not None: + headers["Last-Event-ID"] = last_event_id + + with requests.get( + _build_url(path), + headers=headers, + params=params, + stream=True, + timeout=(30, STREAM_IDLE_TIMEOUT if max_wait is None else max_wait), + ) as response: + _raise_for_error(response, "GET", path) + + def chunks() -> Iterator[bytes]: + try: + for chunk in response.iter_content(chunk_size=None): + yield chunk + if time.monotonic() >= deadline: + return + except ( + requests.exceptions.ConnectionError, + requests.exceptions.ChunkedEncodingError, + ): + # An idle-read timeout or a dropped connection ends the snapshot. + return + + yield from _parse_event_stream(chunks()) diff --git a/runpod/error.py b/runpod/error.py index 35a1a8f3..c9c75b70 100644 --- a/runpod/error.py +++ b/runpod/error.py @@ -39,8 +39,11 @@ def __init__( query: Optional[str] = None, status_code: Optional[int] = None, errors: Optional[list[str]] = None, + retry_after: Optional[float] = None, ): super().__init__(message) self.query = query self.status_code = status_code self.errors = errors or [] + # Seconds to wait before retrying, from a 429's Retry-After header. + self.retry_after = retry_after diff --git a/tests/test_api/test_ctl_commands.py b/tests/test_api/test_ctl_commands.py index d3134abc..75d540b2 100644 --- a/tests/test_api/test_ctl_commands.py +++ b/tests/test_api/test_ctl_commands.py @@ -1,10 +1,12 @@ """Tests for the API wrapper commands.""" from copy import deepcopy +from datetime import datetime, timezone from urllib.parse import unquote from unittest.mock import patch import pytest +import requests from runpod.api import ctl_commands from runpod.error import QueryError @@ -125,6 +127,256 @@ def test_get_pod_propagates_other_api_errors(): ctl_commands.get_pod("pod") +def _log_event(event_id, line, source="container"): + return { + "id": event_id, + "data": f'{{"ts": "2026-06-01T12:00:00Z", "source": "{source}", "line": "{line}"}}', + } + + +def _entry(event_id, line, source="container"): + return {"id": event_id, "ts": "2026-06-01T12:00:00Z", "source": source, "line": line} + + +def _streams(*batches): + """Return a read_event_stream side effect that serves one batch per call. + + A batch is a list of events, optionally ending in an exception to raise. + """ + batches = list(batches) + + def stream(*_args, **_kwargs): + for item in batches.pop(0): + if isinstance(item, Exception): + raise item + yield item + + return stream + + +def test_get_pod_logs_returns_entries_with_ids(): + events = [_log_event("1", "starting"), _log_event("2", "ready", "system")] + with patch( + "runpod.api.ctl_commands.read_event_stream", side_effect=_streams(events) + ) as stream: + logs = ctl_commands.get_pod_logs( + "pod/id", + tail=50, + since=datetime(2026, 6, 1, tzinfo=timezone.utc), + max_wait=2, + api_key="key", + ) + + assert logs == [_entry("1", "starting"), _entry("2", "ready", "system")] + stream.assert_called_once() + assert stream.call_args.args == ("/v2/pods/pod%2Fid/logs",) + kwargs = stream.call_args.kwargs + assert kwargs["api_key"] == "key" + assert kwargs["params"] == {"tail": 50, "since": "2026-06-01T00:00:00+00:00"} + assert kwargs["last_event_id"] is None + assert 0 < kwargs["max_wait"] <= 2 + + +def test_get_pod_logs_omits_unset_params(): + with patch( + "runpod.api.ctl_commands.read_event_stream", side_effect=_streams([]) + ) as stream: + assert ctl_commands.get_pod_logs("pod", source="system") == [] + + assert stream.call_args.kwargs["params"] == {"source": "system"} + + +def test_get_pod_logs_drops_oldest_lines_past_byte_cap(): + events = [_log_event(str(i), "x" * 10) for i in range(5)] + with patch("runpod.api.ctl_commands.read_event_stream", side_effect=_streams(events)): + logs = ctl_commands.get_pod_logs("pod", max_bytes=25) + + assert [entry["id"] for entry in logs] == ["3", "4"] + + +def test_get_pod_logs_does_not_retry(): + with ( + patch( + "runpod.api.ctl_commands.read_event_stream", + side_effect=_streams([QueryError("slow down", status_code=429)]), + ), + pytest.raises(QueryError), + ): + ctl_commands.get_pod_logs("pod") + + +def test_get_endpoint_worker_logs_uses_worker_path(): + with patch( + "runpod.api.ctl_commands.read_event_stream", + side_effect=_streams([_log_event("1", "Worker ready.")]), + ) as stream: + logs = ctl_commands.get_endpoint_worker_logs("ep", "worker/1", tail=10) + + assert logs == [_entry("1", "Worker ready.")] + assert stream.call_args.args == ("/v2/serverless/ep/workers/worker%2F1/logs",) + assert stream.call_args.kwargs["params"] == {"tail": 10} + + +@pytest.mark.parametrize( + "kwargs", + [ + {"tail": -1}, + {"tail": 5001}, + {"source": "both"}, + {"max_wait": 0}, + {"since": datetime(2026, 6, 1)}, + ], +) +@pytest.mark.parametrize( + "call", + [ + lambda **kw: ctl_commands.get_pod_logs("pod", **kw), + lambda **kw: ctl_commands.iter_pod_logs("pod", **kw), + lambda **kw: ctl_commands.get_endpoint_worker_logs("ep", "w", **kw), + lambda **kw: ctl_commands.iter_endpoint_worker_logs("ep", "w", **kw), + ], + ids=["get_pod", "iter_pod", "get_worker", "iter_worker"], +) +def test_log_functions_validate_arguments_eagerly(call, kwargs): + with ( + patch("runpod.api.ctl_commands.read_event_stream") as stream, + pytest.raises(ValueError), + ): + call(**kwargs) + + stream.assert_not_called() + + +def test_get_pod_logs_requires_max_wait(): + with pytest.raises(ValueError): + ctl_commands.get_pod_logs("pod", max_wait=None) + + +def test_iter_pod_logs_resumes_from_last_event_id(): + with ( + patch( + "runpod.api.ctl_commands.read_event_stream", + side_effect=_streams( + [_log_event("1", "a"), _log_event("2", "b")], + [], + [_log_event("3", "c")], + ), + ) as stream, + patch("runpod.api.ctl_commands.time.sleep") as sleep, + ): + logs = ctl_commands.iter_pod_logs("pod", tail=5, source="container") + lines = [next(logs)["line"] for _ in range(3)] + logs.close() + + assert lines == ["a", "b", "c"] + assert [call.kwargs["last_event_id"] for call in stream.call_args_list] == [ + None, + "2", + "2", + ] + assert all(call.kwargs["max_wait"] is None for call in stream.call_args_list) + assert all( + call.kwargs["params"] == {"tail": 5, "source": "container"} + for call in stream.call_args_list + ) + assert [call.args for call in sleep.call_args_list] == [(1,), (1,)] + + +def test_iter_pod_logs_retries_rate_limit_and_network_errors_on_reconnect(): + with ( + patch( + "runpod.api.ctl_commands.read_event_stream", + side_effect=_streams( + [_log_event("1", "a")], + [QueryError("slow down", status_code=429, retry_after=12)], + [requests.exceptions.ConnectionError("reset")], + [_log_event("2", "b")], + ), + ), + patch("runpod.api.ctl_commands.time.sleep") as sleep, + ): + logs = ctl_commands.iter_pod_logs("pod") + lines = [next(logs)["line"] for _ in range(2)] + logs.close() + + assert lines == ["a", "b"] + assert [call.args for call in sleep.call_args_list] == [(1,), (12,), (1,)] + + +@pytest.mark.parametrize( + "failure", + [ + QueryError("pod not found", status_code=404), + QueryError("slow down", status_code=429, retry_after=3), + requests.exceptions.ConnectionError("refused"), + ], +) +def test_iter_pod_logs_raises_on_first_connection_failure(failure): + with ( + patch( + "runpod.api.ctl_commands.read_event_stream", + side_effect=_streams([failure]), + ), + pytest.raises(type(failure)), + ): + next(ctl_commands.iter_pod_logs("pod")) + + +def test_iter_pod_logs_raises_non_retryable_error_on_reconnect(): + with ( + patch( + "runpod.api.ctl_commands.read_event_stream", + side_effect=_streams( + [_log_event("1", "a")], + [QueryError("forbidden", status_code=403)], + ), + ), + patch("runpod.api.ctl_commands.time.sleep"), + pytest.raises(QueryError, match="forbidden"), + ): + list(ctl_commands.iter_pod_logs("pod")) + + +def test_iter_pod_logs_stops_at_max_wait(): + with ( + patch( + "runpod.api.ctl_commands.read_event_stream", + side_effect=_streams([_log_event("1", "a")], [_log_event("2", "b")]), + ) as stream, + patch("runpod.api.ctl_commands.time.sleep"), + patch("runpod.api.ctl_commands.time.monotonic", side_effect=[0, 0, 4, 4, 11, 11]), + ): + logs = list(ctl_commands.iter_pod_logs("pod", max_wait=10)) + + assert [entry["line"] for entry in logs] == ["a", "b"] + assert [call.kwargs["max_wait"] for call in stream.call_args_list] == [10, 6] + + +def test_iter_endpoint_worker_logs_uses_worker_path(): + with patch( + "runpod.api.ctl_commands.read_event_stream", + side_effect=_streams([_log_event("1", "a")]), + ) as stream: + logs = ctl_commands.iter_endpoint_worker_logs("ep", "w") + assert next(logs)["line"] == "a" + logs.close() + + assert stream.call_args.args == ("/v2/serverless/ep/workers/w/logs",) + + +def test_get_endpoint_workers_unwraps_response(): + workers = [{"id": "worker", "status": "RUNNING"}] + with patch( + "runpod.api.ctl_commands.run_rest_request", + return_value={"workers": workers, "summary": {"RUNNING": 1}}, + ) as request: + assert ctl_commands.get_endpoint_workers("ep/1", api_key="key") == workers + + request.assert_called_once_with( + "GET", "/v2/serverless/ep%2F1/workers", api_key="key" + ) + + @pytest.fixture def template_backend(): """Model REST template expansion and its persistent-volume size floor.""" diff --git a/tests/test_api/test_logs_server.py b/tests/test_api/test_logs_server.py new file mode 100644 index 00000000..2da26d30 --- /dev/null +++ b/tests/test_api/test_logs_server.py @@ -0,0 +1,228 @@ +"""Log readers against a local fake of the Runpod REST v2 log endpoints. + +These run over real sockets so the requests/urllib3 streaming behavior (chunked +reads, read timeouts, early close) is exercised, not mocked. +""" + +import json +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import parse_qs, urlsplit + +import pytest + +import runpod +from runpod.api import ctl_commands +from runpod.error import QueryError + +EVENTS = [ + {"ts": f"2026-06-01T12:00:0{i}Z", "source": "container", "line": f"line {i} é"} + for i in range(6) +] + + +def _frame(index): + return f"id: {index}\ndata: {json.dumps(EVENTS[index])}\n\n".encode() + + +class FakeRunpod(BaseHTTPRequestHandler): + """Serves a log stream whose behavior is picked by the resource ID.""" + + protocol_version = "HTTP/1.1" + requests_seen = [] + rate_limited = set() + + def log_message(self, *_args): + pass + + def _chunk(self, data): + self.wfile.write(b"%x\r\n%s\r\n" % (len(data), data)) + self.wfile.flush() + + def _json(self, status, payload, headers=None): + body = json.dumps(payload).encode() + self.send_response(status) + self.send_header("Content-Type", "application/problem+json") + self.send_header("Content-Length", str(len(body))) + for key, value in (headers or {}).items(): + self.send_header(key, value) + self.end_headers() + self.wfile.write(body) + + def _start_stream(self): + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Transfer-Encoding", "chunked") + self.end_headers() + + def _end_stream(self): + self.wfile.write(b"0\r\n\r\n") + self.wfile.flush() + self.close_connection = True + + def do_GET(self): # pylint: disable=invalid-name + url = urlsplit(self.path) + FakeRunpod.requests_seen.append( + { + "path": url.path, + "query": parse_qs(url.query), + "last_event_id": self.headers.get("Last-Event-ID"), + "accept": self.headers.get("Accept"), + } + ) + if url.path == "/v2/serverless/ep/workers": + self._json(200, {"workers": [{"id": "w1"}], "summary": {"RUNNING": 1}}) + return + + mode = url.path.split("/")[-2] + resume_from = int(self.headers.get("Last-Event-ID", "-1")) + 1 + if mode == "missing": + self._json(404, {"title": "Not Found", "status": 404, "detail": "pod not found"}) + return + if mode == "ratelimit" and resume_from > 0 and mode not in self.rate_limited: + FakeRunpod.rate_limited.add(mode) + self._json(429, {"detail": "rate limited"}, {"Retry-After": "0"}) + return + + self._start_stream() + try: + if mode in ("resume", "ratelimit"): + # Two events per connection, then close: the client must resume. + for index in range(resume_from, min(resume_from + 2, len(EVENTS))): + self._chunk(_frame(index)) + self._end_stream() + return + for index in range(3): + self._chunk(_frame(index)) + if mode == "chatty": + index = 3 + while True: + self._chunk(_frame(index % len(EVENTS))) + index += 1 + time.sleep(0.02) + if mode == "partial": + self._chunk(b'id: 9\ndata: {"line": "cut') + time.sleep(10) + except (BrokenPipeError, ConnectionResetError): + pass + + +class QuietServer(ThreadingHTTPServer): + daemon_threads = True + + def handle_error(self, request, client_address): + pass + + +@pytest.fixture(name="server", scope="module") +def fixture_server(): + server = QuietServer(("127.0.0.1", 0), FakeRunpod) + threading.Thread(target=server.serve_forever, daemon=True).start() + yield f"http://127.0.0.1:{server.server_port}" + server.shutdown() + + +@pytest.fixture(autouse=True) +def fixture_api(server, monkeypatch): + monkeypatch.setenv("RUNPOD_API_BASE_URL", server) + monkeypatch.setattr(runpod, "api_key", "key") + monkeypatch.setattr(ctl_commands, "LOG_RECONNECT_DELAY", 0.01) + FakeRunpod.requests_seen.clear() + FakeRunpod.rate_limited.clear() + + +def _lines(logs): + return [entry["line"] for entry in logs] + + +def test_snapshot_returns_backfill_when_stream_goes_idle(): + started = time.monotonic() + logs = runpod.get_pod_logs("idle", tail=3, source="container", max_wait=0.5) + + assert time.monotonic() - started < 2 + assert logs[0] == {"id": "0", **EVENTS[0]} + assert _lines(logs) == ["line 0 é", "line 1 é", "line 2 é"] + assert FakeRunpod.requests_seen == [ + { + "path": "/v2/pods/idle/logs", + "query": {"tail": ["3"], "source": ["container"]}, + "last_event_id": None, + "accept": "text/event-stream", + } + ] + + +def test_snapshot_stops_at_deadline_on_busy_stream(): + started = time.monotonic() + logs = runpod.get_pod_logs("chatty", max_wait=0.5) + + assert time.monotonic() - started < 1.5 + assert len(logs) > 3 + + +def test_snapshot_discards_event_cut_mid_frame(): + assert _lines(runpod.get_pod_logs("partial", max_wait=0.5)) == [ + "line 0 é", + "line 1 é", + "line 2 é", + ] + + +def test_snapshot_keeps_newest_lines_past_byte_cap(): + logs = runpod.get_pod_logs("idle", max_wait=0.5, max_bytes=len("line 0 é") * 2) + + assert _lines(logs) == ["line 1 é", "line 2 é"] + + +def test_missing_pod_raises_query_error(): + with pytest.raises(QueryError, match="pod not found") as raised: + runpod.get_pod_logs("missing", max_wait=0.5) + + assert raised.value.status_code == 404 + with pytest.raises(QueryError): + next(runpod.iter_pod_logs("missing")) + + +def test_follow_resumes_from_last_event_id_without_gaps_or_repeats(): + logs = runpod.iter_pod_logs("resume") + entries = [next(logs) for _ in range(len(EVENTS))] + logs.close() + + assert [entry["id"] for entry in entries] == ["0", "1", "2", "3", "4", "5"] + assert [seen["last_event_id"] for seen in FakeRunpod.requests_seen[:3]] == [ + None, + "1", + "3", + ] + + +def test_follow_waits_out_rate_limit_on_reconnect(): + logs = runpod.iter_pod_logs("ratelimit") + entries = [next(logs) for _ in range(4)] + logs.close() + + assert [entry["id"] for entry in entries] == ["0", "1", "2", "3"] + assert [seen["last_event_id"] for seen in FakeRunpod.requests_seen[:3]] == [ + None, + "1", + "1", + ] + + +def test_follow_stops_at_max_wait(): + started = time.monotonic() + logs = list(runpod.iter_pod_logs("idle", max_wait=0.5)) + + assert time.monotonic() - started < 2 + assert _lines(logs)[:3] == ["line 0 é", "line 1 é", "line 2 é"] + assert all(entry["id"] in {"0", "1", "2"} for entry in logs[3:]) + + +def test_worker_logs_and_worker_listing(): + workers = runpod.get_endpoint_workers("ep") + logs = runpod.get_endpoint_worker_logs("ep", workers[0]["id"], max_wait=0.5) + + assert workers == [{"id": "w1"}] + assert len(logs) == 3 + assert FakeRunpod.requests_seen[-1]["path"] == "/v2/serverless/ep/workers/w1/logs" diff --git a/tests/test_api/test_rest.py b/tests/test_api/test_rest.py index d0027468..86a65da1 100644 --- a/tests/test_api/test_rest.py +++ b/tests/test_api/test_rest.py @@ -1,11 +1,13 @@ """Tests for the REST API transport.""" -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch import pytest import runpod -from runpod.api.rest import run_rest_request +import requests + +from runpod.api.rest import _parse_event_stream, read_event_stream, run_rest_request from runpod.error import AuthenticationError, QueryError from runpod.user_agent import USER_AGENT @@ -115,3 +117,119 @@ def test_request_uses_text_for_non_json_error(): run_rest_request("GET", "/v2/pods", api_key="key") assert raised.value.status_code == 500 + + +def _stream_response(chunks, status_code=200): + response = MagicMock() + response.status_code = status_code + response.headers = {} + response.__enter__.return_value = response + + def iter_content(chunk_size=None): + for chunk in chunks: + if isinstance(chunk, Exception): + raise chunk + yield chunk + + response.iter_content.side_effect = iter_content + return response + + +def test_parse_event_stream_reassembles_split_frames(): + chunks = [ + b"id: 1\r\ndata: {\"line\": \"caf", + "\u00e9\"}\r\n\r\n: keep-alive\n\n".encode()[:1], + "\u00e9\"}\r\n\r\n: keep-alive\n\n".encode()[1:], + b"event: log\ndata: a\ndata: b\n\n", + b"id: 3\ndata: incomplete", + ] + + assert list(_parse_event_stream(iter(chunks))) == [ + {"id": "1", "data": '{"line": "caf\u00e9"}'}, + {"event": "log", "data": "a\nb"}, + ] + + +def test_read_event_stream_requests_sse(): + response = _stream_response([b"id: 1\ndata: x\n\n"]) + with patch("runpod.api.rest.requests.get", return_value=response) as get: + events = list( + read_event_stream( + "/v2/pods/pod/logs", api_key="key", params={"tail": 5}, max_wait=2 + ) + ) + + assert events == [{"id": "1", "data": "x"}] + get.assert_called_once_with( + "https://api.runpod.io/v2/pods/pod/logs", + headers={ + "Accept": "text/event-stream", + "Content-Type": "application/json", + "User-Agent": USER_AGENT, + "Authorization": "Bearer key", + }, + params={"tail": 5}, + stream=True, + timeout=(30, 2), + ) + + +def test_read_event_stream_raises_for_error_response(): + response = _stream_response([], status_code=404) + response.json.return_value = {"detail": "pod not found"} + with ( + patch("runpod.api.rest.requests.get", return_value=response), + pytest.raises(QueryError, match="pod not found") as raised, + ): + list(read_event_stream("/v2/pods/pod/logs", api_key="key")) + + assert raised.value.status_code == 404 + + +def test_read_event_stream_stops_at_deadline(): + response = _stream_response([b"data: 1\n\n", b"data: 2\n\n", b"data: 3\n\n"]) + with ( + patch("runpod.api.rest.requests.get", return_value=response), + patch("runpod.api.rest.time.monotonic", side_effect=[0, 1, 5]), + ): + events = list(read_event_stream("/v2/pods/pod/logs", api_key="key", max_wait=5)) + + assert events == [{"data": "1"}, {"data": "2"}] + + +def test_read_event_stream_ends_on_idle_timeout(): + response = _stream_response( + [b"data: 1\n\n", requests.exceptions.ConnectionError("read timed out")] + ) + with patch("runpod.api.rest.requests.get", return_value=response): + events = list(read_event_stream("/v2/pods/pod/logs", api_key="key")) + + assert events == [{"data": "1"}] + + +def test_read_event_stream_resumes_and_uses_idle_timeout_without_deadline(): + response = _stream_response([b"id: 2\ndata: x\n\n"]) + with patch("runpod.api.rest.requests.get", return_value=response) as get: + events = list( + read_event_stream( + "/v2/pods/pod/logs", api_key="key", max_wait=None, last_event_id="1" + ) + ) + + assert events == [{"id": "2", "data": "x"}] + assert get.call_args.kwargs["headers"]["Last-Event-ID"] == "1" + assert get.call_args.kwargs["timeout"] == (30, 60) + + +@pytest.mark.parametrize(("header", "expected"), [("12", 12.0), ("soon", None)]) +def test_request_exposes_retry_after_on_rate_limit(header, expected): + response = _response(status_code=429, payload={"detail": "rate limited"}) + response.headers = {"Retry-After": header} + with ( + patch("runpod.api.rest.requests.request", return_value=response), + pytest.raises(QueryError, match="rate limited") as raised, + ): + run_rest_request("GET", "/v2/pods", api_key="key") + + assert raised.value.status_code == 429 + assert raised.value.retry_after == expected diff --git a/tests/test_init.py b/tests/test_init.py index 91242328..87721f04 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -23,7 +23,8 @@ def test_api_functions_accessible(self): api_functions = [ 'create_container_registry_auth', 'create_endpoint', 'create_pod', 'create_template', 'delete_container_registry_auth', 'get_endpoints', 'get_gpu', 'get_gpus', - 'get_pod', 'get_pods', 'get_user', 'resume_pod', 'stop_pod', 'terminate_pod', + 'get_endpoint_worker_logs', 'get_endpoint_workers', 'get_pod', 'get_pod_logs', + 'get_pods', 'get_user', 'iter_endpoint_worker_logs', 'iter_pod_logs', 'resume_pod', 'stop_pod', 'terminate_pod', 'update_container_registry_auth', 'update_endpoint_template', 'update_user_settings' ] @@ -92,7 +93,8 @@ def test_all_covers_expected_public_api(self): # API functions 'create_container_registry_auth', 'create_endpoint', 'create_pod', 'create_template', 'delete_container_registry_auth', 'get_endpoints', 'get_gpu', 'get_gpus', - 'get_pod', 'get_pods', 'get_user', 'resume_pod', 'stop_pod', 'terminate_pod', + 'get_endpoint_worker_logs', 'get_endpoint_workers', 'get_pod', 'get_pod_logs', + 'get_pods', 'get_user', 'iter_endpoint_worker_logs', 'iter_pod_logs', 'resume_pod', 'stop_pod', 'terminate_pod', 'update_container_registry_auth', 'update_endpoint_template', 'update_user_settings', # Config functions 'check_credentials', 'get_credentials', 'set_credentials',