From e1cab1a3c628b0f5b6ecea4543ee849c0088ebf8 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Wed, 29 Jul 2026 15:24:16 -0400 Subject: [PATCH 1/4] Validate throughput loss metrics --- macfleet/monitoring/throughput.py | 17 +++++++++++++++- tests/test_monitoring/test_throughput.py | 25 ++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/macfleet/monitoring/throughput.py b/macfleet/monitoring/throughput.py index fac9196..c5b45e5 100644 --- a/macfleet/monitoring/throughput.py +++ b/macfleet/monitoring/throughput.py @@ -35,6 +35,7 @@ def __post_init__(self) -> None: ("compute_time_sec", self.compute_time_sec), ("sync_time_sec", self.sync_time_sec), ("step_time_sec", self.step_time_sec), + ("timestamp", self.timestamp), ): if ( isinstance(value, bool) @@ -43,6 +44,14 @@ def __post_init__(self) -> None: or value < 0 ): raise ValueError(f"{name} must be a non-negative finite number") + if self.loss is not None: + if ( + isinstance(self.loss, bool) + or not isinstance(self.loss, (int, float)) + or not math.isfinite(float(self.loss)) + ): + raise ValueError("loss must be a finite number when recorded") + self.loss = float(self.loss) @property def throughput(self) -> float: @@ -256,4 +265,10 @@ def sync_done(self) -> None: def record_loss(self, loss: float) -> None: """Record the loss value for this step.""" - self._loss = loss + if ( + isinstance(loss, bool) + or not isinstance(loss, (int, float)) + or not math.isfinite(float(loss)) + ): + raise ValueError("loss must be a finite number") + self._loss = float(loss) diff --git a/tests/test_monitoring/test_throughput.py b/tests/test_monitoring/test_throughput.py index 5cc3625..a55e1cf 100644 --- a/tests/test_monitoring/test_throughput.py +++ b/tests/test_monitoring/test_throughput.py @@ -42,6 +42,9 @@ def test_zero_step_time(self): ("compute_time_sec", -0.1), ("sync_time_sec", float("nan")), ("step_time_sec", float("inf")), + ("timestamp", -0.1), + ("loss", float("nan")), + ("loss", True), ], ) def test_rejects_invalid_metrics(self, field, value): @@ -57,6 +60,17 @@ def test_rejects_invalid_metrics(self, field, value): with pytest.raises(ValueError, match=field): StepMetrics(**kwargs) + def test_normalizes_integer_loss(self): + m = StepMetrics( + step=0, + samples=32, + compute_time_sec=0.08, + sync_time_sec=0.02, + step_time_sec=0.1, + loss=1, + ) + assert m.loss == 1.0 + # ------------------------------------------------------------------ # # ThroughputTracker manual recording # @@ -198,3 +212,14 @@ def test_failed_step_is_not_recorded(self): assert tracker.total_steps == 0 assert tracker.total_samples == 0 assert tracker.throughput_history == [] + + @pytest.mark.parametrize("loss", [float("nan"), float("inf"), True, "1.0"]) + def test_rejects_invalid_loss_values(self, loss): + tracker = ThroughputTracker() + + with pytest.raises(ValueError, match="loss"): + with tracker.step(32) as s: + s.record_loss(loss) + + assert tracker.total_steps == 0 + assert tracker.total_samples == 0 From cf692ae8fa337c9d9911e7b286b4de63ba3e5d87 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Wed, 29 Jul 2026 15:25:01 -0400 Subject: [PATCH 2/4] Fix dataset guard remediation advice --- macfleet/training/guards.py | 16 +++++++++++++--- tests/test_training/test_guards.py | 12 ++++++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/macfleet/training/guards.py b/macfleet/training/guards.py index 5bcbabc..6598bb8 100644 --- a/macfleet/training/guards.py +++ b/macfleet/training/guards.py @@ -132,12 +132,22 @@ def check_dataset_sufficient( required = batch_size * min_batches if dataset_len < required: shortfall = required - dataset_len + max_batch_size = dataset_len // min_batches + if max_batch_size >= 1: + remediation = ( + f"Either: (a) collect more data, or " + f"(b) reduce batch_size to {max_batch_size} or smaller." + ) + else: + remediation = ( + f"Either: (a) collect at least {min_batches} samples, or " + f"(b) lower min_batches; no positive batch_size can produce " + f"{min_batches} full batch(es) from {dataset_len} sample(s)." + ) raise DatasetSizeError( f"Dataset has {dataset_len} samples but needs >= {required} to " f"run at least {min_batches} batch(es) of size {batch_size} " f"across {world_size} rank(s). " f"Shortfall: {shortfall} samples. " - f"Either: (a) collect more data, " - f"(b) reduce batch_size to {dataset_len // min_batches} or " - f"smaller, or (c) reduce world_size." + f"{remediation}" ) diff --git a/tests/test_training/test_guards.py b/tests/test_training/test_guards.py index 9543358..bbed7ee 100644 --- a/tests/test_training/test_guards.py +++ b/tests/test_training/test_guards.py @@ -122,6 +122,7 @@ def test_error_suggests_reduced_batch_size(self): with pytest.raises(DatasetSizeError) as exc: check_dataset_sufficient(50, 128, 1) assert "reduce batch_size" in str(exc.value).lower() + assert "reduce batch_size to 50" in str(exc.value) def test_error_suggests_more_data(self): with pytest.raises(DatasetSizeError) as exc: @@ -137,6 +138,17 @@ def test_min_batches_enforced(self): ) assert "10 batch" in str(exc.value) + def test_min_batches_error_never_suggests_zero_batch_size(self): + with pytest.raises(DatasetSizeError) as exc: + check_dataset_sufficient( + dataset_len=5, batch_size=10, world_size=1, min_batches=10, + ) + + msg = str(exc.value) + assert "batch_size to 0" not in msg + assert "lower min_batches" in msg + assert "no positive batch_size" in msg + def test_world_size_must_be_positive(self): with pytest.raises(DatasetSizeError, match="world_size must be a positive"): check_dataset_sufficient(1000, 32, 0) From 49da166d836dd7003cec69058406dfd2258001a1 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Wed, 29 Jul 2026 15:25:36 -0400 Subject: [PATCH 3/4] Guard throughput loss recording lifecycle --- macfleet/monitoring/throughput.py | 4 ++++ tests/test_monitoring/test_throughput.py | 11 +++++++++++ 2 files changed, 15 insertions(+) diff --git a/macfleet/monitoring/throughput.py b/macfleet/monitoring/throughput.py index c5b45e5..d5e1b2a 100644 --- a/macfleet/monitoring/throughput.py +++ b/macfleet/monitoring/throughput.py @@ -265,6 +265,10 @@ def sync_done(self) -> None: def record_loss(self, loss: float) -> None: """Record the loss value for this step.""" + if self._start == 0: + raise RuntimeError("record_loss() must be called inside a step context") + if self._loss is not None: + raise RuntimeError("record_loss() may only be called once per step") if ( isinstance(loss, bool) or not isinstance(loss, (int, float)) diff --git a/tests/test_monitoring/test_throughput.py b/tests/test_monitoring/test_throughput.py index a55e1cf..24c71ef 100644 --- a/tests/test_monitoring/test_throughput.py +++ b/tests/test_monitoring/test_throughput.py @@ -119,6 +119,17 @@ def test_phase_markers_require_ordered_single_calls(self): with pytest.raises(RuntimeError, match="only be called once"): context.sync_done() + def test_loss_must_be_recorded_inside_context_once(self): + context = ThroughputTracker().step(1) + + with pytest.raises(RuntimeError, match="inside a step context"): + context.record_loss(1.0) + + with context: + context.record_loss(1.0) + with pytest.raises(RuntimeError, match="only be called once"): + context.record_loss(0.5) + def test_rolling_average(self): tracker = ThroughputTracker(window_size=3) From 3a42601c549e8f8f199f0b1fe78e78365fc9d7af Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Wed, 29 Jul 2026 15:26:10 -0400 Subject: [PATCH 4/4] Keep optional YAML import mypy-clean --- macfleet/cli/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/macfleet/cli/main.py b/macfleet/cli/main.py index cd4c0da..52bac6d 100644 --- a/macfleet/cli/main.py +++ b/macfleet/cli/main.py @@ -698,7 +698,7 @@ def _load_train_config(config_path: str) -> dict[str, object]: loaded = json.loads(path.read_text()) elif suffix in {".yaml", ".yml"}: try: - import yaml + import yaml # type: ignore[import-untyped] except ImportError as e: raise click.ClickException( "Reading YAML configs requires PyYAML. Install with "