Add get_training_input to storage_utils #438

kmontemayor2-sc · 2026-01-16T18:58:45Z

Scope of work done

Add server-side util so we can remotely fetch the training input.

Since this is kind of a big PR not adding the client-side equivalent in this one :P

Again this is server-side code, and it's really meant to be called by users.

Where is the documentation for this feature?: N/A

Did you add automated tests or write a test plan?

Updated Changelog.md? NO

Ready for code review?: NO

mkolodner-sc

Thanks Kyle! Did a pass here and left some comments/questions

mkolodner-sc · 2026-01-20T22:57:39Z

python/gigl/distributed/graph_store/storage_utils.py

+
+
+def get_training_input(
+    split: Union[Literal["train", "val", "test"], str],


Is this sufficient? What if we want "all" nodes (i.e. dataset.node_ids)

That wouldn't be "training"input" would it? We can add another function to do that in the future (get_all_nodes?)

It could be for the random negative loader for link prediction training, which would need some dataset.node_ids or equivalent.

We can add another function to do that in the future (get_all_nodes?)

We wouldn't need a whole different function, we could just specify some 'all' split and if its that we use _dataset.node_ids.

Hmmmm, I do think that the tuple[Tensor, Tensor, Tensor | None] is important for the ABLP.

I guess I could rename this to get_ablp_input or something? Would that ameliorate your concerns?

Discussed offline, renaming will be fine here for this function, and in a follow-up we will refactor the get_node_ids_on_rank utility so that it can be used to split, making it extendable for the SNC use case, and can be called in this function to reduce the duplicity. Can we add a TODO here in the meantime?

Sure added todos :)

python/gigl/distributed/graph_store/storage_utils.py

python/tests/test_assets/distributed/utils.py

python/tests/unit/distributed/graph_store/storage_utils_test.py

mkolodner-sc · 2026-01-21T18:10:01Z

python/gigl/distributed/graph_store/storage_utils.py

+
+
+def get_training_input(
+    split: Union[Literal["train", "val", "test"], str],


It could be for the random negative loader for link prediction training, which would need some dataset.node_ids or equivalent.

We can add another function to do that in the future (get_all_nodes?)

We wouldn't need a whole different function, we could just specify some 'all' split and if its that we use _dataset.node_ids.

python/gigl/distributed/graph_store/storage_utils.py

mkolodner-sc · 2026-01-21T18:17:34Z

python/gigl/distributed/graph_store/storage_utils.py

+            f"Anchor nodes must be a torch.Tensor or a dict[NodeType, torch.Tensor], got {type(anchors)}"
+        )
+
+    anchors_for_rank = shard_nodes_by_process(anchor_nodes, rank, world_size)


I remember we had a discussion at some point of potentially moving this up to be user-facing instead of done under-the-hood by a GiGL utility/class. Would be curious to get your thoughts on this decision here as well

hmmmm, we did and I think we agreed we should do less of this - though we still do it in GS mode.

I went with this approach in GS mode so we minimize network traffic.

Since this is how we are currently doing this for RemoteDistDataset and friends, should we keep this approach for now and re-visit in the future, for all node fetching?

Adding some get_for_rank bool flag is probably sufficient here? WDYT about that?

I'm fine either way if you'd prefer to do this now or a follow-up. If in follow-up, let's just make sure we add a TODO here

I had the same qq.
My understanding was that there was some agreement where the client should have complicated logic of knowing what data to fetch from where.

Now whether the user prompts the client code to fetch exactly what they want, or uses this sort of utility on the client to implicitly fetch some split of nodes, the logic lives inside the client and the server is dumb and just fetches the data requested.

I see storage layer as a sort of a "db" specific to graph sampling.

If complex code to decide what data to provide to what client brings forward some smells, speicifcally I think it:

Breaks determinism (The operations are implicit, and outside of control from client)

Makes queries non-portable i.e. we are stuck with defining queries inside of the server vs client controlled

Couples the query patterns to runtime topology; which may change in the future and the coupling will induce extra eng effort int the future to circumvent

Potentially break ability to do replica/retry strategies (if needed in future) as the client doesn't know what data to expect and alignment between the two will be difficult unless client hosts logic.

If we do decide to go for this route I think we should prove out that it is for some sane reason i.e. actually proving out the network traffic argument. Would love elaboration here.

python/tests/unit/distributed/graph_store/storage_utils_test.py

mkolodner-sc

Thanks Kyle for the work!

svij-sc · 2026-01-23T01:39:00Z

python/tests/unit/distributed/graph_store/storage_utils_test.py

 class TestRemoteDataset(unittest.TestCase):
    def setUp(self) -> None:
-        """Reset the global dataset before each test."""
+        """Reset the global dataset and initialize process group before each test."""


and initialize process group before each test.

This comment doesn't seem accurate i.e. should we be initializing here in this method?

svij-sc · 2026-01-23T01:40:28Z

python/tests/test_assets/distributed/utils.py

    )


+def destroy_test_process_group() -> None:


nit: Lazy abstraction to a function, esp cause only used in one place.

svij-sc · 2026-01-23T01:45:16Z

python/tests/unit/distributed/graph_store/storage_utils_test.py

+
+        The split ratios are calculated as:
+        - num_val = len(val_user_ids) / total_users
+        - num_test = len(test_user_ids) / total_users


I am a little confused reading this; isn't the user already passing in val_user_ids? why do we need to recompute the val split if one is being provided?

svij-sc · 2026-01-23T01:50:57Z

python/tests/unit/distributed/graph_store/storage_utils_test.py

+
+        Creates a dataset with:
+        - USER nodes: [0, 1, 2, 3, 4]
+        - STORY nodes: [0, 1, 2, 3, 4]


mega nit: STORY --> ITEM
More generic term for OSS

svij-sc · 2026-01-23T01:53:40Z

python/tests/unit/distributed/graph_store/storage_utils_test.py

+
+        # Set up edge partition books and edge indices
+        edge_partition_book = {
+            _USER_TO_STORY: torch.zeros(5, dtype=torch.int64),


magic number 5.
Is this length of edges?

Similar magic numbers below re edge index.

svij-sc · 2026-01-23T01:54:23Z

python/tests/unit/distributed/graph_store/storage_utils_test.py

+
+        Creates a dataset with:
+        - USER nodes: [0, 1, 2, 3, 4]
+        - STORY nodes: [0, 1, 2, 3, 4]


Should we make these arguments? i.e. user_node_ids, item_node_ids, ?
Or atleast default arguments/constants that can be referenced below to make the function more modular?

svij-sc · 2026-01-23T02:02:22Z

python/gigl/distributed/graph_store/storage_utils.py

+            f"Anchor nodes must be a torch.Tensor or a dict[NodeType, torch.Tensor], got {type(anchors)}"
+        )
+
+    anchors_for_rank = shard_nodes_by_process(anchor_nodes, rank, world_size)


I had the same qq.
My understanding was that there was some agreement where the client should have complicated logic of knowing what data to fetch from where.

Now whether the user prompts the client code to fetch exactly what they want, or uses this sort of utility on the client to implicitly fetch some split of nodes, the logic lives inside the client and the server is dumb and just fetches the data requested.

svij-sc · 2026-01-23T02:12:19Z

python/gigl/distributed/graph_store/storage_utils.py

+            f"Anchor nodes must be a torch.Tensor or a dict[NodeType, torch.Tensor], got {type(anchors)}"
+        )
+
+    anchors_for_rank = shard_nodes_by_process(anchor_nodes, rank, world_size)


I see storage layer as a sort of a "db" specific to graph sampling.

If complex code to decide what data to provide to what client brings forward some smells, speicifcally I think it:

Breaks determinism (The operations are implicit, and outside of control from client)

Makes queries non-portable i.e. we are stuck with defining queries inside of the server vs client controlled

Couples the query patterns to runtime topology; which may change in the future and the coupling will induce extra eng effort int the future to circumvent

Potentially break ability to do replica/retry strategies (if needed in future) as the client doesn't know what data to expect and alignment between the two will be difficult unless client hosts logic.

If we do decide to go for this route I think we should prove out that it is for some sane reason i.e. actually proving out the network traffic argument. Would love elaboration here.

kmonte added 2 commits January 16, 2026 18:55

Add get_training_input to storage_utils

b5038e5

fix

3b62990

mkolodner-sc reviewed Jan 20, 2026

View reviewed changes

kmonte added 3 commits January 21, 2026 00:41

Merge branch 'main' into kmonte/storage-util-training-input

b392b77

address comments

c9659fc

address comments

7028fa3

mkolodner-sc reviewed Jan 21, 2026

View reviewed changes

kmonte added 2 commits January 21, 2026 19:09

Merge branch 'main' into kmonte/storage-util-training-input

c2860e4

address comments

6659a82

mkolodner-sc approved these changes Jan 21, 2026

View reviewed changes

kmonte added 3 commits January 21, 2026 19:59

address comments

494b844

Merge branch 'main' into kmonte/storage-util-training-input

559f2ed

Merge branch 'main' into kmonte/storage-util-training-input

b9702ac

svij-sc reviewed Jan 23, 2026

View reviewed changes



		def get_training_input(
		split: Union[Literal["train", "val", "test"], str],

Add get_training_input to storage_utils #438

Are you sure you want to change the base?

Add get_training_input to storage_utils #438

Conversation

kmontemayor2-sc commented Jan 16, 2026

Uh oh!

mkolodner-sc left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mkolodner-sc Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mkolodner-sc left a comment

Choose a reason for hiding this comment

Uh oh!

svij-sc Jan 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

mkolodner-sc Jan 21, 2026 •

edited

Loading

svij-sc Jan 23, 2026 •

edited

Loading