Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/passes/runtime.py: 45%

321 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# Copyright (c) 2026 Huawei Technologies Co., Ltd. 

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

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

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

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

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

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

11# ----------------------------------------------------------------------------------------------------------- 

12 

13"""Runtime for GE Python pass bridge/native artifact set.""" 

14 

15import importlib.util 

16import json 

17import os 

18import secrets 

19import shutil 

20import subprocess 

21import sys 

22import sysconfig 

23from dataclasses import dataclass 

24from pathlib import Path 

25from types import ModuleType 

26from typing import Dict, Iterable, List, Optional, Tuple 

27 

28from ge._internal.artifact_utils import ( 

29 PythonArtifact, 

30 current_platform_tag, 

31 current_python_tag, 

32) 

33 

34from ._artifact_utils import ( 

35 BRIDGE_ABI_VERSION, 

36 NATIVE_MODULE_NAME, 

37 artifacts_root, 

38 find_prebuilt_artifact, 

39 iter_artifacts, 

40 load_artifact_from_dir, 

41 load_native_module, 

42) 

43 

44_FALLBACK_RESOURCES_MODULE = "_sources.py" 

45_MATERIALIZED_CODEGEN_DIR = "fallback_sources" 

46 

47 

48@dataclass(frozen=True) 

49class PythonBuildInfo: 

50 tag: str 

51 executable: str 

52 version: str 

53 include_dir: Path 

54 library: Optional[Path] 

55 pybind_include: Optional[Path] 

56 

57 

58@dataclass(frozen=True) 

59class _BuildInputs: 

60 config: dict 

61 root: Path 

62 src_dir: Path 

63 include_dir: Path 

64 

65 

66@dataclass(frozen=True) 

67class _CompiledArtifactSet: 

68 artifact_paths: Dict[str, Path] 

69 python_info: PythonBuildInfo 

70 

71 

72def _codegen_root() -> Path: 

73 return Path(__file__).resolve().parent / "fallback_codegen" 

74 

75 

76def _fallback_artifact_dir() -> Path: 

77 return artifacts_root() / f"{current_python_tag()}-{current_platform_tag()}" 

78 

79 

80def _resolve_pybind_include() -> Optional[Path]: 

81 try: 

82 import pybind11 

83 

84 include = Path(pybind11.get_include()) 

85 if not include.is_dir(): 

86 return None 

87 return include 

88 except Exception: 

89 return None 

90 

91 

92def _resolve_libpython_dirs() -> List[Path]: 

93 libdirs: List[Path] = [] 

94 for item in [ 

95 Path(sys.prefix) / "lib", 

96 Path(sys.exec_prefix) / "lib", 

97 Path(sys.executable).resolve().parent.parent / "lib", 

98 sysconfig.get_config_var("LIBDIR") or "", 

99 ]: 

100 path = Path(item) 

101 if item and path not in libdirs: 

102 libdirs.append(path) 

103 return libdirs 

104 

105 

106def _resolve_python_library(lib_version: str) -> Optional[Path]: 

107 candidates = [ 

108 sysconfig.get_config_var("LDLIBRARY"), 

109 sysconfig.get_config_var("INSTSONAME"), 

110 sysconfig.get_config_var("LIBRARY"), 

111 f"libpython{lib_version}.so.1.0" if lib_version else "", 

112 f"libpython{lib_version}.so" if lib_version else "", 

113 ] 

114 seen: List[str] = [] 

115 for candidate in candidates: 

116 if candidate and candidate not in seen: 

117 seen.append(candidate) 

118 libdirs = _resolve_libpython_dirs() 

119 matches: List[Path] = [] 

120 for name in seen: 

121 for directory in libdirs: 

122 candidate_path = directory / name 

123 if candidate_path.exists(): 

124 matches.append(candidate_path) 

125 shared = next((item for item in matches if ".so" in item.name), None) 

126 library = shared or (matches[0] if matches else None) 

127 if library is not None and not library.is_file(): 

128 library = None 

129 return library 

130 

131 

132def _query_current_python_build_info() -> Optional[PythonBuildInfo]: 

133 pybind_include = _resolve_pybind_include() 

134 version = ( 

135 sysconfig.get_config_var("VERSION") 

136 or f"{sys.version_info.major}.{sys.version_info.minor}" 

137 ) 

138 lib_version = ( 

139 version + ("m" if sys.version_info[:2] <= (3, 7) else "") if version else "" 

140 ) 

141 library = _resolve_python_library(lib_version) 

142 include_dir = Path( 

143 sysconfig.get_path("include") or sysconfig.get_config_var("INCLUDEPY") or "" 

144 ) 

145 if not include_dir.is_dir(): 

146 return None 

147 return PythonBuildInfo( 

148 tag=current_python_tag(), 

149 executable=sys.executable, 

150 version=sys.version.split()[0], 

151 include_dir=include_dir, 

152 library=library, 

153 pybind_include=pybind_include, 

154 ) 

155 

156 

157def _load_codegen_config(root: Path) -> Optional[dict]: 

158 config_path = root / "build_config.json" 

159 if not config_path.is_file(): 

160 return None 

161 try: 

162 return json.loads(config_path.read_text(encoding="utf-8")) 

163 except (OSError, json.JSONDecodeError): 

164 return None 

165 

166 

167def _build_inputs_from_root(root: Path, config: dict) -> Optional[_BuildInputs]: 

168 src_dir = root / "src" 

169 include_dir = root / "include" 

170 if not src_dir.is_dir() or not include_dir.is_dir(): 

171 return None 

172 return _BuildInputs( 

173 config=config, root=root, src_dir=src_dir, include_dir=include_dir 

174 ) 

175 

176 

177def _resolve_build_inputs() -> Optional[_BuildInputs]: 

178 root = _codegen_root() 

179 config = _load_codegen_config(root) 

180 if config is None: 

181 return None 

182 return _build_inputs_from_root(root, config) 

183 

184 

185def _load_fallback_resources_module(module_path: Path) -> Optional[ModuleType]: 

186 spec = importlib.util.spec_from_file_location( 

187 "_ge_pass_fallback_resources", module_path 

188 ) 

189 if spec is None or spec.loader is None: 

190 return None 

191 module = importlib.util.module_from_spec(spec) 

192 try: 

193 spec.loader.exec_module(module) 

194 except Exception: 

195 return None 

196 return module 

197 

198 

199def _materialize_fallback_resources( 

200 codegen_dir: Path, config: dict, work_dir: Path 

201) -> Optional[_BuildInputs]: 

202 module_path = codegen_dir / _FALLBACK_RESOURCES_MODULE 

203 if not module_path.is_file(): 

204 return None 

205 module = _load_fallback_resources_module(module_path) 

206 if module is None: 

207 return None 

208 materialize = getattr(module, "materialize", None) 

209 if not callable(materialize): 

210 return None 

211 

212 resource_root = work_dir / _MATERIALIZED_CODEGEN_DIR 

213 _remove_tree_quietly(resource_root) 

214 try: 

215 materialize(resource_root) 

216 except Exception: 

217 return None 

218 return _build_inputs_from_root(resource_root, config) 

219 

220 

221def _resolve_fallback_build_inputs(work_dir: Path) -> Optional[_BuildInputs]: 

222 root = _codegen_root() 

223 config = _load_codegen_config(root) 

224 if config is None: 

225 return None 

226 if (root / _FALLBACK_RESOURCES_MODULE).is_file(): 

227 return _materialize_fallback_resources(root, config, work_dir) 

228 return None 

229 

230 

231def _resolve_cann_paths_from_root(root: Path) -> Optional[Tuple[Path, Path, Path]]: 

232 if not root.is_dir(): 

233 return None 

234 include_dir = root / "include" 

235 lib64_dir = root / "lib64" 

236 pkg_inc_dir = root / "pkg_inc" 

237 if include_dir.is_dir() and lib64_dir.is_dir() and pkg_inc_dir.is_dir(): 

238 return include_dir, lib64_dir, pkg_inc_dir 

239 return None 

240 

241 

242def _resolve_cann_paths() -> Optional[Tuple[Path, Path, Path]]: 

243 """ 

244 1. cann run package structure 

245 - xxx/Ascend/cann/ <-- env:ASCEND_HOME_PATH 

246 - include/ 

247 - lib64/ 

248 - pkg_inc/ 

249 - python/site-packages/ge/passes/runtime.py 

250 

251 2. if current Python file is not in the package structure above, you need to source set_env.bash 

252 before execution 

253 """ 

254 parent_paths = Path(__file__).resolve().parents 

255 if len(parent_paths) > 4: 

256 cann_paths = _resolve_cann_paths_from_root(parent_paths[4]) 

257 if cann_paths is not None: 

258 return cann_paths 

259 

260 cann_root = os.environ.get("ASCEND_HOME_PATH", "").strip() 

261 if cann_root: 

262 return _resolve_cann_paths_from_root(Path(cann_root)) 

263 return None 

264 

265 

266def _replace_placeholders(value: str, replacements: Dict[str, str]) -> str: 

267 resolved = value 

268 for placeholder, replacement in replacements.items(): 

269 resolved = resolved.replace(placeholder, replacement) 

270 return resolved 

271 

272 

273def _resolve_build_config( 

274 config: dict, python_info: PythonBuildInfo, fallback_root: Path 

275) -> dict: 

276 if python_info.pybind_include is None: 

277 raise RuntimeError( 

278 "Cannot resolve pybind11 include. Please install pybind11 for this Python." 

279 ) 

280 if python_info.library is None: 

281 raise RuntimeError( 

282 "Cannot resolve libpython shared library for current Python." 

283 ) 

284 cann_paths = _resolve_cann_paths() 

285 if cann_paths is None: 

286 raise RuntimeError("Cannot resolve CANN include/lib64/pkg_inc.") 

287 cann_include, cann_lib64, cann_pkg_inc = cann_paths 

288 replacements = { 

289 "@PYTHON_INCLUDE@": os.fspath(python_info.include_dir), 

290 "@PYTHON_LIBRARY@": os.fspath(python_info.library), 

291 "@PYTHON_LIBDIR@": os.fspath(python_info.library.parent), 

292 "@PYBIND11_INCLUDE@": os.fspath(python_info.pybind_include), 

293 "@CANN_INCLUDE_DIR@": os.fspath(cann_include), 

294 "@CANN_PKG_INC@": os.fspath(cann_pkg_inc), 

295 "@CANN_LIB64@": os.fspath(cann_lib64), 

296 "@FALLBACK_ROOT@": os.fspath(fallback_root), 

297 } 

298 

299 def resolve_obj(obj): 

300 if isinstance(obj, str): 

301 return _replace_placeholders(obj, replacements) 

302 if isinstance(obj, list): 

303 return [resolve_obj(item) for item in obj] 

304 if isinstance(obj, dict): 

305 return {key: resolve_obj(value) for key, value in obj.items()} 

306 return obj 

307 

308 return resolve_obj(config) 

309 

310 

311def _run_command(command: List[str]) -> None: 

312 completed = subprocess.run( 

313 command, 

314 text=True, 

315 stdout=subprocess.PIPE, 

316 stderr=subprocess.STDOUT, 

317 check=False, 

318 ) 

319 if completed.returncode != 0: 

320 raise RuntimeError( 

321 "Command failed: {}\n{}".format(" ".join(command), completed.stdout) 

322 ) 

323 

324 

325def _iter_target_sources( 

326 target_name: str, build_inputs: _BuildInputs 

327) -> Iterable[Path]: 

328 source_dir = build_inputs.root / "src" / target_name 

329 if not source_dir.is_dir(): 

330 raise RuntimeError(f"Cannot find fallback source dir: {source_dir}") 

331 sources = sorted(source_dir.glob("*.cc")) 

332 if not sources: 

333 raise RuntimeError(f"No fallback sources found under: {source_dir}") 

334 yield from sources 

335 

336 

337def _target_compile_base_args(target_config: dict) -> List[str]: 

338 compile_args: List[str] = [] 

339 for key in ("cxx_defines", "cxx_includes", "cxx_flags"): 

340 args = target_config.get(key) 

341 if not isinstance(args, list): 

342 raise RuntimeError(f"Missing fallback {key}.") 

343 compile_args.extend(args) 

344 return compile_args 

345 

346 

347def _target_link_args(target_config: dict) -> List[str]: 

348 link_args = target_config.get("link_args") 

349 if not isinstance(link_args, list): 

350 raise RuntimeError("Missing fallback link args.") 

351 return link_args 

352 

353 

354def _compile_target_objects( 

355 target_name: str, target_config: dict, build_inputs: _BuildInputs, work_dir: Path 

356) -> List[Path]: 

357 compiler = os.environ.get("CXX") or "c++" 

358 obj_dir = work_dir / f"{target_name}_obj" 

359 obj_dir.mkdir(parents=True, exist_ok=True) 

360 base_args = _target_compile_base_args(target_config) 

361 objects: List[Path] = [] 

362 for index, source_path in enumerate( 

363 _iter_target_sources(target_name, build_inputs) 

364 ): 

365 object_path = obj_dir / f"{index}_{source_path.stem}.o" 

366 command = ( 

367 [compiler] 

368 + base_args 

369 + ["-c", os.fspath(source_path), "-o", os.fspath(object_path)] 

370 ) 

371 _run_command(command) 

372 objects.append(object_path) 

373 return objects 

374 

375 

376def _link_target(target_config: dict, objects: List[Path], work_dir: Path) -> Path: 

377 compiler = os.environ.get("CXX") or "c++" 

378 output = work_dir / target_config["output"] 

379 command = [compiler, "-shared", "-o", os.fspath(output)] 

380 command.extend(os.fspath(obj) for obj in objects) 

381 command.extend(_target_link_args(target_config)) 

382 _run_command(command) 

383 return output 

384 

385 

386def _build_target( 

387 target_name: str, target_config: dict, build_inputs: _BuildInputs, work_dir: Path 

388) -> Path: 

389 objects = _compile_target_objects( 

390 target_name, target_config, build_inputs, work_dir 

391 ) 

392 return _link_target(target_config, objects, work_dir) 

393 

394 

395def _build_targets( 

396 config: dict, build_inputs: _BuildInputs, work_dir: Path 

397) -> Dict[str, Path]: 

398 targets = config.get("targets", {}) 

399 if not isinstance(targets, dict) or not targets: 

400 raise RuntimeError("Missing fallback target configs.") 

401 built_targets: Dict[str, Path] = {} 

402 for target_name, target_config in targets.items(): 

403 if not isinstance(target_config, dict): 

404 raise RuntimeError(f"Invalid fallback target config: {target_name}") 

405 output_name = target_config.get("output") 

406 if not isinstance(output_name, str) or not output_name: 

407 raise RuntimeError(f"Missing fallback target output: {target_name}") 

408 for key in ("cxx_defines", "cxx_includes", "cxx_flags", "link_args"): 

409 if not isinstance(target_config.get(key), list): 

410 raise RuntimeError(f"Missing fallback target {key}: {target_name}") 

411 built_targets[output_name] = _build_target( 

412 target_name, target_config, build_inputs, work_dir 

413 ) 

414 return built_targets 

415 

416 

417def _compile_artifact_set( 

418 build_inputs: _BuildInputs, work_dir: Path 

419) -> _CompiledArtifactSet: 

420 python_info = _query_current_python_build_info() 

421 if python_info is None: 

422 raise RuntimeError("Cannot resolve current Python build info.") 

423 config = _resolve_build_config(build_inputs.config, python_info, build_inputs.root) 

424 work_dir.mkdir(parents=True, exist_ok=True) 

425 artifact_paths = _build_targets(config, build_inputs, work_dir) 

426 return _CompiledArtifactSet(artifact_paths=artifact_paths, python_info=python_info) 

427 

428 

429def _atomic_publish_file(src: Path, dst: Path) -> None: 

430 dst.parent.mkdir(parents=True, exist_ok=True) 

431 os.replace(src, dst) 

432 

433 

434def _atomic_write(path: Path, data: bytes) -> None: 

435 tmp = path.with_suffix(path.suffix + f".tmp.{os.getpid()}.{secrets.token_hex(4)}") 

436 tmp.write_bytes(data) 

437 os.replace(tmp, path) 

438 

439 

440def _remove_tree_quietly(path: Path) -> None: 

441 shutil.rmtree(path, ignore_errors=True) 

442 

443 

444def _make_unique_work_dir(final_dir: Path) -> Path: 

445 return final_dir / f".work.{os.getpid()}.{secrets.token_hex(4)}" 

446 

447 

448def _format_optional_path(path: Optional[Path]) -> str: 

449 return os.fspath(path) if path is not None else "not-found" 

450 

451 

452def _build_manifest_json(python_info: PythonBuildInfo) -> bytes: 

453 manifest = { 

454 "python_tag": current_python_tag(), 

455 "python_version": python_info.version, 

456 "platform": current_platform_tag(), 

457 "bridge_abi": BRIDGE_ABI_VERSION, 

458 "build_python": { 

459 "executable": python_info.executable, 

460 "version": python_info.version, 

461 "include": os.fspath(python_info.include_dir), 

462 "libpython": _format_optional_path(python_info.library), 

463 "pybind11_include": _format_optional_path(python_info.pybind_include), 

464 "link_python": True, 

465 }, 

466 "artifacts": { 

467 "bridge": "libge_python_pass_bridge.so", 

468 "native": "_ge_pass_native.so", 

469 }, 

470 } 

471 return (json.dumps(manifest, indent=2, sort_keys=True) + "\n").encode("utf-8") 

472 

473 

474def _iter_load_candidates() -> Iterable[PythonArtifact]: 

475 prebuilt = find_prebuilt_artifact() 

476 if prebuilt is not None: 

477 yield prebuilt 

478 

479 

480def _format_missing_artifact_error(load_errors: List[str]) -> str: 

481 python_tag = current_python_tag() 

482 platform_tag = current_platform_tag() 

483 discovered_artifacts = sorted( 

484 f"{artifact.python_tag}-{artifact.platform_tag}-abi{artifact.abi}" 

485 for artifact in iter_artifacts() 

486 ) 

487 discovered_text = ( 

488 ", ".join(discovered_artifacts) if discovered_artifacts else "none" 

489 ) 

490 expected_wheel = f"ge_py_pass_bridge-*-{python_tag}-{python_tag}-*.whl" 

491 load_error_text = "; ".join(load_errors) if load_errors else "none" 

492 return ( 

493 "Failed to load GE Python pass native artifact for runtime " 

494 f"python tag '{python_tag}', platform '{platform_tag}', " 

495 f"bridge ABI {BRIDGE_ABI_VERSION}. " 

496 f"Searched artifact root: {artifacts_root()}. " 

497 f"Discovered valid artifacts: {discovered_text}. " 

498 f"Load errors: {load_error_text}. " 

499 "Please install the native artifact wheel that matches this Python " 

500 f"runtime, for example '{expected_wheel}', or reinstall the CANN run " 

501 "package that contains the matching ge_py_pass_bridge wheel." 

502 ) 

503 

504 

505def ensure_native_module() -> ModuleType: 

506 loaded_module = sys.modules.get(NATIVE_MODULE_NAME) 

507 if loaded_module is not None: 

508 return loaded_module 

509 

510 load_errors: List[str] = [] 

511 native: Optional[ModuleType] = None 

512 for artifact in _iter_load_candidates(): 

513 try: 

514 native = load_native_module(artifact.native_path) 

515 break 

516 except Exception as err: 

517 load_errors.append( 

518 f"load native artifact '{artifact.native_path}' failed: {err}" 

519 ) 

520 continue 

521 

522 if native is None: 

523 try: 

524 artifact = run_fallback_codegen() 

525 except Exception as err: 

526 load_errors.append(f"fallback codegen failed: {err}") 

527 else: 

528 try: 

529 native = load_native_module(artifact.native_path) 

530 except Exception as err: 

531 load_errors.append( 

532 f"load fallback native artifact '{artifact.native_path}' failed: {err}" 

533 ) 

534 

535 if native is None: 

536 raise ImportError(_format_missing_artifact_error(load_errors)) 

537 return native 

538 

539 

540def run_fallback_codegen() -> PythonArtifact: 

541 final_dir = _fallback_artifact_dir() 

542 final_dir.mkdir(parents=True, exist_ok=True) 

543 work_dir = _make_unique_work_dir(final_dir) 

544 try: 

545 build_inputs = _resolve_fallback_build_inputs(work_dir) 

546 if build_inputs is None: 

547 raise RuntimeError( 

548 "Fallback codegen unavailable: codegen resources are invalid or unavailable." 

549 ) 

550 compiled = _compile_artifact_set(build_inputs, work_dir) 

551 for filename, path in compiled.artifact_paths.items(): 

552 _atomic_publish_file(path, final_dir / filename) 

553 _atomic_write( 

554 final_dir / "manifest.json", _build_manifest_json(compiled.python_info) 

555 ) 

556 finally: 

557 _remove_tree_quietly(work_dir) 

558 

559 artifact = load_artifact_from_dir(final_dir) 

560 if artifact is None: 

561 raise RuntimeError( 

562 f"Fallback codegen completed but published artifact is incomplete: {final_dir}" 

563 ) 

564 return artifact