Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_multicore/shmem/

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

7# http://www.apache.org/licenses/LICENSE-2.0

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"""Reference-counted Torch lifecycle for the process-wide SHMEM Runtime."""

17from __future__ import annotations

19__all__ = [

20 "acquire",

21 "release",

22]

24import threading

25import time

26from collections.abc import Callable

27from dataclasses import dataclass

28from functools import wraps

29from typing import TYPE_CHECKING, Any, NoReturn, Optional

31from .. import _collective

32from ._runtime import _load_native, _torch_modules

34if TYPE_CHECKING:

35 import torch

36 from torch.distributed import ProcessGroup

39_lock = threading.RLock()

40_users = 0

41_root_group: Any | None = None

42_root_uses_distributed: bool | None = None

43_root_size: int | None = None

44_root_ranks: tuple[int, ...] | None = None

45_device_index: int | None = None

46# Blocks SHMEM access after physical, device, or Root convergence failures.

47_runtime_invalidated = False

48_NATIVE_ERROR_PREFIX = "[HYPER_MULTICORE_SHMEM_ERROR]"

49_NATIVE_MESSAGE_FIELD = " message="

50_CANN_DETAIL_FIELD = ", cann_detail=["

51_RETRYABLE_RUNTIME_HINT = "The active SHMEM Runtime remains usable."

52_TERMINAL_RUNTIME_HINT = "The SHMEM Runtime is no longer usable; restart the process."

53_CONVERGENCE_FAILURE_HINT = (

54 "The SHMEM Runtime can no longer be managed safely after Root convergence failed; restart the process."

55)

56_INVALIDATED_RUNTIME_MESSAGE = (

57 "SHMEM Runtime is no longer usable after a shutdown or restart failure; restart the process"

58)

61@dataclass(frozen=True)

62class _AllocationSpec:

63 """Physical properties of one symmetric Allocation after restart."""

65 shape: tuple[int, ...]

66 dtype: torch.dtype

67 alignment_bytes: int

70@dataclass(frozen=True)

71class _AllocationReplacement:

72 """Replace one complete active Allocation with a target specification."""

74 current: torch.Tensor

75 target: _AllocationSpec

78@dataclass(frozen=True)

79class _RestartResult:

80 """Symmetric Tensors and stage timings produced by a successful restart."""

82 tensors: tuple[torch.Tensor, ...]

83 stages_ms: dict[str, float]

86class _RuntimeUnavailableError(RuntimeError):

87 """Mark a stage failure after which the active Runtime cannot be used safely."""

90def _exception_summary(error: Exception) -> str:

91 """Extract display text without changing the exception or its cause."""

92 message = str(error)

93 if message.startswith(_NATIVE_ERROR_PREFIX):

94 _, separator, native_message = message.partition(_NATIVE_MESSAGE_FIELD)

95 if separator:

96 message = native_message

97 message, _, _ = message.partition(_CANN_DETAIL_FIELD)

98 return message

100

101def _raise_lifecycle_failure(message: str, error: Exception, *, hint: str) -> NoReturn:

102 """Raise a lifecycle error with recovery guidance and the available local cause."""

103 message = f"{message.rstrip('.')}. {hint}"

104 if error.__cause__ is not None:

105 raise RuntimeError(message) from error.__cause__

106 if error.__suppress_context__:

107 raise RuntimeError(message) from None

108 raise RuntimeError(message) from error

109

110

111def _reference_count() -> int:

112 """Return the process-local count of acquired SHMEM Runtime references."""

113 with _lock:

114 return _users

115

116

117def _runtime_is_healthy() -> bool:

118 """Return whether the process Runtime accepts SHMEM calls."""

119 with _lock:

120 return _users > 0 and not _runtime_invalidated

121

122

123def _host_barrier() -> None:

124 """Complete a host-synchronous HCCL barrier over the acquired Root group."""

125 if _users <= 0:

126 raise RuntimeError("SHMEM Runtime has no active reference to barrier")

127 if _root_uses_distributed and _root_size > 1:

128 _, dist = _torch_modules()

129 dist.barrier(group=_root_group)

130

131

132def _describe_root(root_group: Any, dist: Any) -> tuple[int, tuple[int, ...]]:

133 """Validate group membership and return group-local CANN coordinates."""

134 world_size = int(dist.get_world_size())

135 root_rank = int(dist.get_rank(root_group))

136 if root_rank < 0:

137 raise RuntimeError(f"current process must belong to the SHMEM Root Group: root_rank={root_rank}")

138 global_ranks = tuple(int(rank) for rank in dist.get_process_group_ranks(root_group))

139 if (not global_ranks or len(set(global_ranks)) != len(global_ranks)

140 or any(rank < 0 or rank >= world_size for rank in global_ranks)):

141 raise RuntimeError(

142 f"SHMEM Root Group has invalid ordered membership: {global_ranks}, world_size={world_size}"

143 )

144 if root_rank >= len(global_ranks):

145 raise RuntimeError(

146 f"current process must belong to the SHMEM Root Group: root_rank={root_rank}, ranks={global_ranks}"

147 )

148 return root_rank, global_ranks

149

150

151def root_members(group: ProcessGroup | None) -> tuple[int, ...]:

152 """Return ordered SHMEM membership without relying on ProcessGroup identity.

153

154 Args:

155 group: Root process group, or None for the default root.

156 """

157 _, dist = _torch_modules()

158 return _describe_root(dist.group.WORLD if group is None else group, dist)[1] if dist.is_initialized() else (0,)

159

160

161def _subgroup_unique_id(

162 native: Any,

163 dist: Any,

164 group: Any,

165 rank: int,

166 ranks: tuple[int, ...],

167 *,

168 failure_prefix: str = "SHMEM bootstrap failed",

169) -> bytes:

170 """Broadcast an independent CANN bootstrap ID without involving other PP stages."""

171 payload = [None, None]

172 root_error = None

173 if rank == 0:

174 # Relay every vendor bootstrap failure so peers cannot wait indefinitely for a root ID.

175 try:

176 payload[0] = native._get_unique_id() # pylint: disable=protected-access

177 except Exception as error: # pylint: disable=broad-exception-caught

178 root_error = error

179 payload[1] = _collective.failure_report(error, message=_exception_summary(error))

180 try:

181 dist.broadcast_object_list(payload, src=ranks[0], group=group)

182 except Exception as error:

183 raise RuntimeError(

184 f"{failure_prefix} on root_rank={rank}: broadcast failed: {_exception_summary(error)}"

185 ) from error

186 if payload[1] is not None:

187 details = _collective.format_collective_failures(

188 [payload[1]],

189 rank_label="root_rank",

190 ranks_label="root_ranks",

191 )

192 raise RuntimeError(f"{failure_prefix} on {details}") from root_error

193 if not isinstance(payload[0], bytes) or not payload[0]:

194 raise RuntimeError(f"{failure_prefix} on root_rank=0: returned an invalid unique ID")

195 return payload[0]

196

197

198def _validate_active_root(selected_group: Any, dist: Any, heap_size_bytes: int | None) -> None:

199 """Validate another owner against the initialized root."""

200 if heap_size_bytes is not None:

201 config = _load_native()._debug_state()["config"] # pylint: disable=protected-access

202 if config["heap_size_bytes"] < heap_size_bytes:

203 raise RuntimeError("active SHMEM heap is smaller than the requested managed layout")

204 if dist.is_initialized() != _root_uses_distributed:

205 raise RuntimeError("torch.distributed state changed during the active SHMEM Runtime lifecycle")

206 if dist.is_initialized() and selected_group is not _root_group:

207 _, selected_ranks = _describe_root(selected_group, dist)

208 if selected_ranks != _root_ranks:

209 raise RuntimeError(

210 "SHMEM Root Group has different ordered membership from the active Runtime: "

211 f"active={_root_ranks}, requested={selected_ranks}"

212 )

213

214

215def acquire(root_group: ProcessGroup | None = None, *, heap_size_bytes: int | None = None) -> None:

216 """Acquire one reference to the Root-only process SHMEM Runtime.

217

218 The first reference initializes the Native Runtime. Later references with the same ordered group membership

219 share it. Every successful call must be paired with one :func:`release` call after the consumer has stopped using

220 SHMEM and freed its Allocations.

221

222 Args:

223 heap_size_bytes: Optional explicit initialization size, overriding the environment for this lifecycle.

224 root_group: The EP ProcessGroup, which may be a subgroup of Torch WORLD. ``None`` selects WORLD when

225 Torch distributed is initialized, or a one-PE Root otherwise. Disjoint groups initialize independently;

226 one process cannot participate in two different active SHMEM roots simultaneously.

227

228 Raises:

229 RuntimeError: If the Root identity or framework state differs from the active lifecycle, or the Runtime was

230 invalidated by a shutdown or restart failure.

231 """

232 global _users, _root_group, _root_uses_distributed, _root_size # pylint: disable=global-statement

233 global _root_ranks, _device_index # pylint: disable=global-statement

234

235 with _lock:

236 if heap_size_bytes is not None and (isinstance(heap_size_bytes, bool)

237 or not isinstance(heap_size_bytes, int) or heap_size_bytes <= 0):

238 raise ValueError("heap_size_bytes must be a positive integer")

239 if _runtime_invalidated:

240 raise RuntimeError(_INVALIDATED_RUNTIME_MESSAGE)

241

242 _, dist = _torch_modules()

243 uses_distributed = bool(dist.is_initialized())

244 if uses_distributed:

245 selected_group = dist.group.WORLD if root_group is None else root_group

246 else:

247 if root_group is not None:

248 raise RuntimeError("a SHMEM Root Group cannot be selected before torch.distributed initialization")

249 selected_group = None

250

251 if _users > 0:

252 _validate_active_root(selected_group, dist, heap_size_bytes)

253 _users += 1

254 return

255

256 if uses_distributed:

257 root_rank, root_ranks = _describe_root(selected_group, dist)

258 root_size = len(root_ranks)

259 dist.barrier(group=selected_group)

260 else:

261 root_rank, root_size = 0, 1

262 root_ranks = None

263

264 # Native initialization failure is retryable and must not create a consumer reference.

265 native = _load_native()

266 if heap_size_bytes is not None:

267 unique_id = (_subgroup_unique_id(native, dist, selected_group, root_rank, root_ranks)

268 if uses_distributed else b"")

269 native._initialize(root_rank, root_size, unique_id, heap_size_bytes) # pylint: disable=protected-access

270 elif uses_distributed and root_ranks != tuple(range(dist.get_world_size())):

271 unique_id = _subgroup_unique_id(native, dist, selected_group, root_rank, root_ranks)

272 native._initialize(root_rank, root_size, unique_id) # pylint: disable=protected-access

273 else:

274 native._initialize(root_rank, root_size) # pylint: disable=protected-access

275 _root_group = selected_group

276 _root_uses_distributed = uses_distributed

277 _root_size = root_size

278 _root_ranks = root_ranks

279 native_state = native._debug_state() # pylint: disable=protected-access

280 _device_index = int(native_state["device_index"])

281 _users = 1

282

283

284def _validate_shutdown_preconditions() -> None:

285 """Validate local shutdown state without modifying the active Runtime."""

286 _, dist = _torch_modules()

287 if bool(dist.is_initialized()) != _root_uses_distributed:

288 raise _RuntimeUnavailableError(

289 "torch.distributed state changed during the active SHMEM Runtime lifecycle"

290 )

291 _load_native()._validate_shutdown() # pylint: disable=protected-access

292

293

294def _synchronize_device(device: Any = None) -> None:

295 """Synchronize NPU work and classify a device failure as Runtime-unavailable."""

296 torch, _ = _torch_modules()

297 try:

298 if device is None:

299 torch.npu.synchronize()

300 else:

301 torch.npu.synchronize(device)

302 except Exception as error: # pylint: disable=broad-exception-caught

303 raise _RuntimeUnavailableError(_exception_summary(error)) from error

304

305

306def _runtime_access(function):

307 """Exclude SHMEM API submissions during a managed Runtime restart."""

308 @wraps(function)

309 def guarded(*args: Any, **kwargs: Any) -> Any:

310 """Serialize one API submission against managed heap teardown."""

311 with _lock:

312 if _runtime_invalidated:

313 raise RuntimeError(_INVALIDATED_RUNTIME_MESSAGE)

314 return function(*args, **kwargs)

315 return guarded

316

317

318_StagePayload = tuple[Any, Optional[_collective.FailureReport], bool]

319

320

321def _execute_local_stage(operation: Callable[[], Any]) -> tuple[_StagePayload, Exception | None]:

322 """Execute one local stage and preserve its non-serializable exception cause."""

323 try:

324 return (operation(), None, False), None

325 except Exception as exception: # pylint: disable=broad-exception-caught

326 runtime_unavailable = isinstance(exception, _RuntimeUnavailableError)

327 diagnostic = exception.__cause__ if runtime_unavailable and exception.__cause__ is not None else exception

328 local_error = exception.__cause__ if runtime_unavailable else exception

329 if runtime_unavailable and diagnostic is exception:

330 report = "RuntimeError", _exception_summary(diagnostic)

331 else:

332 report = _collective.failure_report(diagnostic, message=_exception_summary(diagnostic))

333 return (None, report, runtime_unavailable), local_error

334

335

336def _evaluate_stage_results(peers: list[_StagePayload], local_value: Any, local_error: Exception | None,

337 action: str, stage: str) -> None:

338 """Raise a local diagnostic for peer failures or inconsistent stage values."""

339 if not any(item[1] is not None for item in peers):

340 if any(item[0] != local_value for item in peers):

341 raise RuntimeError(f"{action} failed during {stage}: values differ across Root ranks")

342 return

343 runtime_unavailable = any(item[2] for item in peers)

344 details = _collective.format_collective_failures(

345 [item[1] for item in peers],

346 rank_label="root_rank",

347 ranks_label="root_ranks",

348 )

349 error_type = _RuntimeUnavailableError if runtime_unavailable else RuntimeError

350 raise error_type(f"{action} failed during {stage} on {details}") from local_error

351

352

353def _reconfiguration_stage(

354 operation: Callable[[], Any],

355 *,

356 action: str,

357 stage: str,

358 root_group: Any,

359) -> float:

360 """Execute a stage and agree on its returned manifest; non-manifest stages return None."""

361 started = time.perf_counter()

362 payload, local_error = _execute_local_stage(operation)

363 peers = _collective.exchange(

364 payload,

365 group=root_group,

366 failure_prefix=f"{action} failed during {stage} while converging Root rank results",

367 )

368 _evaluate_stage_results(peers, payload[0], local_error, action, stage)

369 return (time.perf_counter() - started) * 1000

370

371

372def _validate_allocation_spec(specification: _AllocationSpec) -> tuple[tuple[int, ...], str, int, int]:

373 """Validate one target Allocation and return its rank-independent manifest."""

374 if not isinstance(specification.shape, tuple) or not specification.shape:

375 raise ValueError("Allocation shape must be a non-empty tuple")

376 if any(isinstance(dimension, bool) or not isinstance(dimension, int) or dimension <= 0

377 for dimension in specification.shape):

378 raise ValueError(f"Allocation shape must contain positive integers, got {specification.shape}")

379 alignment = specification.alignment_bytes

380 if (isinstance(alignment, bool) or not isinstance(alignment, int)

381 or alignment <= 0 or alignment & (alignment - 1)):

382 raise ValueError(f"Allocation alignment_bytes must be a positive power of two, got {alignment!r}")

383 try:

384 payload_bytes = specification.dtype.itemsize

385 for dimension in specification.shape:

386 payload_bytes *= dimension

387 except (AttributeError, TypeError) as error:

388 raise TypeError(f"Allocation dtype must be a torch dtype, got {specification.dtype!r}") from error

389 return specification.shape, str(specification.dtype), alignment, payload_bytes

390

391

392def _validate_replacements(

393 replacements: tuple[_AllocationReplacement, ...],

394 native_state: dict[str, Any],

395 heap_size_bytes: int,

396) -> tuple[tuple[Any, ...], tuple[Any, ...]]:

397 """Validate complete active Allocations and return current and target manifests."""

398 if not isinstance(replacements, tuple):

399 raise TypeError("replacements must be a tuple")

400 active_records = {

401 int(record["allocation_base"]): int(record["allocation_bytes"])

402 for record in native_state["active_allocations"]

403 }

404 current_manifest = []

405 target_manifest = []

406 current_bases = set()

407 target_footprint = 0

408 for index, replacement in enumerate(replacements):

409 if not isinstance(replacement, _AllocationReplacement):

410 raise TypeError(f"replacements[{index}] must be an _AllocationReplacement")

411 current = replacement.current

412 if current is None:

413 raise ValueError(f"replacements[{index}].current must be an active Tensor")

414 if getattr(current, "_base", None) is not None:

415 raise ValueError(f"replacements[{index}].current must be a complete Allocation, not a Tensor view")

416 try:

417 allocation_base = int(current.data_ptr())

418 allocation_bytes = int(current.numel()) * int(current.element_size())

419 current_shape = tuple(int(dimension) for dimension in current.shape)

420 current_dtype = str(current.dtype)

421 contiguous = bool(current.is_contiguous())

422 except (AttributeError, TypeError, ValueError) as error:

423 raise TypeError(f"replacements[{index}].current must be a Tensor") from error

424 if not contiguous or active_records.get(allocation_base) != allocation_bytes:

425 raise ValueError(f"replacements[{index}].current must be one complete active Allocation")

426 if allocation_base in current_bases:

427 raise ValueError("replacements must not contain duplicate current Allocations")

428 current_bases.add(allocation_base)

429 current_manifest.append((current_shape, current_dtype, allocation_bytes))

430 target_shape, target_dtype, alignment, payload_bytes = _validate_allocation_spec(replacement.target)

431 target_manifest.append((target_shape, target_dtype, alignment, payload_bytes))

432 target_footprint += payload_bytes + alignment - 1

433 if current_bases != set(active_records):

434 raise RuntimeError("replacements must cover every active SHMEM Allocation exactly once")

435 if target_footprint > heap_size_bytes:

436 raise RuntimeError(

437 "replacement Allocations exceed the target SHMEM heap: "

438 f"required={target_footprint}, heap_size_bytes={heap_size_bytes}"

439 )

440 return tuple(current_manifest), tuple(target_manifest)

441

442

443def _restart(

444 *,

445 heap_size_bytes: int,

446 replacements: tuple[_AllocationReplacement, ...],

447 preflight_hook: Callable[[], object],

448) -> _RestartResult:

449 """Collectively rebuild the sole user's symmetric Allocations.

450

451 Args:

452 heap_size_bytes: Positive byte capacity of the restarted symmetric Heap.

453 replacements: Complete active Allocations and their target physical specifications.

454 preflight_hook: Read-only business-state snapshot executed once while the lifecycle lock is held.

455

456 Returns:

457 Replacement Tensors in request order and generic restart stage timings.

458

459 Raises:

460 RuntimeError: If the Runtime is unhealthy, another user exists, rank state differs, or restart fails.

461

462 Note:

463 Ordinary preflight validation failures leave a healthy Runtime usable only when Root results converge.

464 Existing Runtime invalidation, state drift, device synchronization errors and convergence failures require

465 a process restart even before Allocation release. Any failure after release begins invalidates the Runtime.

466 All Root ranks must call this function in the same order.

467 """

468 global _runtime_invalidated # pylint: disable=global-statement

469 with _lock:

470 native = _load_native()

471 torch, dist = _torch_modules()

472 destructive_started = False

473 current_stage = "Preflight"

474

475 def run_stage(stage: str, operation: Callable[[], Any]) -> float:

476 """Run one converged restart stage under a stable diagnostic label."""

477 nonlocal current_stage

478 current_stage = stage

479 return _reconfiguration_stage(

480 operation,

481 action="SHMEM restart",

482 stage=stage,

483 root_group=_root_group,

484 )

485

486 def preflight() -> tuple[Any, ...]:

487 """Capture and validate the complete local pre-destructive state."""

488 if _runtime_invalidated:

489 raise _RuntimeUnavailableError(_INVALIDATED_RUNTIME_MESSAGE)

490 if _users != 1:

491 raise RuntimeError("SHMEM restart requires one healthy Runtime user")

492 if (isinstance(heap_size_bytes, bool)

493 or not isinstance(heap_size_bytes, int) or heap_size_bytes <= 0):

494 raise ValueError("heap_size_bytes must be a positive integer")

495 if bool(dist.is_initialized()) != _root_uses_distributed:

496 raise _RuntimeUnavailableError(

497 "torch.distributed state changed during the active SHMEM Runtime lifecycle"

498 )

499 try:

500 native_state = native._debug_state() # pylint: disable=protected-access

501 except Exception as error: # pylint: disable=broad-exception-caught

502 raise _RuntimeUnavailableError(_exception_summary(error)) from error

503 if native_state.get("state") != "Ready":

504 raise _RuntimeUnavailableError(

505 f"SHMEM Runtime state changed during the active lifecycle: state={native_state.get('state')!r}"

506 )

507 if _device_index is None or native_state.get("device_index") != _device_index:

508 raise _RuntimeUnavailableError("SHMEM Runtime device changed during the active lifecycle")

509 if int(torch.npu.current_device()) != _device_index:

510 raise RuntimeError("current NPU device differs from the SHMEM Runtime device")

511 if torch.npu.is_current_stream_capturing():

512 raise RuntimeError("SHMEM restart is unavailable during graph capture")

513 business_state = preflight_hook()

514 current_manifest, target_manifest = _validate_replacements(

515 replacements,

516 native_state,

517 heap_size_bytes,

518 )

519 return heap_size_bytes, business_state, current_manifest, target_manifest

520

521 try:

522 stages_ms = {

523 "preflight_ms": run_stage("Preflight", preflight),

524 }

525 device = torch.device("npu", _device_index)

526 stages_ms["quiesce_ms"] = run_stage("Quiesce", lambda: _synchronize_device(device))

527

528 # _api imports lifecycle guards, so the local import avoids a cycle while reusing Allocation validation

529 # and diagnostics.

530 from . import _api # pylint: disable=import-outside-toplevel

531

532 replacement_tensors = []

533

534 def release_allocations() -> None:

535 """Release every declared Allocation and verify an empty registry."""

536 for replacement in replacements:

537 _api.free(replacement.current)

538 if native._debug_state()["active_allocations"]: # pylint: disable=protected-access

539 raise RuntimeError("SHMEM restart release left active Allocations")

540

541 destructive_started = True

542 stages_ms["release_ms"] = run_stage(

543 "ReleaseAllocations",

544 release_allocations,

545 )

546 validate_ms = run_stage(

547 "ValidateShutdown",

548 native._validate_shutdown, # pylint: disable=protected-access

549 )

550 shutdown_ms = run_stage(

551 "Shutdown",

552 native._shutdown, # pylint: disable=protected-access

553 )

554 stages_ms["finalize_ms"] = validate_ms + shutdown_ms

555 current_stage = "Bootstrap"

556 started = time.perf_counter()

557 rank = int(dist.get_rank(_root_group)) if _root_uses_distributed else 0

558 unique_id = (_subgroup_unique_id(

559 native,

560 dist,

561 _root_group,

562 rank,

563 _root_ranks,

564 failure_prefix="SHMEM restart failed during Bootstrap",

565 )

566 if _root_uses_distributed else b"")

567 stages_ms["bootstrap_ms"] = (time.perf_counter() - started) * 1000

568

569 def initialize() -> None:

570 """Initialize the target Heap and verify the frozen device."""

571 native._initialize( # pylint: disable=protected-access

572 rank,

573 _root_size,

574 unique_id,

575 heap_size_bytes,

576 )

577 native_state = native._debug_state() # pylint: disable=protected-access

578 if int(native_state["device_index"]) != _device_index:

579 raise RuntimeError("SHMEM Runtime restarted on a different NPU device")

580

581 stages_ms["initialize_ms"] = run_stage(

582 "Initialize",

583 initialize,

584 )

585

586 def restore_allocations() -> None:

587 """Create replacement Tensors in their declared order."""

588 replacement_tensors.extend(

589 _api.empty(

590 replacement.target.shape,

591 dtype=replacement.target.dtype,

592 alignment=replacement.target.alignment_bytes,

593 )

594 for replacement in replacements

595 )

596

597 stages_ms["restore_ms"] = run_stage(

598 "RestoreAllocations",

599 restore_allocations,

600 )

601 except Exception as error: # pylint: disable=broad-exception-caught

602 convergence_failed = isinstance(error, _collective.CollectiveConvergenceError)

603 runtime_unavailable = isinstance(error, _RuntimeUnavailableError)

604 if destructive_started or convergence_failed or runtime_unavailable:

605 _runtime_invalidated = True

606 prefix = f"SHMEM restart failed during {current_stage}"

607 message = str(error)

608 if not message.startswith(prefix):

609 message = f"{prefix}: {_exception_summary(error)}"

610 if convergence_failed:

611 hint = _CONVERGENCE_FAILURE_HINT

612 elif _runtime_invalidated:

613 hint = _TERMINAL_RUNTIME_HINT

614 else:

615 hint = _RETRYABLE_RUNTIME_HINT

616 _raise_lifecycle_failure(message, error, hint=hint)

617 return _RestartResult(tuple(replacement_tensors), stages_ms)

618

619

620def release() -> None:

621 """Release one SHMEM Runtime reference and shut down the last reference.

622

623 The last consumer must free every symmetric Allocation and stop all SHMEM work before calling this function. All

624 Root ranks must release their final reference in the same collective order. Retryable validation failures retain

625 the final reference so the caller can correct the condition. Root state changes, NPU synchronization failures and

626 failures to converge Root results invalidate lifecycle management and require a process restart.

627

628 Raises:

629 RuntimeError: If no reference is active, framework state changed, shutdown preconditions fail, or Native

630 shutdown fails.

631 """

632 global _users, _runtime_invalidated, _root_group, _root_uses_distributed # pylint: disable=global-statement

633 global _root_size # pylint: disable=global-statement

634 global _root_ranks, _device_index # pylint: disable=global-statement

635

636 with _lock:

637 if _runtime_invalidated:

638 raise RuntimeError(_INVALIDATED_RUNTIME_MESSAGE)

639 if _users <= 0:

640 raise RuntimeError("SHMEM Runtime has no active reference to release")

641 if _users > 1:

642 _users -= 1

643 return

644

645 try:

646 _reconfiguration_stage(

647 _validate_shutdown_preconditions,

648 action="SHMEM release",

649 stage="Preflight",

650 root_group=_root_group,

651 )

652 _reconfiguration_stage(

653 _synchronize_device,

654 action="SHMEM release",

655 stage="Quiesce",

656 root_group=_root_group,

657 )

658 except RuntimeError as error:

659 convergence_failed = isinstance(error, _collective.CollectiveConvergenceError)

660 runtime_unavailable = isinstance(error, _RuntimeUnavailableError)

661 if convergence_failed or runtime_unavailable:

662 _runtime_invalidated = True

663 hint = (_CONVERGENCE_FAILURE_HINT if convergence_failed else

664 _TERMINAL_RUNTIME_HINT if runtime_unavailable else _RETRYABLE_RUNTIME_HINT)

665 _raise_lifecycle_failure(str(error), error, hint=hint)

666 _runtime_invalidated = True

667 try:

668 # All peers become terminal if any Native shutdown fails; none may reinitialize alone.

669 _reconfiguration_stage(

670 _load_native()._shutdown, # pylint: disable=protected-access

671 action="SHMEM release",

672 stage="Shutdown",

673 root_group=_root_group,

674 )

675 except RuntimeError as error:

676 _raise_lifecycle_failure(

677 str(error),

678 error,

679 hint=_TERMINAL_RUNTIME_HINT,

680 )

681 finally:

682 _users = 0

683 _root_group = None

684 _root_uses_distributed = None

685 _root_size = None

686 _root_ranks = None

687 _device_index = None

688 _runtime_invalidated = False

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / shmem / _lifecycle.py: 89%

359 statements