include link to Beaker job with Slack noties

allenai · Dec 21, 2024 · c634066 · c634066
1 parent 3505660
commit c634066
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/olmo_core/train/callbacks/slack_notifier.py b/src/olmo_core/train/callbacks/slack_notifier.py
@@ -11,6 +11,7 @@
 from .callback import Callback
 
 SLACK_WEBHOOK_URL_ENV_VAR = "SLACK_WEBHOOK_URL"
+BEAKER_JOB_ID_ENV_VAR = "BEAKER_JOB_ID"
 
 
 class SlackNotificationSetting(StrEnum):
@@ -110,8 +111,13 @@ def _post_message(self, msg: str):
             f"- epoch: {self.trainer.epoch}\n"
             f"- tokens: {self.trainer.global_train_tokens_seen:,d}"
         )
+
         if self.name is not None:
             msg = f"Run `{self.name}` {msg}\n{progress}"
         else:
             msg = f"Run {msg}\n{progress}"
+
+        if BEAKER_JOB_ID_ENV_VAR in os.environ:
+            msg = f"{msg}\n*Beaker job:* https://beaker.org/job/{os.environ[BEAKER_JOB_ID_ENV_VAR]}"
+
         requests.post(webhook_url, json={"text": msg})