Add CONTIGUOUS shard strategy (#545)

Co-authored-by: kmontemayor <kyle.e.montemayor@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: kmontemayor <kmontemayor@snapchat.com>

Author: 4 people

gigl/distributed/dist_ablp_neighborloader.py

@@ -666,7 +666,12 @@ def _setup_for_graph_store(
         # Extract supervision edge types and derive label edge types from the
         # ABLPInputNodes.labels dict (keyed by supervision edge type).
         self._supervision_edge_types = list(first_input.labels.keys())
-        has_negatives = any(neg is not None for _, neg in first_input.labels.values())
+        has_negatives = False
+        for ablp_input in input_nodes.values():
+            for maybe_negative_labels in ablp_input.labels.values():
+                if maybe_negative_labels is not None:
+                    has_negatives = True
+                    break
 
         self._positive_label_edge_types = [
             message_passing_to_positive_label(et) for et in self._supervision_edge_types

gigl/distributed/graph_store/dist_server.py

@@ -39,6 +39,7 @@
     FetchABLPInputRequest,
     FetchNodesRequest,
 )
+from gigl.distributed.graph_store.sharding import ServerSlice
 from gigl.distributed.sampler import ABLPNodeSamplerInput
 from gigl.distributed.sampler_options import PPRSamplerOptions, SamplerOptions
 from gigl.distributed.utils.neighborloader import shard_nodes_by_process
@@ -283,7 +284,7 @@ def get_node_ids(
 
         Args:
             request: The node-fetch request, including split, node type,
-                and round-robin rank/world_size.
+                and either round-robin rank/world_size or a contiguous slice.
 
         Returns:
             The node ids.
@@ -306,6 +307,7 @@ def get_node_ids(
             node_type=request.node_type,
             rank=request.rank,
             world_size=request.world_size,
+            server_slice=request.server_slice,
         )
 
     def _get_node_ids(
@@ -314,6 +316,7 @@ def _get_node_ids(
         node_type: Optional[NodeType],
         rank: Optional[int] = None,
         world_size: Optional[int] = None,
+        server_slice: Optional[ServerSlice] = None,
     ) -> torch.Tensor:
         """Core implementation for fetching node IDs by split, type, and sharding.
 
@@ -366,6 +369,8 @@ def _get_node_ids(
                 f"node_type was not provided, so node ids must be a torch.Tensor (e.g. a homogeneous dataset), got {type(nodes)}."
             )
 
+        if server_slice is not None:
+            return server_slice.slice_tensor(nodes)
         if rank is not None and world_size is not None:
             return shard_nodes_by_process(nodes, rank, world_size)
         return nodes
@@ -420,6 +425,7 @@ def get_ablp_input(
             node_type=request.node_type,
             rank=request.rank,
             world_size=request.world_size,
+            server_slice=request.server_slice,
         )
         positive_label_edge_type, negative_label_edge_type = select_label_edge_types(
             request.supervision_edge_type, self.dataset.get_edge_types()

gigl/distributed/graph_store/messages.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass
 from typing import Literal, Optional, Union
 
+from gigl.distributed.graph_store.sharding import ServerSlice
 from gigl.src.common.types.graph_data import EdgeType, NodeType
 
 
@@ -17,6 +18,9 @@ class FetchNodesRequest:
             Must be provided together with ``rank``.
         split: The split of the dataset to get node ids from.
         node_type: The type of nodes to get node ids for.
+        server_slice: An optional :class:`~gigl.distributed.graph_store.sharding.ServerSlice`
+            describing the fraction of this server's data to return.
+            Cannot be combined with ``rank``/``world_size``.
 
     Examples:
         Fetch all nodes without sharding:
@@ -36,18 +40,25 @@ class FetchNodesRequest:
     world_size: Optional[int] = None
     split: Optional[Union[Literal["train", "val", "test"], str]] = None
     node_type: Optional[NodeType] = None
+    server_slice: Optional[ServerSlice] = None
 
     def validate(self) -> None:
         """Validate that the request has consistent rank/world_size.
 
         Raises:
-            ValueError: If only one of ``rank`` or ``world_size`` is provided.
+            ValueError:
+                If only one of ``rank`` or ``world_size`` is provided.
+                If ``server_slice`` is provided together with ``rank`` or ``world_size``.
         """
         if (self.rank is None) ^ (self.world_size is None):
             raise ValueError(
                 "rank and world_size must be provided together. "
                 f"Received rank={self.rank}, world_size={self.world_size}"
             )
+        if self.server_slice is not None and (
+            self.rank is not None or self.world_size is not None
+        ):
+            raise ValueError("server_slice cannot be combined with rank/world_size.")
 
 
 @dataclass(frozen=True)
@@ -62,6 +73,9 @@ class FetchABLPInputRequest:
             Must be provided together with ``world_size``.
         world_size: The total number of processes in the distributed setup.
             Must be provided together with ``rank``.
+        server_slice: An optional :class:`~gigl.distributed.graph_store.sharding.ServerSlice`
+            describing the fraction of this server's data to return.
+            Cannot be combined with ``rank``/``world_size``.
 
     Examples:
         Fetch training ABLP input without sharding:
@@ -78,15 +92,22 @@ class FetchABLPInputRequest:
     supervision_edge_type: EdgeType
     rank: Optional[int] = None
     world_size: Optional[int] = None
+    server_slice: Optional[ServerSlice] = None
 
     def validate(self) -> None:
         """Validate that the request has consistent rank/world_size.
 
         Raises:
-            ValueError: If only one of ``rank`` or ``world_size`` is provided.
+            ValueError:
+                If only one of ``rank`` or ``world_size`` is provided.
+                If ``server_slice`` is provided together with ``rank`` or ``world_size``.
         """
         if (self.rank is None) ^ (self.world_size is None):
             raise ValueError(
                 "rank and world_size must be provided together. "
                 f"Received rank={self.rank}, world_size={self.world_size}"
             )
+        if self.server_slice is not None and (
+            self.rank is not None or self.world_size is not None
+        ):
+            raise ValueError("server_slice cannot be combined with rank/world_size.")