Skip to content

Commit 25eef43

Browse files
bdicebdice
committed
Merge branch-25.02 into branch-25.04
2 parents b578959 + 3546217 commit 25eef43

24 files changed

+136
-86
lines changed

‎.devcontainer/cuda11.8-pip/devcontainer.json

+1-1
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
"args": {
66
"CUDA": "11.8",
77
"PYTHON_PACKAGE_MANAGER": "pip",
8-
"BASE": "rapidsai/devcontainers:25.04-cpp-cuda11.8-ucx1.17.0-openmpi-ubuntu22.04"
8+
"BASE": "rapidsai/devcontainers:25.04-cpp-cuda11.8-ucx1.18.0-openmpi-ubuntu22.04"
99
}
1010
},
1111
"runArgs": [

‎.devcontainer/cuda12.5-conda/devcontainer.json ‎.devcontainer/cuda12.8-conda/devcontainer.json

+4-4
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,15 @@
33
"context": "${localWorkspaceFolder}/.devcontainer",
44
"dockerfile": "${localWorkspaceFolder}/.devcontainer/Dockerfile",
55
"args": {
6-
"CUDA": "12.5",
6+
"CUDA": "12.8",
77
"PYTHON_PACKAGE_MANAGER": "conda",
88
"BASE": "rapidsai/devcontainers:25.04-cpp-mambaforge-ubuntu22.04"
99
}
1010
},
1111
"runArgs": [
1212
"--rm",
1313
"--name",
14-
"${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.04-cuda12.5-conda"
14+
"${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.04-cuda12.8-conda"
1515
],
1616
"hostRequirements": {"gpu": "optional"},
1717
"features": {
@@ -20,7 +20,7 @@
2020
"overrideFeatureInstallOrder": [
2121
"ghcr.io/rapidsai/devcontainers/features/rapids-build-utils"
2222
],
23-
"initializeCommand": ["/bin/bash", "-c", "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config,conda/pkgs,conda/${localWorkspaceFolderBasename}-cuda12.5-envs}"],
23+
"initializeCommand": ["/bin/bash", "-c", "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config,conda/pkgs,conda/${localWorkspaceFolderBasename}-cuda12.8-envs}"],
2424
"postAttachCommand": ["/bin/bash", "-c", "if [ ${CODESPACES:-false} = 'true' ]; then . devcontainer-utils-post-attach-command; . rapids-post-attach-command; fi"],
2525
"workspaceFolder": "/home/coder",
2626
"workspaceMount": "source=${localWorkspaceFolder},target=/home/coder/cugraph-gnn,type=bind,consistency=consistent",
@@ -29,7 +29,7 @@
2929
"source=${localWorkspaceFolder}/../.cache,target=/home/coder/.cache,type=bind,consistency=consistent",
3030
"source=${localWorkspaceFolder}/../.config,target=/home/coder/.config,type=bind,consistency=consistent",
3131
"source=${localWorkspaceFolder}/../.conda/pkgs,target=/home/coder/.conda/pkgs,type=bind,consistency=consistent",
32-
"source=${localWorkspaceFolder}/../.conda/${localWorkspaceFolderBasename}-cuda12.5-envs,target=/home/coder/.conda/envs,type=bind,consistency=consistent"
32+
"source=${localWorkspaceFolder}/../.conda/${localWorkspaceFolderBasename}-cuda12.8-envs,target=/home/coder/.conda/envs,type=bind,consistency=consistent"
3333
],
3434
"customizations": {
3535
"vscode": {

‎.devcontainer/cuda12.5-pip/devcontainer.json ‎.devcontainer/cuda12.8-pip/devcontainer.json

+6-6
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,20 @@
33
"context": "${localWorkspaceFolder}/.devcontainer",
44
"dockerfile": "${localWorkspaceFolder}/.devcontainer/Dockerfile",
55
"args": {
6-
"CUDA": "12.5",
6+
"CUDA": "12.8",
77
"PYTHON_PACKAGE_MANAGER": "pip",
8-
"BASE": "rapidsai/devcontainers:25.04-cpp-cuda12.5-ucx1.17.0-openmpi-ubuntu22.04"
8+
"BASE": "rapidsai/devcontainers:25.04-cpp-cuda12.8-ucx1.18.0-openmpi-ubuntu22.04"
99
}
1010
},
1111
"runArgs": [
1212
"--rm",
1313
"--name",
14-
"${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.04-cuda12.5-pip"
14+
"${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.04-cuda12.8-pip"
1515
],
1616
"hostRequirements": {"gpu": "optional"},
1717
"features": {
1818
"ghcr.io/rapidsai/devcontainers/features/cuda:25.4": {
19-
"version": "12.5",
19+
"version": "12.8",
2020
"installcuBLAS": true,
2121
"installcuSOLVER": true,
2222
"installcuRAND": true,
@@ -28,15 +28,15 @@
2828
"ghcr.io/rapidsai/devcontainers/features/cuda",
2929
"ghcr.io/rapidsai/devcontainers/features/rapids-build-utils"
3030
],
31-
"initializeCommand": ["/bin/bash", "-c", "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config/pip,local/share/${localWorkspaceFolderBasename}-cuda12.5-venvs}"],
31+
"initializeCommand": ["/bin/bash", "-c", "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config/pip,local/share/${localWorkspaceFolderBasename}-cuda12.8-venvs}"],
3232
"postAttachCommand": ["/bin/bash", "-c", "if [ ${CODESPACES:-false} = 'true' ]; then . devcontainer-utils-post-attach-command; . rapids-post-attach-command; fi"],
3333
"workspaceFolder": "/home/coder",
3434
"workspaceMount": "source=${localWorkspaceFolder},target=/home/coder/cugraph-gnn,type=bind,consistency=consistent",
3535
"mounts": [
3636
"source=${localWorkspaceFolder}/../.aws,target=/home/coder/.aws,type=bind,consistency=consistent",
3737
"source=${localWorkspaceFolder}/../.cache,target=/home/coder/.cache,type=bind,consistency=consistent",
3838
"source=${localWorkspaceFolder}/../.config,target=/home/coder/.config,type=bind,consistency=consistent",
39-
"source=${localWorkspaceFolder}/../.local/share/${localWorkspaceFolderBasename}-cuda12.5-venvs,target=/home/coder/.local/share/venvs,type=bind,consistency=consistent"
39+
"source=${localWorkspaceFolder}/../.local/share/${localWorkspaceFolderBasename}-cuda12.8-venvs,target=/home/coder/.local/share/venvs,type=bind,consistency=consistent"
4040
],
4141
"customizations": {
4242
"vscode": {

‎.github/workflows/pr.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ jobs:
7979
uses: rapidsai/shared-workflows/.github/workflows/[email protected]
8080
with:
8181
arch: '["amd64"]'
82-
cuda: '["12.5"]'
82+
cuda: '["12.8"]'
8383
build_command: |
8484
sccache -z;
8585
build-all --verbose -j$(nproc --ignore=1);

‎conda/environments/all_cuda-118_arch-x86_64.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ dependencies:
3939
- pytest-forked
4040
- pytest-xdist
4141
- pytorch>=2.3
42-
- pytorch_geometric>=2.5,<2.6
42+
- pytorch_geometric>=2.5,<2.7
4343
- rapids-build-backend>=0.3.0,<0.4.0.dev0
4444
- rmm==25.4.*,>=0.0.0a0
4545
- scikit-build-core>=0.10.0

‎conda/environments/all_cuda-121_arch-x86_64.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ dependencies:
4545
- pytest-forked
4646
- pytest-xdist
4747
- pytorch>=2.3
48-
- pytorch_geometric>=2.5,<2.6
48+
- pytorch_geometric>=2.5,<2.7
4949
- rapids-build-backend>=0.3.0,<0.4.0.dev0
5050
- rmm==25.4.*,>=0.0.0a0
5151
- scikit-build-core>=0.10.0

‎conda/environments/all_cuda-124_arch-x86_64.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ dependencies:
4545
- pytest-forked
4646
- pytest-xdist
4747
- pytorch>=2.3
48-
- pytorch_geometric>=2.5,<2.6
48+
- pytorch_geometric>=2.5,<2.7
4949
- rapids-build-backend>=0.3.0,<0.4.0.dev0
5050
- rmm==25.4.*,>=0.0.0a0
5151
- scikit-build-core>=0.10.0

‎conda/recipes/cugraph-pyg/meta.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ requirements:
3636
- cupy >=12.0.0
3737
- cugraph ={{ minor_version }}
3838
- tensordict >=0.1.2
39-
- pytorch_geometric >=2.5,<2.6
39+
- pytorch_geometric >=2.5,<2.7
4040

4141
tests:
4242
imports:

‎dependencies.yaml

+6-2
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,10 @@ dependencies:
229229
cuda: "12.5"
230230
packages:
231231
- cuda-version=12.5
232+
- matrix:
233+
cuda: "12.8"
234+
packages:
235+
- cuda-version=12.8
232236
cuda:
233237
specific:
234238
- output_types: [conda]
@@ -440,10 +444,10 @@ dependencies:
440444
common:
441445
- output_types: [conda]
442446
packages:
443-
- pytorch_geometric>=2.5,<2.6
447+
- pytorch_geometric>=2.5,<2.7
444448
- output_types: [pyproject, requirements]
445449
packages:
446-
- torch-geometric>=2.5,<2.6
450+
- torch-geometric>=2.5,<2.7
447451

448452
depends_on_pylibwholegraph:
449453
common:

‎python/cugraph-pyg/conda/cugraph_pyg_dev_cuda-118.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ dependencies:
1515
- pytest-cov
1616
- pytest-xdist
1717
- pytorch>=2.3
18-
- pytorch_geometric>=2.5,<2.6
18+
- pytorch_geometric>=2.5,<2.7
1919
- tensordict>=0.1.2
2020
- torchdata
2121
name: cugraph_pyg_dev_cuda-118

‎python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_mnmg.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -150,10 +150,10 @@ def load_partitioned_data(
150150
)
151151

152152
# Load features
153-
feature_store["node", "x"] = torch.load(
153+
feature_store["node", "x", None] = torch.load(
154154
os.path.join(feature_path, f"rank={rank}_x.pt")
155155
)
156-
feature_store["node", "y"] = torch.load(
156+
feature_store["node", "y", None] = torch.load(
157157
os.path.join(feature_path, f"rank={rank}_y.pt")
158158
)
159159

‎python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_sg.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -128,8 +128,8 @@ def load_data(
128128
] = data.edge_index
129129

130130
feature_store = cugraph_pyg.data.TensorDictFeatureStore()
131-
feature_store["node", "x"] = data.x
132-
feature_store["node", "y"] = data.y
131+
feature_store["node", "x", None] = data.x
132+
feature_store["node", "y", None] = data.y
133133

134134
return (
135135
(feature_store, graph_store),

‎python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_snmg.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -108,8 +108,8 @@ def run_train(
108108
] = ixr
109109

110110
feature_store = TensorDictFeatureStore()
111-
feature_store["node", "x"] = data.x
112-
feature_store["node", "y"] = data.y
111+
feature_store["node", "x", None] = data.x
112+
feature_store["node", "y", None] = data.y
113113

114114
dist.barrier()
115115

‎python/cugraph-pyg/cugraph_pyg/examples/rgcn_link_class_mnmg.py

+15-9
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2024, NVIDIA CORPORATION.
1+
# Copyright (c) 2024-2025, NVIDIA CORPORATION.
22
# Licensed under the Apache License, Version 2.0 (the "License");
33
# you may not use this file except in compliance with the License.
44
# You may obtain a copy of the License at
@@ -93,7 +93,11 @@ def train(epoch, model, optimizer, train_loader, edge_feature_store, num_steps=N
9393
optimizer.zero_grad()
9494

9595
for i, batch in enumerate(train_loader):
96-
r = edge_feature_store[("n", "e", "n"), "rel"][batch.e_id].flatten().cuda()
96+
r = (
97+
edge_feature_store[("n", "e", "n"), "rel", None][batch.e_id]
98+
.flatten()
99+
.cuda()
100+
)
97101
z = model.encode(batch.edge_index, r)
98102

99103
loss = model.recon_loss(z, batch.edge_index)
@@ -301,13 +305,18 @@ def load_partitioned_data(rank, edge_path, rel_path, pos_path, neg_path, meta_pa
301305
feature_store = TensorDictFeatureStore()
302306
edge_feature_store = WholeFeatureStore()
303307

308+
with open(meta_path, "r") as f:
309+
meta = json.load(f)
310+
311+
print("num nodes:", meta["num_nodes"])
312+
304313
# Load edge index
305-
graph_store[("n", "e", "n"), "coo"] = torch.load(
306-
os.path.join(edge_path, f"rank={rank}.pt")
307-
)
314+
graph_store[
315+
("n", "e", "n"), "coo", False, (meta["num_nodes"], meta["num_nodes"])
316+
] = torch.load(os.path.join(edge_path, f"rank={rank}.pt"))
308317

309318
# Load edge rel type
310-
edge_feature_store[("n", "e", "n"), "rel"] = torch.load(
319+
edge_feature_store[("n", "e", "n"), "rel", None] = torch.load(
311320
os.path.join(rel_path, f"rank={rank}.pt")
312321
)
313322

@@ -333,9 +342,6 @@ def load_partitioned_data(rank, edge_path, rel_path, pos_path, neg_path, meta_pa
333342
splits[stage]["tail_neg"] = tail_neg
334343
splits[stage]["relation"] = relation
335344

336-
with open(meta_path, "r") as f:
337-
meta = json.load(f)
338-
339345
return (feature_store, graph_store), edge_feature_store, splits, meta
340346

341347

‎python/cugraph-pyg/cugraph_pyg/examples/rgcn_link_class_sg.py

+5-3
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2024, NVIDIA CORPORATION.
1+
# Copyright (c) 2024-2025, NVIDIA CORPORATION.
22
# Licensed under the Apache License, Version 2.0 (the "License");
33
# you may not use this file except in compliance with the License.
44
# You may obtain a copy of the License at
@@ -90,8 +90,10 @@ def load_data(
9090
edge_feature_store = TensorDictFeatureStore()
9191
meta = {}
9292

93-
graph_store[("n", "e", "n"), "coo"] = dataset.edge_index
94-
edge_feature_store[("n", "e", "n"), "rel"] = dataset.edge_reltype.pin_memory()
93+
graph_store[
94+
("n", "e", "n"), "coo", False, (dataset.num_nodes, dataset.num_nodes)
95+
] = dataset.edge_index
96+
edge_feature_store[("n", "e", "n"), "rel", None] = dataset.edge_reltype.pin_memory()
9597
meta["num_nodes"] = dataset.num_nodes
9698
meta["num_rels"] = dataset.edge_reltype.max() + 1
9799

‎python/cugraph-pyg/cugraph_pyg/examples/rgcn_link_class_snmg.py

+7-5
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2024, NVIDIA CORPORATION.
1+
# Copyright (c) 2024-2025, NVIDIA CORPORATION.
22
# Licensed under the Apache License, Version 2.0 (the "License");
33
# you may not use this file except in compliance with the License.
44
# You may obtain a copy of the License at
@@ -111,11 +111,13 @@ def load_data(
111111
feature_store = TensorDictFeatureStore() # empty fs required by PyG
112112
edge_feature_store = WholeFeatureStore()
113113

114-
graph_store[("n", "e", "n"), "coo"] = torch.tensor_split(
115-
data.edge_index.cuda(), world_size, dim=1
116-
)[rank]
114+
print("num nodes:", data.num_nodes)
115+
116+
graph_store[
117+
("n", "e", "n"), "coo", False, (data.num_nodes, data.num_nodes)
118+
] = torch.tensor_split(data.edge_index.cuda(), world_size, dim=1)[rank]
117119

118-
edge_feature_store[("n", "e", "n"), "rel"] = torch.tensor_split(
120+
edge_feature_store[("n", "e", "n"), "rel", None] = torch.tensor_split(
119121
data.edge_reltype.cuda(),
120122
world_size,
121123
)[rank]

‎python/cugraph-pyg/cugraph_pyg/tests/data/test_feature_store.py

+5-5
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2024, NVIDIA CORPORATION.
1+
# Copyright (c) 2024-2025, NVIDIA CORPORATION.
22
# Licensed under the Apache License, Version 2.0 (the "License");
33
# you may not use this file except in compliance with the License.
44
# You may obtain a copy of the License at
@@ -30,15 +30,15 @@ def test_tensordict_feature_store_basic_api():
3030

3131
other_features = torch.randint(1024, (10, 5))
3232

33-
feature_store["node", "feat0"] = node_features_0
34-
feature_store["node", "feat1"] = node_features_1
35-
feature_store["other", "feat"] = other_features
33+
feature_store["node", "feat0", None] = node_features_0
34+
feature_store["node", "feat1", None] = node_features_1
35+
feature_store["other", "feat", None] = other_features
3636

3737
assert (feature_store["node"]["feat0"][:] == node_features_0).all()
3838
assert (feature_store["node"]["feat1"][:] == node_features_1).all()
3939
assert (feature_store["other"]["feat"][:] == other_features).all()
4040

4141
assert len(feature_store.get_all_tensor_attrs()) == 3
4242

43-
del feature_store["node", "feat0"]
43+
del feature_store["node", "feat0", None]
4444
assert len(feature_store.get_all_tensor_attrs()) == 2

‎python/cugraph-pyg/cugraph_pyg/tests/data/test_feature_store_mg.py

+9-7
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2024, NVIDIA CORPORATION.
1+
# Copyright (c) 2024-2025, NVIDIA CORPORATION.
22
# Licensed under the Apache License, Version 2.0 (the "License");
33
# you may not use this file except in compliance with the License.
44
# You may obtain a copy of the License at
@@ -46,22 +46,24 @@ def run_test_wholegraph_feature_store_basic_api(rank, world_size, dtype):
4646
features = features.reshape((features.numel() // 100, 100)).to(torch_dtype)
4747

4848
tensordict_store = TensorDictFeatureStore()
49-
tensordict_store["node", "fea"] = features
49+
tensordict_store["node", "fea", None] = features
5050

5151
whole_store = WholeFeatureStore()
52-
whole_store["node", "fea"] = torch.tensor_split(features, world_size)[rank]
52+
whole_store["node", "fea", None] = torch.tensor_split(features, world_size)[rank]
5353

5454
ix = torch.arange(features.shape[0])
5555
assert (
56-
whole_store["node", "fea"][ix].cpu() == tensordict_store["node", "fea"][ix]
56+
whole_store["node", "fea", None][ix].cpu()
57+
== tensordict_store["node", "fea", None][ix]
5758
).all()
5859

5960
label = torch.arange(0, features.shape[0]).reshape((features.shape[0], 1))
60-
tensordict_store["node", "label"] = label
61-
whole_store["node", "label"] = torch.tensor_split(label, world_size)[rank]
61+
tensordict_store["node", "label", None] = label
62+
whole_store["node", "label", None] = torch.tensor_split(label, world_size)[rank]
6263

6364
assert (
64-
whole_store["node", "fea"][ix].cpu() == tensordict_store["node", "fea"][ix]
65+
whole_store["node", "fea", None][ix].cpu()
66+
== tensordict_store["node", "fea", None][ix]
6567
).all()
6668

6769
pylibwholegraph.torch.initialize.finalize()

‎python/cugraph-pyg/cugraph_pyg/tests/data/test_graph_store.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2024, NVIDIA CORPORATION.
1+
# Copyright (c) 2024-2025, NVIDIA CORPORATION.
22
# Licensed under the Apache License, Version 2.0 (the "License");
33
# you may not use this file except in compliance with the License.
44
# You may obtain a copy of the License at
@@ -30,8 +30,12 @@ def test_graph_store_basic_api():
3030

3131
ei = torch.stack([dst, src])
3232

33+
num_nodes = karate.number_of_nodes()
34+
3335
graph_store = GraphStore()
34-
graph_store.put_edge_index(ei, ("person", "knows", "person"), "coo")
36+
graph_store.put_edge_index(
37+
ei, ("person", "knows", "person"), "coo", False, (num_nodes, num_nodes)
38+
)
3539

3640
rei = graph_store.get_edge_index(("person", "knows", "person"), "coo")
3741

‎python/cugraph-pyg/cugraph_pyg/tests/data/test_graph_store_mg.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2024, NVIDIA CORPORATION.
1+
# Copyright (c) 2024-2025, NVIDIA CORPORATION.
22
# Licensed under the Apache License, Version 2.0 (the "License");
33
# you may not use this file except in compliance with the License.
44
# You may obtain a copy of the License at
@@ -30,8 +30,12 @@ def test_graph_store_basic_api_mg():
3030

3131
ei = torch.stack([dst, src])
3232

33+
num_nodes = karate.number_of_nodes()
34+
3335
graph_store = GraphStore(is_multi_gpu=True)
34-
graph_store.put_edge_index(ei, ("person", "knows", "person"), "coo")
36+
graph_store.put_edge_index(
37+
ei, ("person", "knows", "person"), "coo", False, (num_nodes, num_nodes)
38+
)
3539

3640
rei = graph_store.get_edge_index(("person", "knows", "person"), "coo")
3741

0 commit comments

Comments
 (0)