Support --since arg for dstack logs command (#3258)

Author: r4victor

src/dstack/_internal/cli/commands/logs.py

@@ -1,9 +1,12 @@
 import argparse
 import sys
+from datetime import datetime
+from typing import Optional
 
 from dstack._internal.cli.commands import APIBaseCommand
 from dstack._internal.cli.services.completion import RunNameCompleter
 from dstack._internal.core.errors import CLIError
+from dstack._internal.utils.common import parse_since
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -30,14 +33,25 @@ def _register(self):
             type=int,
             default=0,
         )
+        self._parser.add_argument(
+            "--since",
+            help=(
+                "Show only logs newer than the specified date."
+                " Can be a duration (e.g. 10s, 5m, 1d) or an RFC 3339 string (e.g. 2023-09-24T15:30:00Z)."
+            ),
+            type=str,
+        )
         self._parser.add_argument("run_name").completer = RunNameCompleter(all=True)  # type: ignore[attr-defined]
 
     def _command(self, args: argparse.Namespace):
         super()._command(args)
         run = self.api.runs.get(args.run_name)
         if run is None:
             raise CLIError(f"Run {args.run_name} not found")
+
+        start_time = _get_start_time(args.since)
         logs = run.logs(
+            start_time=start_time,
             diagnose=args.diagnose,
             replica_num=args.replica,
             job_num=args.job,
@@ -48,3 +62,12 @@ def _command(self, args: argparse.Namespace):
                 sys.stdout.buffer.flush()
         except KeyboardInterrupt:
             pass
+
+
+def _get_start_time(since: Optional[str]) -> Optional[datetime]:
+    if since is None:
+        return None
+    try:
+        return parse_since(since)
+    except ValueError as e:
+        raise CLIError(e.args[0])

src/dstack/_internal/utils/common.py

@@ -144,22 +144,34 @@ def pretty_resources(
     return " ".join(parts)
 
 
-def since(timestamp: str) -> datetime:
+def parse_since(value: str) -> datetime:
+    """
+    Returns a timestamp given an RFC 3339 string (e.g. 2023-09-24T15:30:00Z)
+    or a duration (e.g. 10s, 5m, 1d) between the timestamp and now.
+    """
     try:
-        seconds = parse_pretty_duration(timestamp)
+        seconds = parse_pretty_duration(value)
         return get_current_datetime() - timedelta(seconds=seconds)
     except ValueError:
         pass
     try:
-        return datetime.fromisoformat(timestamp)
+        res = datetime.fromisoformat(value)
     except ValueError:
         pass
+    else:
+        return check_time_offset_aware(res)
     try:
-        return datetime.fromtimestamp(int(timestamp))
+        return datetime.fromtimestamp(int(value), tz=timezone.utc)
     except Exception:
         raise ValueError("Invalid datetime format")
 
 
+def check_time_offset_aware(time: datetime) -> datetime:
+    if time.tzinfo is None:
+        raise ValueError("Timestamp is not offset-aware. Specify timezone.")
+    return time
+
+
 def parse_pretty_duration(duration: str) -> int:
     regex = re.compile(r"(?P<amount>\d+)(?P<unit>s|m|h|d|w)$")
     re_match = regex.match(duration)