Catch and ignore CommError during W&B cancel check

allenai · Apr 23, 2024 · 40b6853 · 40b6853
1 parent 7be71cd
commit 40b6853
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/olmo/train.py b/olmo/train.py
@@ -914,6 +914,7 @@ def check_if_cancelled(self) -> Tuple[bool, int]:
                 # Finally, check if someone canceled the run from W&B by adding the 'cancel' / 'canceled' tag..
                 # We won't see it in the run object. So we have to use the import/export API to check.
                 from requests.exceptions import RequestException
+                from wandb.errors import CommError
 
                 try:
                     api = wandb.Api(api_key=api_key)
@@ -924,8 +925,8 @@ def check_if_cancelled(self) -> Tuple[bool, int]:
                             cancel_reason = "Weights & Biases tag"
                             extra_steps = self.cfg.extra_steps_after_cancel
                             break
-                except RequestException:
-                    pass
+                except (RequestException, CommError):
+                    log.info("Failed to check if W&B run is cancelled, continuing run.")
 
         run_canceled = synchronize_flag(should_cancel, self.device)
         if run_canceled: