Improve workload allocation safety

macfleet/pool/scheduler.py

@@ -13,6 +13,17 @@
 from macfleet.pool.registry import ClusterRegistry
 
 
+def _positive_finite_capacity(value: object) -> float:
+    """Return a usable scheduling capacity, or zero for invalid telemetry."""
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        return 0.0
+    try:
+        capacity = float(value)
+    except (OverflowError, ValueError):
+        return 0.0
+    return capacity if math.isfinite(capacity) and capacity > 0 else 0.0
+
+
 @dataclass
 class WorkloadAssignment:
     """Workload assignment for a single node."""
@@ -83,25 +94,24 @@ def compute_weights(self) -> dict[str, float]:
 
         raw_weights: dict[str, float] = {}
         for node in nodes:
-            if (
-                self.config.use_throughput
-                and math.isfinite(node.throughput_samples_sec)
-                and node.throughput_samples_sec > 0
-            ):
-                base = node.throughput_samples_sec
-            else:
-                base = float(node.hardware.gpu_cores)
+            throughput = _positive_finite_capacity(node.throughput_samples_sec)
+            gpu_capacity = _positive_finite_capacity(node.hardware.gpu_cores)
+            base = throughput if self.config.use_throughput and throughput else gpu_capacity
 
             factor = node.hardware.thermal_pressure.workload_multiplier
             raw_weights[node.node_id] = base * factor
 
-        total = sum(raw_weights.values())
-        if total <= 0:
+        max_weight = max(raw_weights.values())
+        if max_weight <= 0:
             # Equal split
             n = len(nodes)
             return {nid: 1.0 / n for nid in raw_weights}
 
-        return {nid: w / total for nid, w in raw_weights.items()}
+        # Normalize after scaling by the largest capacity. Summing raw finite
+        # values can still overflow (for example, two 1e308 throughput samples).
+        scaled_weights = {nid: weight / max_weight for nid, weight in raw_weights.items()}
+        total = sum(scaled_weights.values())
+        return {nid: weight / total for nid, weight in scaled_weights.items()}
 
     def assign(self, global_batch_size: int) -> list[WorkloadAssignment]:
         """Produce workload assignments for all alive nodes.
@@ -122,24 +132,32 @@ def assign(self, global_batch_size: int) -> list[WorkloadAssignment]:
         ranks = self.registry.get_ranks()
 
         assignments: list[WorkloadAssignment] = []
-        remaining = global_batch_size
-
         sorted_nodes = sorted(weights.keys(), key=lambda nid: ranks.get(nid, 999))
 
+        ideal_batches = {
+            node_id: global_batch_size * weights[node_id]
+            for node_id in sorted_nodes
+        }
+        batches = {
+            node_id: int(ideal_batches[node_id])
+            for node_id in sorted_nodes
+        }
+        remaining = global_batch_size - sum(batches.values())
+        remainder_order = sorted(
+            sorted_nodes,
+            key=lambda node_id: (
+                -(ideal_batches[node_id] - batches[node_id]),
+                ranks.get(node_id, 999),
+                node_id,
+            ),
+        )
+        for node_id in remainder_order[:remaining]:
+            batches[node_id] += 1
+
         for i, node_id in enumerate(sorted_nodes):
             weight = weights[node_id]
             rank = ranks.get(node_id, i)
-
-            if i == len(sorted_nodes) - 1:
-                # Last node gets remainder (clamped: when global_batch_size <
-                # node count, earlier max(1, ...) bumps can exhaust the budget).
-                batch = max(0, remaining)
-            elif remaining <= 0:
-                # Budget already exhausted — don't hand out batches we can't back.
-                batch = 0
-            else:
-                batch = max(1, int(global_batch_size * weight))
-                remaining -= batch
+            batch = batches[node_id]
 
             is_viable = batch >= self.config.min_batch_per_node
 

macfleet/training/sampler.py

@@ -195,21 +195,56 @@ def set_epoch(self, epoch: int) -> None:
 # --------------------------------------------------------------------------- #
 
 
+def _normalize_capacities(
+    values: list[int] | list[float],
+    *,
+    name: str,
+    integers_only: bool,
+) -> list[float]:
+    """Validate non-negative capacities and normalize without overflow."""
+    if not values:
+        return []
+
+    capacities: list[float] = []
+    for value in values:
+        valid_type = (
+            isinstance(value, int)
+            if integers_only
+            else isinstance(value, (int, float))
+        )
+        if isinstance(value, bool) or not valid_type:
+            kind = "non-negative integers" if integers_only else "finite non-negative numbers"
+            raise ValueError(f"{name} must contain {kind}")
+        try:
+            capacity = float(value)
+        except (OverflowError, ValueError) as exc:
+            raise ValueError(f"{name} must contain finite non-negative numbers") from exc
+        if not math.isfinite(capacity) or capacity < 0:
+            raise ValueError(f"{name} must contain finite non-negative numbers")
+        capacities.append(capacity)
+
+    maximum = max(capacities)
+    if maximum == 0:
+        return [1.0 / len(capacities)] * len(capacities)
+
+    scaled = [capacity / maximum for capacity in capacities]
+    total = sum(scaled)
+    return [capacity / total for capacity in scaled]
+
+
 def compute_weights_from_gpu_cores(gpu_cores: list[int]) -> list[float]:
     """Compute workload weights proportional to GPU core counts."""
-    if not gpu_cores:
-        return []
-    total = sum(gpu_cores)
-    if total == 0:
-        return [1.0 / len(gpu_cores)] * len(gpu_cores)
-    return [c / total for c in gpu_cores]
+    return _normalize_capacities(
+        gpu_cores,
+        name="gpu_cores",
+        integers_only=True,
+    )
 
 
 def compute_weights_from_throughput(throughputs: list[float]) -> list[float]:
     """Compute workload weights proportional to measured throughput."""
-    if not throughputs:
-        return []
-    total = sum(throughputs)
-    if total == 0:
-        return [1.0 / len(throughputs)] * len(throughputs)
-    return [t / total for t in throughputs]
+    return _normalize_capacities(
+        throughputs,
+        name="throughputs",
+        integers_only=False,
+    )

tests/test_pool/test_scheduler.py

@@ -106,6 +106,44 @@ def test_non_finite_throughput_falls_back_to_gpu_cores(self):
 
         assert weights == pytest.approx({"air": 1 / 3, "pro": 2 / 3})
 
+    @pytest.mark.parametrize("invalid_cores", [-1, float("nan"), float("inf"), True])
+    def test_invalid_gpu_capacity_cannot_create_invalid_weights(self, invalid_cores):
+        reg = ClusterRegistry("invalid")
+        invalid = _make_node("invalid", gpu_cores=10)
+        invalid.hardware.gpu_cores = invalid_cores
+        reg.register(invalid)
+        reg.register(_make_node("valid", gpu_cores=20))
+
+        weights = Scheduler(reg, SchedulerConfig(use_throughput=False)).compute_weights()
+
+        assert weights == {"invalid": 0.0, "valid": 1.0}
+
+    def test_all_invalid_gpu_capacities_fall_back_to_equal_weights(self):
+        reg = ClusterRegistry("a")
+        a = _make_node("a")
+        b = _make_node("b")
+        a.hardware.gpu_cores = -1
+        b.hardware.gpu_cores = float("nan")
+        reg.register(a)
+        reg.register(b)
+
+        weights = Scheduler(reg, SchedulerConfig(use_throughput=False)).compute_weights()
+
+        assert weights == {"a": 0.5, "b": 0.5}
+
+    def test_extreme_finite_throughput_normalizes_without_overflow(self):
+        reg = ClusterRegistry("a")
+        a = _make_node("a")
+        b = _make_node("b")
+        a.throughput_samples_sec = 1e308
+        b.throughput_samples_sec = 1e308
+        reg.register(a)
+        reg.register(b)
+
+        weights = Scheduler(reg).compute_weights()
+
+        assert weights == {"a": 0.5, "b": 0.5}
+
 
 class TestSchedulerAssignment:
     @pytest.mark.parametrize("batch_size", [0, -1, True, 1.5])
@@ -129,6 +167,38 @@ def test_batch_split(self):
         by_id = {a.node_id: a for a in assignments}
         assert by_id["b"].batch_size > by_id["a"].batch_size
 
+    def test_remainder_favors_largest_fractional_allocation(self):
+        reg = ClusterRegistry("a")
+        reg.register(_make_node("a", gpu_cores=34))
+        reg.register(_make_node("b", gpu_cores=33))
+        reg.register(_make_node("c", gpu_cores=33))
+
+        assignments = Scheduler(
+            reg,
+            SchedulerConfig(use_throughput=False),
+        ).assign(global_batch_size=10)
+
+        by_id = {assignment.node_id: assignment.batch_size for assignment in assignments}
+        assert by_id == {"a": 4, "b": 3, "c": 3}
+
+    def test_batch_allocation_is_independent_of_registration_order(self):
+        allocations = []
+        for registration_order in (("a", "b", "c"), ("c", "b", "a")):
+            reg = ClusterRegistry("a")
+            core_counts = {"a": 34, "b": 33, "c": 33}
+            for node_id in registration_order:
+                reg.register(_make_node(node_id, gpu_cores=core_counts[node_id]))
+
+            assignments = Scheduler(
+                reg,
+                SchedulerConfig(use_throughput=False),
+            ).assign(global_batch_size=10)
+            allocations.append(
+                {assignment.node_id: assignment.batch_size for assignment in assignments}
+            )
+
+        assert allocations == [{"a": 4, "b": 3, "c": 3}] * 2
+
     def test_minimum_batch_guard(self):
         reg = ClusterRegistry("strong")
         reg.register(_make_node("strong", gpu_cores=100))

tests/test_training/test_sampler.py

@@ -191,3 +191,28 @@ def test_zero_gpu_cores(self):
     def test_zero_throughput(self):
         weights = compute_weights_from_throughput([0.0, 0.0])
         assert weights == [0.5, 0.5]
+
+    @pytest.mark.parametrize(
+        "gpu_cores",
+        [[-1, 2], [1.5, 2], [True, 2], [10**10000, 2]],
+    )
+    def test_invalid_gpu_capacities_are_rejected(self, gpu_cores):
+        with pytest.raises(ValueError, match="gpu_cores"):
+            compute_weights_from_gpu_cores(gpu_cores)
+
+    @pytest.mark.parametrize(
+        "throughputs",
+        [
+            [-1.0, 2.0],
+            [float("nan"), 2.0],
+            [float("inf"), 2.0],
+            [True, 2.0],
+            ["1", 2.0],
+        ],
+    )
+    def test_invalid_throughput_capacities_are_rejected(self, throughputs):
+        with pytest.raises(ValueError, match="throughputs"):
+            compute_weights_from_throughput(throughputs)
+
+    def test_extreme_finite_throughputs_normalize_without_overflow(self):
+        assert compute_weights_from_throughput([1e308, 1e308]) == [0.5, 0.5]