Improve training telemetry guards

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 "

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,14 @@ def sync_done(self) -> None:
 
     def record_loss(self, loss: float) -> None:
         """Record the loss value for this step."""
-        self._loss = loss
+        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))
+            or not math.isfinite(float(loss))
+        ):
+            raise ValueError("loss must be a finite number")
+        self._loss = float(loss)

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}"
         )

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                                 #
@@ -105,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)
 
@@ -198,3 +223,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

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)