feat(sdk): extend cli job logs for streaming and filtering

python/lightning_sdk/cli/job/logs.py

@@ -6,9 +6,6 @@
 
 from lightning_sdk.cli.legacy.job_and_mmt_action import _JobAndMMTAction
 from lightning_sdk.cli.utils.logging import LightningCommand
-from lightning_sdk.status import Status
-
-_TERMINAL = (Status.Completed, Status.Failed, Status.Stopped)
 
 
 @click.command("logs", cls=LightningCommand)
@@ -22,17 +19,33 @@
         "If not specified can be selected interactively."
     ),
 )
-def logs_job(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+@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("--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.")
+def logs_job(
+    name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    follow: bool = False,
+    tail: Optional[int] = None,
+    rank: Optional[int] = None,
+    timestamps: bool = False,
+) -> None:
     """Print the logs for a job.
 
-    Logs are available once the job reaches a terminal state (Completed, Failed or
-    Stopped). While the job is still pending or running this prints its current status
-    instead of erroring — re-run once it has finished.
+    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.
     """
     job = _JobAndMMTAction().job(name=name, teamspace=teamspace)
-    if job.status not in _TERMINAL:
-        raise click.ClickException(
-            f"Job '{job.name}' is {job.status}; logs are only available once it reaches a "
-            "terminal state (Completed/Failed/Stopped). Re-run this command after it finishes."
-        )
-    click.echo(job.logs)
+
+    try:
+        logs = job.logs(follow=follow, tail=tail, rank=rank, timestamps=timestamps)
+        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

python/tests/cli/job/test_logs.py

@@ -2,7 +2,6 @@
 
 from click.testing import CliRunner
 
-from lightning_sdk.status import Status
 from tests.cli.help import assert_help_contains, mock_command_logging
 
 
@@ -12,6 +11,10 @@ def test_job_logs_help() -> None:
         "lightning job logs --help",
         "Usage: lightning job logs",
         "Print the logs for a job.",
+        "--follow",
+        "--tail",
+        "--rank",
+        "--timestamps",
     )
 
 
@@ -40,14 +43,12 @@ def job(self, name=None, teamspace=None):
 
 
 @mock_command_logging
-def test_job_logs_prints_logs_when_terminal(monkeypatch) -> None:
+def test_job_logs_prints_snapshot(monkeypatch) -> None:
     from lightning_sdk.cli.job.logs import logs_job
 
     captured: dict = {}
     job = MagicMock()
-    job.name = "my-job"
-    job.status = Status.Completed
-    job.logs = "hello from the job\n42"
+    job.logs.return_value = "hello from the job\n42"
     _patch_action(monkeypatch, job, captured)
 
     result = CliRunner().invoke(logs_job, ["my-job", "--teamspace", "org/teamspace"])
@@ -56,21 +57,39 @@ def test_job_logs_prints_logs_when_terminal(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)
 
 
 @mock_command_logging
-def test_job_logs_errors_while_not_terminal(monkeypatch) -> None:
+def test_job_logs_follows_with_options(monkeypatch) -> None:
     from lightning_sdk.cli.job.logs import logs_job
 
     captured: dict = {}
     job = MagicMock()
-    job.name = "my-job"
-    job.status = Status.Pending
+    job.logs.return_value = iter(["line 1", "line 2"])
     _patch_action(monkeypatch, job, captured)
 
-    result = CliRunner().invoke(logs_job, ["my-job", "--teamspace", "org/teamspace"])
+    result = CliRunner().invoke(
+        logs_job,
+        ["my-job", "--follow", "--tail", "10", "--rank", "2", "--timestamps"],
+    )
+
+    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)
+
+
+@mock_command_logging
+def test_job_logs_reports_sdk_errors_cleanly(monkeypatch) -> None:
+    from lightning_sdk.cli.job.logs import logs_job
+
+    captured: dict = {}
+    job = MagicMock()
+    job.logs.side_effect = RuntimeError("Logs are not available while the job is Pending.")
+    _patch_action(monkeypatch, job, captured)
+
+    result = CliRunner().invoke(logs_job, ["my-job"])
 
-    # errors cleanly (mentioning the status) rather than raising the raw SDK RuntimeError
     assert result.exit_code != 0
     assert "Pending" in result.output
     assert not isinstance(result.exception, RuntimeError)