Fix run resuming in W&B (#1028)

* Fix run resuming in W&B * Trigger CI * Revert "Trigger CI" This reverts commit fb7632a. * Trigger CI * Remove resume argument * Revert "Trigger CI" This reverts commit 805736a.
mozilla · Feb 5, 2025 · e8dfd75 · e8dfd75
1 parent 05e6d9a
commit e8dfd75
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 12 deletions.
diff --git a/pipeline/eval/eval.py b/pipeline/eval/eval.py
@@ -350,7 +350,7 @@ def main(args_list: Optional[list[str]] = None) -> None:
             return
 
         logger.info(f"Publishing metrics to Weight & Biases ({run_client.extra_kwargs})")
-        run_client.open(resume=True)
+        run_client.open()
         run_client.handle_metrics(metrics=[metric])
         run_client.close()
 
@@ -362,7 +362,7 @@ def main(args_list: Optional[list[str]] = None) -> None:
             suffix=run_client.suffix,
         )
         logger.info("Adding metric row to the 'group_logs' run")
-        group_logs_client.open(resume=True)
+        group_logs_client.open()
 
         # Restore existing metrics data
         data = list_existing_group_logs_metrics(group_logs_client.wandb)

diff --git a/tracking/translations_parser/publishers.py b/tracking/translations_parser/publishers.py
@@ -1,6 +1,5 @@
 import csv
 import logging
-import os
 import sys
 from abc import ABC
 from collections import defaultdict
@@ -135,28 +134,25 @@ def close(self) -> None:
 
         self.wandb.finish()
 
-    def open(self, parser=None, resume: bool = False) -> None:
+    def open(self, parser=None) -> None:
         self.parser = parser
         config = getattr(parser, "config", {}).copy()
         config.update(self.extra_kwargs.pop("config", {}))
         # Publish datasets stats directly in the dashboard
         datasets = config.pop("datasets", None)
 
-        # Avoid overriding an existing run on a first training, this should not happen
-        if resume is False and int(os.environ.get("RUN_ID", 0)) > 0:
-            logger.warning(
-                "Training has been resumed but resume option has been set to False, skipping publication."
-            )
-            return
-
         try:
             self.wandb = wandb.init(
                 project=self.project,
                 group=self.group,
                 name=self.run,
                 id=self.run,
                 config=config,
-                resume=resume,
+                # Since we use unique run names based on group ID (e.g. finetune-student_MjcJG),
+                # we can use "allow" mode for resuming a stopped Taskcluster run in case of preemption.
+                # It will continue logging to the same run if it exists.
+                # Offline publication should handle run deletion separately (use --override-runs).
+                resume="allow",
                 **self.extra_kwargs,
             )
             if self.wandb.resumed: