Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / heap_manager.py: 91%
327 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 05:03 +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"""MegaMoe workspace capacity and SHMEM Heap policy."""
17from __future__ import annotations
19import math
20import os
21import threading
22import time
23from collections.abc import Callable, Iterator, Mapping
24from contextlib import contextmanager
25from dataclasses import dataclass, replace
26from typing import Any, NoReturn
28import torch
30from hyper_multicore import _collective, shmem
31from hyper_multicore.shmem import _lifecycle
33from .spec import (
34 _DEFAULT_CAPACITY_FACTOR,
35 MegaMoeSpec,
36 _align_capacity,
37 _grow_receive_capacity,
38 _resolve_receive_capacity,
39)
40from .workspace import MegaMoeWorkspace, _spec_workspace_bytes
43_HEAP_GRANULARITY_BYTES = 2 * 1024 * 1024
44_HEAP_MANAGER = None
47def _round_heap_bytes(value: int) -> int:
48 """Round the aggregate payload size to physical Heap pages."""
49 return (value + _HEAP_GRANULARITY_BYTES - 1) // _HEAP_GRANULARITY_BYTES * _HEAP_GRANULARITY_BYTES
52@dataclass(frozen=True)
53class _WorkspaceLayout:
54 """One reservation's bytes, collective identity and receive capacity."""
56 required_bytes: int
57 collective_signature: tuple[Any, ...]
58 receive_capacity: int
59 token_capacity: int | None = None
62@dataclass
63class _HeapEntry:
64 """One ordered lazy reservation and its optional bound workspace."""
66 specification: Mapping[str, Any]
67 layout: _WorkspaceLayout
68 workspace: MegaMoeWorkspace | None = None
69 bound_spec: MegaMoeSpec | None = None
72@dataclass(frozen=True)
73class _RestartEntryAgreement:
74 """Rank-independent state for one entry in a restart transaction."""
76 layout: _WorkspaceLayout
77 is_bound: bool
78 is_materialized: bool
81@dataclass(frozen=True)
82class _RestartAgreement:
83 """Rank-independent entry state for one restart epoch."""
85 epoch: int
86 entries: tuple[_RestartEntryAgreement, ...]
89@dataclass(frozen=True)
90class _RestartBinding:
91 """Map one materialized workspace to a contiguous restart result slice."""
93 entry: _HeapEntry
94 start: int
95 stop: int
98def _workspace_layout(
99 specification: Mapping[str, Any],
100 dtype: torch.dtype,
101 capacity: int | None = None,
102 token_capacity: int | None = None,
103) -> _WorkspaceLayout:
104 """Describe one MegaMoe symmetric workspace reservation."""
105 physical = dict(specification)
106 if token_capacity is not None:
107 physical["local_num_tokens"] = token_capacity
108 if capacity is None:
109 capacity = _resolve_receive_capacity(
110 specification["initial_capacity_factor"],
111 physical["local_num_tokens"] * specification["top_k"],
112 specification["ep_size"],
113 )
114 signature = tuple(sorted((key, value) for key, value in specification.items() if key != "ep_group"))
115 return _WorkspaceLayout(
116 required_bytes=_spec_workspace_bytes(physical, dtype, receive_capacity=capacity),
117 collective_signature=signature,
118 receive_capacity=capacity,
119 token_capacity=physical["local_num_tokens"] if specification.get("auto_tokens", False) else None,
120 )
123class _HeapManager:
124 """Own MegaMoe Heap policy and the sole SHMEM Runtime user reference."""
126 def __init__(self, group: Any, tensor: torch.Tensor) -> None:
127 """Freeze one Root and tensor binding without acquiring the Runtime."""
128 self.root_group = group
129 self.members = _lifecycle.root_members(group)
130 self.device = tensor.device
131 self.dtype = tensor.dtype
132 self.element_size = tensor.element_size()
133 self.entries: list[_HeapEntry] = []
134 self.epoch = 0
135 self.state = "ready"
136 self.growth_records: list[dict[str, Any]] = []
137 self.heap_limit: int | None = None
138 self.heap_bytes = 0
139 self._configured = False
140 self._runtime_attached = False
141 self._lock = threading.RLock()
143 @property
144 def workspaces(self) -> list[MegaMoeWorkspace]:
145 """Return bound workspaces in deterministic Allocation order."""
146 return [entry.workspace for entry in self.entries if entry.workspace is not None]
148 @property
149 def runtime_attached(self) -> bool:
150 """Return whether this manager owns the process Runtime reference."""
151 return self._runtime_attached
153 @contextmanager
154 def access(self) -> Iterator[None]:
155 """Exclude concurrent submissions while allowing same-thread restart callbacks."""
156 if not self._lock.acquire(blocking=False):
157 raise RuntimeError("SHMEM root does not support concurrent Heap operations")
158 try:
159 if self.state == "failed":
160 raise RuntimeError(
161 "SHMEM Heap is no longer usable after a failed restart; restart the process"
162 )
163 yield
164 finally:
165 self._lock.release()
167 def reserve(
168 self,
169 reservations: tuple[tuple[Mapping[str, Any], _WorkspaceLayout], ...],
170 *,
171 allow_resize: bool,
172 ) -> tuple[Mapping[str, Any], ...]:
173 """Add lazy reservations and resize an attached Runtime when required."""
174 with self.access():
175 previous_count = len(self.entries)
176 previous_configuration = self._configured, self.heap_limit, self.heap_bytes
177 added_specifications = self._reserve(reservations)
178 layouts = [entry.layout for entry in self.entries]
179 try:
180 configured = os.getenv("HYPER_MULTICORE_SHMEM_HEAP_SIZE")
181 if not self._configured:
182 self._run_and_check_consistency(
183 lambda: self._configure(configured),
184 stage="Configuration",
185 value=(configured, str(self.dtype), tuple(layouts)),
186 )
187 self._configured = True
188 else:
189 self._run_and_check_consistency(
190 lambda: None,
191 stage="LayoutAgreement",
192 value=tuple(layouts),
193 )
194 required = self._required_bytes(layouts)
195 if required > self.heap_bytes:
196 if self.workspaces and not allow_resize:
197 raise RuntimeError("Create all resources before initializing a fixed SHMEM layout")
198 heap_bytes = self._budget(required)
199 if self._runtime_attached:
200 self._restart(layouts, heap_bytes)
201 else:
202 self.heap_bytes = heap_bytes
203 return added_specifications
204 except Exception:
205 if self.state == "ready":
206 del self.entries[previous_count:]
207 self._configured, self.heap_limit, self.heap_bytes = previous_configuration
208 raise
210 def _run_and_check_consistency(self, operation: Callable[[], None], *, stage: str, value: Any) -> None:
211 """Run a local operation, then check failures and values across the Root group.
213 One result exchange checks that every rank succeeded and supplied the same value. A single-rank configuration
214 ValueError is propagated unchanged.
215 """
216 local_error = None
217 try:
218 operation()
219 except Exception as error: # pylint: disable=broad-exception-caught
220 local_error = error
221 report = None if local_error is None else _collective.failure_report(local_error)
222 prefix = f"SHMEM Heap failed during {stage}"
223 peers = _collective.exchange(
224 (value, report),
225 group=self.root_group,
226 failure_prefix=f"{prefix} while converging root_rank results",
227 )
228 failures = [item[1] for item in peers]
229 if any(failure is not None for failure in failures):
230 if len(peers) == 1 and isinstance(local_error, ValueError):
231 raise local_error
232 details = _collective.format_collective_failures(
233 failures,
234 rank_label="root_rank",
235 ranks_label="root_ranks",
236 )
237 raise RuntimeError(f"{prefix} on {details}") from local_error
238 if any(item[0] != value for item in peers):
239 raise RuntimeError(f"{prefix}: values differ across Root ranks")
241 def rollback_reservations(
242 self,
243 added_specifications: tuple[Mapping[str, Any], ...],
244 workspace: MegaMoeWorkspace | None,
245 ) -> None:
246 """Remove only uncommitted reservations from a collectively failed construction."""
247 with self._lock:
248 if workspace is not None:
249 for entry in self.entries:
250 if entry.workspace is workspace:
251 entry.workspace = None
252 entry.bound_spec = None
253 self.entries[:] = [
254 entry
255 for entry in self.entries
256 if entry.workspace is not None
257 or not any(entry.specification is specification for specification in added_specifications)
258 ]
260 def attach_runtime(self) -> bool:
261 """Acquire the manager's sole Runtime reference if it is not already attached."""
262 with self._lock:
263 if self._runtime_attached:
264 return False
265 if not self._configured or self.heap_bytes <= 0:
266 raise RuntimeError("SHMEM Heap must reserve a valid layout before attaching the Runtime")
267 shmem.acquire(self.root_group, heap_size_bytes=self.heap_bytes)
268 self._runtime_attached = True
269 return True
271 def detach_runtime(self) -> None:
272 """Release the sole Runtime reference after all workspaces are gone."""
273 with self._lock:
274 if not self._runtime_attached:
275 return
276 if self.workspaces:
277 raise RuntimeError("cannot detach SHMEM Runtime while MegaMoe workspaces remain bound")
278 shmem.release()
279 self._runtime_attached = False
281 def bind(
282 self,
283 workspace: MegaMoeWorkspace,
284 bound_spec: MegaMoeSpec,
285 specification: Mapping[str, Any],
286 ) -> int:
287 """Bind a workspace and its physical specification to one reservation."""
288 with self._lock:
289 for entry in self.entries:
290 if entry.specification is specification:
291 if entry.workspace is not None or entry.bound_spec is not None:
292 raise RuntimeError("SHMEM Heap owner is already bound")
293 entry.workspace = workspace
294 entry.bound_spec = bound_spec
295 return entry.layout.receive_capacity
296 raise RuntimeError("SHMEM Heap owner was not reserved")
298 def close_workspace(self, workspace: MegaMoeWorkspace) -> None:
299 """Close and remove one bound workspace without releasing the Runtime user."""
300 with self._lock:
301 workspace.close()
302 self.entries[:] = [entry for entry in self.entries if entry.workspace is not workspace]
304 def ensure_capacity(self, workspace: MegaMoeWorkspace, maximum_received_slots: int) -> None:
305 """Grow one workspace while preserving its committed source capacity.
307 Args:
308 workspace: Owner of the receive reservation.
309 maximum_received_slots: Largest current owner load in the EP group.
310 """
311 if maximum_received_slots <= workspace.capacity_floor:
312 return
313 with self.access():
314 entry = self._entry_for_workspace(workspace)
315 if entry.bound_spec is None:
316 raise RuntimeError("SHMEM Heap workspace has no bound MegaMoe specification")
317 capacity, minimum = _grow_receive_capacity(
318 entry.bound_spec,
319 entry.layout.receive_capacity,
320 maximum_received_slots,
321 )
322 layout = _workspace_layout(entry.specification, self.dtype, capacity, entry.layout.token_capacity)
323 layouts = [layout if item is entry else item.layout for item in self.entries]
324 if self.heap_limit is not None and self._required_bytes(layouts) > self.heap_limit:
325 layout = _workspace_layout(entry.specification, self.dtype, minimum, entry.layout.token_capacity)
326 layouts = [layout if item is entry else item.layout for item in self.entries]
327 self._restart(layouts, self._budget(self._required_bytes(layouts)))
329 def ensure_token_capacity(self, workspace: MegaMoeWorkspace, tokens: int, received: int) -> MegaMoeSpec:
330 """Grow source and destination storage in one transaction after route agreement.
332 Args:
333 workspace: Automatic execution resource receiving this route.
334 tokens: Group-maximum actual token rows, aligned for scheduling.
335 received: Group-maximum actual destination rows.
337 Returns:
338 Committed physical specification, independent of retained execution plans.
339 """
340 with self.access():
341 entry = self._entry_for_workspace(workspace)
342 spec = entry.bound_spec
343 if spec is None or not spec.auto_tokens:
344 raise RuntimeError("Automatic token growth requires an automatic workspace")
345 grow_source = tokens > spec.local_num_tokens
346 grow_receive = spec.dispatch_mode == "push" and received > entry.layout.receive_capacity
347 if not grow_source and not grow_receive:
348 return spec
349 factor = spec.capacity_growth_factor or _DEFAULT_CAPACITY_FACTOR
350 token_limit = ((2**31 - 1) // max(spec.hidden_size, spec.top_k) // 128) * 128
351 token_capacity = (_align_capacity(max(tokens, math.ceil(min(token_limit, spec.local_num_tokens * factor))))
352 if grow_source else spec.local_num_tokens)
353 receive_capacity = entry.layout.receive_capacity
354 if grow_receive:
355 receive_capacity = _align_capacity(max(received, math.ceil(min(2**31 - 1, receive_capacity * factor))))
356 layout = _workspace_layout(entry.specification, self.dtype, receive_capacity, token_capacity)
357 layouts = [layout if item is entry else item.layout for item in self.entries]
358 if self.heap_limit is not None and self._required_bytes(layouts) > self.heap_limit:
359 layout = _workspace_layout(entry.specification, self.dtype,
360 max(entry.layout.receive_capacity, _align_capacity(received)),
361 max(spec.local_num_tokens, tokens))
362 layouts = [layout if item is entry else item.layout for item in self.entries]
363 self._restart(layouts, self._budget(self._required_bytes(layouts)))
364 return entry.bound_spec
366 @staticmethod
367 def _layout_spec(entry: _HeapEntry, layout: _WorkspaceLayout) -> MegaMoeSpec:
368 """Project candidate capacities without modifying a committed specification."""
369 if layout.token_capacity is None:
370 return entry.bound_spec
371 return replace(entry.bound_spec, local_num_tokens=layout.token_capacity,
372 receive_capacity=layout.receive_capacity)
374 @staticmethod
375 def _required_bytes(layouts: list[_WorkspaceLayout]) -> int:
376 """Return the aggregate payload rounded once to physical Heap pages."""
377 return _round_heap_bytes(sum(layout.required_bytes for layout in layouts))
379 def _budget(self, required: int) -> int:
380 """Apply the explicit limit or preserve the grow-only physical Heap size."""
381 if self.heap_limit is not None and required > self.heap_limit:
382 raise RuntimeError(
383 f"SHMEM growth exceeds explicit Heap budget: required={required}, limit={self.heap_limit}"
384 )
385 return self.heap_limit if self.heap_limit is not None else max(self.heap_bytes, required)
387 def _configure(self, configured: str | None) -> None:
388 """Validate the explicit Heap limit and select the initial budget."""
389 required = self._required_bytes([entry.layout for entry in self.entries])
390 if configured is not None:
391 try:
392 self.heap_limit = int(configured)
393 except ValueError as error:
394 raise ValueError(
395 "HYPER_MULTICORE_SHMEM_HEAP_SIZE must be a positive integer number of bytes"
396 ) from error
397 if self.heap_limit <= 0:
398 raise ValueError("HYPER_MULTICORE_SHMEM_HEAP_SIZE must be positive")
399 if self.heap_limit < required:
400 raise RuntimeError(
401 f"HYPER_MULTICORE_SHMEM_HEAP_SIZE is too small: requires at least {required} bytes"
402 )
403 if self.heap_limit % _HEAP_GRANULARITY_BYTES:
404 raise ValueError(
405 "HYPER_MULTICORE_SHMEM_HEAP_SIZE must be a multiple of the 2 MiB physical page size"
406 )
407 self.heap_bytes = self._budget(required)
409 def _reserve(
410 self,
411 reservations: tuple[tuple[Mapping[str, Any], _WorkspaceLayout], ...],
412 ) -> tuple[Mapping[str, Any], ...]:
413 """Append specifications by identity without disturbing existing order."""
414 added = []
415 for specification, layout in reservations:
416 if not any(entry.specification is specification for entry in self.entries):
417 self.entries.append(_HeapEntry(specification, layout))
418 added.append(specification)
419 return tuple(added)
421 def _entry_for_workspace(self, workspace: MegaMoeWorkspace) -> _HeapEntry:
422 """Find one bound workspace by identity."""
423 for entry in self.entries:
424 if entry.workspace is workspace:
425 return entry
426 raise RuntimeError("SHMEM Heap owner was not bound")
428 def _describe_restart(
429 self,
430 layouts: list[_WorkspaceLayout],
431 ) -> tuple[
432 tuple[_lifecycle._AllocationReplacement, ...],
433 tuple[_RestartBinding, ...],
434 _RestartAgreement,
435 ]:
436 """Describe physical replacements and result slices without validating business state."""
437 replacements = []
438 bindings = []
439 agreements = []
440 for entry, layout in zip(self.entries, layouts):
441 is_bound = entry.workspace is not None
442 materialized = False
443 if is_bound:
444 start = len(replacements)
445 workspace_replacements = entry.workspace.describe_restart_allocations(
446 self._layout_spec(entry, layout),
447 self.dtype,
448 layout.receive_capacity,
449 )
450 replacements.extend(workspace_replacements)
451 materialized = bool(workspace_replacements)
452 if materialized:
453 bindings.append(_RestartBinding(entry, start, len(replacements)))
454 agreements.append(_RestartEntryAgreement(layout, is_bound, materialized))
455 return (
456 tuple(replacements),
457 tuple(bindings),
458 _RestartAgreement(self.epoch, tuple(agreements)),
459 )
461 def _restart_preflight(
462 self,
463 layouts: list[_WorkspaceLayout],
464 replacements: tuple[_lifecycle._AllocationReplacement, ...],
465 bindings: tuple[_RestartBinding, ...],
466 agreement: _RestartAgreement,
467 ) -> _RestartAgreement:
468 """Validate the frozen business snapshot and return its rank-independent agreement."""
469 if self.state != "ready":
470 raise RuntimeError("SHMEM Heap is not ready for restart")
471 if len(layouts) != len(self.entries):
472 raise RuntimeError("SHMEM Heap entries changed during restart preflight")
473 binding_by_entry = {id(binding.entry): binding for binding in bindings}
474 for entry, entry_agreement in zip(self.entries, agreement.entries):
475 if entry.workspace is None:
476 continue
477 if entry.bound_spec is None:
478 raise RuntimeError("SHMEM Heap workspace has no bound MegaMoe specification")
479 binding = binding_by_entry.get(id(entry))
480 workspace_replacements = () if binding is None else replacements[binding.start:binding.stop]
481 materialized = entry.workspace.validate_restart_allocations(
482 entry.bound_spec,
483 self.dtype,
484 entry.layout.receive_capacity,
485 workspace_replacements,
486 )
487 if materialized != entry_agreement.is_materialized:
488 raise RuntimeError("SHMEM Heap materialization state changed during preflight")
489 return agreement
491 def _validate_restart_result(
492 self,
493 result: _lifecycle._RestartResult,
494 bindings: tuple[_RestartBinding, ...],
495 layouts: list[_WorkspaceLayout],
496 ) -> None:
497 """Validate every result slice before assigning any workspace field."""
498 expected_count = sum(binding.stop - binding.start for binding in bindings)
499 if len(result.tensors) != expected_count:
500 raise RuntimeError("SHMEM restart returned an unexpected Allocation count")
501 layout_by_entry = {id(entry): layout for entry, layout in zip(self.entries, layouts)}
502 for binding in bindings:
503 entry = binding.entry
504 if entry.workspace is None or entry.bound_spec is None:
505 raise RuntimeError("SHMEM Heap workspace binding changed after restart")
506 entry.workspace.validate_restarted_allocations(
507 self._layout_spec(entry, layout_by_entry[id(entry)]),
508 self.dtype,
509 layout_by_entry[id(entry)].receive_capacity,
510 result.tensors[binding.start:binding.stop],
511 )
513 def _bind_restart_result(
514 self,
515 result: _lifecycle._RestartResult,
516 bindings: tuple[_RestartBinding, ...],
517 layouts: list[_WorkspaceLayout],
518 ) -> None:
519 """Bind validated result slices back to their original workspace objects."""
520 layout_by_entry = {id(entry): layout for entry, layout in zip(self.entries, layouts)}
521 for binding in bindings:
522 entry = binding.entry
523 entry.workspace.bind_restarted_allocations(
524 self._layout_spec(entry, layout_by_entry[id(entry)]),
525 self.dtype,
526 self.device,
527 layout_by_entry[id(entry)].receive_capacity,
528 result.tensors[binding.start:binding.stop],
529 )
531 def _restart(self, layouts: list[_WorkspaceLayout], heap_bytes: int) -> None:
532 """Restart SHMEM, validate and bind local results, then publish business state."""
533 started = time.perf_counter()
534 replacements, bindings, agreement = self._describe_restart(layouts)
535 old_bytes = self.heap_bytes
537 def raise_post_restart_failure(error: Exception, *, context: str) -> NoReturn:
538 """Invalidate the manager and report a post-restart business-state failure."""
539 self.state = "failed"
540 raise RuntimeError(
541 f"{context}: {error}. The restarted SHMEM Allocations cannot be rolled back; "
542 "terminate the process and let the launcher restart the job."
543 ) from error
545 try:
546 result = shmem._restart( # pylint: disable=protected-access
547 heap_size_bytes=heap_bytes,
548 replacements=replacements,
549 preflight_hook=lambda: self._restart_preflight(layouts, replacements, bindings, agreement),
550 )
551 except Exception:
552 if not _lifecycle._runtime_is_healthy(): # pylint: disable=protected-access
553 self.state = "failed"
554 raise
556 context = "MegaMoe restart result validation failed due to a local implementation error"
557 try:
558 self._validate_restart_result(result, bindings, layouts)
559 context = "MegaMoe workspace binding failed after SHMEM restart due to a local implementation error"
560 self._bind_restart_result(result, bindings, layouts)
561 except Exception as error: # pylint: disable=broad-exception-caught
562 raise_post_restart_failure(error, context=context)
564 stages_ms = dict(result.stages_ms)
565 for entry, layout in zip(self.entries, layouts):
566 if entry.bound_spec is not None:
567 entry.bound_spec = self._layout_spec(entry, layout)
568 entry.layout = layout
569 if entry.workspace is not None:
570 entry.workspace.capacity_floor = layout.receive_capacity
571 self.heap_bytes = heap_bytes
572 self.epoch += 1
573 self.growth_records.append({
574 "epoch": self.epoch,
575 "old_heap_bytes": old_bytes,
576 "heap_bytes": heap_bytes,
577 "stages_ms": stages_ms,
578 "total_ms": (time.perf_counter() - started) * 1000,
579 "capacities": [entry.layout.receive_capacity for entry in self.entries],
580 })
583def _get_heap_manager(group: Any, tensor: torch.Tensor) -> _HeapManager:
584 """Return the compatible process MegaMoe Heap Manager."""
585 global _HEAP_MANAGER # pylint: disable=global-statement
586 if (_HEAP_MANAGER is None
587 or not _HEAP_MANAGER.workspaces and not _HEAP_MANAGER.runtime_attached):
588 _HEAP_MANAGER = _HeapManager(group, tensor)
589 return _HEAP_MANAGER
590 if _HEAP_MANAGER.members != _lifecycle.root_members(group) or _HEAP_MANAGER.device != tensor.device:
591 raise RuntimeError("SHMEM Heap requires one device and ordered Root per process")
592 if _HEAP_MANAGER.dtype != tensor.dtype or _HEAP_MANAGER.element_size != tensor.element_size():
593 raise ValueError("SHMEM Heap owners must use the same dtype and element size")
594 return _HEAP_MANAGER