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

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 backend discovery and construction.""" 

17 

18__all__ = [ 

19 "CHECKPOINTER_REGISTRY", 

20 "build_checkpointer", 

21] 

22 

23from typing import Any, Callable, Iterator, List, MutableMapping, Optional, Type, Union 

24 

25from hyper_parallel.components.checkpoint.base import CheckpointerBase 

26 

27 

28class Registry(MutableMapping): 

29 """Domain registry for checkpointer backends. 

30 

31 Inlined from the former ``auto_models/components/utils/registry.py`` in 

32 stage 7 (05 §10.3 row 565): the generic Registry was merged into the 

33 domain registries that consume it instead of living in a utils package. 

34 """ 

35 

36 # Class instance object, so that a call to `register` can be reflected into all other files correctly, even if 

37 # a new instance is created (in order to locally override a given function) 

38 registry = [] 

39 

40 def __init__(self, name: str) -> None: 

41 """Initialize the registry. 

42 

43 Args: 

44 name: Human-readable name of the registered category, used in 

45 error messages. 

46 """ 

47 self._name = name 

48 self.registry.append(name) 

49 self._local_mapping = {} 

50 self._global_mapping = {} 

51 

52 def __getitem__(self, key: str) -> Union[Type, Callable]: 

53 """Look up a registered class or function by key. 

54 

55 Args: 

56 key: Registration key. 

57 

58 Returns: 

59 The locally overridden value when present, otherwise the globally 

60 registered value. 

61 

62 Raises: 

63 ValueError: If ``key`` is not registered. 

64 """ 

65 # First check if instance has a local override 

66 if key not in self.valid_keys(): 

67 raise ValueError(f"Unknown {self._name} name: {key}. No {self._name} registered for this source.") 

68 if key in self._local_mapping: 

69 return self._local_mapping[key] 

70 return self._global_mapping[key] 

71 

72 def __setitem__(self, key: str, value: Union[Type, Callable]) -> None: 

73 """Set a local override for ``key`` without affecting other instances.""" 

74 # Allow local update of the default functions without impacting other instances 

75 self._local_mapping.update({key: value}) 

76 

77 def __delitem__(self, key: str) -> None: 

78 """Delete the local override for ``key``.""" 

79 del self._local_mapping[key] 

80 

81 def __iter__(self) -> Iterator[str]: 

82 """Iterate over all valid keys, local overrides taking precedence.""" 

83 # Ensure we use all keys, with the overwritten ones on top 

84 return iter({**self._global_mapping, **self._local_mapping}) 

85 

86 def __len__(self) -> int: 

87 """Return the number of distinct registered keys.""" 

88 return len(self._global_mapping.keys() | self._local_mapping.keys()) 

89 

90 def register( 

91 self, 

92 key: str, 

93 cls_or_func: Optional[Union[Type, Callable]] = None, 

94 ) -> Union[Type, Callable]: 

95 """Register a class or function under ``key``. 

96 

97 Can be used either as a plain call or as a decorator:: 

98 

99 registry.register("name", MyClass) 

100 

101 @registry.register("name") 

102 class MyClass: ... 

103 

104 Args: 

105 key: Registration key. 

106 cls_or_func: The class or function to register. When ``None``, a 

107 decorator is returned instead. 

108 

109 Returns: 

110 The registered class/function, or a decorator when ``cls_or_func`` 

111 is ``None``. 

112 

113 Raises: 

114 ValueError: If ``key`` is already registered (decorator form). 

115 """ 

116 if cls_or_func is not None: 

117 self._global_mapping[key] = cls_or_func 

118 return cls_or_func 

119 

120 def decorator(cls_or_func: Union[Type, Callable]) -> Union[Type, Callable]: 

121 """Register the decorated class or function under ``key``.""" 

122 if key in self._global_mapping: 

123 raise ValueError( 

124 f"{self._name} for '{key}' is already registered. Cannot register duplicate {self._name}." 

125 ) 

126 self._global_mapping.update({key: cls_or_func}) 

127 return cls_or_func 

128 

129 return decorator 

130 

131 def valid_keys(self) -> List[str]: 

132 """Return the list of all registered keys.""" 

133 return list(self.keys()) 

134 

135 

136CHECKPOINTER_REGISTRY = Registry("checkpointer") 

137 

138 

139def build_checkpointer( 

140 ckpt_manager: str = "dcp", 

141 **kwargs: Any, 

142) -> "CheckpointerBase": 

143 """Instantiate the checkpointer backend named by ``ckpt_manager``. 

144 

145 Args: 

146 ckpt_manager: Registered backend name. 

147 **kwargs: Forwarded to the backend constructor. 

148 

149 Returns: 

150 The constructed checkpointer. 

151 """ 

152 return CHECKPOINTER_REGISTRY[ckpt_manager](**kwargs) 

153 

154 

155@CHECKPOINTER_REGISTRY.register("dcp") 

156def _dcp_checkpointer(**kwargs: Any) -> CheckpointerBase: 

157 """Construct the distributed-checkpoint (DCP) backend.""" 

158 # Imported lazily: dcp_checkpointer imports CheckpointerBase from the 

159 # base module, so a top-level import here would be circular. 

160 # pylint: disable=import-outside-toplevel,cyclic-import 

161 from hyper_parallel.components.checkpoint.dcp_checkpointer import ( 

162 DistributedCheckpointer, 

163 ) 

164 

165 return DistributedCheckpointer(**kwargs)