Apply and verify local patches on the pinned GraphLearn-Torch build

CHANGELOG.md

@@ -6,6 +6,19 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [Unreleased]
 
+### Added
+
+- Apply local patches to the pinned GraphLearn-Torch build in `install_glt.sh` (`gigl/scripts/patches/`): an exact
+  bitmap distinct-count in CPU graph init replacing `at::_unique`'s ~3x transient allocation, int32 CSR column-id
+  support in the CPU samplers, and shared-memory queue unpin/cleanup on teardown. `verify_glt_patches.py` gates the
+  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
+
 ### Removed
 
 - Remove the deprecated `RESOURCE_CONFIG_PATH` environment variable; use `GIGL_RESOURCE_CONFIG_URI` instead by

gigl/scripts/install_glt.sh

@@ -44,11 +44,25 @@ then
     # * https://github.com/alibaba/graphlearn-for-pytorch/pull/153
     # * https://github.com/alibaba/graphlearn-for-pytorch/pull/151
     # Thus, checking out a specific commit instead of a tagged version.
+    # Resolve this script's directory BEFORE cd'ing, so the patch path survives the cd below.
+    GIGL_SCRIPTS_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
     git clone https://github.com/alibaba/graphlearn-for-pytorch.git \
         && cd graphlearn-for-pytorch \
         && git checkout 88ff111ac0d9e45c6c9d2d18cfc5883dca07e9f9 \
         && git submodule update --init \
         && bash install_dependencies.sh
+    # Local patches applied on top of the pinned commit, in order. 0001 replaces the
+    # at::_unique distinct-count in InitCPUGraphFromCSR with an exact bitmap count: _unique's
+    # sort allocates ~3x the size of `indices` transiently (measured 3.00x at 200M edges), which
+    # is ~94 GiB at a 4.2B-edge partition and OOM-kills the graph init on hosts whose budget
+    # assumes init holds no copy of the topology. `patch` rather than `git apply` so failure
+    # modes are the familiar reject files; `--forward` + explicit failure so a re-run or an
+    # upstream change that absorbs the fix stops the build LOUDLY instead of shipping an
+    # unpatched wheel that dies 3 hours into a 16-GPU job.
+    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
     if has_cuda_driver;
     then
         echo "Will use CUDA for GLT..."
@@ -72,6 +86,13 @@ then
     uv pip install dist/*.whl \
         && cd .. \
         && rm -rf graphlearn-for-pytorch
+    # Applying a patch and SHIPPING it are different guarantees: the loop above proves the source
+    # tree changed, this proves the installed .so behaves. A stale build dir or a second wheel on
+    # the path would otherwise produce an image that silently OOMs (0001) or rejects the int32
+    # topology the trainer builds (0002), hours into a multi-GPU job.
+    echo "Verifying the GLT patches took effect in the installed wheel"
+    python "${GIGL_SCRIPTS_DIR}/verify_glt_patches.py" \
+        || { echo "FATAL: GLT patches did not take effect in the installed wheel"; exit 1; }
 else
     echo "Skipping install of GraphLearn-Torch on Mac"
 fi

gigl/scripts/patches/0001-glt-cpu-graph-col-count-bitmap.patch

@@ -0,0 +1,74 @@
+--- a/graphlearn_torch/csrc/cpu/graph.cc
++++ b/graphlearn_torch/csrc/cpu/graph.cc
+@@ -16,6 +16,10 @@
+ #include "graphlearn_torch/include/graph.h"
+ #include "graphlearn_torch/include/common.h"
+
++#include <algorithm>
++#include <cstdint>
++#include <vector>
++
+ namespace graphlearn_torch {
+
+ void Graph::InitCPUGraphFromCSR(
+@@ -30,7 +34,59 @@
+   col_idx_ = indices.data_ptr<int64_t>();
+   row_count_ = indptr.size(0) - 1;
+   edge_count_ = indices.size(0);
+-  col_count_ = std::get<0>(at::_unique(indices)).size(0);
++  // 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;
++  } else {
++    const uint64_t bitmap_bytes = (static_cast<uint64_t>(max_col) >> 3) + 1;
++    if (bitmap_bytes >
++        static_cast<uint64_t>(edge_count_) * static_cast<uint64_t>(sizeof(int64_t))) {
++      // Domain so sparse that the bitmap would exceed what `indices` itself occupies; the sort
++      // is the cheaper peak here.
++      col_count_ = std::get<0>(at::_unique(indices)).size(0);
++    } else {
++      std::vector<uint8_t> seen(bitmap_bytes, 0);
++      int64_t distinct_cols = 0;
++      for (int64_t i = 0; i < edge_count_; ++i) {
++        const uint64_t col = static_cast<uint64_t>(col_idx_[i]);
++        uint8_t& byte = seen[col >> 3];
++        const uint8_t bit = static_cast<uint8_t>(1u << (col & 7));
++        if (!(byte & bit)) {
++          byte |= bit;
++          ++distinct_cols;
++        }
++      }
++      col_count_ = distinct_cols;
++    }
++  }
+
+   if (edge_ids.numel()) {
+     CheckEq<int64_t>(edge_ids.dim(), 1);