diff --git a/python/lightning_sdk/api/job_api.py b/python/lightning_sdk/api/job_api.py index 878044c8..12c0b7ea 100644 --- a/python/lightning_sdk/api/job_api.py +++ b/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. diff --git a/python/lightning_sdk/job.py b/python/lightning_sdk/job.py index 6a1dfa48..d742669f 100644 --- a/python/lightning_sdk/job.py +++ b/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: diff --git a/python/tests/api/test_job_api.py b/python/tests/api/test_job_api.py index 2c2f5bcc..0630c0cc 100644 --- a/python/tests/api/test_job_api.py +++ b/python/tests/api/test_job_api.py @@ -1,9 +1,17 @@ +import json +import sys +import types from typing import List from unittest import mock import pytest -from lightning_sdk.api.job_api import JobApiV2 +from lightning_sdk.api.job_api import ( + JobApiV2, + _decode_log_messages, + _format_log_timestamp, + _job_logs_ws_url, +) from lightning_sdk.lightning_cloud.openapi import ( JobsServiceUpdateJobBody, V1Job, @@ -236,3 +244,200 @@ def test_jobv2_delete(mocker_auth, cloudspace_id, expected_cloudspace_id): job_api.delete_job("test-job-id", "ts-abc", cloudspace_id) delete_job_mock.assert_called_once_with(id="test-job-id", project_id="ts-abc", cloudspace_id=expected_cloudspace_id) + + +# --------------------------------------------------------------------------- +# live log streaming +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ({"seconds": 0, "nanos": 0}, "1970-01-01T00:00:00+00:00"), + ({"seconds": 1_700_000_000}, "2023-11-14T22:13:20+00:00"), + ({"nanos": 5}, None), # no seconds -> cannot build a timestamp + ("2023-11-14T22:13:20Z", "2023-11-14T22:13:20Z"), # already-formatted string passes through + (None, None), + (123, None), # unexpected type + ], +) +def test_format_log_timestamp(value, expected): + assert _format_log_timestamp(value) == expected + + +def test_decode_log_messages_plain_and_json(): + # a JSON array of entries -> one line per entry, using the `message` field + frame = json.dumps([{"message": "hello"}, {"Message": "world"}]) + assert list(_decode_log_messages(frame)) == ["hello", "world"] + + # non-JSON payloads are yielded verbatim + assert list(_decode_log_messages("not json")) == ["not json"] + + +def test_decode_log_messages_with_timestamps(): + frame = json.dumps( + [ + {"message": "hello", "timestamp": {"seconds": 0, "nanos": 0}}, + {"message": "no-ts"}, # missing timestamp -> line unchanged + ] + ) + assert list(_decode_log_messages(frame, timestamps=True)) == [ + "1970-01-01T00:00:00+00:00 hello", + "no-ts", + ] + + +def test_job_logs_ws_url_builds_websocket_url(): + url = _job_logs_ws_url("ts-abc", "job-1", follow=True, tail=50, rank=2) + + assert url.startswith("wss://") + assert "/v1/projects/ts-abc/jobs/job-1/logs" in url + assert "follow=true" in url + assert "direction=forward" in url + assert "tail=50" in url + assert "rank=2" in url + + # optional params are omitted when not provided + minimal = _job_logs_ws_url("ts-abc", "job-1", follow=False, tail=None, rank=None) + assert "follow=false" in minimal + assert "tail=" not in minimal + assert "rank=" not in minimal + + +class _FakeWebSocket: + """Minimal stand-in for a websocket-client connection.""" + + def __init__(self, frames): + self._frames = list(frames) + self.closed = False + + def recv(self): + if not self._frames: + return "" # empty frame signals the server closed the stream + item = self._frames.pop(0) + if isinstance(item, Exception): + raise item + return item + + def close(self): + self.closed = True + + +def _install_fake_websocket(mocker, connections): + """Register a fake ``websocket`` module whose ``create_connection`` returns the given connections.""" + module = types.ModuleType("websocket") + + # names mirror websocket-client exactly, since stream_logs imports them by name + class WebSocketConnectionClosedException(Exception): # noqa: N818 + pass + + class WebSocketTimeoutException(Exception): # noqa: N818 + pass + + module.WebSocketConnectionClosedException = WebSocketConnectionClosedException + module.WebSocketTimeoutException = WebSocketTimeoutException + module.create_connection = mock.MagicMock(side_effect=list(connections)) + + mocker.patch.dict(sys.modules, {"websocket": module}) + return module + + +def test_stream_logs_yields_lines_and_stops(mocker, mocker_auth): + mocker.patch("lightning_sdk.api.job_api.Auth") + frame = json.dumps([{"message": "line-1"}, {"message": "line-2"}]) + module = _install_fake_websocket(mocker, [_FakeWebSocket([frame])]) + + job_api = JobApiV2() + lines = list(job_api.stream_logs("job-1", "ts-abc", follow=False)) + + assert lines == ["line-1", "line-2"] + # no reconnect when follow is off -> exactly one connection + assert module.create_connection.call_count == 1 + + +def test_stream_logs_passes_timestamps_flag(mocker, mocker_auth): + mocker.patch("lightning_sdk.api.job_api.Auth") + frame = json.dumps([{"message": "hi", "timestamp": {"seconds": 0, "nanos": 0}}]) + _install_fake_websocket(mocker, [_FakeWebSocket([frame])]) + + job_api = JobApiV2() + lines = list(job_api.stream_logs("job-1", "ts-abc", follow=False, timestamps=True)) + + assert lines == ["1970-01-01T00:00:00+00:00 hi"] + + +def test_stream_logs_reconnects_on_transient_drop(mocker, mocker_auth): + mocker.patch("lightning_sdk.api.job_api.Auth") + + module_holder = {} + + def _closed_exc(): + return module_holder["module"].WebSocketConnectionClosedException() + + # first connection yields a line then drops; second connection yields another line then ends + first = _FakeWebSocket([json.dumps([{"message": "before-drop"}])]) + second = _FakeWebSocket([json.dumps([{"message": "after-reconnect"}])]) + module = _install_fake_websocket(mocker, [first, second]) + module_holder["module"] = module + # make the first connection drop after its frame + first._frames.append(_closed_exc()) + + # stop the loop after the second connection by raising out of the reconnect backoff + class _StopLoopError(Exception): + pass + + mocker.patch("lightning_sdk.api.job_api.time.sleep", side_effect=[None, _StopLoopError()]) + + job_api = JobApiV2() + # job is still running, so drops are treated as transient and we reconnect + job_api.get_job = mock.MagicMock(return_value=V1Job(id="job-1", state="running")) + collected = [] + + def _drain() -> None: + for line in job_api.stream_logs("job-1", "ts-abc", follow=True): + collected.append(line) + + with pytest.raises(_StopLoopError): + _drain() + + assert collected == ["before-drop", "after-reconnect"] + assert module.create_connection.call_count == 2 + + +def test_stream_logs_follow_stops_when_job_finishes(mocker, mocker_auth): + mocker.patch("lightning_sdk.api.job_api.Auth") + ws = _FakeWebSocket([json.dumps([{"message": "last-line"}])]) + module = _install_fake_websocket(mocker, [ws]) + # after the final line the socket goes quiet (heartbeats only) -> recv times out + ws._frames.append(module.WebSocketTimeoutException()) + + job_api = JobApiV2() + job_api.get_job = mock.MagicMock(return_value=V1Job(id="job-1", state="completed")) + + lines = list(job_api.stream_logs("job-1", "ts-abc", follow=True)) + + assert lines == ["last-line"] + assert module.create_connection.call_count == 1 # finished -> no reconnect + job_api.get_job.assert_called() + + +def test_stream_logs_follow_keeps_waiting_while_running(mocker, mocker_auth): + mocker.patch("lightning_sdk.api.job_api.Auth") + ws = _FakeWebSocket([json.dumps([{"message": "l1"}])]) + module = _install_fake_websocket(mocker, [ws]) + # quiet (still running -> keep waiting), another line, then quiet again (now finished -> stop) + ws._frames.append(module.WebSocketTimeoutException()) + ws._frames.append(json.dumps([{"message": "l2"}])) + ws._frames.append(module.WebSocketTimeoutException()) + + job_api = JobApiV2() + job_api.get_job = mock.MagicMock( + side_effect=[V1Job(id="job-1", state="running"), V1Job(id="job-1", state="completed")] + ) + + lines = list(job_api.stream_logs("job-1", "ts-abc", follow=True)) + + assert lines == ["l1", "l2"] + assert module.create_connection.call_count == 1 + assert job_api.get_job.call_count == 2 diff --git a/python/tests/core/test_job.py b/python/tests/core/test_job.py index 1ed2d740..4ca74638 100644 --- a/python/tests/core/test_job.py +++ b/python/tests/core/test_job.py @@ -3,7 +3,7 @@ import pytest -from lightning_sdk.job import Job +from lightning_sdk.job import _RUNNING_LOGS_IDLE_TIMEOUT, Job from lightning_sdk.lightning_cloud.openapi import ( JobsServiceUpdateJobBody, V1Job, @@ -951,3 +951,104 @@ def test_submit_job_from_running_studio( org="org-abc", ) assert keeping_alive_mock.call_count == 0 + + +@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): + 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")) + logs_mock = mock.MagicMock(return_value="line 1\nline 2\n") + job._job_api.get_logs_finished = logs_mock + + # the proxy behaves like the log text: equality, str(), .splitlines(), iteration + assert job.logs == "line 1\nline 2\n" + assert str(job.logs) == "line 1\nline 2\n" + assert job.logs.splitlines() == ["line 1", "line 2"] + assert list(job.logs) == ["line 1", "line 2"] + 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_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")) + job._job_api.get_logs_finished = mock.MagicMock(return_value="a\nb\nc\nd\n") + + assert job.logs(tail=2) == "c\nd" + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_job_logs_rank_warns_when_finished(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")) + job._job_api.get_logs_finished = mock.MagicMock(return_value="done") + + with pytest.warns(UserWarning, match="rank"): + _ = job.logs(rank=1) + + +@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): + 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 + + # 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, + ) + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_job_logs_follow_on_finished_returns_saved_lines(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")) + job._job_api.get_logs_finished = mock.MagicMock(return_value="a\nb\n") + stream_mock = mock.MagicMock() + job._job_api.stream_logs = stream_mock + + # following an already-finished job must NOT open a websocket (which would hang); it + # returns the saved lines as an iterator and stops. + assert list(job.logs(follow=True)) == ["a", "b"] + stream_mock.assert_not_called() + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_job_logs_follow_delegates_to_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 = mock.MagicMock(return_value=iter(["line-1", "line-2"])) + job._job_api.stream_logs = stream_mock + + result = list(job.logs(follow=True, tail=10, rank=1, timestamps=True)) + + assert result == ["line-1", "line-2"] + stream_mock.assert_called_once_with( + job_id="test-job-id", + teamspace_id=job.teamspace.id, + follow=True, + tail=10, + rank=1, + idle_timeout=None, + timestamps=True, + )