Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / shmem / _runtime.py: 85%

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"""Lazy access to the private Torch SHMEM Runtime binding.""" 

16 

17from __future__ import annotations 

18 

19__all__ = [ 

20 "_load_native", 

21 "_torch_modules", 

22] 

23 

24import ctypes 

25import importlib.util 

26import sys 

27from pathlib import Path 

28from typing import Any, Final 

29 

30from hyper_multicore._loader import ( 

31 NativeComponentUnavailableError, load_torch_npu, native_loading, native_context, native_ready, 

32 require_multicore_environment, validate_native_compatibility, 

33 select_native_shmem, 

34) 

35from hyper_multicore._native_manifest import SHMEM_DEPENDENCIES, resource_path 

36 

37 

38_NATIVE_MODULE: Final = "hyper_multicore_shmem_torch" 

39 

40_native: Any | None = None 

41_native_libraries: list[Any] = [] 

42 

43 

44def _load_native() -> Any: 

45 """Import and cache the private binding after the first successful load.""" 

46 global _native # pylint: disable=global-statement 

47 

48 if native_ready("shmem"): 

49 return _native 

50 require_multicore_environment() 

51 load_torch_npu() 

52 selected = validate_native_compatibility() 

53 backend, resources = select_native_shmem(selected) 

54 path = resource_path(selected["root"], selected["variant"]["shmem_so"][backend]) 

55 with native_loading("shmem") as needed: 

56 if needed: 

57 for dependency in SHMEM_DEPENDENCIES: 

58 library = resource_path(selected["root"], f"{resources['shmem_root']}/{dependency}") 

59 _native_libraries.append(ctypes.CDLL(str(library), mode=ctypes.RTLD_LOCAL)) 

60 cached = sys.modules.get(_NATIVE_MODULE) 

61 if cached is not None: 

62 if Path(getattr(cached, "__file__", "")).resolve() != path: 

63 raise NativeComponentUnavailableError(f"Unexpected SHMEM module already loaded: {cached!r}") 

64 _native = cached 

65 else: 

66 try: 

67 spec = importlib.util.spec_from_file_location(_NATIVE_MODULE, path) 

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

69 raise ImportError(f"No extension loader for {path}") 

70 native = importlib.util.module_from_spec(spec) 

71 sys.modules[_NATIVE_MODULE] = native 

72 spec.loader.exec_module(native) 

73 _native = native 

74 except BaseException as error: 

75 sys.modules.pop(_NATIVE_MODULE, None) 

76 raise NativeComponentUnavailableError( 

77 f"[HYPER-MULTICORE-SHMEM-LOAD-FAILED] path={path} {native_context()}: {error}" 

78 ) from error 

79 return _native 

80 

81 

82def _torch_modules() -> tuple[Any, Any]: 

83 """Load Torch only when the internal SHMEM boundary is used.""" 

84 import torch # pylint: disable=C0415 

85 import torch.distributed as dist # pylint: disable=C0415 

86 

87 return torch, dist