diff toolbox/GPUselect.py @ 0:6bec4fef6b2e draft

"planemo upload for repository https://github.com/ohsu-comp-bio/unmicst commit 73e4cae15f2d7cdc86719e77470eb00af4b6ebb7-dirty"
author perssond
date Fri, 12 Mar 2021 00:17:29 +0000
parents
children
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/toolbox/GPUselect.py	Fri Mar 12 00:17:29 2021 +0000
@@ -0,0 +1,19 @@
+import subprocess, re
+import numpy as np
+
+def pick_gpu_lowest_memory():
+    output = subprocess.Popen("nvidia-smi", stdout=subprocess.PIPE, shell=True).communicate()[0]
+    output=output.decode("ascii")
+    gpu_output = output[output.find("Memory-Usage"):]
+        # lines of the form
+        # |    0      8734    C   python                                       11705MiB |
+    memory_regex = re.compile(r"[|]\s+?\D+?.+[ ](?P<gpu_memory>\d+)MiB /")
+    rows = gpu_output.split("\n")
+    result=[]
+    for row in gpu_output.split("\n"):
+        m = memory_regex.search(row)
+        if not m:
+            continue
+        gpu_memory = int(m.group("gpu_memory"))
+        result.append(gpu_memory)
+    return np.argsort(result)[0]
\ No newline at end of file