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 return f"[DebugCall depth={self.call_depth}]" 

44 

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

46 prefix = " " * indent 

47 lines = [f"{prefix}{self._render_self()}"] 

48 for child in self.children: 

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

50 return "\n".join(lines) 

51 

52 

53@dataclass 

54class OpCall(DebugCall): 

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

56 op_name: str = "" 

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

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

59 

60 def _render_self(self) -> str: 

61 inputs_str = ", ".join( 

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

63 for t in self.input_infos 

64 ) 

65 outputs_str = ", ".join( 

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

67 for t in self.output_infos 

68 ) 

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

70 

71 

72@dataclass 

73class CollectiveCall(DebugCall): 

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

75 collective_type: str = "" 

76 group_size: int = 0 

77 group: Optional[str] = None 

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

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

80 input_dtype: str = "" 

81 

82 def _render_self(self) -> str: 

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

84 return ( 

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

86 f"{group_str} " 

87 f"input_shape={self.input_shape} " 

88 f"output_shape={self.output_shape}" 

89 ) 

90 

91 

92@dataclass 

93class RedistributeCall(DebugCall): 

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

95 src_placements: Optional[Tuple] = None 

96 dst_placements: Optional[Tuple] = None 

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

98 

99 def _render_self(self) -> str: 

100 return ( 

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

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

103 ) 

104 

105 

106@dataclass 

107class AnnotateCall(DebugCall): 

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

109 annotation: str = "" 

110 module_fqn: str = "" 

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

112 

113 def _render_self(self) -> str: 

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