Check environment supports target device in Dataset constructor (#243)

* Use `HAS_GPU` as part of Dataset device choice

* Update import of HAS_GPU in dataset

* Raise RuntimeError in Dataset if we can't run on GPU when cpu=False

* Update error message in case of failed Dataset initialization

* Set default value of cpu to None instead of False in fixture

* Add test of expected Dataset cpu argument behaviour

---------

Co-authored-by: Karl Higley <karlb@nvidia.com>

merlin/io/dataset.py

@@ -244,9 +244,22 @@ def __init__(
         # Cache for "real" (sampled) metadata
         self._real_meta = {}
 
-        # Check if we are keeping data in cpu memory
+        # Check if we are keeping data in host or gpu device memory
         self.cpu = cpu
-        if not self.cpu:
+        if self.cpu is False:
+            if not HAS_GPU:
+                raise RuntimeError(
+                    "Cannot initialize Dataset on GPU. "
+                    "No devices detected (with pynvml). "
+                    "Check that pynvml can be initialized. "
+                )
+            if cudf is None:
+                raise RuntimeError(
+                    "Cannot initialize Dataset on GPU. "
+                    "cudf package not found. "
+                    "Check that cudf is installed in this environment and can be imported.  "
+                )
+        if self.cpu is None:
             self.cpu = cudf is None or not HAS_GPU
 
         # Keep track of base dataset (optional)

tests/conftest.py

@@ -221,7 +221,7 @@ def dataset(request, paths, engine):
     try:
         cpu = request.getfixturevalue("cpu")
     except Exception:  # pylint: disable=broad-except
-        cpu = False
+        cpu = None
 
     kwargs = {}
     if engine == "csv-no-header":

tests/unit/io/test_dataset.py

@@ -0,0 +1,51 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import pandas as pd
+import pytest
+
+from merlin.core.compat import HAS_GPU, cudf
+from merlin.core.dispatch import make_df
+from merlin.io import Dataset
+
+
+class TestDatasetCpu:
+    def test_true(self):
+        dataset = Dataset(make_df({"a": [1, 2, 3]}), cpu=True)
+        assert dataset.cpu is True
+        assert isinstance(dataset.compute(), pd.DataFrame)
+
+    @pytest.mark.skipif(not (cudf and HAS_GPU), reason="requires cuDF and GPU")
+    def test_default_cudf(self):
+        dataset = Dataset(make_df({"a": [1, 2, 3]}))
+        assert dataset.cpu is False
+        assert isinstance(dataset.compute(), cudf.DataFrame)
+
+    @pytest.mark.skipif(cudf and HAS_GPU, reason="requires environment without cuDF and GPU")
+    def test_default_pandas(self):
+        dataset = Dataset(make_df({"a": [1, 2, 3]}))
+        assert dataset.cpu is True
+        assert isinstance(dataset.compute(), pd.DataFrame)
+
+    @pytest.mark.skipif(not (cudf and HAS_GPU), reason="requires cuDF and GPU")
+    def test_false_with_cudf_and_gpu(self):
+        dataset = Dataset(make_df({"a": [1, 2, 3]}), cpu=False)
+        assert dataset.cpu is False
+        assert isinstance(dataset.compute(), cudf.DataFrame)
+
+    @pytest.mark.skipif(cudf or HAS_GPU, reason="requires environment without cuDF or GPU")
+    def test_false_missing_cudf_or_gpu(self):
+        with pytest.raises(RuntimeError):
+            Dataset(make_df({"a": [1, 2, 3]}), cpu=False)