Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / planner.py: 95%
130 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"""Estimate frozen samples, propose assignments, and accept beneficial plans."""
17from __future__ import annotations
19import math
20from collections.abc import Mapping, Sequence
21from copy import copy
22from itertools import chain
23from typing import Any
25from hyper_parallel.distributed_data.balancing_algorithm import (
26 BalancingAlgorithm,
27 FeasibleLPTBalancingAlgorithm,
28 LPTBalancingAlgorithm,
29 resolve_balancing_algorithm,
30)
31from hyper_parallel.distributed_data.cost_model import CostModel, resolve_cost_model
32from hyper_parallel.distributed_data.schema import (
33 DistributedPackingPlan,
34 OversizedPolicy,
35 PackingBinPlan,
36 PackingConstraints,
37 SampleKey,
38 SampleMetadata,
39 WorkloadCost,
40)
43class DynamicPackingPlanner:
44 """Keep cost estimation and placement independently replaceable.
46 Every step is evaluated using the supplied or default workload model and
47 assignment algorithm. The original rank-local grouping is retained unless
48 the candidate's relative objective reduction strictly exceeds the threshold.
49 """
51 def __init__(
52 self,
53 *,
54 data_parallel_size: int,
55 seq_len: int,
56 local_batch_size: int,
57 oversized_policy: OversizedPolicy = "error",
58 packing_budgets: Mapping[str, float] | None = None,
59 model_config: Any = None,
60 cost_model: CostModel | None = None,
61 balancing_algorithm: BalancingAlgorithm | None = None,
62 validate: bool = True,
63 min_balance_gain: float = 0.0,
64 ) -> None:
65 """Configure packing boundaries and independent planning policies.
67 Args:
68 data_parallel_size: Number of target DP ranks.
69 seq_len: Token capacity of one packed sequence.
70 local_batch_size: Non-empty packed sequences required per rank.
71 oversized_policy: ``error`` or singleton token overflow.
72 packing_budgets: Additive per-bin stage limits, independent of costs.
73 model_config: Effective architecture for the default workload model.
74 cost_model: Optional deterministic SampleMetadata-to-WorkloadCost callback.
75 balancing_algorithm: Optional assignment and objective implementation.
76 validate: Audit trusted input/output membership. Custom assignments
77 always receive basic shape, conservation, and capacity checks.
78 min_balance_gain: Minimum relative objective reduction in [0, 1).
79 Equal, worse, or below-threshold proposals retain the reference.
80 """
81 for name, value in (
82 ("data_parallel_size", data_parallel_size),
83 ("seq_len", seq_len),
84 ("local_batch_size", local_batch_size),
85 ):
86 if not isinstance(value, int) or isinstance(value, bool) or value < 1:
87 raise ValueError(f"{name} must be a positive integer, but got {value!r}.")
88 if (
89 not isinstance(min_balance_gain, (int, float))
90 or isinstance(min_balance_gain, bool)
91 or not math.isfinite(min_balance_gain)
92 or not 0.0 <= min_balance_gain < 1.0
93 ):
94 raise ValueError("min_balance_gain must be in [0, 1).")
95 self.data_parallel_size = data_parallel_size
96 self.seq_len = seq_len
97 self.local_batch_size = local_batch_size
98 self.oversized_policy = oversized_policy
99 self._validate = validate
100 self._constraints = PackingConstraints(seq_len, oversized_policy, packing_budgets)
101 self.cost_model = resolve_cost_model(cost_model, model_config)
102 self.balancing_algorithm = resolve_balancing_algorithm(balancing_algorithm)
103 if callable(getattr(self.cost_model, "padding_cost", None)) and isinstance(
104 self.balancing_algorithm, LPTBalancingAlgorithm,
105 ):
106 self.balancing_algorithm = copy(self.balancing_algorithm)
107 self.balancing_algorithm.padding_cost = self.padding_cost
108 self.objective = getattr(
109 self.balancing_algorithm, "objective_name", type(self.balancing_algorithm).__qualname__,
110 )
111 self.min_balance_gain = float(min_balance_gain)
112 self.last_sample_costs: dict[SampleKey, WorkloadCost] = {}
113 self.last_balance_decision: dict[str, float | bool] = {}
115 @property
116 def distributed_bin_count(self) -> int:
117 """Return the number of bins required for one distributed yield."""
118 return self.data_parallel_size * self.local_batch_size
120 @property
121 def distributed_token_budget(self) -> int:
122 """Return the aggregate non-oversized token capacity."""
123 return self.distributed_bin_count * self.seq_len
125 def padding_cost(self, padding_tokens: int) -> WorkloadCost:
126 """Estimate constructor padding in the same units as source-sample costs.
128 Args:
129 padding_tokens: Number of tokens appended to one packed row.
131 Returns:
132 Padding workload, or zero when the model has no padding estimate.
133 """
134 estimate = getattr(self.cost_model, "padding_cost", lambda _: WorkloadCost())
135 if not callable(estimate):
136 raise ValueError("cost_model.padding_cost must be callable")
137 cost = estimate(padding_tokens)
138 if not isinstance(cost, WorkloadCost):
139 raise ValueError("cost_model.padding_cost must return WorkloadCost")
140 return cost
142 def plan(
143 self,
144 samples: Sequence[SampleMetadata],
145 *,
146 reference_bins: Sequence[Sequence[SampleKey]],
147 step: int,
148 target_slot_count: int | None = None,
149 ) -> DistributedPackingPlan:
150 """Evaluate one frozen sample set and preserve non-beneficial layouts.
152 Args:
153 samples: Exactly the selected occurrences, each with a bound SampleKey.
154 reference_bins: Known-feasible original grouping, ordered by
155 micro-batch, DP rank, and local bin. Every selected key occurs once.
156 step: Zero-based index of the first distributed micro-batch.
157 target_slot_count: Number of rank-by-micro-batch destinations.
158 Defaults to data_parallel_size for one round. Must be a positive
159 integer multiple of data_parallel_size; each slot holds
160 local_batch_size bins.
162 Returns:
163 A canonical plan with one local batch per target slot. Window slots
164 are ordered by micro-batch, then DP rank.
165 """
166 target_slot_count = self._resolve_target_slot_count(target_slot_count)
167 samples = self._estimate_samples(samples)
168 for item in samples:
169 self._constraints.validate_sample(item)
170 samples_by_key = {item.key: item for item in samples}
171 original, reference_costs = self._build_rank_packing_plans(
172 reference_bins, samples_by_key, target_slot_count=target_slot_count, audit=self._validate,
173 )
174 assignment_options = {}
175 if (
176 isinstance(self.balancing_algorithm, FeasibleLPTBalancingAlgorithm)
177 and type(self.balancing_algorithm).assign is FeasibleLPTBalancingAlgorithm.assign
178 ):
179 # Custom assign implementations keep the existing policy contract.
180 assignment_options["objective_group_size"] = self.data_parallel_size
181 candidate_keys = self.balancing_algorithm.assign(
182 samples,
183 reference_bins=reference_bins,
184 constraints=self._constraints,
185 data_parallel_size=target_slot_count,
186 local_batch_size=self.local_batch_size,
187 **assignment_options,
188 )
189 trusted_assignment = (
190 isinstance(self.balancing_algorithm, LPTBalancingAlgorithm)
191 and type(self.balancing_algorithm).assign is LPTBalancingAlgorithm.assign
192 )
193 audit_candidate = self._validate or not trusted_assignment
194 candidate, candidate_costs = self._build_rank_packing_plans(
195 candidate_keys, samples_by_key, target_slot_count=target_slot_count, audit=audit_candidate,
196 )
197 if (
198 self.data_parallel_size < target_slot_count and self.local_batch_size == 1
199 and isinstance(self.balancing_algorithm, LPTBalancingAlgorithm)
200 ):
201 # Validate assignments before scheduling so a feasible reference
202 # cannot hide an invalid candidate from a custom LPT subclass.
203 scheduled = (
204 self._schedule_window_bins(batches, costs)
205 for batches, costs in ((candidate, candidate_costs), (original, reference_costs))
206 )
207 # Cost sorting need not improve normalized variance or subclass
208 # objectives, so retain the original candidate as an option.
209 candidate, candidate_costs = min(
210 ((candidate, candidate_costs), *scheduled), key=lambda result: self._objective_value(result[1]),
211 )
212 accepted = self._accept_candidate(reference_costs, candidate_costs)
213 local_batches, rank_costs = (candidate, candidate_costs) if accepted else (original, reference_costs)
214 return DistributedPackingPlan(
215 step=step,
216 seq_len=self.seq_len,
217 local_batches=local_batches,
218 rank_costs=rank_costs,
219 validate=False,
220 )
222 def plan_window(
223 self,
224 samples: Sequence[SampleMetadata],
225 *,
226 reference_bins: Sequence[Sequence[SampleKey]],
227 step: int,
228 ) -> tuple[DistributedPackingPlan, ...]:
229 """Balance native rounds jointly over rank-by-micro-batch destinations.
231 Args:
232 samples: The frozen union of the window's native sample occurrences.
233 reference_bins: Original bins ordered by micro-batch, DP rank, then local bin.
234 step: First distributed micro-batch index in the window.
236 Returns:
237 One ordinary distributed plan per micro-batch. The benefit gate sums
238 per-round objectives so improvements do not hide a slower window.
239 """
240 rounds, remainder = divmod(len(reference_bins), self.distributed_bin_count)
241 if not rounds or remainder:
242 raise ValueError("A balancing window must contain complete distributed micro-batches.")
243 plan = self.plan(
244 samples, reference_bins=reference_bins, step=step, target_slot_count=self.data_parallel_size * rounds,
245 )
246 return tuple(
247 DistributedPackingPlan(
248 step=step + offset,
249 seq_len=self.seq_len,
250 local_batches=plan.local_batches[start:start + self.data_parallel_size],
251 rank_costs=plan.rank_costs[start:start + self.data_parallel_size],
252 validate=False,
253 )
254 for offset, start in enumerate(range(0, len(plan.local_batches), self.data_parallel_size))
255 )
257 def _resolve_target_slot_count(self, target_slot_count: int | None) -> int:
258 """Default to one round and reject incomplete or empty target windows."""
259 if target_slot_count is None:
260 return self.data_parallel_size
261 if (
262 not isinstance(target_slot_count, int) or isinstance(target_slot_count, bool)
263 or target_slot_count < 1 or target_slot_count % self.data_parallel_size
264 ):
265 raise ValueError("target_slot_count must be a positive integer multiple of data_parallel_size.")
266 return target_slot_count
268 @staticmethod
269 def _schedule_window_bins(
270 local_batches: tuple[tuple[PackingBinPlan, ...], ...],
271 rank_costs: tuple[WorkloadCost, ...],
272 ) -> tuple[tuple[tuple[PackingBinPlan, ...], ...], tuple[WorkloadCost, ...]]:
273 """Group intact, validated packs by cost within the frozen window."""
274 order = sorted(
275 range(len(local_batches)),
276 key=lambda index: (-rank_costs[index].llm, local_batches[index][0].sample_keys),
277 )
278 return tuple(local_batches[index] for index in order), tuple(rank_costs[index] for index in order)
280 def _estimate_samples(
281 self, samples: Sequence[SampleMetadata],
282 ) -> tuple[SampleMetadata, ...]:
283 """Estimate keyed sample costs without mutating physical source metadata."""
284 self.last_sample_costs = {}
285 estimated_samples = []
286 for item in samples:
287 if not isinstance(item.key, SampleKey):
288 raise ValueError("Planner samples must have a SampleKey.")
289 cost = self.cost_model(item)
290 if not isinstance(cost, WorkloadCost):
291 raise ValueError(f"cost_model must return WorkloadCost for sample {item.key}, but got {type(cost)}.")
292 # Preserve physical metadata without replaying recursive feature validation.
293 estimated_item = copy(item)
294 object.__setattr__(estimated_item, "cost", cost)
295 estimated_samples.append(estimated_item)
296 self.last_sample_costs[item.key] = cost
297 return tuple(estimated_samples)
299 def _build_rank_packing_plans(
300 self,
301 key_bins: Sequence[Sequence[SampleKey]],
302 samples_by_key: Mapping[SampleKey, SampleMetadata],
303 *,
304 target_slot_count: int,
305 audit: bool,
306 ) -> tuple[tuple[tuple[PackingBinPlan, ...], ...], tuple[WorkloadCost, ...]]:
307 """Validate assignments and reconstruct per-slot bins and costs, including new PAD."""
308 key_bins = tuple(tuple(keys) for keys in key_bins)
309 if len(key_bins) != target_slot_count * self.local_batch_size or any(not keys for keys in key_bins):
310 raise ValueError("Assignments must contain target_slot_count * local_batch_size non-empty bins.")
311 if audit:
312 planned_keys = list(chain.from_iterable(key_bins))
313 if len(planned_keys) != len(samples_by_key) or set(planned_keys) != set(samples_by_key):
314 raise ValueError("Assignments must conserve every selected sample key exactly once.")
315 local_batches = [[] for _ in range(target_slot_count)]
316 rank_costs = [WorkloadCost() for _ in range(target_slot_count)]
317 for index, keys in enumerate(key_bins):
318 items = tuple(samples_by_key[key] for key in keys)
319 self._constraints.validate_bin(items)
320 tokens = sum(item.pack_tokens for item in items)
321 slot = index // self.local_batch_size
322 rank_costs[slot] += sum((item.cost for item in items), WorkloadCost())
323 rank_costs[slot] += self.padding_cost(max(0, self.seq_len - tokens))
324 local_batches[slot].append(PackingBinPlan(
325 sample_keys=keys,
326 pack_tokens=tokens,
327 oversized=tokens > self.seq_len,
328 validate=False,
329 ))
330 return tuple(tuple(bins) for bins in local_batches), tuple(rank_costs)
332 def _accept_candidate(
333 self, reference_costs: Sequence[WorkloadCost], candidate_costs: Sequence[WorkloadCost],
334 ) -> bool:
335 """Record the objective comparison and enforce the minimum relative-gain threshold."""
336 before = self._objective_value(reference_costs)
337 after = self._objective_value(candidate_costs)
338 gain = (before - after) / before if before > 0.0 else 0.0
339 accepted = gain > self.min_balance_gain
340 self.last_balance_decision = {
341 "original_objective": before,
342 "candidate_objective": after,
343 "accepted_objective": after if accepted else before,
344 "relative_gain": gain,
345 "min_balance_gain": self.min_balance_gain,
346 "accepted": accepted,
347 }
348 return accepted
350 def _objective_value(self, rank_costs: Sequence[WorkloadCost]) -> float:
351 scores = (
352 self._round_objective(rank_costs[start:start + self.data_parallel_size])
353 for start in range(0, len(rank_costs), self.data_parallel_size)
354 )
355 return math.fsum(scores)
357 def _round_objective(self, rank_costs: Sequence[WorkloadCost]) -> float:
358 score = self.balancing_algorithm.objective(rank_costs)
359 if (
360 not isinstance(score, (int, float)) or isinstance(score, bool)
361 or not math.isfinite(score) or score < 0.0
362 ):
363 raise ValueError("balancing_algorithm.objective must return a finite, non-negative number.")
364 return float(score)
367__all__ = ["DynamicPackingPlanner", "OversizedPolicy"]