diff --git a/python/lightning_sdk/cli/job/logs.py b/python/lightning_sdk/cli/job/logs.py index 07590f0d..de003c4b 100644 --- a/python/lightning_sdk/cli/job/logs.py +++ b/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 diff --git a/python/tests/cli/job/test_logs.py b/python/tests/cli/job/test_logs.py index f94e2512..2f470769 100644 --- a/python/tests/cli/job/test_logs.py +++ b/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)