Merge pull request #547 from pentschev/fix-debug-tests

Fix debug tests and remove duplicate functions
rapidsai · Jun 16, 2020 · bac9f1c · bac9f1c
2 parents 391403a + 47e9977
commit bac9f1c
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 15 deletions.
diff --git a/debug-tests/debug_utils.py b/debug-tests/debug_utils.py
@@ -6,11 +6,9 @@
 import cupy
 import rmm
 
-ITERATIONS = 100
-
+from .utils import get_num_gpus
 
-def cuda_array(size):
-    return rmm.DeviceBuffer(size=size)
+ITERATIONS = 100
 
 
 def set_rmm():
@@ -41,15 +39,6 @@ def parse_args(args=None):
     return parser.parse_args()
 
 
-def get_num_gpus():
-    import pynvml
-
-    pynvml.nvmlInit()
-    ngpus = pynvml.nvmlDeviceGetCount()
-    pynvml.nvmlShutdown()
-    return ngpus
-
-
 def get_cuda_devices():
     if "CUDA_VISIBLE_DEVICES" in os.environ:
         return os.environ["CUDA_VISIBLE_DEVICES"].split(",")

diff --git a/debug-tests/recv.py b/debug-tests/recv.py
@@ -5,7 +5,8 @@
 import pynvml
 import pytest
 import ucp
-from utils import ITERATIONS, recv, send, set_rmm
+from debug_utils import ITERATIONS, set_rmm
+from utils import recv, send
 
 pynvml.nvmlInit()
 

diff --git a/debug-tests/send.py b/debug-tests/send.py
@@ -7,7 +7,8 @@
 import cloudpickle
 import pytest
 import ucp
-from utils import ITERATIONS, recv, send, set_rmm
+from debug_utils import ITERATIONS, set_rmm
+from utils import recv, send
 
 cmd = "nvidia-smi nvlink --setcontrol 0bz"  # Get output in bytes
 # subprocess.check_call(cmd, shell=True)