From 83e7f8af5490af5a5d15d063b069c73a9b4fe1e6 Mon Sep 17 00:00:00 2001 From: Ivan Zaitsev Date: Tue, 31 Mar 2026 12:46:05 -0700 Subject: [PATCH 1/3] feat: read AI advisor verdicts from CH and use in autorevert decisions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When advisor verdicts are available in misc.autorevert_advisor_verdicts, autorevert now uses them to make faster decisions: - revert → produce AutorevertPattern immediately (skip restarts) - not_related → early exit (Ineligible) - garbage → suppress signal for 2 hours since verdict timestamp - unsure → continue normal restart-to-confirm flow Key changes: - AIAdvisorResult dataclass + AdvisorVerdict enum in signal.py - SignalCommit gets optional advisor_result field - fetch_advisor_verdicts() query in signal_extraction_datasource.py - _attach_advisor_verdicts() in signal_extraction.py wires CH data to SignalCommits - Flakiness check moved AFTER advisor check/dispatch so advisor can still evaluate flaky-looking signals - 8 new tests covering all verdict types, expiry, signal key matching --- .../pytorch_auto_revert/signal.py | 155 ++++++++-- .../pytorch_auto_revert/signal_extraction.py | 80 ++++- .../signal_extraction_datasource.py | 58 ++++ .../pytorch_auto_revert/tests/test_signal.py | 290 +++++++++++++++++- .../tests/test_signal_actions.py | 132 ++++++++ .../grants.sql | 2 + 6 files changed, 686 insertions(+), 31 deletions(-) create mode 100644 clickhouse_db_schema/misc.autorevert_advisor_verdicts/grants.sql diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal.py index 296dc21b01..040185c595 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal.py @@ -6,6 +6,29 @@ from .bisection_planner import GapBisectionPlanner +class AdvisorVerdict(Enum): + """Verdict from the AI advisor workflow.""" + + REVERT = "revert" + UNSURE = "unsure" + NOT_RELATED = "not_related" + GARBAGE = "garbage" + + +@dataclass(frozen=True) +class AIAdvisorResult: + """Result from a previously-dispatched AI advisor run, read from ClickHouse. + + Attached to SignalCommit so that process_valid_autorevert_pattern() can + use the advisor's analysis to make faster revert/skip decisions. + """ + + verdict: AdvisorVerdict + confidence: float + timestamp: datetime # when the verdict was produced + signal_key: str # the signal key this verdict applies to + + class SignalStatus(Enum): """signal status enum""" @@ -76,6 +99,8 @@ class IneligibleReason(Enum): "insufficient_successes" # not enough successes to make call ) PENDING_GAP = "pending_gap" # unknown/pending commits present + ADVISOR_NOT_RELATED = "advisor_not_related" # AI advisor says not related + ADVISOR_GARBAGE = "advisor_garbage" # AI advisor says signal is garbage @dataclass @@ -128,7 +153,13 @@ def is_failure(self) -> bool: class SignalCommit: """All events for a single commit, ordered oldest → newest by start time.""" - def __init__(self, head_sha: str, timestamp: datetime, events: List[SignalEvent]): + def __init__( + self, + head_sha: str, + timestamp: datetime, + events: List[SignalEvent], + advisor_result: Optional[AIAdvisorResult] = None, + ): self.head_sha = head_sha self.timestamp = timestamp # enforce events ordered by time, then by wf_run_id (oldest first) @@ -139,6 +170,8 @@ def __init__(self, head_sha: str, timestamp: datetime, events: List[SignalEvent] self.statuses = {} for e in self.events: self.statuses[e.status] = self.statuses.get(e.status, 0) + 1 + # Optional AI advisor result for this (commit, signal) pair + self.advisor_result = advisor_result @property def has_pending(self) -> bool: @@ -384,8 +417,10 @@ def partition_by_autorevert_pattern(self) -> Optional[PartitionedCommits]: picking_failed = True # simple state machine # first broadly partition into failed and successful + # A commit with both success and failure (flaky) stays in the failed + # partition — only pure successes trigger the transition. for c in self.commits: - if c.has_success: + if c.has_success and not c.has_failure: picking_failed = False elif c.has_failure and not picking_failed: # encountered a failure after the streak of successes @@ -409,24 +444,93 @@ def partition_by_autorevert_pattern(self) -> Optional[PartitionedCommits]: return PartitionedCommits(failed=failed, unknown=unknown, successful=successful) + # Minimum confidence threshold for acting on advisor verdicts + ADVISOR_CONFIDENCE_THRESHOLD = 0.9 + + def _build_autorevert_pattern( + self, partition: "PartitionedCommits" + ) -> AutorevertPattern: + """Build an AutorevertPattern from a validated partition.""" + suspected = partition.failed[-1] + newer_failures = [c.head_sha for c in partition.failed[:-1]] + failure_event = next( + (e for e in suspected.events if e.is_failure and e.job_id is not None), + None, + ) + return AutorevertPattern( + workflow_name=self.workflow_name, + newer_failing_commits=newer_failures, + suspected_commit=suspected.head_sha, + older_successful_commit=partition.successful[0].head_sha, + wf_run_id=failure_event.wf_run_id if failure_event else None, + job_id=failure_event.job_id if failure_event else None, + ) + + def _check_advisor_verdict( + self, partition: "PartitionedCommits" + ) -> Optional[Union[AutorevertPattern, Ineligible]]: + """Check if an AI advisor verdict is available for the suspect commit. + + Returns: + - AutorevertPattern if advisor says "revert" with sufficient confidence + - Ineligible if advisor says "not_related" or "garbage" (within 2h window) + - None if no verdict, verdict is "unsure", or confidence below threshold + """ + suspected = partition.failed[-1] + result = suspected.advisor_result + if result is None: + return None + + # Only act on the verdict if it matches this signal + if result.signal_key != self.key: + return None + + # Ignore low-confidence verdicts + if result.confidence < self.ADVISOR_CONFIDENCE_THRESHOLD: + return None + + if result.verdict == AdvisorVerdict.REVERT: + return self._build_autorevert_pattern(partition) + + if result.verdict == AdvisorVerdict.NOT_RELATED: + return Ineligible( + IneligibleReason.ADVISOR_NOT_RELATED, + f"AI advisor says not related (confidence={result.confidence:.2f})", + ) + + if result.verdict == AdvisorVerdict.GARBAGE: + # Garbage verdict blocks the signal for 2 hours since the verdict timestamp + from datetime import timezone + + now = datetime.now(timezone.utc) + verdict_age = now - result.timestamp.replace(tzinfo=timezone.utc) + if verdict_age.total_seconds() < 2 * 3600: + return Ineligible( + IneligibleReason.ADVISOR_GARBAGE, + f"AI advisor says garbage signal, suppressed for 2h " + f"(confidence={result.confidence:.2f}, " + f"age={int(verdict_age.total_seconds() / 60)}min)", + ) + # Garbage verdict expired — fall through to normal processing + + # "unsure" or expired garbage → continue normally + return None + def process_valid_autorevert_pattern( self, *, bisection_limit: Optional[int] = None ) -> Union[AutorevertPattern, RestartCommits, Ineligible]: """ Detect valid autorevert pattern in the Signal. - Validates all invariants before checking for the pattern. + Validates invariants, checks AI advisor verdicts, and determines + appropriate action. Returns one of: - AutorevertPattern: a confirmed pattern ready for action - RestartCommits: a suggested set of commits to restart to reduce uncertainty - Ineligible: reason + optional message when no pattern is actionable yet """ - if self.detect_flaky(): - return Ineligible( - IneligibleReason.FLAKY, - "signal is flaky (mixed outcomes on same commit)", - ) + # Early exits that prevent partition computation if self.detect_fixed(): return Ineligible( IneligibleReason.FIXED, "signal appears recovered at head" @@ -443,7 +547,12 @@ def process_valid_autorevert_pattern( "insufficient history to form failed/unknown/successful partitions", ) - # Advisor eligibility: partition exists with sufficient data and + # --- AI advisor verdict check (before flakiness and other heuristics) --- + advisor_decision = self._check_advisor_verdict(partition) + if advisor_decision is not None: + return advisor_decision + + # Advisor dispatch eligibility: partition exists with sufficient data and # no unknown commits between failed and successful partitions. # Unknown commits mean the partition boundary is uncertain — # wait for restarts to resolve before dispatching advisor. @@ -459,6 +568,15 @@ def process_valid_autorevert_pattern( successful_commits=tuple(c.head_sha for c in partition.successful), ) + # Flakiness check — placed after advisor check/dispatch so that + # advisor can still evaluate flaky-looking signals + if self.detect_flaky(): + return Ineligible( + IneligibleReason.FLAKY, + "signal is flaky (mixed outcomes on same commit)", + advisor=advisor, + ) + restart_commits = set() # Cover the unknown gap (between failed and successful partitions) @@ -554,21 +672,4 @@ def process_valid_autorevert_pattern( ) # all invariants validated, confirmed not infra, pattern exists - # failed is newest -> older; the last element is the suspected commit - suspected = partition.failed[-1] - newer_failures = [c.head_sha for c in partition.failed[:-1]] - - # Extract job_id and wf_run_id from a failing event on the suspected commit - failure_event = next( - (e for e in suspected.events if e.is_failure and e.job_id is not None), - None, - ) - - return AutorevertPattern( - workflow_name=self.workflow_name, - newer_failing_commits=newer_failures, - suspected_commit=suspected.head_sha, - older_successful_commit=partition.successful[0].head_sha, - wf_run_id=failure_event.wf_run_id if failure_event else None, - job_id=failure_event.job_id if failure_event else None, - ) + return self._build_autorevert_pattern(partition) diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_extraction.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_extraction.py index 0008efb5b2..c96068c519 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_extraction.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_extraction.py @@ -12,7 +12,15 @@ from typing import Dict, Iterable, List, Optional, Set, Tuple from .job_agg_index import JobAggIndex, JobMeta, SignalStatus as AggStatus -from .signal import Signal, SignalCommit, SignalEvent, SignalSource, SignalStatus +from .signal import ( + AdvisorVerdict, + AIAdvisorResult, + Signal, + SignalCommit, + SignalEvent, + SignalSource, + SignalStatus, +) from .signal_extraction_datasource import SignalExtractionDatasource from .signal_extraction_types import ( JobBaseName, @@ -107,7 +115,10 @@ def extract(self) -> List[Signal]: # Inject synthetic PENDING events for workflow runs that are known to be # pending but have no events in a given signal (e.g. multi-stage workflows). - return self._inject_pending_workflow_events(signals, jobs) + signals = self._inject_pending_workflow_events(signals, jobs) + + # Attach AI advisor verdicts to SignalCommit objects + return self._attach_advisor_verdicts(signals, commits) # ----------------------------- # Deduplication (GitHub-specific) @@ -230,6 +241,71 @@ def _inject_pending_workflow_events( ) return out + # ----------------------------- + # Advisor verdict attachment + # ----------------------------- + def _attach_advisor_verdicts( + self, + signals: List[Signal], + commits: List[Tuple[Sha, datetime]], + ) -> List[Signal]: + """Fetch advisor verdicts from CH and attach to SignalCommit objects. + + For each (commit_sha, signal_key) pair that has a verdict in + misc.autorevert_advisor_verdicts, sets the advisor_result field + on the corresponding SignalCommit. + """ + head_shas = [sha for sha, _ in commits] + signal_keys = list({s.key for s in signals}) + verdicts = self._datasource.fetch_advisor_verdicts( + repo_full_name=self.repo_full_name, + head_shas=head_shas, + signal_keys=signal_keys, + lookback_hours=self.lookback_hours, + ) + if not verdicts: + return signals + + out: List[Signal] = [] + for s in signals: + new_commits: List[SignalCommit] = [] + for c in s.commits: + key = (c.head_sha.strip(), s.key) + v = verdicts.get(key) + if v is not None: + verdict_str, confidence, ts = v + try: + advisor_verdict = AdvisorVerdict(verdict_str) + except ValueError: + advisor_verdict = AdvisorVerdict.UNSURE + advisor_result = AIAdvisorResult( + verdict=advisor_verdict, + confidence=confidence, + timestamp=ts, + signal_key=s.key, + ) + new_commits.append( + SignalCommit( + head_sha=c.head_sha, + timestamp=c.timestamp, + events=c.events, + advisor_result=advisor_result, + ) + ) + else: + new_commits.append(c) + out.append( + Signal( + key=s.key, + workflow_name=s.workflow_name, + commits=new_commits, + job_base_name=s.job_base_name, + test_module=s.test_module, + source=s.source, + ) + ) + return out + # ----------------------------- # Phase B — Tests (tests.all_test_runs only) # ----------------------------- diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_extraction_datasource.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_extraction_datasource.py index d4f43b9960..9d491ee8ea 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_extraction_datasource.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_extraction_datasource.py @@ -313,6 +313,64 @@ def fetch_tests_for_job_ids( ) return rows + def fetch_advisor_verdicts( + self, + *, + repo_full_name: str, + head_shas: List[Sha], + signal_keys: List[str], + lookback_hours: int, + ) -> Dict[tuple[str, str], tuple[str, float, datetime]]: + """Fetch AI advisor verdicts from misc.autorevert_advisor_verdicts. + + Queries by both commit SHAs AND signal keys to minimize data transferred. + + Returns a dict keyed by (commit_sha, signal_key) → (verdict, confidence, timestamp). + When multiple verdicts exist for the same (commit, signal), the most recent is used. + """ + if not head_shas or not signal_keys: + return {} + + log = logging.getLogger(__name__) + t0 = time.perf_counter() + query = """ + SELECT + toString(suspect_commit) AS suspect_commit, + signal_key, + verdict, + confidence, + timestamp + FROM misc.autorevert_advisor_verdicts + WHERE repo = {repo:String} + AND suspect_commit IN {shas:Array(String)} + AND signal_key IN {keys:Array(String)} + AND timestamp > now() - INTERVAL {hours:UInt32} HOUR + ORDER BY suspect_commit, signal_key, timestamp DESC + """ + params = { + "repo": repo_full_name, + "shas": [str(s) for s in head_shas], + "keys": signal_keys, + "hours": lookback_hours, + } + results: Dict[tuple[str, str], tuple[str, float, datetime]] = {} + for attempt in RetryWithBackoff(): + with attempt: + res = CHCliFactory().client.query(query, parameters=params) + for r in res.result_rows: + key = (str(r[0]).strip(), str(r[1])) + if key not in results: # keep most recent (ORDER BY ... DESC) + results[key] = (str(r[2]), float(r[3]), r[4]) + + dt = time.perf_counter() - t0 + log.info( + "[extract] Advisor verdicts fetched: %d for %d commits in %.2fs", + len(results), + len(head_shas), + dt, + ) + return results + def fetch_autorevert_state_rows( self, *, diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal.py index 0382b12606..6b18bbe477 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal.py @@ -1,7 +1,9 @@ import unittest -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from pytorch_auto_revert.signal import ( + AdvisorVerdict, + AIAdvisorResult, AutorevertPattern, DispatchAdvisor, Ineligible, @@ -772,7 +774,8 @@ def test_advisor_not_emitted_for_early_returns(self): s = Signal(key="job", workflow_name="wf", commits=[c_flaky, c_base]) res = s.process_valid_autorevert_pattern() self.assertIsInstance(res, Ineligible) - self.assertEqual(res.reason, IneligibleReason.FLAKY) + # Flaky commit has success → detect_fixed() returns True (early exit before partition) + self.assertEqual(res.reason, IneligibleReason.FIXED) self.assertIsNone(res.advisor) def test_advisor_has_correct_partition_shas(self): @@ -815,5 +818,288 @@ def test_advisor_has_correct_partition_shas(self): self.assertEqual(advisor.suspect_commit, "fail_2") +class TestAdvisorVerdictIntegration(unittest.TestCase): + """Tests for AI advisor verdict handling in process_valid_autorevert_pattern.""" + + def setUp(self) -> None: + self.t0 = datetime(2025, 8, 19, 12, 0, 0) + + def _ev(self, name: str, status: SignalStatus, minute: int) -> SignalEvent: + return SignalEvent( + name=name, + status=status, + started_at=ts(self.t0, minute), + wf_run_id=1, + ) + + def _make_signal_with_advisor( + self, verdict: AdvisorVerdict, confidence: float = 0.95 + ) -> Signal: + """Build a signal with 3 failures and 2 successes where the suspect + commit has an advisor result.""" + advisor_result = AIAdvisorResult( + verdict=verdict, + confidence=confidence, + timestamp=self.t0, + signal_key="job", + ) + c_newest = SignalCommit( + head_sha="sha_newest", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 7)], + ) + c_newer = SignalCommit( + head_sha="sha_newer", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 5)], + ) + c_suspected = SignalCommit( + head_sha="sha_mid", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 4)], + advisor_result=advisor_result, + ) + c_base = SignalCommit( + head_sha="sha_old", + timestamp=ts(self.t0, 0), + events=[ + self._ev("job", SignalStatus.SUCCESS, 3), + self._ev("job", SignalStatus.SUCCESS, 6), + ], + ) + return Signal( + key="job", + workflow_name="wf", + commits=[c_newest, c_newer, c_suspected, c_base], + ) + + def test_advisor_revert_produces_autorevert_pattern(self): + """When advisor says 'revert', produce AutorevertPattern immediately.""" + s = self._make_signal_with_advisor(AdvisorVerdict.REVERT) + res = s.process_valid_autorevert_pattern() + self.assertIsInstance(res, AutorevertPattern) + self.assertEqual(res.suspected_commit, "sha_mid") + self.assertEqual(res.older_successful_commit, "sha_old") + self.assertEqual(res.newer_failing_commits, ["sha_newest", "sha_newer"]) + + def test_advisor_not_related_produces_ineligible(self): + """When advisor says 'not_related', return Ineligible.""" + s = self._make_signal_with_advisor(AdvisorVerdict.NOT_RELATED) + res = s.process_valid_autorevert_pattern() + self.assertIsInstance(res, Ineligible) + self.assertEqual(res.reason, IneligibleReason.ADVISOR_NOT_RELATED) + + def test_advisor_garbage_produces_ineligible_within_2h(self): + """When advisor says 'garbage' and verdict is < 2h old, suppress signal.""" + # Use a recent timestamp + advisor_result = AIAdvisorResult( + verdict=AdvisorVerdict.GARBAGE, + confidence=0.95, + timestamp=datetime.now(tz=timezone.utc), + signal_key="job", + ) + c_fail = SignalCommit( + head_sha="sha_fail", + timestamp=ts(self.t0, 0), + events=[ + self._ev("job", SignalStatus.FAILURE, 4), + self._ev("job", SignalStatus.FAILURE, 5), + self._ev("job", SignalStatus.FAILURE, 6), + ], + advisor_result=advisor_result, + ) + c_base = SignalCommit( + head_sha="sha_base", + timestamp=ts(self.t0, 0), + events=[ + self._ev("job", SignalStatus.SUCCESS, 2), + self._ev("job", SignalStatus.SUCCESS, 3), + ], + ) + s = Signal(key="job", workflow_name="wf", commits=[c_fail, c_base]) + res = s.process_valid_autorevert_pattern() + self.assertIsInstance(res, Ineligible) + self.assertEqual(res.reason, IneligibleReason.ADVISOR_GARBAGE) + + def test_advisor_garbage_expires_after_2h(self): + """When garbage verdict is > 2h old, it expires and normal processing resumes.""" + old_timestamp = datetime.now(tz=timezone.utc) - timedelta(hours=3) + advisor_result = AIAdvisorResult( + verdict=AdvisorVerdict.GARBAGE, + confidence=0.95, + timestamp=old_timestamp, + signal_key="job", + ) + c_newest = SignalCommit( + head_sha="sha_newest", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 7)], + ) + c_newer = SignalCommit( + head_sha="sha_newer", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 5)], + ) + c_suspected = SignalCommit( + head_sha="sha_mid", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 4)], + advisor_result=advisor_result, + ) + c_base = SignalCommit( + head_sha="sha_old", + timestamp=ts(self.t0, 0), + events=[ + self._ev("job", SignalStatus.SUCCESS, 3), + self._ev("job", SignalStatus.SUCCESS, 6), + ], + ) + s = Signal( + key="job", + workflow_name="wf", + commits=[c_newest, c_newer, c_suspected, c_base], + ) + res = s.process_valid_autorevert_pattern() + # Garbage expired — should proceed to AutorevertPattern + self.assertIsInstance(res, AutorevertPattern) + + def test_advisor_unsure_continues_normal_processing(self): + """When advisor says 'unsure', continue with normal autorevert logic.""" + s = self._make_signal_with_advisor(AdvisorVerdict.UNSURE) + res = s.process_valid_autorevert_pattern() + # With 3 failures and 2 successes, should produce AutorevertPattern + self.assertIsInstance(res, AutorevertPattern) + + def test_advisor_low_confidence_ignored(self): + """Advisor verdict below confidence threshold is ignored.""" + s = self._make_signal_with_advisor(AdvisorVerdict.NOT_RELATED, confidence=0.5) + res = s.process_valid_autorevert_pattern() + # Low confidence NOT_RELATED should be ignored → normal AutorevertPattern + self.assertIsInstance(res, AutorevertPattern) + + def test_advisor_high_confidence_revert_acts(self): + """Advisor 'revert' at exactly the threshold acts.""" + s = self._make_signal_with_advisor(AdvisorVerdict.REVERT, confidence=0.9) + res = s.process_valid_autorevert_pattern() + self.assertIsInstance(res, AutorevertPattern) + self.assertEqual(res.suspected_commit, "sha_mid") + + def test_advisor_wrong_signal_key_ignored(self): + """Advisor result for a different signal_key is ignored.""" + advisor_result = AIAdvisorResult( + verdict=AdvisorVerdict.NOT_RELATED, + confidence=0.99, + timestamp=self.t0, + signal_key="different_signal", # doesn't match "job" + ) + c_newest = SignalCommit( + head_sha="sha_newest", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 7)], + ) + c_newer = SignalCommit( + head_sha="sha_newer", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 5)], + ) + c_suspected = SignalCommit( + head_sha="sha_mid", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 4)], + advisor_result=advisor_result, + ) + c_base = SignalCommit( + head_sha="sha_old", + timestamp=ts(self.t0, 0), + events=[ + self._ev("job", SignalStatus.SUCCESS, 3), + self._ev("job", SignalStatus.SUCCESS, 6), + ], + ) + s = Signal( + key="job", + workflow_name="wf", + commits=[c_newest, c_newer, c_suspected, c_base], + ) + res = s.process_valid_autorevert_pattern() + # NOT_RELATED verdict is for wrong signal, should be ignored + self.assertIsInstance(res, AutorevertPattern) + + def test_flaky_check_after_advisor(self): + """Flaky check happens after advisor, so advisor can still fire on flaky signals.""" + # Signal with a flaky commit (suspect) that has an advisor revert verdict + advisor_result = AIAdvisorResult( + verdict=AdvisorVerdict.REVERT, + confidence=0.95, + timestamp=self.t0, + signal_key="job", + ) + c_newest = SignalCommit( + head_sha="sha_newest", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 5)], + ) + # Suspect commit is flaky (both success and failure) — with partition fix, + # it stays in the failed partition because it has failures + c_suspect_flaky = SignalCommit( + head_sha="sha_suspect", + timestamp=ts(self.t0, -5), + events=[ + self._ev("job", SignalStatus.SUCCESS, 2), + self._ev("job", SignalStatus.FAILURE, 3), + ], + advisor_result=advisor_result, + ) + c_base = SignalCommit( + head_sha="sha_base", + timestamp=ts(self.t0, -10), + events=[self._ev("job", SignalStatus.SUCCESS, 1)], + ) + s = Signal( + key="job", + workflow_name="wf", + commits=[c_newest, c_suspect_flaky, c_base], + ) + res = s.process_valid_autorevert_pattern() + # Advisor said revert on the suspect → AutorevertPattern + # even though the signal has a flaky commit + self.assertIsInstance(res, AutorevertPattern) + self.assertEqual(res.suspected_commit, "sha_suspect") + + def test_no_advisor_result_continues_normally(self): + """Without advisor result, processing is unchanged.""" + c_newest = SignalCommit( + head_sha="sha_newest", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 7)], + ) + c_newer = SignalCommit( + head_sha="sha_newer", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 5)], + ) + c_suspected = SignalCommit( + head_sha="sha_mid", + timestamp=ts(self.t0, 0), + events=[self._ev("job", SignalStatus.FAILURE, 4)], + # No advisor_result + ) + c_base = SignalCommit( + head_sha="sha_old", + timestamp=ts(self.t0, 0), + events=[ + self._ev("job", SignalStatus.SUCCESS, 3), + self._ev("job", SignalStatus.SUCCESS, 6), + ], + ) + s = Signal( + key="job", + workflow_name="wf", + commits=[c_newest, c_newer, c_suspected, c_base], + ) + res = s.process_valid_autorevert_pattern() + self.assertIsInstance(res, AutorevertPattern) + + if __name__ == "__main__": unittest.main() diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_actions.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_actions.py index 20f0f5be59..31d8816122 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_actions.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_actions.py @@ -1101,5 +1101,137 @@ def test_dispatch_metadata_format(self, mock_gh): self.assertEqual(d["mode"], "log") +class TestAttachAdvisorVerdicts(unittest.TestCase): + """Tests for SignalExtractor._attach_advisor_verdicts.""" + + def test_attaches_verdict_to_matching_commit(self): + from pytorch_auto_revert.signal import ( + AdvisorVerdict, + Signal, + SignalCommit, + SignalEvent, + SignalSource, + SignalStatus, + ) + from pytorch_auto_revert.signal_extraction import SignalExtractor + from pytorch_auto_revert.signal_extraction_types import Sha + + t0 = datetime(2025, 8, 19, 12, 0, 0) + c1 = SignalCommit("sha_aaa", t0, [ + SignalEvent("j", SignalStatus.FAILURE, t0, wf_run_id=1, job_id=10), + ]) + c2 = SignalCommit("sha_bbb", t0, [ + SignalEvent("j", SignalStatus.SUCCESS, t0, wf_run_id=2, job_id=20), + ]) + signal = Signal( + key="test_key", workflow_name="trunk", + commits=[c1, c2], source=SignalSource.TEST, + ) + + extractor = SignalExtractor(workflows=["trunk"], lookback_hours=16) + # Mock datasource to return a verdict for (sha_aaa, test_key) + extractor._datasource = Mock() + extractor._datasource.fetch_advisor_verdicts.return_value = { + ("sha_aaa", "test_key"): ("revert", 0.95, t0), + } + + commits = [(Sha("sha_aaa"), t0), (Sha("sha_bbb"), t0)] + result = extractor._attach_advisor_verdicts([signal], commits) + + self.assertEqual(len(result), 1) + s = result[0] + # sha_aaa should have advisor_result + self.assertIsNotNone(s.commits[0].advisor_result) + self.assertEqual(s.commits[0].advisor_result.verdict, AdvisorVerdict.REVERT) + self.assertAlmostEqual(s.commits[0].advisor_result.confidence, 0.95) + self.assertEqual(s.commits[0].advisor_result.signal_key, "test_key") + # sha_bbb should NOT have advisor_result + self.assertIsNone(s.commits[1].advisor_result) + + def test_no_verdicts_returns_signals_unchanged(self): + from pytorch_auto_revert.signal import ( + Signal, SignalCommit, SignalEvent, SignalSource, SignalStatus, + ) + from pytorch_auto_revert.signal_extraction import SignalExtractor + from pytorch_auto_revert.signal_extraction_types import Sha + + t0 = datetime(2025, 8, 19, 12, 0, 0) + c1 = SignalCommit("sha_aaa", t0, [ + SignalEvent("j", SignalStatus.FAILURE, t0, wf_run_id=1), + ]) + signal = Signal( + key="k", workflow_name="wf", commits=[c1], source=SignalSource.TEST, + ) + + extractor = SignalExtractor(workflows=["wf"], lookback_hours=16) + extractor._datasource = Mock() + extractor._datasource.fetch_advisor_verdicts.return_value = {} + + result = extractor._attach_advisor_verdicts([signal], [(Sha("sha_aaa"), t0)]) + + # Should return same signals (no modification) + self.assertIs(result[0], signal) + self.assertIsNone(result[0].commits[0].advisor_result) + + def test_passes_signal_keys_to_datasource(self): + from pytorch_auto_revert.signal import ( + Signal, SignalCommit, SignalEvent, SignalSource, SignalStatus, + ) + from pytorch_auto_revert.signal_extraction import SignalExtractor + from pytorch_auto_revert.signal_extraction_types import Sha + + t0 = datetime(2025, 8, 19, 12, 0, 0) + s1 = Signal( + key="key_a", workflow_name="wf", + commits=[SignalCommit("sha1", t0, [])], + source=SignalSource.TEST, + ) + s2 = Signal( + key="key_b", workflow_name="wf", + commits=[SignalCommit("sha1", t0, [])], + source=SignalSource.TEST, + ) + + extractor = SignalExtractor(workflows=["wf"], lookback_hours=16) + extractor._datasource = Mock() + extractor._datasource.fetch_advisor_verdicts.return_value = {} + + extractor._attach_advisor_verdicts([s1, s2], [(Sha("sha1"), t0)]) + + call_kwargs = extractor._datasource.fetch_advisor_verdicts.call_args[1] + self.assertIn("signal_keys", call_kwargs) + self.assertCountEqual(call_kwargs["signal_keys"], ["key_a", "key_b"]) + self.assertEqual(call_kwargs["head_shas"], [Sha("sha1")]) + + def test_invalid_verdict_string_defaults_to_unsure(self): + from pytorch_auto_revert.signal import ( + AdvisorVerdict, Signal, SignalCommit, SignalEvent, + SignalSource, SignalStatus, + ) + from pytorch_auto_revert.signal_extraction import SignalExtractor + from pytorch_auto_revert.signal_extraction_types import Sha + + t0 = datetime(2025, 8, 19, 12, 0, 0) + c1 = SignalCommit("sha_aaa", t0, [ + SignalEvent("j", SignalStatus.FAILURE, t0, wf_run_id=1), + ]) + signal = Signal( + key="k", workflow_name="wf", commits=[c1], source=SignalSource.TEST, + ) + + extractor = SignalExtractor(workflows=["wf"], lookback_hours=16) + extractor._datasource = Mock() + extractor._datasource.fetch_advisor_verdicts.return_value = { + ("sha_aaa", "k"): ("bogus_verdict", 0.5, t0), + } + + result = extractor._attach_advisor_verdicts([signal], [(Sha("sha_aaa"), t0)]) + + self.assertIsNotNone(result[0].commits[0].advisor_result) + self.assertEqual( + result[0].commits[0].advisor_result.verdict, AdvisorVerdict.UNSURE + ) + + if __name__ == "__main__": unittest.main() diff --git a/clickhouse_db_schema/misc.autorevert_advisor_verdicts/grants.sql b/clickhouse_db_schema/misc.autorevert_advisor_verdicts/grants.sql new file mode 100644 index 0000000000..a1d0fc4c8e --- /dev/null +++ b/clickhouse_db_schema/misc.autorevert_advisor_verdicts/grants.sql @@ -0,0 +1,2 @@ +-- Grant SELECT to revert_lambda (used by the autorevert lambda to read advisor verdicts) +GRANT SELECT ON misc.autorevert_advisor_verdicts TO revert_lambda; From bd297fa92d91fa6763789ab555267014b97422ac Mon Sep 17 00:00:00 2001 From: Ivan Zaitsev Date: Tue, 31 Mar 2026 15:49:54 -0700 Subject: [PATCH 2/3] feat: propagate advisor verdict through revert flow and render in HUD MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. AutorevertPattern carries optional advisor_verdict field. When a revert is advisor-accelerated, the revert comment includes "Note: This revert was accelerated by the AI advisor" with verdict and confidence. 2. State JSON includes advisor data in two places (both forward/backward compatible — absent in older states, gracefully ignored): - outcomes: advisor_verdict on AutorevertPattern data - columns: advisor_results map per (commit_sha → verdict/confidence) 3. HUD renderer shows: - "AI:revert" / "AI:not_related" / etc badges in table cells - "[AI: revert @95%]" in outcome notes for advisor-accelerated reverts --- .../pytorch_auto_revert/hud_renderer.py | 54 ++++++++++-- .../pytorch_auto_revert/run_state_logger.py | 18 +++- .../pytorch_auto_revert/signal.py | 10 ++- .../pytorch_auto_revert/signal_actions.py | 18 ++++ .../tests/test_signal_actions.py | 82 ++++++++++++++----- .../tests/test_signal_extraction.py | 3 + 6 files changed, 154 insertions(+), 31 deletions(-) diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/hud_renderer.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/hud_renderer.py index 0723b0a09c..dcd5ebb1ce 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/hud_renderer.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/hud_renderer.py @@ -2,7 +2,7 @@ import re from datetime import datetime -from typing import Any, Dict, List, Mapping, Optional, Sequence, Union +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union from .signal import SignalStatus from .utils import build_pytorch_hud_url @@ -53,6 +53,13 @@ td.cell.hl-baseline { background: #e6f7ff; } td.cell.hl-newer-fail { background: #fdecea; } td.cell.hl-restart { outline: 2px dashed #888; outline-offset: -2px; } + .advisor-cell { font-size: 10px; display: block; margin-top: 2px; } + .advisor-cell.adv-revert { color: #a40000; } + .advisor-cell.adv-not_related { color: #1a73e8; } + .advisor-cell.adv-garbage { color: #7a5a00; } + .advisor-cell.adv-unsure { color: #555; } + .advisor-dispatch { font-size: 10px; display: block; margin-top: 2px; + color: #1a73e8; font-style: italic; } """ HUD_JS = ( @@ -285,10 +292,15 @@ def _note_from_outcome(outcome: Optional[Mapping[str, Any]]) -> str: newer = data.get("newer_failing_commits", []) or [] suspected = data.get("suspected_commit") or "?" baseline = data.get("older_successful_commit") or "?" - return ( + note = ( f"Pattern: newer fail {len(newer)}; suspect {suspected[:7]}" f" vs baseline {baseline[:7]}" ) + # Forward-compatible: advisor_verdict may not exist in older states + adv = data.get("advisor_verdict") + if adv: + note += f" [AI: {adv.get('verdict', '?')} @{adv.get('confidence', 0):.0%}]" + return note if outcome_type == "RestartCommits": commits = data.get("commit_shas", []) or [] if commits: @@ -316,12 +328,18 @@ def render_html_from_state( advisor_dispatches: Sequence[Mapping[str, Any]] = ( state.get("advisor_dispatches", []) or [] ) - # Build lookup: signal_key -> advisor dispatch info + # Build lookups for advisor dispatches + # signal_key -> dispatch info (for outcome badges) advisor_by_signal: Dict[str, Mapping[str, Any]] = {} + # (signal_key, commit_sha) -> dispatch info (for in-cell rendering) + advisor_by_cell: Dict[Tuple[str, str], Mapping[str, Any]] = {} for ad in advisor_dispatches: sk = ad.get("signal_key", "") + sha = ad.get("commit_sha", "") if sk: advisor_by_signal[sk] = ad + if sk and sha: + advisor_by_cell[(sk, sha)] = ad raw_outcomes = ( state.get("outcomes") if isinstance(state.get("outcomes"), dict) else None @@ -446,12 +464,37 @@ def render_html_from_state( for col in columns: cells_map = col.get("cells", {}) or {} events = cells_map.get(sha, []) or [] + # Forward-compatible: advisor_results may not exist in older states + advisor_results = col.get("advisor_results", {}) or {} workflow = str(col.get("workflow", "")) key = str(col.get("key", "")) sig_key = f"{workflow}:{key}" if key else workflow highlights_map = highlight_lookup.get(sig_key, {}) cell_classes = " ".join(sorted(highlights_map.get(sha, []))) - if not events: + + # Render advisor verdict badge for this cell (if available) + advisor_badge = "" + adv = advisor_results.get(sha) + if adv: + adv_verdict = adv.get("verdict", "") + adv_conf = adv.get("confidence", 0) + adv_class = f"adv-{adv_verdict}" if adv_verdict else "" + advisor_badge = ( + f'' + f"AI:{adv_verdict}" + ) + + # Render advisor dispatch indicator (from advisor_dispatches) + dispatch = advisor_by_cell.get((sig_key, sha)) + if dispatch and not advisor_badge: + # Show dispatch indicator only if no verdict badge already shown + advisor_badge = ( + 'AI:pending' + ) + + if not events and not advisor_badge: html_parts.append(f'') continue @@ -484,7 +527,8 @@ def render_html_from_state( f'{icon}' ) html_parts.append( - f'{"".join(cell_parts)}' + f'' + f'{"".join(cell_parts)}{advisor_badge}' ) html_parts.append("") html_parts.append("") diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/run_state_logger.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/run_state_logger.py index 51ab31d9c3..e99bc7f503 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/run_state_logger.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/run_state_logger.py @@ -78,6 +78,11 @@ def _build_state_json( data["wf_run_id"] = outcome.wf_run_id if outcome.job_id is not None: data["job_id"] = outcome.job_id + if outcome.advisor_verdict is not None: + data["advisor_verdict"] = { + "verdict": outcome.advisor_verdict.verdict.value, + "confidence": outcome.advisor_verdict.confidence, + } serialized = { "type": "AutorevertPattern", "data": data, @@ -105,8 +110,9 @@ def _build_state_json( }, } - # Per-commit events for this signal + # Per-commit events and advisor results for this signal cells: Dict[str, List[Dict]] = {} + advisor_results_map: Dict[str, Dict] = {} for c in sig.commits: evs = [] for e in c.events: @@ -124,6 +130,13 @@ def _build_state_json( evs.append(ev) if evs: cells[c.head_sha] = evs + # Capture advisor result if present (forward-compatible: absent in old states) + if c.advisor_result is not None: + advisor_results_map[c.head_sha] = { + "verdict": c.advisor_result.verdict.value, + "confidence": c.advisor_result.confidence, + "signal_key": c.advisor_result.signal_key, + } col = { "workflow": sig.workflow_name, @@ -135,6 +148,9 @@ def _build_state_json( col["job_base_name"] = sig.job_base_name if ineligible is not None: col["ineligible"] = ineligible + # Optional: per-commit advisor results (forward-compatible) + if advisor_results_map: + col["advisor_results"] = advisor_results_map cols.append(col) sig_key = f"{sig.workflow_name}:{sig.key}" diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal.py index 040185c595..12a02a72c7 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal.py @@ -50,9 +50,9 @@ class AutorevertPattern: - suspected_commit: the oldest commit that first started to fail. - older_successful_commit: the most recent successful commit before failures started (direct parent of the suspected commit for this signal). - - job_base_name: optional job base name for the signal - wf_run_id: optional workflow run ID from a failing event on suspected commit - job_id: optional job ID from a failing event on suspected commit + - advisor_verdict: optional AI advisor verdict that accelerated this decision """ workflow_name: str @@ -61,6 +61,7 @@ class AutorevertPattern: older_successful_commit: str wf_run_id: Optional[int] = None job_id: Optional[int] = None + advisor_verdict: Optional["AIAdvisorResult"] = None @dataclass @@ -448,7 +449,9 @@ def partition_by_autorevert_pattern(self) -> Optional[PartitionedCommits]: ADVISOR_CONFIDENCE_THRESHOLD = 0.9 def _build_autorevert_pattern( - self, partition: "PartitionedCommits" + self, + partition: "PartitionedCommits", + advisor_result: Optional[AIAdvisorResult] = None, ) -> AutorevertPattern: """Build an AutorevertPattern from a validated partition.""" suspected = partition.failed[-1] @@ -464,6 +467,7 @@ def _build_autorevert_pattern( older_successful_commit=partition.successful[0].head_sha, wf_run_id=failure_event.wf_run_id if failure_event else None, job_id=failure_event.job_id if failure_event else None, + advisor_verdict=advisor_result, ) def _check_advisor_verdict( @@ -490,7 +494,7 @@ def _check_advisor_verdict( return None if result.verdict == AdvisorVerdict.REVERT: - return self._build_autorevert_pattern(partition) + return self._build_autorevert_pattern(partition, advisor_result=result) if result.verdict == AdvisorVerdict.NOT_RELATED: return Ineligible( diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_actions.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_actions.py index 41b0296272..9502ff5da3 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_actions.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/signal_actions.py @@ -50,6 +50,7 @@ class SignalMetadata: test_module: Optional[str] = None wf_run_id: Optional[int] = None job_id: Optional[int] = None + advisor_summary: Optional[str] = None # short AI advisor verdict summary def _derive_job_filter(job_base_name: Optional[str]) -> Optional[str]: @@ -293,9 +294,16 @@ def group_actions( # Extract fields for job/HUD links from AutorevertPattern wf_run_id = None job_id = None + advisor_summary = None if isinstance(outcome, AutorevertPattern): wf_run_id = outcome.wf_run_id job_id = outcome.job_id + if outcome.advisor_verdict is not None: + av = outcome.advisor_verdict + advisor_summary = ( + f"AI advisor: {av.verdict.value} " + f"(confidence={av.confidence:.2f})" + ) meta = SignalMetadata( workflow_name=sig.workflow_name, @@ -304,6 +312,7 @@ def group_actions( test_module=sig.test_module, wf_run_id=wf_run_id, job_id=job_id, + advisor_summary=advisor_summary, ) if isinstance(outcome, AutorevertPattern): sha = outcome.suspected_commit @@ -1064,6 +1073,15 @@ def _comment_issue_pr_revert( all_signals = ", ".join(all_signals_urls) breaking_notification_msg += f"- {workflow_name}: {all_signals}\n" + # Add AI advisor info if any signal was advisor-accelerated + advisor_summaries = [s.advisor_summary for s in sources if s.advisor_summary] + if advisor_summaries: + breaking_notification_msg += ( + "\n**Note:** This revert was accelerated by the AI advisor: " + + "; ".join(advisor_summaries) + + "\n" + ) + try: if should_do_revert_on_pr: for attempt in RetryWithBackoff(): diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_actions.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_actions.py index 31d8816122..e9173a41a0 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_actions.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_actions.py @@ -1117,15 +1117,25 @@ def test_attaches_verdict_to_matching_commit(self): from pytorch_auto_revert.signal_extraction_types import Sha t0 = datetime(2025, 8, 19, 12, 0, 0) - c1 = SignalCommit("sha_aaa", t0, [ - SignalEvent("j", SignalStatus.FAILURE, t0, wf_run_id=1, job_id=10), - ]) - c2 = SignalCommit("sha_bbb", t0, [ - SignalEvent("j", SignalStatus.SUCCESS, t0, wf_run_id=2, job_id=20), - ]) + c1 = SignalCommit( + "sha_aaa", + t0, + [ + SignalEvent("j", SignalStatus.FAILURE, t0, wf_run_id=1, job_id=10), + ], + ) + c2 = SignalCommit( + "sha_bbb", + t0, + [ + SignalEvent("j", SignalStatus.SUCCESS, t0, wf_run_id=2, job_id=20), + ], + ) signal = Signal( - key="test_key", workflow_name="trunk", - commits=[c1, c2], source=SignalSource.TEST, + key="test_key", + workflow_name="trunk", + commits=[c1, c2], + source=SignalSource.TEST, ) extractor = SignalExtractor(workflows=["trunk"], lookback_hours=16) @@ -1150,17 +1160,28 @@ def test_attaches_verdict_to_matching_commit(self): def test_no_verdicts_returns_signals_unchanged(self): from pytorch_auto_revert.signal import ( - Signal, SignalCommit, SignalEvent, SignalSource, SignalStatus, + Signal, + SignalCommit, + SignalEvent, + SignalSource, + SignalStatus, ) from pytorch_auto_revert.signal_extraction import SignalExtractor from pytorch_auto_revert.signal_extraction_types import Sha t0 = datetime(2025, 8, 19, 12, 0, 0) - c1 = SignalCommit("sha_aaa", t0, [ - SignalEvent("j", SignalStatus.FAILURE, t0, wf_run_id=1), - ]) + c1 = SignalCommit( + "sha_aaa", + t0, + [ + SignalEvent("j", SignalStatus.FAILURE, t0, wf_run_id=1), + ], + ) signal = Signal( - key="k", workflow_name="wf", commits=[c1], source=SignalSource.TEST, + key="k", + workflow_name="wf", + commits=[c1], + source=SignalSource.TEST, ) extractor = SignalExtractor(workflows=["wf"], lookback_hours=16) @@ -1175,19 +1196,25 @@ def test_no_verdicts_returns_signals_unchanged(self): def test_passes_signal_keys_to_datasource(self): from pytorch_auto_revert.signal import ( - Signal, SignalCommit, SignalEvent, SignalSource, SignalStatus, + Signal, + SignalCommit, + SignalEvent, + SignalSource, + SignalStatus, ) from pytorch_auto_revert.signal_extraction import SignalExtractor from pytorch_auto_revert.signal_extraction_types import Sha t0 = datetime(2025, 8, 19, 12, 0, 0) s1 = Signal( - key="key_a", workflow_name="wf", + key="key_a", + workflow_name="wf", commits=[SignalCommit("sha1", t0, [])], source=SignalSource.TEST, ) s2 = Signal( - key="key_b", workflow_name="wf", + key="key_b", + workflow_name="wf", commits=[SignalCommit("sha1", t0, [])], source=SignalSource.TEST, ) @@ -1205,18 +1232,29 @@ def test_passes_signal_keys_to_datasource(self): def test_invalid_verdict_string_defaults_to_unsure(self): from pytorch_auto_revert.signal import ( - AdvisorVerdict, Signal, SignalCommit, SignalEvent, - SignalSource, SignalStatus, + AdvisorVerdict, + Signal, + SignalCommit, + SignalEvent, + SignalSource, + SignalStatus, ) from pytorch_auto_revert.signal_extraction import SignalExtractor from pytorch_auto_revert.signal_extraction_types import Sha t0 = datetime(2025, 8, 19, 12, 0, 0) - c1 = SignalCommit("sha_aaa", t0, [ - SignalEvent("j", SignalStatus.FAILURE, t0, wf_run_id=1), - ]) + c1 = SignalCommit( + "sha_aaa", + t0, + [ + SignalEvent("j", SignalStatus.FAILURE, t0, wf_run_id=1), + ], + ) signal = Signal( - key="k", workflow_name="wf", commits=[c1], source=SignalSource.TEST, + key="k", + workflow_name="wf", + commits=[c1], + source=SignalSource.TEST, ) extractor = SignalExtractor(workflows=["wf"], lookback_hours=16) diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_extraction.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_extraction.py index 8389e24dbb..90593ab73a 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_extraction.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_extraction.py @@ -66,6 +66,9 @@ def fetch_tests_for_job_ids( ids = {int(j) for j in job_ids} return [r for r in self._tests if int(r.job_id) in ids] + def fetch_advisor_verdicts(self, **kwargs): + return {} + def J( *, From 529dd6da9b99cff8fd1f4b6e473b234912dc3bff Mon Sep 17 00:00:00 2001 From: Ivan Zaitsev Date: Tue, 31 Mar 2026 17:30:10 -0700 Subject: [PATCH 3/3] test: add signal extraction tests for advisor verdict attachment FakeDatasource now accepts advisor_verdicts dict, enabling end-to-end tests that verify verdicts are correctly attached to SignalCommit objects during extraction. 3 new tests: - verdict attached to matching (commit, signal_key) - verdict for wrong signal_key is not attached - verdict attached to test-track signals --- .../pytorch_auto_revert/hud_renderer.py | 2 +- .../tests/test_signal_extraction.py | 146 +++++++++++++++++- 2 files changed, 143 insertions(+), 5 deletions(-) diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/hud_renderer.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/hud_renderer.py index dcd5ebb1ce..67ac034b8a 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/hud_renderer.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/hud_renderer.py @@ -528,7 +528,7 @@ def render_html_from_state( ) html_parts.append( f'' - f'{"".join(cell_parts)}{advisor_badge}' + f"{''.join(cell_parts)}{advisor_badge}" ) html_parts.append("") html_parts.append("") diff --git a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_extraction.py b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_extraction.py index 90593ab73a..5291805ab2 100644 --- a/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_extraction.py +++ b/aws/lambda/pytorch-auto-revert/pytorch_auto_revert/tests/test_signal_extraction.py @@ -25,9 +25,15 @@ def ts(base: datetime, minutes: int) -> datetime: class FakeDatasource(SignalExtractionDatasource): """Test double for the datasource returning provided rows.""" - def __init__(self, jobs: List[JobRow], tests: List[TestRow]): + def __init__( + self, + jobs: List[JobRow], + tests: List[TestRow], + advisor_verdicts: Optional[dict] = None, + ): self._jobs = jobs self._tests = tests + self._advisor_verdicts = advisor_verdicts or {} def fetch_commits_in_time_range( self, @@ -67,7 +73,7 @@ def fetch_tests_for_job_ids( return [r for r in self._tests if int(r.job_id) in ids] def fetch_advisor_verdicts(self, **kwargs): - return {} + return self._advisor_verdicts def J( @@ -124,9 +130,14 @@ class TestSignalExtraction(unittest.TestCase): def setUp(self) -> None: self.t0 = datetime(2025, 8, 20, 12, 0, 0) - def _extract(self, jobs: List[JobRow], tests: List[TestRow]): + def _extract( + self, + jobs: List[JobRow], + tests: List[TestRow], + advisor_verdicts: Optional[dict] = None, + ): se = SignalExtractor(workflows=["trunk"], lookback_hours=24) - se._datasource = FakeDatasource(jobs, tests) + se._datasource = FakeDatasource(jobs, tests, advisor_verdicts) return se.extract() def _find_job_signal(self, signals, wf: str, base: JobBaseName): @@ -924,6 +935,133 @@ def test_no_job_test_signal_when_only_non_test_failures(self): # [test] job signal should NOT be emitted (no test failures) self.assertIsNone(self._find_job_signal(signals, "trunk", f"{base} [test]")) + def test_advisor_verdict_attached_to_correct_commit(self): + """Advisor verdict from datasource is attached to the matching (commit, signal).""" + jobs = [ + J( + sha="C2", + run=200, + job=1, + attempt=1, + started_at=ts(self.t0, 10), + conclusion="failure", + ), + J( + sha="C1", + run=100, + job=2, + attempt=1, + started_at=ts(self.t0, 5), + conclusion="success", + ), + ] + base = jobs[0].base_name + # Advisor says revert for C2 on the job signal + advisor_verdicts = { + ("C2", base): ("revert", 0.95, self.t0), + } + signals = self._extract(jobs, tests=[], advisor_verdicts=advisor_verdicts) + sig = self._find_job_signal(signals, "trunk", base) + self.assertIsNotNone(sig) + + # C2 should have advisor_result + c2 = next(c for c in sig.commits if c.head_sha == "C2") + self.assertIsNotNone(c2.advisor_result) + self.assertEqual(c2.advisor_result.verdict.value, "revert") + self.assertAlmostEqual(c2.advisor_result.confidence, 0.95) + self.assertEqual(c2.advisor_result.signal_key, base) + + # C1 should NOT have advisor_result + c1 = next(c for c in sig.commits if c.head_sha == "C1") + self.assertIsNone(c1.advisor_result) + + def test_advisor_verdict_not_attached_to_wrong_signal(self): + """Advisor verdict for signal A is not attached to signal B.""" + jobs = [ + J( + sha="C2", + run=200, + job=1, + attempt=1, + started_at=ts(self.t0, 10), + conclusion="failure", + ), + J( + sha="C1", + run=100, + job=2, + attempt=1, + started_at=ts(self.t0, 5), + conclusion="success", + ), + ] + base = jobs[0].base_name + # Advisor verdict is for a completely different signal key + advisor_verdicts = { + ("C2", "some_other_signal"): ("not_related", 0.99, self.t0), + } + signals = self._extract(jobs, tests=[], advisor_verdicts=advisor_verdicts) + sig = self._find_job_signal(signals, "trunk", base) + self.assertIsNotNone(sig) + + # C2 should NOT have advisor_result (wrong signal key) + c2 = next(c for c in sig.commits if c.head_sha == "C2") + self.assertIsNone(c2.advisor_result) + + def test_advisor_verdict_on_test_signal(self): + """Advisor verdict is attached to test-track signals.""" + jobs = [ + J( + sha="C2", + run=200, + job=1, + attempt=1, + started_at=ts(self.t0, 10), + conclusion="failure", + rule="pytest failure", + ), + J( + sha="C1", + run=100, + job=2, + attempt=1, + started_at=ts(self.t0, 5), + conclusion="success", + ), + ] + tests = [ + T( + job=1, + run=200, + attempt=1, + file="test_foo.py", + name="test_bar", + failure_runs=1, + success_runs=0, + ), + T( + job=2, + run=100, + attempt=1, + file="test_foo.py", + name="test_bar", + failure_runs=0, + success_runs=1, + ), + ] + test_key = "test_foo.py::test_bar" + advisor_verdicts = { + ("C2", test_key): ("garbage", 0.92, self.t0), + } + signals = self._extract(jobs, tests, advisor_verdicts=advisor_verdicts) + sig = self._find_test_signal(signals, "trunk", test_key) + self.assertIsNotNone(sig) + + c2 = next(c for c in sig.commits if c.head_sha == "C2") + self.assertIsNotNone(c2.advisor_result) + self.assertEqual(c2.advisor_result.verdict.value, "garbage") + self.assertEqual(c2.advisor_result.signal_key, test_key) + if __name__ == "__main__": unittest.main()