[CRCR] Fix multi-job state key collision by adding job_name to Redis key (#8228)

## Summary

- PR #8170 changed the Redis state key from `check_run_id` (unique per
job) to `run_id:run_attempt` (unique per workflow run), causing state
machine collisions when multiple jobs in the same workflow run sent
callbacks concurrently
- This fix preserves PR #8170's re-run correlation intent while adding
per-job uniqueness via `job_name` (`github.job`, already present in
callback payloads)
- The state key is now
`crcr:state:{delivery_id}:{repo}:{run_id}:{run_attempt}:{job_name}`,
with `job_name` omitted for DISPATCHED records (backward-compatible)

## Test plan

- [ ] All 78 unit tests pass locally (verified)
- [ ] Trigger 3 successful runs on `TorchedHat/pytorch-redhat-ci` L2
edge case workflow (multi-job `ec05-a`/`ec05-b` jobs no longer collide)
- [ ] Verify single-job workflows still work (no regression from
optional `job_name`)

Author: subinz1

aws/lambda/cross_repo_ci_relay/callback/callback_handler.py

@@ -60,10 +60,11 @@ def _verify_access(
     return allowlist, repo_level
 
 
-def _parse_callback_body(body: dict) -> tuple[str, str, int, int, str]:
-    """Return (delivery_id, status, run_id, run_attempt, workflow_name) from ``body``.
+def _parse_callback_body(body: dict) -> tuple[str, str, int, int, str, str | None]:
+    """Return (delivery_id, status, run_id, run_attempt, workflow_name, job_name).
 
-    run_id and run_attempt uniquely identify a workflow run execution.
+    run_id and run_attempt identify a workflow run execution.
+    job_name (``github.job``) disambiguates multiple jobs within the same run.
     run_attempt defaults to 1 when not present in the callback body.
 
     Raises HTTPException(400) on any missing or mis-typed field.
@@ -72,17 +73,16 @@ def _parse_callback_body(body: dict) -> tuple[str, str, int, int, str]:
         delivery_id = body["delivery_id"]
         workflow_dict = body["workflow"]
         status = workflow_dict["status"]
-        run_id = int(workflow_dict["run_id"])  # Required for HUD grouping
-        run_attempt = int(
-            workflow_dict.get("run_attempt", 1)
-        )  # Default to 1 if not present
-        workflow_name = workflow_dict["name"]  # Required for HUD grouping
+        run_id = int(workflow_dict["run_id"])
+        run_attempt = int(workflow_dict.get("run_attempt", 1))
+        workflow_name = workflow_dict["name"]
+        job_name = workflow_dict.get("job_name")
     except (KeyError, TypeError) as exc:
         logger.warning(f"missing required field in callback body: {exc}")
         raise HTTPException(
             400, f"callback body missing required field: {exc}"
         ) from exc
-    return delivery_id, status, run_id, run_attempt, workflow_name
+    return delivery_id, status, run_id, run_attempt, workflow_name, job_name
 
 
 def _update_state_and_compute_metrics(
@@ -96,6 +96,7 @@ def _update_state_and_compute_metrics(
     dispatch_record: CallbackStateRecord,
     workflow_record: CallbackStateRecord | None,
     payload: dict | None = None,
+    job_name: str | None = None,
 ) -> dict:
     """Persist the new workflow state to Redis and return CI timing metrics.
 
@@ -129,6 +130,7 @@ def _update_state_and_compute_metrics(
             current_timestamp,
             workflow_name,
             payload=payload,
+            job_name=job_name,
         )
     except RedisError:
         raise HTTPException(
@@ -144,7 +146,7 @@ def _update_state_and_compute_metrics(
         raise
 
     updated_workflow_record = redis_helper.get_callback_state(
-        config, delivery_id, verified_repo, run_id, run_attempt
+        config, delivery_id, verified_repo, run_id, run_attempt, job_name=job_name
     )
     if updated_workflow_record is None:
         return ci_metrics
@@ -188,7 +190,9 @@ def handle(config: RelayConfig, body: dict, verified_repo: str) -> dict:
         return {"ok": True, "status": "ignored"}
     _, repo_level = result
 
-    delivery_id, status, run_id, run_attempt, workflow_name = _parse_callback_body(body)
+    delivery_id, status, run_id, run_attempt, workflow_name, job_name = (
+        _parse_callback_body(body)
+    )
 
     dispatch_record = redis_helper.get_callback_state(
         config, delivery_id, verified_repo, DISPATCH_RUN_ID, DISPATCH_RUN_ATTEMPT
@@ -202,7 +206,7 @@ def handle(config: RelayConfig, body: dict, verified_repo: str) -> dict:
         raise HTTPException(400, "callback rejected: no matching dispatch record")
 
     workflow_record = redis_helper.get_callback_state(
-        config, delivery_id, verified_repo, run_id, run_attempt
+        config, delivery_id, verified_repo, run_id, run_attempt, job_name=job_name
     )
 
     payload = None
@@ -230,16 +234,16 @@ def handle(config: RelayConfig, body: dict, verified_repo: str) -> dict:
         dispatch_record,
         workflow_record,
         payload=payload,
+        job_name=job_name,
     )
 
-    # Track in-progress jobs for zombie detection.
     if status == "in_progress":
         redis_helper.add_in_progress_tracker(
-            config, delivery_id, verified_repo, run_id, run_attempt
+            config, delivery_id, verified_repo, run_id, run_attempt, job_name=job_name
         )
     elif status == "completed":
         redis_helper.remove_in_progress_tracker(
-            config, delivery_id, verified_repo, run_id, run_attempt
+            config, delivery_id, verified_repo, run_id, run_attempt, job_name=job_name
         )
 
     trusted = {

aws/lambda/cross_repo_ci_relay/callback/cleanup_handler.py

@@ -68,29 +68,30 @@ def _cleanup_one(
     repo = zombie["downstream_repo"]
     run_id = zombie["run_id"]
     run_attempt = zombie["run_attempt"]
+    job_name = zombie.get("job_name")
     hud_ok = True
 
-    # 1. Build and forward timeout payload to HUD
     try:
         trusted, untrusted = _build_timeout_payload(zombie, completed_at)
         forward_to_hud(config, trusted, untrusted)
         logger.info(
-            "zombie HUD forward succeeded repo=%s run_id=%s run_attempt=%s",
+            "zombie HUD forward succeeded repo=%s run_id=%s run_attempt=%s job_name=%s",
             repo,
             run_id,
             run_attempt,
+            job_name,
         )
     except Exception:
         logger.exception(
-            "zombie HUD forward failed repo=%s run_id=%s run_attempt=%s",
+            "zombie HUD forward failed repo=%s run_id=%s run_attempt=%s job_name=%s",
             repo,
             run_id,
             run_attempt,
+            job_name,
         )
         hud_ok = False
 
     if hud_ok:
-        # 2. Mark state as COMPLETED in Redis (best-effort)
         try:
             redis_helper.set_callback_state(
                 config,
@@ -100,22 +101,20 @@ def _cleanup_one(
                 run_attempt,
                 CallbackState.COMPLETED,
                 time.time(),
+                job_name=job_name,
             )
         except (AssertionError, RedisError):
-            # Race: another process already resolved this record, or Redis
-            # was unavailable.
             logger.warning(
                 "zombie state transition failed (may already be resolved) "
-                "delivery_id=%s repo=%s run_id=%s run_attempt=%s",
+                "delivery_id=%s repo=%s run_id=%s run_attempt=%s job_name=%s",
                 delivery_id,
                 repo,
                 run_id,
                 run_attempt,
+                job_name,
             )
-        # Erase the records from Redis only when HUD was successfully
-        # updated, keeping the two systems in sync.
         redis_helper.remove_in_progress_tracker(
-            config, delivery_id, repo, run_id, run_attempt
+            config, delivery_id, repo, run_id, run_attempt, job_name=job_name
         )
 
     return {"ok": hud_ok}

aws/lambda/cross_repo_ci_relay/tests/test_callback_handler.py

@@ -19,23 +19,32 @@ def _cfg():
     return cfg
 
 
-def _body(status="completed", workflow_name="default", run_id=99999, run_attempt=1):
+def _body(
+    status="completed",
+    workflow_name="default",
+    run_id=99999,
+    run_attempt=1,
+    job_name=None,
+):
+    wf = {
+        "status": status,
+        "conclusion": "success" if status == "completed" else None,
+        "name": "CI",
+        "url": "http://ci.example.com/run/1",
+        "workflow_name": workflow_name,
+        "run_id": run_id,
+        "run_attempt": run_attempt,
+    }
+    if job_name is not None:
+        wf["job_name"] = job_name
     return {
         "event_type": "pull_request",
         "delivery_id": "del-123",
         "payload": {
             "pull_request": {"number": 42, "head": {"sha": "abc123"}},
             "repository": {"full_name": "pytorch/pytorch"},
         },
-        "workflow": {
-            "status": status,
-            "conclusion": "success" if status == "completed" else None,
-            "name": "CI",
-            "url": "http://ci.example.com/run/1",
-            "workflow_name": workflow_name,
-            "run_id": run_id,
-            "run_attempt": run_attempt,
-        },
+        "workflow": wf,
     }
 
 
@@ -54,7 +63,13 @@ def setUp(self):
 
         # Setup default: dispatch exists, workflow state is None (in_progress not yet reported)
         def default_get_state(
-            cfg, delivery_id, repo, run_id_arg, run_attempt_arg, client=None
+            cfg,
+            delivery_id,
+            repo,
+            run_id_arg,
+            run_attempt_arg,
+            client=None,
+            job_name=None,
         ):
             if run_id_arg == DISPATCH_RUN_ID:
                 return CallbackStateRecord(
@@ -234,6 +249,30 @@ def test_redis_error_fetching_workflow_record_proceeds(self):
         _, trusted_arg, _ = self.mock_hud.call_args[0]
         self.assertIsNone(trusted_arg["ci_metrics"]["execution_time"])
 
+    def test_job_name_passed_to_redis_state_calls(self):
+        """job_name from callback body is forwarded to all Redis state calls."""
+        dispatch_record = CallbackStateRecord(CallbackState.DISPATCHED, 1000.0, {})
+        in_progress_record = CallbackStateRecord(CallbackState.IN_PROGRESS, 1030.0, {})
+        self.mock_redis.get_callback_state.side_effect = [
+            dispatch_record,
+            None,
+            in_progress_record,
+        ]
+
+        handle(
+            _cfg(),
+            _body(status="in_progress", job_name="build"),
+            verified_repo="org/repo",
+        )
+
+        # set_callback_state should have been called with job_name="build"
+        call_kwargs = self.mock_redis.set_callback_state.call_args
+        self.assertEqual(call_kwargs.kwargs.get("job_name"), "build")
+
+        # add_in_progress_tracker should have been called with job_name="build"
+        tracker_kwargs = self.mock_redis.add_in_progress_tracker.call_args
+        self.assertEqual(tracker_kwargs.kwargs.get("job_name"), "build")
+
 
 if __name__ == "__main__":
     unittest.main()

aws/lambda/cross_repo_ci_relay/tests/test_cleanup_handler.py

@@ -54,6 +54,7 @@ def _zombie_entry(
     repo="org/repo",
     run_id=99999,
     run_attempt=1,
+    job_name="my-job",
     in_progress_ts=None,
     body_overrides=None,
 ):
@@ -79,6 +80,7 @@ def _zombie_entry(
         "downstream_repo": repo,
         "run_id": run_id,
         "run_attempt": run_attempt,
+        "job_name": job_name,
         "state_record": CallbackStateRecord(
             CallbackState.IN_PROGRESS, in_progress_ts, stored_payload
         ),
@@ -186,22 +188,21 @@ def test_cleans_single_zombie(self):
             untrusted["callback_payload"]["workflow"]["conclusion"], "timed_out"
         )
 
-        # Redis state was updated to COMPLETED
+        # Redis state was updated to COMPLETED with job_name
         self.mock_redis.set_callback_state.assert_called_once()
         call_args = self.mock_redis.set_callback_state.call_args[0]
         self.assertEqual(call_args[1], "del-123")
         self.assertEqual(call_args[2], "org/repo")
         self.assertEqual(call_args[3], 99999)
         self.assertEqual(call_args[4], 1)
         self.assertEqual(call_args[5], CallbackState.COMPLETED)
+        call_kwargs = self.mock_redis.set_callback_state.call_args[1]
+        self.assertEqual(call_kwargs.get("job_name"), "my-job")
 
-        # ZSET entry was removed
+        # ZSET entry was removed with job_name
         self.mock_redis.remove_in_progress_tracker.assert_called_once()
-        rm_args = self.mock_redis.remove_in_progress_tracker.call_args[0]
-        self.assertEqual(rm_args[1], "del-123")
-        self.assertEqual(rm_args[2], "org/repo")
-        self.assertEqual(rm_args[3], 99999)
-        self.assertEqual(rm_args[4], 1)
+        rm_kwargs = self.mock_redis.remove_in_progress_tracker.call_args[1]
+        self.assertEqual(rm_kwargs.get("job_name"), "my-job")
 
     def test_cleans_multiple_zombies(self):
         """Multiple zombies are all cleaned independently."""