feat(sdk): live logs for jobs

python/lightning_sdk/api/job_api.py

@@ -1,5 +1,9 @@
+import json
 import time
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from contextlib import suppress
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Union
+from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
 from urllib.request import urlopen
 
 from lightning_sdk.api.utils import _get_cloud_url as _cloud_url
@@ -9,6 +13,7 @@
     resolve_path_mappings,
 )
 from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__
+from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
     JobsServiceCreateJobBody,
     JobsServiceUpdateJobBody,
@@ -26,6 +31,80 @@
 if TYPE_CHECKING:
     from lightning_sdk.status import Status
 
+# While following, recv() is given this timeout so we can periodically check whether the job has
+# finished. It must stay below the server's log-socket heartbeat (10s) so recv() actually wakes up
+# during quiet periods.
+_FOLLOW_POLL_INTERVAL = 5.0
+
+
+def _job_logs_ws_url(
+    teamspace_id: str,
+    job_id: str,
+    *,
+    follow: bool,
+    tail: Optional[int],
+    rank: Optional[int],
+) -> str:
+    """Build the ``wss://`` URL for following a job's logs over the controlplane websocket."""
+    parsed = urlparse(f"/v1/projects/{teamspace_id}/jobs/{job_id}/logs")
+    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)
+    path = urlunparse(parsed._replace(query=urlencode(query)))
+
+    absolute = f"{_cloud_url().rstrip('/')}/{path.lstrip('/')}"
+    if absolute.startswith("https://"):
+        return "wss://" + absolute[len("https://") :]
+    if absolute.startswith("http://"):
+        return "ws://" + absolute[len("http://") :]
+    return absolute
+
+
+def _format_log_timestamp(value: object) -> Optional[str]:
+    """Best-effort convert a ``LogEntry`` timestamp into an ISO-8601 string.
+
+    The controlplane serialises the protobuf ``Timestamp`` as ``{"seconds": ..., "nanos": ...}``,
+    but we also accept an already-formatted string in case that ever changes.
+    """
+    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).isoformat()
+    if isinstance(value, str) and value:
+        return value
+    return None
+
+
+def _decode_log_messages(message: str, *, timestamps: bool = False) -> Iterator[str]:
+    """Decode a websocket log frame into individual log lines.
+
+    Args:
+        message: A raw websocket text frame (usually a JSON array of log entries).
+        timestamps: When ``True``, prepend each line with its ISO-8601 timestamp if available.
+    """
+    try:
+        payload = json.loads(message)
+    except json.JSONDecodeError:
+        yield message
+        return
+
+    entries = payload if isinstance(payload, list) else [payload]
+    for entry in entries:
+        if isinstance(entry, dict):
+            line = entry.get("message") or entry.get("Message") or json.dumps(entry)
+            if timestamps:
+                ts = _format_log_timestamp(entry.get("timestamp") or entry.get("Timestamp"))
+                if ts:
+                    line = f"{ts} {line}"
+            yield line
+        else:
+            yield str(entry)
+
 
 class JobApiV2:
     """Native (v2 jobs-service) API client for single-machine Jobs."""
@@ -290,6 +369,88 @@ def get_logs_finished(self, job_id: str, teamspace_id: str) -> str:
         data = urlopen(resp.url).read().decode("utf-8")
         return remove_datetime_prefix(str(data))
 
+    def stream_logs(
+        self,
+        job_id: str,
+        teamspace_id: str,
+        *,
+        follow: bool = True,
+        tail: Optional[int] = None,
+        rank: Optional[int] = None,
+        idle_timeout: Optional[float] = None,
+        reconnect: bool = True,
+        timestamps: bool = False,
+    ) -> Iterator[str]:
+        """Yield job log lines live over the controlplane websocket.
+
+        Reconnects automatically on transient drops while ``follow`` is set, so this can
+        stay open for the full lifetime of a long-running job.
+
+        Args:
+            job_id: The unique identifier of the job whose logs are streamed.
+            teamspace_id: The ID of the teamspace that owns the job.
+            follow: Keep the stream open and yield new lines as they are produced.
+            tail: Number of recent lines to emit before following.
+            rank: Distributed job rank to stream from.
+            idle_timeout: If set, stop the stream after this many seconds without data.
+            reconnect: Reconnect on transient socket drops while following.
+            timestamps: When ``True``, prepend each line with its ISO-8601 timestamp.
+
+        Yields:
+            Individual decoded log lines.
+        """
+        try:
+            import websocket
+            from websocket import (
+                WebSocketConnectionClosedException,
+                WebSocketTimeoutException,
+            )
+        except ImportError as ex:
+            raise RuntimeError(
+                "Streaming logs requires the 'websocket-client' package (pip install websocket-client)."
+            ) from ex
+
+        auth_header = Auth().authenticate()
+        url = _job_logs_ws_url(teamspace_id, job_id, follow=follow, tail=tail, rank=rank)
+
+        # For follow we poll recv() with a timeout so we can periodically check whether the job has
+        # finished: the server keeps the socket open (heartbeats only) after a job ends rather than
+        # closing it, so a plain blocking recv() would hang forever once the job is done.
+        recv_timeout = idle_timeout if idle_timeout is not None else (_FOLLOW_POLL_INTERVAL if follow else None)
+
+        while True:
+            ws = None
+            try:
+                kwargs: Dict = {"header": [f"Authorization: {auth_header}"]}
+                if recv_timeout is not None:
+                    kwargs["timeout"] = recv_timeout
+                ws = websocket.create_connection(url, **kwargs)
+                while True:
+                    try:
+                        message = ws.recv()
+                    except WebSocketTimeoutException:
+                        if follow:
+                            # quiet socket (heartbeats only): stop if the job is done, else keep waiting
+                            if self._is_job_finished(job_id, teamspace_id):
+                                return
+                            continue
+                        return  # idle timeout on a snapshot: treat the stream as finished
+                    except WebSocketConnectionClosedException:
+                        break  # fall through to reconnect logic
+                    if message == "":
+                        break
+                    yield from _decode_log_messages(message, timestamps=timestamps)
+            finally:
+                if ws is not None:
+                    with suppress(Exception):
+                        ws.close()
+
+            if not (follow and reconnect):
+                return
+            if self._is_job_finished(job_id, teamspace_id):
+                return  # socket dropped and the job is done — nothing more to stream
+            time.sleep(1)  # brief backoff before reconnecting, matches the CLI
+
     def get_studio_name(self, job: V1Job) -> Optional[str]:
         """Return the name of the Studio linked to this job, or ``None`` if none is attached.
 
@@ -345,6 +506,29 @@ def get_mmt_name(self, job: V1Job) -> str:
                 return splits[0]
         return ""
 
+    def _is_job_finished(self, job_id: str, teamspace_id: str) -> bool:
+        """Return whether the job has reached a terminal state.
+
+        Used while following logs to decide when to stop: the server keeps the log socket open
+        (heartbeats only) after a job ends, so we poll the job state to detect completion.
+
+        Args:
+            job_id: The unique identifier of the job.
+            teamspace_id: The ID of the teamspace that owns the job.
+
+        Returns:
+            ``True`` if the job is Completed/Failed/Stopped, else ``False`` (including when the
+            state cannot be determined, so a transient status-check failure never kills a stream).
+        """
+        from lightning_sdk.status import Status
+
+        try:
+            job = self.get_job(job_id, teamspace_id)
+        except Exception:
+            # a status-check failure must not kill an active stream
+            return False
+        return self._job_state_to_external(job.state) in (Status.Stopped, Status.Completed, Status.Failed)
+
     def _job_state_to_external(self, state: str) -> "Status":
         """Convert a raw v2 job state string to the public ``Status`` enum.
 

python/lightning_sdk/job.py

@@ -1,6 +1,6 @@
 import warnings
 from pathlib import PurePath
-from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, TypedDict, Union
 
 from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.api.job_api import JobApiV2
@@ -29,6 +29,71 @@
     "Job",
 ]
 
+# A running job's log stream never sends a clean end-of-stream frame, so a bounded
+# snapshot ("give me the logs up to now") is read until this many seconds of silence.
+_RUNNING_LOGS_IDLE_TIMEOUT = 5.0
+
+
+class _Logs:
+    """A logs handle that is both a value and callable.
+
+    ``job.logs`` behaves like the log text (a snapshot), so ``print(job.logs)``,
+    ``job.logs.splitlines()`` and ``for line in job.logs`` all work. Calling it,
+    ``job.logs(follow=True, tail=..., rank=...)``, fetches logs with options and
+    returns an iterator of lines while ``follow=True``.
+
+    Note: this is a str-like proxy, not an actual ``str`` (``isinstance(job.logs, str)``
+    is ``False``), and static type checkers see ``_Logs`` rather than ``str``.
+    """
+
+    def __init__(self, fetch: Callable[..., Union[str, Iterator[str]]]) -> None:
+        self._fetch = fetch
+        self._cached: Optional[str] = None
+
+    def __call__(
+        self,
+        *,
+        follow: bool = False,
+        tail: Optional[int] = None,
+        rank: Optional[int] = None,
+        timestamps: bool = False,
+    ) -> Union[str, Iterator[str]]:
+        return self._fetch(follow=follow, tail=tail, rank=rank, timestamps=timestamps)
+
+    def _text(self) -> str:
+        if self._cached is None:
+            self._cached = self._fetch(follow=False)  # type: ignore[assignment]
+        return self._cached
+
+    def __str__(self) -> str:
+        return self._text()
+
+    def __repr__(self) -> str:
+        return repr(self._text())
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._text().splitlines())
+
+    def __len__(self) -> int:
+        return len(self._text())
+
+    def __contains__(self, item: object) -> bool:
+        return item in self._text()
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, _Logs):
+            other = other._text()
+        return self._text() == other
+
+    __hash__ = None  # type: ignore[assignment]
+
+    def __getattr__(self, name: str) -> Any:
+        # only invoked for attributes not found normally; delegate to the log text.
+        # guard private/dunder lookups to avoid recursing through _fetch/_cached.
+        if name.startswith("_"):
+            raise AttributeError(name)
+        return getattr(self._text(), name)
+
 
 class JobDict(TypedDict):
     name: str
@@ -439,11 +504,92 @@ def share_path(self) -> Optional[str]:
         raise NotImplementedError("Not implemented yet")
 
     @property
-    def logs(self) -> str:
-        if self.status not in (Status.Failed, Status.Completed, Status.Stopped):
-            raise RuntimeError("Getting jobs logs while the job is pending or running is not supported yet!")
+    def logs(self) -> _Logs:
+        """The job's logs.
+
+        Use it as a value for a snapshot of the logs up to now::
 
-        return self._job_api.get_logs_finished(job_id=self._guaranteed_job.id, teamspace_id=self.teamspace.id)
+            print(job.logs)
+
+        or call it to pass options and/or follow the logs live::
+
+            recent = job.logs(tail=100)            # snapshot of the last 100 lines
+            for line in job.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.
+        - ``rank``: Distributed job rank to read from (running jobs only).
+        - ``timestamps``: Prepend each line with its ISO-8601 timestamp.
+        """
+        return _Logs(self._compute_logs)
+
+    def _compute_logs(
+        self,
+        *,
+        follow: bool = False,
+        tail: Optional[int] = None,
+        rank: Optional[int] = None,
+        timestamps: bool = False,
+    ) -> Union[str, Iterator[str]]:
+        """Fetch the logs, dispatching on job state. See :attr:`logs` for the public API."""
+        status = self.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).
+        if follow and status == Status.Running:
+            return self._stream_logs(follow=True, tail=tail, rank=rank, timestamps=timestamps)
+
+        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,
+                )
+            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,
+                )
+            )
+        else:
+            raise RuntimeError(f"Logs are not available while the job is {status}.")
+
+        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
+
+    def _stream_logs(
+        self,
+        *,
+        follow: bool = True,
+        tail: Optional[int] = None,
+        rank: Optional[int] = None,
+        idle_timeout: Optional[float] = None,
+        timestamps: bool = False,
+    ) -> Iterator[str]:
+        """Stream the job's logs live over the websocket (internal; see :attr:`logs`)."""
+        return self._job_api.stream_logs(
+            job_id=self._guaranteed_job.id,
+            teamspace_id=self.teamspace.id,
+            follow=follow,
+            tail=tail,
+            rank=rank,
+            idle_timeout=idle_timeout,
+            timestamps=timestamps,
+        )
 
     @property
     def link(self) -> str: