From 7f29e9e22e1cfc34c79ffbf7594ed99874dbbfd7 Mon Sep 17 00:00:00 2001 From: epwalsh Date: Fri, 31 May 2024 13:12:33 -0700 Subject: [PATCH] improve handling of `since` option --- CHANGELOG.md | 8 ++++++++ beaker/services/experiment.py | 13 ++++++++++--- beaker/services/job.py | 8 ++++---- beaker/util.py | 12 ++++++------ 4 files changed, 28 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 90698bf..7f1c8f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,14 @@ use patch releases for compatibility fixes instead. ## Unreleased +### Added + +- Added `since` argument to `Beaker.experiment.follow()`. + +### Fixed + +- Fixed an issue with using `timedelta` objects for the `since` argument with `Beaker.(experiment|job).(follow|logs)`. + ## [v1.27.1](https://github.com/allenai/beaker-py/releases/tag/v1.27.1) - 2024-05-31 ### Added diff --git a/beaker/services/experiment.py b/beaker/services/experiment.py index f427da8..bde5426 100644 --- a/beaker/services/experiment.py +++ b/beaker/services/experiment.py @@ -349,8 +349,8 @@ def logs( :param quiet: If ``True``, progress won't be displayed. :param since: Only show logs since a particular time. Could be a :class:`~datetime.datetime` object (naive datetimes will be treated as UTC), a timestamp string in the form of RFC 3339 - (e.g. "2013-01-02T13:23:37Z"), or a relative time - (e.g. a :class:`~datetime.timedelta` or a string like "42m"). + (e.g. "2013-01-02T13:23:37Z"), or a :class:`~datetime.timedelta` + (e.g. `timedelta(seconds=60)`, which will show you the logs beginning 60 seconds ago). :raises ValueError: The experiment has no tasks or jobs, or the experiment has multiple tasks but ``task`` is not specified. @@ -699,6 +699,7 @@ def follow( timeout: Optional[float] = None, strict: bool = False, include_timestamps: bool = True, + since: Optional[Union[str, datetime, timedelta]] = None, ) -> Generator[bytes, None, Experiment]: """ Follow an experiment live, creating a generator that produces log lines @@ -726,6 +727,10 @@ def follow( :class:`~beaker.exceptions.JobFailedError` will be raised for non-zero exit codes. :param include_timestamps: If ``True`` (the default) timestamps from the Beaker logs will be included in the output. + :param since: Only show logs since a particular time. Could be a :class:`~datetime.datetime` object + (naive datetimes will be treated as UTC), a timestamp string in the form of RFC 3339 + (e.g. "2013-01-02T13:23:37Z"), or a :class:`~datetime.timedelta` + (e.g. `timedelta(seconds=60)`, which will show you the logs beginning 60 seconds ago). :raises ExperimentNotFound: If any experiment can't be found. :raises ValueError: The experiment has no tasks or jobs, or the experiment has multiple tasks but @@ -779,7 +784,9 @@ def follow( time.sleep(2.0) assert job is not None # for mypy - yield from self.beaker.job.follow(job, strict=strict, include_timestamps=include_timestamps) + yield from self.beaker.job.follow( + job, strict=strict, include_timestamps=include_timestamps, since=since + ) return self.get(experiment.id if isinstance(experiment, Experiment) else experiment) def url( diff --git a/beaker/services/job.py b/beaker/services/job.py index dc418f3..2e930d6 100644 --- a/beaker/services/job.py +++ b/beaker/services/job.py @@ -129,8 +129,8 @@ def logs( :param quiet: If ``True``, progress won't be displayed. :param since: Only show logs since a particular time. Could be a :class:`~datetime.datetime` object (naive datetimes will be treated as UTC), a timestamp string in the form of RFC 3339 - (e.g. "2013-01-02T13:23:37Z"), or a relative time - (e.g. a :class:`~datetime.timedelta` or a string like "42m"). + (e.g. "2013-01-02T13:23:37Z"), or a :class:`~datetime.timedelta` + (e.g. `timedelta(seconds=60)`, which will show you the logs beginning 60 seconds ago). :raises JobNotFound: If the job can't be found. :raises BeakerError: Any other :class:`~beaker.exceptions.BeakerError` type that can occur. @@ -422,8 +422,8 @@ def follow( will be included in the output. :param since: Only show logs since a particular time. Could be a :class:`~datetime.datetime` object (naive datetimes will be treated as UTC), a timestamp string in the form of RFC 3339 - (e.g. "2013-01-02T13:23:37Z"), or a relative time - (e.g. a :class:`~datetime.timedelta` or a string like "42m"). + (e.g. "2013-01-02T13:23:37Z"), or a :class:`~datetime.timedelta` + (e.g. `timedelta(seconds=60)`, which will show you the logs beginning 60 seconds ago). :raises JobNotFound: If any job can't be found. :raises JobTimeoutError: If the ``timeout`` expires. diff --git a/beaker/util.py b/beaker/util.py index 41578f5..7b652fc 100644 --- a/beaker/util.py +++ b/beaker/util.py @@ -3,7 +3,7 @@ import time import warnings from collections import OrderedDict -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from functools import wraps from pathlib import Path from typing import Any, Callable, Optional, Set, Tuple, Type, TypeVar, Union @@ -107,12 +107,12 @@ def prop_with_cache(self): def format_since(since: Union[datetime, timedelta, str]) -> str: if isinstance(since, datetime): - if since.tzinfo is None: - return since.strftime("%Y-%m-%dT%H:%M:%S.%fZ") - else: - return since.strftime("%Y-%m-%dT%H:%M:%S.%f%z") + if since.tzinfo is not None: + # Convert to UTC. + since = since.astimezone(timezone.utc) + return since.strftime("%Y-%m-%dT%H:%M:%S.%fZ") elif isinstance(since, timedelta): - return f"{since.total_seconds()}s" + return format_since(datetime.now(tz=timezone.utc) - since) else: return since