Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / collectives / cc.py: 100%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

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 

6# 

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

8# 

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"""PyTorch distributed process group API.""" 

16from datetime import timedelta 

17from typing import Any, Optional, Union 

18 

19import torch.distributed as dist 

20from torch.distributed import ProcessGroup 

21 

22from hyper_parallel.core.utils.communication import EXISTING_COMM_GROUPS 

23 

24 

25def _group_key(ranks: list[int]) -> str: 

26 """Build a stable cache key from process-group ranks.""" 

27 return str(tuple(sorted(ranks))) 

28 

29 

30def init_process_group( 

31 backend: Optional[str] = None, 

32 *, 

33 init_method: Optional[str] = None, 

34 timeout: Optional[timedelta] = None, 

35 world_size: int = -1, 

36 rank: int = -1, 

37 store: Any = None, 

38 pg_options: Any = None, 

39 device_id: Any = None 

40) -> None: 

41 """ 

42 Init global process group, this is the start of distributed job. 

43 

44 Args: 

45 backend: The backend used for distributed communication. 

46 init_method: The method to initialize the process group. 

47 timeout: Timeout for operations executed against the process group. 

48 world_size: Number of processes participating in the job 

49 rank: Rank of the current process 

50 store: Key/value store for exchanging connection information 

51 pg_options: Process group options for backend-specific configurations 

52 device_id: Specific device this process will work on 

53 """ 

54 dist.init_process_group( 

55 backend=backend, 

56 init_method=init_method, 

57 timeout=timeout, 

58 world_size=world_size, 

59 rank=rank, 

60 store=store, 

61 pg_options=pg_options, 

62 device_id=device_id, 

63 ) 

64 

65 

66def destroy_process_group(group: Optional[ProcessGroup] = None) -> None: 

67 """ 

68 Destroy a given process group. 

69 

70 Args: 

71 group: The process group to be destroyed. If None, destroys the default group. 

72 

73 """ 

74 if group is None: 

75 EXISTING_COMM_GROUPS.clear() 

76 else: 

77 keys_to_destroy = [key for key, cached_group in EXISTING_COMM_GROUPS.items() if cached_group == group] 

78 for key in keys_to_destroy: 

79 del EXISTING_COMM_GROUPS[key] 

80 dist.destroy_process_group(group) 

81 

82 

83def get_process_group_ranks(group: Optional[ProcessGroup] = None) -> list[int]: 

84 """ 

85 Get rank list of the given process group. 

86 

87 Args: 

88 group: The process group to get ranks from. If None, uses the default group. 

89 

90 Returns: 

91 List of ranks in the specified process group. 

92 

93 """ 

94 resolved_group = group if group is not None else dist.group.WORLD 

95 return dist.get_process_group_ranks(resolved_group) 

96 

97 

98def get_backend(group: Optional[ProcessGroup] = None) -> str: 

99 """ 

100 Get the backend of the given process group. 

101 Args: 

102 group: The process group to get backend from. If None, uses the default group. 

103 

104 Returns: 

105 The backend name of the specified process group. 

106 

107 """ 

108 return dist.get_backend(group) 

109 

110 

111def split_group(parent_pg: Optional[ProcessGroup] = None, 

112 split_ranks: Optional[list] = None, 

113 timeout: Optional[timedelta] = None, 

114 pg_options: Optional[Any] = None, 

115 group_desc: Optional[str] = None, 

116 ) -> Optional[ProcessGroup]: 

117 """ 

118 Create split group relative to the parent process group. 

119 """ 

120 del parent_pg, timeout, pg_options, group_desc 

121 if not split_ranks: 

122 raise ValueError("split_ranks cannot be None or empty") 

123 

124 current_rank = dist.get_rank() 

125 current_group = None 

126 for ranks in split_ranks: 

127 key = _group_key(ranks) 

128 group = EXISTING_COMM_GROUPS.get(key) 

129 if group is None: 

130 group = dist.new_group(ranks=ranks) 

131 EXISTING_COMM_GROUPS[key] = group 

132 if current_rank in ranks: 

133 current_group = group 

134 return current_group 

135 

136 

137def get_group_local_rank(group: Optional[ProcessGroup] = None) -> int: 

138 """get group local rank id""" 

139 return dist.get_rank(group) 

140 

141 

142def mark_created_groups(process_group: Union[ProcessGroup, list[ProcessGroup]]) -> None: 

143 """ 

144 mark created groups 

145 

146 Args: 

147 process_group (Union[Any, list[Any]]): A process group or a list of process groups. 

148 

149 """ 

150 groups = process_group if isinstance(process_group, list) else [process_group] 

151 for group in groups: 

152 ranks = dist.get_process_group_ranks(group) 

153 EXISTING_COMM_GROUPS[_group_key(ranks)] = group