Skip to content

Commit 7cf31b6

Browse files
committed
start integrating pynvml instead of nvidia-smi
1 parent 81fe552 commit 7cf31b6

File tree

1 file changed

+19
-15
lines changed

1 file changed

+19
-15
lines changed

nvdocker/nvdocker.py

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,24 @@
33
from subprocess import check_output
44
import re
55
import docker
6+
from pynvml import *
67

78
class NVDockerClient:
89

10+
nvml_initialized = False
11+
912
def __init__(self):
1013
self.docker_client = docker.from_env(version="auto")
14+
NVDockerClient.__check_nvml_init()
15+
16+
"""
17+
Private method to check if nvml is loaded (and load the library if it isn't loaded)
18+
"""
19+
def __check_nvml_init():
20+
if not NVDockerClient.nvml_initialized:
21+
nvmlInit()
22+
print("NVIDIA Driver Version:", nvmlSystemGetDriverVersion())
23+
NVDockerClient.nvml_initialized = True
1124

1225
#TODO: Testing on MultiGPU
1326
def create_container(self, image, **kwargs):
@@ -152,22 +165,13 @@ def exec_run(self, cid, cmd):
152165

153166
@staticmethod
154167
def gpu_info():
155-
#output = check_output(["nvidia-smi", "-L"]).decode("utf-8")
156-
keys = ['memory_free', 'memory_used', 'memory_total']
157-
query_gpu = check_output(["nvidia-smi", "--query-gpu=memory.free,memory.used,memory.total","--format=csv,noheader"]).decode("utf-8")
158-
#regex = re.compile(r"GPU (?P<id>\d+):")
159-
query_gpu = query_gpu.strip()
168+
NVDockerClient.__check_nvml_init()
160169
gpus = {}
161-
id = 0
162-
for gpu in query_gpu.split("\n"):
163-
gpu_info = {}
164-
key_id = 0
165-
for info in gpu.split(","):
166-
info = info.strip()
167-
gpu_info[keys[key_id]] = info.split(" ")[0];
168-
key_id += 1
169-
gpus[id] = gpu_info;
170-
id += 1
170+
num_gpus = nvmlDeviceGetCount()
171+
for i in range(num_gpus):
172+
gpu_handle = nvmlDeviceGetHandleByIndex(i)
173+
gpu_name = nvmlDeviceGetName(gpu_handle)
174+
gpus[i] = {"gpu_handle": gpu_handle, "gpu_name": gpu_name}
171175
return gpus
172176

173177
@staticmethod

0 commit comments

Comments
 (0)