robustify get_gpu_info

just4jc · Feb 5, 2018 · 12fa4f8 · 12fa4f8
1 parent 14bcc85
commit 12fa4f8
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/interface_py/h2o4gpu/util/gpu.py b/src/interface_py/h2o4gpu/util/gpu.py
@@ -84,6 +84,11 @@ def get_gpu_info_subprocess(return_usage=False):
         import os
         cudavis = os.getenv("CUDA_VISIBLE_DEVICES")
         if cudavis is not None:
+            # prune away white-space, non-numerics, except commas for simple checking
+            cudavis = "".join(cudavis.split())
+            import re
+            cudavis = re.sub("[^0-9,]", "", cudavis)
+
             lencudavis = len(cudavis)
             if lencudavis == 0:
                 total_gpus = 0