Skip to content

Concurrently Issue All dist_feature.async_get()Requests in Base Sampler _collate_fn() #690

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jchmura-sc wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Concurrently Issue All dist_feature.async_get()Requests in Base Sampler _collate_fn() #690

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e0ef423
Fire of all async get at once
jchmura-sc Jul 8, 2026
8913797
Revert comment delete
jchmura-sc Jul 8, 2026
2110fff
Unify naming to result_key
jchmura-sc Jul 8, 2026
a1ceb39
Extract shared async get logic in both homogeneous and heterogenous b…
jchmura-sc Jul 8, 2026
67728d4
Update docstring
jchmura-sc Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 44 additions & 46 deletions gigl/distributed/base_sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import asyncio
from collections import defaultdict
from dataclasses import dataclass
from typing import Optional, Union
Expand Down Expand Up @@ -238,6 +239,10 @@ async def _collate_fn(
re-fetch approach would require. The non-``DistFeature`` path (plain
``torch.Tensor`` labels) is unchanged — it never applied ``.T[0]``.

In non-all2all mode, this method also issues all independent
``async_get`` requests before awaiting them, so label, node-feature, and
edge-feature fetches can overlap.

# TODO (mkolodner-sc): Now that GiGL owns this method, investigate whether
# post-processing steps in DistNeighborLoader._collate_fn can be folded in
# here and simplified — e.g. set_missing_features (populating empty tensors
Expand All @@ -260,6 +265,9 @@ async def _collate_fn(
for k, v in output.metadata.items():
result_map[f"#META.{k}"] = v

futs: dict[str, asyncio.Future[torch.Tensor]] = {}
label_keys: set[str] = set()

if is_hetero:
for ntype, nodes in output.node.items():
result_map[f"{as_str(ntype)}.ids"] = nodes
Expand All @@ -284,20 +292,13 @@ async def _collate_fn(
if not isinstance(input_type, tuple):
if self.dist_node_labels is not None:
if isinstance(self.dist_node_labels, DistFeature):
fut = self.dist_node_labels.async_get(
output.node[input_type], input_type
)
nlabels = await wrap_torch_future(fut)
# DistFeature always returns [N, K]. We collapse K=1 to 1-D
# [N] to match GLT's convention and what downstream code
# (e.g. CrossEntropyLoss) expects for data.y. Multi-label
# (K>1) keeps the full 2-D matrix.
# TODO (mkolodner-sc): Consider investigating always returning
# 2-D — this may be a breaking change for single-label
# training pipelines (e.g. CrossEntropyLoss expects 1-D data.y).
result_map[f"{as_str(input_type)}.nlabels"] = (
nlabels if nlabels.shape[1] > 1 else nlabels.T[0]
result_key = f"{as_str(input_type)}.nlabels"
futs[result_key] = wrap_torch_future(
self.dist_node_labels.async_get(
output.node[input_type], input_type
)
)
label_keys.add(result_key)
else:
node_labels = self.dist_node_labels.get(input_type, None)
if node_labels is not None:
Expand All @@ -313,17 +314,12 @@ async def _collate_fn(
for ntype, nfeats in nfeat_dict.items():
result_map[f"{as_str(ntype)}.nfeats"] = nfeats
else:
nfeat_fut_dict = {}
for ntype, nodes in output.node.items():
nodes = nodes.to(torch.long)
nfeat_fut_dict[ntype] = self.dist_node_feature.async_get(
nodes, ntype
futs[f"{as_str(ntype)}.nfeats"] = wrap_torch_future(
self.dist_node_feature.async_get(nodes, ntype)
)
for ntype, fut in nfeat_fut_dict.items():
nfeats = await wrap_torch_future(fut)
result_map[f"{as_str(ntype)}.nfeats"] = nfeats
if self.dist_edge_feature is not None and self.with_edge:
efeat_fut_dict = {}
for etype in self.edge_types:
if self.edge_dir == "in":
eids = result_map.get(
Expand All @@ -333,16 +329,13 @@ async def _collate_fn(
eids = result_map.get(f"{as_str(etype)}.eids", None)
if eids is not None:
eids = eids.to(torch.long)
efeat_fut_dict[etype] = self.dist_edge_feature.async_get(
eids, etype
result_key = (
f"{as_str(etype)}.efeats"
if self.edge_dir == "out"
else f"{as_str(reverse_edge_type(etype))}.efeats"
)
for etype, fut in efeat_fut_dict.items():
efeats = await wrap_torch_future(fut)
if self.edge_dir == "out":
result_map[f"{as_str(etype)}.efeats"] = efeats
elif self.edge_dir == "in":
result_map[f"{as_str(reverse_edge_type(etype))}.efeats"] = (
efeats
futs[result_key] = wrap_torch_future(
self.dist_edge_feature.async_get(eids, etype)
)
if output.batch is not None:
for ntype, batch in output.batch.items():
Expand All @@ -362,34 +355,39 @@ async def _collate_fn(
result_map["eids"] = output.edge
if self.dist_node_labels is not None:
if isinstance(self.dist_node_labels, DistFeature):
fut = self.dist_node_labels.async_get(output.node)
nlabels = await wrap_torch_future(fut)
# DistFeature always returns [N, K]. We collapse K=1 to 1-D
# [N] to match GLT's convention and what downstream code
# (e.g. CrossEntropyLoss) expects for data.y. Multi-label
# (K>1) keeps the full 2-D matrix.
# TODO (mkolodner-sc): Consider investigating always returning
# 2-D — this may be a breaking change for single-label
# training pipelines (e.g. CrossEntropyLoss expects 1-D data.y).
result_map["nlabels"] = (
nlabels if nlabels.shape[1] > 1 else nlabels.T[0]
futs["nlabels"] = wrap_torch_future(
self.dist_node_labels.async_get(output.node)
)
label_keys.add("nlabels")
else:
result_map["nlabels"] = self.dist_node_labels[
output.node.to(self.dist_node_labels.device)
]
if self.dist_node_feature is not None:
fut = self.dist_node_feature.async_get(output.node)
nfeats = await wrap_torch_future(fut)
result_map["nfeats"] = nfeats
futs["nfeats"] = wrap_torch_future(
self.dist_node_feature.async_get(output.node)
)
if self.dist_edge_feature is not None:
eids = result_map["eids"]
fut = self.dist_edge_feature.async_get(eids)
efeats = await wrap_torch_future(fut)
result_map["efeats"] = efeats
futs["efeats"] = wrap_torch_future(
self.dist_edge_feature.async_get(eids)
)
if output.batch is not None:
result_map["batch"] = output.batch

values = await asyncio.gather(*futs.values())
for result_key, value in zip(futs, values):
if result_key in label_keys:
# DistFeature always returns [N, K]. We collapse K=1 to 1-D
# [N] to match GLT's convention and what downstream code
# (e.g. CrossEntropyLoss) expects for data.y. Multi-label
# (K>1) keeps the full 2-D matrix.
# TODO (mkolodner-sc): Consider investigating always returning
# 2-D — this may be a breaking change for single-label
# training pipelines (e.g. CrossEntropyLoss expects 1-D data.y).
value = value if value.shape[1] > 1 else value.T[0]
result_map[result_key] = value

return result_map

async def _sample_from_nodes(
Expand Down