diff --git a/src/olmo_core/train/callbacks/slack_notifier.py b/src/olmo_core/train/callbacks/slack_notifier.py index 57c42d56..a4fb5feb 100644 --- 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})