Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / base.py: 79%

1# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates 

2# Copyright 2026 Huawei Technologies Co., Ltd 

3# 

4# Licensed under the Apache License, Version 2.0 (the "License"); 

5# you may not use this file except in compliance with the License. 

6# You may obtain a copy of the License at 

7# 

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

9# 

10# Unless required by applicable law or agreed to in writing, software 

11# distributed under the License is distributed on an "AS IS" BASIS, 

12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

13# See the License for the specific language governing permissions and 

14# limitations under the License. 

15# ============================================================================ 

16"""Checkpointer interface --- owns the on-disk checkpoint format. 

17 

18The split against ``CheckpointerCallback`` is deliberate: 

19 

20* the **callback** decides *when* to save, *what* goes into the payload, and how 

21 a restored payload maps back onto trainer objects (progress counters, LR 

22 scheduler, dataloader position, RNG); 

23* the **checkpointer** knows only how to write that payload to a directory and 

24 read it back --- storage layout, marker files, async plumbing, and the 

25 backend-specific load semantics. 

26 

27Swapping the persistence backend is therefore a matter of registering another 

28:class:`CheckpointerBase`, with no change to the callback. 

29""" 

30 

31__all__ = [ 

32 "CheckpointerBase", 

33] 

34 

35from abc import ABC, abstractmethod 

36from typing import Any, Dict, Optional 

37 

38 

39class CheckpointerBase(ABC): 

40 """Persist and restore a training payload for one checkpoint directory. 

41 

42 The payload is a plain dict whose ``model`` / ``optimizer`` entries are 

43 state dicts, plus an optional ``extra_state`` bundle of non-tensor training 

44 progress. Implementations own where each part physically lands. 

45 """ 

46 

47 @abstractmethod 

48 def save( 

49 self, 

50 path: str, 

51 state: Dict[str, Any], 

52 *, 

53 global_step: int, 

54 save_async: bool = False, 

55 ) -> None: 

56 """Write ``state`` into the checkpoint directory ``path``. 

57 

58 Args: 

59 path: Directory for this checkpoint; created if absent. 

60 state: Payload to persist (``model`` / ``optimizer`` / 

61 optionally ``extra_state``). 

62 global_step: Step this checkpoint represents, recorded so a later 

63 run can tell a complete checkpoint from a half-written one. 

64 save_async: Return before persistence finishes. The caller must 

65 reach :meth:`maybe_wait_for_async_save` before relying on the 

66 checkpoint or exiting the process. 

67 """ 

68 

69 @abstractmethod 

70 def load( 

71 self, 

72 path: str, 

73 state: Dict[str, Any], 

74 *, 

75 strict_model: bool = True, 

76 extra_state_skeleton: Optional[Dict[str, Any]] = None, 

77 ) -> Dict[str, Any]: 

78 """Fill ``state`` in place from the checkpoint at ``path``. 

79 

80 Args: 

81 path: Directory holding the checkpoint. 

82 state: Payload to fill; its existing keys define what is read. 

83 strict_model: Require every requested ``model`` key to exist in the 

84 checkpoint. Disable for adapter-only (PEFT) checkpoints, which 

85 legitimately omit frozen base weights. 

86 extra_state_skeleton: A dict shaped like the ``extra_state`` that 

87 was saved. Required to restore a bundle that was embedded in 

88 the payload, because a backend fills only the keys the caller 

89 already declares. Pass ``None`` to skip ``extra_state`` restore. 

90 

91 Returns: 

92 The same ``state`` object, filled. 

93 """ 

94 

95 def maybe_wait_for_async_save(self) -> None: 

96 """Block until an in-flight async save is fully persisted. 

97 

98 Safe to call when nothing is pending. Backends without async support 

99 inherit this no-op. 

100 """ 

101 return 

102 

103 def raise_for_failed_async_save(self) -> None: 

104 """Re-raise an in-flight async save's failure, without waiting for it. 

105 

106 Callers use this to learn about a failure at their next convenient point 

107 rather than at their next save, which may be thousands of steps away. A 

108 save still running is left alone, so it is cheap enough to call every 

109 step. Backends without async support inherit this no-op. 

110 """ 

111 return 

112 

113 def find_latest_checkpoint(self, checkpoint_dir: str) -> Optional[str]: 

114 """Return the newest resumable checkpoint under ``checkpoint_dir``. 

115 

116 An instance method like the rest of this interface, so a backend whose 

117 discovery needs configuration (an object-store client, credentials, a 

118 bucket) has somewhere to read it from. Returns ``None`` when the 

119 directory holds no usable checkpoint. 

120 """ 

121 raise NotImplementedError