Skip to content

Commit ed34107

Browse files
committed
new: support use_gpu algorithm parameter
1 parent 81aba1a commit ed34107

File tree

5 files changed

+23
-12
lines changed

5 files changed

+23
-12
lines changed

README.md

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -166,13 +166,16 @@ import nx_arangodb as nxadb
166166

167167
G = nxadb.Graph(name="MyGraph")
168168

169+
# Option 1: Use Global Config
169170
nx.config.backends.arangodb.use_gpu = False
170-
171171
nx.pagerank(G)
172172
nx.betweenness_centrality(G)
173173
# ...
174-
175174
nx.config.backends.arangodb.use_gpu = True
175+
176+
# Option 2: Use Local Config
177+
nx.pagerank(G, use_gpu=False)
178+
nx.betweenness_centrality(G, use_gpu=False)
176179
```
177180

178181
<p align="center">

doc/algorithms/index.rst

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,14 +43,17 @@ You can also force-run algorithms on CPU even if ``nx-cugraph`` is installed:
4343
4444
G = nxadb.Graph(name="MyGraph")
4545
46+
# Option 1: Use Global Config
4647
nx.config.backends.arangodb.use_gpu = False
47-
4848
nx.pagerank(G)
4949
nx.betweenness_centrality(G)
5050
# ...
51-
5251
nx.config.backends.arangodb.use_gpu = True
5352
53+
# Option 2: Use Local Config
54+
nx.pagerank(G, use_gpu=False)
55+
nx.betweenness_centrality(G, use_gpu=False)
56+
5457
5558
.. image:: ../_static/dispatch.png
5659
:align: center

doc/nx_arangodb.ipynb

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -236,9 +236,7 @@
236236
"outputs": [],
237237
"source": [
238238
"# 5. Run an algorithm (CPU)\n",
239-
"nx.config.backends.arangodb.use_gpu = False # Optional\n",
240-
"\n",
241-
"res = nx.pagerank(G)"
239+
"res = nx.pagerank(G, use_gpu=False)"
242240
]
243241
},
244242
{
@@ -357,8 +355,6 @@
357355
"source": [
358356
"# 4. Run an algorithm (GPU)\n",
359357
"# See *Package Installation* to install nx-cugraph ^\n",
360-
"nx.config.backends.arangodb.use_gpu = True\n",
361-
"\n",
362358
"res = nx.pagerank(G)"
363359
]
364360
},

nx_arangodb/interface.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,9 @@ def _auto_func(func_name: str, /, *args: Any, **kwargs: Any) -> Any:
6363
dfunc = _registered_algorithms[func_name]
6464

6565
backend_priority: list[str] = []
66-
if nxadb.convert.GPU_AVAILABLE and nx.config.backends.arangodb.use_gpu:
66+
67+
use_gpu = bool(kwargs.pop("use_gpu", nx.config.backends.arangodb.use_gpu))
68+
if nxadb.convert.GPU_AVAILABLE and use_gpu:
6769
backend_priority.append("cugraph")
6870

6971
for backend in backend_priority:

tests/test.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -447,7 +447,12 @@ def test_gpu_pagerank(graph_cls: type[nxadb.Graph]) -> None:
447447
assert gpu_cached_time < gpu_no_cache_time
448448
assert_pagerank(res_gpu_cached, res_gpu_no_cache, 10)
449449

450-
# 4. CPU
450+
# 4. CPU (with use_gpu=False)
451+
start_cpu_force_no_gpu = time.time()
452+
res_cpu_force_no_gpu = nx.pagerank(graph, use_gpu=False)
453+
cpu_force_no_gpu_time = time.time() - start_cpu_force_no_gpu
454+
455+
# 5. CPU
451456
assert graph.nxcg_graph is not None
452457
graph.clear_nxcg_cache()
453458
assert graph.nxcg_graph is None
@@ -456,12 +461,14 @@ def test_gpu_pagerank(graph_cls: type[nxadb.Graph]) -> None:
456461
start_cpu = time.time()
457462
res_cpu = nx.pagerank(graph)
458463
cpu_time = time.time() - start_cpu
464+
assert_pagerank(res_cpu, res_cpu_force_no_gpu, 10)
459465

460466
assert graph.nxcg_graph is None
461-
462467
m = "GPU execution should be faster than CPU execution"
463468
assert gpu_time < cpu_time, m
469+
assert gpu_time < cpu_force_no_gpu_time, m
464470
assert gpu_no_cache_time < cpu_time, m
471+
assert gpu_no_cache_time < cpu_force_no_gpu_time, m
465472
assert_pagerank(res_gpu_no_cache, res_cpu, 10)
466473

467474

0 commit comments

Comments
 (0)