added systemMetricsMonitor callback (#1260)

* added systemMetricsMonitor callback * added unit testing for building system_metrics_monitor * temporarily removing flag check in unit test for system_metrics_monitor build * removed end of line formatting issue * fixed end of line error in __init__.py * fixed extra line in init --------- Co-authored-by: Daniel King <[email protected]> Co-authored-by: Mihir Patel <[email protected]>
mosaicml · Jun 19, 2024 · 4b1fecb · 4b1fecb
1 parent c23be4a
commit 4b1fecb
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/llmfoundry/callbacks/__init__.py b/llmfoundry/callbacks/__init__.py
@@ -11,6 +11,7 @@
     OptimizerMonitor,
     RuntimeEstimator,
     SpeedMonitor,
+    SystemMetricsMonitor,
 )
 
 from llmfoundry.callbacks.async_eval_callback import AsyncEval
@@ -35,6 +36,7 @@
 from llmfoundry.callbacks.scheduled_gc_callback import ScheduledGarbageCollector
 from llmfoundry.registry import callbacks, callbacks_with_config
 
+callbacks.register('system_metrics_monitor', func=SystemMetricsMonitor)
 callbacks.register('lr_monitor', func=LRMonitor)
 callbacks.register('memory_monitor', func=MemoryMonitor)
 callbacks.register('memory_snapshot', func=MemorySnapshot)

diff --git a/tests/callbacks/test_system_metrics_monitor.py b/tests/callbacks/test_system_metrics_monitor.py
@@ -0,0 +1,15 @@
+# Copyright 2024 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+
+from composer.callbacks import SystemMetricsMonitor
+
+from llmfoundry.utils.builders import build_callback
+
+
+def test_system_metrics_monitor_callback_builds():
+    callback = build_callback(
+        'system_metrics_monitor',
+        kwargs={},
+        train_config={'train_loader': {}},
+    )
+    assert isinstance(callback, SystemMetricsMonitor)