Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/platform/torch/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 typing import Any, Callable, DefaultDict, Dict, Generator, Iterator, List, Optional, Tuple

31import torch

32from torch.utils._pytree import tree_map

33from torch.utils.checkpoint import DefaultDeviceType

34from torch.utils.checkpoint import checkpoint as torch_checkpoint

35from torch.utils.checkpoint import set_checkpoint_early_stop

38_DEFAULT_DETERMINISM_MODE = "default"

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

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

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

42_SESSION_FRAMES_LOCK = threading.RLock()

45class CheckpointError(RuntimeError):

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

49class _Handle:

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

53class _Holder:

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

56 def __init__(self) -> None:

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

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

61class _StopRecomputationError(Exception):

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

65class _SessionActivation:

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

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

69 """Initialize one scoped session activation."""

70 self.session_id = session_id

71 self.retain_on_unpack = retain_on_unpack

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

75class _NoopSaveInputs(torch.autograd.Function):

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

78 @staticmethod

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

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

81 del args

82 return torch.empty((0,))

84 @staticmethod

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

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

87 del output

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

89 tensor_indices, tensors = zip(*tensor_pairs)

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

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

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

94 """Reconstruct the original checkpoint arguments."""

95 restored_args = [

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

97 for index, value in enumerate(stored_args)

98 ]

99 return restored_args[1:]

100

101 ctx.get_args = get_args

102 ctx.save_for_backward(*tensors)

103

104 @staticmethod

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

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

107 del ctx, grad_outputs

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

109

110

111class _CheckpointFrame:

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

113

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

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

116 self.recompute_fn = recompute_fn

117 self.input_saver = None

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

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

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

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

122 self.early_stop = early_stop

123 self.metadata_fn = metadata_fn

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

125 self.forward_completed = False

126 self.ignore_saved_mismatch = False

127 self.active_session: Optional[_SessionActivation] = None

128

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

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

131 if self.ignore_saved_mismatch:

132 return

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

134 raise CheckpointError(

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

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

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

138 )

139

140 mismatches = []

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

142 holder = weak_holder()

143 if holder is None:

144 continue

145 handle = holder.handles.get(session_id)

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

147 _internal_assert(

148 handle in self.recomputed[session_id],

149 "Missing recomputed tensor during metadata validation.",

150 )

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

152 recomputed_metadata = self.metadata_fn(recomputed_tensor)

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

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

155

156 if mismatches:

157 details = "\n".join(

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

159 for index, forward_metadata, recomputed_metadata in mismatches

160 )

161 raise CheckpointError(

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

163 )

164

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

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

167 for weak_holder in self.weak_holders:

168 holder = weak_holder()

169 if holder is not None:

170 holder.handles.pop(session_id, None)

171 self.recomputed.pop(session_id, None)

172 self.recomp_counter.pop(session_id, None)

173 self.is_recomputed.pop(session_id, None)

174

175

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

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

178 if frame.active_session is activation:

179 return

180 if frame.active_session is not None:

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

182 frame.active_session = activation

183 activation.frames.add(frame)

184

185

186def _register_session_frame(

187 frame: _CheckpointFrame,

188 session_id: Any,

189 activation: Optional[_SessionActivation] = None,

190) -> None:

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

192 with _SESSION_FRAMES_LOCK:

193 _SESSION_FRAMES[session_id].add(frame)

194 if activation is not None:

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

196 _bind_session_activation(frame, activation)

197

198

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

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

201 with _SESSION_FRAMES_LOCK:

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

203 _bind_session_activation(frame, activation)

204

205

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

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

208 with _SESSION_FRAMES_LOCK:

209 for frame in list(activation.frames):

210 if frame.active_session is activation:

211 frame.active_session = None

212 activation.frames.clear()

213

214

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

216 if not condition:

217 raise CheckpointError(message)

218

219

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

221 return contextlib.nullcontext(), contextlib.nullcontext()

222

223

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

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

226

227

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

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

230 device_types = []

231

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

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

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

235 device_types.append(value.device.type)

236

237 tree_map(add_device_type, args)

238 device_types_set = set(device_types)

239 if len(device_types_set) > 1:

240 warnings.warn(

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

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

243 stacklevel=3,

244 )

245 if not device_types:

246 return DefaultDeviceType.get_device_type()

247 if "cuda" in device_types_set:

248 return "cuda"

249 return device_types[0]

250

251

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

253 if device_type == "meta":

254 return torch.device("meta")

255 return getattr(torch, device_type)

256

257

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

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

260 device_ids = []

261

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

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

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

265 device_ids.append(value.get_device())

266

267 tree_map(add_device_id, args)

268 device_module = _get_device_module(device_type)

269 states = []

270 for device_id in device_ids:

271 with device_module.device(device_id):

272 states.append(device_module.get_rng_state())

273 return device_ids, states

274

275

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

277 if device_type == "meta":

278 return

279 device_module = _get_device_module(device_type)

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

281 with device_module.device(device):

282 device_module.set_rng_state(state)

283

284

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

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

287 device_kwargs = None

288 if torch.amp.is_autocast_available(device_type):

289 device_kwargs = {

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

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

292 "cache_enabled": torch.is_autocast_cache_enabled(),

293 }

294 cpu_kwargs = {

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

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

297 "cache_enabled": torch.is_autocast_cache_enabled(),

298 }

299 return device_kwargs, cpu_kwargs

300

301

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

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

304 frame_ref = weakref.ref(frame)

305

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

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

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

309 target_frame = frame_ref()

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

311 recompute_index = target_frame.recomp_counter[session_id]

312 target_frame.recomp_counter[session_id] += 1

313

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

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

316 target_frame.ignore_saved_mismatch = True

317 return tensor

318 raise CheckpointError(

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

320 )

321

322 holder = target_frame.weak_holders[recompute_index]()

323 if holder is not None:

324 _internal_assert(

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

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

327 )

328 handle = _Handle()

329 holder.handles[session_id] = handle

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

331

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

333 raise _StopRecomputationError

334 return tensor

335

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

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

338 return tensor

339

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

341

342

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

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

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

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

347 return function(*args, **kwargs)

348

349

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

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

352 if frame.is_recomputed[session_id]:

353 return

354

355 activation = frame.active_session

356 if activation is not None:

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

358 previous_activation = _RECOMPUTE_SESSION.get()

359 token = None

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

361 token = _RECOMPUTE_SESSION.set(activation)

362 try:

363 input_context = frame.input_saver.grad_fn

364 args = input_context.get_args(input_context.saved_tensors)

365 try:

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

367 _run_fn_with_dynamo_disabled(frame.recompute_fn, *args)

368 except _StopRecomputationError:

369 pass

370 finally:

371 if token is not None:

372 _RECOMPUTE_SESSION.reset(token)

373 frame.is_recomputed[session_id] = True

374 frame.check_recomputed_tensors_match(session_id)

375

376

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

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

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

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

381 holder = _Holder()

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

383 if frame.metadata_fn is not None:

384 with torch.no_grad():

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

386 return holder

387

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

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

390 activation = frame.active_session

391 if activation is not None:

392 session_id = activation.session_id

393 retain_on_unpack = activation.retain_on_unpack

394 else:

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

396 if session_id == -1:

397 session_id = int(uuid.uuid4())

398 retain_on_unpack = False

399

400 _run_recomputation(frame, session_id)

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

402 handle = holder.handles[session_id]

403 if handle is None:

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

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

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

407 if not retain_on_unpack:

408 holder.handles[session_id] = None

409 return tensor

410

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

412

413

414def _is_compiling() -> bool:

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

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

417

418

419def _native_checkpoint(

420 function: Callable,

421 *args: Any,

422 context_fn: Callable,

423 preserve_rng_state: bool,

424 determinism_check: str,

425 debug: bool,

426 early_stop: bool,

427 **kwargs: Any,

428) -> Any:

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

430 with set_checkpoint_early_stop(early_stop):

431 return torch_checkpoint(

432 function,

433 *args,

434 use_reentrant=False,

435 context_fn=context_fn,

436 preserve_rng_state=preserve_rng_state,

437 determinism_check=determinism_check,

438 debug=debug,

439 **kwargs,

440 )

441

442

443def _checkpoint_without_reentrant_generator(

444 function: Callable,

445 preserve_rng_state: bool,

446 context_fn: Callable,

447 determinism_check: str,

448 early_stop: bool,

449 *args: Any,

450 **kwargs: Any,

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

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

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

454 if determinism_check not in metadata_functions:

455 raise ValueError(

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

457 )

458 metadata_fn = metadata_functions[determinism_check]

459

460 device_type = _infer_device_type(*args)

461 device_module = _get_device_module(device_type)

462 contexts = context_fn()

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

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

465 forward_context, recompute_context = contexts

466 device_autocast_kwargs, cpu_autocast_kwargs = _get_autocast_kwargs(device_type)

467

468 had_device_in_forward = False

469 forward_devices: List[int] = []

470 forward_device_states: List[Any] = []

471 forward_cpu_state = None

472 if preserve_rng_state:

473 forward_cpu_state = torch.get_rng_state()

474 if getattr(device_module, "_initialized", False):

475 had_device_in_forward = True

476 forward_devices, forward_device_states = _get_device_states(device_type, *args)

477

478 def recompute_fn(*inputs: Any) -> None:

479 """Restore execution state and rerun the checkpointed function."""

480 function_kwargs, *function_args = inputs

481 rng_devices = forward_devices if preserve_rng_state and had_device_in_forward else []

482 with torch.random.fork_rng(

483 devices=rng_devices,

484 enabled=preserve_rng_state,

485 device_type=device_type,

486 ):

487 if preserve_rng_state:

488 torch.set_rng_state(forward_cpu_state)

489 if had_device_in_forward:

490 _set_device_states(device_type, forward_devices, forward_device_states)

491

492 device_autocast_context = contextlib.nullcontext()

493 if device_autocast_kwargs is not None:

494 device_autocast_context = torch.amp.autocast(device_type=device_type, **device_autocast_kwargs)

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

496 function(*function_args, **function_kwargs)

497

498 frame = _CheckpointFrame(recompute_fn, early_stop, metadata_fn)

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

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

501

502 if frame.input_saver.grad_fn is None:

503 yield

504 return

505

506 activation = _RECOMPUTE_SESSION.get()

507 if activation is not None:

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

509

510 collector = _RECOMPUTE_COLLECTOR.get()

511 if collector is not None:

512 collector.append(frame)

513 try:

514 with _create_checkpoint_hooks(frame), forward_context:

515 yield

516 frame.forward_completed = True

517

518 if getattr(device_module, "_initialized", False) and preserve_rng_state and not had_device_in_forward:

519 raise RuntimeError(

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

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

522 )

523 except BaseException:

524 if collector is not None and frame in collector:

525 collector.remove(frame)

526 raise

527

528

529def checkpoint(

530 function: Callable,

531 *args: Any,

532 use_reentrant: bool = False,

533 context_fn: Callable = _noop_context_fn,

534 preserve_rng_state: bool = True,

535 determinism_check: str = _DEFAULT_DETERMINISM_MODE,

536 debug: bool = False,

537 early_stop: bool = True,

538 **kwargs: Any,

539) -> Any:

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

541

542 Eager execution always uses this implementation. Compile execution falls

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

544 """

545 if use_reentrant is not False:

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

547 if not isinstance(early_stop, bool):

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

549 if not isinstance(preserve_rng_state, bool):

550 raise ValueError(

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

552 )

553 if not callable(context_fn):

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

555 if _is_compiling():

556 return _native_checkpoint(

557 function,

558 *args,

559 context_fn=context_fn,

560 preserve_rng_state=preserve_rng_state,

561 determinism_check=determinism_check,

562 debug=debug,

563 early_stop=early_stop,

564 **kwargs,

565 )

566 if debug:

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

568

569 generator = _checkpoint_without_reentrant_generator(

570 function,

571 preserve_rng_state,

572 context_fn,

573 determinism_check,

574 early_stop,

575 *args,

576 **kwargs,

577 )

578 next(generator)

579 try:

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

581 except BaseException:

582 generator.close()

583 raise

584 try:

585 next(generator)

586 except StopIteration:

587 return result

588 generator.close()

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

590

591

592@contextlib.contextmanager

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

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

595 handles = []

596 token = _RECOMPUTE_COLLECTOR.set(handles)

597 try:

598 yield handles

599 finally:

600 _RECOMPUTE_COLLECTOR.reset(token)

601

602

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

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

605 if not isinstance(handle, _CheckpointFrame):

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

607 _validate_session_id(session_id)

608 activation = _RECOMPUTE_SESSION.get()

609 if (

610 activation is not None

611 and activation.session_id == session_id

612 and activation.retain_on_unpack

613 ):

614 _register_session_frame(handle, session_id, activation)

615 _run_recomputation(handle, session_id)

616 return

617 if activation is not None:

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

619

620 _register_session_frame(handle, session_id)

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

622 _run_recomputation(handle, session_id)

623

624

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

626 if session_id is None:

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

628 try:

629 hash(session_id)

630 except TypeError as error:

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

632

633

634@contextlib.contextmanager

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

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

637 _validate_session_id(session_id)

638 if not isinstance(retain_on_unpack, bool):

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

640 if _RECOMPUTE_SESSION.get() is not None:

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

642 activation = _SessionActivation(session_id, retain_on_unpack)

643 token = _RECOMPUTE_SESSION.set(activation)

644 try:

645 _activate_registered_frames(activation)

646 yield session_id

647 finally:

648 try:

649 _deactivate_session(activation)

650 finally:

651 _RECOMPUTE_SESSION.reset(token)

652

653

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

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

656 _validate_session_id(session_id)

657 with _SESSION_FRAMES_LOCK:

658 registered_frames = _SESSION_FRAMES.pop(session_id, None)

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

660 for frame in frames:

661 frame.clear_session(session_id)

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / activation_checkpoint / checkpoint.py: 78%

381 statements