Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / schema.py: 67%
215 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"""Stable contracts for sample-level distributed dynamic packing."""
17from __future__ import annotations
19import math
20from collections.abc import Mapping, Sequence
21from dataclasses import InitVar, dataclass, field
22from typing import Any, Literal
25OversizedPolicy = Literal["error", "single"]
28def _validate_nonnegative_costs(values: Mapping[str, float], name: str) -> None:
29 """Validate named finite, non-negative scalar costs or budgets."""
30 if not isinstance(values, Mapping):
31 raise ValueError(f"{name} must be a mapping of non-empty strings to finite, non-negative numbers.")
32 for key, value in values.items():
33 if not isinstance(key, str) or not key:
34 raise ValueError(f"{name} keys must be non-empty strings, but got {key!r}.")
35 if (
36 type(value) not in (int, float)
37 or (isinstance(value, float) and not math.isfinite(value))
38 or value < 0
39 ):
40 raise ValueError(f"{name}[{key!r}] must be finite and non-negative, but got {value!r}.")
43def _validate_feature(value: Any, name: str, ancestors: set[int]) -> None:
44 """Reject tensors, arbitrary objects, cycles, and non-finite feature values."""
45 if value is None:
46 return
47 if type(value) not in (bool, int, float, str, dict, list, tuple):
48 raise ValueError(f"{name} must contain only CPU-serializable basic types, but got {type(value)}.")
49 if isinstance(value, (bool, int, str)):
50 return
51 if isinstance(value, float):
52 if not math.isfinite(value):
53 raise ValueError(f"{name} must be finite, but got {value!r}.")
54 return
55 if id(value) in ancestors:
56 raise ValueError(f"{name} must not contain reference cycles.")
57 ancestors.add(id(value))
58 if isinstance(value, dict):
59 for key, child in value.items():
60 if not isinstance(key, str):
61 raise ValueError(f"{name} keys must be strings, but got {key!r}.")
62 _validate_feature(child, f"{name}[{key!r}]", ancestors)
63 else:
64 for index, child in enumerate(value):
65 _validate_feature(child, f"{name}[{index}]", ancestors)
66 ancestors.remove(id(value))
69@dataclass(frozen=True)
70class WorkloadCost:
71 """Normalized sample costs used to balance multimodal work across DP ranks.
73 The planner keeps stage costs separate so a user-supplied metadata callback
74 can represent samples that stress different parts of a multimodal model.
75 """
77 io: float = 0.0
78 transform: float = 0.0
79 encoder: float = 0.0
80 llm: float = 0.0
81 memory: float = 0.0
82 communication: float = 0.0
84 def __post_init__(self) -> None:
85 """Validate every cost component."""
86 for name, value in self.as_dict().items():
87 if (
88 not isinstance(value, (int, float))
89 or isinstance(value, bool)
90 or not math.isfinite(value)
91 or value < 0
92 ):
93 raise ValueError(f"WorkloadCost.{name} must be finite and non-negative, but got {value!r}.")
95 def __add__(self, other: "WorkloadCost") -> "WorkloadCost":
96 """Add corresponding stage costs."""
97 if not isinstance(other, WorkloadCost):
98 return NotImplemented
99 return WorkloadCost(
100 io=self.io + other.io,
101 transform=self.transform + other.transform,
102 encoder=self.encoder + other.encoder,
103 llm=self.llm + other.llm,
104 memory=self.memory + other.memory,
105 communication=self.communication + other.communication,
106 )
108 @property
109 def dominant(self) -> float:
110 """Return the largest stage cost."""
111 return max(self.as_tuple())
113 @property
114 def total(self) -> float:
115 """Return the sum of stage costs."""
116 return sum(self.as_tuple())
118 def as_tuple(self) -> tuple[float, ...]:
119 """Return components in deterministic planner order."""
120 return (self.io, self.transform, self.encoder, self.llm, self.memory, self.communication)
122 def as_dict(self) -> dict[str, float]:
123 """Return a serializable component mapping."""
124 return {
125 "io": self.io,
126 "transform": self.transform,
127 "encoder": self.encoder,
128 "llm": self.llm,
129 "memory": self.memory,
130 "communication": self.communication,
131 }
134@dataclass(frozen=True)
135class SampleMetadata:
136 """Lightweight metadata used to plan one Dataset sample.
138 Args:
139 pack_tokens: Tokens occupied by the sample in a packed sequence. This
140 must use the same accounting as ``pack_fn``.
141 cost: Optional normalized multimodal workload estimate.
142 sample_id: Optional user-facing identifier used for diagnostics. The
143 runtime uses :class:`SampleKey` for routing and uniqueness.
144 features: CPU metadata for workload estimation. Values may recursively
145 contain ``None``, booleans, integers, finite floats, strings, lists,
146 tuples, and dictionaries with string keys. Tensors are rejected.
147 packing_costs: Named finite, non-negative physical packing footprints.
148 These are independent of ``cost`` and enforce per-bin hard budgets.
149 key: Runtime sample occurrence assigned by the Reader. Callbacks and
150 precomputed metadata may leave it unset; direct planner inputs require it.
151 """
153 pack_tokens: int
154 cost: WorkloadCost = WorkloadCost()
155 sample_id: int | str | None = None
156 features: dict[str, Any] = field(default_factory=dict)
157 packing_costs: dict[str, float] = field(default_factory=dict)
158 key: SampleKey | None = None
160 def __post_init__(self) -> None:
161 """Validate identity, workload, and lightweight packing metadata."""
162 if not isinstance(self.pack_tokens, int) or isinstance(self.pack_tokens, bool) or self.pack_tokens < 1:
163 raise ValueError(f"SampleMetadata.pack_tokens must be a positive integer, but got {self.pack_tokens!r}.")
164 if not isinstance(self.cost, WorkloadCost):
165 raise ValueError(f"SampleMetadata.cost must be WorkloadCost, but got {type(self.cost)}.")
166 if self.sample_id is not None:
167 valid_id = isinstance(self.sample_id, (int, str)) and not isinstance(self.sample_id, bool)
168 if not valid_id or (isinstance(self.sample_id, int) and self.sample_id < 0):
169 raise ValueError("SampleMetadata.sample_id must be a non-negative integer, non-empty string, or None.")
170 if isinstance(self.sample_id, str) and not self.sample_id:
171 raise ValueError("SampleMetadata.sample_id string must not be empty.")
172 if not isinstance(self.features, dict):
173 raise ValueError("SampleMetadata.features must be a dictionary.")
174 _validate_feature(self.features, "SampleMetadata.features", set())
175 if not isinstance(self.packing_costs, dict):
176 raise ValueError("SampleMetadata.packing_costs must be a dictionary.")
177 _validate_nonnegative_costs(self.packing_costs, "SampleMetadata.packing_costs")
178 if not isinstance(self.key, (SampleKey, type(None))):
179 raise ValueError("SampleMetadata.key must be a SampleKey or None.")
182@dataclass(frozen=True, order=True)
183class SampleKey:
184 """Stable identity for one Dataset sample occurrence owned by a Reader.
186 ``global_sample_position`` identifies an occurrence in the source stream.
187 Native samplers and Dataset Readers use global stream positions; local
188 step sources use a rank-local position within the step.
189 ``sequence_index`` identifies an independent sequence inside a selected
190 packed record; whole-record loading leaves it at zero.
191 """
193 reader_rank: int
194 dataset_index: int
195 global_sample_position: int = 0
196 sequence_index: int = 0
198 def __post_init__(self) -> None:
199 """Validate non-negative routing coordinates."""
200 for name in ("reader_rank", "dataset_index", "global_sample_position", "sequence_index"):
201 value = getattr(self, name)
202 if not isinstance(value, int) or isinstance(value, bool) or value < 0:
203 raise ValueError(f"SampleKey.{name} must be a non-negative integer, but got {value!r}.")
206@dataclass(frozen=True)
207class PackingConstraints:
208 """Shared packing feasibility contract for selection and balanced placement.
210 Args:
211 seq_len: Token capacity of a regular bin.
212 oversized_policy: Permit token overflow only for a singleton when
213 ``single``. Named stage budgets never permit singleton overflow.
214 packing_budgets: Per-bin additive hard caps. Every configured name must
215 be explicitly present in each sample's ``packing_costs``, including
216 zero-cost stages. ``None`` retains token-only packing.
217 """
219 seq_len: int
220 oversized_policy: OversizedPolicy = "error"
221 packing_budgets: Mapping[str, float] | None = None
223 def __post_init__(self) -> None:
224 """Validate token and stage capacities and copy the caller's mapping."""
225 if not isinstance(self.seq_len, int) or isinstance(self.seq_len, bool) or self.seq_len < 1:
226 raise ValueError(f"seq_len must be a positive integer, but got {self.seq_len!r}.")
227 if self.oversized_policy not in ("error", "single"):
228 raise ValueError("oversized_policy must be 'error' or 'single'.")
229 budgets = {} if self.packing_budgets is None else self.packing_budgets
230 _validate_nonnegative_costs(budgets, "packing_budgets")
231 object.__setattr__(self, "packing_budgets", dict(budgets))
233 def validate_sample(self, item: SampleMetadata) -> None:
234 """Reject samples that cannot occupy even an otherwise empty bin.
236 Args:
237 item: Sample metadata with physical packing footprints.
238 """
239 for name, budget in self.packing_budgets.items():
240 if name not in item.packing_costs:
241 raise ValueError(f"Sample {item.key} is missing packing_costs[{name!r}] required by packing_budgets.")
242 value = item.packing_costs[name]
243 if value > budget:
244 raise ValueError(
245 f"Sample {item.key} requires packing_costs[{name!r}]={value}, "
246 f"exceeding packing_budgets[{name!r}]={budget}. Stage budgets do not permit singleton overflow."
247 )
248 if item.pack_tokens > self.seq_len and self.oversized_policy == "error":
249 raise ValueError(
250 f"Sample {item.key} requires {item.pack_tokens} tokens, exceeding seq_len={self.seq_len}. "
251 "Set oversized_policy='single' only when the packer supports singleton overflow."
252 )
254 def fits(
255 self,
256 pack_tokens: int,
257 packing_costs: Mapping[str, float],
258 item: SampleMetadata,
259 ) -> bool:
260 """Return whether an individually valid sample fits the current bin.
262 Args:
263 pack_tokens: Tokens already assigned to the bin.
264 packing_costs: Physical footprints already assigned to the bin.
265 item: Validated metadata for the sample being considered.
266 """
267 if item.pack_tokens > self.seq_len:
268 token_fit = pack_tokens == 0 and self.oversized_policy == "single"
269 else:
270 token_fit = pack_tokens + item.pack_tokens <= self.seq_len
271 return token_fit and all(
272 packing_costs.get(name, 0) + item.packing_costs[name] <= budget
273 for name, budget in self.packing_budgets.items()
274 )
276 def add_costs(
277 self,
278 packing_costs: Mapping[str, float],
279 item: SampleMetadata,
280 ) -> dict[str, float]:
281 """Return accumulated physical costs for the constrained stages.
283 Args:
284 packing_costs: Physical footprints already assigned to the bin.
285 item: Sample metadata whose footprints will be added.
286 """
287 return {
288 name: packing_costs.get(name, 0) + item.packing_costs[name]
289 for name in self.packing_budgets
290 }
292 def validate_bin(self, items: Sequence[SampleMetadata]) -> None:
293 """Validate a reference or final bin against the same placement rules.
295 Args:
296 items: Metadata for the samples assigned to one bin.
297 """
298 pack_tokens = 0
299 packing_costs: dict[str, float] = {}
300 for item in items:
301 self.validate_sample(item)
302 if not self.fits(pack_tokens, packing_costs, item):
303 raise ValueError(
304 f"Packing bin cannot admit sample {item.key} within seq_len={self.seq_len} "
305 f"and packing_budgets={dict(self.packing_budgets)}."
306 )
307 pack_tokens += item.pack_tokens
308 packing_costs = self.add_costs(packing_costs, item)
311@dataclass(frozen=True)
312class PackingBinPlan:
313 """Ordered raw samples that one Data Constructor passes to ``pack_fn``.
315 ``validate=False`` skips repeated scans for internally constructed, trusted
316 bins. The option is construction-only and is not serialized.
317 """
319 sample_keys: tuple[SampleKey, ...]
320 pack_tokens: int
321 oversized: bool = False
322 validate: InitVar[bool] = True
324 def __post_init__(self, validate: bool) -> None:
325 """Validate one non-empty sequence-packing bin."""
326 if not validate:
327 return
328 if not self.sample_keys:
329 raise ValueError("PackingBinPlan.sample_keys must not be empty.")
330 if any(not isinstance(key, SampleKey) for key in self.sample_keys):
331 raise ValueError("PackingBinPlan.sample_keys must contain SampleKey values.")
332 if len(self.sample_keys) != len(set(self.sample_keys)):
333 raise ValueError("PackingBinPlan.sample_keys must not contain duplicates.")
334 if not isinstance(self.pack_tokens, int) or isinstance(self.pack_tokens, bool) or self.pack_tokens < 1:
335 raise ValueError(f"PackingBinPlan.pack_tokens must be positive, but got {self.pack_tokens!r}.")
336 if not isinstance(self.oversized, bool):
337 raise ValueError(f"PackingBinPlan.oversized must be boolean, but got {self.oversized!r}.")
340@dataclass(frozen=True)
341class DistributedPackingPlan:
342 """Deterministic sample-to-constructor plan for one distributed yield.
344 ``validate=False`` skips repeated scans for internally constructed, trusted
345 plans. The option is construction-only and is not serialized.
346 """
348 step: int
349 seq_len: int
350 local_batches: tuple[tuple[PackingBinPlan, ...], ...]
351 rank_costs: tuple[WorkloadCost, ...]
352 validate: InitVar[bool] = True
354 def __post_init__(self, validate: bool) -> None:
355 """Validate plan dimensions, slots, and unique sample assignments."""
356 if not validate:
357 return
358 self._validate_dimensions()
359 keys = self._validate_local_batches()
360 if len(keys) != len(set(keys)):
361 raise ValueError("A sample key may appear only once in a distributed packing plan.")
363 def _validate_dimensions(self) -> None:
364 """Require valid step and sequence length plus matching non-empty rank batches."""
365 if not isinstance(self.step, int) or isinstance(self.step, bool) or self.step < 0:
366 raise ValueError(f"DistributedPackingPlan.step must be non-negative, but got {self.step!r}.")
367 if not isinstance(self.seq_len, int) or isinstance(self.seq_len, bool) or self.seq_len < 1:
368 raise ValueError(f"DistributedPackingPlan.seq_len must be positive, but got {self.seq_len!r}.")
369 if not self.local_batches:
370 raise ValueError("DistributedPackingPlan.local_batches must not be empty.")
371 if len(self.rank_costs) != len(self.local_batches):
372 raise ValueError("DistributedPackingPlan.rank_costs must match local_batches.")
373 if any(not isinstance(cost, WorkloadCost) for cost in self.rank_costs):
374 raise ValueError("DistributedPackingPlan.rank_costs must contain WorkloadCost values.")
375 batch_sizes = {len(local_batch) for local_batch in self.local_batches}
376 if len(batch_sizes) != 1 or 0 in batch_sizes:
377 raise ValueError("DistributedPackingPlan local batches must have equal non-zero sizes.")
379 @property
380 def data_parallel_size(self) -> int:
381 """Return the number of target Data Constructor ranks."""
382 return len(self.local_batches)
384 @property
385 def local_batch_size(self) -> int:
386 """Return the number of packed sequences per target rank."""
387 return len(self.local_batches[0])
389 def _validate_local_batches(self) -> list[SampleKey]:
390 keys = []
391 for local_batch in self.local_batches:
392 self._validate_local_batch(local_batch)
393 keys.extend(key for packing_bin in local_batch for key in packing_bin.sample_keys)
394 return keys
396 def _validate_local_batch(self, local_batch: tuple[PackingBinPlan, ...]) -> None:
397 if any(not isinstance(packing_bin, PackingBinPlan) for packing_bin in local_batch):
398 raise ValueError("DistributedPackingPlan local batches must contain PackingBinPlan values.")
399 for packing_bin in local_batch:
400 if packing_bin.pack_tokens > self.seq_len and not packing_bin.oversized:
401 raise ValueError("A non-oversized packing bin exceeds seq_len.")
402 if packing_bin.oversized and len(packing_bin.sample_keys) != 1:
403 raise ValueError("An oversized sample must occupy its packing bin alone.")
405 @property
406 def selected_keys(self) -> tuple[SampleKey, ...]:
407 """Return all selected keys in deterministic constructor order."""
408 return tuple(
409 key
410 for local_batch in self.local_batches
411 for packing_bin in local_batch
412 for key in packing_bin.sample_keys
413 )
415 def local_batch_for(self, data_rank: int) -> tuple[PackingBinPlan, ...]:
416 """Return the ordered packed sequences for one data-parallel rank.
418 Args:
419 data_rank: Rank within the data-parallel domain.
421 Returns:
422 Target rank's local packing bins.
423 """
424 valid_rank = isinstance(data_rank, int) and not isinstance(data_rank, bool)
425 if not valid_rank or not 0 <= data_rank < self.data_parallel_size:
426 raise ValueError(f"data_rank must be in [0, {self.data_parallel_size}), but got {data_rank!r}.")
427 return self.local_batches[data_rank]
429 def local_sample_keys(self, data_rank: int) -> tuple[SampleKey, ...]:
430 """Return the sample keys assigned to one target data rank.
432 Args:
433 data_rank: Rank within the data-parallel domain.
435 Returns:
436 Sample keys in local pack order.
437 """
438 return tuple(
439 key
440 for packing_bin in self.local_batch_for(data_rank)
441 for key in packing_bin.sample_keys
442 )