|
| 1 | +import ctypes |
| 2 | +import os |
| 3 | +from typing import Iterable, List |
| 4 | + |
| 5 | + |
| 6 | +def _candidate_prefixes(path: str) -> List[str]: |
| 7 | + """ |
| 8 | + Return HPCC install prefixes to search for libs. |
| 9 | + Prefer HPCC_PATH; if absent and explicitly opted-in, fall back to /opt/hpcc. |
| 10 | + """ |
| 11 | + prefixes: List[str] = [] |
| 12 | + if path: |
| 13 | + prefixes.append(path) |
| 14 | + |
| 15 | + seen = set() |
| 16 | + unique: List[str] = [] |
| 17 | + for p in prefixes: |
| 18 | + if p and p not in seen: |
| 19 | + seen.add(p) |
| 20 | + unique.append(p) |
| 21 | + return unique |
| 22 | + |
| 23 | + |
| 24 | +def _try_load(paths: Iterable[str], name: str) -> bool: |
| 25 | + """Try to load a shared library from given paths or system search path.""" |
| 26 | + for path in paths: |
| 27 | + full = os.path.join(path, "lib", name) |
| 28 | + if os.path.exists(full): |
| 29 | + try: |
| 30 | + ctypes.CDLL(full, mode=ctypes.RTLD_GLOBAL) |
| 31 | + return True |
| 32 | + except OSError: |
| 33 | + # Try next candidate |
| 34 | + continue |
| 35 | + # Last resort: rely on loader search path |
| 36 | + try: |
| 37 | + ctypes.CDLL(name, mode=ctypes.RTLD_GLOBAL) |
| 38 | + return True |
| 39 | + except OSError: |
| 40 | + return False |
| 41 | + |
| 42 | + |
| 43 | +def preload_hpcc() -> None: |
| 44 | + """ |
| 45 | + Best-effort preload of key HPCC runtime libs with RTLD_GLOBAL. |
| 46 | +
|
| 47 | + This mirrors the behavior of torch's HPCC build that loads libtorch_global_deps.so, |
| 48 | + but avoids introducing a hard torch dependency. All failures are swallowed. |
| 49 | + """ |
| 50 | + hpcc_path = os.getenv("HPCC_PATH") |
| 51 | + if not hpcc_path: |
| 52 | + return |
| 53 | + |
| 54 | + prefixes = _candidate_prefixes(hpcc_path) |
| 55 | + libs = [ |
| 56 | + "libhcruntime.so", |
| 57 | + "libhcToolsExt.so", |
| 58 | + "libruntime_cu.so", |
| 59 | + "libhccompiler.so", |
| 60 | + ] |
| 61 | + |
| 62 | + for lib in libs: |
| 63 | + _try_load(prefixes, lib) |
| 64 | + |
| 65 | + |
| 66 | +def _should_preload_device(device_type: str) -> bool: |
| 67 | + """ |
| 68 | + Check if preload is needed for a specific device type. |
| 69 | + """ |
| 70 | + device_env_map = { |
| 71 | + "METAX": ["HPCC_PATH", "INFINICORE_PRELOAD_HPCC"], # HPCC/METAX |
| 72 | + # Add other device types here as needed: |
| 73 | + # "ASCEND": ["ASCEND_PATH"], |
| 74 | + # "CAMBRICON": ["NEUWARE_HOME"], |
| 75 | + } |
| 76 | + |
| 77 | + env_vars = device_env_map.get(device_type, []) |
| 78 | + for env_var in env_vars: |
| 79 | + if os.getenv(env_var): |
| 80 | + return True |
| 81 | + return False |
| 82 | + |
| 83 | + |
| 84 | +def preload_device(device_type: str) -> None: |
| 85 | + """ |
| 86 | + Preload runtime libraries for a specific device type if needed. |
| 87 | +
|
| 88 | + Args: |
| 89 | + device_type: Device type name (e.g., "METAX", "ASCEND", etc.) |
| 90 | + """ |
| 91 | + if device_type == "METAX": |
| 92 | + preload_hpcc() |
| 93 | + # Add other device preload functions here as needed: |
| 94 | + # elif device_type == "ASCEND": |
| 95 | + # preload_ascend() |
| 96 | + # etc. |
| 97 | + |
| 98 | + |
| 99 | +def preload() -> None: |
| 100 | + """ |
| 101 | + Universal preload function that loops through device types and preloads when required. |
| 102 | +
|
| 103 | + This function detects available device types and preloads their runtime libraries |
| 104 | + if the environment indicates they are needed. |
| 105 | + """ |
| 106 | + # Device types that may require preload |
| 107 | + device_types = [ |
| 108 | + "METAX", # HPCC/METAX |
| 109 | + # Add other device types here as they are implemented: |
| 110 | + # "ASCEND", |
| 111 | + # "CAMBRICON", |
| 112 | + # etc. |
| 113 | + ] |
| 114 | + |
| 115 | + for device_type in device_types: |
| 116 | + if _should_preload_device(device_type): |
| 117 | + try: |
| 118 | + preload_device(device_type) |
| 119 | + except Exception: |
| 120 | + # Swallow all errors - preload is best-effort |
| 121 | + pass |
0 commit comments