Make preassigned WandB run IDs resumable (#3030)

miles/utils/test_wandb_utils.py

@@ -40,6 +40,20 @@ def test_primary_wandb_init_uses_extended_init_timeout(monkeypatch):
     assert args.wandb_run_id == "run-id"
 
 
+def test_primary_wandb_init_resumes_a_preassigned_run(monkeypatch):
+    """An explicitly assigned run ID lets a replacement primary resume the run."""
+    init_calls = []
+
+    monkeypatch.setattr(wandb_utils.wandb, "init", lambda **kwargs: init_calls.append(kwargs))
+    monkeypatch.setattr(wandb_utils.wandb, "define_metric", lambda *args, **kwargs: None)
+    monkeypatch.setattr(wandb_utils.wandb, "run", SimpleNamespace(id="run-id"), raising=False)
+
+    wandb_utils.init_wandb_primary(_args(wandb_run_id="assigned-run"))
+
+    assert init_calls[0]["id"] == "assigned-run"
+    assert init_calls[0]["resume"] == "allow"
+
+
 def test_secondary_wandb_init_uses_extended_init_timeout(monkeypatch):
     init_calls = []
 

miles/utils/tracking_utils/wandb_utils.py

@@ -64,6 +64,8 @@ def init_wandb_primary(args):
         "name": run_name,
         "config": _compute_config_for_logging(args),
     }
+    if args.wandb_run_id is not None:
+        init_kwargs |= {"id": args.wandb_run_id, "resume": "allow"}
 
     # Configure settings based on offline/online mode
     if offline:

tests/fast/utils/tracking_utils/test_wandb_utils.py

@@ -1,11 +1,42 @@
 import json
 from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
 
 from tests.fast.utils.env_report.conftest import make_args
 
+from miles.utils.tracking_utils import wandb_utils
 from miles.utils.tracking_utils.wandb_utils import _compute_config_for_logging
 
 
+class TestPrimaryWandbInit:
+    def test_primary_wandb_init_resumes_a_preassigned_run(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """A replacement primary resumes the explicitly assigned WandB run."""
+        init_calls: list[dict[str, object]] = []
+        args = make_args(
+            rank=0,
+            use_wandb=True,
+            wandb_dir=None,
+            wandb_group="group",
+            wandb_host=None,
+            wandb_key=None,
+            wandb_mode=None,
+            wandb_project="project",
+            wandb_random_suffix=False,
+            wandb_run_id="assigned-run",
+            wandb_team=None,
+        )
+        monkeypatch.setattr(wandb_utils.wandb, "init", lambda **kwargs: init_calls.append(kwargs))
+        monkeypatch.setattr(wandb_utils.wandb, "define_metric", lambda *args, **kwargs: None)
+        monkeypatch.setattr(wandb_utils.wandb, "run", SimpleNamespace(id="assigned-run"), raising=False)
+
+        wandb_utils.init_wandb_primary(args)
+
+        assert init_calls[0]["id"] == "assigned-run"
+        assert init_calls[0]["resume"] == "allow"
+
+
 class TestTheLoggedConfig:
     def test_carries_the_record_of_the_launch_that_started_this_run(self, tmp_path: Path) -> None:
         """The launcher's record is what an algorithm engineer looks for first, and wandb is where they look."""