Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/core/activation

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"""Eager non-reentrant activation checkpointing for the Torch backend.

17The saved-tensor hook algorithm is adapted from ``torch.utils.checkpoint`` in

18PyTorch release/2.9. Hyper owns the scheduling and session extensions here so

19the implementation can run consistently on PyTorch 2.6, 2.7, and 2.9 without

20patching the installed framework.

21"""

22import contextlib

23import contextvars

24import threading

25import uuid

26import warnings

27import weakref

28from collections import defaultdict

29from dataclasses import dataclass, field

30from functools import partial

31from typing import Any, Callable, DefaultDict, Dict, Generator, Iterator, List, Optional, Tuple

33import torch

34from torch.utils._pytree import tree_map

35from torch.utils.checkpoint import DefaultDeviceType

36from torch.utils.checkpoint import checkpoint as torch_checkpoint

37from torch.utils.checkpoint import set_checkpoint_early_stop

40_DEFAULT_DETERMINISM_MODE = "default"

41_RECOMPUTE_COLLECTOR = contextvars.ContextVar("hyper_recompute_collector", default=None)

42_RECOMPUTE_SESSION = contextvars.ContextVar("hyper_recompute_session", default=None)

43_SESSION_FRAMES: DefaultDict[Any, weakref.WeakSet] = defaultdict(weakref.WeakSet)

44_SESSION_FRAMES_LOCK = threading.RLock()

47class CheckpointError(RuntimeError):

48 """Raised when checkpoint forward and recomputation are inconsistent."""

51class _Handle:

52 """Identity key for one recomputed saved tensor."""

55class _Holder:

56 """Saved-tensor placeholder containing handles keyed by recompute session."""

58 def __init__(self) -> None:

59 """Initialize an empty per-session handle mapping."""

60 self.handles: Dict[Any, Optional[_Handle]] = {}

63class _StopRecomputationError(Exception):

64 """Internal control-flow exception used by early-stop recomputation."""

67class _SessionActivation:

68 """Control-plane state shared by checkpoint frames in one session scope."""

70 def __init__(self, session_id: Any, retain_on_unpack: bool) -> None:

71 """Initialize one scoped session activation."""

72 self.session_id = session_id

73 self.retain_on_unpack = retain_on_unpack

74 self.frames: weakref.WeakSet = weakref.WeakSet()

77class _NoopSaveInputs(torch.autograd.Function):

78 """Save checkpoint inputs without adding a meaningful forward operation."""

80 @staticmethod

81 def forward(*args: Any) -> Any:

82 """Return a dummy output whose grad function retains checkpoint inputs."""

83 del args

84 return torch.empty((0,))

86 @staticmethod

87 def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:

88 """Save tensor inputs while retaining non-tensor input structure."""

89 del output

90 tensor_pairs = [(index, value) for index, value in enumerate(inputs) if isinstance(value, torch.Tensor)]

91 tensor_indices, tensors = zip(*tensor_pairs)

92 index_to_saved_index = {input_index: saved_index for saved_index, input_index in enumerate(tensor_indices)}

93 stored_args = [None if isinstance(value, torch.Tensor) else value for value in inputs]

95 def get_args(saved_tensors: Tuple[Any, ...]) -> List[Any]:

96 """Reconstruct the original checkpoint arguments."""

97 restored_args = [

98 saved_tensors[index_to_saved_index[index]] if index in tensor_indices else value

99 for index, value in enumerate(stored_args)

100 ]

101 return restored_args[1:]

102

103 ctx.get_args = get_args

104 ctx.save_for_backward(*tensors)

105

106 @staticmethod

107 def backward(ctx: Any, *grad_outputs: Any) -> None:

108 """Reject direct backward through the internal input saver."""

109 del ctx, grad_outputs

110 raise CheckpointError("The internal checkpoint input saver must not be backwarded directly.")

111

112

113class _CheckpointFrame:

114 """State shared by one checkpoint forward and its recomputations."""

115

116 def __init__(self, recompute_fn: Callable, early_stop: bool, metadata_fn: Optional[Callable]) -> None:

117 """Initialize frame state captured by the saved-tensor hooks."""

118 self.recompute_fn = recompute_fn

119 self.input_saver = None

120 self.weak_holders: List[weakref.ReferenceType] = []

121 self.recomputed: DefaultDict[Any, weakref.WeakKeyDictionary] = defaultdict(weakref.WeakKeyDictionary)

122 self.recomp_counter: DefaultDict[Any, int] = defaultdict(int)

123 self.is_recomputed: DefaultDict[Any, bool] = defaultdict(bool)

124 self.early_stop = early_stop

125 self.metadata_fn = metadata_fn

126 self.x_metadatas: List[Any] = []

127 self.forward_completed = False

128 self.ignore_saved_mismatch = False

129 self.active_session: Optional[_SessionActivation] = None

130

131 def check_recomputed_tensors_match(self, session_id: Any) -> None:

132 """Validate saved-tensor count and metadata after recomputation."""

133 if self.ignore_saved_mismatch:

134 return

135 if len(self.weak_holders) != self.recomp_counter[session_id]:

136 raise CheckpointError(

137 "Hyper checkpoint saved a different number of tensors during forward and recomputation. "

138 f"Forward saved {len(self.weak_holders)} tensors, but recomputation saved "

139 f"{self.recomp_counter[session_id]} tensors."

140 )

141

142 mismatches = []

143 for index, weak_holder in enumerate(self.weak_holders):

144 holder = weak_holder()

145 if holder is None:

146 continue

147 handle = holder.handles.get(session_id)

148 _internal_assert(handle is not None, "Missing recomputed tensor handle during metadata validation.")

149 _internal_assert(

150 handle in self.recomputed[session_id],

151 "Missing recomputed tensor during metadata validation.",

152 )

153 recomputed_tensor = self.recomputed[session_id][handle]

154 recomputed_metadata = self.metadata_fn(recomputed_tensor)

155 if self.x_metadatas[index] != recomputed_metadata:

156 mismatches.append((index, self.x_metadatas[index], recomputed_metadata))

157

158 if mismatches:

159 details = "\n".join(

160 f"tensor {index}: forward={forward_metadata}, recompute={recomputed_metadata}"

161 for index, forward_metadata, recomputed_metadata in mismatches

162 )

163 raise CheckpointError(

164 "Hyper checkpoint detected different tensor metadata during recomputation:\n" + details

165 )

166

167 def clear_session(self, session_id: Any) -> None:

168 """Release all tensors and handles associated with one session."""

169 for weak_holder in self.weak_holders:

170 holder = weak_holder()

171 if holder is not None:

172 holder.handles.pop(session_id, None)

173 self.recomputed.pop(session_id, None)

174 self.recomp_counter.pop(session_id, None)

175 self.is_recomputed.pop(session_id, None)

176

177

178def _bind_session_activation(frame: _CheckpointFrame, activation: _SessionActivation) -> None:

179 """Bind one activation to a frame outside the unpack hot path."""

180 if frame.active_session is activation:

181 return

182 if frame.active_session is not None:

183 raise CheckpointError("Concurrent recompute sessions on the same checkpoint frame are not supported.")

184 frame.active_session = activation

185 activation.frames.add(frame)

186

187

188def _register_session_frame(

189 frame: _CheckpointFrame,

190 session_id: Any,

191 activation: Optional[_SessionActivation] = None,

192) -> None:

193 """Register a frame for cleanup and bind its current activation when present."""

194 with _SESSION_FRAMES_LOCK:

195 _SESSION_FRAMES[session_id].add(frame)

196 if activation is not None:

197 _internal_assert(activation.session_id == session_id, "Session activation key does not match its frame.")

198 _bind_session_activation(frame, activation)

199

200

201def _activate_registered_frames(activation: _SessionActivation) -> None:

202 """Install an activation on every frame already registered for its session."""

203 with _SESSION_FRAMES_LOCK:

204 for frame in list(_SESSION_FRAMES.get(activation.session_id, ())):

205 _bind_session_activation(frame, activation)

206

207

208def _deactivate_session(activation: _SessionActivation) -> None:

209 """Remove one activation from every frame bound at context entry."""

210 with _SESSION_FRAMES_LOCK:

211 for frame in list(activation.frames):

212 if frame.active_session is activation:

213 frame.active_session = None

214 activation.frames.clear()

215

216

217def _internal_assert(condition: bool, message: str) -> None:

218 if not condition:

219 raise CheckpointError(message)

220

221

222def _noop_context_fn() -> Tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]:

223 return contextlib.nullcontext(), contextlib.nullcontext()

224

225

226def _default_metadata_fn(tensor: Any) -> Dict[str, Any]:

227 return {"shape": tensor.shape, "dtype": tensor.dtype, "device": tensor.device}

228

229

230def _infer_device_type(*args: Any) -> str:

231 """Return the preferred non-CPU device type found in checkpoint inputs."""

232 device_types = []

233

234 def add_device_type(value: Any) -> None:

235 """Record one non-CPU tensor device type."""

236 if isinstance(value, torch.Tensor) and value.device.type != "cpu":

237 device_types.append(value.device.type)

238

239 tree_map(add_device_type, args)

240 device_types_set = set(device_types)

241 if len(device_types_set) > 1:

242 warnings.warn(

243 "Hyper checkpoint received tensors on multiple non-CPU device types. RNG state is preserved only for "

244 "one device type; CUDA is preferred when present.",

245 stacklevel=3,

246 )

247 if not device_types:

248 return DefaultDeviceType.get_device_type()

249 if "cuda" in device_types_set:

250 return "cuda"

251 return device_types[0]

252

253

254def _get_device_module(device_type: str) -> Any:

255 if device_type == "meta":

256 return torch.device("meta")

257 return getattr(torch, device_type)

258

259

260def _get_device_states(device_type: str, *args: Any) -> Tuple[List[int], List[Any]]:

261 """Capture RNG states for non-CPU input devices of the requested type."""

262 device_ids = []

263

264 def add_device_id(value: Any) -> None:

265 """Record one non-CPU tensor device index."""

266 if isinstance(value, torch.Tensor) and value.device.type not in {"cpu", "meta"}:

267 device_ids.append(value.get_device())

268

269 tree_map(add_device_id, args)

270 device_module = _get_device_module(device_type)

271 states = []

272 for device_id in device_ids:

273 with device_module.device(device_id):

274 states.append(device_module.get_rng_state())

275 return device_ids, states

276

277

278def _set_device_states(device_type: str, devices: List[int], states: List[Any]) -> None:

279 if device_type == "meta":

280 return

281 device_module = _get_device_module(device_type)

282 for device, state in zip(devices, states):

283 with device_module.device(device):

284 device_module.set_rng_state(state)

285

286

287def _get_autocast_kwargs(device_type: str) -> Tuple[Optional[Dict[str, Any]], Dict[str, Any]]:

288 """Return active autocast settings for the selected device and CPU."""

289 device_kwargs = None

290 if torch.amp.is_autocast_available(device_type):

291 device_kwargs = {

292 "enabled": torch.is_autocast_enabled(device_type),

293 "dtype": torch.get_autocast_dtype(device_type),

294 "cache_enabled": torch.is_autocast_cache_enabled(),

295 }

296 cpu_kwargs = {

297 "enabled": torch.is_autocast_enabled("cpu"),

298 "dtype": torch.get_autocast_dtype("cpu"),

299 "cache_enabled": torch.is_autocast_cache_enabled(),

300 }

301 return device_kwargs, cpu_kwargs

302

303

304def _create_recomputation_hooks(frame: _CheckpointFrame, session_id: Any) -> Any:

305 """Create saved-tensor hooks that retain tensors from one recomputation."""

306 frame_ref = weakref.ref(frame)

307

308 def pack_hook(tensor: Any) -> Any:

309 """Store recomputed tensors in their forward holders."""

310 tensor = tensor.detach() if tensor.requires_grad else tensor

311 target_frame = frame_ref()

312 _internal_assert(target_frame is not None, "Checkpoint frame was released during recomputation.")

313 recompute_index = target_frame.recomp_counter[session_id]

314 target_frame.recomp_counter[session_id] += 1

315

316 if recompute_index >= len(target_frame.weak_holders):

317 if not target_frame.early_stop and not target_frame.forward_completed:

318 target_frame.ignore_saved_mismatch = True

319 return tensor

320 raise CheckpointError(

321 "Hyper checkpoint tried to save more tensors during recomputation than during forward."

322 )

323

324 holder = target_frame.weak_holders[recompute_index]()

325 if holder is not None:

326 _internal_assert(

327 holder.handles.get(session_id) is None,

328 "A recomputed tensor handle already exists for this session.",

329 )

330 handle = _Handle()

331 holder.handles[session_id] = handle

332 target_frame.recomputed[session_id][handle] = tensor

333

334 if target_frame.early_stop and target_frame.recomp_counter[session_id] == len(target_frame.weak_holders):

335 raise _StopRecomputationError

336 return tensor

337

338 def unpack_hook(tensor: Any) -> Any:

339 """Return tensors saved by operations inside the recomputation."""

340 return tensor

341

342 return torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook)

343

344

345# PyTorch exposes this tracing guard only as a private decorator.

346@torch._disable_dynamo # pylint: disable=protected-access

347def _run_fn_with_dynamo_disabled(function: Callable, *args: Any, **kwargs: Any) -> Any:

348 """Run recomputation without tracing the saved-tensor unpack hook with Dynamo."""

349 return function(*args, **kwargs)

350

351

352def _run_recomputation(frame: _CheckpointFrame, session_id: Any) -> None:

353 """Run and validate a frame recomputation for the given session."""

354 if frame.is_recomputed[session_id]:

355 return

356

357 activation = frame.active_session

358 if activation is not None:

359 _internal_assert(activation.session_id == session_id, "Active session key does not match recomputation key.")

360 previous_activation = _RECOMPUTE_SESSION.get()

361 token = None

362 if activation is not None and previous_activation is not activation:

363 token = _RECOMPUTE_SESSION.set(activation)

364 try:

365 input_context = frame.input_saver.grad_fn

366 args = input_context.get_args(input_context.saved_tensors)

367 try:

368 with _create_recomputation_hooks(frame, session_id), torch.autograd.enable_grad():

369 _run_fn_with_dynamo_disabled(frame.recompute_fn, *args)

370 except _StopRecomputationError:

371 pass

372 finally:

373 if token is not None:

374 _RECOMPUTE_SESSION.reset(token)

375 frame.is_recomputed[session_id] = True

376 frame.check_recomputed_tensors_match(session_id)

377

378

379def _create_checkpoint_hooks(frame: _CheckpointFrame) -> Any:

380 """Create hooks that lazily recompute tensors saved during forward."""

381 def pack_hook(tensor: Any) -> _Holder:

382 """Replace a forward saved tensor with an opaque holder."""

383 holder = _Holder()

384 frame.weak_holders.append(weakref.ref(holder))

385 if frame.metadata_fn is not None:

386 with torch.no_grad():

387 frame.x_metadatas.append(frame.metadata_fn(tensor))

388 return holder

389

390 def unpack_hook(holder: _Holder) -> Any:

391 """Return the corresponding tensor from lazy or prefired recomputation."""

392 activation = frame.active_session

393 if activation is not None:

394 session_id = activation.session_id

395 retain_on_unpack = activation.retain_on_unpack

396 else:

397 session_id = torch._C._current_graph_task_id() # pylint: disable=W0212

398 if session_id == -1:

399 session_id = int(uuid.uuid4())

400 retain_on_unpack = False

401

402 _run_recomputation(frame, session_id)

403 _internal_assert(session_id in holder.handles, "No recomputed tensor was saved for this checkpoint value.")

404 handle = holder.handles[session_id]

405 if handle is None:

406 raise CheckpointError("A checkpoint tensor was unpacked more than once in the same recompute session.")

407 _internal_assert(handle in frame.recomputed[session_id], "The recomputed tensor has already been released.")

408 tensor = frame.recomputed[session_id][handle]

409 if not retain_on_unpack:

410 holder.handles[session_id] = None

411 return tensor

412

413 return torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook)

414

415

416def _is_compiling() -> bool:

417 compiler = getattr(torch, "compiler", None)

418 return bool(compiler is not None and compiler.is_compiling())

419

420

421def _native_checkpoint(

422 function: Callable,

423 *args: Any,

424 context_fn: Callable,

425 preserve_rng_state: bool,

426 determinism_check: str,

427 debug: bool,

428 early_stop: bool,

429 **kwargs: Any,

430) -> Any:

431 """Use the public native API for compile, adapting 2.6/2.7 early-stop."""

432 with set_checkpoint_early_stop(early_stop):

433 return torch_checkpoint(

434 function,

435 *args,

436 use_reentrant=False,

437 context_fn=context_fn,

438 preserve_rng_state=preserve_rng_state,

439 determinism_check=determinism_check,

440 debug=debug,

441 **kwargs,

442 )

443

444

445@dataclass

446class _ExecutionState:

447 """Device, RNG, and autocast state captured before a checkpointed forward.

448

449 Bundles the values that both the forward pass and every later recomputation

450 must agree on, so the generator that owns the checkpoint lifetime can hand

451 them around as one object instead of a dozen loose locals.

452 """

453

454 device_type: str

455 device_module: Any

456 forward_context: Any

457 recompute_context: Any

458 device_autocast_kwargs: Optional[Dict[str, Any]]

459 cpu_autocast_kwargs: Dict[str, Any]

460 device_was_initialized: bool = False

461 device_ids: List[int] = field(default_factory=list)

462 device_states: List[Any] = field(default_factory=list)

463 cpu_state: Any = None

464

465 @classmethod

466 def capture(cls, context_fn: Callable, preserve_rng_state: bool, args: Tuple[Any, ...]) -> "_ExecutionState":

467 """Resolve contexts and capture RNG state before the forward pass runs."""

468 device_type = _infer_device_type(*args)

469 device_module = _get_device_module(device_type)

470 forward_context, recompute_context = _resolve_contexts(context_fn)

471 device_autocast_kwargs, cpu_autocast_kwargs = _get_autocast_kwargs(device_type)

472

473 # Device RNG is captured only when the device is already initialized;

474 # otherwise the device indexes its state lazily at first use.

475 device_initialized = preserve_rng_state and getattr(device_module, "_initialized", False)

476 device_ids, device_states = _get_device_states(device_type, *args) if device_initialized else ([], [])

477

478 return cls(

479 device_type=device_type,

480 device_module=device_module,

481 forward_context=forward_context,

482 recompute_context=recompute_context,

483 device_autocast_kwargs=device_autocast_kwargs,

484 cpu_autocast_kwargs=cpu_autocast_kwargs,

485 device_was_initialized=device_initialized,

486 device_ids=device_ids,

487 device_states=device_states,

488 cpu_state=torch.get_rng_state() if preserve_rng_state else None,

489 )

490

491 def check_device_state_preserved(self, preserve_rng_state: bool) -> None:

492 """Fail when the device was first initialized inside the checkpointed forward."""

493 device_initialized = getattr(self.device_module, "_initialized", False)

494 if preserve_rng_state and not self.device_was_initialized and device_initialized:

495 raise RuntimeError(

496 "The device state was initialized inside a Hyper checkpoint forward, so its initial RNG state "

497 "could not be preserved. Initialize the device before entering checkpoint."

498 )

499

500 def recompute_fn(self, function: Callable, *inputs: Any) -> None:

501 """Restore execution state and rerun the checkpointed ``function``."""

502 function_kwargs, *function_args = inputs

503 with torch.random.fork_rng(

504 devices=self.device_ids,

505 enabled=self.cpu_state is not None,

506 device_type=self.device_type,

507 ):

508 if self.cpu_state is not None:

509 torch.set_rng_state(self.cpu_state)

510 if self.device_was_initialized:

511 _set_device_states(self.device_type, self.device_ids, self.device_states)

512

513 device_autocast_context = contextlib.nullcontext()

514 if self.device_autocast_kwargs is not None:

515 device_autocast_context = torch.amp.autocast(

516 device_type=self.device_type, **self.device_autocast_kwargs

517 )

518 with device_autocast_context, torch.amp.autocast("cpu", **self.cpu_autocast_kwargs), self.recompute_context:

519 function(*function_args, **function_kwargs)

520

521

522def _resolve_contexts(context_fn: Callable) -> Tuple[Any, Any]:

523 """Validate and unpack the (forward, recompute) context tuple from ``context_fn``."""

524 contexts = context_fn()

525 if not isinstance(contexts, tuple) or len(contexts) != 2:

526 raise ValueError("context_fn must return a (forward_context, recompute_context) tuple.")

527 return contexts[0], contexts[1]

528

529

530def _checkpoint_without_reentrant_generator(

531 function: Callable,

532 preserve_rng_state: bool,

533 context_fn: Callable,

534 determinism_check: str,

535 early_stop: bool,

536 *args: Any,

537 **kwargs: Any,

538) -> Generator[None, None, None]:

539 """Set up eager checkpoint state around the caller's forward execution."""

540 metadata_fn = _resolve_metadata_fn(determinism_check)

541 state = _ExecutionState.capture(context_fn, preserve_rng_state, args)

542

543 frame = _CheckpointFrame(partial(state.recompute_fn, function), early_stop, metadata_fn)

544 dummy = torch.empty((0,), requires_grad=True)

545 frame.input_saver = _NoopSaveInputs.apply(dummy, kwargs, *args)

546

547 if frame.input_saver.grad_fn is None:

548 yield

549 return

550

551 if _RECOMPUTE_SESSION.get() is not None:

552 raise CheckpointError("Nested checkpoint is not supported during scheduled recomputation.")

553

554 collector = _RECOMPUTE_COLLECTOR.get()

555 if collector is not None:

556 collector.append(frame)

557 try:

558 with _create_checkpoint_hooks(frame), state.forward_context:

559 yield

560 frame.forward_completed = True

561 state.check_device_state_preserved(preserve_rng_state)

562 except BaseException:

563 if collector is not None and frame in collector:

564 collector.remove(frame)

565 raise

566

567

568def _resolve_metadata_fn(determinism_check: str) -> Callable:

569 """Return the metadata extractor implementing ``determinism_check``."""

570 metadata_functions = {_DEFAULT_DETERMINISM_MODE: _default_metadata_fn, "none": lambda tensor: None}

571 if determinism_check not in metadata_functions:

572 raise ValueError(

573 f"determinism_check must be one of {list(metadata_functions)}, but got {determinism_check!r}."

574 )

575 return metadata_functions[determinism_check]

576

577

578def checkpoint(

579 function: Callable,

580 *args: Any,

581 use_reentrant: bool = False,

582 context_fn: Callable = _noop_context_fn,

583 preserve_rng_state: bool = True,

584 determinism_check: str = _DEFAULT_DETERMINISM_MODE,

585 debug: bool = False,

586 early_stop: bool = True,

587 **kwargs: Any,

588) -> Any:

589 """Run Hyper's non-reentrant checkpoint implementation.

590

591 Eager execution always uses this implementation. Compile execution falls

592 back to PyTorch's public non-reentrant checkpoint API.

593 """

594 if use_reentrant is not False:

595 raise ValueError("Hyper checkpoint only supports use_reentrant=False.")

596 if not isinstance(early_stop, bool):

597 raise ValueError(f"early_stop must be bool, but got {type(early_stop).__name__}.")

598 if not isinstance(preserve_rng_state, bool):

599 raise ValueError(

600 f"preserve_rng_state must be bool, but got {type(preserve_rng_state).__name__}."

601 )

602 if not callable(context_fn):

603 raise ValueError("context_fn must be callable.")

604 if _is_compiling():

605 return _native_checkpoint(

606 function,

607 *args,

608 context_fn=context_fn,

609 preserve_rng_state=preserve_rng_state,

610 determinism_check=determinism_check,

611 debug=debug,

612 early_stop=early_stop,

613 **kwargs,

614 )

615 if debug:

616 raise ValueError("debug=True is not supported by Hyper eager checkpoint yet.")

617

618 generator = _checkpoint_without_reentrant_generator(

619 function,

620 preserve_rng_state,

621 context_fn,

622 determinism_check,

623 early_stop,

624 *args,

625 **kwargs,

626 )

627 next(generator)

628 try:

629 result = function(*args, **kwargs)

630 except BaseException:

631 generator.close()

632 raise

633 try:

634 next(generator)

635 except StopIteration:

636 return result

637 generator.close()

638 raise CheckpointError("The internal checkpoint generator yielded more than once.")

639

640

641@contextlib.contextmanager

642def recompute_handle_collector_ctx() -> Iterator[List[Any]]:

643 """Collect opaque checkpoint handles created in this context."""

644 handles = []

645 token = _RECOMPUTE_COLLECTOR.set(handles)

646 try:

647 yield handles

648 finally:

649 _RECOMPUTE_COLLECTOR.reset(token)

650

651

652def recompute_handle(handle: Any, session_id: Any) -> None:

653 """Run one collected checkpoint recomputation ahead of backward."""

654 if not isinstance(handle, _CheckpointFrame):

655 raise ValueError("handle must be produced by recompute_handle_collector_ctx().")

656 _validate_session_id(session_id)

657 activation = _RECOMPUTE_SESSION.get()

658 if (

659 activation is not None

660 and activation.session_id == session_id

661 and activation.retain_on_unpack

662 ):

663 _register_session_frame(handle, session_id, activation)

664 _run_recomputation(handle, session_id)

665 return

666 if activation is not None:

667 raise CheckpointError("recompute_handle cannot enter another active recompute session.")

668

669 _register_session_frame(handle, session_id)

670 with recompute_session_ctx(session_id=session_id, retain_on_unpack=True):

671 _run_recomputation(handle, session_id)

672

673

674def _validate_session_id(session_id: Any) -> None:

675 if session_id is None:

676 raise ValueError("session_id must not be None.")

677 try:

678 hash(session_id)

679 except TypeError as error:

680 raise ValueError("session_id must be hashable.") from error

681

682

683@contextlib.contextmanager

684def recompute_session_ctx(session_id: Any, retain_on_unpack: bool = False) -> Iterator[Any]:

685 """Select the key and retention policy used by checkpoint unpack hooks."""

686 _validate_session_id(session_id)

687 if not isinstance(retain_on_unpack, bool):

688 raise ValueError(f"retain_on_unpack must be bool, but got {type(retain_on_unpack).__name__}.")

689 if _RECOMPUTE_SESSION.get() is not None:

690 raise CheckpointError("Nested recompute session contexts are not supported.")

691 activation = _SessionActivation(session_id, retain_on_unpack)

692 token = _RECOMPUTE_SESSION.set(activation)

693 try:

694 _activate_registered_frames(activation)

695 yield session_id

696 finally:

697 try:

698 _deactivate_session(activation)

699 finally:

700 _RECOMPUTE_SESSION.reset(token)

701

702

703def clear_recompute_session(session_id: Any) -> None:

704 """Release retained recomputation data for a session; repeated calls are safe."""

705 _validate_session_id(session_id)

706 with _SESSION_FRAMES_LOCK:

707 registered_frames = _SESSION_FRAMES.pop(session_id, None)

708 frames = list(registered_frames) if registered_frames is not None else []

709 for frame in frames:

710 frame.clear_session(session_id)

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / checkpoint.py: 82%

397 statements