Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / __init__.py: 100%

1# Copyright 2026 Huawei Technologies Co., Ltd. All rights reserved. 

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

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

4# You may obtain a copy of the License at 

5# 

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

7# 

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

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

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

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

12# limitations under the License. 

13# ============================================================================ 

14""" 

15Hyper Parallel Checkpoint Package. 

16 

17This package provides distributed checkpoint saving and loading capabilities, 

18including support for tensor sharding, resharding, and layout management. 

19""" 

20 

21__all__ = [ 

22 # Main API 

23 "save", 

24 "async_save", 

25 "AsyncSaveResponse", 

26 "load", 

27 # Metadata 

28 "Metadata", 

29 "MetadataIndex", 

30 "TensorStorageMetadata", 

31 "BytesStorageMetadata", 

32 "ChunkStorageMetadata", 

33 "TensorProperties", 

34 "CHUNK_INFO", 

35 "ChunkInfo", 

36 # Planner interfaces 

37 "SavePlanner", 

38 "LoadPlanner", 

39 "SavePlan", 

40 "LoadPlan", 

41 "WriteItem", 

42 "ReadItem", 

43 "WriteItemType", 

44 "LoadItemType", 

45 # Standard planners 

46 "StandardSavePlanner", 

47 "StandardLoadPlanner", 

48 # Storage interfaces 

49 "StorageWriter", 

50 "StorageReader", 

51 "StorageInfo", 

52 "WriteResult", 

53 # File system storage 

54 "FileSystemWriter", 

55 "FileSystemReader", 

56 # Optimizer state dict 

57 "get_optim_state_dict", 

58 "set_optim_state_dict", 

59] 

60 

61# Main API 

62from hyper_parallel.core.distributed_checkpoint.api import ( 

63 AsyncSaveResponse, 

64 async_save, 

65 load, 

66 save, 

67) 

68 

69# Metadata structures 

70from hyper_parallel.core.distributed_checkpoint.metadata import ( 

71 BytesStorageMetadata, 

72 ChunkStorageMetadata, 

73 Metadata, 

74 MetadataIndex, 

75 TensorProperties, 

76 TensorStorageMetadata, 

77 CHUNK_INFO, 

78 ChunkInfo 

79) 

80 

81# Planner interfaces and data structures 

82from hyper_parallel.core.distributed_checkpoint.planner import ( 

83 LoadItemType, 

84 LoadPlan, 

85 LoadPlanner, 

86 ReadItem, 

87 SavePlan, 

88 SavePlanner, 

89 WriteItem, 

90 WriteItemType, 

91) 

92 

93# Standard planner implementations 

94from hyper_parallel.core.distributed_checkpoint.standard_planner import ( 

95 StandardLoadPlanner, 

96 StandardSavePlanner, 

97) 

98 

99# Storage interfaces and data structures 

100from hyper_parallel.core.distributed_checkpoint.storage import ( 

101 StorageInfo, 

102 StorageReader, 

103 StorageWriter, 

104 WriteResult, 

105) 

106 

107# File system storage implementations 

108from hyper_parallel.core.distributed_checkpoint.filesystem_storage import ( 

109 FileSystemReader, 

110 FileSystemWriter, 

111) 

112 

113# Optimizer state dict 

114from hyper_parallel.core.distributed_checkpoint.state_dict import ( 

115 get_optim_state_dict, 

116 set_optim_state_dict, 

117)