gigl.distributed.graph_store.dist_server

GiGL implementation of GLT DistServer.

Uses GiGL’s DistSamplingProducer which supports neighbor sampling and ABLP (anchor-based link prediction) via BaseGiGLSampler subclasses (DistNeighborSampler for k-hop, DistPPRNeighborSampler for PPR).

Based on alibaba/graphlearn-for-pytorch

Teardown protocol

Cluster teardown is a strict three-phase sequence:

1. Per-loader (compute side). Each DistLoader.shutdown() issues DistServer.destroy_sampling_input(channel_id) against every storage server it registered with. This call performs the actual teardown work on the server: runtime.unregister_input for the channel and, when the channel is the last on its backend, runtime.shutdown for the backend.

2. Per-compute-process. After all loaders are torn down, the compute process calls gigl.distributed.graph_store.compute.shutdown_compute_process, which calls glt.distributed.shutdown_client and tears down the compute torch process group.

3. Per-storage-process. wait_and_shutdown_server blocks until DistServer.exit flips, then runs DistServer.shutdown() (strict validation — see below), barrier(), shutdown_rpc().

An example fo the teardown protocol: ```python def compute_process():

train_loader = … val_loader = … test_loader = …

loaders = [train_loader, val_loader, test_loader] for data in train_loader:

train(data) if should_val():

for data in val_loader:

val(data)

# Shutdown the loaders after training and validation. train_loader.shutdown() val_loader.shutdown()

for data in test_loader:

test(data)

# Shutdown the loader after testing. test_loader.shutdown()

# Step 2: Per-compute-process shutdown_compute_process()

# Step 3: Per-storage-process wait_and_shutdown_server() ```

DistServer.shutdown() does no teardown work itself. It validates that phase 1 ran for every channel/backend, raises RuntimeError if state remains, and drops residual tombstone/stats bookkeeping. The actual runtime.shutdown() work happens in destroy_sampling_input when the last channel on a backend is destroyed. wait_and_shutdown_server catches DistServer.shutdown exceptions so a buggy compute client cannot wedge healthy storage peers on the barrier; the failing storage process re-raises after the barrier so the orchestrator sees a non-zero exit.

Attributes

FETCH_SLOW_LOG_SECS
R
SERVER_EXIT_STATUS_CHECK_INTERVAL	Interval (in seconds) to check exit status of server.
logger

Classes

ChannelState	Per-channel state for a registered sampling input.
DistServer	A server that supports launching remote sampling workers for
SamplingBackendState	Per-backend state for a shared sampling backend.

Functions

get_server()	Get the DistServer instance on the current process.
init_server(num_servers, server_rank, dataset, ...[, ...])	Initialize the current process as a server and establish connections
wait_and_shutdown_server()	Block until all clients have shut down, then shut down the

Module Contents

class gigl.distributed.graph_store.dist_server.ChannelState

Per-channel state for a registered sampling input.

Parameters:

• backend_id – The ID of the backend this channel belongs to.

• worker_key – The unique key identifying this compute-rank channel.

• channel – The shared-memory channel for passing sampled messages.

• epoch – The last epoch started on this channel.

• lock – A reentrant lock guarding channel-level operations.

• tombstoned – Terminal-state tombstone. Set to True once destroy_sampling_input has finished cleaning up this channel. The channel object may briefly outlive its registry entry — an in-flight fetch_one_sampled_message on another RPC thread reads this flag to break out of its recv loop. Past-participle name: this is “done” state, not “in progress” state. Asymmetric with SamplingBackendState.tearing_down by design — backends are removed from the registry when teardown completes, so no terminal flag is needed there; channels are not removed until the in-flight fetch observes the tombstone, so the flag has to live on the object.

backend_id: int

channel: graphlearn_torch.channel.ShmChannel

epoch: int = -1

lock: threading.RLock

tombstoned: bool = False

worker_key: str

class gigl.distributed.graph_store.dist_server.DistServer(dataset, log_every_n=50)

A server that supports launching remote sampling workers for training clients.

Note that this server is enabled only when the distribution mode is a server-client framework, and the graph and feature store will be partitioned and managed by all server nodes.

Parameters:

• dataset (DistDataset) – The DistDataset object of a partition of graph data and feature data, along with distributed partition books.

• log_every_n (int) – Log aggregated fetch_one_sampled_message timing stats after every N fetch attempts per channel.

destroy_sampling_input(channel_id)

Destroy one registered sampling channel and maybe its backend.

If this is the last channel on the backend, the backend is shut down and removed.

Parameters:

channel_id (int) – The ID of the channel to destroy.

Return type:

None

exit()

Set the exit flag to True.

Return type:

bool

fetch_one_sampled_message(channel_id)

Fetch one sampled message from a registered channel.

Parameters:

channel_id (int) – The ID of the channel to fetch from.

Returns:

A tuple of (message, is_done). If is_done is True, no more messages will be produced for this epoch.

Return type:

tuple[Optional[graphlearn_torch.channel.SampleMessage], bool]

get_ablp_input(request)

Get the ABLP (Anchor Based Link Prediction) input for distributed processing.

The caller always provides supervision_edge_type in outward form (anchor_node_type, relation, supervision_node_type), matching the taskMetadata convention.

When the dataset is in-sampled (edge_dir == "in") the stored positive/negative label edges are reversed relative to that outward form, so we reverse the supervision edge type before resolving the stored label edge types. This mirrors the colocated loader’s handling in _setup_for_colocated().

Parameters:

request (gigl.distributed.graph_store.messages.FetchABLPInputRequest) – The ABLP fetch request, including split, node type, supervision edge type, and an optional contiguous server slice.

Returns:

A tuple containing the anchor nodes for the rank, the positive labels, and the negative labels. The positive labels are of shape [N, M], where N is the number of anchor nodes and M is the number of positive labels. The negative labels are of shape [N, M], where N is the number of anchor nodes and M is the number of negative labels. The negative labels may be None if no negative labels are available.

Raises:

ValueError – If the split is invalid.

Return type:

tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]

get_dataset_meta()

Get the meta info of the distributed dataset managed by the current server, including partition info and graph types.

Return type:

tuple[int, int, Optional[list[gigl.src.common.types.graph_data.NodeType]], Optional[list[gigl.src.common.types.graph_data.EdgeType]]]

get_edge_dir()

Get the edge direction from the dataset.

Returns:

The edge direction.

Return type:

Literal[‘in’, ‘out’]

get_edge_feature_info()

Get edge feature information from the dataset.

Returns:

• A single FeatureInfo object for homogeneous graphs

• A dict mapping EdgeType to FeatureInfo for heterogeneous graphs

• None if no edge features are available

Return type:

Edge feature information, which can be

get_edge_index(edge_type, layout)

Parameters:

• edge_type (Optional[gigl.src.common.types.graph_data.EdgeType])

• layout (str)

Return type:

tuple[torch.Tensor, torch.Tensor]

get_edge_partition_book(edge_type)

Gets the partition book for the specified edge type. :param edge_type: The edge type to look up. Must be None for

homogeneous datasets and non-None for heterogeneous ones.

Returns:

The partition book for the requested edge type, or None if no partition book is available.

Raises:

ValueError – If edge_type is mismatched with the dataset type.

Parameters:

edge_type (Optional[gigl.src.common.types.graph_data.EdgeType])

Return type:

Optional[graphlearn_torch.partition.PartitionBook]

get_edge_size(edge_type, layout)

Parameters:

• edge_type (Optional[gigl.src.common.types.graph_data.EdgeType])

• layout (str)

Return type:

tuple[int, int]

get_edge_types()

Get the edge types from the dataset.

Returns:

The edge types in the dataset, None if the dataset is homogeneous.

Return type:

Optional[list[gigl.src.common.types.graph_data.EdgeType]]

get_edge_weights_registered()

Return whether edge weights were registered in the dataset.

Returns:

True if edge weights were registered via DistPartitioner.register_edge_weights().

Return type:

bool

get_node_feature(node_type, index)

Parameters:

• node_type (Optional[gigl.src.common.types.graph_data.NodeType])

• index (torch.Tensor)

Return type:

torch.Tensor

get_node_feature_info()

Get node feature information from the dataset.

Returns:

• A single FeatureInfo object for homogeneous graphs

• A dict mapping NodeType to FeatureInfo for heterogeneous graphs

• None if no node features are available

Return type:

Node feature information, which can be

get_node_ids(request)

Get the node ids from the dataset.

Parameters:

request (gigl.distributed.graph_store.messages.FetchNodesRequest) – The node-fetch request, including split, node type, and an optional contiguous server slice.

Returns:

The node ids.

Raises:

• ValueError –

  • If the split is invalid

  • If the node ids are not a torch.Tensor or a dict[NodeType, torch.Tensor]

  • If the node type is provided for a homogeneous dataset

  • If the node ids are not a dict[NodeType, torch.Tensor] when no node type is provided

• Note – When split=None, all nodes are queryable. This means nodes from any

• split (train, val, or test) may be returned. This is useful when you need –

• to sample neighbors during inference, as neighbor nodes may belong to any split. –

Return type:

torch.Tensor

get_node_label(node_type, index)

Parameters:

• node_type (Optional[gigl.src.common.types.graph_data.NodeType])

• index (torch.Tensor)

Return type:

torch.Tensor

get_node_partition_book(node_type)

Gets the partition book for the specified node type.

Parameters:

node_type (Optional[gigl.src.common.types.graph_data.NodeType]) – The node type to look up. Must be None for homogeneous datasets and non-None for heterogeneous ones.

Returns:

The partition book for the requested node type, or None if no partition book is available.

Raises:

ValueError – If node_type is mismatched with the dataset type.

Return type:

Optional[graphlearn_torch.partition.PartitionBook]

get_node_partition_id(node_type, index)

Parameters:

• node_type (Optional[gigl.src.common.types.graph_data.NodeType])

• index (torch.Tensor)

Return type:

Optional[torch.Tensor]

get_node_types()

Get the node types from the dataset.

Returns:

The node types in the dataset, None if the dataset is homogeneous.

Return type:

Optional[list[gigl.src.common.types.graph_data.NodeType]]

get_tensor_size(node_type)

Parameters:

node_type (Optional[gigl.src.common.types.graph_data.NodeType])

Return type:

torch.Size

init_sampling_backend(opts)

Create or reuse a shared sampling backend for one loader instance.

If a backend with the same backend_key already exists, blocks until the original initializer has finished, then returns its ID. If the original initializer failed, re-raises that failure so the second caller does not see a half-initialized backend.

Parameters:

opts (gigl.distributed.graph_store.messages.InitSamplingBackendRequest) – The initialization request containing the backend key, worker options, sampler options, and sampling config.

Returns:

The unique backend ID.

Raises:

RuntimeError – If a prior concurrent initialization for the same backend_key failed, or if the backend is currently being torn down by destroy_sampling_input.

Return type:

int

register_sampling_input(opts)

Register one compute-rank input channel on an existing backend.

Parameters:

opts (gigl.distributed.graph_store.messages.RegisterBackendRequest) – The registration request containing the backend ID, worker key, sampler input, sampling config, and buffer settings.

Returns:

The unique channel ID for this input.

Raises:

• KeyError – If opts.backend_id does not refer to a registered backend (caller bug; backend must be registered first).

• Exception – Re-raises any failure from runtime.register_input after rolling back the partial channel state.

Return type:

int

shutdown()

Final post-teardown bookkeeping cleanup — does not perform any shutdown.

Lenient contract: callers are expected to have torn down every channel and backend via destroy_sampling_input before invoking this method. The actual runtime.shutdown() work happens in destroy_sampling_input when the last channel on a backend is destroyed; this method only drops residual tombstone/stats bookkeeping and warns if any sampling state is still registered.

Return type:

None

start_new_epoch_sampling(channel_id, epoch)

Start one new epoch on one registered channel.

No-op if this channel has already started epoch or a later epoch (idempotent — safe to call repeatedly from retries).

No-op if the channel is in the tombstoned set: this handles the legitimate destroy/start race window where a compute peer’s start RPC arrives after destroy has landed.

Parameters:

• channel_id (int) – The ID of the channel to start the epoch on.

• epoch (int) – The epoch number to start.

Raises:

RuntimeError – If channel_id was never registered on this server (vs. legitimately tombstoned — tombstoned ids are treated as a silent no-op).

Return type:

None

wait_for_exit()

Block until the exit flag been set to True.

Return type:

None

dataset

class gigl.distributed.graph_store.dist_server.SamplingBackendState

Per-backend state for a shared sampling backend.

Parameters:

• backend_id – The unique ID of this backend.

• backend_key – The key identifying this backend (e.g. "dist_neighbor_loader_0").

• runtime – The shared sampling backend runtime.

• active_channels – Set of channel IDs currently registered on this backend.

• lock – A reentrant lock guarding backend-level operations.

• init_complete – Whether runtime.init_backend() has completed successfully.

• init_error – If runtime.init_backend() raised, the exception; otherwise None.

• tearing_down – In-progress marker. Set to True while destroy_sampling_input is tearing down this backend’s runtime. A second-caller of init_sampling_backend that blocks on lock and observes this flag raises RuntimeError so it does not reuse a half-shutdown runtime. Present-participle name: this is “in progress” state, not “done” state. Asymmetric with ChannelState.tombstoned by design — the registry entry is removed once teardown completes, so no terminal flag is needed; after removal nothing can find this object.

active_channels: set[int]

backend_id: int

backend_key: str

init_complete: bool = False

init_error: BaseException | None = None

lock: threading.RLock

runtime: gigl.distributed.graph_store.shared_dist_sampling_producer.SharedDistSamplingBackend

tearing_down: bool = False

gigl.distributed.graph_store.dist_server.get_server()

Get the DistServer instance on the current process.

Return type:

DistServer

gigl.distributed.graph_store.dist_server.init_server(num_servers, server_rank, dataset, master_addr, master_port, num_clients=0, num_rpc_threads=16, request_timeout=180, server_group_name=None, is_dynamic=False)

Initialize the current process as a server and establish connections with all other servers and clients. Note that this method should be called only in the server-client distribution mode.

Parameters:

• num_servers (int) – Number of processes participating in the server group.

• server_rank (int) – Rank of the current process withing the server group (it should be a number between 0 and num_servers-1).

• dataset (DistDataset) – The DistDataset object of a partition of graph data and feature data, along with distributed partition book info.

• master_addr (str) – The master TCP address for RPC connection between all servers and clients, the value of this parameter should be same for all servers and clients.

• master_port (int) – The master TCP port for RPC connection between all servers and clients, the value of this parameter should be same for all servers and clients.

• num_clients (int) – Number of processes participating in the client group. if is_dynamic is True, this parameter will be ignored.

• num_rpc_threads (int) – The number of RPC worker threads used for the current server to respond remote requests. (Default: 16).

• request_timeout (int) – The max timeout seconds for remote requests, otherwise an exception will be raised. (Default: 16).

• server_group_name (str) – A unique name of the server group that current process belongs to. If set to None, a default name will be used. (Default: None).

• is_dynamic (bool) – Whether the world size is dynamic. (Default: False).

Return type:

None

gigl.distributed.graph_store.dist_server.wait_and_shutdown_server()

Block until all clients have shut down, then shut down the server on the current process and destroy all RPC connections.

Best-effort cluster liveness: if DistServer.shutdown raises (e.g. because a client crashed leaving state behind) we capture the exception, run barrier() + shutdown_rpc() so healthy storage peers do not hang on the barrier, then re-raise so the orchestrator sees a non-zero exit on the failing storage process.

Step 3 of the three-phase teardown described in the module docstring.

Return type:

None

gigl.distributed.graph_store.dist_server.FETCH_SLOW_LOG_SECS = 1.0

gigl.distributed.graph_store.dist_server.R

gigl.distributed.graph_store.dist_server.SERVER_EXIT_STATUS_CHECK_INTERVAL = 5.0

Interval (in seconds) to check exit status of server.

gigl.distributed.graph_store.dist_server.logger