fix compute invariants

batch/batch/driver/main.py

@@ -1193,28 +1193,25 @@ async def check(tx):
 
         attempt_by_job_group_resources = fold(attempt_resources, lambda k: (k[0], k[1]))
         attempt_by_job_resources = fold(attempt_resources, lambda k: (k[0], k[2]))
+        job_by_job_resources = fold(agg_job_resources, lambda k: (k[0], k[2]))
         job_by_job_group_resources = fold(agg_job_resources, lambda k: (k[0], k[1]))
-        job_group_by_billing_project_resources = fold(agg_job_group_resources, lambda k: k[1])
+        job_group_by_job_group_resources = fold(agg_job_group_resources, lambda k: (k[0], k[1]))
+        job_group_by_billing_project_resources = fold(agg_job_group_resources, lambda k: k[2])
 
-        agg_job_group_resources_2 = {
-            (batch_id, job_group_id): resources
-            for (batch_id, job_group_id, _), resources in agg_job_group_resources.items()
-        }
-
-        assert attempt_by_job_group_resources == agg_job_group_resources_2, (
-            dictdiffer.diff(attempt_by_job_group_resources, agg_job_group_resources_2),
+        assert attempt_by_job_group_resources == job_group_by_job_group_resources, (
+            dictdiffer.diff(attempt_by_job_group_resources, job_group_by_job_group_resources),
             attempt_by_job_group_resources,
-            agg_job_group_resources_2,
+            job_group_by_job_group_resources,
         )
-        assert attempt_by_job_resources == agg_job_resources, (
+        assert attempt_by_job_resources == job_by_job_resources, (
             dictdiffer.diff(attempt_by_job_resources, agg_job_resources),
             attempt_by_job_resources,
             agg_job_resources,
         )
-        assert job_by_job_group_resources == agg_job_group_resources_2, (
-            dictdiffer.diff(job_by_job_group_resources, agg_job_group_resources_2),
+        assert job_by_job_group_resources == job_group_by_job_group_resources, (
+            dictdiffer.diff(job_by_job_group_resources, job_group_by_job_group_resources),
             job_by_job_group_resources,
-            agg_job_group_resources_2,
+            job_group_by_job_group_resources,
         )
         assert job_group_by_billing_project_resources == agg_billing_project_resources, (
             dictdiffer.diff(job_group_by_billing_project_resources, agg_billing_project_resources),