diff --git a/python/docs/cli/index.rst b/python/docs/cli/index.rst index 6b864f4d..d7a60cf4 100644 --- a/python/docs/cli/index.rst +++ b/python/docs/cli/index.rst @@ -59,6 +59,8 @@ Common Workflows * Develop interactively with :doc:`studio`. * Submit and inspect training or batch work with :doc:`job` and :doc:`mmt`. +* Tail, follow, and search a run's logs with its own ``logs`` command — see + :doc:`job`, :doc:`mmt`, :doc:`deployment`, and :doc:`sandbox`. * Build and operate inference services with :doc:`deployment` and :doc:`model`. * Move data and artifacts with :doc:`file`, :doc:`folder`, :doc:`container`, and :doc:`cp`. * Configure accounts, organizations, teamspaces, cloud accounts, and SSH with diff --git a/python/docs/sdk/examples.rst b/python/docs/sdk/examples.rst index ac7b654b..5c1bfd68 100644 --- a/python/docs/sdk/examples.rst +++ b/python/docs/sdk/examples.rst @@ -29,4 +29,5 @@ CLI examples examples/mmts-cli examples/teamspaces-cli examples/sandboxes-cli + examples/logs-cli examples/api-cli diff --git a/python/docs/sdk/examples/logs-cli.rst b/python/docs/sdk/examples/logs-cli.rst new file mode 100644 index 00000000..5ec033af --- /dev/null +++ b/python/docs/sdk/examples/logs-cli.rst @@ -0,0 +1 @@ +.. include:: ../../../examples/logs_cli.rst diff --git a/python/examples/README.md b/python/examples/README.md index 3b9506dc..95cdadf4 100644 --- a/python/examples/README.md +++ b/python/examples/README.md @@ -69,6 +69,7 @@ export LIGHTNING_SANDBOX_API_KEY="..." | [`mmts_cli.rst`](mmts_cli.rst) | Launch and inspect multi-machine training runs from the CLI | | [`teamspaces_cli.rst`](teamspaces_cli.rst) | Set CLI context and pass explicit teamspace scope to resource commands | | [`sandboxes_cli.rst`](sandboxes_cli.rst) | Create sandboxes, run commands, inspect logs, stop, resume, and delete from the CLI | +| [`logs_cli.rst`](logs_cli.rst) | Tail, follow, search, and page logs for jobs, MMTs, deployments, and sandboxes | | [`api_cli.rst`](api_cli.rst) | Call authenticated Lightning API endpoints directly from shell scripts | # Running examples diff --git a/python/examples/jobs.rst b/python/examples/jobs.rst index b14cebaa..174bdbae 100644 --- a/python/examples/jobs.rst +++ b/python/examples/jobs.rst @@ -59,7 +59,18 @@ Operational notes - ``Job.run`` creates a new job; ``Job("name", teamspace=...)`` fetches an existing one. -- ``job.logs`` is available after the job reaches a terminal state. +- ``job.logs`` reads the logs saved so far, whether the job is running or + finished. Call it to pass options: ``job.logs(follow=True)`` streams new lines + from a running job, and ``tail``, ``timestamps``, ``query`` (only lines + containing every whitespace-separated term) and ``severity`` (``error``, + ``warning``, ``info`` or ``debug`` and above) narrow what comes back. +- ``mmt.logs`` reads every machine of a multi-machine job, merged into one + timeline and labelled with the machine each line came from. +- The same reads are available from the CLI, one command per resource: + ``lightning job logs ``, ``lightning mmt logs ``, + ``lightning deployment logs `` and ``lightning sandbox logs ``. Each + takes ``--follow``, ``--tail``, ``--since``/``--until``, ``--query`` and + ``--severity``. - Studio-backed jobs must run in the same teamspace and cloud account as the Studio. - Container-backed jobs cannot also pass ``studio=``. diff --git a/python/examples/logs_cli.rst b/python/examples/logs_cli.rst new file mode 100644 index 00000000..42c5afe9 --- /dev/null +++ b/python/examples/logs_cli.rst @@ -0,0 +1,90 @@ +Logs CLI examples +================= + +Every resource that produces logs — jobs, multi-machine jobs, deployments, and +sandboxes — has its own ``logs`` command, and they share the same flags. Use +them to tail a run from a terminal, search a finished run for an error, or watch +only the lines that matter. + +Prerequisites +------------- + +.. code-block:: console + + $ pip install lightning-sdk -U + $ lightning login + $ lightning config set teamspace owner/teamspace + +Every example below also accepts ``--teamspace owner/teamspace`` explicitly; pass +it when you work across several teamspaces or run in CI. + +Pick what to read +----------------- + +One command per resource, addressed by name (or id for sandboxes): + +.. code-block:: console + + $ lightning job logs sdk-tutorial-job + $ lightning mmt logs sdk-tutorial-mmt + $ lightning deployment logs sdk-tutorial-api + $ lightning sandbox logs sbx-42 + +A multi-machine job and a multi-replica deployment merge every machine or replica +into one timeline, and each line is labelled with where it came from. To read a +single machine of an MMT, read it as a job: ``lightning job logs ``. + +Tail and follow +--------------- + +``--tail`` prints the last N lines; ``--follow`` (``-f``) keeps the stream open +until the resource finishes or you press Ctrl-C: + +.. code-block:: console + + $ lightning job logs sdk-tutorial-job --tail 100 + $ lightning deployment logs sdk-tutorial-api --follow + $ lightning mmt logs sdk-tutorial-mmt --tail 50 --follow + +Search and narrow +----------------- + +``--query`` keeps only lines containing every whitespace-separated term, +``--severity`` sets the minimum level (``error`` > ``warning`` > ``info`` > +``debug``), and ``--since``/``--until`` bound the time range. Time bounds take a +duration — ``30s``, ``5m``, ``2h``, ``3d``, ``1w`` — or an RFC3339 timestamp: + +.. code-block:: console + + $ lightning job logs sdk-tutorial-job --query "CUDA out of memory" + $ lightning deployment logs sdk-tutorial-api --severity warning --since 2h + $ lightning mmt logs sdk-tutorial-mmt --query error --since 3d --until 1d + $ lightning job logs sdk-tutorial-job --since 2026-01-01T00:00:00Z + +Filters combine with ``--tail`` and ``--follow``, so you can watch only what +matters, and ``--timestamps`` prefixes each line with its ISO-8601 timestamp: + +.. code-block:: console + + $ lightning deployment logs sdk-tutorial-api --severity error --follow + $ lightning job logs sdk-tutorial-job --tail 200 --timestamps + +Sandboxes +--------- + +A sandbox reads by id, and an optional command id narrows to a single detached +command; omit it for the merged logs of every command in the sandbox: + +.. code-block:: console + + $ lightning sandbox logs sbx-42 + $ lightning sandbox logs sbx-42 cmd-abc123 + $ lightning sandbox logs sbx-42 --query error + +See also +-------- + +- ``jobs_cli.rst`` for submitting and inspecting the jobs these logs come from. +- ``sandboxes_cli.rst`` for running detached sandbox commands. +- ``lightning job logs --help`` (also ``mmt``, ``deployment``, and ``sandbox``) + for the full option reference. diff --git a/python/lightning_sdk/api/deployment_api.py b/python/lightning_sdk/api/deployment_api.py index bc4698a1..6cc44ba4 100644 --- a/python/lightning_sdk/api/deployment_api.py +++ b/python/lightning_sdk/api/deployment_api.py @@ -7,12 +7,17 @@ from pathlib import Path, PurePosixPath from tempfile import TemporaryDirectory from time import sleep -from typing import Any, Dict, List, Literal, Optional, Sequence, TextIO, Tuple, Union +from typing import Any, Dict, Iterator, List, Literal, Optional, Sequence, TextIO, Tuple, Union +from urllib.parse import urlparse import requests from lightning_sdk.api import lightning_storage_upload as lightning_storage_upload_api -from lightning_sdk.api.utils import _BlobUploader, _machine_to_compute_name, resolve_path_mappings +from lightning_sdk.api.logs_api import LogEntry, parse_log_entries +from lightning_sdk.api.utils import _BlobUploader, _get_cloud_url, _machine_to_compute_name, resolve_path_mappings + +# `Auth` is the deployment endpoint auth type in this module, so the login client is aliased. +from lightning_sdk.lightning_cloud.login import Auth as _LoginAuth from lightning_sdk.lightning_cloud.openapi import ( JobsServiceCreateDeploymentBody, JobsServiceReloadDeploymentWeightsBody, @@ -621,6 +626,54 @@ def get_job_logs( kwargs = {k: v for k, v in kwargs.items() if v is not None} return self._client.jobs_service_get_job_logs(project_id=teamspace_id, id=job_id, **kwargs) + def iter_job_log_entries( + self, + teamspace_id: str, + job_id: str, + *, + deployment_id: Optional[str] = None, + since: Optional[str] = None, + until: Optional[str] = None, + rank: Optional[int] = None, + ) -> Iterator[LogEntry]: + """Yield a job's saved log lines from the legacy per-job log pages. + + :meth:`get_job_logs` returns page metadata only, so each page is downloaded from its + signed URL and parsed here. Prefer :class:`~lightning_sdk.api.logs_api.LogsApi`, which + returns lines inline; this path exists for reads the merged logs API cannot serve yet + (currently only ``rank``). + + Args: + teamspace_id: The teamspace that owns the job. + job_id: The job ID. + deployment_id: Optional deployment ID filter. + since: Optional start timestamp. + until: Optional end timestamp. + rank: Optional distributed job rank. + + Yields: + LogEntry: The saved log lines, page by page. + """ + logs = self.get_job_logs( + teamspace_id, + job_id, + deployment_id=deployment_id, + since=since, + until=until, + rank=rank, + ) + session = requests.Session() + session.headers.update({"Authorization": _LoginAuth().authenticate()}) + for page in logs.pages or []: + url = getattr(page, "url", None) + if not url: + continue + if not urlparse(url).netloc: + url = f"{_get_cloud_url().rstrip('/')}/{url.lstrip('/')}" + response = session.get(url, timeout=60) + response.raise_for_status() + yield from parse_log_entries(response.text) + def stop(self, deployment: V1Deployment) -> V1Deployment: """Scale a deployment to zero replicas and wait until all replicas have stopped. diff --git a/python/lightning_sdk/api/logs_api.py b/python/lightning_sdk/api/logs_api.py new file mode 100644 index 00000000..82ef535b --- /dev/null +++ b/python/lightning_sdk/api/logs_api.py @@ -0,0 +1,525 @@ +"""Client for the unified logs API (``GET /v1/projects/{id}/page-logs``). + +A single endpoint serves every log-bearing resource: one or many jobs, a whole +deployment (all replicas), every rank of a multi-machine job, and sandbox +commands. Log lines are returned inline and filtered server-side (substring +``query`` and minimum ``severity``), so a caller never fetches and parses log +pages itself. The response also carries a ``follow_url`` websocket -- populated +only while the resource is still active -- for tailing the live stream. + +Historical reads are cursor-paginated forward in time; there is no server-side +``tail``, so a bounded tail is applied client-side after paging. +""" + +import json +import time +from collections import deque +from contextlib import suppress +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from functools import partial +from typing import Any, Callable, Deque, Dict, Iterator, List, Optional, Sequence, Set, Tuple, Union +from urllib.parse import urlparse + +from lightning_sdk.api.utils import _get_cloud_url +from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.lightning_cloud.rest_client import LightningClient + +# While tailing, ``recv()`` is given this timeout so a quiet socket can be re-checked (the server +# holds the connection open with heartbeats rather than closing it). It must stay below the +# server's log-socket heartbeat interval (10s) so ``recv()`` actually wakes up during a lull. +_FOLLOW_POLL_INTERVAL = 5.0 + +# Budget for the websocket handshake, which can take longer than a read poll. +_CONNECT_TIMEOUT = 30.0 + +# Look-back windows tried, newest first, when serving a `tail` without an explicit `since`. A +# long-lived resource (a deployment with months of replicas) would otherwise have its whole +# history paged just to keep the last few lines, so start narrow and widen only if needed. +_TAIL_WINDOWS = ( + 10 * 60, + 60 * 60, + 6 * 60 * 60, + 24 * 60 * 60, + 7 * 24 * 60 * 60, +) + +# Overlap window used to de-duplicate the live tail against the history that was just printed: +# the follow socket starts at "now", so lines written between the last page read and the socket +# being established can arrive twice. Only the most recent history keys can collide, so the set +# is bounded rather than holding every line of a long job. +_MAX_DEDUP_KEYS = 5000 + +SEVERITIES = ("error", "warning", "info", "debug") + + +@dataclass(frozen=True) +class LogEntry: + """One log line, as returned by the logs API or its follow socket.""" + + message: str + timestamp: Optional[datetime] = None + line: int = 0 + severity: str = "" + # The job the line came from: the replica for a deployment, the rank for a multi-machine job. + resource_id: str = "" + + def format(self, *, timestamps: bool = False, prefix: Optional[str] = None) -> str: + """Render the entry as a printable line. + + Args: + timestamps: Prepend the entry's ISO-8601 timestamp when it has one. + prefix: Optional label (e.g. a replica name) to bracket in front of the message. + """ + parts = [] + if timestamps and self.timestamp is not None: + parts.append(self.timestamp.isoformat()) + if prefix: + parts.append(f"[{prefix}]") + parts.append(self.message) + return " ".join(parts) + + @property + def dedup_key(self) -> Tuple[int, Optional[datetime], str]: + """Identity used to drop a live line that history already produced. + + ``line`` is per-resource, so it is not unique across replicas; combining it with the + timestamp and message is what the console UI does too. + """ + return (self.line, self.timestamp, self.message) + + +@dataclass +class LogsPage: + """One page of history, plus the follow socket for whatever is still running.""" + + entries: List[LogEntry] = field(default_factory=list) + next_page_token: Optional[str] = None + # Absent once the resource is finished: there is nothing left to tail. + follow_url: Optional[str] = None + + +def _parse_timestamp(value: Any) -> Optional[datetime]: + """Normalise the timestamp shapes the two transports use into a ``datetime``. + + The REST client deserialises the protobuf ``Timestamp`` into a ``datetime`` already; the + websocket sends it as ``{"seconds": ..., "nanos": ...}``. An ISO-8601 string is accepted in + case either representation changes. + """ + if value is None or isinstance(value, datetime): + return value + if isinstance(value, dict): + seconds = value.get("seconds") + if seconds is None: + return None + epoch = float(seconds) + float(value.get("nanos", 0)) / 1e9 + return datetime.fromtimestamp(epoch, tz=timezone.utc) + if isinstance(value, str) and value: + with suppress(ValueError): + return datetime.fromisoformat(value.replace("Z", "+00:00")) + return None + + +def _to_int(value: Any) -> int: + try: + return int(value) + except (TypeError, ValueError): + return 0 + + +def _entry_from_model(entry: Any) -> LogEntry: + """Build a :class:`LogEntry` from a generated ``V1JobLogEntry``.""" + return LogEntry( + message=str(getattr(entry, "message", "") or ""), + timestamp=_parse_timestamp(getattr(entry, "timestamp", None)), + line=_to_int(getattr(entry, "line", 0)), + severity=str(getattr(entry, "severity", "") or ""), + resource_id=str(getattr(entry, "resource_id", "") or ""), + ) + + +def parse_log_entries(payload_text: str) -> List[LogEntry]: + """Decode a JSON array of log entries into :class:`LogEntry` objects. + + Used for both websocket frames and the bodies of legacy log pages, which share this shape. + Keys are snake_case (``resource_id``) because the server serialises its protobuf struct + directly; camelCase is accepted too. Text that is not JSON at all is surfaced line by line so + nothing is silently dropped. + """ + try: + payload = json.loads(payload_text) + except json.JSONDecodeError: + return [LogEntry(message=line) for line in payload_text.splitlines()] + + raw_entries = payload if isinstance(payload, list) else [payload] + entries = [] + for raw in raw_entries: + if not isinstance(raw, dict): + entries.append(LogEntry(message=str(raw))) + continue + entries.append( + LogEntry( + message=str(raw.get("message") or raw.get("Message") or ""), + timestamp=_parse_timestamp(raw.get("timestamp") or raw.get("Timestamp")), + line=_to_int(raw.get("line")), + severity=str(raw.get("severity") or ""), + resource_id=str(raw.get("resource_id") or raw.get("resourceId") or ""), + ) + ) + return entries + + +def _websocket_url(url: str) -> str: + """Return ``url`` as an absolute ``ws(s)://`` URL. + + The server already hands back an absolute ``wss://`` follow URL; a relative or ``http(s)`` + URL is still normalised so the caller does not depend on that. + """ + if not urlparse(url).netloc: + url = f"{_get_cloud_url().rstrip('/')}/{url.lstrip('/')}" + if url.startswith("https://"): + return "wss://" + url[len("https://") :] + if url.startswith("http://"): + return "ws://" + url[len("http://") :] + return url + + +class _RecentKeys: + """Bounded set of the most recently seen entry keys.""" + + def __init__(self, maxlen: int = _MAX_DEDUP_KEYS) -> None: + self._order: Deque[Tuple[int, Optional[datetime], str]] = deque() + self._keys: Set[Tuple[int, Optional[datetime], str]] = set() + self._maxlen = maxlen + + def add(self, entry: LogEntry) -> None: + key = entry.dedup_key + if key in self._keys: + return + self._order.append(key) + self._keys.add(key) + while len(self._order) > self._maxlen: + self._keys.discard(self._order.popleft()) + + def __contains__(self, entry: LogEntry) -> bool: + return entry.dedup_key in self._keys + + +class LogsApi: + """Read logs for jobs, deployments, multi-machine jobs and sandbox commands.""" + + def __init__(self, client: Optional[LightningClient] = None) -> None: + self._client = client or LightningClient(max_tries=7) + + def get_page( + self, + teamspace_id: str, + *, + job_ids: Sequence[str] = (), + deployment_id: Optional[str] = None, + mmt_id: Optional[str] = None, + sandbox_id: Optional[str] = None, + sandbox_command_ids: Sequence[str] = (), + since: Optional[str] = None, + until: Optional[str] = None, + query: Optional[str] = None, + severity: Optional[str] = None, + page_size: Optional[int] = None, + page_token: Optional[str] = None, + ) -> LogsPage: + """Fetch a single page of logs. + + At least one selector (``job_ids``, ``deployment_id``, ``mmt_id``, ``sandbox_id`` or + ``sandbox_command_ids``) must be given; the server rejects a request without one. + + Args: + teamspace_id: The teamspace that owns the resource. + job_ids: Read these jobs, merged into one timeline. + deployment_id: Read every replica of this deployment. + mmt_id: Read every rank of this multi-machine job. + sandbox_id: Read every recorded command of this sandbox. + sandbox_command_ids: Narrow to specific sandbox commands. + since: Only include lines at or after this RFC3339 timestamp. + until: Only include lines at or before this RFC3339 timestamp. + query: Only include lines containing every whitespace-separated term. + severity: Only include lines at or above this level (see :data:`SEVERITIES`). + page_size: Lines per page; the server defaults to 1000. + page_token: Cursor from a previous page's ``next_page_token``. + + Returns: + LogsPage: The page's entries, its continuation token and the follow URL. + """ + kwargs: Dict[str, Any] = {} + if job_ids: + kwargs["job_ids"] = list(job_ids) + if deployment_id: + kwargs["deployment_id"] = deployment_id + if mmt_id: + kwargs["mmt_id"] = mmt_id + if sandbox_id: + kwargs["sandbox_id"] = sandbox_id + if sandbox_command_ids: + kwargs["sandbox_command_ids"] = list(sandbox_command_ids) + if since: + kwargs["since"] = since + if until: + kwargs["until"] = until + if query: + kwargs["query"] = query + if severity: + kwargs["severity"] = severity + if page_size is not None: + kwargs["page_size"] = str(page_size) + if page_token: + kwargs["page_token"] = page_token + + response = self._client.jobs_service_get_logs(project_id=teamspace_id, **kwargs) + return LogsPage( + entries=[_entry_from_model(entry) for entry in (getattr(response, "entries", None) or [])], + next_page_token=getattr(response, "next_page_token", None) or None, + follow_url=getattr(response, "follow_url", None) or None, + ) + + def stream( + self, + teamspace_id: str, + *, + job_ids: Sequence[str] = (), + deployment_id: Optional[str] = None, + mmt_id: Optional[str] = None, + sandbox_id: Optional[str] = None, + sandbox_command_ids: Sequence[str] = (), + since: Optional[str] = None, + until: Optional[str] = None, + query: Optional[str] = None, + severity: Optional[str] = None, + page_size: Optional[int] = None, + follow: bool = False, + tail: Optional[int] = None, + tail_anchor: Optional[Union[datetime, str]] = None, + idle_timeout: Optional[float] = None, + fallback_to_live: bool = False, + stop: Optional[Callable[[], bool]] = None, + ) -> Iterator[LogEntry]: + """Yield the resource's history, then optionally tail its live stream. + + Args: + teamspace_id: The teamspace that owns the resource. + job_ids: Read these jobs, merged into one timeline. + deployment_id: Read every replica of this deployment. + mmt_id: Read every rank of this multi-machine job. + sandbox_id: Read every recorded command of this sandbox. + sandbox_command_ids: Narrow to specific sandbox commands. + since: Only include lines at or after this RFC3339 timestamp. + until: Only include lines at or before this RFC3339 timestamp. + query: Only include lines containing every whitespace-separated term. + severity: Only include lines at or above this level (see :data:`SEVERITIES`). + page_size: Lines per page; the server defaults to 1000. + follow: After the history, tail new lines until ``stop`` or interruption. + tail: Only yield the last N historical lines. The API pages forward only, so this + reads recent windows first and widens until N lines are found (see + :data:`_TAIL_WINDOWS`) instead of reading a long history in full. + tail_anchor: Where that window search ends; defaults to ``until``, else now. Set it to + a finished resource's stop time so its last lines are found without reading the + whole history. It bounds the search only, never the lines returned. + idle_timeout: While tailing, stop after this many seconds without a line. + fallback_to_live: When the history is empty, tail the live stream even if + ``follow`` is not set. Used to still show something for a running resource whose + logs are not in the new storage format yet. + stop: Called while the tail is quiet; return ``True`` to end the stream. + + Yields: + LogEntry: History entries in time order, then live entries as they arrive. + """ + selectors: Dict[str, Any] = { + "job_ids": job_ids, + "deployment_id": deployment_id, + "mmt_id": mmt_id, + "sandbox_id": sandbox_id, + "sandbox_command_ids": sandbox_command_ids, + } + if not any(selectors.values()): + raise ValueError("One of job_ids, deployment_id, mmt_id, sandbox_id or sandbox_command_ids is required.") + + filters: Dict[str, Any] = { + "until": until, + "query": query, + "severity": severity, + "page_size": page_size, + **selectors, + } + recent = _RecentKeys() + state: Dict[str, Any] = {} + printed = 0 + + if tail: + tail_entries = self._tail_history( + teamspace_id, since=since, tail=tail, tail_anchor=tail_anchor, state=state, **filters + ) + for entry in tail_entries: + recent.add(entry) + printed += 1 + yield entry + else: + for entry in self._iter_history(teamspace_id, since=since, state=state, **filters): + recent.add(entry) + printed += 1 + yield entry + + # Set while paging above, which the loops just drained. + follow_url = state.get("follow_url") + if not follow_url: + return + if not follow and not (fallback_to_live and printed == 0): + return + + yield from self.follow( + follow_url, + recent=recent, + idle_timeout=idle_timeout, + stop=stop, + reconnect=follow, + ) + + def _iter_history( + self, + teamspace_id: str, + *, + state: Dict[str, Any], + **kwargs: Any, + ) -> Iterator[LogEntry]: + """Page through the history in time order, recording the first page's follow URL in ``state``.""" + page_token: Optional[str] = None + while True: + page = self.get_page(teamspace_id, page_token=page_token, **kwargs) + if page_token is None: + state["follow_url"] = page.follow_url + yield from page.entries + page_token = page.next_page_token + if not page_token: + return + + def _tail_history( + self, + teamspace_id: str, + *, + tail: int, + since: Optional[str], + until: Optional[str] = None, + tail_anchor: Optional[Union[datetime, str]] = None, + state: Dict[str, Any], + **kwargs: Any, + ) -> List[LogEntry]: + """Return the last ``tail`` historical lines. + + The API pages forward from the start of the window, so reading a long-lived resource's + whole history just to keep its last few lines is wasteful. Instead, read a window ending + at ``until`` (or now) and widen it until enough lines turn up, falling back to the full + history last. An explicit ``since`` is honoured as-is: the caller already bounded the read. + """ + page = partial(self._iter_history, teamspace_id, until=until, state=state, **kwargs) + if since is not None: + return list(deque(page(since=since), maxlen=tail)) + + # Windows end where the logs do, so the last lines of a resource that finished a while + # ago are found in the first, narrowest window rather than by reading everything. Only + # the window's lower bound is sent, so a line written after the anchor still shows up. + anchor = _parse_timestamp(tail_anchor) or _parse_timestamp(until) or datetime.now(timezone.utc) + if anchor.tzinfo is None: + # The server only accepts RFC3339, which needs an offset; timestamps here are UTC. + anchor = anchor.replace(tzinfo=timezone.utc) + for seconds in (*_TAIL_WINDOWS, None): + window_start = None if seconds is None else (anchor - timedelta(seconds=seconds)).isoformat() + entries = deque(page(since=window_start), maxlen=tail) + if len(entries) >= tail or seconds is None: + return list(entries) + return [] + + def follow( + self, + follow_url: str, + *, + recent: Optional[_RecentKeys] = None, + idle_timeout: Optional[float] = None, + stop: Optional[Callable[[], bool]] = None, + reconnect: bool = True, + ) -> Iterator[LogEntry]: + """Yield log lines from a ``follow_url`` websocket as they arrive. + + The socket starts at the present moment and stays open (heartbeats only) once the + resource finishes rather than closing, so a quiet socket is re-checked against ``stop`` + instead of blocking forever. Any ``query``/``severity`` filter is already encoded in the + URL and applied server-side. + + Args: + follow_url: The ``follow_url`` from a :class:`LogsPage`. + recent: Keys of already-yielded history, to drop the replay overlap. + idle_timeout: Stop after this many seconds without a line. + stop: Called while the socket is quiet; return ``True`` to end the stream. + reconnect: Re-establish the socket after a transient drop. + + Yields: + LogEntry: Live entries in arrival order. + """ + try: + import websocket + from websocket import ( + WebSocketConnectionClosedException, + WebSocketTimeoutException, + ) + except ImportError as ex: + raise RuntimeError( + "Following logs requires the 'websocket-client' package (pip install websocket-client)." + ) from ex + + auth_header = Auth().authenticate() + url = _websocket_url(follow_url) + # Poll `recv()` rather than blocking on it, so `stop` and `idle_timeout` get a chance to + # run. The poll has to stay below the server's heartbeat interval: a heartbeat resets the + # socket's read timeout, so a longer one would never fire. Silence is therefore measured + # here rather than left to `recv()`. + recv_timeout = _FOLLOW_POLL_INTERVAL + if idle_timeout is not None: + recv_timeout = min(idle_timeout, _FOLLOW_POLL_INTERVAL) + last_line_at = time.monotonic() + + while True: + ws = None + try: + # The handshake gets a longer budget than the deliberately short read poll. + ws = websocket.create_connection( + url, header=[f"Authorization: {auth_header}"], timeout=_CONNECT_TIMEOUT + ) + ws.settimeout(recv_timeout) + while True: + try: + message = ws.recv() + except WebSocketTimeoutException: + if idle_timeout is not None and time.monotonic() - last_line_at >= idle_timeout: + return # quiet for long enough: treat the stream as finished + if stop is not None and stop(): + return + continue + except WebSocketConnectionClosedException: + break # fall through to the reconnect decision + if message == "": + break + entries = parse_log_entries(message) + if entries: + last_line_at = time.monotonic() + for entry in entries: + if recent is not None: + if entry in recent: + continue + recent.add(entry) + yield entry + finally: + if ws is not None: + with suppress(Exception): + ws.close() + + if not reconnect or idle_timeout is not None: + return + if stop is not None and stop(): + return + time.sleep(1) # brief backoff, then reconnect and keep tailing diff --git a/python/lightning_sdk/cli/deployment/logs.py b/python/lightning_sdk/cli/deployment/logs.py index 99bc5b1b..ca80a59c 100644 --- a/python/lightning_sdk/cli/deployment/logs.py +++ b/python/lightning_sdk/cli/deployment/logs.py @@ -1,284 +1,175 @@ -"""Deployment logs command.""" +"""Deployment logs command. -import json -import threading -import time -from contextlib import suppress -from typing import Any, Iterable, List, Optional, Sequence -from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse +Reads every replica of a deployment, merged into one timeline. The reading and rendering live +in :mod:`lightning_sdk.cli.utils.logs`, shared with the other per-resource log commands. +""" + +from typing import Optional, Sequence -import requests import rich_click as click from lightning_sdk.api.deployment_api import DeploymentApi -from lightning_sdk.api.utils import _get_cloud_url -from lightning_sdk.cli.deployment.common import resolve_deployment +from lightning_sdk.api.job_api import JobApiV2 +from lightning_sdk.api.logs_api import SEVERITIES +from lightning_sdk.cli.deployment.common import resolve_deployment, resolve_teamspace from lightning_sdk.cli.utils.logging import LightningCommand -from lightning_sdk.cli.utils.teamspace_option import resolve_teamspace -from lightning_sdk.lightning_cloud.login import Auth -from lightning_sdk.lightning_cloud.openapi import V1Job +from lightning_sdk.cli.utils.logs import ( + LIVE_FALLBACK_IDLE_TIMEOUT, + LogSelection, + deployment_replica_labels, + read_logs, + resolve_time, +) -_LIVE_FALLBACK_IDLE_TIMEOUT = 8 -_LIVE_FALLBACK_TAIL = 100 +_DEFAULT_TAIL = 100 @click.command("logs", cls=LightningCommand) @click.argument("name") @click.option("--teamspace", help="Override default teamspace (format: owner/teamspace).") @click.option("--job-id", "job_ids", multiple=True, help="Specific deployment job ID. Can be repeated.") -@click.option("--since", help="Only include logs after this timestamp.") -@click.option("--until", help="Only include logs before this timestamp.") -@click.option("--rank", type=int, help="Distributed job rank.") -@click.option("--follow", "-f", is_flag=True, default=False, help="Follow live logs after printing available pages.") -@click.option("--tail", type=int, help="Number of recent live log lines to show when following.") +@click.option("--since", help='Only include lines at or after this time (e.g. "2h", RFC3339).') +@click.option("--until", help='Only include lines at or before this time (e.g. "30m", RFC3339).') +@click.option("--query", help="Only include lines containing every whitespace-separated term.") +@click.option( + "--severity", + type=click.Choice(SEVERITIES), + help="Only include lines at or above this severity.", +) +@click.option("--rank", type=int, help="Machine of a single replica to read (legacy log path).") +@click.option("--follow", "-f", is_flag=True, default=False, help="Stream new log lines as they are produced.") +@click.option("--tail", type=int, help=f"Only show the last N lines. Defaults to {_DEFAULT_TAIL}.") +@click.option("--timestamps", is_flag=True, default=False, help="Prepend each line with its ISO-8601 timestamp.") def deployment_logs( name: str, teamspace: Optional[str] = None, job_ids: Sequence[str] = (), since: Optional[str] = None, until: Optional[str] = None, + query: Optional[str] = None, + severity: Optional[str] = None, rank: Optional[int] = None, follow: bool = False, tail: Optional[int] = None, + timestamps: bool = False, ) -> None: - """Print deployment logs.""" + """Print deployment logs. + + Reads every replica by default, merged into one timeline and labelled with the replica each + line came from. Pass --job-id (repeatable) to read specific replicas. + """ resolved_teamspace = resolve_teamspace(teamspace) api = DeploymentApi() deployment = resolve_deployment(api, resolved_teamspace.id, name) - jobs = _resolve_jobs(api, resolved_teamspace.id, deployment.id, job_ids) - if not jobs: - click.echo("No jobs found for this deployment.") - return - - auth_header = Auth().authenticate() - session = requests.Session() - session.headers.update({"Authorization": auth_header}) - follow_targets = [] - live_fallback_targets = [] - prefix = len(jobs) > 1 - for job in jobs: - logs = api.get_job_logs( + if rank is not None: + _ranked_logs( + api, resolved_teamspace.id, - job.id, - deployment_id=deployment.id, + deployment.id, + job_ids, since=since, until=until, rank=rank, - ) - printed_lines = _print_pages(session, job, logs.pages or [], prefix=prefix) - if follow: - follow_targets.append((job, logs.follow_url)) - elif printed_lines == 0: - live_fallback_targets.append((job, logs.follow_url)) - - if follow: - _stream_jobs( - resolved_teamspace.id, - follow_targets, - auth_header, - follow=True, - idle_timeout=None, - rank=rank, + follow=follow, tail=tail, - prefix=prefix, - ) - elif live_fallback_targets: - _stream_jobs( - resolved_teamspace.id, - live_fallback_targets, - auth_header, - follow=True, - idle_timeout=_LIVE_FALLBACK_IDLE_TIMEOUT, - rank=rank, - tail=tail or _LIVE_FALLBACK_TAIL, - prefix=prefix, + timestamps=timestamps, ) + return - -def _resolve_jobs( + selected = list(job_ids) + if not selected: + jobs = api.list_deployment_jobs(resolved_teamspace.id, deployment.id, limit=100) + if not jobs: + click.echo("No jobs found for this deployment.") + return + labels = {job.id: job.name or job.id for job in jobs} if len(jobs) > 1 else {} + else: + labels = deployment_replica_labels(resolved_teamspace.id, deployment.id) if len(selected) > 1 else {} + + read_logs( + LogSelection( + teamspace_id=resolved_teamspace.id, + job_ids=selected, + # Selecting the deployment picks up replicas that start later; a job id list is fixed. + deployment_id=None if selected else deployment.id, + labels=labels, + ), + query=query, + severity=severity, + since=resolve_time(since, "--since"), + until=resolve_time(until, "--until"), + # A deployment's history can span months of replicas. With no tail and no range asked + # for, show the recent tail rather than paging from the beginning of time. + tail=_DEFAULT_TAIL if tail is None and since is None and until is None else tail, + follow=follow, + timestamps=timestamps, + ) + + +def _ranked_logs( api: DeploymentApi, teamspace_id: str, deployment_id: str, job_ids: Sequence[str], -) -> List[V1Job]: - if job_ids: - all_jobs = api.list_deployment_jobs(teamspace_id, deployment_id, limit=100) - jobs_by_id = {job.id: job for job in all_jobs} - return [ - jobs_by_id.get(job_id) or V1Job(id=job_id, name=job_id, deployment_id=deployment_id) for job_id in job_ids - ] - return api.list_deployment_jobs(teamspace_id, deployment_id, limit=100) - - -def _print_pages(session: requests.Session, job: V1Job, pages: Iterable[Any], *, prefix: bool) -> int: - lines = 0 - for page in pages: - url = getattr(page, "url", None) - if not url: - continue - response = session.get(_absolute_url(url), timeout=60) - response.raise_for_status() - lines += _print_page_text(job, response.text, prefix=prefix) - return lines - - -def _print_page_text(job: V1Job, text: str, *, prefix: bool) -> int: - try: - payload = json.loads(text) - except json.JSONDecodeError: - lines = 0 - for line in text.splitlines(): - _echo_log_line(job, line, prefix=prefix) - lines += 1 - return lines - - entries = payload if isinstance(payload, list) else [payload] - lines = 0 - for entry in entries: - if isinstance(entry, dict): - _echo_log_line(job, entry.get("message") or entry.get("Message") or json.dumps(entry), prefix=prefix) - else: - _echo_log_line(job, str(entry), prefix=prefix) - lines += 1 - return lines - - -def _stream_jobs( - teamspace_id: str, - jobs: Sequence[tuple[V1Job, Optional[str]]], - auth_header: str, *, + since: Optional[str], + until: Optional[str], + rank: int, follow: bool, - idle_timeout: Optional[float], - rank: Optional[int], tail: Optional[int], - prefix: bool, + timestamps: bool, ) -> None: - try: - import websocket - from websocket import WebSocketConnectionClosedException, WebSocketTimeoutException - except ImportError as ex: - raise click.ClickException("Following logs requires the websocket-client package.") from ex - - stop_event = threading.Event() - sockets = [] - threads = [] - errors = [] + """Read one machine of one replica over the legacy per-job log path. - def worker(job: V1Job, follow_url: Optional[str]) -> None: - url = _follow_url(follow_url, teamspace_id, job.id, follow=follow, rank=rank, tail=tail) - while not stop_event.is_set(): - ws = None - kwargs = {"header": [f"Authorization: {auth_header}"]} - if idle_timeout is not None: - kwargs["timeout"] = idle_timeout - try: - ws = websocket.create_connection(_websocket_url(url), **kwargs) - sockets.append(ws) - while not stop_event.is_set(): - try: - message = ws.recv() - except WebSocketTimeoutException: - if idle_timeout is not None: - stop_event.set() - break - raise - except WebSocketConnectionClosedException: - if idle_timeout is not None: - stop_event.set() - break - _print_websocket_message(job, message, prefix=prefix) - if idle_timeout is not None: - break - except WebSocketConnectionClosedException: - if idle_timeout is not None: - stop_event.set() - break - except Exception as ex: - if not stop_event.is_set(): - errors.append(ex) - stop_event.set() - finally: - if ws is not None: - with suppress(Exception): - ws.close() + --rank selects a machine inside a replica, which the merged logs API has no equivalent for: + it returns every machine's lines tagged with the replica instead. Until it does, a ranked + read uses the older per-job endpoints, and those serve a single replica at a time. + """ + if len(job_ids) > 1: + raise click.ClickException("--rank reads one replica at a time; pass a single --job-id.") - if follow and not stop_event.is_set(): - time.sleep(1) + if job_ids: + job_id = job_ids[0] + else: + jobs = api.list_deployment_jobs(teamspace_id, deployment_id, limit=100) + if not jobs: + click.echo("No jobs found for this deployment.") + return + if len(jobs) > 1: + raise click.ClickException("This deployment has several replicas; pass --job-id to pick one for --rank.") + job_id = jobs[0].id + + entries = list( + api.iter_job_log_entries( + teamspace_id, + job_id, + deployment_id=deployment_id, + since=resolve_time(since, "--since"), + until=resolve_time(until, "--until"), + rank=rank, + ) + ) + if tail is not None: + entries = entries[-tail:] + for entry in entries: + click.echo(entry.format(timestamps=timestamps)) - for job, follow_url in jobs: - thread = threading.Thread(target=worker, args=(job, follow_url), daemon=True) - thread.start() - threads.append(thread) + if not follow and entries: + return try: - while any(thread.is_alive() for thread in threads): - if errors: - for ws in sockets: - ws.close() - for thread in threads: - thread.join(timeout=0.2) + for line in JobApiV2().stream_logs( + job_id=job_id, + teamspace_id=teamspace_id, + follow=follow, + tail=tail, + rank=rank, + idle_timeout=None if follow else LIVE_FALLBACK_IDLE_TIMEOUT, + timestamps=timestamps, + ): + click.echo(line) except KeyboardInterrupt: - stop_event.set() - for ws in sockets: - ws.close() - if errors: - raise click.ClickException(str(errors[0])) - - -def _print_websocket_message(job: V1Job, message: str, *, prefix: bool) -> None: - try: - payload = json.loads(message) - except json.JSONDecodeError: - _echo_log_line(job, message, prefix=prefix) - return - - entries = payload if isinstance(payload, list) else [payload] - for entry in entries: - if isinstance(entry, dict): - _echo_log_line(job, entry.get("message") or entry.get("Message") or json.dumps(entry), prefix=prefix) - else: - _echo_log_line(job, str(entry), prefix=prefix) - - -def _echo_log_line(job: V1Job, line: str, *, prefix: bool) -> None: - if prefix: - click.echo(f"[{job.name or job.id}] {line}") - else: - click.echo(line) - - -def _absolute_url(url: str) -> str: - if url.startswith(("http://", "https://", "ws://", "wss://")): - return url - return f"{_get_cloud_url().rstrip('/')}/{url.lstrip('/')}" - - -def _websocket_url(url: str) -> str: - absolute = _absolute_url(url) - if absolute.startswith("https://"): - return "wss://" + absolute[len("https://") :] - if absolute.startswith("http://"): - return "ws://" + absolute[len("http://") :] - return absolute - - -def _follow_url( - follow_url: Optional[str], - teamspace_id: str, - job_id: str, - *, - follow: bool, - rank: Optional[int], - tail: Optional[int], -) -> str: - url = follow_url or f"/v1/projects/{teamspace_id}/jobs/{job_id}/logs" - parsed = urlparse(url) - query = dict(parse_qsl(parsed.query, keep_blank_values=True)) - query.update({"follow": str(follow).lower(), "direction": "forward"}) - if rank is not None: - query["rank"] = str(rank) - if tail is not None: - query["tail"] = str(tail) - return urlunparse(parsed._replace(query=urlencode(query))) + pass + except RuntimeError as ex: + raise click.ClickException(str(ex)) from ex diff --git a/python/lightning_sdk/cli/job/logs.py b/python/lightning_sdk/cli/job/logs.py index de003c4b..467ea2ab 100644 --- a/python/lightning_sdk/cli/job/logs.py +++ b/python/lightning_sdk/cli/job/logs.py @@ -4,8 +4,10 @@ import rich_click as click +from lightning_sdk.api.logs_api import SEVERITIES from lightning_sdk.cli.legacy.job_and_mmt_action import _JobAndMMTAction from lightning_sdk.cli.utils.logging import LightningCommand +from lightning_sdk.cli.utils.logs import resolve_time @click.command("logs", cls=LightningCommand) @@ -23,6 +25,15 @@ @click.option("--tail", type=int, default=None, help="Only show the last N lines.") @click.option("--rank", type=int, default=None, help="Distributed job rank to read from (running jobs only).") @click.option("--timestamps", is_flag=True, default=False, help="Prepend each line with its ISO-8601 timestamp.") +@click.option("--since", default=None, help='Only include lines at or after this time (e.g. "2h", RFC3339).') +@click.option("--until", default=None, help='Only include lines at or before this time (e.g. "30m", RFC3339).') +@click.option("--query", default=None, help="Only include lines containing every whitespace-separated term.") +@click.option( + "--severity", + type=click.Choice(SEVERITIES), + default=None, + help="Only include lines at or above this severity.", +) def logs_job( name: Optional[str] = None, teamspace: Optional[str] = None, @@ -30,16 +41,30 @@ def logs_job( tail: Optional[int] = None, rank: Optional[int] = None, timestamps: bool = False, + since: Optional[str] = None, + until: Optional[str] = None, + query: Optional[str] = None, + severity: Optional[str] = None, ) -> None: """Print the logs for a job. Prints a snapshot of the logs available so far. Pass --follow to stream new - lines from a running job until it finishes or you press Ctrl-C. + lines from a running job until it finishes or you press Ctrl-C. --query and + --severity are applied by the server, to both the snapshot and the stream. """ job = _JobAndMMTAction().job(name=name, teamspace=teamspace) try: - logs = job.logs(follow=follow, tail=tail, rank=rank, timestamps=timestamps) + logs = job.logs( + follow=follow, + tail=tail, + rank=rank, + timestamps=timestamps, + since=resolve_time(since, "--since"), + until=resolve_time(until, "--until"), + query=query, + severity=severity, + ) if follow: for line in logs: click.echo(line) diff --git a/python/lightning_sdk/cli/mmt/__init__.py b/python/lightning_sdk/cli/mmt/__init__.py index db01a2d1..ac773317 100644 --- a/python/lightning_sdk/cli/mmt/__init__.py +++ b/python/lightning_sdk/cli/mmt/__init__.py @@ -8,11 +8,13 @@ def register_commands(group: click.Group) -> None: from lightning_sdk.cli.mmt.delete import delete_mmt from lightning_sdk.cli.mmt.inspect import inspect_mmt from lightning_sdk.cli.mmt.list import list_mmts + from lightning_sdk.cli.mmt.logs import logs_mmt from lightning_sdk.cli.mmt.run import run_mmt from lightning_sdk.cli.mmt.stop import stop_mmt group.add_command(run_mmt, name="run") group.add_command(list_mmts, name="list") group.add_command(inspect_mmt, name="inspect") + group.add_command(logs_mmt, name="logs") group.add_command(stop_mmt, name="stop") group.add_command(delete_mmt, name="delete") diff --git a/python/lightning_sdk/cli/mmt/logs.py b/python/lightning_sdk/cli/mmt/logs.py new file mode 100644 index 00000000..9113469e --- /dev/null +++ b/python/lightning_sdk/cli/mmt/logs.py @@ -0,0 +1,73 @@ +"""MMT logs command.""" + +from typing import Optional + +import rich_click as click + +from lightning_sdk.api.logs_api import SEVERITIES +from lightning_sdk.cli.legacy.job_and_mmt_action import _JobAndMMTAction +from lightning_sdk.cli.utils.logging import LightningCommand +from lightning_sdk.cli.utils.logs import resolve_time + + +@click.command("logs", cls=LightningCommand) +@click.argument("name", required=False) +@click.option( + "--teamspace", + default=None, + help=( + "the name of the teamspace the multi-machine job lives in. " + "Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace). " + "If not specified can be selected interactively." + ), +) +@click.option("--follow", "-f", is_flag=True, default=False, help="Stream new log lines as they are produced.") +@click.option("--tail", type=int, default=None, help="Only show the last N lines.") +@click.option("--timestamps", is_flag=True, default=False, help="Prepend each line with its ISO-8601 timestamp.") +@click.option("--since", default=None, help='Only include lines at or after this time (e.g. "2h", RFC3339).') +@click.option("--until", default=None, help='Only include lines at or before this time (e.g. "30m", RFC3339).') +@click.option("--query", default=None, help="Only include lines containing every whitespace-separated term.") +@click.option( + "--severity", + type=click.Choice(SEVERITIES), + default=None, + help="Only include lines at or above this severity.", +) +def logs_mmt( + name: Optional[str] = None, + teamspace: Optional[str] = None, + follow: bool = False, + tail: Optional[int] = None, + timestamps: bool = False, + since: Optional[str] = None, + until: Optional[str] = None, + query: Optional[str] = None, + severity: Optional[str] = None, +) -> None: + """Print the logs for a multi-machine job. + + Reads every machine, merged into one timeline and labelled with the machine each line came + from. Pass --follow to stream new lines until the job finishes or you press Ctrl-C. To read a + single machine, use `lightning job logs `. + """ + mmt = _JobAndMMTAction().mmt(name=name, teamspace=teamspace) + + try: + logs = mmt.logs( + follow=follow, + tail=tail, + timestamps=timestamps, + since=resolve_time(since, "--since"), + until=resolve_time(until, "--until"), + query=query, + severity=severity, + ) + if follow: + for line in logs: + click.echo(line) + elif logs: + click.echo(logs) + except KeyboardInterrupt: + pass + except RuntimeError as ex: + raise click.ClickException(str(ex)) from ex diff --git a/python/lightning_sdk/cli/utils/logs.py b/python/lightning_sdk/cli/utils/logs.py new file mode 100644 index 00000000..7118f947 --- /dev/null +++ b/python/lightning_sdk/cli/utils/logs.py @@ -0,0 +1,150 @@ +"""Shared pieces for the per-resource log commands. + +Logs are read one resource at a time — ``lightning job logs``, ``lightning mmt logs``, +``lightning deployment logs``, ``lightning sandbox logs`` — and this module holds what more than +one of them needs: time-bound parsing, query highlighting, and the read-and-print loop for a +command that merges several resources into one stream. + +The client itself is :class:`~lightning_sdk.api.logs_api.LogsApi`, and single-resource commands +go through their SDK object (``Job.logs``, ``MMT.logs``) rather than this reader. +""" + +import re +from dataclasses import dataclass, field, replace +from datetime import datetime, timedelta, timezone +from typing import Dict, Optional, Sequence + +import rich_click as click + +from lightning_sdk.api.logs_api import LogsApi + +# Lightning brand purple (#a78bfa) as an RGB tuple for truecolor styling. +_MATCH_COLOR = (167, 139, 250) + +# A resource whose logs are not in the current storage format has no saved lines to print. +# Rather than show nothing, briefly tail the live stream and stop once it goes quiet. +LIVE_FALLBACK_IDLE_TIMEOUT = 8 + +_RELATIVE_TIME = re.compile(r"^(\d+)([smhdw])$") +_RELATIVE_UNITS = {"s": "seconds", "m": "minutes", "h": "hours", "d": "days", "w": "weeks"} + + +def highlight(text: str, query: Optional[str]) -> str: + """Wrap case-insensitive occurrences of ``query`` in ``text`` with a match style. + + click.echo strips these ANSI codes automatically when stdout is not a terminal, so + piped/redirected output stays plain. + """ + if not query: + return text + return re.sub( + re.escape(query), + lambda m: click.style(m.group(0), fg=_MATCH_COLOR, bold=True), + text, + flags=re.IGNORECASE, + ) + + +def resolve_time(value: Optional[str], flag: str) -> Optional[str]: + """Turn a relative bound like ``2h`` into the RFC3339 timestamp the API expects. + + An RFC3339 value is passed through. Anything else is rejected here: the server ignores a + bound it cannot parse, which would silently widen the read instead of failing. + """ + if not value: + return None + match = _RELATIVE_TIME.match(value.strip()) + if match: + amount, unit = match.groups() + delta = timedelta(**{_RELATIVE_UNITS[unit]: int(amount)}) + return (datetime.now(timezone.utc) - delta).isoformat() + try: + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + raise click.UsageError( + f"{flag} must be a duration like 30s/5m/2h/3d/1w, or an RFC3339 timestamp. Got {value!r}." + ) from None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.isoformat() + + +@dataclass +class LogSelection: + """A resolved set of log sources, plus display names for the resources behind them.""" + + teamspace_id: str + job_ids: Sequence[str] = () + deployment_id: Optional[str] = None + mmt_id: Optional[str] = None + sandbox_id: Optional[str] = None + sandbox_command_ids: Sequence[str] = () + # resource id -> label, used to mark which replica or machine a line came from. Empty when + # only one resource can appear, so single-resource output stays unprefixed. + labels: Dict[str, str] = field(default_factory=dict) + + def selectors(self) -> Dict[str, object]: + return { + "job_ids": list(self.job_ids), + "deployment_id": self.deployment_id, + "mmt_id": self.mmt_id, + "sandbox_id": self.sandbox_id, + "sandbox_command_ids": list(self.sandbox_command_ids), + } + + +def read_logs( + selection: LogSelection, + *, + query: Optional[str] = None, + severity: Optional[str] = None, + since: Optional[str] = None, + until: Optional[str] = None, + tail: Optional[int] = None, + follow: bool = False, + timestamps: bool = False, + tail_anchor: Optional[object] = None, +) -> None: + """Read and print logs for ``selection``, labelling lines by the resource they came from. + + The history is paginated automatically, then the live stream is tailed when following. + ``tail`` searches back in widening windows instead of reading a long history in full. + """ + printed = 0 + try: + entries = LogsApi().stream( + selection.teamspace_id, + since=since, + until=until, + query=query, + severity=severity, + follow=follow, + tail=tail, + tail_anchor=tail_anchor, + idle_timeout=None if follow else LIVE_FALLBACK_IDLE_TIMEOUT, + # Nothing saved yet for a running resource: tail its live stream so a snapshot still + # shows something. + fallback_to_live=not follow, + **selection.selectors(), + ) + for entry in entries: + label = selection.labels.get(entry.resource_id) if selection.labels else None + # Highlighting the message before formatting keeps the timestamp and label out of it. + line = replace(entry, message=highlight(entry.message, query)).format(timestamps=timestamps, prefix=label) + click.echo(line) + printed += 1 + except KeyboardInterrupt: + pass + except RuntimeError as ex: + raise click.ClickException(str(ex)) from ex + + if not printed: + click.echo("No logs matched.", err=True) + + +def deployment_replica_labels(teamspace_id: str, deployment_id: str) -> Dict[str, str]: + """Map a deployment's replica job ids to their names, for labelling merged output.""" + from lightning_sdk.api.deployment_api import DeploymentApi + + jobs = DeploymentApi().list_deployment_jobs(teamspace_id, deployment_id, limit=100) + return {job.id: job.name or job.id for job in jobs} if len(jobs) > 1 else {} diff --git a/python/lightning_sdk/job.py b/python/lightning_sdk/job.py index efe634cd..a77cc426 100644 --- a/python/lightning_sdk/job.py +++ b/python/lightning_sdk/job.py @@ -4,6 +4,7 @@ from lightning_sdk.api.cloud_account_api import CloudAccountApi from lightning_sdk.api.job_api import JobApiV2 +from lightning_sdk.api.logs_api import LogsApi from lightning_sdk.api.utils import AccessibleResource, _get_cloud_url, raise_access_error_if_not_allowed from lightning_sdk.status import Status from lightning_sdk.utils.logging import TrackCallsMeta @@ -57,8 +58,21 @@ def __call__( tail: Optional[int] = None, rank: Optional[int] = None, timestamps: bool = False, + since: Optional[str] = None, + until: Optional[str] = None, + query: Optional[str] = None, + severity: Optional[str] = None, ) -> Union[str, Iterator[str]]: - return self._fetch(follow=follow, tail=tail, rank=rank, timestamps=timestamps) + return self._fetch( + follow=follow, + tail=tail, + rank=rank, + timestamps=timestamps, + since=since, + until=until, + query=query, + severity=severity, + ) def _text(self) -> str: if self._cached is None: @@ -149,6 +163,7 @@ def __init__( self._prevent_refetch_latest = False self._cloud_account_api = CloudAccountApi() self._job_api = JobApiV2() + self._logs_api = LogsApi() if _fetch_job: from lightning_sdk.lightning_cloud.openapi.rest import ApiException @@ -536,6 +551,13 @@ def logs(self) -> _Logs: - ``tail``: Only include the last N lines. - ``rank``: Distributed job rank to read from (running jobs only). - ``timestamps``: Prepend each line with its ISO-8601 timestamp. + - ``since``/``until``: Only include lines within this RFC3339 time range. + - ``query``: Only include lines containing every whitespace-separated term. + - ``severity``: Only include lines at or above this level (``error``, ``warning``, + ``info`` or ``debug``). + + ``since``, ``until``, ``query`` and ``severity`` are applied by the server, to both the + saved logs and the live stream. """ return _Logs(self._compute_logs) @@ -546,42 +568,116 @@ def _compute_logs( tail: Optional[int] = None, rank: Optional[int] = None, timestamps: bool = False, + since: Optional[str] = None, + until: Optional[str] = None, + query: Optional[str] = None, + severity: Optional[str] = None, ) -> Union[str, Iterator[str]]: """Fetch the logs, dispatching on job state. See :attr:`logs` for the public API.""" status = self.status + if rank is not None: + # Reading one machine of a multi-machine job has no equivalent on the logs API, so a + # ranked read stays on the older per-job log path. + return self._compute_logs_ranked(status=status, follow=follow, tail=tail, rank=rank, timestamps=timestamps) + + if status not in (Status.Running, Status.Failed, Status.Completed, Status.Stopped): + raise RuntimeError(f"Logs are not available while the job is {status}.") + # Live-follow only makes sense while the job is running. For a finished job there is - # nothing more to stream, so we fall through and return the saved lines rather than - # opening a websocket that would never receive anything (and would reconnect forever). + # nothing more to stream, so we return the saved lines rather than opening a websocket + # that would never receive anything (and would reconnect forever). if follow and status == Status.Running: - return self._stream_logs(follow=True, tail=tail, rank=rank, timestamps=timestamps) + return self._stream_entries(follow=True, tail=tail, timestamps=timestamps, query=query, severity=severity) - if status in (Status.Failed, Status.Completed, Status.Stopped): - if rank is not None: - warnings.warn( - "`rank` is only supported for running jobs; ignoring it for finished-job logs.", - stacklevel=2, - ) + lines = list( + self._stream_entries(follow=False, tail=tail, timestamps=timestamps, query=query, severity=severity) + ) + if not lines and status != Status.Running and query is None and severity is None: + # Nothing in the current log format for this job: fall back to the saved log file. text = self._job_api.get_logs_finished(job_id=self._guaranteed_job.id, teamspace_id=self.teamspace.id) - elif status == Status.Running: - # a running job's stream has no clean end signal, so read until it goes idle - text = "\n".join( - self._stream_logs( - follow=False, - tail=tail, - rank=rank, - timestamps=timestamps, - idle_timeout=_RUNNING_LOGS_IDLE_TIMEOUT, - ) + if tail is not None: + text = "\n".join(text.splitlines()[-tail:]) + return iter(text.splitlines()) if follow else text + + # `follow` on an already-finished job returns the saved lines as an iterator (and stops), + # keeping the return type consistent with the live-follow path. + return iter(lines) if follow else "\n".join(lines) + + def _compute_logs_ranked( + self, + *, + status: Status, + follow: bool, + tail: Optional[int], + rank: int, + timestamps: bool, + ) -> Union[str, Iterator[str]]: + """Read a single machine's logs over the legacy per-job websocket.""" + if status in (Status.Failed, Status.Completed, Status.Stopped): + warnings.warn( + "`rank` is only supported for running jobs; ignoring it for finished-job logs.", + stacklevel=3, ) - else: + text = self._job_api.get_logs_finished(job_id=self._guaranteed_job.id, teamspace_id=self.teamspace.id) + if tail is not None: + text = "\n".join(text.splitlines()[-tail:]) + return iter(text.splitlines()) if follow else text + + if status != Status.Running: raise RuntimeError(f"Logs are not available while the job is {status}.") + if follow: + return self._stream_logs(follow=True, tail=tail, rank=rank, timestamps=timestamps) + + # a running job's stream has no clean end signal, so read until it goes idle + text = "\n".join( + self._stream_logs( + follow=False, + tail=tail, + rank=rank, + timestamps=timestamps, + idle_timeout=_RUNNING_LOGS_IDLE_TIMEOUT, + ) + ) if tail is not None: text = "\n".join(text.splitlines()[-tail:]) - # `follow` on an already-finished job returns the saved lines as an iterator (and stops), - # keeping the return type consistent with the live-follow path. - return iter(text.splitlines()) if follow else text + return text + + def _stream_entries( + self, + *, + follow: bool, + tail: Optional[int], + timestamps: bool, + since: Optional[str] = None, + until: Optional[str] = None, + query: Optional[str] = None, + severity: Optional[str] = None, + ) -> Iterator[str]: + """Yield formatted log lines from the logs API (saved logs, then the live tail).""" + job = self._guaranteed_job + job_id = job.id + entries = self._logs_api.stream( + self.teamspace.id, + job_ids=[job_id], + since=since, + until=until, + query=query, + severity=severity, + follow=follow, + tail=tail, + # A finished job's last lines sit at its stop time, so start the tail search there + # instead of walking back from now through a job that ran days ago. + tail_anchor=getattr(job, "stopped_at", None), + idle_timeout=None if follow else _RUNNING_LOGS_IDLE_TIMEOUT, + # A running job whose logs are not in the current storage format yet has no saved + # history; tail its live stream so a snapshot still shows something. + fallback_to_live=not follow, + stop=lambda: self._job_api._is_job_finished(job_id, self.teamspace.id), + ) + for entry in entries: + yield entry.format(timestamps=timestamps) def _stream_logs( self, diff --git a/python/lightning_sdk/mmt.py b/python/lightning_sdk/mmt.py index 2bccd524..46749392 100644 --- a/python/lightning_sdk/mmt.py +++ b/python/lightning_sdk/mmt.py @@ -1,9 +1,11 @@ import warnings -from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol, Tuple, TypedDict, Union +from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Protocol, Tuple, TypedDict, Union from lightning_sdk.api.cloud_account_api import CloudAccountApi +from lightning_sdk.api.logs_api import LogsApi from lightning_sdk.api.mmt_api import MMTApiV2 from lightning_sdk.api.utils import AccessibleResource, _get_cloud_url, raise_access_error_if_not_allowed +from lightning_sdk.job import _RUNNING_LOGS_IDLE_TIMEOUT, _Logs from lightning_sdk.status import Status from lightning_sdk.utils.logging import TrackCallsMeta from lightning_sdk.utils.resolve import ( @@ -154,6 +156,7 @@ def __init__( self._prevent_refetch_latest = False self._cloud_account_api = CloudAccountApi() self._job_api = MMTApiV2() + self._logs_api = LogsApi() if _fetch_job: self._update_internal_job() @@ -480,6 +483,10 @@ def _update_internal_job(self) -> None: def name(self) -> str: return self._name + @property + def resource_id(self) -> Optional[str]: + return self._guaranteed_job.id + @property def teamspace(self) -> "Teamspace": return self._teamspace @@ -510,8 +517,102 @@ def num_machines(self) -> int: return self._job_api.get_num_machines(self._guaranteed_job) @property - def logs(self) -> str: - raise NotImplementedError + def logs(self) -> _Logs: + """The logs of every machine, merged into one timeline. + + Use it as a value for a snapshot of the logs up to now:: + + print(mmt.logs) + + or call it to pass options and/or follow the logs live:: + + recent = mmt.logs(tail=100) # snapshot of the last 100 lines + for line in mmt.logs(follow=True): # stream new lines as they arrive + print(line) + + Options: + + - ``follow``: Keep the stream open and yield new lines as they are produced. + Returns an iterator of lines instead of a string. + - ``tail``: Only include the last N lines. + - ``timestamps``: Prepend each line with its ISO-8601 timestamp. + - ``since``/``until``: Only include lines within this RFC3339 time range. + - ``query``: Only include lines containing every whitespace-separated term. + - ``severity``: Only include lines at or above this level (``error``, ``warning``, + ``info`` or ``debug``). + + Every line is labelled with the machine it came from. To read a single machine, use + ``mmt.machines[rank].logs``. + """ + return _Logs(self._compute_logs) + + def _compute_logs( + self, + *, + follow: bool = False, + tail: Optional[int] = None, + rank: Optional[int] = None, + timestamps: bool = False, + since: Optional[str] = None, + until: Optional[str] = None, + query: Optional[str] = None, + severity: Optional[str] = None, + ) -> Union[str, Iterator[str]]: + """Fetch the merged logs of every machine. See :attr:`logs` for the public API.""" + if rank is not None: + raise ValueError("`rank` is not supported here; read a single machine with `mmt.machines[rank].logs`.") + + status = self.status + if status not in (Status.Running, Status.Failed, Status.Completed, Status.Stopped): + raise RuntimeError(f"Logs are not available while the job is {status}.") + + lines = self._stream_entries( + follow=follow and status == Status.Running, + tail=tail, + timestamps=timestamps, + since=since, + until=until, + query=query, + severity=severity, + ) + if follow and status == Status.Running: + return lines + collected = list(lines) + return iter(collected) if follow else "\n".join(collected) + + def _stream_entries( + self, + *, + follow: bool, + tail: Optional[int], + timestamps: bool, + since: Optional[str] = None, + until: Optional[str] = None, + query: Optional[str] = None, + severity: Optional[str] = None, + ) -> Iterator[str]: + """Yield formatted log lines for every machine, labelled with the machine they came from.""" + names = {machine._guaranteed_job.id: machine.name for machine in self.machines} + entries = self._logs_api.stream( + self.teamspace.id, + mmt_id=self._guaranteed_job.id, + since=since, + until=until, + query=query, + severity=severity, + follow=follow, + tail=tail, + # A finished job's last lines sit at its stop time, so start the tail search there + # instead of walking back from now through a job that ran days ago. + tail_anchor=getattr(self._guaranteed_job, "stopped_at", None), + idle_timeout=None if follow else _RUNNING_LOGS_IDLE_TIMEOUT, + # A running job whose logs are not in the current storage format yet has no saved + # history; tail its live stream so a snapshot still shows something. + fallback_to_live=not follow, + stop=lambda: self.status in (Status.Stopped, Status.Completed, Status.Failed), + ) + for entry in entries: + yield entry.format(timestamps=timestamps, prefix=names.get(entry.resource_id, entry.resource_id)) def dict(self) -> Dict[str, object]: studio = self.studio diff --git a/python/tests/api/test_logs_api.py b/python/tests/api/test_logs_api.py new file mode 100644 index 00000000..e0e974c5 --- /dev/null +++ b/python/tests/api/test_logs_api.py @@ -0,0 +1,360 @@ +from datetime import datetime, timedelta, timezone +from unittest import mock + +import pytest + +from lightning_sdk.api.logs_api import ( + _CONNECT_TIMEOUT, + _FOLLOW_POLL_INTERVAL, + _TAIL_WINDOWS, + LogEntry, + LogsApi, + _websocket_url, + parse_log_entries, +) +from lightning_sdk.lightning_cloud.openapi import V1GetLogsResponse, V1JobLogEntry + + +def _api(*pages): + """A LogsApi whose client returns ``pages`` in order.""" + client = mock.MagicMock() + client.jobs_service_get_logs.side_effect = list(pages) + return LogsApi(client=client), client + + +def _entry(message, *, line=0, resource_id="", severity="", timestamp=None): + return V1JobLogEntry(message=message, line=line, resource_id=resource_id, severity=severity, timestamp=timestamp) + + +def test_parse_log_entries_reads_snake_case_frames() -> None: + entries = parse_log_entries( + '[{"message":"ready","line":3,"resource_id":"job-1","severity":"info",' + '"timestamp":{"seconds":1784000000,"nanos":500000000}}]' + ) + + assert len(entries) == 1 + assert entries[0].message == "ready" + assert entries[0].line == 3 + assert entries[0].resource_id == "job-1" + assert entries[0].severity == "info" + assert entries[0].timestamp == datetime.fromtimestamp(1784000000.5, tz=timezone.utc) + + +def test_parse_log_entries_accepts_camel_case_and_single_objects() -> None: + entries = parse_log_entries('{"message":"ready","resourceId":"job-1"}') + + assert [(e.message, e.resource_id) for e in entries] == [("ready", "job-1")] + + +def test_parse_log_entries_falls_back_to_raw_lines() -> None: + # a frame that is not JSON must still surface rather than being dropped + assert [e.message for e in parse_log_entries("plain\ntext")] == ["plain", "text"] + + +def test_parse_log_entries_tolerates_missing_timestamp() -> None: + assert parse_log_entries('[{"message":"ready"}]')[0].timestamp is None + + +def test_log_entry_format() -> None: + entry = LogEntry(message="ready", timestamp=datetime(2026, 7, 27, 9, 0, tzinfo=timezone.utc)) + + assert entry.format() == "ready" + assert entry.format(prefix="replica-0") == "[replica-0] ready" + assert entry.format(timestamps=True) == "2026-07-27T09:00:00+00:00 ready" + assert entry.format(timestamps=True, prefix="replica-0") == "2026-07-27T09:00:00+00:00 [replica-0] ready" + # an entry with no timestamp is printed as-is rather than gaining an empty column + assert LogEntry(message="ready").format(timestamps=True) == "ready" + + +def test_websocket_url_preserves_absolute_wss_url() -> None: + url = "wss://lightning.ai/v1/projects/project-id/logs?follow=true" + + assert _websocket_url(url) == url + + +def test_websocket_url_upgrades_scheme_and_relative_paths() -> None: + assert _websocket_url("https://lightning.ai/v1/logs") == "wss://lightning.ai/v1/logs" + assert _websocket_url("http://localhost:8080/v1/logs") == "ws://localhost:8080/v1/logs" + assert _websocket_url("/v1/projects/p/logs").endswith("/v1/projects/p/logs") + assert _websocket_url("/v1/projects/p/logs").startswith("ws") + + +def test_get_page_maps_selectors_and_filters() -> None: + api, client = _api(V1GetLogsResponse(entries=[], follow_url="wss://host/logs")) + + page = api.get_page( + "project-id", + job_ids=["job-1", "job-2"], + query="boom", + severity="error", + since="2026-07-27T00:00:00Z", + page_size=250, + ) + + assert page.follow_url == "wss://host/logs" + kwargs = client.jobs_service_get_logs.call_args.kwargs + assert kwargs["project_id"] == "project-id" + assert kwargs["job_ids"] == ["job-1", "job-2"] + assert kwargs["query"] == "boom" + assert kwargs["severity"] == "error" + assert kwargs["since"] == "2026-07-27T00:00:00Z" + assert kwargs["page_size"] == "250" + # unset options must not be sent at all + assert "until" not in kwargs + assert "deployment_id" not in kwargs + assert "page_token" not in kwargs + + +def test_stream_requires_a_selector() -> None: + api, _ = _api() + + with pytest.raises(ValueError, match="job_ids"): + list(api.stream("project-id")) + + +def test_stream_follows_page_tokens() -> None: + api, client = _api( + V1GetLogsResponse(entries=[_entry("one")], next_page_token="cursor-1"), + V1GetLogsResponse(entries=[_entry("two")], next_page_token=""), + ) + + entries = list(api.stream("project-id", deployment_id="dep-id")) + + assert [e.message for e in entries] == ["one", "two"] + assert client.jobs_service_get_logs.call_count == 2 + assert client.jobs_service_get_logs.call_args_list[1].kwargs["page_token"] == "cursor-1" + + +def test_stream_tail_keeps_the_last_lines_across_pages() -> None: + api, _ = _api( + V1GetLogsResponse(entries=[_entry("a"), _entry("b")], next_page_token="cursor-1"), + V1GetLogsResponse(entries=[_entry("c"), _entry("d")]), + ) + + entries = list(api.stream("project-id", mmt_id="mmt-id", tail=3)) + + # the API only pages forward, so the tail is applied once the history is exhausted + assert [e.message for e in entries] == ["b", "c", "d"] + + +def test_stream_tail_widens_the_window_until_enough_lines() -> None: + api, client = _api( + # nothing in the most recent window... + V1GetLogsResponse(entries=[]), + # ...so the next, wider one is tried + V1GetLogsResponse(entries=[_entry("a"), _entry("b")]), + ) + + entries = list(api.stream("project-id", deployment_id="dep-id", tail=2)) + + assert [e.message for e in entries] == ["a", "b"] + windows = [call.kwargs["since"] for call in client.jobs_service_get_logs.call_args_list] + assert len(windows) == 2 + # the second attempt reaches further back than the first + assert windows[1] < windows[0] + + +def test_stream_tail_falls_back_to_the_full_history() -> None: + # every bounded window comes up short, so the last attempt is unbounded + api, client = _api(*[V1GetLogsResponse(entries=[]) for _ in range(5)], V1GetLogsResponse(entries=[_entry("old")])) + + entries = list(api.stream("project-id", deployment_id="dep-id", tail=5)) + + assert [e.message for e in entries] == ["old"] + assert "since" not in client.jobs_service_get_logs.call_args_list[-1].kwargs + + +def test_stream_tail_anchor_places_the_window_at_the_stop_time() -> None: + api, client = _api(V1GetLogsResponse(entries=[_entry("last line")])) + stopped = datetime(2026, 7, 24, 21, 12, tzinfo=timezone.utc) + + entries = list(api.stream("project-id", job_ids=["job-1"], tail=1, tail_anchor=stopped)) + + assert [e.message for e in entries] == ["last line"] + kwargs = client.jobs_service_get_logs.call_args.kwargs + # the first window ends at the stop time rather than now, so an old job needs one call + assert kwargs["since"] == (stopped - timedelta(seconds=_TAIL_WINDOWS[0])).isoformat() + # the anchor bounds the search only: nothing written after the stop time is filtered out + assert "until" not in kwargs + + +def test_stream_tail_honours_an_explicit_since() -> None: + api, client = _api(V1GetLogsResponse(entries=[_entry("a"), _entry("b"), _entry("c")])) + + entries = list(api.stream("project-id", deployment_id="dep-id", tail=2, since="2026-07-01T00:00:00Z")) + + assert [e.message for e in entries] == ["b", "c"] + # the caller bounded the read themselves, so no window search happens + assert client.jobs_service_get_logs.call_count == 1 + assert client.jobs_service_get_logs.call_args.kwargs["since"] == "2026-07-01T00:00:00Z" + + +def test_stream_does_not_follow_without_a_follow_url() -> None: + api, _ = _api(V1GetLogsResponse(entries=[_entry("done")], follow_url="")) + api.follow = mock.MagicMock() + + entries = list(api.stream("project-id", job_ids=["job-1"], follow=True)) + + assert [e.message for e in entries] == ["done"] + # a finished resource has nothing to tail, so no socket is opened + api.follow.assert_not_called() + + +def test_stream_tails_after_history_when_following() -> None: + api, _ = _api(V1GetLogsResponse(entries=[_entry("saved")], follow_url="wss://host/logs")) + api.follow = mock.MagicMock(return_value=iter([LogEntry(message="live")])) + + entries = list(api.stream("project-id", job_ids=["job-1"], follow=True)) + + assert [e.message for e in entries] == ["saved", "live"] + assert api.follow.call_args.args == ("wss://host/logs",) + assert api.follow.call_args.kwargs["reconnect"] is True + + +def test_stream_live_fallback_only_when_history_is_empty() -> None: + api, _ = _api( + V1GetLogsResponse(entries=[], follow_url="wss://host/logs"), + V1GetLogsResponse(entries=[_entry("saved")], follow_url="wss://host/logs"), + ) + api.follow = mock.MagicMock(return_value=iter([LogEntry(message="live")])) + + empty_history = list(api.stream("project-id", job_ids=["job-1"], fallback_to_live=True, idle_timeout=1)) + assert [e.message for e in empty_history] == ["live"] + # a one-shot fallback must not reconnect once the socket goes quiet + assert api.follow.call_args.kwargs["reconnect"] is False + + api.follow.reset_mock() + with_history = list(api.stream("project-id", job_ids=["job-1"], fallback_to_live=True, idle_timeout=1)) + assert [e.message for e in with_history] == ["saved"] + api.follow.assert_not_called() + + +class _FakeSocket: + def __init__(self, frames): + self._frames = list(frames) + self.closed = False + self.timeout = None + + def settimeout(self, timeout): + self.timeout = timeout + + def recv(self): + if not self._frames: + raise _FakeTimeoutError + frame = self._frames.pop(0) + if isinstance(frame, Exception): + raise frame + return frame + + def close(self): + self.closed = True + + +class _FakeTimeoutError(Exception): + pass + + +class _FakeClosedError(Exception): + pass + + +@pytest.fixture() +def fake_websocket(monkeypatch): + """Stand in for the websocket-client module used by ``LogsApi.follow``.""" + module = mock.MagicMock() + module.WebSocketTimeoutException = _FakeTimeoutError + module.WebSocketConnectionClosedException = _FakeClosedError + monkeypatch.setitem(__import__("sys").modules, "websocket", module) + monkeypatch.setattr("lightning_sdk.api.logs_api.Auth", mock.MagicMock()) + return module + + +def test_follow_yields_frames_until_idle(fake_websocket) -> None: + socket = _FakeSocket(['[{"message":"live-1"},{"message":"live-2"}]']) + fake_websocket.create_connection.return_value = socket + + entries = list(LogsApi(client=mock.MagicMock()).follow("wss://host/logs", idle_timeout=1)) + + assert [e.message for e in entries] == ["live-1", "live-2"] + assert socket.closed + # the handshake gets its own budget; the read poll is set on the socket afterwards + assert fake_websocket.create_connection.call_args.kwargs["timeout"] == _CONNECT_TIMEOUT + assert socket.timeout == 1 + + +def test_follow_polls_below_the_server_heartbeat(fake_websocket) -> None: + fake_websocket.create_connection.return_value = _FakeSocket([]) + + # a heartbeat resets the socket read timeout, so the poll must stay short and the silence + # be measured by the caller instead + with mock.patch("lightning_sdk.api.logs_api.time.monotonic", side_effect=[0.0, 0.0, 61.0]): + list(LogsApi(client=mock.MagicMock()).follow("wss://host/logs", idle_timeout=60)) + + assert fake_websocket.create_connection.return_value.timeout == _FOLLOW_POLL_INTERVAL + + +def test_follow_keeps_waiting_while_lines_keep_arriving(fake_websocket) -> None: + fake_websocket.create_connection.return_value = _FakeSocket( + ['[{"message":"one"}]', _FakeTimeoutError(), '[{"message":"two"}]'] + ) + clock = [0.0, 0.0, 1.0, 1.0, 1.0, 99.0, 99.0, 99.0] + + with mock.patch("lightning_sdk.api.logs_api.time.monotonic", side_effect=clock): + entries = list(LogsApi(client=mock.MagicMock()).follow("wss://host/logs", idle_timeout=5)) + + # the gap between the two lines is under the idle timeout, so the stream survives it + assert [e.message for e in entries] == ["one", "two"] + + +def test_follow_stops_when_the_resource_finishes(fake_websocket) -> None: + fake_websocket.create_connection.return_value = _FakeSocket(['[{"message":"live"}]']) + + entries = list(LogsApi(client=mock.MagicMock()).follow("wss://host/logs", stop=lambda: True)) + + # the socket stays open (heartbeats only) after the resource ends, so `stop` ends the stream + assert [e.message for e in entries] == ["live"] + + +def test_follow_drops_lines_already_seen_in_history(fake_websocket) -> None: + api, _ = _api( + V1GetLogsResponse(entries=[_entry("saved", line=1)], follow_url="wss://host/logs"), + ) + # the live socket starts at "now" and replays the line history just produced + fake_websocket.create_connection.return_value = _FakeSocket( + ['[{"message":"saved","line":1},{"message":"new","line":2}]'] + ) + + entries = list(api.stream("project-id", job_ids=["job-1"], follow=True, idle_timeout=1)) + + assert [e.message for e in entries] == ["saved", "new"] + + +def test_follow_reconnects_after_a_drop(fake_websocket) -> None: + sockets = [ + _FakeSocket(['[{"message":"before"}]', _FakeClosedError()]), + _FakeSocket(['[{"message":"after"}]']), + ] + fake_websocket.create_connection.side_effect = sockets + stop = mock.MagicMock(side_effect=[False, True]) + + with mock.patch("lightning_sdk.api.logs_api.time.sleep"): + entries = list(LogsApi(client=mock.MagicMock()).follow("wss://host/logs", stop=stop)) + + assert [e.message for e in entries] == ["before", "after"] + assert fake_websocket.create_connection.call_count == 2 + + +def test_follow_requires_websocket_client(monkeypatch) -> None: + import builtins + + real_import = builtins.__import__ + + def fake_import(name, *args, **kwargs): + if name == "websocket": + raise ImportError("no websocket-client") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + with pytest.raises(RuntimeError, match="websocket-client"): + list(LogsApi(client=mock.MagicMock()).follow("wss://host/logs")) diff --git a/python/tests/cli/deployment/test_deployment.py b/python/tests/cli/deployment/test_deployment.py index 06510ecb..fa76841f 100644 --- a/python/tests/cli/deployment/test_deployment.py +++ b/python/tests/cli/deployment/test_deployment.py @@ -6,9 +6,10 @@ import pytest from click.testing import CliRunner +from lightning_sdk.api.logs_api import LogEntry from lightning_sdk.cli.deployment.create import create_deployment from lightning_sdk.cli.deployment.list import list_deployments -from lightning_sdk.cli.deployment.logs import _follow_url, _print_page_text, _resolve_jobs, _websocket_url +from lightning_sdk.cli.deployment.logs import deployment_logs from lightning_sdk.cli.deployment.reload_weights import reload_weights from lightning_sdk.lightning_cloud.openapi import ( V1BYOMSpec, @@ -386,42 +387,145 @@ def test_create_old_byom_image_variant_rejected(monkeypatch) -> None: assert result.exit_code != 0 # flag renamed to --serving-image-variant +def _patch_logs_command(monkeypatch, api, entries): + """Wire the deployment logs command to ``api`` and make the logs API yield ``entries``.""" + teamspace = SimpleNamespace(id="project-id") + monkeypatch.setattr( + "lightning_sdk.cli.deployment.logs.resolve_teamspace", + MagicMock(return_value=teamspace), + ) + monkeypatch.setattr("lightning_sdk.cli.deployment.logs.DeploymentApi", MagicMock(return_value=api)) + # the command delegates to the shared reader, which owns the API client + stream = MagicMock(return_value=list(entries)) + monkeypatch.setattr( + "lightning_sdk.cli.utils.logs.LogsApi", + MagicMock(return_value=SimpleNamespace(stream=stream, get_page=MagicMock())), + ) + return stream + + +def _deployment_api_with_replicas(*names): + api = MagicMock() + api.get_deployment_by_name.return_value = V1Deployment(name="my-deployment", id="dep-id", project_id="project-id") + api.list_deployment_jobs.return_value = [ + V1Job(id=f"job-{i}", name=name, deployment_id="dep-id") for i, name in enumerate(names) + ] + return api + + @mock_command_logging -def test_follow_url_preserves_existing_query() -> None: - url = _follow_url( - "/v1/projects/project-id/jobs/job-id/logs?token=abc", - "project-id", - "job-id", - follow=True, - rank=3, - tail=50, +def test_deployment_logs_reads_whole_deployment_and_labels_replicas(monkeypatch) -> None: + api = _deployment_api_with_replicas("replica-0", "replica-1") + stream = _patch_logs_command( + monkeypatch, + api, + [ + LogEntry(message="ready", resource_id="job-0"), + LogEntry(message="serving", resource_id="job-1"), + ], ) - assert url.startswith("/v1/projects/project-id/jobs/job-id/logs?") - assert "token=abc" in url - assert "follow=true" in url - assert "deploymentId" not in url - assert "rank=3" in url - assert "tail=50" in url + result = CliRunner().invoke(deployment_logs, ["my-deployment", "--teamspace", "ecorp/test"]) + + assert result.exit_code == 0, result.output + # one call for every replica, each line labelled with the replica it came from + assert result.output == "[replica-0] ready\n[replica-1] serving\n" + kwargs = stream.call_args.kwargs + assert kwargs["deployment_id"] == "dep-id" + assert kwargs["job_ids"] == [] + # months of replicas can sit behind a deployment, so a plain read shows the recent tail + assert kwargs["tail"] == 100 @mock_command_logging -def test_websocket_url_preserves_absolute_wss_url() -> None: - url = _websocket_url("wss://lightning.ai/v1/projects/project-id/jobs/job-id/logs?follow=true") +def test_deployment_logs_single_replica_is_not_labelled(monkeypatch) -> None: + api = _deployment_api_with_replicas("replica-0") + _patch_logs_command(monkeypatch, api, [LogEntry(message="ready", resource_id="job-0")]) - assert url == "wss://lightning.ai/v1/projects/project-id/jobs/job-id/logs?follow=true" + result = CliRunner().invoke(deployment_logs, ["my-deployment"]) + + assert result.exit_code == 0, result.output + assert result.output == "ready\n" @mock_command_logging -def test_print_page_text_renders_json_log_entries(capsys) -> None: - job = SimpleNamespace(id="job-id", name="replica-0") +def test_deployment_logs_selected_job_ids_and_filters(monkeypatch) -> None: + api = _deployment_api_with_replicas("replica-0", "replica-1") + stream = _patch_logs_command(monkeypatch, api, []) - rendered = _print_page_text(job, '[{"message":"ready"},{"Message":"serving"}]', prefix=False) - empty = _print_page_text(job, "[]", prefix=False) + result = CliRunner().invoke( + deployment_logs, + [ + "my-deployment", + "--job-id", + "job-1", + "--query", + "timeout", + "--severity", + "error", + "--since", + "2026-07-27T00:00:00Z", + "--follow", + ], + ) - assert rendered == 2 - assert empty == 0 - assert capsys.readouterr().out == "ready\nserving\n" + assert result.exit_code == 0, result.output + kwargs = stream.call_args.kwargs + assert kwargs["job_ids"] == ["job-1"] + # a fixed job id list replaces the deployment selector + assert kwargs["deployment_id"] is None + assert kwargs["query"] == "timeout" + assert kwargs["severity"] == "error" + assert kwargs["since"] == "2026-07-27T00:00:00+00:00" + assert kwargs["follow"] is True + assert kwargs["idle_timeout"] is None + + +@mock_command_logging +def test_deployment_logs_rejects_bad_severity(monkeypatch) -> None: + api = _deployment_api_with_replicas("replica-0") + _patch_logs_command(monkeypatch, api, []) + + result = CliRunner().invoke(deployment_logs, ["my-deployment", "--severity", "critical"]) + + assert result.exit_code != 0 + assert "critical" in result.output + + +@mock_command_logging +def test_deployment_logs_reports_no_jobs(monkeypatch) -> None: + api = _deployment_api_with_replicas() + _patch_logs_command(monkeypatch, api, []) + + result = CliRunner().invoke(deployment_logs, ["my-deployment"]) + + assert result.exit_code == 0, result.output + assert "No jobs found for this deployment." in result.output + + +@mock_command_logging +def test_deployment_logs_rank_uses_legacy_path(monkeypatch) -> None: + api = _deployment_api_with_replicas("replica-0") + api.iter_job_log_entries.return_value = iter([LogEntry(message="from rank 2")]) + _patch_logs_command(monkeypatch, api, []) + + result = CliRunner().invoke(deployment_logs, ["my-deployment", "--rank", "2"]) + + assert result.exit_code == 0, result.output + assert result.output == "from rank 2\n" + assert api.iter_job_log_entries.call_args.kwargs["rank"] == 2 + + +@mock_command_logging +def test_deployment_logs_rank_needs_a_single_replica(monkeypatch) -> None: + api = _deployment_api_with_replicas("replica-0", "replica-1") + _patch_logs_command(monkeypatch, api, []) + + result = CliRunner().invoke(deployment_logs, ["my-deployment", "--rank", "0"]) + + assert result.exit_code != 0 + assert "--job-id" in result.output + api.iter_job_log_entries.assert_not_called() @mock_command_logging @@ -447,13 +551,14 @@ def test_reload_weights_calls_api_and_prints_version(monkeypatch) -> None: @mock_command_logging -def test_resolve_jobs_filters_specific_job_id() -> None: - api = MagicMock() - api.list_deployment_jobs.return_value = [ - V1Job(id="job-1", name="replica-0", deployment_id="dep-id"), - V1Job(id="job-2", name="replica-1", deployment_id="dep-id"), - ] - - jobs = _resolve_jobs(api, "project-id", "dep-id", ["job-2"]) - - assert [job.id for job in jobs] == ["job-2"] +def test_deployment_logs_help() -> None: + assert_help_contains( + "lightning deployment logs --help", + "Usage: lightning deployment logs", + "--job-id", + "--query", + "--severity", + "--follow", + "--tail", + "--timestamps", + ) diff --git a/python/tests/cli/job/test_logs.py b/python/tests/cli/job/test_logs.py index 2f470769..ad3cfc32 100644 --- a/python/tests/cli/job/test_logs.py +++ b/python/tests/cli/job/test_logs.py @@ -15,6 +15,8 @@ def test_job_logs_help() -> None: "--tail", "--rank", "--timestamps", + "--query", + "--severity", ) @@ -57,7 +59,9 @@ def test_job_logs_prints_snapshot(monkeypatch) -> None: assert captured == {"name": "my-job", "teamspace": "org/teamspace"} assert "hello from the job" in result.output assert "42" in result.output - job.logs.assert_called_once_with(follow=False, tail=None, rank=None, timestamps=False) + job.logs.assert_called_once_with( + follow=False, tail=None, rank=None, timestamps=False, since=None, until=None, query=None, severity=None + ) @mock_command_logging @@ -76,7 +80,39 @@ def test_job_logs_follows_with_options(monkeypatch) -> None: assert result.exit_code == 0 assert result.output == "line 1\nline 2\n" - job.logs.assert_called_once_with(follow=True, tail=10, rank=2, timestamps=True) + job.logs.assert_called_once_with( + follow=True, tail=10, rank=2, timestamps=True, since=None, until=None, query=None, severity=None + ) + + +@mock_command_logging +def test_job_logs_passes_filters(monkeypatch) -> None: + from lightning_sdk.cli.job.logs import logs_job + + captured: dict = {} + job = MagicMock() + job.logs.return_value = "boom" + _patch_action(monkeypatch, job, captured) + + result = CliRunner().invoke(logs_job, ["my-job", "--query", "boom", "--severity", "error"]) + + assert result.exit_code == 0, result.output + job.logs.assert_called_once_with( + follow=False, tail=None, rank=None, timestamps=False, since=None, until=None, query="boom", severity="error" + ) + + +@mock_command_logging +def test_job_logs_rejects_unknown_severity(monkeypatch) -> None: + from lightning_sdk.cli.job.logs import logs_job + + job = MagicMock() + _patch_action(monkeypatch, job, {}) + + result = CliRunner().invoke(logs_job, ["my-job", "--severity", "critical"]) + + assert result.exit_code != 0 + job.logs.assert_not_called() @mock_command_logging diff --git a/python/tests/cli/mmt/test_logs.py b/python/tests/cli/mmt/test_logs.py new file mode 100644 index 00000000..2812665c --- /dev/null +++ b/python/tests/cli/mmt/test_logs.py @@ -0,0 +1,94 @@ +from unittest.mock import MagicMock + +from click.testing import CliRunner + +from tests.cli.help import assert_help_contains, mock_command_logging + + +@mock_command_logging +def test_mmt_logs_help() -> None: + assert_help_contains( + "lightning mmt logs --help", + "Usage: lightning mmt logs", + "Print the logs for a multi-machine job.", + "--follow", + "--tail", + "--timestamps", + "--query", + "--severity", + ) + + +@mock_command_logging +def test_mmts_logs_help() -> None: + assert_help_contains( + "lightning mmts logs --help", + "Usage: lightning mmts logs", + "Print the logs for a multi-machine job.", + ) + + +def _patch_action(monkeypatch, mmt: MagicMock, captured: dict) -> None: + class _FakeJobAndMMTAction: + def mmt(self, name=None, teamspace=None): + captured["name"] = name + captured["teamspace"] = teamspace + return mmt + + monkeypatch.setattr("lightning_sdk.cli.mmt.logs._JobAndMMTAction", _FakeJobAndMMTAction) + + +@mock_command_logging +def test_mmt_logs_prints_merged_snapshot(monkeypatch) -> None: + from lightning_sdk.cli.mmt.logs import logs_mmt + + captured: dict = {} + mmt = MagicMock() + mmt.logs.return_value = "[my-mmt-0] rank 0 up\n[my-mmt-1] rank 1 up" + _patch_action(monkeypatch, mmt, captured) + + result = CliRunner().invoke(logs_mmt, ["my-mmt", "--teamspace", "org/teamspace"]) + + assert result.exit_code == 0, result.output + assert captured == {"name": "my-mmt", "teamspace": "org/teamspace"} + assert "[my-mmt-0] rank 0 up" in result.output + assert "[my-mmt-1] rank 1 up" in result.output + mmt.logs.assert_called_once_with( + follow=False, tail=None, timestamps=False, since=None, until=None, query=None, severity=None + ) + + +@mock_command_logging +def test_mmt_logs_follows_with_options(monkeypatch) -> None: + from lightning_sdk.cli.mmt.logs import logs_mmt + + captured: dict = {} + mmt = MagicMock() + mmt.logs.return_value = iter(["line 1", "line 2"]) + _patch_action(monkeypatch, mmt, captured) + + result = CliRunner().invoke( + logs_mmt, + ["my-mmt", "--follow", "--tail", "10", "--timestamps", "--query", "loss", "--severity", "error"], + ) + + assert result.exit_code == 0, result.output + assert result.output == "line 1\nline 2\n" + mmt.logs.assert_called_once_with( + follow=True, tail=10, timestamps=True, since=None, until=None, query="loss", severity="error" + ) + + +@mock_command_logging +def test_mmt_logs_reports_sdk_errors_cleanly(monkeypatch) -> None: + from lightning_sdk.cli.mmt.logs import logs_mmt + + mmt = MagicMock() + mmt.logs.side_effect = RuntimeError("Logs are not available while the job is Pending.") + _patch_action(monkeypatch, mmt, {}) + + result = CliRunner().invoke(logs_mmt, ["my-mmt"]) + + assert result.exit_code != 0 + assert "Pending" in result.output + assert not isinstance(result.exception, RuntimeError) diff --git a/python/tests/cli/utils/test_logs.py b/python/tests/cli/utils/test_logs.py new file mode 100644 index 00000000..a68c9542 --- /dev/null +++ b/python/tests/cli/utils/test_logs.py @@ -0,0 +1,172 @@ +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import rich_click as click +from click.testing import CliRunner + +from lightning_sdk.api.logs_api import LogEntry +from lightning_sdk.cli.utils.logs import ( + LIVE_FALLBACK_IDLE_TIMEOUT, + LogSelection, + deployment_replica_labels, + highlight, + read_logs, + resolve_time, +) + + +def test_resolve_time_accepts_durations() -> None: + resolved = resolve_time("2h", "--since") + + # the server only parses RFC3339 and silently ignores anything else, so durations are + # resolved before the request goes out + delta = datetime.now(timezone.utc) - datetime.fromisoformat(resolved) + assert 1.9 * 3600 < delta.total_seconds() < 2.1 * 3600 + for value, unit_seconds in (("30s", 30), ("5m", 300), ("3d", 3 * 86400), ("1w", 7 * 86400)): + delta = datetime.now(timezone.utc) - datetime.fromisoformat(resolve_time(value, "--since")) + assert unit_seconds * 0.95 < delta.total_seconds() < unit_seconds * 1.05 + 1 + + +def test_resolve_time_passes_rfc3339_through_and_assumes_utc() -> None: + assert resolve_time("2026-07-27T00:00:00Z", "--since") == "2026-07-27T00:00:00+00:00" + assert resolve_time("2026-07-27T00:00:00+02:00", "--since") == "2026-07-27T00:00:00+02:00" + # a bare timestamp is not valid RFC3339 without an offset, so it is read as UTC + assert resolve_time("2026-07-27T00:00:00", "--since") == "2026-07-27T00:00:00+00:00" + assert resolve_time(None, "--since") is None + + +def test_resolve_time_rejects_anything_else() -> None: + with pytest.raises(click.UsageError, match="duration like 30s/5m/2h/3d/1w"): + resolve_time("last tuesday", "--since") + + +def test_highlight_marks_case_insensitive_matches() -> None: + assert ( + highlight("connection ERROR here", "error") + == f"connection {click.style('ERROR', fg=(167, 139, 250), bold=True)} here" + ) + assert highlight("plain", None) == "plain" + + +def _patch_api(monkeypatch, entries): + stream = MagicMock(return_value=list(entries)) + monkeypatch.setattr( + "lightning_sdk.cli.utils.logs.LogsApi", + MagicMock(return_value=SimpleNamespace(stream=stream)), + ) + return stream + + +def _run(fn): + """Invoke ``fn`` through a click command so click.echo output is captured.""" + + @click.command() + def _cmd() -> None: + fn() + + return CliRunner().invoke(_cmd) + + +def test_read_logs_labels_lines_by_resource(monkeypatch) -> None: + _patch_api( + monkeypatch, + [LogEntry(message="ready", resource_id="job-0"), LogEntry(message="serving", resource_id="job-1")], + ) + selection = LogSelection( + teamspace_id="ts-id", + deployment_id="dep-id", + labels={"job-0": "replica-0", "job-1": "replica-1"}, + ) + + result = _run(lambda: read_logs(selection)) + + assert result.exit_code == 0, result.output + assert result.output == "[replica-0] ready\n[replica-1] serving\n" + + +def test_read_logs_leaves_single_resource_output_unlabelled(monkeypatch) -> None: + _patch_api(monkeypatch, [LogEntry(message="ready", resource_id="job-0")]) + selection = LogSelection(teamspace_id="ts-id", job_ids=["job-0"]) + + result = _run(lambda: read_logs(selection)) + + assert result.output == "ready\n" + + +def test_read_logs_passes_filters_and_stream_mode(monkeypatch) -> None: + stream = _patch_api(monkeypatch, []) + selection = LogSelection(teamspace_id="ts-id", mmt_id="mmt-1") + + result = _run( + lambda: read_logs(selection, query="boom", severity="error", tail=25, since="2026-07-27T00:00:00+00:00") + ) + + assert result.exit_code == 0, result.output + kwargs = stream.call_args.kwargs + assert kwargs["mmt_id"] == "mmt-1" + assert kwargs["query"] == "boom" + assert kwargs["severity"] == "error" + assert kwargs["tail"] == 25 + assert kwargs["since"] == "2026-07-27T00:00:00+00:00" + assert kwargs["follow"] is False + # a resource with nothing saved still shows its live stream, bounded by the idle timeout + assert kwargs["fallback_to_live"] is True + assert kwargs["idle_timeout"] == LIVE_FALLBACK_IDLE_TIMEOUT + assert "No logs matched." in result.output + + +def test_read_logs_follow_tails_without_an_idle_timeout(monkeypatch) -> None: + stream = _patch_api(monkeypatch, [LogEntry(message="live")]) + selection = LogSelection(teamspace_id="ts-id", job_ids=["job-0"]) + + result = _run(lambda: read_logs(selection, follow=True)) + + assert result.output == "live\n" + kwargs = stream.call_args.kwargs + assert kwargs["follow"] is True + assert kwargs["idle_timeout"] is None + assert kwargs["fallback_to_live"] is False + + +def test_read_logs_highlights_matches(monkeypatch) -> None: + _patch_api(monkeypatch, [LogEntry(message="connection ERROR here")]) + selection = LogSelection(teamspace_id="ts-id", job_ids=["job-0"]) + + @click.command() + def _cmd() -> None: + read_logs(selection, query="error") + + result = CliRunner().invoke(_cmd, color=True) + + assert click.style("ERROR", fg=(167, 139, 250), bold=True) in result.output + + +def test_read_logs_reports_a_missing_websocket_client_cleanly(monkeypatch) -> None: + def _boom(*args, **kwargs): + raise RuntimeError("Following logs requires the 'websocket-client' package") + + monkeypatch.setattr( + "lightning_sdk.cli.utils.logs.LogsApi", + MagicMock(return_value=SimpleNamespace(stream=_boom)), + ) + selection = LogSelection(teamspace_id="ts-id", job_ids=["job-0"]) + + result = _run(lambda: read_logs(selection, follow=True)) + + assert result.exit_code == 1 + assert not isinstance(result.exception, RuntimeError) + + +def test_deployment_replica_labels_only_labels_when_several_replicas(monkeypatch) -> None: + api = MagicMock() + api.list_deployment_jobs.return_value = [ + SimpleNamespace(id="job-0", name="replica-0"), + SimpleNamespace(id="job-1", name="replica-1"), + ] + monkeypatch.setattr("lightning_sdk.api.deployment_api.DeploymentApi", lambda: api) + assert deployment_replica_labels("ts-id", "dep-id") == {"job-0": "replica-0", "job-1": "replica-1"} + + api.list_deployment_jobs.return_value = [SimpleNamespace(id="job-0", name="replica-0")] + assert deployment_replica_labels("ts-id", "dep-id") == {} diff --git a/python/tests/conftest.py b/python/tests/conftest.py index 69c081d3..c005159a 100644 --- a/python/tests/conftest.py +++ b/python/tests/conftest.py @@ -33,6 +33,7 @@ V1ExternalCluster, V1ExternalClusterSpec, V1GetCloudSpaceInstanceStatusResponse, + V1GetLogsResponse, V1GetUserResponse, V1Job, V1JobSpec, @@ -2618,7 +2619,17 @@ def mmt_api_get_job_by_name_mocker(mocker): @pytest.fixture() -def internal_job_logs_mocker(mocker): +def internal_get_logs_mocker(mocker): + """Serve the logs API with no saved lines, so log reads take the legacy fallback.""" + return mocker.patch( + "lightning_sdk.lightning_cloud.openapi.api.jobs_service_api.JobsServiceApi.jobs_service_get_logs", + autospec=True, + return_value=V1GetLogsResponse(entries=[]), + ) + + +@pytest.fixture() +def internal_job_logs_mocker(mocker, internal_get_logs_mocker): log_msg = "[2025-01-08T14:15:03.797142418Z] ⚡ ~ echo Hello\n[2025-01-08T14:15:03.803077717Z] Hello\n" dummy_url = "http://dummy-url.com/logs" diff --git a/python/tests/core/test_job.py b/python/tests/core/test_job.py index 4ca74638..f660fb32 100644 --- a/python/tests/core/test_job.py +++ b/python/tests/core/test_job.py @@ -1,8 +1,10 @@ import os +from datetime import datetime, timezone from unittest import mock import pytest +from lightning_sdk.api.logs_api import LogEntry from lightning_sdk.job import _RUNNING_LOGS_IDLE_TIMEOUT, Job from lightning_sdk.lightning_cloud.openapi import ( JobsServiceUpdateJobBody, @@ -953,12 +955,61 @@ def test_submit_job_from_running_studio( assert keeping_alive_mock.call_count == 0 +def _stub_logs_api(job, entries=()): + """Point the job's logs API at ``entries`` and return the stub.""" + stream_mock = mock.MagicMock(return_value=list(entries)) + job._logs_api.stream = stream_mock + return stream_mock + + @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) -def test_job_logs_finished_snapshot(job_api_get_job_by_name_mocker, internal_studio_init_mocker): +def test_job_logs_reads_logs_api(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) job._job_api.get_job = mock.MagicMock(return_value=V1Job(id="test-job-id", state="completed")) + stream_mock = _stub_logs_api( + job, + [ + LogEntry(message="line 1", timestamp=datetime(2026, 7, 27, 9, 0, tzinfo=timezone.utc)), + LogEntry(message="line 2"), + ], + ) + fallback = mock.MagicMock() + job._job_api.get_logs_finished = fallback + + assert job.logs == "line 1\nline 2" + fallback.assert_not_called() + assert stream_mock.call_args.kwargs["job_ids"] == ["test-job-id"] + assert stream_mock.call_args.kwargs["follow"] is False + assert stream_mock.call_args.args == (job.teamspace.id,) + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_job_logs_timestamps_and_filters(job_api_get_job_by_name_mocker, internal_studio_init_mocker): + studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") + job = Job("test-job", studio.teamspace) + + job._job_api.get_job = mock.MagicMock(return_value=V1Job(id="test-job-id", state="running")) + stream_mock = _stub_logs_api( + job, + [LogEntry(message="boom", severity="error", timestamp=datetime(2026, 7, 27, 9, 0, tzinfo=timezone.utc))], + ) + + assert job.logs(timestamps=True, query="boom", severity="error") == "2026-07-27T09:00:00+00:00 boom" + kwargs = stream_mock.call_args.kwargs + assert kwargs["query"] == "boom" + assert kwargs["severity"] == "error" + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_job_logs_finished_falls_back_to_saved_file(job_api_get_job_by_name_mocker, internal_studio_init_mocker): + studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") + job = Job("test-job", studio.teamspace) + + job._job_api.get_job = mock.MagicMock(return_value=V1Job(id="test-job-id", state="completed")) + # no lines in the current log format (e.g. an older job): the saved log file is used instead + _stub_logs_api(job, []) logs_mock = mock.MagicMock(return_value="line 1\nline 2\n") job._job_api.get_logs_finished = logs_mock @@ -970,15 +1021,33 @@ def test_job_logs_finished_snapshot(job_api_get_job_by_name_mocker, internal_stu logs_mock.assert_called_with(job_id="test-job-id", teamspace_id=job.teamspace.id) +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_job_logs_filtered_read_does_not_fall_back(job_api_get_job_by_name_mocker, internal_studio_init_mocker): + studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") + job = Job("test-job", studio.teamspace) + + job._job_api.get_job = mock.MagicMock(return_value=V1Job(id="test-job-id", state="completed")) + _stub_logs_api(job, []) + fallback = mock.MagicMock(return_value="everything\n") + job._job_api.get_logs_finished = fallback + + # a filtered read matching nothing must stay empty rather than dump the unfiltered file + assert job.logs(query="nope") == "" + fallback.assert_not_called() + + @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) def test_job_logs_finished_tail(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) job._job_api.get_job = mock.MagicMock(return_value=V1Job(id="test-job-id", state="completed")) + stream_mock = _stub_logs_api(job, []) job._job_api.get_logs_finished = mock.MagicMock(return_value="a\nb\nc\nd\n") assert job.logs(tail=2) == "c\nd" + # the API pages forward only, so the tail is requested from it too + assert stream_mock.call_args.kwargs["tail"] == 2 @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) @@ -994,25 +1063,37 @@ def test_job_logs_rank_warns_when_finished(job_api_get_job_by_name_mocker, inter @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) -def test_job_logs_running_snapshot_reads_until_idle(job_api_get_job_by_name_mocker, internal_studio_init_mocker): +def test_job_logs_running_snapshot_falls_back_to_live(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) job._job_api.get_job = mock.MagicMock(return_value=V1Job(id="test-job-id", state="running")) - stream_mock = mock.MagicMock(return_value=iter(["a", "b"])) - job._job_api.stream_logs = stream_mock + stream_mock = _stub_logs_api(job, [LogEntry(message="a"), LogEntry(message="b")]) - # a running snapshot reads the live stream until it goes idle, then joins the lines assert job.logs() == "a\nb" - stream_mock.assert_called_once_with( - job_id="test-job-id", - teamspace_id=job.teamspace.id, - follow=False, - tail=None, - rank=None, - idle_timeout=_RUNNING_LOGS_IDLE_TIMEOUT, - timestamps=False, - ) + kwargs = stream_mock.call_args.kwargs + # a running job with no saved lines still shows its live stream, bounded by the idle timeout + assert kwargs["fallback_to_live"] is True + assert kwargs["idle_timeout"] == _RUNNING_LOGS_IDLE_TIMEOUT + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_job_logs_follow_streams_from_logs_api(job_api_get_job_by_name_mocker, internal_studio_init_mocker): + studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") + job = Job("test-job", studio.teamspace) + + job._job_api.get_job = mock.MagicMock(return_value=V1Job(id="test-job-id", state="running")) + stream_mock = _stub_logs_api(job, [LogEntry(message="live-1"), LogEntry(message="live-2")]) + legacy = mock.MagicMock() + job._job_api.stream_logs = legacy + + assert list(job.logs(follow=True)) == ["live-1", "live-2"] + legacy.assert_not_called() + kwargs = stream_mock.call_args.kwargs + assert kwargs["follow"] is True + assert kwargs["idle_timeout"] is None + # the tail keeps running until the job itself is done + assert callable(kwargs["stop"]) @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) @@ -1022,6 +1103,7 @@ def test_job_logs_follow_on_finished_returns_saved_lines(job_api_get_job_by_name job._job_api.get_job = mock.MagicMock(return_value=V1Job(id="test-job-id", state="completed")) job._job_api.get_logs_finished = mock.MagicMock(return_value="a\nb\n") + _stub_logs_api(job, []) stream_mock = mock.MagicMock() job._job_api.stream_logs = stream_mock diff --git a/python/tests/core/test_mmt.py b/python/tests/core/test_mmt.py index eb8670c9..6195d816 100644 --- a/python/tests/core/test_mmt.py +++ b/python/tests/core/test_mmt.py @@ -2,6 +2,7 @@ import pytest +from lightning_sdk.api.logs_api import LogEntry from lightning_sdk.lightning_cloud.openapi import ( JobsServiceUpdateMultiMachineJobBody, V1Job, @@ -87,6 +88,14 @@ def test_mmt_exposes_placement_group_id(mmt_api_get_job_by_name_mocker, internal assert job.placement_group_id == "pg-1" +def test_mmt_exposes_resource_id(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker): + studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") + job = MMT("test-job", studio.teamspace) + job._job = V1MultiMachineJob(id="mmt-123") + + assert job.resource_id == "mmt-123" + + @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) def test_mmt_members_have_stable_rank_identity(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -455,3 +464,67 @@ def test_mmtv2_delete(mmt_api_get_job_by_name_mocker, internal_studio_init_mocke job.delete() delete_job_mock.assert_called_once_with(job_id="test-job-id", teamspace_id="ts-abc001") + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_mmtv2_logs_merges_ranks(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker, monkeypatch): + studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") + job = MMT("test-job", studio.teamspace) + + monkeypatch.setattr(type(job), "status", mock.PropertyMock(return_value=Status.Completed)) + machines = [mock.MagicMock(), mock.MagicMock()] + for i, machine in enumerate(machines): + machine.name = f"test-job-{i}" + machine._guaranteed_job = V1Job(id=f"job-{i}") + monkeypatch.setattr(type(job), "machines", mock.PropertyMock(return_value=tuple(machines))) + + stream_mock = mock.MagicMock( + return_value=[ + LogEntry(message="rank 0 up", resource_id="job-0"), + LogEntry(message="rank 1 up", resource_id="job-1"), + ] + ) + job._logs_api.stream = stream_mock + + # every rank is read through the one mmt selector, each line labelled with its machine + assert job.logs == "[test-job-0] rank 0 up\n[test-job-1] rank 1 up" + assert stream_mock.call_args.kwargs["mmt_id"] == "test-job-id" + assert stream_mock.call_args.kwargs["follow"] is False + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_mmtv2_logs_follow_and_filters(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker, monkeypatch): + studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") + job = MMT("test-job", studio.teamspace) + + monkeypatch.setattr(type(job), "status", mock.PropertyMock(return_value=Status.Running)) + monkeypatch.setattr(type(job), "machines", mock.PropertyMock(return_value=())) + stream_mock = mock.MagicMock(return_value=[LogEntry(message="boom", resource_id="job-0")]) + job._logs_api.stream = stream_mock + + assert list(job.logs(follow=True, query="boom", severity="error")) == ["[job-0] boom"] + kwargs = stream_mock.call_args.kwargs + assert kwargs["follow"] is True + assert kwargs["query"] == "boom" + assert kwargs["severity"] == "error" + assert kwargs["idle_timeout"] is None + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_mmtv2_logs_rejects_rank(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker, monkeypatch): + studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") + job = MMT("test-job", studio.teamspace) + + with pytest.raises(ValueError, match="mmt.machines"): + job.logs(rank=1) + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_mmtv2_logs_unavailable_while_pending(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker, monkeypatch): + studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") + job = MMT("test-job", studio.teamspace) + + monkeypatch.setattr(type(job), "status", mock.PropertyMock(return_value=Status.Pending)) + + with pytest.raises(RuntimeError, match="Pending"): + str(job.logs)