From a485a0751e9f5adfe92ce2eff06d13e0b599b0e7 Mon Sep 17 00:00:00 2001
From: Alexandria Barghi <abarghi@nvidia.com>
Date: Tue, 26 Sep 2023 07:26:32 -0700
Subject: [PATCH] commit

---
 .../all_cuda-118_arch-x86_64.yaml             |  1 +
 .../all_cuda-120_arch-x86_64.yaml             |  1 +
 python/cugraph-service/server/pyproject.toml  |  1 +
 .../tests/data_store/test_gnn_feat_storage.py |  2 --
 .../test_gnn_feat_storage_wholegraph.py       | 27 +++++++++++++++----
 python/cugraph/pyproject.toml                 |  1 +
 6 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/conda/environments/all_cuda-118_arch-x86_64.yaml b/conda/environments/all_cuda-118_arch-x86_64.yaml
index c66890f8ae5..ea3a931cc43 100644
--- a/conda/environments/all_cuda-118_arch-x86_64.yaml
+++ b/conda/environments/all_cuda-118_arch-x86_64.yaml
@@ -53,6 +53,7 @@ dependencies:
 - pydata-sphinx-theme
 - pylibcugraphops==23.10.*
 - pylibraft==23.10.*
+- pylibwholegraph=23.10.*
 - pytest
 - pytest-benchmark
 - pytest-cov
diff --git a/conda/environments/all_cuda-120_arch-x86_64.yaml b/conda/environments/all_cuda-120_arch-x86_64.yaml
index 3afb1415572..8dc038432bb 100644
--- a/conda/environments/all_cuda-120_arch-x86_64.yaml
+++ b/conda/environments/all_cuda-120_arch-x86_64.yaml
@@ -52,6 +52,7 @@ dependencies:
 - pydata-sphinx-theme
 - pylibcugraphops==23.10.*
 - pylibraft==23.10.*
+- pylibwholegraph=23.10.*
 - pytest
 - pytest-benchmark
 - pytest-cov
diff --git a/python/cugraph-service/server/pyproject.toml b/python/cugraph-service/server/pyproject.toml
index f25ea6c46e5..538acf8af51 100644
--- a/python/cugraph-service/server/pyproject.toml
+++ b/python/cugraph-service/server/pyproject.toml
@@ -46,6 +46,7 @@ test = [
     "networkx>=2.5.1",
     "numpy>=1.21",
     "pandas",
+    "pylibwholegraph=23.10.*",
     "pytest",
     "pytest-benchmark",
     "pytest-cov",
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 5fa522e733d..2b0ec4b11d0 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
@@ -10,7 +10,6 @@
 # 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 FeatureStore class
 
 import pytest
 import numpy as np
@@ -107,4 +106,3 @@ def test_feature_storage_pytorch_backend():
     expected = ar3[indices_to_fetch]
     assert isinstance(output_fs, torch.Tensor)
     np.testing.assert_array_equal(output_fs.numpy(), expected)
-
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
index 58c1fc23181..7694b259fda 100644
--- 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
@@ -1,7 +1,19 @@
+# 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 pytest
 import numpy as np
 
-import cudf
 from cugraph.gnn import FeatureStore
 
 import pylibwholegraph.binding.wholememory_binding as wmb
@@ -10,9 +22,13 @@
 
 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,
+        world_rank,
+        world_size,
+        world_rank,
+        world_size,
     )
     wm_comm = wm_comm.wmb_comm
 
@@ -27,13 +43,14 @@ def func(world_rank: int, world_size: int):
     expected = ar3[indices_to_fetch]
     np.testing.assert_array_equal(output_fs.cpu().numpy(), expected)
 
-
     wmb.finalize()
 
+
+@pytest.mark.sg
 def test_feature_storage_wholegraph_backend():
     gpu_count = wmb.fork_get_gpu_count()
-    print('gpu count:', 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
+    multiprocess_run(1, func)
diff --git a/python/cugraph/pyproject.toml b/python/cugraph/pyproject.toml
index cadf6879e23..006584e7f69 100644
--- a/python/cugraph/pyproject.toml
+++ b/python/cugraph/pyproject.toml
@@ -54,6 +54,7 @@ test = [
     "networkx>=2.5.1",
     "numpy>=1.21",
     "pandas",
+    "pylibwholegraph=23.10.*",
     "pytest",
     "pytest-benchmark",
     "pytest-cov",