gigl.distributed.utils.dist_sampler#

Sampler factory helpers shared across sampling producers.

Attributes#

SamplerInput

Union of all supported sampler input types.

SamplerRuntime

Union of all supported GiGL sampler runtime types.

Functions#

create_dist_sampler(*, data, sampling_config, ...)

Create a GiGL sampler runtime for one channel on one worker.

Module Contents#

gigl.distributed.utils.dist_sampler.create_dist_sampler(*, data, sampling_config, worker_options, channel, sampler_options, degree_tensors, current_device)[source]#

Create a GiGL sampler runtime for one channel on one worker.

Parameters:

  • data (graphlearn_torch.distributed.DistDataset) – The distributed dataset containing graph topology and features.

  • sampling_config (graphlearn_torch.sampler.SamplingConfig) – Configuration for sampling behavior (neighbors, edges, etc.).

  • worker_options (Union[graphlearn_torch.distributed.MpDistSamplingWorkerOptions, graphlearn_torch.distributed.RemoteDistSamplingWorkerOptions]) – Worker-level options (RPC settings, device placement, concurrency).

  • channel (graphlearn_torch.channel.ChannelBase) – The communication channel for passing sampled messages.

  • sampler_options (gigl.distributed.sampler_options.SamplerOptions) – Algorithm-specific options (k-hop or PPR).

  • degree_tensors (Optional[Union[torch.Tensor, dict[graphlearn_torch.typing.NodeType, torch.Tensor]]]) – Pre-computed degree tensors required by PPR sampling. Must not be None when sampler_options is PPRSamplerOptions.

  • current_device (torch.device) – The device on which sampling will run.

Returns:

A configured sampler runtime, either DistNeighborSampler or DistPPRNeighborSampler.

Raises:

NotImplementedError – If sampler_options is an unsupported type.

Return type:

SamplerRuntime

Note

When sampling_config.seed is unset we draw one here, per worker. The seed belongs to the process that samples rather than to the config: the config is shared across ranks in Graph Store mode and compared for equality, so a value drawn while building it would differ per rank and be rejected.

A seed is a large performance win, not just a determinism knob. GLT’s CPURandomSampler::UniformSample reads RandomSeedManager::getInstance().getSeed() on every call, and that call happens once per source row of a batch:

The std::mt19937 it feeds is thread_local static, so after the first call the read is discarded – but it still runs. With no seed set, getSeed() constructs a std::random_device and draws from it, about 5 us of real work per source row for a value that is thrown away:

For the production use case we see up to 7x speed up, and 29x speedup in local testing.

Passing the seed to the sampler is what sets it: GLT calls RandomSeedManager::getInstance().setSeed from NeighborSampler.__init__ only when the seed is not None. That manager is process-global and the generator it feeds is thread_local, so the first sampler built on a worker thread fixes that thread’s stream. Per-channel seeds therefore are not a determinism guarantee.

TODO(kmonte): Drop this workaround if GLT ever hoists the getSeed() call out of CPURandomSampler::UniformSample and into the engine initializer, so the unseeded path stops paying per-row entropy.

gigl.distributed.utils.dist_sampler.SamplerInput[source]#

Union of all supported sampler input types.

gigl.distributed.utils.dist_sampler.SamplerRuntime[source]#

Union of all supported GiGL sampler runtime types.