Log exception on inactivity callback

llmfoundry/callbacks/__init__.py

@@ -19,6 +19,7 @@
 from llmfoundry.callbacks.eval_output_logging_callback import EvalOutputLogging
 from llmfoundry.callbacks.fdiff_callback import FDiffMetrics
 from llmfoundry.callbacks.hf_checkpointer import HuggingFaceCheckpointer
+from llmfoundry.callbacks.inactivity_callback import InactivityCallback
 from llmfoundry.callbacks.log_mbmoe_tok_per_expert_callback import \
     MegaBlocksMoE_TokPerExpert
 from llmfoundry.callbacks.monolithic_ckpt_callback import \
@@ -47,6 +48,7 @@
 callbacks.register('oom_observer', func=OOMObserver)
 callbacks.register('eval_output_logging', func=EvalOutputLogging)
 callbacks.register('mbmoe_tok_per_expert', func=MegaBlocksMoE_TokPerExpert)
+callbacks.register('inactivity', func=InactivityCallback)
 
 callbacks_with_config.register('async_eval', func=AsyncEval)
 callbacks_with_config.register('curriculum_learning', func=CurriculumLearning)

llmfoundry/callbacks/inactivity_callback.py

@@ -0,0 +1,56 @@
+# Copyright 2024 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+import signal
+import threading
+from typing import Optional
+
+from composer import Callback, Logger, State
+from composer.loggers import MosaicMLLogger
+
+from llmfoundry.utils.exceptions import RunTimeoutError
+
+log = logging.getLogger(__name__)
+
+
+def _timeout(timeout: int, mosaicml_logger: Optional[MosaicMLLogger] = None):
+    log.error(
+        f'Timeout after no Trainer events were triggered for {timeout} seconds.',
+    )
+    if mosaicml_logger is not None:
+        mosaicml_logger.log_exception(RunTimeoutError(timeout=timeout))
+    os.kill(os.getpid(), signal.SIGINT)
+
+
+class InactivityCallback(Callback):
+
+    def __init__(
+        self,
+        timeout: int = 1800,
+        mosaicml_logger: Optional[MosaicMLLogger] = None,
+    ):
+        self.timeout = timeout
+        self.mosaicml_logger = mosaicml_logger
+        self.timer: Optional[threading.Timer] = None
+
+    def _reset(self):
+        if self.timer is not None:
+            self.timer.cancel()
+        self.timer = None
+
+    def _timeout(self):
+        self._reset()
+        self.timer = threading.Timer(
+            self.timeout,
+            _timeout,
+            [self.timeout, self.mosaicml_logger],
+        )
+        self.timer.daemon = True
+        self.timer.start()
+
+    def fit_end(self, state: State, logger: Logger):
+        del state
+        del logger
+        self._timeout()

llmfoundry/utils/exceptions.py

@@ -246,3 +246,12 @@ def __init__(self, dataset_name: str, split: str) -> None:
         self.split = split
         message = f'Your dataset (name={dataset_name}, split={split}) is misconfigured. Please check your dataset format and make sure you can load your dataset locally.'
         super().__init__(message)
+
+
+class RunTimeoutError(RuntimeError):
+    """Error thrown when a run times out."""
+
+    def __init__(self, timeout: int) -> None:
+        self.timeout = timeout
+        message = f'Run timed out after {timeout} seconds.'
+        super().__init__(message)

scripts/train/train.py

@@ -365,6 +365,11 @@ def main(cfg: DictConfig) -> Trainer:
     callback_configs = train_cfg.callbacks or {}
 
     # Callbacks
+    if callback_configs is not None:
+        for name, callback_cfg in callback_configs.items():
+            if name == 'inactivity':
+                callback_cfg['mosaicml_logger'] = mosaicml_logger
+
     callbacks: List[Callback] = [
         build_callback(
             name=str(name),