Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/offline_compile/offline_compile.py: 88%

129 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-07-27 10:03 +0800

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3# ------------------------------------------------------------------- 

4# ----------------------------------------------------------------------------------------------------------- 

5# Copyright (c) 2026 Huawei Technologies Co., Ltd. 

6# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 

7# CANN Open Software License Agreement Version 2.0 (the "License"). 

8# Please refer to the License for details. You may not use this file except in compliance with the License. 

9# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 

10# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. 

11# See LICENSE in the root of the software repository for the full text of the License. 

12# ----------------------------------------------------------------------------------------------------------- 

13 

14"""Offline graph compilation module.""" 

15 

16import ctypes 

17from dataclasses import dataclass, field 

18from typing import Dict, List, Optional, Tuple 

19 

20from ge._capi.pyoffline_compile_wrapper import ModelBufferDataPtr, offline_compile_lib 

21from ge.error import raise_ge_error 

22from ge.graph.graph import Graph 

23 

24 

25@dataclass 

26class GraphWithOptions: 

27 """One graph together with its build options for bundle compilation.""" 

28 

29 graph: Graph 

30 build_options: Dict[str, str] = field(default_factory=dict) 

31 

32 def __post_init__(self) -> None: 

33 if not isinstance(self.graph, Graph): 

34 raise TypeError("graph must be a Graph") 

35 self.build_options = _normalize_options(self.build_options, "build_options") 

36 

37 

38class ModelBuffer: 

39 """ModelBuffer class for offline build operations. 

40 

41 This class provides a Pythonic interface for offline build operations 

42 using the GraphEngine C API. 

43 

44 Example: 

45 >>> model = build_model(graph, {"input_format": "ND"}) 

46 >>> size = model.length 

47 >>> save_model("sample", model) 

48 """ 

49 

50 def __init__(self) -> None: 

51 """Prevent direct instantiation of ModelBuffer objects.""" 

52 self._owns_handle = False 

53 self._handle = None 

54 raise RuntimeError("ModelBuffer objects should not be created directly") 

55 

56 def __del__(self) -> None: 

57 """Clean up resources.""" 

58 if self._owns_handle: 

59 offline_compile_lib.GeApiWrapper_ModelBuffer_Destroy(self._handle) 

60 self._handle = None 

61 

62 def __copy__(self) -> None: 

63 """Copy is not supported.""" 

64 raise RuntimeError("ModelBuffer does not support copy") 

65 

66 def __deepcopy__(self, memodict) -> None: 

67 """Deep copy is not supported.""" 

68 raise RuntimeError("ModelBuffer does not support deepcopy") 

69 

70 @property 

71 def length(self) -> int: 

72 """Get model buffer length. 

73 

74 Returns: 

75 Model buffer length in bytes. 

76 """ 

77 return self.get_length() 

78 

79 @classmethod 

80 def _create_from(cls, handle: ModelBufferDataPtr) -> "ModelBuffer": 

81 """Create ModelBuffer object from C++ handle. (internal use only by e.g 

82 build_model(graph, build_options), do not use this method directly) 

83 

84 Args: 

85 handle: C++ ModelBufferData object handle. 

86 

87 Returns: 

88 ModelBuffer object. 

89 

90 Raises: 

91 ValueError: If handle is None. 

92 """ 

93 if not handle: 

94 raise ValueError("Failed to create ModelBuffer") 

95 instance = cls.__new__(cls) 

96 instance._handle = handle 

97 instance._owns_handle = True 

98 return instance 

99 

100 def get_length(self) -> int: 

101 """Get model buffer length. 

102 

103 Returns: 

104 Model buffer length in bytes. 

105 """ 

106 return int(offline_compile_lib.GeApiWrapper_ModelBuffer_GetLength(self._handle)) 

107 

108 

109def _normalize_bundle_options( 

110 graph_with_options: List[GraphWithOptions], 

111) -> List[GraphWithOptions]: 

112 if not isinstance(graph_with_options, list): 

113 raise TypeError("graph_with_options must be a list") 

114 if len(graph_with_options) <= 1: 

115 raise ValueError("graph_with_options size must be larger than 1") 

116 

117 normalized_items = [] 

118 for item in graph_with_options: 

119 if not isinstance(item, GraphWithOptions): 

120 raise TypeError("Each item in graph_with_options must be a GraphWithOptions") 

121 normalized_items.append(item) 

122 return normalized_items 

123 

124 

125def _normalize_options(options: Optional[dict], arg_name: str) -> Dict[str, str]: 

126 if options is None: 

127 return {} 

128 if not isinstance(options, dict): 

129 raise TypeError(f"{arg_name} must be a dictionary") 

130 normalized = {} 

131 for key, value in options.items(): 

132 if not isinstance(key, str) or not isinstance(value, str): 

133 raise TypeError(f"{arg_name} keys and values must be strings") 

134 normalized[key] = value 

135 return normalized 

136 

137 

138def _dict_to_c_arrays( 

139 options: Dict[str, str], 

140) -> Tuple[Optional[ctypes.Array], Optional[ctypes.Array], int]: 

141 if not options: 

142 return None, None, 0 

143 

144 size = len(options) 

145 key_array = (ctypes.c_char_p * size)() 

146 value_array = (ctypes.c_char_p * size)() 

147 for index, (key, value) in enumerate(options.items()): 

148 key_array[index] = key.encode("utf-8") 

149 value_array[index] = value.encode("utf-8") 

150 return key_array, value_array, size 

151 

152 

153def _cast_char_array(array: Optional[ctypes.Array]): 

154 if array is None: 

155 return None 

156 return ctypes.cast(array, ctypes.POINTER(ctypes.c_char_p)) 

157 

158 

159def build_initialize(global_options: Optional[dict] = None) -> None: 

160 """Initialize resources required for offline graph build. 

161 

162 Args: 

163 global_options: Optional global build configuration map. Keys and 

164 values must both be strings. 

165 

166 Raises: 

167 TypeError: If arguments have incorrect types. 

168 GeError: If GE fails to initialize build resources. 

169 """ 

170 options = _normalize_options(global_options, "global_options") 

171 key_array, value_array, size = _dict_to_c_arrays(options) 

172 ret = offline_compile_lib.GeApiWrapper_OfflineCompile_BuildInitialize( 

173 _cast_char_array(key_array), 

174 _cast_char_array(value_array), 

175 size, 

176 ) 

177 if ret != 0: 

178 raise_ge_error("BuildInitialize", ret) 

179 

180 

181def build_finalize() -> None: 

182 """Release all resources.""" 

183 offline_compile_lib.GeApiWrapper_OfflineCompile_BuildFinalize() 

184 

185 

186def build_model(graph: Graph, build_options: Optional[dict] = None) -> ModelBuffer: 

187 """Compile a graph into an offline model kept in memory. 

188 

189 Args: 

190 graph: The Graph object to compile. 

191 build_options: Optional graph-level build configuration map. Keys and 

192 values must both be strings. If an option is configured both here 

193 and in func:`build_initialize`, the value passed here takes 

194 precedence. 

195 

196 Returns: 

197 ModelBuffer containing the compiled offline model. 

198 

199 Raises: 

200 TypeError: If arguments have incorrect types. 

201 GeError: If GE fails to compile the graph. 

202 """ 

203 if not isinstance(graph, Graph): 

204 raise TypeError("graph must be a Graph") 

205 options = _normalize_options(build_options, "build_options") 

206 key_array, value_array, size = _dict_to_c_arrays(options) 

207 model_buffer_handle = ModelBufferDataPtr() 

208 ret = offline_compile_lib.GeApiWrapper_OfflineCompile_BuildModel( 

209 graph._handle, 

210 _cast_char_array(key_array), 

211 _cast_char_array(value_array), 

212 size, 

213 ctypes.byref(model_buffer_handle), 

214 ) 

215 if ret != 0: 

216 raise_ge_error("BuildModel", ret) 

217 return ModelBuffer._create_from(model_buffer_handle) 

218 

219 

220def save_model(output_file: str, model: ModelBuffer) -> None: 

221 """Serialize an in-memory offline model to an ``.om`` file. 

222 

223 Args: 

224 output_file: Base name of the output model file. The generated offline 

225 model file name automatically ends with the ``.om`` suffix. If the 

226 OM file name contains the operating system and architecture, the 

227 OM file can only be used in the runtime environment for that 

228 operating system and architecture. 

229 model: Offline model buffer. 

230 

231 Raises: 

232 TypeError: If arguments have incorrect types. 

233 GeError: If GE fails to save the model. 

234 """ 

235 if not isinstance(output_file, str): 

236 raise TypeError("output_file must be a string") 

237 if not isinstance(model, ModelBuffer): 

238 raise TypeError("model must be a ModelBuffer") 

239 ret = offline_compile_lib.GeApiWrapper_OfflineCompile_SaveModel(output_file.encode("utf-8"), model._handle) 

240 if ret != 0: 

241 raise_ge_error("SaveModel", ret, output_file=output_file) 

242 

243 

244def bundle_build_model(graph_with_options: List[GraphWithOptions]) -> ModelBuffer: 

245 """Compile a group of graphs into a bundle offline model. 

246 

247 Args: 

248 graph_with_options: A list of GraphWithOptions items to 

249 compile together. Each item contains one graph and its build 

250 options. The list must contain more than one graph. 

251 

252 Returns: 

253 ModelBuffer containing the compiled bundle model. 

254 

255 Raises: 

256 TypeError: If arguments have incorrect types. 

257 ValueError: If fewer than two graphs are provided. 

258 GeError: If GE fails to compile the bundle model. 

259 """ 

260 normalized_items = _normalize_bundle_options(graph_with_options) 

261 graph_count = len(normalized_items) 

262 graph_handles = (ctypes.c_void_p * graph_count)(*[item.graph._handle for item in normalized_items]) 

263 size_array = (ctypes.c_int * graph_count)() 

264 # Hold references to each graph's option ctypes arrays so pointers in keys/values stay valid for the C call. 

265 option_buffers: List[Tuple[Optional[ctypes.Array], Optional[ctypes.Array]]] = [] 

266 keys = (ctypes.POINTER(ctypes.c_char_p) * graph_count)() 

267 values = (ctypes.POINTER(ctypes.c_char_p) * graph_count)() 

268 for index, item in enumerate(normalized_items): 

269 key_array, value_array, size = _dict_to_c_arrays(item.build_options) 

270 option_buffers.append((key_array, value_array)) 

271 keys[index] = _cast_char_array(key_array) 

272 values[index] = _cast_char_array(value_array) 

273 size_array[index] = size 

274 

275 model_buffer_handle = ModelBufferDataPtr() 

276 ret = offline_compile_lib.GeApiWrapper_OfflineCompile_BundleBuildModel( 

277 graph_handles, 

278 keys, 

279 values, 

280 size_array, 

281 graph_count, 

282 ctypes.byref(model_buffer_handle), 

283 ) 

284 if ret != 0: 

285 raise_ge_error("BundleBuildModel", ret, graph_count=graph_count) 

286 return ModelBuffer._create_from(model_buffer_handle) 

287 

288 

289def bundle_save_model(output_file: str, model: ModelBuffer) -> None: 

290 """Serialize an in-memory bundle model to an ``.om`` file. 

291 

292 Args: 

293 output_file: Base name of the output model file. The generated offline 

294 model file name automatically ends with the ``.om`` suffix. If the 

295 OM file name contains the operating system and architecture, the 

296 OM file can only be used in the runtime environment for that 

297 operating system and architecture. 

298 model: Bundle model buffer. 

299 

300 Raises: 

301 TypeError: If arguments have incorrect types. 

302 GeError: If GE fails to save the bundle model. 

303 """ 

304 if not isinstance(output_file, str): 

305 raise TypeError("output_file must be a string") 

306 if not isinstance(model, ModelBuffer): 

307 raise TypeError("model must be a ModelBuffer") 

308 ret = offline_compile_lib.GeApiWrapper_OfflineCompile_BundleSaveModel(output_file.encode("utf-8"), model._handle) 

309 if ret != 0: 

310 raise_ge_error("BundleSaveModel", ret, output_file=output_file)