Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / memory_estimation / size.py: 87%

1# Copyright 2025 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"""Memory arithmetic""" 

16 

17from enum import Enum, auto 

18 

19from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger 

20 

21 

22class Unit(Enum): 

23 """Memory units""" 

24 

25 B = auto() 

26 KB = auto() 

27 MB = auto() 

28 GB = auto() 

29 

30 @classmethod 

31 def from_string(cls, string): 

32 """Constructor from string""" 

33 unit = string.strip().upper() 

34 if unit == "GB": 

35 return cls.GB 

36 if unit == "MB": 

37 return cls.MB 

38 if unit == "KB": 

39 return cls.KB 

40 if unit == "B": 

41 return cls.B 

42 logger.warning("Memory unit was not specified. Byte is taken") 

43 return cls.B 

44 

45 def __str__(self): 

46 if self == self.GB: 

47 return "GB" 

48 if self == self.MB: 

49 return "MB" 

50 if self == self.KB: 

51 return "KB" 

52 return "B" 

53 

54 def __lt__(self, other): 

55 if self.__class__ is other.__class__: 

56 return self.value < other.value 

57 return NotImplemented 

58 

59 

60class Memory: 

61 """Memory units""" 

62 

63 size: float 

64 unit: Unit 

65 

66 def __init__(self, size, unit): 

67 self.size = size 

68 self.unit = unit 

69 

70 @classmethod 

71 def from_string(cls, string): 

72 """Constructor from string""" 

73 if not string or not string.strip(): 

74 raise ValueError("Empty string is not a valid memory representation") 

75 numeric = "0123456789-." 

76 stripped_string = string.strip() 

77 numeric_len = 0 

78 for c in stripped_string: 

79 if c not in numeric: 

80 break 

81 numeric_len += 1 

82 if numeric_len <= 0: 

83 raise ValueError( 

84 "Memory string must start with a numeric memory size" 

85 ) 

86 memory = cls( 

87 float(stripped_string[:numeric_len]), 

88 Unit.from_string(stripped_string[numeric_len:]), 

89 ) 

90 if numeric_len >= len(stripped_string): 

91 logger.critical("Memory seems to be wrong: %s", str(memory)) 

92 return memory 

93 

94 @classmethod 

95 def from_b(cls, size: float): 

96 """Constructor from bytes""" 

97 return cls(size, Unit.B) 

98 

99 @classmethod 

100 def from_kb(cls, size: float): 

101 """Constructor from kilo bytes""" 

102 return cls(size, Unit.KB) 

103 

104 @classmethod 

105 def from_mb(cls, size: float): 

106 """Constructor from mega bytes""" 

107 return cls(size, Unit.MB) 

108 

109 @classmethod 

110 def from_gb(cls, size: float): 

111 """Constructor from giga bytes""" 

112 return cls(size, Unit.GB) 

113 

114 @classmethod 

115 def zero(cls): 

116 """Constructor for memory size 0""" 

117 return cls(0, Unit.B) 

118 

119 def __str__(self): 

120 if self.unit == Unit.GB: 

121 return f"{self.size:.2f}{self.unit}" 

122 return f"{self.size:.0f}{self.unit}" 

123 

124 def to(self, unit: Unit): 

125 """Convert to the given unit""" 

126 diff = self.unit.value - unit.value 

127 self.size *= pow(1024, diff) 

128 self.unit = unit 

129 return self 

130 

131 def to_gb(self): 

132 """Convert to GB""" 

133 return self.to(Unit.GB) 

134 

135 def to_mb(self): 

136 """Convert to MB""" 

137 return self.to(Unit.MB) 

138 

139 def to_kb(self): 

140 """Convert to KB""" 

141 return self.to(Unit.KB) 

142 

143 def to_b(self): 

144 """Convert to Byte""" 

145 return self.to(Unit.B) 

146 

147 def increase(self, mem): 

148 """Addition of 2 memory sizes""" 

149 if self.__class__ is not mem.__class__: 

150 return NotImplemented 

151 converted_mem = Memory(mem.size, mem.unit).to(self.unit) 

152 self.size += converted_mem.size 

153 return self 

154 

155 def __add__(self, mem): 

156 """Addition of 2 memory sizes""" 

157 if self.__class__ is not mem.__class__: 

158 return NotImplemented 

159 new_mem = Memory(self.size, self.unit) 

160 new_mem.to(mem.unit) 

161 new_mem.size += mem.size 

162 if new_mem.unit < self.unit: 

163 return new_mem 

164 return new_mem.to(self.unit) 

165 

166 def decrease(self, mem): 

167 """Subtraction of 2 memory sizes""" 

168 if self.__class__ is not mem.__class__: 

169 return NotImplemented 

170 converted_mem = Memory(mem.size, mem.unit).to(self.unit) 

171 self.size -= converted_mem.size 

172 return self 

173 

174 def __sub__(self, mem): 

175 """Subtraction of 2 memory sizes""" 

176 if self.__class__ is not mem.__class__: 

177 return NotImplemented 

178 new_mem = Memory(self.size, self.unit) 

179 new_mem.to(mem.unit) 

180 new_mem.size -= mem.size 

181 if new_mem.unit < self.unit: 

182 return new_mem 

183 return new_mem.to(self.unit) 

184 

185 def __lt__(self, mem): 

186 """Comparison of 2 memory sizes""" 

187 if self.__class__ is not mem.__class__: 

188 return NotImplemented 

189 converted_mem = Memory(mem.size, mem.unit).to(self.unit) 

190 return self.size < converted_mem.size 

191 

192 def __le__(self, mem): 

193 """Comparison of 2 memory sizes""" 

194 if self.__class__ is not mem.__class__: 

195 return NotImplemented 

196 converted_mem = Memory(mem.size, mem.unit).to(self.unit) 

197 return self.size <= converted_mem.size 

198 

199 def __abs__(self): 

200 """Absolute value""" 

201 if self.size < 0: 

202 self.size = -self.size 

203 return self