From 27bc8a0538afef8c246d33046ca65850dfa2224c Mon Sep 17 00:00:00 2001 From: epwalsh Date: Fri, 31 May 2024 10:58:59 -0700 Subject: [PATCH 1/4] Add `--tail` option to `gantry follow` --- CHANGELOG.md | 4 +++ gantry/commands/follow.py | 7 ++-- gantry/util.py | 73 +++++++++------------------------------ pyproject.toml | 2 +- 4 files changed, 26 insertions(+), 60 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5cdb933..8b7a4f7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased +### Added + +- Added `--tail` option to `gantry follow` to only tail the logs as opposed to dumping the entire history. + ## [v1.2.0](https://github.com/allenai/beaker-gantry/releases/tag/v1.2.0) - 2024-05-30 ### Changed diff --git a/gantry/commands/follow.py b/gantry/commands/follow.py index 23f53e2..900577c 100644 --- a/gantry/commands/follow.py +++ b/gantry/commands/follow.py @@ -7,11 +7,14 @@ @main.command(**CLICK_COMMAND_DEFAULTS) @click.argument("experiment", nargs=1, required=True, type=str) -def follow(experiment: str): +@click.option( + "-t", "--tail", is_flag=True, help="Only tail the logs as opposed to printing all logs so far." +) +def follow(experiment: str, tail: bool = False): """ Follow the logs for a running experiment. """ beaker = Beaker.from_env(session=True) exp = beaker.experiment.get(experiment) - job = util.follow_experiment(beaker, exp) + job = util.follow_experiment(beaker, exp, tail=tail) util.display_results(beaker, exp, job) diff --git a/gantry/util.py b/gantry/util.py index 066f0e7..fb74d58 100644 --- a/gantry/util.py +++ b/gantry/util.py @@ -1,9 +1,9 @@ import tempfile import time -from datetime import datetime +from datetime import datetime, timedelta from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING, Iterable, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Optional, Tuple, cast import requests import rich @@ -15,7 +15,6 @@ Digest, Experiment, Job, - JobTimeoutError, SecretNotFound, WorkspaceNotSet, ) @@ -74,43 +73,10 @@ def parse_git_remote_url(url: str) -> Tuple[str, str]: return account, repo -def display_logs(logs: Iterable[bytes], ignore_timestamp: Optional[str] = None) -> Optional[str]: +def follow_experiment( + beaker: Beaker, experiment: Experiment, timeout: int = 0, tail: bool = False +) -> Job: console = rich.get_console() - latest_timestamp: Optional[str] = None - - def print_line(line: str): - if not line: - return - nonlocal latest_timestamp - # Remove timestamp - try: - timestamp, line = line.split("Z ", maxsplit=1) - latest_timestamp = f"{timestamp}Z" - if ignore_timestamp is not None and latest_timestamp == ignore_timestamp: - return - except ValueError: - pass - console.print(line, highlight=False, markup=False) - - line_buffer = "" - for bytes_chunk in logs: - chunk = line_buffer + bytes_chunk.decode(errors="ignore") - chunk = chunk.replace("\r", "\n") - lines = chunk.split("\n") - if chunk.endswith("\n"): - line_buffer = "" - else: - # Last line chunk is probably incomplete. - lines, line_buffer = lines[:-1], lines[-1] - for line in lines: - print_line(line) - - print_line(line_buffer) - return latest_timestamp - - -def follow_experiment(beaker: Beaker, experiment: Experiment, timeout: int = 0) -> Job: - start = time.monotonic() # Wait for job to start... job: Optional[Job] = beaker.experiment.tasks(experiment.id)[0].latest_job # type: ignore @@ -123,27 +89,20 @@ def follow_experiment(beaker: Beaker, experiment: Experiment, timeout: int = 0) exit_code: Optional[int] = job.status.exit_code - stream_logs = exit_code is None + stream_logs = exit_code is None and not job.is_finalized if stream_logs: print() rich.get_console().rule("Logs") - - last_timestamp: Optional[str] = None - since: Optional[Union[str, datetime]] = datetime.utcnow() - while stream_logs and exit_code is None and not job.is_finalized: - job = beaker.experiment.tasks(experiment.id)[0].latest_job # type: ignore - assert job is not None - exit_code = job.status.exit_code - last_timestamp = display_logs( - beaker.job.logs(job, quiet=True, since=since), - ignore_timestamp=last_timestamp, - ) - since = last_timestamp or since - time.sleep(2.0) - if timeout > 0 and time.monotonic() - start >= timeout: - raise JobTimeoutError(f"Job did not finish within {timeout} seconds") - - if stream_logs: + for line_bytes in beaker.job.follow( + job, + timeout=timeout if timeout > 0 else None, + include_timestamps=False, + since=datetime.utcnow() - timedelta(seconds=5) if tail else None, + ): + line = line_bytes.decode(errors="ignore") + if line.endswith("\n"): + line = line[:-1] + console.print(line, highlight=False, markup=False) rich.get_console().rule("End logs") print() diff --git a/pyproject.toml b/pyproject.toml index 7ff41c8..5c4b26d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ authors = [ license = {file = "LICENSE"} requires-python = ">3.7" dependencies = [ - "beaker-py>=1.26.13,<2.0", + "beaker-py>=1.27.0,<2.0", "GitPython>=3.0,<4.0", "rich", "click", From 774e258e0aba92951d2fe73c5f62e70740303810 Mon Sep 17 00:00:00 2001 From: epwalsh Date: Fri, 31 May 2024 11:00:24 -0700 Subject: [PATCH 2/4] clean up --- gantry/util.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/gantry/util.py b/gantry/util.py index fb74d58..61dad6b 100644 --- a/gantry/util.py +++ b/gantry/util.py @@ -3,7 +3,7 @@ from datetime import datetime, timedelta from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING, Optional, Tuple, cast +from typing import Optional, Tuple, cast import requests import rich @@ -26,9 +26,6 @@ from .exceptions import * from .version import VERSION -if TYPE_CHECKING: - from datetime import timedelta - class StrEnum(str, Enum): def __str__(self) -> str: From 6355d60592eb1ccdeeb2d21ac295b2c3694ab527 Mon Sep 17 00:00:00 2001 From: epwalsh Date: Fri, 31 May 2024 11:07:45 -0700 Subject: [PATCH 3/4] make sure to refresh the job --- gantry/util.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/gantry/util.py b/gantry/util.py index 61dad6b..f598ed9 100644 --- a/gantry/util.py +++ b/gantry/util.py @@ -103,6 +103,9 @@ def follow_experiment( rich.get_console().rule("End logs") print() + # Refresh the job. + job = beaker.job.get(job) + return job From 0fcf2c42ba26bfefcec0cd6193d207ba677260cb Mon Sep 17 00:00:00 2001 From: epwalsh Date: Fri, 31 May 2024 11:09:43 -0700 Subject: [PATCH 4/4] fix --- gantry/util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gantry/util.py b/gantry/util.py index f598ed9..c53791b 100644 --- a/gantry/util.py +++ b/gantry/util.py @@ -104,7 +104,7 @@ def follow_experiment( print() # Refresh the job. - job = beaker.job.get(job) + job = beaker.job.get(job.id) return job