From 57662d818c8d822e6a24d12abf99bb4711f056a7 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi <abarghi@nvidia.com>
Date: Tue, 26 Sep 2023 07:11:35 -0700
Subject: [PATCH] testing

---
 dependencies.yaml                             |  1 +
 .../tests/data_store/test_gnn_feat_storage.py | 36 -----------------
 .../test_gnn_feat_storage_wholegraph.py       | 39 +++++++++++++++++++
 3 files changed, 40 insertions(+), 36 deletions(-)
 create mode 100644 python/cugraph/cugraph/tests/data_store/test_gnn_feat_storage_wholegraph.py

diff --git a/dependencies.yaml b/dependencies.yaml
index 04ec1b6e957..8f9dc1c7fd7 100644
--- a/dependencies.yaml
+++ b/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]
diff --git a/python/cugraph/cugraph/tests/data_store/test_gnn_feat_storage.py b/python/cugraph/cugraph/tests/data_store/test_gnn_feat_storage.py
index 38b701ae8ae..5fa522e733d 100644
--- a/python/cugraph/cugraph/tests/data_store/test_gnn_feat_storage.py
+++ b/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)
\ No newline at end of file
diff --git a/python/cugraph/cugraph/tests/data_store/test_gnn_feat_storage_wholegraph.py b/python/cugraph/cugraph/tests/data_store/test_gnn_feat_storage_wholegraph.py
new file mode 100644
index 00000000000..58c1fc23181
--- /dev/null
+++ b/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)
\ No newline at end of file