Skip to content

Derive with_edge from edge feature presence (perf)#694

Merged
kmontemayor2-sc merged 4 commits into
mainfrom
with-edge-optimization
Jul 15, 2026
Merged

Derive with_edge from edge feature presence (perf)#694
kmontemayor2-sc merged 4 commits into
mainfrom
with-edge-optimization

Conversation

@kmontemayor2-sc

Copy link
Copy Markdown
Collaborator

What

BaseDistLoader.create_sampling_config hardcoded with_edge=True for every
loader. This derives it from whether the dataset actually has edge features:

  with_edge = dataset_schema.edge_feature_info is not None

DatasetSchema.edge_feature_info is already populated in both colocated and
graph-store modes (fetched from storage rank 0 in graph-store mode), so the
derivation is correct on compute nodes as well.

Why

For datasets with no edge features anywhere, with_edge=True forces GLT to
collect sampled edge ids per hop and transport them in every SampleMessage.
Those eids are pure overhead when there are no edge features to gather: the
sampler writes an int64 eid tensor of one element per sampled edge (gated on
self.with_edge). With with_edge=False, GLT takes the sample path instead
of sample_with_edge, output.edge is None, no eids are collected or
serialized, and edge-feature-store construction is skipped.

Behavior change (changelog)

Batches produced from featureless datasets no longer carry GLT's data.edge
(homogeneous) or data[etype].edge (heterogeneous) sampled-edge-id attribute.
A repo-wide grep found no consumer that reads this attribute off a produced
batch — the only in-repo references are the sampler's own write sites. External
code that reads batch.edge on a featureless dataset would now find it
absent/None. Datasets that do have edge features are unaffected.

Tests

  • Direct unit tests for create_sampling_config: with_edge is False when
    edge_feature_info is None, True for homogeneous and heterogeneous edge
    features.
  • End-to-end: a featureless homogeneous dataset yields batches with no sampled
    edge ids (data.edge is None).
  • Weighted sampling without edge features still avoids zero-weight edges.
  • PPR unaffected (outputs come from metadata, not sampled eids).
  • Graph-store provenance: compute-side edge_feature_info is populated from the
    storage cluster.

All affected files pass; make type_check passes.

@kmontemayor2-sc

Copy link
Copy Markdown
Collaborator Author

/all_test

@github-actions

github-actions Bot commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 00:38:38UTC : 🔄 Lint Test started.

@ 00:46:08UTC : ✅ Workflow completed successfully.

@github-actions

github-actions Bot commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 00:38:38UTC : 🔄 E2E Test started.

@ 02:13:17UTC : ❌ Workflow failed.
Please check the logs for more details.

@github-actions

github-actions Bot commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 00:38:39UTC : 🔄 C++ Unit Test started.

@ 00:40:22UTC : ✅ Workflow completed successfully.

@github-actions

github-actions Bot commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 00:38:41UTC : 🔄 Python Unit Test started.

@ 01:54:03UTC : ✅ Workflow completed successfully.

@github-actions

github-actions Bot commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 00:38:42UTC : 🔄 Scala Unit Test started.

@ 00:48:27UTC : ✅ Workflow completed successfully.

@github-actions

github-actions Bot commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 00:38:42UTC : 🔄 Integration Test started.

@ 02:02:27UTC : ✅ Workflow completed successfully.

kmontemayor and others added 3 commits July 14, 2026 15:15
… hardcoding True

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ph-store modes

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@kmontemayor2-sc
kmontemayor2-sc force-pushed the with-edge-optimization branch from 5cd3258 to 93c3291 Compare July 14, 2026 15:17

@mkolodner-sc mkolodner-sc left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks Kyle for this optimization! Generally LGTM with some minor comments

Comment thread tests/unit/distributed/distributed_neighborloader_test.py Outdated
Comment thread tests/unit/distributed/graph_store/remote_dist_dataset_test.py
…ith_edge test docs

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@kmontemayor2-sc
kmontemayor2-sc marked this pull request as ready for review July 14, 2026 23:35
@kmontemayor2-sc
kmontemayor2-sc added this pull request to the merge queue Jul 14, 2026
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Jul 15, 2026
@kmontemayor2-sc
kmontemayor2-sc added this pull request to the merge queue Jul 15, 2026
Merged via the queue into main with commit 5c3a329 Jul 15, 2026
7 checks passed
@kmontemayor2-sc
kmontemayor2-sc deleted the with-edge-optimization branch July 15, 2026 17:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants