From f9557a854f5215cb85031f856c4eb814e9b52ee8 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Sat, 1 Aug 2026 22:01:29 -0400 Subject: [PATCH 1/3] Allocate scheduler remainders fairly --- macfleet/pool/scheduler.py | 34 +++++++++++++++++++------------ tests/test_pool/test_scheduler.py | 32 +++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 13 deletions(-) diff --git a/macfleet/pool/scheduler.py b/macfleet/pool/scheduler.py index b27a7b9..621595b 100644 --- a/macfleet/pool/scheduler.py +++ b/macfleet/pool/scheduler.py @@ -122,24 +122,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 diff --git a/tests/test_pool/test_scheduler.py b/tests/test_pool/test_scheduler.py index 4c94aa6..5cd0806 100644 --- a/tests/test_pool/test_scheduler.py +++ b/tests/test_pool/test_scheduler.py @@ -129,6 +129,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)) From 9b353a6f436db996f12cbd351543a738100a26c3 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Sat, 1 Aug 2026 22:02:40 -0400 Subject: [PATCH 2/3] Harden scheduler capacity weights --- macfleet/pool/scheduler.py | 32 +++++++++++++++++--------- tests/test_pool/test_scheduler.py | 38 +++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 11 deletions(-) diff --git a/macfleet/pool/scheduler.py b/macfleet/pool/scheduler.py index 621595b..bdc0046 100644 --- a/macfleet/pool/scheduler.py +++ b/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. diff --git a/tests/test_pool/test_scheduler.py b/tests/test_pool/test_scheduler.py index 5cd0806..bace55b 100644 --- a/tests/test_pool/test_scheduler.py +++ b/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]) From 8a32f06492d684f2bdbc37cb005f798ef8543103 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Sat, 1 Aug 2026 22:03:47 -0400 Subject: [PATCH 3/3] Validate sampler capacity inputs --- macfleet/training/sampler.py | 59 +++++++++++++++++++++++------ tests/test_training/test_sampler.py | 25 ++++++++++++ 2 files changed, 72 insertions(+), 12 deletions(-) diff --git a/macfleet/training/sampler.py b/macfleet/training/sampler.py index ebc8298..28b3ddd 100644 --- a/macfleet/training/sampler.py +++ b/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, + ) diff --git a/tests/test_training/test_sampler.py b/tests/test_training/test_sampler.py index 3143a61..7736116 100644 --- a/tests/test_training/test_sampler.py +++ b/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]