Enhancement/select device

CHANGELOG.md

@@ -10,6 +10,8 @@
 - [#496](https://github.com/helmholtz-analytics/heat/pull/496) New feature: flipud()
 - [#498](https://github.com/helmholtz-analytics/heat/pull/498) Feature: flip()
 - [#499](https://github.com/helmholtz-analytics/heat/pull/499) Bugfix: MPI datatype mapping: `torch.int16` now maps to `MPI.SHORT` instead of `MPI.SHORT_INT`
+- [#488](https://github.com/helmholtz-analytics/heat/pull/488) Enhancement: Rework of the test device selection. Environment variable is now called HEAT_USE_DEVICE and has more options.
+
 - [#501](https://github.com/helmholtz-analytics/heat/pull/501) New Feature: flatten
 - [#506](https://github.com/helmholtz-analytics/heat/pull/506) Bugfix: setup.py has correct version parsing
 - [#507](https://github.com/helmholtz-analytics/heat/pull/507) Bugfix: sanitize_axis changes axis of 0-dim scalars to None

heat/cluster/tests/test_kmeans.py

@@ -3,20 +3,14 @@
 
 import heat as ht
 
-if os.environ.get("DEVICE") == "gpu" and ht.torch.cuda.is_available():
-    ht.use_device("gpu")
-    ht.torch.cuda.set_device(ht.torch.device(ht.get_device().torch_device))
-else:
-    ht.use_device("cpu")
-device = ht.get_device().torch_device
-ht_device = None
-if os.environ.get("DEVICE") == "lgpu" and ht.torch.cuda.is_available():
-    device = ht.gpu.torch_device
-    ht_device = ht.gpu
-    ht.torch.cuda.set_device(device)
+from heat.core.tests.test_suites.basic_test import BasicTest
 
 
-class TestKMeans(unittest.TestCase):
+class TestKMeans(BasicTest):
+    @classmethod
+    def setUpClass(cls):
+        super(TestKMeans, cls).setUpClass()
+
     def test_clusterer(self):
         kmeans = ht.cluster.KMeans()
         self.assertTrue(ht.is_estimator(kmeans))
@@ -38,7 +32,9 @@ def test_get_and_set_params(self):
     def test_fit_iris_unsplit(self):
         for split in [None, 0]:
             # get some test data
-            iris = ht.load("heat/datasets/data/iris.csv", sep=";", split=split)
+            iris = ht.load(
+                "heat/datasets/data/iris.csv", sep=";", split=split, device=self.ht_device
+            )
 
             # fit the clusters
             k = 3
@@ -58,7 +54,7 @@ def test_fit_iris_unsplit(self):
 
     def test_exceptions(self):
         # get some test data
-        iris_split = ht.load("heat/datasets/data/iris.csv", sep=";", split=1)
+        iris_split = ht.load("heat/datasets/data/iris.csv", sep=";", split=1, device=self.ht_device)
 
         # build a clusterer
         k = 3

heat/cluster/tests/test_spectral.py

@@ -3,20 +3,14 @@
 
 import heat as ht
 
-if os.environ.get("DEVICE") == "gpu" and ht.torch.cuda.is_available():
-    ht.use_device("gpu")
-    ht.torch.cuda.set_device(ht.torch.device(ht.get_device().torch_device))
-else:
-    ht.use_device("cpu")
-device = ht.get_device().torch_device
-ht_device = None
-if os.environ.get("DEVICE") == "lgpu" and ht.torch.cuda.is_available():
-    device = ht.gpu.torch_device
-    ht_device = ht.gpu
-    ht.torch.cuda.set_device(device)
+from heat.core.tests.test_suites.basic_test import BasicTest
 
 
-class TestSpectral(unittest.TestCase):
+class TestSpectral(BasicTest):
+    @classmethod
+    def setUpClass(cls):
+        super(TestSpectral, cls).setUpClass()
+
     def test_clusterer(self):
         spectral = ht.cluster.Spectral()
         self.assertTrue(ht.is_estimator(spectral))
@@ -46,7 +40,7 @@ def test_get_and_set_params(self):
 
     def test_fit_iris(self):
         # get some test data
-        iris = ht.load("heat/datasets/data/iris.csv", sep=";", split=0)
+        iris = ht.load("heat/datasets/data/iris.csv", sep=";", split=0, device=self.ht_device)
         m = 10
 
         # fit the clusters
@@ -84,7 +78,7 @@ def test_fit_iris(self):
         with self.assertRaises(NotImplementedError):
             spectral = ht.cluster.Spectral(metric="ahalanobis", n_lanczos=m)
 
-        iris_split = ht.load("heat/datasets/data/iris.csv", sep=";", split=1)
+        iris_split = ht.load("heat/datasets/data/iris.csv", sep=";", split=1, device=self.ht_device)
         spectral = ht.cluster.Spectral(n_lanczos=20)
         with self.assertRaises(NotImplementedError):
             spectral.fit(iris_split)