From e0ef42326b8abda2d3ca5a61772ca5a79c13c1f2 Mon Sep 17 00:00:00 2001 From: jchmura Date: Wed, 8 Jul 2026 18:32:57 +0000 Subject: [PATCH 1/5] Fire of all async get at once --- gigl/distributed/base_sampler.py | 90 ++++++++++++++++---------------- 1 file changed, 44 insertions(+), 46 deletions(-) diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 986ba5d58..9f5e3215a 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -1,3 +1,4 @@ +import asyncio from collections import defaultdict from dataclasses import dataclass from typing import Optional, Union @@ -61,6 +62,12 @@ def _stable_unique_preserve_order(nodes: torch.Tensor) -> torch.Tensor: return unique_nodes[stable_order] +def _collapse_single_label_dim(labels: torch.Tensor) -> torch.Tensor: + # DistFeature always returns [N, K]. Collapse K=1 to 1-D [N] to match + # GLT's convention and downstream losses such as CrossEntropyLoss. + return labels if labels.shape[1] > 1 else labels.T[0] + + @dataclass class SampleLoopInputs: """Inputs prepared for the neighbor sampling loop in _sample_from_nodes. @@ -281,23 +288,18 @@ async def _collate_fn( ) input_type = output.input_type assert input_type is not None + futs = {} + label_keys = set() 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] + key = f"{as_str(input_type)}.nlabels" + futs[key] = wrap_torch_future( + self.dist_node_labels.async_get( + output.node[input_type], input_type + ) ) + label_keys.add(key) else: node_labels = self.dist_node_labels.get(input_type, None) if node_labels is not None: @@ -313,17 +315,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( @@ -333,17 +330,19 @@ 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) ) + values = await asyncio.gather(*futs.values()) + for key, value in zip(futs, values): + result_map[key] = ( + _collapse_single_label_dim(value) if key in label_keys else value + ) if output.batch is not None: for ntype, batch in output.batch.items(): result_map[f"{as_str(ntype)}.batch"] = batch @@ -360,33 +359,32 @@ async def _collate_fn( ) if self.with_edge: result_map["eids"] = output.edge + futs = {} + label_keys = set() 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) + ) + values = await asyncio.gather(*futs.values()) + for key, value in zip(futs, values): + result_map[key] = ( + _collapse_single_label_dim(value) if key in label_keys else value + ) if output.batch is not None: result_map["batch"] = output.batch From 89137973ac24f66d8c9ca8adb2ff6b4b591e55d5 Mon Sep 17 00:00:00 2001 From: jchmura Date: Wed, 8 Jul 2026 18:38:03 +0000 Subject: [PATCH 2/5] Revert comment delete --- gigl/distributed/base_sampler.py | 32 ++++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 9f5e3215a..867e0534c 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -62,12 +62,6 @@ def _stable_unique_preserve_order(nodes: torch.Tensor) -> torch.Tensor: return unique_nodes[stable_order] -def _collapse_single_label_dim(labels: torch.Tensor) -> torch.Tensor: - # DistFeature always returns [N, K]. Collapse K=1 to 1-D [N] to match - # GLT's convention and downstream losses such as CrossEntropyLoss. - return labels if labels.shape[1] > 1 else labels.T[0] - - @dataclass class SampleLoopInputs: """Inputs prepared for the neighbor sampling loop in _sample_from_nodes. @@ -340,9 +334,16 @@ async def _collate_fn( ) values = await asyncio.gather(*futs.values()) for key, value in zip(futs, values): - result_map[key] = ( - _collapse_single_label_dim(value) if key in label_keys else value - ) + if 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[key] = value if output.batch is not None: for ntype, batch in output.batch.items(): result_map[f"{as_str(ntype)}.batch"] = batch @@ -382,9 +383,16 @@ async def _collate_fn( ) values = await asyncio.gather(*futs.values()) for key, value in zip(futs, values): - result_map[key] = ( - _collapse_single_label_dim(value) if key in label_keys else value - ) + if 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[key] = value if output.batch is not None: result_map["batch"] = output.batch From 2110fff99e07227de2cb13d468a5f57eb664f8ff Mon Sep 17 00:00:00 2001 From: jchmura Date: Wed, 8 Jul 2026 20:03:09 +0000 Subject: [PATCH 3/5] Unify naming to result_key --- gigl/distributed/base_sampler.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 867e0534c..5ff0f5eba 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -287,13 +287,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): - key = f"{as_str(input_type)}.nlabels" - futs[key] = wrap_torch_future( + 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(key) + label_keys.add(result_key) else: node_labels = self.dist_node_labels.get(input_type, None) if node_labels is not None: @@ -333,8 +333,8 @@ async def _collate_fn( self.dist_edge_feature.async_get(eids, etype) ) values = await asyncio.gather(*futs.values()) - for key, value in zip(futs, values): - if key in label_keys: + 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 @@ -343,7 +343,7 @@ async def _collate_fn( # 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[key] = value + result_map[result_key] = value if output.batch is not None: for ntype, batch in output.batch.items(): result_map[f"{as_str(ntype)}.batch"] = batch @@ -382,8 +382,8 @@ async def _collate_fn( self.dist_edge_feature.async_get(eids) ) values = await asyncio.gather(*futs.values()) - for key, value in zip(futs, values): - if key in label_keys: + 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 @@ -392,7 +392,7 @@ async def _collate_fn( # 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[key] = value + result_map[result_key] = value if output.batch is not None: result_map["batch"] = output.batch From a1ceb39a58e3fc8a216466212daf59f8cafd8f23 Mon Sep 17 00:00:00 2001 From: jchmura Date: Wed, 8 Jul 2026 20:48:10 +0000 Subject: [PATCH 4/5] Extract shared async get logic in both homogeneous and heterogenous blocks --- gigl/distributed/base_sampler.py | 57 ++++++++++++-------------------- 1 file changed, 22 insertions(+), 35 deletions(-) diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 5ff0f5eba..b0d7b3a7a 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -231,13 +231,12 @@ async def _collate_fn( (GLT 0.2.4). The method name preserves GLT's original typo so that this override is matched correctly at runtime. - The only behavioural change from the GLT original is in the ``DistFeature`` - label-fetch paths (both homogeneous and heterogeneous): GLT writes - ``nlabels.T[0]``, which silently discards all label columns beyond the first - and breaks multi-label node classification. This override writes the full - ``nlabels`` tensor instead, avoiding the extra RPC call that a super()-then- - re-fetch approach would require. The non-``DistFeature`` path (plain - ``torch.Tensor`` labels) is unchanged — it never applied ``.T[0]``. + Behavioural changes from the GLT original: + 1. For ``DistFeature`` label fetches, preserve multi-label tensors by + collapsing only K=1 labels from [N, K] to [N]. + 2. In non-all2all mode, issue 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 @@ -261,6 +260,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 @@ -282,8 +284,6 @@ async def _collate_fn( ) input_type = output.input_type assert input_type is not None - futs = {} - label_keys = set() if not isinstance(input_type, tuple): if self.dist_node_labels is not None: if isinstance(self.dist_node_labels, DistFeature): @@ -332,18 +332,6 @@ async def _collate_fn( futs[result_key] = wrap_torch_future( self.dist_edge_feature.async_get(eids, etype) ) - 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 if output.batch is not None: for ntype, batch in output.batch.items(): result_map[f"{as_str(ntype)}.batch"] = batch @@ -360,8 +348,6 @@ async def _collate_fn( ) if self.with_edge: result_map["eids"] = output.edge - futs = {} - label_keys = set() if self.dist_node_labels is not None: if isinstance(self.dist_node_labels, DistFeature): futs["nlabels"] = wrap_torch_future( @@ -381,21 +367,22 @@ async def _collate_fn( futs["efeats"] = wrap_torch_future( self.dist_edge_feature.async_get(eids) ) - 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 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( From 67728d4278e9fb651e78a25883399cbd7f4f1987 Mon Sep 17 00:00:00 2001 From: jchmura Date: Wed, 8 Jul 2026 20:50:13 +0000 Subject: [PATCH 5/5] Update docstring --- gigl/distributed/base_sampler.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index b0d7b3a7a..4530e292c 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -231,12 +231,17 @@ async def _collate_fn( (GLT 0.2.4). The method name preserves GLT's original typo so that this override is matched correctly at runtime. - Behavioural changes from the GLT original: - 1. For ``DistFeature`` label fetches, preserve multi-label tensors by - collapsing only K=1 labels from [N, K] to [N]. - 2. In non-all2all mode, issue all independent ``async_get`` requests - before awaiting them, so label, node-feature, and edge-feature - fetches can overlap. + The only behavioural change from the GLT original is in the ``DistFeature`` + label-fetch paths (both homogeneous and heterogeneous): GLT writes + ``nlabels.T[0]``, which silently discards all label columns beyond the first + and breaks multi-label node classification. This override writes the full + ``nlabels`` tensor instead, avoiding the extra RPC call that a super()-then- + 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