Skip to content

Commit 380b817

Browse files
thxCodethxCode
committed
refactor: tools
Signed-off-by: thxCode <thxcode0824@gmail.com>
1 parent b710664 commit 380b817

9 files changed

Lines changed: 253 additions & 183 deletions

File tree

‎gpustack_runtime/detector/pyacl/__init__.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -384,13 +384,20 @@ def wrapper(*args, **kwargs):
384384

385385

386386
def _LoadAclLibrary():
387+
"""
388+
Load the library if it isn't loaded already.
389+
"""
387390
global aclLib
391+
388392
if aclLib is None:
393+
# lock to ensure only one caller loads the library
389394
libLoadLock.acquire()
395+
390396
try:
397+
# ensure the library still isn't loaded
391398
if aclLib is None:
392399
if sys.platform.startswith("win"):
393-
# ACL not supported on Windows
400+
# Do not support Windows yet.
394401
raise ACLError(ACL_ERROR_LIBRARY_NOT_FOUND)
395402
# Linux path
396403
locs = [
@@ -419,6 +426,7 @@ def _LoadAclLibrary():
419426
if aclLib is None:
420427
raise ACLError(ACL_ERROR_LIBRARY_NOT_FOUND)
421428
finally:
429+
# lock is always freed
422430
libLoadLock.release()
423431

424432

‎gpustack_runtime/detector/pyamdgpu/__init__.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,10 +212,17 @@ class c_amdgpu_gpu_info(_PrintableStructure):
212212

213213

214214
def _LoadAMDGPULibrary():
215+
"""
216+
Load the library if it isn't loaded already.
217+
"""
215218
global amdgpuLib
219+
216220
if amdgpuLib is None:
221+
# lock to ensure only one caller loads the library
217222
libLoadLock.acquire()
223+
218224
try:
225+
# ensure the library still isn't loaded
219226
if amdgpuLib is None:
220227
if sys.platform.startswith("win"):
221228
# Do not support Windows yet.
@@ -235,6 +242,7 @@ def _LoadAMDGPULibrary():
235242
if amdgpuLib is None:
236243
raise AMDGPUError(AMDGPU_ERROR_LIBRARY_NOT_FOUND)
237244
finally:
245+
# lock is always freed
238246
libLoadLock.release()
239247

240248

‎gpustack_runtime/detector/pycuda/__init__.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
from __future__ import annotations
22

3-
import contextlib
43
import string
54
import sys
65
import threading
76
from ctypes import *
7+
from functools import wraps
88
from typing import ClassVar
99

1010
## C Type mappings ##
@@ -390,10 +390,17 @@ def wrapper(*args, **kwargs):
390390

391391

392392
def _LoadCudaLibrary():
393+
"""
394+
Load the library if it isn't loaded already.
395+
"""
393396
global cudaLib
397+
394398
if cudaLib is None:
399+
# lock to ensure only one caller loads the library
395400
libLoadLock.acquire()
401+
396402
try:
403+
# ensure the library still isn't loaded
397404
if cudaLib is None:
398405
if sys.platform.startswith("win"):
399406
# Do not support Windows yet.
@@ -411,6 +418,7 @@ def _LoadCudaLibrary():
411418
if cudaLib is None:
412419
raise CUDAError(CUDA_ERROR_LIBRARY_NOT_FOUND)
413420
finally:
421+
# lock is always freed
414422
libLoadLock.release()
415423

416424

‎gpustack_runtime/detector/pydcmi/__init__.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -733,14 +733,20 @@ def wrapper(*args, **kwargs):
733733

734734

735735
def _LoadDcmiLibrary():
736+
"""
737+
Load the library if it isn't loaded already.
738+
"""
736739
global dcmiLib
740+
737741
if dcmiLib is None:
742+
# lock to ensure only one caller loads the library
738743
libLoadLock.acquire()
744+
739745
try:
746+
# ensure the library still isn't loaded
740747
if dcmiLib is None:
741748
if sys.platform.startswith("win"):
742-
# DCMI is typically used on Linux, but for completeness,
743-
# Windows support would require different path handling.
749+
# Do not support Windows yet.
744750
raise DCMIError(DCMI_ERROR_LIBRARY_NOT_FOUND)
745751
# Linux path
746752
locs = [
@@ -757,6 +763,7 @@ def _LoadDcmiLibrary():
757763
if dcmiLib is None:
758764
raise DCMIError(DCMI_ERROR_LIBRARY_NOT_FOUND)
759765
finally:
766+
# lock is always freed
760767
libLoadLock.release()
761768

762769

‎gpustack_runtime/detector/pyhgml/__init__.py‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,12 @@
11
##
22
# Python bindings for the HGML library
33
##
4-
5-
import os
64
import string
75
import sys
86
import threading
97
from ctypes import *
108
from functools import wraps
9+
from pathlib import Path
1110

1211
## C Type mappings ##
1312
## Enums
@@ -2403,21 +2402,25 @@ def _LoadHgmlLibrary():
24032402
if hgmlLib is None:
24042403
# lock to ensure only one caller loads the library
24052404
libLoadLock.acquire()
2406-
24072405
try:
24082406
# ensure the library still isn't loaded
24092407
if hgmlLib is None:
2410-
try:
2411-
if sys.platform.startswith("win"):
2412-
# HGML is typically used on Linux, but for completeness,
2413-
# Windows support would require different path handling.
2414-
_hgmlCheckReturn(HGML_ERROR_LIBRARY_NOT_FOUND)
2415-
else:
2416-
hgmlLib = CDLL("libhgml.so")
2417-
except OSError:
2418-
_hgmlCheckReturn(HGML_ERROR_LIBRARY_NOT_FOUND)
2408+
if sys.platform.startswith("win"):
2409+
# Do not support Windows yet.
2410+
raise HGMLError(HGML_ERROR_LIBRARY_NOT_FOUND)
2411+
# Linux path
2412+
locs = [
2413+
"libhgml.so",
2414+
str(Path(__file__).resolve().parent / "libhgml.so"),
2415+
]
2416+
for loc in locs:
2417+
try:
2418+
hgmlLib = CDLL(loc)
2419+
break
2420+
except OSError:
2421+
pass
24192422
if hgmlLib is None:
2420-
_hgmlCheckReturn(HGML_ERROR_LIBRARY_NOT_FOUND)
2423+
raise HGMLError(HGML_ERROR_LIBRARY_NOT_FOUND)
24212424
finally:
24222425
# lock is always freed
24232426
libLoadLock.release()

‎gpustack_runtime/detector/pyhsa/__init__.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -331,10 +331,16 @@ def wrapper(*args, **kwargs):
331331

332332

333333
def _LoadHsaLibrary():
334+
"""
335+
Load the library if it isn't loaded already.
336+
"""
334337
global hsaLib
338+
335339
if hsaLib is None:
340+
# lock to ensure only one caller loads the library
336341
libLoadLock.acquire()
337342
try:
343+
# ensure the library still isn't loaded
338344
if hsaLib is None:
339345
if sys.platform.startswith("win"):
340346
# Do not support Windows yet.
@@ -360,7 +366,10 @@ def _LoadHsaLibrary():
360366
break
361367
except OSError:
362368
pass
369+
if hsaLib is None:
370+
raise HSAError(HSA_STATUS_ERROR_LIBRARY_NOT_FOUND)
363371
finally:
372+
# lock is always released
364373
libLoadLock.release()
365374

366375

‎gpustack_runtime/detector/pyixml/__init__.py‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
##
22
# Python bindings for the IXML library
33
##
4-
54
#####
65
# Copyright (c) 2011-2023, NVIDIA Corporation. All rights reserved.
76
#
@@ -36,9 +35,9 @@
3635
from ctypes import *
3736
from functools import wraps
3837
import sys
39-
import os
4038
import threading
4139
import string
40+
from pathlib import Path
4241

4342
## C Type mappings ##
4443
## Enums
@@ -2146,21 +2145,25 @@ def _LoadNvmlLibrary():
21462145
if nvmlLib is None:
21472146
# lock to ensure only one caller loads the library
21482147
libLoadLock.acquire()
2149-
21502148
try:
21512149
# ensure the library still isn't loaded
21522150
if nvmlLib is None:
2153-
try:
2154-
if sys.platform.startswith("win"):
2155-
# IXML is typically used on Linux, but for completeness,
2156-
# Windows support would require different path handling.
2157-
_nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND)
2158-
else:
2159-
nvmlLib = CDLL("libixml.so")
2160-
except OSError:
2161-
_nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND)
2151+
if sys.platform.startswith("win"):
2152+
# Do not support Windows yet.
2153+
raise NVMLError(NVML_ERROR_LIBRARY_NOT_FOUND)
2154+
# Linux path
2155+
locs = [
2156+
"libixml.so",
2157+
str(Path(__file__).resolve().parent / "libixml.so"),
2158+
]
2159+
for loc in locs:
2160+
try:
2161+
nvmlLib = CDLL(loc)
2162+
break
2163+
except OSError:
2164+
pass
21622165
if nvmlLib is None:
2163-
_nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND)
2166+
raise NVMLError(NVML_ERROR_LIBRARY_NOT_FOUND)
21642167
finally:
21652168
# lock is always freed
21662169
libLoadLock.release()

‎gpustack_runtime/detector/pyrocmcore/__init__.py‎

Lines changed: 42 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,42 +1,58 @@
11
from __future__ import annotations
22

33
import os
4+
import sys
45
import threading
56
from ctypes import *
67
from pathlib import Path
78

9+
# Example ROCM_SMI_LIB_PATH
10+
# - /opt/dtk-24.04.3/rocm_smi/lib
11+
# - /opt/rocm/rocm_smi/lib
12+
rocmcore_lib_path = os.getenv("ROCM_CORE_LIB_PATH")
13+
if not rocmcore_lib_path:
14+
# Example ROCM_PATH/ROCM_HOME
15+
# - /opt/dtk-24.04.3
16+
# - /opt/rocm
17+
rocm_path = Path(os.getenv("ROCM_HOME", os.getenv("ROCM_PATH") or "/opt/rocm"))
18+
rocmcore_lib_path = str(rocm_path / "lib")
19+
else:
20+
rocm_path = Path(
21+
os.getenv(
22+
"ROCM_HOME",
23+
os.getenv("ROCM_PATH") or str(Path(rocmcore_lib_path).parent.parent),
24+
)
25+
)
26+
27+
rocmcore_lib_loc = Path(rocmcore_lib_path) / "librocm-core.so"
28+
829
## Lib loading ##
930
rocmcoreLib = None
1031
libLoadLock = threading.Lock()
1132

12-
if rocmcoreLib is None:
13-
# Example ROCM_SMI_LIB_PATH
14-
# - /opt/dtk-24.04.3/rocm_smi/lib
15-
# - /opt/rocm/rocm_smi/lib
16-
rocmcore_lib_path = os.getenv("ROCM_CORE_LIB_PATH")
17-
if not rocmcore_lib_path:
18-
# Example ROCM_PATH/ROCM_HOME
19-
# - /opt/dtk-24.04.3
20-
# - /opt/rocm
21-
rocm_path = Path(os.getenv("ROCM_HOME", os.getenv("ROCM_PATH") or "/opt/rocm"))
22-
rocmcore_lib_path = str(rocm_path / "lib")
23-
else:
24-
rocm_path = Path(
25-
os.getenv(
26-
"ROCM_HOME",
27-
os.getenv("ROCM_PATH") or str(Path(rocmcore_lib_path).parent.parent),
28-
)
29-
)
3033

31-
rocmcore_lib_loc = Path(rocmcore_lib_path) / "librocm-core.so"
32-
if rocmcore_lib_loc.exists():
34+
def _LoadRocmCoreLibrary():
35+
"""
36+
Load the library if it isn't loaded already.
37+
"""
38+
global rocmcoreLib
39+
40+
if rocmcoreLib is None:
41+
# lock to ensure only one caller loads the library
3342
libLoadLock.acquire()
3443
try:
35-
if not rocmcoreLib:
36-
rocmcoreLib = CDLL(rocmcore_lib_loc)
37-
except OSError:
38-
pass
44+
# ensure the library still isn't loaded
45+
if (
46+
rocmcoreLib is None
47+
and not sys.platform.startswith("win")
48+
and rocmcore_lib_loc.is_file()
49+
):
50+
try:
51+
rocmcoreLib = CDLL(str(rocmcore_lib_loc))
52+
except OSError:
53+
pass
3954
finally:
55+
# lock is always released
4056
libLoadLock.release()
4157

4258

@@ -59,6 +75,8 @@ def getROCmVersion() -> str | None:
5975
except OSError:
6076
continue
6177

78+
_LoadRocmCoreLibrary()
79+
6280
if rocmcoreLib:
6381
try:
6482
major = c_uint32()

0 commit comments

Comments
 (0)