Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / balancing_algorithm.py: 95%
219 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Replaceable capacity-aware assignment algorithms and their objectives."""
17from __future__ import annotations
19import math
20import sys
21from collections.abc import Callable, Iterator, Sequence
22from dataclasses import dataclass, field
23from typing import Literal, Protocol
25from hyper_parallel.distributed_data.schema import (
26 PackingConstraints,
27 SampleKey,
28 SampleMetadata,
29 WorkloadCost,
30)
33class BalancingAlgorithm(Protocol):
34 """Deterministic CPU assignment, independent of workload estimation.
36 Samples already contain estimated costs. Implementations must not mutate
37 metadata, read future samples, or issue collectives. Return every selected
38 key exactly once in non-empty bins, ordered by rank then local bin index.
39 The planner applies capacities and the benefit gate to the returned layout.
40 """
42 def assign(
43 self,
44 samples: Sequence[SampleMetadata],
45 *,
46 reference_bins: Sequence[Sequence[SampleKey]],
47 constraints: PackingConstraints,
48 data_parallel_size: int,
49 local_batch_size: int,
50 ) -> Sequence[Sequence[SampleKey]]:
51 """Return candidate bins, or the original bins if no candidate is feasible.
53 Args:
54 samples: Metadata with bound keys and estimated workload costs.
55 reference_bins: Original rank-major grouping of the selected keys.
56 constraints: Per-bin token capacity and physical stage budgets.
57 data_parallel_size: Number of target DP ranks.
58 local_batch_size: Number of non-empty bins required per rank.
59 """
61 def objective(self, rank_costs: Sequence[WorkloadCost]) -> float:
62 """Return a finite, non-negative score; a lower score is better.
64 Args:
65 rank_costs: Additive workloads for one synchronous round.
66 """
69@dataclass
70class _AssignmentBin:
71 data_rank: int
72 pack_index: int
73 keys: list[SampleKey] = field(default_factory=list)
74 tokens: int = 0
75 packing_costs: dict[str, float] = field(default_factory=dict)
78class LPTBalancingAlgorithm:
79 """Longest-processing-time placement with explicit packing constraints.
81 Assignment sorts by backbone cost and picks the least-loaded feasible rank.
82 ``makespan`` minimizes maximum rank cost. ``balance`` minimizes squared
83 coefficient of variation, equivalent to variance for a frozen sample set.
84 Other workload components remain available to user-defined algorithms.
85 """
87 def __init__(self, objective: Literal["makespan", "balance"] = "makespan") -> None:
88 """Select the scalar objective used by the planner's relative-gain gate.
90 Args:
91 objective: Maximum backbone load, or relative backbone-load variance.
92 """
93 if objective not in ("makespan", "balance"):
94 raise ValueError("objective must be 'makespan' or 'balance'.")
95 self.objective_name = objective
96 self.padding_cost: Callable[[int], WorkloadCost] | None = None
98 def assign(
99 self,
100 samples: Sequence[SampleMetadata],
101 *,
102 reference_bins: Sequence[Sequence[SampleKey]],
103 constraints: PackingConstraints,
104 data_parallel_size: int,
105 local_batch_size: int,
106 ) -> Sequence[Sequence[SampleKey]]:
107 """Propose cost-first bins without changing the selected sample set.
109 Args:
110 samples: Individually feasible samples with estimated workload costs.
111 reference_bins: Known-feasible original rank-major bins.
112 constraints: Token capacity and additive stage budgets.
113 data_parallel_size: Number of target ranks in this balancing domain.
114 local_batch_size: Number of non-empty bins required on each rank.
116 Returns:
117 Candidate rank-major bins, or the reference on a greedy dead end.
118 """
119 ordered = sorted(samples, key=lambda item: (-item.cost.llm, -item.pack_tokens, item.key))
120 bins = [
121 _AssignmentBin(data_rank, pack_index)
122 for data_rank in range(data_parallel_size)
123 for pack_index in range(local_batch_size)
124 ]
125 rank_loads = [local_batch_size * self._padding_llm_cost(constraints.seq_len)] * data_parallel_size
126 rank_tokens = [0] * data_parallel_size
127 for index, item in enumerate(ordered):
128 feasible = [
129 packing_bin for packing_bin in bins
130 if (index >= len(bins) or not packing_bin.keys)
131 and constraints.fits(packing_bin.tokens, packing_bin.packing_costs, item)
132 ]
133 if not feasible:
134 return reference_bins
135 selected = min(
136 feasible,
137 key=lambda packing_bin, sample=item: self._placement_score(
138 packing_bin, sample, constraints.seq_len, rank_loads, rank_tokens,
139 ),
140 )
141 rank_loads[selected.data_rank] += item.cost.llm + self._padding_delta(
142 selected.tokens, item.pack_tokens, constraints.seq_len,
143 )
144 selected.keys.append(item.key)
145 selected.tokens += item.pack_tokens
146 selected.packing_costs = constraints.add_costs(selected.packing_costs, item)
147 rank_tokens[selected.data_rank] += min(item.pack_tokens, constraints.seq_len)
148 return tuple(tuple(packing_bin.keys) for packing_bin in bins)
150 def _padding_llm_cost(self, tokens: int) -> float:
151 if self.padding_cost is None:
152 return 0.0
153 # The planner installs the callback after policy construction; static inference sees only None.
154 return self.padding_cost(max(0, tokens)).llm # pylint: disable=not-callable
156 def _padding_delta(self, tokens: int, added_tokens: int, capacity: int) -> float:
157 return self._padding_llm_cost(capacity - tokens - added_tokens) - self._padding_llm_cost(capacity - tokens)
159 def _placement_score(
160 self, packing_bin: _AssignmentBin,
161 item: SampleMetadata,
162 seq_len: int,
163 rank_loads: Sequence[float],
164 rank_tokens: Sequence[int],
165 ) -> tuple[float | int, ...]:
166 """Rank feasible placements by workload, remaining capacity and deterministic ties."""
167 remaining_capacity = seq_len - min(seq_len, packing_bin.tokens + item.pack_tokens)
168 score = rank_loads[packing_bin.data_rank]
169 if self.padding_cost is not None:
170 loads = list(rank_loads)
171 loads[packing_bin.data_rank] += item.cost.llm + self._padding_delta(
172 packing_bin.tokens, item.pack_tokens, seq_len,
173 )
174 score = self.objective([WorkloadCost(llm=load) for load in loads])
175 return (
176 score,
177 remaining_capacity,
178 rank_tokens[packing_bin.data_rank],
179 packing_bin.data_rank,
180 packing_bin.pack_index,
181 )
183 def objective(self, rank_costs: Sequence[WorkloadCost]) -> float:
184 """Score completed rank loads using the configured optimization target.
186 Args:
187 rank_costs: Additive workloads reconstructed by the planner.
188 """
189 maximum = max((cost.llm for cost in rank_costs), default=0.0)
190 if self.objective_name == "makespan" or not maximum:
191 return maximum
192 # Scaling before squaring avoids overflow with FLOPs-sized estimates.
193 loads = [cost.llm / maximum for cost in rank_costs]
194 mean = math.fsum(loads) / len(loads)
195 return math.fsum((load / mean - 1.0) ** 2 for load in loads) / len(loads)
198class _WindowLPT(LPTBalancingAlgorithm):
199 """Score placements after scheduling complete packs into synchronous rounds."""
201 def __init__(self, group_size: int) -> None:
202 """Bind the physical DP size without mutating the caller's policy."""
203 super().__init__()
204 self.group_size = group_size
206 def _placement_score(
207 self, packing_bin: _AssignmentBin, item: SampleMetadata, seq_len: int,
208 rank_loads: Sequence[float], rank_tokens: Sequence[int],
209 ) -> tuple[float | int, ...]:
210 loads = list(rank_loads)
211 loads[packing_bin.data_rank] += item.cost.llm + self._padding_delta(
212 packing_bin.tokens, item.pack_tokens, seq_len,
213 )
214 score = math.fsum(sorted(loads, reverse=True)[::self.group_size])
215 return (
216 score, seq_len - min(seq_len, packing_bin.tokens + item.pack_tokens),
217 rank_tokens[packing_bin.data_rank], packing_bin.data_rank, packing_bin.pack_index,
218 )
221class FeasibleLPTBalancingAlgorithm(LPTBalancingAlgorithm):
222 """Refine token-only greedy assignments within a fixed pair budget.
224 Suffix subset sums preserve feasible completion for two bins. Larger
225 domains prioritize round bottlenecks. A failed greedy pass starts from
226 the known-feasible reference; successful assignments can also improve.
227 Repairs require one bin per rank and no extra packing budgets.
228 Each pair's bitset table is limited to an estimated 256 MiB, not process RSS.
229 The planner still applies its window objective and minimum-gain gate.
230 """
232 def __init__(
233 self, objective: Literal["makespan", "balance"] = "makespan", *, max_repair_pairs: int = 64,
234 ) -> None:
235 """Bound repair attempts independently of the number of packs.
237 Args:
238 objective: Maximum backbone load, or relative backbone-load variance.
239 max_repair_pairs: Positive maximum number of pair checks per assignment.
240 A window shares one budget across all its rounds.
241 """
242 super().__init__(objective)
243 if not isinstance(max_repair_pairs, int) or isinstance(max_repair_pairs, bool) or max_repair_pairs < 1:
244 raise ValueError("max_repair_pairs must be a positive integer.")
245 self.max_repair_pairs = max_repair_pairs
247 def assign(
248 self,
249 samples: Sequence[SampleMetadata],
250 *,
251 reference_bins: Sequence[Sequence[SampleKey]],
252 constraints: PackingConstraints,
253 data_parallel_size: int,
254 local_batch_size: int,
255 objective_group_size: int | None = None,
256 ) -> Sequence[Sequence[SampleKey]]:
257 """Propose window-aware placements and bounded capacity-feasible refinements.
259 Args:
260 samples: Selected occurrences with estimated workload costs.
261 reference_bins: Known-feasible original grouping.
262 constraints: Token capacity and optional physical stage budgets.
263 data_parallel_size: Number of target ranks, including window rounds.
264 local_batch_size: Required non-empty bins per target rank.
265 objective_group_size: Physical DP ranks per window round. Defaults
266 to data_parallel_size for an ordinary single-round assignment.
267 """
268 group_size = data_parallel_size if objective_group_size is None else objective_group_size
269 if (
270 not isinstance(group_size, int) or isinstance(group_size, bool) or group_size < 1
271 or data_parallel_size % group_size
272 ):
273 raise ValueError("objective_group_size must be a positive divisor of data_parallel_size.")
274 placement = self
275 if (
276 self._uses_scheduled_makespan() and group_size < data_parallel_size and local_batch_size == 1
277 and type(self)._placement_score is LPTBalancingAlgorithm._placement_score
278 ):
279 placement = _WindowLPT(group_size)
280 placement.padding_cost = self.padding_cost
281 candidate = LPTBalancingAlgorithm.assign(
282 placement,
283 samples, reference_bins=reference_bins, constraints=constraints,
284 data_parallel_size=data_parallel_size, local_batch_size=local_batch_size,
285 )
286 if (
287 local_batch_size != 1 or constraints.packing_budgets or len(reference_bins) < 2
288 ):
289 return candidate
290 return self._refine_candidate(samples, candidate, reference_bins, constraints.seq_len, group_size)
292 def _refine_candidate(
293 self,
294 samples: Sequence[SampleMetadata],
295 candidate: Sequence[Sequence[SampleKey]],
296 reference_bins: Sequence[Sequence[SampleKey]],
297 capacity: int,
298 group_size: int,
299 ) -> Sequence[Sequence[SampleKey]]:
300 """Refine one-bin-per-slot assignments without token overflow or stage budgets."""
301 lookup = {item.key: item for item in samples}
302 if any(sum(lookup[key].pack_tokens for key in keys) > capacity for keys in candidate):
303 return candidate
304 if self.padding_cost is not None:
305 candidate = min(
306 (candidate, reference_bins),
307 key=lambda bins: self._repair_objective(
308 [self._bin_cost(lookup, keys, capacity) for keys in bins], group_size,
309 ),
310 )
311 if not self._uses_scheduled_makespan() and (
312 candidate is not reference_bins
313 or any(sum(lookup[key].pack_tokens for key in keys) != capacity for keys in candidate)
314 ):
315 return candidate
316 return self._repair_bin_pairs(lookup, candidate, capacity, group_size)
318 def _uses_scheduled_makespan(self) -> bool:
319 return self.objective_name == "makespan" and type(self).objective is LPTBalancingAlgorithm.objective
321 def _bin_cost(
322 self, lookup: dict[SampleKey, SampleMetadata], keys: Sequence[SampleKey], capacity: int,
323 ) -> WorkloadCost:
324 cost = sum((lookup[key].cost for key in keys), WorkloadCost())
325 if self.padding_cost is not None:
326 cost += self.padding_cost(max(0, capacity - sum(lookup[key].pack_tokens for key in keys)))
327 return cost
329 def _repair_bin_pairs(
330 self,
331 lookup: dict[SampleKey, SampleMetadata],
332 reference_bins: Sequence[Sequence[SampleKey]],
333 capacity: int,
334 group_size: int,
335 ) -> Sequence[Sequence[SampleKey]]:
336 """Try bounded pair repairs, accepting only improvements to the window objective."""
337 bins = [tuple(keys) for keys in reference_bins]
338 costs = [self._bin_cost(lookup, keys, capacity) for keys in bins]
339 score = self._repair_objective(costs, group_size)
340 checked: set[tuple[int, int, int, int]] = set()
341 versions = [0] * len(bins)
342 changed = False
343 scheduled_makespan = self._uses_scheduled_makespan()
344 while len(checked) < self.max_repair_pairs:
345 pairs = {pair for pair in self._repair_candidates(costs, group_size)
346 if (*pair, versions[pair[0]], versions[pair[1]]) not in checked}
347 if scheduled_makespan:
348 def _potential(pair: tuple[int, int]) -> float:
349 """Estimate candidate gain without treating PAD cost as conserved."""
350 # PAD cost changes with the partition. Source costs alone
351 # give an optimistic lower bound for candidate ordering.
352 total = (sum(lookup[key].cost.llm for index in pair for key in bins[index])
353 if self.padding_cost is not None else None)
354 return self._pair_potential(costs, pair, group_size, total)
356 pairs = sorted(pairs, key=lambda pair: (-_potential(pair), pair))
357 else:
358 pairs = sorted(pairs)
359 for left, right in pairs:
360 version = (left, right, versions[left], versions[right])
361 if len(checked) == self.max_repair_pairs:
362 break
363 # Failed and memory-limited attempts consume the same budget.
364 checked.add(version)
365 items = [lookup[key] for index in (left, right) for key in bins[index]]
366 items.sort(key=lambda item: (-item.cost.llm, -item.pack_tokens, item.key))
367 candidate = self._repair_bins(items, capacity, (bins[left], bins[right]), costs, (left, right),
368 group_size)
369 candidate_costs = costs.copy()
370 candidate_costs[left], candidate_costs[right] = (
371 self._bin_cost(lookup, keys, capacity) for keys in candidate
372 )
373 candidate_score = self._repair_objective(candidate_costs, group_size)
374 if candidate_score < score:
375 bins[left], bins[right] = candidate
376 costs, score = candidate_costs, candidate_score
377 changed = True
378 versions[left] += 1
379 versions[right] += 1
380 break
381 else:
382 break
383 return tuple(bins) if changed else reference_bins
385 def _repair_candidates(
386 self, costs: Sequence[WorkloadCost], group_size: int,
387 ) -> Iterator[tuple[int, int]]:
388 """Pair round leaders with heavy, light and nearby-cost packs."""
389 order = sorted(range(len(costs)), key=lambda index: (-costs[index].llm, index))
390 if self._uses_scheduled_makespan():
391 leaders = order[::group_size]
392 else:
393 leaders = [
394 max(range(start, start + group_size), key=lambda index: costs[index].llm)
395 for start in range(0, len(costs), group_size)
396 ]
397 positions = {index: position for position, index in enumerate(order)}
398 # Heavy, light and nearby packs offer different capacity-feasible swaps.
399 for left in leaders:
400 partners = order[:4] + order[-4:]
401 partners += [order[(positions[left] + offset + 1) % len(order)] for offset in range(4)]
402 for right in partners:
403 if left != right:
404 yield min(left, right), max(left, right)
406 @staticmethod
407 def _pair_potential(
408 costs: Sequence[WorkloadCost], pair: tuple[int, int], group_size: int, total: float | None = None,
409 ) -> float:
410 """Rank sparse candidates by a continuous two-pack relaxation."""
411 left, right = pair
412 if total is None:
413 total = costs[left].llm + costs[right].llm
414 others = sorted((cost.llm for index, cost in enumerate(costs) if index not in pair), reverse=True)
415 # Only ranks around a round maximum can change the scheduled objective.
416 # This bounds the number of splits by the round count, even at large DP.
417 breakpoints = {0, (group_size - 1) % group_size, (group_size - 2) % group_size}
418 boundaries = [cost for index, cost in enumerate(others) if index % group_size in breakpoints]
419 splits = {0.0, total, total / 2}
420 splits.update(cost for cost in boundaries if cost <= total)
421 splits.update(total - cost for cost in boundaries if cost <= total)
422 best = min(math.fsum(sorted((*others, split, total - split), reverse=True)[::group_size])
423 for split in splits)
424 before = math.fsum(sorted((cost.llm for cost in costs), reverse=True)[::group_size])
425 return before - best
427 def _repair_objective(self, costs: Sequence[WorkloadCost], group_size: int) -> float:
428 if self._uses_scheduled_makespan():
429 costs = sorted(costs, key=lambda cost: cost.llm, reverse=True)
430 return math.fsum(
431 self.objective(costs[start:start + group_size]) for start in range(0, len(costs), group_size)
432 )
434 def _repair_bins(
435 self,
436 ordered: Sequence[SampleMetadata],
437 capacity: int,
438 reference_bins: Sequence[Sequence[SampleKey]],
439 window_costs: Sequence[WorkloadCost],
440 pair: tuple[int, int],
441 group_size: int,
442 ) -> Sequence[Sequence[SampleKey]]:
443 """Greedily repartition a pair while suffix reachability preserves capacity feasibility."""
444 # Python integers store bits in digits; account for their object overhead.
445 digits = (capacity + 1 + sys.int_info.bits_per_digit - 1) // sys.int_info.bits_per_digit
446 table_bytes = (len(ordered) + 1) * (sys.getsizeof(0) + digits * sys.int_info.sizeof_digit)
447 # Allow document-rich 1M-token packs while bounding each repair's workspace.
448 if table_bytes > 256 * 1024 * 1024:
449 return reference_bins
450 suffix = [1] * (len(ordered) + 1)
451 mask = (1 << (capacity + 1)) - 1
452 for index in range(len(ordered) - 1, -1, -1):
453 reachable = suffix[index + 1]
454 suffix[index] = (reachable | reachable << ordered[index].pack_tokens) & mask
455 tokens = [0, 0]
456 costs = [0.0, 0.0]
457 loads = [cost.llm for cost in window_costs]
458 bins: list[list[SampleKey]] = [[], []]
459 remaining = sum(item.pack_tokens for item in ordered)
460 scheduled_makespan = self._uses_scheduled_makespan()
461 for index, item in enumerate(ordered):
462 remaining -= item.pack_tokens
463 candidates = []
464 for rank in range(2):
465 trial_tokens = tokens.copy()
466 trial_tokens[rank] += item.pack_tokens
467 low = max(0, 1 - trial_tokens[0], remaining + trial_tokens[1] - capacity)
468 high = min(remaining, capacity - trial_tokens[0], remaining + trial_tokens[1] - 1)
469 # Reachability within this interval guarantees two non-empty bins,
470 # including when their combined tokens leave room for padding.
471 if low > high or not (suffix[index + 1] >> low) & ((1 << (high - low + 1)) - 1):
472 continue
473 trial = loads.copy()
474 trial[pair[0]], trial[pair[1]] = costs
475 trial[pair[rank]] += item.cost.llm
476 for side in range(2):
477 trial[pair[side]] += self._padding_llm_cost(capacity - trial_tokens[side])
478 if scheduled_makespan:
479 score = math.fsum(sorted(trial, reverse=True)[::group_size])
480 elif self.padding_cost is not None:
481 score = self._repair_objective([WorkloadCost(llm=value) for value in trial], group_size)
482 else:
483 score = costs[rank]
484 candidates.append((score, capacity - trial_tokens[rank], rank))
485 if not candidates:
486 return reference_bins
487 rank = min(candidates)[-1]
488 bins[rank].append(item.key)
489 tokens[rank] += item.pack_tokens
490 costs[rank] += item.cost.llm
491 return tuple(tuple(keys) for keys in bins)
494def resolve_balancing_algorithm(algorithm: BalancingAlgorithm | None = None) -> BalancingAlgorithm:
495 """Default to LPT with bounded feasibility repair or validate an explicit policy.
497 Args:
498 algorithm: Optional deterministic assignment and objective implementation.
499 """
500 resolved = FeasibleLPTBalancingAlgorithm() if algorithm is None else algorithm
501 if not callable(getattr(resolved, "assign", None)) or not callable(getattr(resolved, "objective", None)):
502 raise ValueError("balancing_algorithm must provide callable assign() and objective() methods.")
503 return resolved
506__all__ = [
507 "BalancingAlgorithm", "FeasibleLPTBalancingAlgorithm", "LPTBalancingAlgorithm", "resolve_balancing_algorithm",
508]