Enabled Weighted Sampling

[mkolodner-sc]

Summary

Adds native weighted edge sampling to GiGL's distributed training pipeline via GLT's CPUWeightedSampler. When enabled, neighbors are sampled proportionally to edge weights rather than uniformly.

New API

• DistPartitioner.register_edge_weights(edge_weights) — registers a 1D per-edge weight tensor (homogeneous or dict[EdgeType, Tensor] for heterogeneous) before calling partition_edge_index_and_edge_features(). Weights are partitioned alongside edge features in the same pass (co-partitioned, mirroring the node features + labels pattern).
• load_torch_tensors_from_tf_record(weight_edge_feat_name=...) — accepts the name of an existing edge feature column to extract as sampling weights during TFRecord loading. The column is sliced out of the feature tensor and stored in LoadedGraphTensors.edge_weights; it is never duplicated in memory.
• build_dataset(weight_edge_feat_name=...) — threads weight_edge_feat_name through to TFRecord loading and then calls register_edge_weights() with the extracted weights.
• DistNeighborLoader(with_weight=True) / DistABLPLoader(with_weight=True) — enables weighted sampling. Defaults to False; must be set explicitly.
• BaseDistLoader.validate_with_weight() — shared validation: raises ValueError if with_weight=True but no weights are registered in the dataset; raises NotImplementedError if used with PPRSamplerOptions (weight-proportional PPR residual propagation is deferred to a future PR).

Implementation notes

• LoadedGraphTensors.edge_weights — new field carrying extracted weights from TFRecord loading through to register_edge_weights().
• GraphPartitionData.weights (field already existed) carries the partitioned weight tensor to DistDataset._initialize_graph(), which forwards it to GLT's init_graph(edge_weights=...).
• DistDataset.has_edge_weights property reflects whether weights were registered at construction time.
• SamplingConfig.with_weight is now threaded through from the loader rather than hardcoded to False.
• Graph Store mode: DistServer.get_edge_weights_registered() and RemoteDistDataset.fetch_edge_weights_registered() propagate has_edge_weights across the RPC boundary so compute nodes can validate with_weight against the remote dataset.

Tests

• tests/unit/distributed/distributed_weighted_sampling_test.py (8 new tests):
  • Correctness (homogeneous + heterogeneous): weight=0 edges to "bad" nodes are never traversed in sampled subgraphs — verified by encoding node class in features and asserting no bad node appears after weighted sampling.
  • Partitioner edge cases: features only, weights only, neither, both (with consistency check that GraphPartitionData.edge_ids == FeaturePartitionData.ids), and heterogeneous partial weights (one edge type weighted, another not).
• tests/unit/common/data/dataloaders_test.py (1 new test): test_load_edge_weights_from_tf_record — verifies that load_torch_tensors_from_tf_record correctly extracts a named column into edge_weights, removes it from edge_features, and returns the right shapes and values.

[kmontemayor2-sc]

Thanks Matt! Me and the robots did a first pass, it's possible they're imagining some of the issues here but I figured I'd flag :)

[mkolodner-sc]

/unit_test

[github-actions]

GiGL Automation

@ 23:03:34UTC : 🔄 Python Unit Test started.

@ 24:11:45UTC : ❌ Workflow failed.
Please check the logs for more details.

[github-actions]

GiGL Automation

@ 23:03:35UTC : 🔄 C++ Unit Test started.

@ 23:05:34UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 23:03:35UTC : 🔄 Scala Unit Test started.

@ 23:13:48UTC : ✅ Workflow completed successfully.

[mkolodner-sc]

/unit_test

[github-actions]

GiGL Automation

@ 06:09:08UTC : 🔄 Scala Unit Test started.

@ 06:19:33UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 06:09:09UTC : 🔄 C++ Unit Test started.

@ 06:13:19UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 06:09:09UTC : 🔄 Python Unit Test started.

@ 07:13:24UTC : ✅ Workflow completed successfully.

[kmontemayor2-sc]

Thanks Matt!

[mkolodner-sc]

/unit_test

[github-actions]

GiGL Automation

@ 22:16:50UTC : 🔄 Python Unit Test started.

@ 23:36:03UTC : ❌ Workflow failed.
Please check the logs for more details.

[github-actions]

GiGL Automation

@ 22:16:51UTC : 🔄 Scala Unit Test started.

@ 22:26:16UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 22:16:51UTC : 🔄 C++ Unit Test started.

@ 22:18:35UTC : ✅ Workflow completed successfully.