experiment: serialize crossover launch locks

1 files changed, 18 insertions, 5 deletions

diff --git a/experiments/validate_crossover_artifact.py b/experiments/validate_crossover_artifact.py
index 8b82ec4..360c654 100644
--- a/experiments/validate_crossover_artifact.py
+++ b/experiments/validate_crossover_artifact.py
@@ -22,13 +22,23 @@ def validate_audit(report, count):
     assert report["test_policy"] == "none"
 
 
+def validate_launch(report, stage, count, expected_commit):
+    assert expected_commit
+    assert report["stage"] == stage
+    assert report["source"]["git_commit"] == expected_commit
+    assert report["num_jobs"] == count
+    assert report["allowed_physical_gpus"] == [5, 7]
+
+
 def validate(report, kind, expected_commit=None):
     if kind == "resnet_p1_launch":
-        assert expected_commit
-        assert report["stage"] == "p1"
-        assert report["source"]["git_commit"] == expected_commit
-        assert report["num_jobs"] == 19
-        assert report["allowed_physical_gpus"] == [5, 7]
+        validate_launch(report, "p1", 19, expected_commit)
+    elif kind == "transformer_p1_launch":
+        validate_launch(report, "p1", 27, expected_commit)
+    elif kind == "resnet_r2_launch":
+        validate_launch(report, "r2", 27, expected_commit)
+    elif kind == "transformer_t2_launch":
+        validate_launch(report, "t2", 27, expected_commit)
     elif kind == "resnet_p1_selector":
         validate_selector(report, "resnet_crossover_p1", 19)
     elif kind == "transformer_p1_selector":
@@ -54,6 +64,9 @@ def main():
         required=True,
         choices=(
             "resnet_p1_launch",
+            "transformer_p1_launch",
+            "resnet_r2_launch",
+            "transformer_t2_launch",
             "resnet_p1_selector",
             "transformer_p1_selector",
             "resnet_r2_audit",