From 48a2e3e346ecae5a5e2d486cc295c3d083693cd0 Mon Sep 17 00:00:00 2001 From: maocheng23 Date: Mon, 24 Aug 2026 15:55:02 -0700 Subject: [PATCH] fix: create the metrics tracker only on trainer rank zero Every trainer rank previously constructed its own W&B/MLflow run or wrote the same TensorBoard directory, duplicating one logical metric stream world-size times. Non-zero ranks now receive no logger; the trainer controller already treats a missing logger as log-nothing, and producer roles keep the console logger. Co-Authored-By: Claude Fable 5 --- specforge/training/assembly.py | 10 +++++++++- tests/test_config/test_unified_feature_reachability.py | 9 +++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/specforge/training/assembly.py b/specforge/training/assembly.py index 5d773ccb9..8e4b89f92 100644 --- a/specforge/training/assembly.py +++ b/specforge/training/assembly.py @@ -292,7 +292,15 @@ def _logger(metrics, step): def _configured_logger(cfg: Config): - """Create an external tracker only for a trainer-bearing run.""" + """Create an external tracker only on the trainer's global rank zero.""" + if cfg.training.role != "producer": + import torch.distributed as dist + + # A distributed run has one logical metric stream. Letting every rank + # create W&B/MLflow runs or write the same TensorBoard directory both + # duplicates metrics and makes large jobs increasingly fragile. + if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0: + return None if cfg.tracking.report_to == "none" or cfg.training.role == "producer": return _logger diff --git a/tests/test_config/test_unified_feature_reachability.py b/tests/test_config/test_unified_feature_reachability.py index 680069324..d4d109ec6 100644 --- a/tests/test_config/test_unified_feature_reachability.py +++ b/tests/test_config/test_unified_feature_reachability.py @@ -205,6 +205,15 @@ def test_tracking_config_reaches_the_existing_tracker_adapter(self): self.assertEqual(output_dir, "/tmp/output") self.assertIs(create.call_args.kwargs["console_logger"], _logger) + def test_only_global_rank_zero_creates_a_training_logger(self): + cfg = Config.model_validate(OFFLINE_EAGLE3) + with ( + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.get_rank", return_value=3), + ): + self.assertIsNone(_configured_logger(cfg)) + def test_tracking_backend_is_strictly_typed(self): with self.assertRaises(ValidationError): Config.model_validate(