[autorevert] Add option to simulate run at a specific timestamp (#7609)

Adds an optional param `--as-of <UTC date[time]>` that will set internal
clock (NOW) in CH queries (specifically, commits and jobs), allowing to
virtually go back in time, and perform restarts / reverts.

When set ` --hours N` lookback window becomes relative to the `as-of`
timestamp.


Testing:

```
 1049  python -m pytorch_auto_revert --dry-run autorevert-checker pull --hours 12  --hud-html 
 1049  python -m pytorch_auto_revert --dry-run autorevert-checker pull --hours 12  --hud-html --as-of "2025-12-18 15:31"
 1052  python -m pytorch_auto_revert --dry-run autorevert-checker pull --hours 12  --hud-html --as-of "2025-12-18 06:13"
```

Author: izaitsevfb

aws/lambda/pytorch-auto-revert/pytorch_auto_revert/__main__.py

@@ -30,7 +30,7 @@
 from .testers.autorevert_v2 import autorevert_v2
 from .testers.hud import render_hud_html_from_clickhouse, write_hud_html_from_cli
 from .testers.restart_checker import workflow_restart_checker
-from .utils import RestartAction, RetryWithBackoff, RevertAction
+from .utils import parse_datetime, RestartAction, RetryWithBackoff, RevertAction
 
 
 # Special constant to indicate --hud-html was passed as a flag (without a value)
@@ -326,6 +326,16 @@ def get_opts(default_config: DefaultConfig) -> argparse.Namespace:
         default=default_config.notify_issue_number,
         help="Issue number to notify",
     )
+    workflow_parser.add_argument(
+        "--as-of",
+        type=str,
+        default=None,
+        help=(
+            "Run as if current time is this timestamp (UTC). "
+            "Accepts ISO 8601 or 'YYYY-MM-DD HH:MM[:SS]' format. "
+            "Useful for testing autorevert logic at a specific point in time."
+        ),
+    )
 
     # workflow-restart-checker subcommand
     workflow_restart_parser = subparsers.add_parser(
@@ -422,6 +432,7 @@ def _get(attr: str, default=None):
         restart_action=_get("restart_action", None),
         revert_action=_get("revert_action", None),
         bisection_limit=_get("bisection_limit", None),
+        as_of=parse_datetime(_get("as_of")) if _get("as_of") else None,
         # Application Settings
         log_level=_get("log_level", DEFAULT_LOG_LEVEL),
         dry_run=_get("dry_run", False),
@@ -477,6 +488,7 @@ def build_config_from_event(
         "restart_action": default_config.restart_action,
         "revert_action": default_config.revert_action,
         "bisection_limit": default_config.bisection_limit,
+        "as_of": None,  # Not supported in Lambda invocation
         # Application Settings
         "log_level": default_config.log_level,
         # Force subcommand to autorevert-checker for Lambda
@@ -674,6 +686,7 @@ def main_run(
                 else (config.revert_action or RevertAction.LOG)
             ),
             bisection_limit=config.bisection_limit,
+            as_of=config.as_of,
         )
         write_hud_html_from_cli(config.hud_html, HUD_HTML_NO_VALUE_FLAG, state_json)
     elif config.subcommand == "workflow-restart-checker":

aws/lambda/pytorch-auto-revert/pytorch_auto_revert/config.py

@@ -6,6 +6,7 @@
 """
 
 from dataclasses import dataclass, field
+from datetime import datetime
 from typing import Optional
 
 from .utils import RestartAction, RevertAction
@@ -71,6 +72,7 @@ class AutorevertConfig:
     restart_action: Optional[RestartAction] = None
     revert_action: Optional[RevertAction] = None
     bisection_limit: Optional[int] = None
+    as_of: Optional[datetime] = None
 
     # -------------------------------------------------------------------------
     # Application Settings

aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_extraction.py

@@ -41,10 +41,12 @@ def __init__(
         workflows: Iterable[str],
         lookback_hours: int = 24,
         repo_full_name: str = "pytorch/pytorch",
+        as_of: Optional[datetime] = None,
     ) -> None:
         self.workflows = list(workflows)
         self.lookback_hours = lookback_hours
         self.repo_full_name = repo_full_name
+        self.as_of = as_of
         # Datasource for DB access
         self._datasource = SignalExtractionDatasource()
 
@@ -69,6 +71,7 @@ def extract(self) -> List[Signal]:
         commits = self._datasource.fetch_commits_in_time_range(
             repo_full_name=self.repo_full_name,
             lookback_hours=self.lookback_hours,
+            as_of=self.as_of,
         )
 
         # Fetch jobs for these commits
@@ -77,6 +80,7 @@ def extract(self) -> List[Signal]:
             workflows=self.workflows,
             lookback_hours=self.lookback_hours,
             head_shas=[sha for sha, _ in commits],
+            as_of=self.as_of,
         )
 
         # Select jobs to participate in test-track details fetch

aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_extraction_datasource.py

@@ -25,7 +25,11 @@ class SignalExtractionDatasource:
     """
 
     def fetch_commits_in_time_range(
-        self, *, repo_full_name: str, lookback_hours: int
+        self,
+        *,
+        repo_full_name: str,
+        lookback_hours: int,
+        as_of: Optional[datetime] = None,
     ) -> List[tuple[Sha, datetime]]:
         """
         Fetch all commits pushed to main within the lookback window.
@@ -35,34 +39,49 @@ def fetch_commits_in_time_range(
         which can contain multiple commits in a single push event.
         For commits with identical timestamps (from the same ghstack push),
         orders by array index descending to match HUD/torchci ordering.
+
+        Args:
+            as_of: If set, use this as the reference time instead of now.
         """
-        lookback_time = datetime.now() - timedelta(hours=lookback_hours)
+        reference_time = as_of if as_of else datetime.now()
+        lookback_time = reference_time - timedelta(hours=lookback_hours)
 
-        query = """
+        # Add upper bound filter only when as_of is specified
+        upper_bound_clause = ""
+        if as_of:
+            upper_bound_clause = "AND head_commit.timestamp <= {as_of:DateTime}"
+
+        query = f"""
         SELECT
             commit.id AS sha,
             max(commit.timestamp) AS ts,
             arrayMax(groupArray(arrayFirstIndex(c -> c.'id' = commit.id, commits))) as array_index
         FROM default.push
         ARRAY JOIN commits as commit, commits as c
-        WHERE head_commit.timestamp >= {lookback_time:DateTime}
+        WHERE head_commit.timestamp >= {{lookback_time:DateTime}}
+          {upper_bound_clause}
           AND ref = 'refs/heads/main'
-          AND dynamoKey like {repo:String}
+          AND dynamoKey like {{repo:String}}
         GROUP BY sha
         ORDER BY ts DESC, array_index DESC
         """
 
-        params = {
+        params: Dict[str, Any] = {
             "lookback_time": lookback_time,
             "repo": f"{repo_full_name}%",
         }
+        if as_of:
+            params["as_of"] = as_of
 
         log = logging.getLogger(__name__)
         log.info(
-            "[extract] Fetching commits in time range: repo=%s lookback=%sh",
+            "[extract] Fetching commits in time range: repo=%s lookback=%sh as_of=%s",
             repo_full_name,
             lookback_hours,
+            as_of.isoformat() if as_of else "none",
         )
+        log.debug("[extract] Query params: %s", params)
+        log.debug("[extract] Query: %s", query)
         t0 = time.perf_counter()
         for attempt in RetryWithBackoff():
             with attempt:
@@ -79,20 +98,32 @@ def fetch_jobs_for_workflows(
         workflows: Iterable[str],
         lookback_hours: int,
         head_shas: List[Sha],
+        as_of: Optional[datetime] = None,
     ) -> List[JobRow]:
         """
         Fetch workflow job rows for the given head_shas and workflows.
 
         Returns rows ordered by head_sha (following the order of head_shas), then by started_at ASC.
+
+        Args:
+            as_of: If set, use this as the reference time instead of now.
         """
-        lookback_time = datetime.now() - timedelta(hours=lookback_hours)
+        reference_time = as_of if as_of else datetime.now()
+        lookback_time = reference_time - timedelta(hours=lookback_hours)
+
+        # Add upper bound filter only when as_of is specified
+        upper_bound_clause = ""
+        if as_of:
+            upper_bound_clause = "AND wf.created_at <= {as_of:DateTime}"
 
         workflow_filter = ""
         params: Dict[str, Any] = {
             "lookback_time": lookback_time,
             "repo": repo_full_name,
             "head_shas": [str(s) for s in head_shas],
         }
+        if as_of:
+            params["as_of"] = as_of
         workflow_list = list(workflows)
         if workflow_list:
             workflow_filter = "AND wf.workflow_name IN {workflows:Array(String)}"
@@ -126,6 +157,7 @@ def fetch_jobs_for_workflows(
         WHERE wf.repository_full_name = {{repo:String}}
           AND wf.head_sha IN {{head_shas:Array(String)}}
           AND wf.created_at >= {{lookback_time:DateTime}}
+          {upper_bound_clause}
           AND (
                 wf.name NOT LIKE '%mem_leak_check%'
                 AND wf.name NOT LIKE '%rerun_disabled_tests%'

aws/lambda/pytorch-auto-revert/pytorch_auto_revert/testers/autorevert_v2.py

@@ -19,6 +19,7 @@ def autorevert_v2(
     restart_action: RestartAction = RestartAction.RUN,
     revert_action: RevertAction = RevertAction.LOG,
     bisection_limit: Optional[int] = None,
+    as_of: Optional[datetime] = None,
 ) -> Tuple[List[Signal], List[Tuple[Signal, SignalProcOutcome]], str]:
     """Run the Signals-based autorevert flow end-to-end.
 
@@ -35,19 +36,23 @@ def autorevert_v2(
 
     logging.info(
         "[v2] Start: workflows=%s hours=%s repo=%s restart_action=%s"
-        " revert_action=%s notify_issue_number=%s bisection=%s",
+        " revert_action=%s notify_issue_number=%s bisection=%s as_of=%s",
         ",".join(workflows),
         hours,
         repo_full_name,
         restart_action,
         revert_action,
         notify_issue_number,
         ("unlimited" if bisection_limit is None else f"limit={bisection_limit}"),
+        (as_of.isoformat() if as_of else "now"),
     )
     logging.info("[v2] Run timestamp (CH log ts) = %s", ts.isoformat())
 
     extractor = SignalExtractor(
-        workflows=workflows, lookback_hours=hours, repo_full_name=repo_full_name
+        workflows=workflows,
+        lookback_hours=hours,
+        repo_full_name=repo_full_name,
+        as_of=as_of,
     )
     signals = extractor.extract()
     logging.info("[v2] Extracted %d signals", len(signals))

aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_parse_datetime.py

@@ -0,0 +1,68 @@
+import sys
+import unittest
+from datetime import datetime, timezone
+
+
+# Ensure package import when running from repo root
+sys.path.insert(0, "aws/lambda/pytorch-auto-revert")
+
+from pytorch_auto_revert.utils import parse_datetime  # noqa: E402
+
+
+class TestParseDatetime(unittest.TestCase):
+    def test_iso8601_with_seconds(self):
+        result = parse_datetime("2025-12-18T15:31:45")
+        expected = datetime(2025, 12, 18, 15, 31, 45, tzinfo=timezone.utc)
+        self.assertEqual(result, expected)
+
+    def test_iso8601_without_seconds(self):
+        result = parse_datetime("2025-12-18T15:31")
+        expected = datetime(2025, 12, 18, 15, 31, 0, tzinfo=timezone.utc)
+        self.assertEqual(result, expected)
+
+    def test_space_separated_with_seconds(self):
+        result = parse_datetime("2025-12-18 15:31:45")
+        expected = datetime(2025, 12, 18, 15, 31, 45, tzinfo=timezone.utc)
+        self.assertEqual(result, expected)
+
+    def test_space_separated_without_seconds(self):
+        result = parse_datetime("2025-12-18 15:31")
+        expected = datetime(2025, 12, 18, 15, 31, 0, tzinfo=timezone.utc)
+        self.assertEqual(result, expected)
+
+    def test_date_only(self):
+        result = parse_datetime("2025-12-18")
+        expected = datetime(2025, 12, 18, 0, 0, 0, tzinfo=timezone.utc)
+        self.assertEqual(result, expected)
+
+    def test_result_is_utc_aware(self):
+        result = parse_datetime("2025-12-18T15:31")
+        self.assertIsNotNone(result.tzinfo)
+        self.assertEqual(result.tzinfo, timezone.utc)
+
+    def test_invalid_format_raises_valueerror(self):
+        with self.assertRaises(ValueError) as ctx:
+            parse_datetime("not-a-date")
+        self.assertIn("Cannot parse datetime", str(ctx.exception))
+
+    def test_partial_date_raises_valueerror(self):
+        with self.assertRaises(ValueError):
+            parse_datetime("2025-12")
+
+    def test_empty_string_raises_valueerror(self):
+        with self.assertRaises(ValueError):
+            parse_datetime("")
+
+    def test_midnight_boundary(self):
+        result = parse_datetime("2025-01-01T00:00:00")
+        expected = datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
+        self.assertEqual(result, expected)
+
+    def test_end_of_day(self):
+        result = parse_datetime("2025-12-31T23:59:59")
+        expected = datetime(2025, 12, 31, 23, 59, 59, tzinfo=timezone.utc)
+        self.assertEqual(result, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()

aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_extraction.py

@@ -1,6 +1,6 @@
 import unittest
 from datetime import datetime, timedelta
-from typing import Iterable, List
+from typing import Iterable, List, Optional
 
 from pytorch_auto_revert.signal import SignalStatus
 from pytorch_auto_revert.signal_extraction import SignalExtractor
@@ -30,7 +30,11 @@ def __init__(self, jobs: List[JobRow], tests: List[TestRow]):
         self._tests = tests
 
     def fetch_commits_in_time_range(
-        self, *, repo_full_name: str, lookback_hours: int
+        self,
+        *,
+        repo_full_name: str,
+        lookback_hours: int,
+        as_of: Optional[datetime] = None,
     ) -> List[tuple[Sha, datetime]]:
         # Extract unique commits from jobs in the order they appear
         seen = set()
@@ -48,6 +52,7 @@ def fetch_jobs_for_workflows(
         lookback_hours: int,
         repo_full_name: str,
         head_shas: List[Sha],
+        as_of: Optional[datetime] = None,
     ) -> List[JobRow]:
         return list(self._jobs)
 
@@ -394,7 +399,11 @@ def test_commits_without_jobs_are_included(self):
 
         class FakeDatasourceWithExtraCommit(FakeDatasource):
             def fetch_commits_in_time_range(
-                self, *, repo_full_name: str, lookback_hours: int
+                self,
+                *,
+                repo_full_name: str,
+                lookback_hours: int,
+                as_of: Optional[datetime] = None,
             ):
                 # Return commits C2, C3 (no jobs), C1 in newest->older order
                 return [

aws/lambda/pytorch-auto-revert/pytorch_auto_revert/utils.py

@@ -1,6 +1,7 @@
 import random
 import time
 import urllib.parse
+from datetime import datetime, timezone
 from enum import Enum
 
 import github
@@ -172,3 +173,40 @@ def proper_workflow_create_dispatch(
     if status != 204:
         raise ValueError(f"Error dispatching workflow: {status}, {headers}, {body}")
     return True
+
+
+def parse_datetime(s: str) -> datetime:
+    """Parse datetime string in various formats (ISO 8601 or common formats).
+
+    The input is interpreted as UTC time and returns a timezone-aware datetime.
+
+    Supported formats:
+        - 2025-12-18T15:31:00 (ISO 8601 with seconds)
+        - 2025-12-18T15:31 (ISO 8601 without seconds)
+        - 2025-12-18 15:31:00 (space-separated with seconds)
+        - 2025-12-18 15:31 (space-separated without seconds)
+        - 2025-12-18 (date only, time defaults to 00:00:00)
+
+    Args:
+        s: Datetime string to parse.
+
+    Returns:
+        Timezone-aware datetime in UTC.
+
+    Raises:
+        ValueError: If the string cannot be parsed in any supported format.
+    """
+    for fmt in (
+        "%Y-%m-%dT%H:%M:%S",
+        "%Y-%m-%dT%H:%M",
+        "%Y-%m-%d %H:%M:%S",
+        "%Y-%m-%d %H:%M",
+        "%Y-%m-%d",
+    ):
+        try:
+            naive = datetime.strptime(s, fmt)
+            # Attach UTC timezone since input is expected to be UTC
+            return naive.replace(tzinfo=timezone.utc)
+        except ValueError:
+            continue
+    raise ValueError(f"Cannot parse datetime: {s!r}")