Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / simulator / utils.py: 98%

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"""Helpers used by the pipeline simulator: numeric coercion, colouring, decorators.""" 

16import time 

17from functools import wraps 

18from typing import Any, Callable, List, Tuple, Union 

19 

20import numpy as np 

21from matplotlib import colors 

22 

23from hyper_parallel.auto_parallel.sapp_ppb.utils.logger import logger 

24 

25ScalarOrMatrix = Union[int, float, List[List[Union[int, float]]], Tuple[Tuple[Union[int, float], ...], ...]] 

26 

27 

28def format_2d_inputs(a: ScalarOrMatrix, raw: int, col: int) -> np.ndarray: 

29 """Coerce ``a`` into a 2-D :class:`numpy.ndarray` of shape ``(raw, col)``. 

30 

31 Args: 

32 a: Scalar broadcast to ``(raw, col)``, a flat sequence treated as one row, 

33 or a nested sequence interpreted as a 2-D matrix. 

34 raw: Number of rows when broadcasting a scalar. 

35 col: Number of columns when broadcasting a scalar. 

36 

37 Returns: 

38 A 2-D array matching the supplied data. 

39 

40 Raises: 

41 ValueError: If ``a`` does not match any of the supported shapes. 

42 """ 

43 if isinstance(a, (int, float)): 

44 return np.broadcast_to(a, (raw, col)) 

45 if isinstance(a, (list, tuple)): 

46 if all(isinstance(item, (list, tuple)) for item in a): 

47 return np.array(a) 

48 if all(isinstance(item, (int, float)) for item in a): 

49 return np.array([a]) 

50 raise ValueError(f"Unsupported inputs: {a}") 

51 raise ValueError(f"Unsupported inputs: {a}") 

52 

53 

54def apply_color(target_list: list, c: List[str]) -> list: 

55 """Wrap each element of ``target_list`` with an ANSI colour escape from ``c``. 

56 

57 Args: 

58 target_list: Values to colour (floats are formatted to four decimals). 

59 c: One ANSI colour code per target element. 

60 

61 Returns: 

62 The same list with each element wrapped in the matching colour escape. 

63 """ 

64 for i, target in enumerate(target_list): 

65 target = f'{target:.4f}' if isinstance(target, float) else target 

66 target_list[i] = f"\033[{c[i]}m{target}\033[0m" 

67 return target_list 

68 

69 

70def apply_format(target_list: list) -> str: 

71 """Join a sequence of pre-coloured values into the single-line bubble report. 

72 

73 Args: 

74 target_list: Coloured strings produced by :func:`apply_color`. 

75 

76 Returns: 

77 The formatted single-line string. 

78 """ 

79 s = f'{target_list[0]:^22}' 

80 symbol = ['=', '+', '+', '+', '+', '+'] 

81 for i in range(len(target_list) - 1): 

82 s = f'{s}{symbol[i]}{target_list[i + 1]:^22}' 

83 return s 

84 

85 

86def color_mix(c1: Any, c2: Any, w1: float = 0.5, w2: float = 0.5) -> Tuple[float, float, float, float]: 

87 """Blend two matplotlib colours with weights ``w1`` and ``w2``. 

88 

89 Args: 

90 c1: First colour in any format understood by :func:`matplotlib.colors.to_rgba`. 

91 c2: Second colour. 

92 w1: Weight for ``c1``. Default: 0.5. 

93 w2: Weight for ``c2``. Default: 0.5. 

94 

95 Returns: 

96 A ``(r, g, b, a)`` tuple with values in ``[0, 1]``. 

97 """ 

98 rgb = (np.array(colors.to_rgba(c1, 1)) * w1 + np.array(colors.to_rgba(c2, 1)) * w2) / (w1 + w2) 

99 return colors.to_rgba(rgb) 

100 

101 

102def dfs_builder(comm: bool = False) -> Callable[[Callable[..., Any]], Callable[..., Any]]: 

103 """Build a decorator that guards a DFS visit against re-entry and unmet dependencies. 

104 

105 Args: 

106 comm: When ``True``, use the communication-aware ``depend_pre``/``depend_left`` 

107 attributes; otherwise use the compute-only ``pre``/``left`` attributes. 

108 

109 Returns: 

110 A decorator wrapping a DFS visit method on :class:`BlockSim`-like objects. 

111 """ 

112 

113 def decorator(func: Callable[..., Any]) -> Callable[..., Any]: 

114 """Attach the DFS visit guards to ``func``.""" 

115 

116 @wraps(func) 

117 def wrapper(*args: Any, **kwargs: Any) -> Any: 

118 """Run ``func`` exactly once per node after asserting dependencies.""" 

119 self = args[0] 

120 pre, left = (self.depend_pre, self.depend_left) if comm else (self.pre, self.left) 

121 if self.finish: 

122 return None 

123 if pre is None or left is None: 

124 raise NotImplementedError 

125 if self.in_queue: 

126 raise ValueError("Dependency loop detected during DFS traversal") 

127 self.in_queue = True 

128 res = func(*args, **kwargs) 

129 self.finish = True 

130 self.in_queue = False 

131 return res 

132 return wrapper 

133 

134 return decorator 

135 

136 

137def timer(func: Callable[..., Any]) -> Callable[..., Any]: 

138 """Log the wall-clock time a function takes. 

139 

140 Args: 

141 func: Callable to time. 

142 

143 Returns: 

144 A wrapper that logs the elapsed time at INFO level after ``func`` returns. 

145 """ 

146 

147 @wraps(func) 

148 def wrapper(*args: Any, **kwargs: Any) -> Any: 

149 """Time one call to ``func`` and log the elapsed wall clock.""" 

150 t0 = time.time() 

151 res = func(*args, **kwargs) 

152 t1 = time.time() - t0 

153 logger.info("function `%s` time used: %.4f s", func.__name__, t1) 

154 return res 

155 

156 return wrapper