Adding more tests to wrapper executor

Signed-off-by: Julio Faracco <[email protected]>

Author: jcfaracco

dasf/pipeline/executors/wrapper.py

@@ -1,8 +1,11 @@
 #!/usr/bin/env python3
 
+import numpy as np
+
 try:
     import cupy as cp
     import rmm
+    from rmm.allocators.cupy import rmm_cupy_allocator
 except ImportError:  # pragma: no cover
     pass
 
@@ -30,7 +33,7 @@ def __init__(self,
 
         if gpu_allocator == "rmm" and self.dtype == TaskExecutorType.single_gpu:
             rmm.reinitialize(managed_memory=True)
-            cp.cuda.set_allocator(rmm.rmm_cupy_allocator)
+            cp.cuda.set_allocator(rmm_cupy_allocator)
 
     @property
     def ngpus(self) -> int:
@@ -47,11 +50,10 @@ def post_run(self, pipeline):
         pass
 
     def get_backend(self):
-        if self.backend == "numpy" and \
-           self.dtype == TaskExecutorType.single_gpu:
-            return eval("cupy")
+        if self.dtype == TaskExecutorType.single_gpu:
+            return eval("cp")
 
-        return eval(self.backend)
+        return eval("np")
 
     def execute(self, fn, *args, **kwargs):
         if get_backend_supported(fn):

tests/pipeline/executors/test_wrapper.py

@@ -3,8 +3,14 @@
 import os
 import unittest
 
+import numpy as np
 from mock import Mock, patch
 
+try:
+    import cupy as cp
+except ImportError:
+    pass
+
 from dasf.pipeline.executors import LocalExecutor
 from dasf.pipeline.types import TaskExecutorType
 from dasf.utils.funcs import is_gpu_supported
@@ -15,4 +21,48 @@ class TestLocalExecutor(unittest.TestCase):
     def test_local_executor_no_gpu(self):
         local = LocalExecutor()
 
-        self.assertTrue(local.dtype, TaskExecutorType.single_cpu)
+        self.assertEqual(local.dtype, TaskExecutorType.single_cpu)
+        self.assertEqual(local.get_backend(), np)
+
+    @patch('dasf.pipeline.executors.wrapper.is_gpu_supported', Mock(return_value=False))
+    def test_local_executor_no_gpu_but_use_gpu(self):
+        local = LocalExecutor(use_gpu=True)
+
+        self.assertEqual(local.dtype, TaskExecutorType.single_cpu)
+        self.assertEqual(local.get_backend(), np)
+
+    @unittest.skipIf(not is_gpu_supported(),
+                     "not supported CUDA in this platform")
+    def test_local_executor_use_gpu(self):
+        local = LocalExecutor(use_gpu=True)
+
+        self.assertEqual(local.dtype, TaskExecutorType.single_gpu)
+        self.assertEqual(local.get_backend(), cp)
+
+    @unittest.skipIf(not is_gpu_supported(),
+                     "not supported CUDA in this platform")
+    def test_local_executor_use_gpu_backend_cupy(self):
+        local = LocalExecutor(use_gpu=True, backend="cupy")
+
+        self.assertEqual(local.dtype, TaskExecutorType.single_gpu)
+        self.assertEqual(local.get_backend(), cp)
+
+    @unittest.skipIf(not is_gpu_supported(),
+                     "not supported CUDA in this platform")
+    def test_local_executor_use_gpu_backend_cupy(self):
+        local = LocalExecutor(backend="cupy")
+
+        self.assertEqual(local.dtype, TaskExecutorType.single_gpu)
+        self.assertEqual(local.backend, "cupy")
+        self.assertEqual(local.get_backend(), cp)
+
+    @unittest.skipIf(not is_gpu_supported(),
+                     "not supported CUDA in this platform")
+    @patch('dasf.pipeline.executors.wrapper.rmm.reinitialize')
+    def test_local_executor_with_rmm(self, rmm):
+        local = LocalExecutor(gpu_allocator="rmm")
+
+        self.assertEqual(local.dtype, TaskExecutorType.single_gpu)
+        self.assertEqual(local.get_backend(), cp)
+
+        rmm.assert_called_once_with(managed_memory=True)