Added live slice count

DannyLiCom · ycchenzheng · commit 5c5e9b187a0d · 2026-06-08T22:09:08.000Z
(cherry picked from commit dc08dde04de59e4c7f8a6d67d16107deb31fc642)
diff --git a/src/maxtext/common/metric_logger.py b/src/maxtext/common/metric_logger.py
@@ -35,6 +35,7 @@
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
+from maxtext.utils import elastic_utils
 from collections import defaultdict
 
 mldiag, _ = mldiagnostics_modules()
@@ -166,6 +167,7 @@ def _log_training_metrics(self, metrics, step):
       log_parts.extend(
           [
               f"completed step: {step}",
+              f"live slice count: {len(elastic_utils.live_slice_indices(self.config))}",
               f"seconds: {scalars['perf/step_time_seconds']:.3f}",
           ]
       )

Original file line number	Diff line number	Diff line change
`@@ -35,6 +35,7 @@`
`35`	`35`	`from maxtext.utils import max_logging`
`36`	`36`	`from maxtext.utils import max_utils`
`37`	`37`	`from maxtext.utils import maxtext_utils`
	`38`	`+from maxtext.utils import elastic_utils`
`38`	`39`	`from collections import defaultdict`
`39`	`40`
`40`	`41`	`mldiag, _ = mldiagnostics_modules()`
`@@ -166,6 +167,7 @@ def _log_training_metrics(self, metrics, step):`
`166`	`167`	`log_parts.extend(`
`167`	`168`	`[`
`168`	`169`	`f"completed step: {step}",`
	`170`	`+ f"live slice count: {len(elastic_utils.live_slice_indices(self.config))}",`
`169`	`171`	`f"seconds: {scalars['perf/step_time_seconds']:.3f}",`
`170`	`172`	`]`
`171`	`173`	`)`