Fix PPR sampler memory and labeled homogeneous ABLP#645
Conversation
… into mkolodner-sc/ppr_gs_memory # Conflicts: # gigl/distributed/dist_ppr_sampler.py
|
/all_tests |
GiGL Automation@ 19:17:28UTC : 🔄 @ 19:19:34UTC : ✅ Workflow completed successfully. |
GiGL Automation@ 19:17:29UTC : 🔄 @ 20:14:17UTC : ✅ Workflow completed successfully. |
GiGL Automation@ 19:17:30UTC : 🔄 @ 20:47:24UTC : ❌ Workflow failed. |
GiGL Automation@ 19:17:30UTC : 🔄 @ 19:26:20UTC : ✅ Workflow completed successfully. |
GiGL Automation@ 19:17:31UTC : 🔄 @ 20:30:59UTC : ✅ Workflow completed successfully. |
GiGL Automation@ 19:17:31UTC : 🔄 @ 19:28:22UTC : ✅ Workflow completed successfully. |
|
/e2e_test |
GiGL Automation@ 20:48:41UTC : 🔄 @ 22:34:48UTC : ❌ Workflow failed. |
|
/e2e_test |
GiGL Automation@ 21:47:12UTC : 🔄 @ 23:07:38UTC : ❌ Workflow failed. |
|
/e2e_test |
GiGL Automation@ 22:42:25UTC : 🔄 @ 24:51:17UTC : ❌ Workflow failed. |
|
/e2e_test |
GiGL Automation@ 24:24:11UTC : 🔄 @ 01:47:47UTC : ❌ Workflow failed. |
|
/e2e_test |
GiGL Automation@ 24:38:06UTC : 🔄 @ 02:04:12UTC : ✅ Workflow completed successfully. |
…emory # Conflicts: # tests/unit/distributed/utils/degree_test.py
mkolodner-sc
changed the title
kmontemayor2-sc
left a comment
kmontemayor2-sc
left a comment
There was a problem hiding this comment.
Thanks Matt! Did a first pass here, fwiw I feel like this could have been multiple PRs for the different fixes / etc but this pr is fine as-is.
| self._degree_tensor = compute_and_broadcast_degree_tensor( | ||
| self.graph, self._edge_dir | ||
| ) | ||
| share_memory(entity=self._degree_tensor) |
There was a problem hiding this comment.
ugh sorry to go back and forth on this, it may be weird to do this here instead of in share_ipc? Or do we not call degree_tensor until the subprocess launch?
In fact we do already share memory in share_ipc 1, is the problem here that we don't call degree_tensor until after we're in the subprocesses already?
Can you remind me where we first call degree_tensor and where that's located in the process tree?
There was a problem hiding this comment.
Can you remind me where we first call degree_tensor and where that's located in the process tree?
Currently, the degree tensor is not built when we call build_dataset. Instead, it is first built after we are inside one of the data loaders and know we are doing PPR sampling (specified by the SamplingOptions).
In fact we do already share memory in share_ipc 1, is the problem here that we don't call degree_tensor until after we're in the subprocesses already?
Yes exactly, this was calling failures specifically on large-scale graph store cases since it was creating copies of the tensor before we had an opportunity to call share_ipc (and thereby share the memory of the degree tensor). The original fix I had here handled this by explicitly calling share memory before handing the degree tensor to the sampling workers for the GS setting, where otherwise it was creating many copies of this degree tensor.
I think the current solution to call share_memory immediately after building the degree tensor makes the most sense, since it doesn't need to rely on share_ipc and provides the most confidence for putting the degree tensor onto shared memory.
There was a problem hiding this comment.
We'd still have local world size copies of this object in shared memory then right? Unless we want to create the degree tensor in the top-process?
IIRC our goal here was to abstract this away from the user s.t. they don't need to pass in this implementation detail into the loaders right?
And in that vein it's probably weird if we ask them to call dataset.degree_tensor naked in the top-level process? I think ideally we can poke at /dev/shm directly to setup the shared memory properly across the local nodes here.
My concern with calling share_memory inside the property is that it may be surprising to users share memory tensors do have additional restrictions on them iirc - I guess the distdataset already does that under the hood so maybe it's not too bad.
I guess I have a slight preference here for calling share_memory in our loaders after we get this tensor, but that's up to you. WDYT?
There was a problem hiding this comment.
We'd still have local world size copies of this object in shared memory then right? Unless we want to create the degree tensor in the top-process?
IIRC our goal here was to abstract this away from the user s.t. they don't need to pass in this implementation detail into the loaders right?
And in that vein it's probably weird if we ask them to call dataset.degree_tensor naked in the top-level process? I think ideally we can poke at /dev/shm directly to setup the shared memory properly across the local nodes here.
Correct, this was a consideration made in the initial design. The tradeoff with the current approach is that there may still be local_world_size copies of the degree tensor, but since the degree tensor isn't impacting the memory bottleneck, we proceeded with this approach. Since this is the intention of the design and isn't blocking training or inference with PPR, I'd prefer if we save any further optimization as a follow-up if that becomes a blocker. One note is that the graph_store case doesn't have this constraint I believe, since it makes the share_memory call from the DistServer, so there are not local world size copies in that setting.
My concern with calling share_memory inside the property is that it may be surprising to users share memory tensors do have additional restrictions on them iirc - I guess the distdataset already does that under the hood so maybe it's not too bad.
I guess I have a slight preference here for calling share_memory in our loaders after we get this tensor, but that's up to you. WDYT?
That makes sense too, I can move the share_memory call back to the loaders in that case.
3 participants
Summary
Makes the PPR sampler cheaper and fixes the labeled-homogeneous ABLP edge case that surfaced while exercising PPR through Graph Store.
Changes include:
DistDataset.degree_tensor.int32and share them across sampling workers instead of rebuilding/copying per worker.DistPPRNeighborSamplerto consume precomputed degree tensors directly.etype=Nonefor true homogeneous graphs.Databatches.