Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / _collective.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"""Cold-path result exchange and failure diagnostics for Multicore owners.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Callable 

20from typing import Any 

21 

22import torch.distributed as dist 

23 

24FailureReport = tuple[str, str] 

25 

26 

27class CollectiveConvergenceError(RuntimeError): 

28 """Raised when peers cannot exchange collective lifecycle results.""" 

29 

30 

31def exchange(value: Any, *, group: Any, failure_prefix: str) -> list[Any]: 

32 """Exchange picklable results, preserving the original communication error as cause. 

33 

34 Args: 

35 value: Picklable local result. 

36 group: Explicit group; None selects WORLD when initialized. 

37 failure_prefix: Caller-owned operation context for communication failures. 

38 

39 Returns: 

40 Results in group rank order, or the local result when distributed is uninitialized. 

41 """ 

42 try: 

43 if not dist.is_initialized(): 

44 return [value] 

45 values = [None] * dist.get_world_size(group) 

46 dist.all_gather_object(values, value, group=group) 

47 return values 

48 except Exception as error: # pylint: disable=broad-exception-caught 

49 exception_type, message = failure_report(error) 

50 summary = message if exception_type == "RuntimeError" else f"{exception_type}: {message}" 

51 raise CollectiveConvergenceError(f"{failure_prefix}: {summary}") from error 

52 

53 

54def failure_report(error: Exception, *, message: str | None = None) -> FailureReport: 

55 """Build the picklable exception type and message for one local failure. 

56 

57 Args: 

58 error: Local exception raised by the collective stage. 

59 message: Optional user-facing replacement for the exception text. 

60 """ 

61 return type(error).__name__, str(error) if message is None else message 

62 

63 

64def format_collective_failures( 

65 failures: list[FailureReport | None], 

66 *, 

67 rank_label: str = "rank", 

68 ranks_label: str = "ranks", 

69) -> str: 

70 """Format only failed ranks, grouping peers with identical diagnostics. 

71 

72 Args: 

73 failures: Failure report or ``None`` in collective-group rank order. 

74 rank_label: Label for one failing rank. 

75 ranks_label: Label for multiple ranks with the same failure. 

76 

77 Returns: 

78 Compact failure descriptions in first-rank order. 

79 """ 

80 grouped: dict[FailureReport, list[int]] = {} 

81 for rank, failure in enumerate(failures): 

82 if failure is not None: 

83 grouped.setdefault(failure, []).append(rank) 

84 descriptions = [] 

85 for (exception_type, message), ranks in grouped.items(): 

86 diagnostic = message if exception_type == "RuntimeError" else f"{exception_type}: {message}" 

87 owner = f"{rank_label}={ranks[0]}" if len(ranks) == 1 else f"{ranks_label}={ranks}" 

88 descriptions.append(f"{owner}: {diagnostic}") 

89 return "; ".join(descriptions) 

90 

91 

92def collective_call( 

93 callback: Callable[[], None], 

94 *, 

95 stage: str, 

96 rank_label: str = "rank", 

97 ranks_label: str = "ranks", 

98 group: Any, 

99) -> None: 

100 """Run one lifecycle stage and converge local failures across peers. 

101 

102 Args: 

103 callback: Operation executed by every rank in the same order. 

104 stage: User-facing description of that stage. 

105 rank_label: Label for one failing collective-group rank. 

106 ranks_label: Label for multiple ranks with the same failure. 

107 group: Explicit collective group; None selects WORLD when initialized. 

108 """ 

109 failure = None 

110 try: 

111 callback() 

112 # Peers must see ordinary local errors before entering the next collective. 

113 except Exception as error: # pylint: disable=broad-exception-caught 

114 failure = error 

115 report = None if failure is None else failure_report(failure) 

116 failures = exchange(report, group=group, failure_prefix=f"{stage} failed while converging {rank_label} results") 

117 if any(report is not None for report in failures): 

118 details = format_collective_failures(failures, rank_label=rank_label, ranks_label=ranks_label) 

119 raise RuntimeError(f"{stage} failed on {details}") from failure