Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / debug / _call_records.py: 95%

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"""Structured call record classes for CommDebugMode tracing.""" 

16import time 

17from dataclasses import dataclass, field 

18from typing import List, Optional, Tuple 

19 

20 

21@dataclass 

22class TensorInfo: 

23 """Metadata snapshot of a tensor at trace time.""" 

24 shape: Tuple[int, ...] 

25 dtype: str 

26 is_dtensor: bool = False 

27 placements: Optional[Tuple] = None 

28 mesh_shape: Optional[Tuple[int, ...]] = None 

29 

30 

31@dataclass 

32class DebugCall: 

33 """Base class for all traced call records. 

34 

35 Maintains a tree structure via ``children`` so that collective calls 

36 appear nested under the operator that triggered them. 

37 """ 

38 call_depth: int = 0 

39 timestamp: float = field(default_factory=time.time) 

40 children: List["DebugCall"] = field(default_factory=list) 

41 

42 def render_self(self) -> str: 

43 """Render this record's own one-line summary, without its children. 

44 

45 Overridden by every subclass to describe its own payload; the base 

46 version only reports the depth of a plain ``DebugCall``. 

47 

48 Returns: 

49 str: A single line, typically wrapped in a class marker such as 

50 ``Op(...)`` or ``Collective(...)``. 

51 """ 

52 return f"[DebugCall depth={self.call_depth}]" 

53 

54 def render(self, indent: int = 0) -> str: 

55 """Render this record and all of its descendants as an indented tree. 

56 

57 Args: 

58 indent (int): Current nesting level; each level is rendered with two 

59 extra leading spaces. Callers normally leave this at ``0`` and let 

60 the recursion increment it. 

61 

62 Returns: 

63 str: Newline-joined lines, one per record in the subtree. 

64 """ 

65 prefix = " " * indent 

66 lines = [f"{prefix}{self.render_self()}"] 

67 for child in self.children: 

68 lines.append(child.render(indent + 1)) 

69 return "\n".join(lines) 

70 

71 

72@dataclass 

73class OpCall(DebugCall): 

74 """Record of a DTensor operator dispatch.""" 

75 op_name: str = "" 

76 input_infos: List[TensorInfo] = field(default_factory=list) 

77 output_infos: List[TensorInfo] = field(default_factory=list) 

78 

79 def render_self(self) -> str: 

80 """Render the dispatched operator's name and its tensor signatures. 

81 

82 Each tensor is tagged ``DTensor`` or ``Tensor`` and shown by shape, so a 

83 plain tensor mixed into a DTensor op is visible at a glance. 

84 

85 Returns: 

86 str: A line of the form ``Op(<name>) inputs=[...] outputs=[...]``. 

87 """ 

88 inputs_str = ", ".join( 

89 f"{'DTensor' if t.is_dtensor else 'Tensor'}{list(t.shape)}" 

90 for t in self.input_infos 

91 ) 

92 outputs_str = ", ".join( 

93 f"{'DTensor' if t.is_dtensor else 'Tensor'}{list(t.shape)}" 

94 for t in self.output_infos 

95 ) 

96 return f"Op({self.op_name}) inputs=[{inputs_str}] outputs=[{outputs_str}]" 

97 

98 

99@dataclass 

100class CollectiveCall(DebugCall): 

101 """Record of a collective communication operation.""" 

102 collective_type: str = "" 

103 group_size: int = 0 

104 group: Optional[str] = None 

105 input_shape: Optional[Tuple[int, ...]] = None 

106 output_shape: Optional[Tuple[int, ...]] = None 

107 input_dtype: str = "" 

108 

109 def render_self(self) -> str: 

110 """Render the collective's type, its process group and its shapes. 

111 

112 ``group`` holds the traced group description when one was available at 

113 trace time; when it is ``None`` the group's rank count is printed instead. 

114 

115 Returns: 

116 str: A line of the form ``Collective(<type>) group=... input_shape=... 

117 output_shape=...``. 

118 """ 

119 group_str = f"group={self.group}" if self.group is not None else f"group_size={self.group_size}" 

120 return ( 

121 f"Collective({self.collective_type}) " 

122 f"{group_str} " 

123 f"input_shape={self.input_shape} " 

124 f"output_shape={self.output_shape}" 

125 ) 

126 

127 

128@dataclass 

129class RedistributeCall(DebugCall): 

130 """Record of a DTensor redistribute operation.""" 

131 src_placements: Optional[Tuple] = None 

132 dst_placements: Optional[Tuple] = None 

133 tensor_shape: Optional[Tuple[int, ...]] = None 

134 

135 def render_self(self) -> str: 

136 """Render the redistribution's tensor shape and its placement change. 

137 

138 Returns: 

139 str: A line of the form ``Redistribute(shape=...) <src> -> <dst>``. 

140 """ 

141 return ( 

142 f"Redistribute(shape={self.tensor_shape}) " 

143 f"{self.src_placements} -> {self.dst_placements}" 

144 ) 

145 

146 

147@dataclass 

148class AnnotateCall(DebugCall): 

149 """Record of a module boundary event (enter/exit).""" 

150 annotation: str = "" 

151 module_fqn: str = "" 

152 event_type: str = "" # "enter" or "exit" 

153 

154 def render_self(self) -> str: 

155 """Render the module boundary event and the module it belongs to. 

156 

157 Returns: 

158 str: A line of the form ``Module(<fqn>) [<enter|exit>]``. 

159 """ 

160 return f"Module({self.module_fqn}) [{self.event_type}]"