testing

dependencies.yaml

@@ -488,6 +488,7 @@ dependencies:
           - *numpy
           - python-louvain
           - scikit-learn>=0.23.1
+          - pylibwholegraph=23.10.*
   test_python_pylibcugraph:
     common:
       - output_types: [conda, pyproject]

python/cugraph/cugraph/tests/data_store/test_gnn_feat_storage.py

@@ -108,39 +108,3 @@ def test_feature_storage_pytorch_backend():
     assert isinstance(output_fs, torch.Tensor)
     np.testing.assert_array_equal(output_fs.numpy(), expected)
 
-
-@pytest.mark.sg
-def test_feature_storage_wholegraph_backend():
-    try:
-        import torch
-    except ModuleNotFoundError:
-        pytest.skip('pytorch not installed')
-
-    fs = FeatureStore(backend='wholegraph')
-
-    ar1 = np.random.randint(low=0, high=100, size=100_000)
-    ar2 = np.random.randint(low=0, high=100, size=100_000)
-    ar3 = np.random.randint(low=0, high=100, size=100_000).reshape(-1, 10)
-    fs = FeatureStore(backend="torch")
-    fs.add_data(ar1, "type1", "feat1")
-    fs.add_data(ar2, "type1", "feat2")
-    fs.add_data(ar3, "type2", "feat1")
-
-
-    indices_to_fetch = np.random.randint(low=0, high=len(ar1), size=1024)
-    output_fs = fs.get_data(indices_to_fetch, type_name="type1", feat_name="feat1")
-    expected = ar1[indices_to_fetch]
-    assert isinstance(output_fs, torch.Tensor)
-    np.testing.assert_array_equal(output_fs.numpy(), expected)
-
-    indices_to_fetch = np.random.randint(low=0, high=len(ar2), size=1024)
-    output_fs = fs.get_data(indices_to_fetch, type_name="type1", feat_name="feat2")
-    expected = ar2[indices_to_fetch]
-    assert isinstance(output_fs, torch.Tensor)
-    np.testing.assert_array_equal(output_fs.numpy(), expected)
-
-    indices_to_fetch = np.random.randint(low=0, high=len(ar3), size=1024)
-    output_fs = fs.get_data(indices_to_fetch, type_name="type2", feat_name="feat1")
-    expected = ar3[indices_to_fetch]
-    assert isinstance(output_fs, torch.Tensor)
-    np.testing.assert_array_equal(output_fs.numpy(), expected)

python/cugraph/cugraph/tests/data_store/test_gnn_feat_storage_wholegraph.py

@@ -0,0 +1,39 @@
+import pytest
+import numpy as np
+
+import cudf
+from cugraph.gnn import FeatureStore
+
+import pylibwholegraph.binding.wholememory_binding as wmb
+from pylibwholegraph.torch.initialize import init_torch_env_and_create_wm_comm
+from pylibwholegraph.utils.multiprocess import multiprocess_run
+
+import torch
+
+def func(world_rank: int, world_size: int):
+    wm_comm, _ = init_torch_env_and_create_wm_comm(
+        world_rank, world_size, world_rank, world_size,
+    )
+    wm_comm = wm_comm.wmb_comm
+
+    ar3 = np.random.randint(low=0, high=100, size=100_000).reshape(10_000, -1)
+    fs = FeatureStore(backend="wholegraph")
+    fs.add_data(ar3, "type2", "feat1")
+
+    indices_to_fetch = np.random.randint(low=0, high=len(ar3), size=1024)
+    output_fs = fs.get_data(indices_to_fetch, type_name="type2", feat_name="feat1")
+    assert isinstance(output_fs, torch.Tensor)
+    assert output_fs.is_cuda
+    expected = ar3[indices_to_fetch]
+    np.testing.assert_array_equal(output_fs.cpu().numpy(), expected)
+
+
+    wmb.finalize()
+
+def test_feature_storage_wholegraph_backend():
+    gpu_count = wmb.fork_get_gpu_count()
+    print('gpu count:', gpu_count)
+    assert gpu_count > 0
+
+    # FIXME make this work in an MG environment
+    multiprocess_run(1, func)