Apply and verify local patches on the pinned GraphLearn-Torch build - #761
Open
dsaini2-sc wants to merge 1 commit into
Open
Apply and verify local patches on the pinned GraphLearn-Torch build#761dsaini2-sc wants to merge 1 commit into
dsaini2-sc wants to merge 1 commit into
Conversation
install_glt.sh applies gigl/scripts/patches/*.patch in order on top of the pinned GLT commit with patch --forward, hard-failing on any non-clean apply, and after pip install runs verify_glt_patches.py, which fails the build unless every patched behavior is observable in the installed wheel. Patches: 0001 exact bitmap distinct-count in CPU graph init (replaces at::_unique's ~3x transient allocation), 0002 int32 CSR column-id support in the CPU samplers (byte-identical sampling to int64; untaught paths reject loudly), 0003 shm-queue unpin/cleanup on teardown. tests/unit/utils/glt_int32_indices_test.py re-establishes the int32 parity guarantees per wheel and skips wholesale on an unpatched wheel. Also fixes post_install.py swallowing install_glt.sh's exit status when invoked as a file, which image builds do; a failed GLT build/install/verify now fails the Docker layer, with a subprocess regression test.
dsaini2-sc
requested review from
kmontemayor2-sc,
mkolodner-sc,
nshah-sc,
svij-sc,
xgao4-sc,
yliu2-sc and
zfan3-sc
as code owners
August 27, 2026 20:39
![semgrep-code-snapchat[bot]](https://avatars.githubusercontent.com/u/40505220?s=60&v=4){kind=small}
kmontemayor2-sc
left a comment
kmontemayor2-sc
left a comment
Collaborator
There was a problem hiding this comment.
Thanks Deepak!
I would prefer if we can leave some TODOs here so we can migrate to our own fork of GLT once that's up
Comment on lines
+19
to
+22
| + // Exact distinct-count of the column ids, replacing at::_unique. _unique is sort-based and | ||
| + // allocates ~3x the size of `indices` in transient anonymous memory (measured 3.00x at 200M | ||
| + // edges) -- ~94 GiB at a 4.2B-edge partition, which OOM-kills billion-edge CPU deployments | ||
| + // whose host budget assumes graph init holds no second copy of the topology. |
Collaborator
There was a problem hiding this comment.
BTW, have we measured this change in isolation? I remember I did some experimentation here before and I found that it didn't really move the needle.
Collaborator
There was a problem hiding this comment.
To be clear, I meant "in isolation, at prod scale"
Collaborator
There was a problem hiding this comment.
Can we add some description at the top of each patch for what it's doing?
Comment on lines
+117
to
+147
| - // Exact distinct-count of the column ids, replacing at::_unique. _unique is sort-based and | ||
| - // allocates ~3x the size of `indices` in transient anonymous memory (measured 3.00x at 200M | ||
| - // edges) -- ~94 GiB at a 4.2B-edge partition, which OOM-kills billion-edge CPU deployments | ||
| - // whose host budget assumes graph init holds no second copy of the topology. | ||
| - // | ||
| - // A seen-bitmap over [0, max_col] gives the identical count in two sequential passes using | ||
| - // max_col/8 + 1 bytes (~115 MiB at 1e9 nodes). It is only VALID for non-negative ids in | ||
| - // contiguous storage, and only CHEAPER when the id domain is dense relative to the edge count. | ||
| - // Both preconditions are checked rather than assumed: an unchecked negative id would cast to a | ||
| - // huge uint64_t and write outside the bitmap, and a sparse-but-huge id would size the bitmap | ||
| - // from the id domain instead of the edge count. The sparse case falls back to _unique, so this | ||
| - // is never worse than upstream; the invalid cases raise instead of corrupting memory. | ||
| - // | ||
| - // Note on strides: `col_idx_` is flat storage and every sampler already reads it that way, so | ||
| - // requiring contiguity here makes col_count_ agree with what sampling actually traverses -- | ||
| - // upstream's _unique(indices) followed the view's strides and could disagree with its own | ||
| - // samplers. | ||
| - TORCH_CHECK(indices.is_contiguous(), | ||
| - "InitCPUGraphFromCSR requires contiguous indices: col_idx_ is read as flat " | ||
| - "storage by the samplers, so a strided view would sample the wrong columns"); | ||
| - int64_t max_col = -1; | ||
| - for (int64_t i = 0; i < edge_count_; ++i) { | ||
| - const int64_t col = col_idx_[i]; | ||
| - TORCH_CHECK(col >= 0, | ||
| - "InitCPUGraphFromCSR requires non-negative column ids (they are node ids), got ", | ||
| - col, " at position ", i); | ||
| - max_col = std::max(max_col, col); | ||
| - } | ||
| - if (max_col < 0) { | ||
| - // No edges. at::_unique of an empty tensor is also empty. | ||
| - col_count_ = 0; |
Collaborator
There was a problem hiding this comment.
Do you think it's preferable to have just one patch here? We're overriding the 0001 patch here right?
Might make it more maintable / readable?
WDYT?
Comment on lines
+232
to
+233
| + uint32_t seed = RandomSeedManager::getInstance().getSeed(); | ||
| + thread_local static std::mt19937 engine(seed); |
Collaborator
There was a problem hiding this comment.
If we're doing this patch can we also apply alibaba/graphlearn-for-pytorch#167? It's a pretty big sampling speedup.
Comment on lines
+62
to
+65
| for glt_patch in "${GIGL_SCRIPTS_DIR}"/patches/*.patch; do | ||
| echo "Applying ${glt_patch}" | ||
| patch -p1 --forward < "${glt_patch}" || { echo "FATAL: ${glt_patch} did not apply"; exit 1; } | ||
| done |
Collaborator
There was a problem hiding this comment.
Can you add a TODO to migrate the patches to our own fork of GLT once we get that up?
Comment on lines
+14
to
+20
| build on the patches being live in the installed wheel by @dsaini2 in https://github.com/Snapchat/GiGL/pull/PENDING | ||
|
|
||
| ### Fixed | ||
|
|
||
| - `gigl/scripts/post_install.py` now propagates `install_glt.sh`'s exit status as its own process exit code; previously | ||
| a failed GLT build/install exited 0 when the file was invoked directly, as image builds do by @dsaini2 in | ||
| https://github.com/Snapchat/GiGL/pull/PENDING |
Collaborator
There was a problem hiding this comment.
nit update pending?
Collaborator
There was a problem hiding this comment.
There's a robot review that this is unreachable - I assume we're planning on wiring this in as a follow up ?
GLT's public
Topology.__init__/Dataset.init_graph upcast edge indices to int64 (data.type(torch.int64) in GLT's python/data/graph.py:53), and
DistDataset._initialize_graph() goes through that path. Only the tests and verify_glt_patches.py reach the int32 C++ code, via
Topology.__new__. So today, production training silently gets zero benefit from patch 0002. Codex confirmed this empirically through the
installed wheel. Action: if wiring int32 topology into DistDataset is a planned follow-up, say so in the CHANGELOG/patch header; otherwise this
patch is dead weight. When you do wire it, add a test that goes through the real Dataset.init_graph path — the silent upcast means a mis-wired
integration would pass every current gate.
zfan3-sc
approved these changes
Aug 28, 2026
zfan3-sc
left a comment
zfan3-sc
left a comment
Collaborator
There was a problem hiding this comment.
Giving a 2nd stamp to unblock merging; Please merge after addressing Kyle's comments. Thanks for the work!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Scope of work done
GiGL already builds GraphLearn-Torch from a pinned commit (
88ff111ac) ingigl/scripts/install_glt.sh, precisely because needed fixes (GLT PRs 151/153/154) were merged upstream but never released. This PR extends that mechanism one step: local patch files applied on top of the pin before the wheel builds, plus a post-install gate proving the patches are live in the installed wheel.Three patches in
gigl/scripts/patches/, applied in order:0001(74 lines): replacesat::_unique's distinct-count inInitCPUGraphFromCSRwith an exact seen-bitmap count._uniqueis sort-based and transiently allocates ~3x the size ofindices(measured 3.00x at 200M edges), which OOM-kills CPU graph init at billion-edge scale. The bitmap is exact, two sequential passes, ~max_col/8 bytes; it falls back to_uniquewhen the id domain is sparse relative to the edge count, rejects negative ids (which would index outside the bitmap), and rejects non-contiguous storage. Never worse than upstream.0002(491 lines): int32 CSR column-id support in the CPU samplers. When node ids fit in int32 the CSC column array halves. Exactly one ofcol_idx_/col_idx32_is set; samplers taught the layout sample byte-identically to int64 and still emit int64 ids; samplers not taught it (weighted, CUDA init) reject int32 loudly viaTORCH_CHECKinstead of reading a null pointer. An int64-only deployment is untouched.0003(131 lines): shared-memory queue unpin/cleanup on teardown -- the SysV segment is detached and removed (andcudaHostUnregistercalled for pinned mappings) from the shared_ptr deleter, so repeated loader construction/teardown cycles stop leaking segments and pinned registrations. An in-flight zero-copy message keeps its mapping alive until released.The apply loop uses
patch -p1 --forwardwith an explicit failure exit, so a re-run or an upstream change that absorbs a patch stops the build loudly rather than shipping an unpatched wheel.gigl/scripts/verify_glt_patches.pythen runs afterpip installand fails the build unless every patched behavior is observable in the installed wheel -- applying is a property of the source tree, shipping is a property of the compiled.so, and the gap between them (stale build dir, second wheel on the path) is silent.tests/unit/utils/glt_int32_indices_test.pyre-establishes the int32 parity guarantees against whatever wheel is installed; on a wheel without the patches every test skips (with a self-explaining skip reason) so existing CI stays green.One pre-existing bug fixed along the way, because it undermines the gate:
post_install.pyreturnedinstall_glt.sh's exit status frommain()but never passed it tosys.exit, so when invoked as a file -- which is how image builds run it -- a failed GLT build/install/verification still exited 0 and produced a successful Docker layer. The__main__guard is nowsys.exit(main()), and aNonestatus (the runner's failure-to-execute path) maps to 1.These patches carry no GiGL-side consumers yet: nothing in
gigl/builds an int32 topology today. Landing the wheel change first keeps the follow-up PR (int32 edge-index plumbing) pure-Python and separately revertable.Where is the documentation for this feature?: rationale comments in
install_glt.sh, the patch files themselves, theverify_glt_patches.pymodule docstring, plus the CHANGELOG entryDid you add automated tests or write a test plan?
Yes, and the change was exercised end-to-end locally: the three patches apply cleanly in sequence to a fresh clone of the pinned GLT commit; a
WITH_CUDA=ONwheel built from that tree passes all 12 checks inverify_glt_patches.py(including the pinned-memory teardown path, on a T4 host); the newglt_int32_indices_test.pyruns 11/11 green against the patched wheel and skips 11/11 against an unpatched one; the fulldistributed_neighborloader_test.pysuite passes against the patched wheel, confirming the default int64 path is unchanged. The newpost_install_test.pyruns the real file in a subprocess with a shimmedbashand asserts the child's exit code (7 and 0) becomes the wrapper's.Updated Changelog.md? YES
Ready for code review?: YES