diff --git a/CHANGELOG.md b/CHANGELOG.md index 50cc83cee..c363bd9de 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/gigl/scripts/install_glt.sh b/gigl/scripts/install_glt.sh index f044c6d0c..dee3212bb 100755 --- a/gigl/scripts/install_glt.sh +++ b/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 diff --git a/gigl/scripts/patches/0001-glt-cpu-graph-col-count-bitmap.patch b/gigl/scripts/patches/0001-glt-cpu-graph-col-count-bitmap.patch new file mode 100644 index 000000000..3b09396c3 --- /dev/null +++ b/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 ++#include ++#include ++ + namespace graphlearn_torch { + + void Graph::InitCPUGraphFromCSR( +@@ -30,7 +34,59 @@ + col_idx_ = indices.data_ptr(); + 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(max_col) >> 3) + 1; ++ if (bitmap_bytes > ++ static_cast(edge_count_) * static_cast(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 seen(bitmap_bytes, 0); ++ int64_t distinct_cols = 0; ++ for (int64_t i = 0; i < edge_count_; ++i) { ++ const uint64_t col = static_cast(col_idx_[i]); ++ uint8_t& byte = seen[col >> 3]; ++ const uint8_t bit = static_cast(1u << (col & 7)); ++ if (!(byte & bit)) { ++ byte |= bit; ++ ++distinct_cols; ++ } ++ } ++ col_count_ = distinct_cols; ++ } ++ } + + if (edge_ids.numel()) { + CheckEq(edge_ids.dim(), 1); diff --git a/gigl/scripts/patches/0002-glt-int32-csr-indices.patch b/gigl/scripts/patches/0002-glt-int32-csr-indices.patch new file mode 100644 index 000000000..1d6e47d9f --- /dev/null +++ b/gigl/scripts/patches/0002-glt-int32-csr-indices.patch @@ -0,0 +1,491 @@ +--- a/graphlearn_torch/include/graph.h ++++ b/graphlearn_torch/include/graph.h +@@ -79,6 +79,18 @@ + return col_idx_; + } + ++ // int32 column-id support, CPU sampling only. Column ids are node ids; when they all fit in ++ // int32 the CSC indices array halves (node id domain < 2^31, checked at init). Exactly one of ++ // col_idx_ / col_idx32_ is set. Samplers that have not been taught the int32 layout must call ++ // ColIsInt32() and reject it LOUDLY rather than read reinterpreted garbage. ++ const int32_t* GetColIdx32() const { ++ return col_idx32_; ++ } ++ ++ bool ColIsInt32() const { ++ return col_is_int32_; ++ } ++ + const int64_t* GetEdgeId() const { + return edge_id_; + } +@@ -122,6 +134,8 @@ + private: + int64_t* row_ptr_; + int64_t* col_idx_; ++ int32_t* col_idx32_{nullptr}; ++ bool col_is_int32_{false}; + int64_t* edge_id_; + float* edge_weight_; + std::vector registered_ptrs_; +--- a/graphlearn_torch/csrc/cpu/graph.cc ++++ b/graphlearn_torch/csrc/cpu/graph.cc +@@ -22,6 +22,76 @@ + + namespace graphlearn_torch { + ++namespace { ++ ++// 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 unsigned value 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: the column array is read as flat storage here and by every sampler, so ++// requiring contiguity 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. ++template ++int64_t CountDistinctColsImpl(const torch::Tensor& indices, ++ const ColT* col, ++ int64_t edge_count) { ++ int64_t max_col = -1; ++ for (int64_t i = 0; i < edge_count; ++i) { ++ const int64_t c = static_cast(col[i]); ++ TORCH_CHECK(c >= 0, ++ "InitCPUGraphFromCSR requires non-negative column ids (they are node ids), " ++ "got ", c, " at position ", i); ++ max_col = std::max(max_col, c); ++ } ++ if (max_col < 0) { ++ // No edges. at::_unique of an empty tensor is also empty. ++ return 0; ++ } ++ const uint64_t bitmap_bytes = (static_cast(max_col) >> 3) + 1; ++ if (bitmap_bytes > ++ static_cast(edge_count) * static_cast(sizeof(ColT))) { ++ // Domain so sparse that the bitmap would exceed what `indices` itself occupies; the sort ++ // is the cheaper peak here. ++ return std::get<0>(at::_unique(indices)).size(0); ++ } ++ std::vector seen(bitmap_bytes, 0); ++ int64_t distinct_cols = 0; ++ for (int64_t i = 0; i < edge_count; ++i) { ++ const uint64_t c = static_cast(col[i]); ++ uint8_t& byte = seen[c >> 3]; ++ const uint8_t bit = static_cast(1u << (c & 7)); ++ if (!(byte & bit)) { ++ byte |= bit; ++ ++distinct_cols; ++ } ++ } ++ return distinct_cols; ++} ++ ++int64_t CountDistinctCols(const torch::Tensor& indices, int64_t edge_count) { ++ TORCH_CHECK(indices.is_contiguous(), ++ "InitCPUGraphFromCSR requires contiguous indices: the column array is read as " ++ "flat storage by the samplers, so a strided view would sample the wrong columns"); ++ if (indices.scalar_type() == torch::kInt32) { ++ return CountDistinctColsImpl( ++ indices, indices.data_ptr(), edge_count); ++ } ++ return CountDistinctColsImpl( ++ indices, indices.data_ptr(), edge_count); ++} ++ ++} // namespace ++ + void Graph::InitCPUGraphFromCSR( + const torch::Tensor& indptr, + const torch::Tensor& indices, +@@ -31,62 +101,23 @@ + CheckEq(indices.dim(), 1); + + row_ptr_ = indptr.data_ptr(); +- col_idx_ = indices.data_ptr(); + row_count_ = indptr.size(0) - 1; + edge_count_ = 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; ++ // Column ids may be int32 (node id domain < 2^31 -- checked below) or int64. Exactly one of ++ // col_idx_ / col_idx32_ is set; samplers dispatch on ColIsInt32(), and samplers that have not ++ // been taught the int32 layout reject it loudly at their own entry points. ++ col_is_int32_ = (indices.scalar_type() == torch::kInt32); ++ if (col_is_int32_) { ++ col_idx32_ = indices.data_ptr(); ++ col_idx_ = nullptr; + } else { +- const uint64_t bitmap_bytes = (static_cast(max_col) >> 3) + 1; +- if (bitmap_bytes > +- static_cast(edge_count_) * static_cast(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 seen(bitmap_bytes, 0); +- int64_t distinct_cols = 0; +- for (int64_t i = 0; i < edge_count_; ++i) { +- const uint64_t col = static_cast(col_idx_[i]); +- uint8_t& byte = seen[col >> 3]; +- const uint8_t bit = static_cast(1u << (col & 7)); +- if (!(byte & bit)) { +- byte |= bit; +- ++distinct_cols; +- } +- } +- col_count_ = distinct_cols; +- } ++ TORCH_CHECK(indices.scalar_type() == torch::kInt64, ++ "InitCPUGraphFromCSR requires int32 or int64 indices, got ", ++ indices.scalar_type()); ++ col_idx_ = indices.data_ptr(); ++ col_idx32_ = nullptr; + } ++ col_count_ = CountDistinctCols(indices, edge_count_); + + if (edge_ids.numel()) { + CheckEq(edge_ids.dim(), 1); +--- a/graphlearn_torch/csrc/cpu/random_sampler.h ++++ b/graphlearn_torch/csrc/cpu/random_sampler.h +@@ -37,21 +37,9 @@ + const int32_t req_num, const int64_t row_count, + const int64_t* row_ptr, int64_t* out_nbr_num); + +- void CSRRowWiseSample(const int64_t* nodes, const int64_t* nbrs_offset, +- const int32_t bs, const int32_t req_num, const int64_t row_count, +- const int64_t* row_ptr, const int64_t* col_idx, int64_t* out_nbrs); +- +- void CSRRowWiseSample(const int64_t* nodes, const int64_t* nbrs_offset, +- const int32_t bs, const int32_t req_num, const int64_t row_count, +- const int64_t* row_ptr, const int64_t* col_idx, const int64_t* edge_id, +- int64_t* out_nbrs, int64_t* out_eid); +- +- void UniformSample(const int64_t* col_begin, const int64_t* col_end, +- const int32_t req_num, int64_t* out_nbrs); +- +- void UniformSample(const int64_t* col_begin, const int64_t* col_end, +- const int64_t* eid_begin, const int64_t* eid_end, +- const int32_t req_num, int64_t* out_nbrs, int64_t* out_eid); ++ // The row-wise sampling helpers are file-local templates in random_sampler.cc, generic over ++ // the column-id storage type (int64 or int32; see Graph::ColIsInt32). Outputs are always ++ // int64 -- the storage dtype is an implementation detail of the graph, not of the API. + + }; + +--- a/graphlearn_torch/csrc/cpu/random_sampler.cc ++++ b/graphlearn_torch/csrc/cpu/random_sampler.cc +@@ -21,6 +21,100 @@ + + namespace graphlearn_torch { + ++namespace { ++ ++// Generic over the column-id storage type ColT (int32 or int64). Neighbor OUTPUTS are always ++// int64: the narrow dtype is a property of the graph's storage, and it must not leak into the ++// sampled message format, which every consumer reads as int64. ++template ++void UniformSampleImpl(const ColT* col_begin, ++ const ColT* col_end, ++ const int32_t req_num, ++ int64_t* out_nbrs) { ++ // with replacement ++ const auto cap = col_end - col_begin; ++ if (req_num < cap) { ++ uint32_t seed = RandomSeedManager::getInstance().getSeed(); ++ thread_local static std::mt19937 engine(seed); ++ std::uniform_int_distribution<> dist(0, cap-1); ++ for (int32_t i = 0; i < req_num; ++i) { ++ out_nbrs[i] = static_cast(col_begin[dist(engine)]); ++ } ++ } else { ++ std::copy(col_begin, col_end, out_nbrs); ++ } ++} ++ ++template ++void UniformSampleImpl(const ColT* col_begin, ++ const ColT* col_end, ++ const int64_t* eid_begin, ++ const int64_t* eid_end, ++ const int32_t req_num, ++ int64_t* out_nbrs, ++ int64_t* out_eid) { ++ // with replacement ++ const auto cap = col_end - col_begin; ++ if (req_num < cap) { ++ uint32_t seed = RandomSeedManager::getInstance().getSeed(); ++ thread_local static std::mt19937 engine(seed); ++ std::uniform_int_distribution<> dist(0, cap-1); ++ for (int32_t i = 0; i < req_num; ++i) { ++ auto idx = dist(engine); ++ out_nbrs[i] = static_cast(col_begin[idx]); ++ out_eid[i] = eid_begin[idx]; ++ } ++ } else { ++ std::copy(col_begin, col_end, out_nbrs); ++ std::copy(eid_begin, eid_end, out_eid); ++ } ++} ++ ++template ++void CSRRowWiseSampleImpl(const int64_t* nodes, ++ const int64_t* nbrs_offset, ++ const int32_t bs, ++ const int32_t req_num, ++ const int64_t row_count, ++ const int64_t* row_ptr, ++ const ColT* col_idx, ++ int64_t* out_nbrs) { ++ at::parallel_for(0, bs, 1, [&](int32_t start, int32_t end){ ++ for(int32_t i = start; i < end; ++i) { ++ auto v = nodes[i]; ++ if (v < row_count) { ++ UniformSampleImpl(col_idx + row_ptr[v], col_idx + row_ptr[v+1], req_num, ++ out_nbrs + nbrs_offset[i]); ++ } ++ } ++ }); ++} ++ ++template ++void CSRRowWiseSampleImpl(const int64_t* nodes, ++ const int64_t* nbrs_offset, ++ const int32_t bs, ++ const int32_t req_num, ++ const int64_t row_count, ++ const int64_t* row_ptr, ++ const ColT* col_idx, ++ const int64_t* edge_ids, ++ int64_t* out_nbrs, ++ int64_t* out_eid) { ++ at::parallel_for(0, bs, 1, [&](int32_t start, int32_t end){ ++ for(int32_t i = start; i < end; ++i) { ++ auto v = nodes[i]; ++ if (v < row_count) { ++ UniformSampleImpl(col_idx + row_ptr[v], col_idx + row_ptr[v+1], ++ edge_ids + row_ptr[v], edge_ids + row_ptr[v+1], ++ req_num, out_nbrs + nbrs_offset[i], out_eid + nbrs_offset[i]); ++ } ++ } ++ }); ++} ++ ++} // namespace ++ + std::tuple + CPURandomSampler::Sample(const torch::Tensor& nodes, int32_t req_num) { + if (req_num < 0) req_num = std::numeric_limits::max(); +@@ -40,8 +134,13 @@ + } + + torch::Tensor nbrs = torch::empty(nbrs_offset[bs], nodes.options()); +- CSRRowWiseSample(nodes_ptr, nbrs_offset, bs, req_num, row_count, +- row_ptr, col_idx, nbrs.data_ptr()); ++ if (graph_->ColIsInt32()) { ++ CSRRowWiseSampleImpl(nodes_ptr, nbrs_offset, bs, req_num, row_count, ++ row_ptr, graph_->GetColIdx32(), nbrs.data_ptr()); ++ } else { ++ CSRRowWiseSampleImpl(nodes_ptr, nbrs_offset, bs, req_num, row_count, ++ row_ptr, col_idx, nbrs.data_ptr()); ++ } + return std::make_tuple(nbrs, nbrs_num); + } + +@@ -66,9 +165,15 @@ + + torch::Tensor nbrs = torch::empty(nbrs_offset[bs], nodes.options()); + torch::Tensor out_eid = torch::empty(nbrs_offset[bs], nodes.options()); +- CSRRowWiseSample(nodes_ptr, nbrs_offset, bs, req_num, row_count, +- row_ptr, col_idx, edge_ids, +- nbrs.data_ptr(), out_eid.data_ptr()); ++ if (graph_->ColIsInt32()) { ++ CSRRowWiseSampleImpl(nodes_ptr, nbrs_offset, bs, req_num, row_count, ++ row_ptr, graph_->GetColIdx32(), edge_ids, ++ nbrs.data_ptr(), out_eid.data_ptr()); ++ } else { ++ CSRRowWiseSampleImpl(nodes_ptr, nbrs_offset, bs, req_num, row_count, ++ row_ptr, col_idx, edge_ids, ++ nbrs.data_ptr(), out_eid.data_ptr()); ++ } + return std::make_tuple(nbrs, nbrs_num, out_eid); + } + +@@ -91,89 +196,4 @@ + }); + } + +-void CPURandomSampler::CSRRowWiseSample( +- const int64_t* nodes, +- const int64_t* nbrs_offset, +- const int32_t bs, +- const int32_t req_num, +- const int64_t row_count, +- const int64_t* row_ptr, +- const int64_t* col_idx, +- int64_t* out_nbrs) { +- at::parallel_for(0, bs, 1, [&](int32_t start, int32_t end){ +- for(int32_t i = start; i < end; ++i) { +- auto v = nodes[i]; +- if (v < row_count) { +- UniformSample(col_idx + row_ptr[v], col_idx + row_ptr[v+1], req_num, +- out_nbrs + nbrs_offset[i]); +- } +- } +- }); +-} +- +-void CPURandomSampler::CSRRowWiseSample( +- const int64_t* nodes, +- const int64_t* nbrs_offset, +- const int32_t bs, +- const int32_t req_num, +- const int64_t row_count, +- const int64_t* row_ptr, +- const int64_t* col_idx, +- const int64_t* edge_ids, +- int64_t* out_nbrs, +- int64_t* out_eid) { +- at::parallel_for(0, bs, 1, [&](int32_t start, int32_t end){ +- for(int32_t i = start; i < end; ++i) { +- auto v = nodes[i]; +- if (v < row_count) { +- UniformSample(col_idx + row_ptr[v], col_idx + row_ptr[v+1], +- edge_ids + row_ptr[v], edge_ids + row_ptr[v+1], +- req_num, out_nbrs + nbrs_offset[i], out_eid + nbrs_offset[i]); +- } +- } +- }); +-} +- +-void CPURandomSampler::UniformSample(const int64_t* col_begin, +- const int64_t* col_end, +- const int32_t req_num, +- int64_t* out_nbrs) { +- // with replacement +- const auto cap = col_end - col_begin; +- if (req_num < cap) { +- uint32_t seed = RandomSeedManager::getInstance().getSeed(); +- thread_local static std::mt19937 engine(seed); +- std::uniform_int_distribution<> dist(0, cap-1); +- for (int32_t i = 0; i < req_num; ++i) { +- out_nbrs[i] = col_begin[dist(engine)]; +- } +- } else { +- std::copy(col_begin, col_end, out_nbrs); +- } +-} +- +-void CPURandomSampler::UniformSample(const int64_t* col_begin, +- const int64_t* col_end, +- const int64_t* eid_begin, +- const int64_t* eid_end, +- const int32_t req_num, +- int64_t* out_nbrs, +- int64_t* out_eid) { +- // with replacement +- const auto cap = col_end - col_begin; +- if (req_num < cap) { +- uint32_t seed = RandomSeedManager::getInstance().getSeed(); +- thread_local static std::mt19937 engine(seed); +- std::uniform_int_distribution<> dist(0, cap-1); +- for (int32_t i = 0; i < req_num; ++i) { +- auto idx = dist(engine); +- out_nbrs[i] = col_begin[idx]; +- out_eid[i] = eid_begin[idx]; +- } +- } else { +- std::copy(col_begin, col_end, out_nbrs); +- std::copy(eid_begin, eid_end, out_eid); +- } +-} +- + } // namespace graphlearn_torch +\ No newline at end of file +--- a/graphlearn_torch/csrc/cpu/weighted_sampler.cc ++++ b/graphlearn_torch/csrc/cpu/weighted_sampler.cc +@@ -28,6 +28,10 @@ + const int64_t* nodes_ptr = nodes.data_ptr(); + int64_t bs = nodes.size(0); + const auto row_ptr = graph_->GetRowPtr(); ++ TORCH_CHECK(!graph_->ColIsInt32(), ++ "CPUWeightedSampler has not been taught the int32 column-id layout; " ++ "initialize this graph with int64 indices or extend the sampler " ++ "(see Graph::ColIsInt32)"); + const auto col_idx = graph_->GetColIdx(); + const auto row_count = graph_->GetRowCount(); + const auto edge_weights = graph_->GetEdgeWeight(); +@@ -55,6 +59,10 @@ + const int64_t* nodes_ptr = nodes.data_ptr(); + int64_t bs = nodes.size(0); + const auto row_ptr = graph_->GetRowPtr(); ++ TORCH_CHECK(!graph_->ColIsInt32(), ++ "CPUWeightedSampler has not been taught the int32 column-id layout; " ++ "initialize this graph with int64 indices or extend the sampler " ++ "(see Graph::ColIsInt32)"); + const auto col_idx = graph_->GetColIdx(); + const auto edge_ids = graph_->GetEdgeId(); + const auto row_count = graph_->GetRowCount(); +--- a/graphlearn_torch/csrc/cpu/random_negative_sampler.cc ++++ b/graphlearn_torch/csrc/cpu/random_negative_sampler.cc +@@ -24,6 +24,10 @@ + std::tuple CPURandomNegativeSampler::Sample( + int32_t req_num, int32_t trials_num, bool padding) { + const int64_t* row_ptr = graph_->GetRowPtr(); ++ TORCH_CHECK(!graph_->ColIsInt32(), ++ "CPURandomNegativeSampler has not been taught the int32 column-id layout; " ++ "initialize this graph with int64 indices or extend the sampler " ++ "(see Graph::ColIsInt32)"); + const int64_t* col_idx = graph_->GetColIdx(); + int64_t row_num = graph_->GetRowCount(); + int64_t col_num = graph_->GetColCount(); +--- a/graphlearn_torch/csrc/cpu/subgraph_op.cc ++++ b/graphlearn_torch/csrc/cpu/subgraph_op.cc +@@ -64,6 +64,10 @@ + std::vector& out_cols, + std::vector& out_eids) { + const auto indptr = graph_->GetRowPtr(); ++ TORCH_CHECK(!graph_->ColIsInt32(), ++ "CPUSubGraphOp has not been taught the int32 column-id layout; " ++ "initialize this graph with int64 indices or extend the sampler " ++ "(see Graph::ColIsInt32)"); + const auto indices = graph_->GetColIdx(); + const auto edge_ids = graph_->GetEdgeId(); + const auto row_count = graph_->GetRowCount(); diff --git a/gigl/scripts/patches/0003-glt-unpin-shm-queue-on-teardown.patch b/gigl/scripts/patches/0003-glt-unpin-shm-queue-on-teardown.patch new file mode 100644 index 000000000..9e550bc19 --- /dev/null +++ b/gigl/scripts/patches/0003-glt-unpin-shm-queue-on-teardown.patch @@ -0,0 +1,131 @@ +diff --git a/graphlearn_torch/csrc/shm_queue.cc b/graphlearn_torch/csrc/shm_queue.cc +index 11302e7..b7d3c75 100644 +--- a/graphlearn_torch/csrc/shm_queue.cc ++++ b/graphlearn_torch/csrc/shm_queue.cc +@@ -18,6 +18,7 @@ limitations under the License. + #include + + #include ++#include + #include + #include + +@@ -241,19 +242,67 @@ bool ShmQueue::Empty() { + + void ShmQueue::PinMemory() { + #ifdef WITH_CUDA ++ auto* deleter = std::get_deleter(meta_); ++ if (deleter == nullptr || deleter->pinned) { ++ // Already pinned by this process: cudaHostRegister would reject the ++ // re-registration of the same range, so a second call is a no-op rather ++ // than an error. A null deleter cannot happen with the constructors ++ // above; refusing to register in that case keeps the invariant that ++ // every registration has a recorded owner to unregister it. ++ return; ++ } + cudaHostRegister(meta_.get(), shm_size_, cudaHostRegisterMapped); + CUDACheckError(); ++ deleter->pinned = true; + #endif + } + + void ShmQueue::ShmQueueMetaDeleter::operator()(ShmQueueMeta* meta_ptr) { ++ bool detach = true; + if (meta_ptr) { +- if (shmid > 0) { ++#ifdef WITH_CUDA ++ if (pinned) { ++ // cudaHostRegister (PinMemory) is per-process driver state keyed on ++ // this mapping's address range, and the shmdt below does NOT clear it. ++ // A dangling registration poisons whatever mapping later reuses these ++ // addresses: the next overlapping registration or mapped allocation ++ // fails with "resource already mapped", asynchronously and far from ++ // here (observed as torch.AcceleratorError in the first backward pass ++ // after tearing down a pinned loader channel). ++ // ++ // Destructor context, so nothing here may throw or exit. The sticky ++ // per-thread error is cleared BEFORE the call so an unrelated earlier ++ // failure cannot masquerade as ours, and the DIRECT status decides: ++ // NotRegistered (someone beat us to it) and CudartUnloading (process ++ // exit, context already gone) proceed to detach; any other failure ++ // keeps the mapping attached, because a live registration on a live ++ // mapping is harmless while a stale one on a reused range is the ++ // exact bug this patch removes. The kernel detaches at process exit ++ // and IPC_RMID below still marks the segment for destruction, so ++ // nothing outlives the process either way. ++ (void)cudaGetLastError(); ++ cudaError_t unregister_status = cudaHostUnregister(meta_ptr); ++ if (unregister_status != cudaSuccess && ++ unregister_status != cudaErrorHostMemoryNotRegistered && ++ unregister_status != cudaErrorCudartUnloading) { ++ fprintf(stderr, ++ "ShmQueue teardown: cudaHostUnregister failed (%s); keeping " ++ "the mapping attached rather than leaving a stale pinned " ++ "registration on a reusable address range.\n", ++ cudaGetErrorString(unregister_status)); ++ detach = false; ++ } ++ (void)cudaGetLastError(); ++ } ++#endif ++ if (shmid >= 0) { + meta_ptr->Finalize(); + } +- Check(shmdt(meta_ptr) != -1, "shmdt failed!"); ++ if (detach) { ++ Check(shmdt(meta_ptr) != -1, "shmdt failed!"); ++ } + } +- if (shmid > 0) { ++ if (shmid >= 0) { + Check(shmctl(shmid, IPC_RMID, 0) != -1, "shmctl(IPC_RMID) failed!"); + } + } +diff --git a/graphlearn_torch/include/shm_queue.h b/graphlearn_torch/include/shm_queue.h +index 5eefe97..bd00916 100644 +--- a/graphlearn_torch/include/shm_queue.h ++++ b/graphlearn_torch/include/shm_queue.h +@@ -227,12 +227,21 @@ private: + /// Deleter to release shared memory when the underlying meta is shared + /// by multiple shared_ptrs. + struct ShmQueueMetaDeleter { +- /// If "shmid > 0", which means current instance of `ShmQueue` is the +- /// creator of the underlying shm, thus, `shmctl` should be called after +- /// detaching the shm in this deleter. Otherwise, current instance is not +- /// the creator and only need to detach the shm with `shmdt`. ++ /// If "shmid >= 0", current instance of `ShmQueue` is the creator of the ++ /// underlying shm, thus, `shmctl` should be called after detaching the ++ /// shm in this deleter. Otherwise (-1, the attach constructor's marker), ++ /// current instance is not the creator and only needs to detach with ++ /// `shmdt`. The creator test MUST include zero: shmget returns ids from ++ /// 0 in a fresh IPC namespace, so the first segment every containerized ++ /// process creates has shmid 0 -- under the old `> 0` test its creator ++ /// never removed it (no IPC_RMID, no Finalize). + int shmid; +- explicit ShmQueueMetaDeleter(int shmid) : shmid(shmid) {} ++ /// Whether this process pinned this mapping via `PinMemory`. The deleter ++ /// is the one place guaranteed to run exactly when the LAST local user of ++ /// the mapping goes away (a queue and its in-flight `ShmData` messages ++ /// share one meta_), so it owns the matching `cudaHostUnregister`. ++ bool pinned; ++ explicit ShmQueueMetaDeleter(int shmid) : shmid(shmid), pinned(false) {} + void operator()(ShmQueueMeta* meta_ptr); + }; + std::shared_ptr meta_; +diff --git a/graphlearn_torch/python/py_export_glt.cc b/graphlearn_torch/python/py_export_glt.cc +index e70ab36..12c013e 100644 +--- a/graphlearn_torch/python/py_export_glt.cc ++++ b/graphlearn_torch/python/py_export_glt.cc +@@ -129,6 +129,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + py::class_(m, "SampleQueue") + .def(py::init(), py::arg("capacity"), py::arg("buf_size")) + .def("pin_memory", &SampleQueue::PinMemory) ++ // Compiled capability marker for the teardown-unpin patch. Runtime behaviour cannot be ++ // probed without a CUDA device (image builds have none), and a wheel built from an ++ // UNPATCHED tree is otherwise indistinguishable at build time: every observable CPU-side ++ // behaviour matches. The marker only exists when this file was compiled from the patched ++ // tree, which is the same .so that carries the patched deleter. ++ .def_static("supports_unpin_on_teardown", []() { return true; }) + .def("empty", &SampleQueue::Empty, + py::call_guard()) + .def("send", &SampleQueue::Enqueue, py::arg("msg"), diff --git a/gigl/scripts/post_install.py b/gigl/scripts/post_install.py index cd13c4b0b..262661d49 100644 --- a/gigl/scripts/post_install.py +++ b/gigl/scripts/post_install.py @@ -54,7 +54,10 @@ def main(): print(f"Executing {cmd}...") result = run_command_and_stream_stdout(cmd) print("Post-install script finished running, with return code: ", result) - return result + # The status must live in this process's exit code, not just the log line above: image + # builds invoke this file directly, and a swallowed failure here ships an image whose + # GLT wheel never built, installed, or verified + return 1 if result is None else result except subprocess.CalledProcessError as e: print(f"Error running install_glt.sh: {e}") @@ -65,4 +68,4 @@ def main(): if __name__ == "__main__": - main() + sys.exit(main()) diff --git a/gigl/scripts/verify_glt_patches.py b/gigl/scripts/verify_glt_patches.py new file mode 100644 index 000000000..e046fc914 --- /dev/null +++ b/gigl/scripts/verify_glt_patches.py @@ -0,0 +1,285 @@ +"""Assert that the patches in ``gigl/scripts/patches/`` are live in the INSTALLED graphlearn_torch. + +Run by ``install_glt.sh`` immediately after ``uv pip install dist/*.whl``, and exits non-zero if +any patched behaviour is missing. + +WHY THIS EXISTS + ``install_glt.sh`` already fails loudly if a patch file does not APPLY. That is a weaker + guarantee than it looks: applying is a property of the source tree, while what ships is a + compiled ``.so``. A patch can apply to a file the build then excludes, a stale build directory + can be reused, or a wheel can be installed from somewhere other than the tree that was + patched -- and every one of those failures is silent, because the unpatched code paths work. + They just work at 3x the memory (patch 0001) or reject the int32 topology the trainer is + about to build (patch 0002), several hours into a 16-GPU job. + + The unit tests in ``tests/unit/utils/glt_int32_indices_test.py`` cover the same ground in far + more detail, but they SKIP on an unpatched wheel by design -- GiGL's CI runs against the + released graphlearn_torch, where failing would be wrong. So the test suite cannot be the gate + for the image, and this script cannot be the detailed test. Both exist, deliberately. + + Deliberately dependency-free (torch and graphlearn_torch only): it runs inside the image build + before the rest of the repo's test dependencies are necessarily importable. + +Usage: + python gigl/scripts/verify_glt_patches.py +""" + +import os +import sys + +import torch +from graphlearn_torch.data import Graph, Topology + +# 0001 replaces at::_unique with a bitmap count; a column id beyond the id domain the bitmap can +# describe must fall back rather than allocate, and a negative id must raise rather than write out +# of bounds. Both are patch-only behaviours: upstream accepts them silently. +_NEGATIVE_COLUMN_ID = -1 + + +def _cpu_graph(indptr: torch.Tensor, indices: torch.Tensor) -> Graph: + topology = Topology.__new__(Topology) + topology._layout = "CSR" + topology._indptr = indptr + topology._indices = indices + topology._edge_ids = None + topology._edge_weights = None + graph = Graph.__new__(Graph) + graph.topo = topology + graph.mode = "CPU" + graph.device = None + graph._graph = None + graph.lazy_init() + return graph + + +def _check(name: str, passed: bool, detail: str = "") -> bool: + print(f" [{'OK ' if passed else 'FAIL'}] {name}{': ' + detail if detail else ''}") + return passed + + +def verify_0001_bitmap_col_count() -> bool: + """The hardened distinct-count: exact on valid input, loud on invalid input.""" + indptr = torch.tensor([0, 2, 3, 5], dtype=torch.int64) + indices = torch.tensor([7, 3, 7, 1, 4], dtype=torch.int64) + graph = _cpu_graph(indptr, indices) + exact = _check( + "0001 col_count is exact", + graph.col_count == int(torch.unique(indices).numel()), + f"{graph.col_count} vs {int(torch.unique(indices).numel())}", + ) + + # Upstream's _unique accepts a negative id; the patched count must reject it, because casting + # it to an unsigned bitmap offset would write outside the allocation. + rejects_negative = False + try: + _cpu_graph( + torch.tensor([0, 1], dtype=torch.int64), + torch.tensor([_NEGATIVE_COLUMN_ID], dtype=torch.int64), + ) + except RuntimeError as error: + rejects_negative = "non-negative" in str(error) + rejected = _check( + "0001 rejects a negative column id", + rejects_negative, + "an unpatched wheel accepts this", + ) + + # Same reasoning for a strided view: every sampler reads the column array as flat storage. + strided = torch.arange(8, dtype=torch.int64)[::2] + rejects_strided = False + try: + _cpu_graph(torch.tensor([0, 2, 4], dtype=torch.int64), strided) + except RuntimeError as error: + rejects_strided = "contiguous" in str(error) + contiguous = _check( + "0001 rejects non-contiguous indices", + rejects_strided, + "an unpatched wheel accepts this", + ) + return exact and rejected and contiguous + + +def verify_0002_int32_indices() -> bool: + """int32 columns must be accepted AND sample identically to int64.""" + indptr = torch.tensor([0, 3, 3, 6], dtype=torch.int64) + columns = [5, 2, 9, 1, 7, 4] + indices64 = torch.tensor(columns, dtype=torch.int64) + indices32 = torch.tensor(columns, dtype=torch.int32) + + try: + graph32 = _cpu_graph(indptr, indices32) + except RuntimeError as error: + return _check("0002 accepts int32 indices", False, str(error).splitlines()[0]) + accepted = _check("0002 accepts int32 indices", True) + + graph64 = _cpu_graph(indptr, indices64) + counts_match = _check( + "0002 col_count matches the int64 graph", + graph32.col_count == graph64.col_count, + f"{graph32.col_count} vs {graph64.col_count}", + ) + + # Full fanout (req_num > max degree) makes UniformSample copy rather than draw, so the two + # graphs are comparable id-by-id. This is the property whose failure is SILENT. + from graphlearn_torch import py_graphlearn_torch as pywrap + + seeds = torch.tensor([0, 2], dtype=torch.int64) + neighbors64, degrees64 = pywrap.CPURandomSampler(graph64.graph_handler).sample( + seeds, 8 + ) + neighbors32, degrees32 = pywrap.CPURandomSampler(graph32.graph_handler).sample( + seeds, 8 + ) + identical = _check( + "0002 samples identically to the int64 graph", + bool(torch.equal(neighbors64, neighbors32)) + and bool(torch.equal(degrees64, degrees32)), + f"{neighbors32.tolist()} vs {neighbors64.tolist()}", + ) + int64_out = _check( + "0002 sampled ids stay int64", + neighbors32.dtype == torch.int64, + str(neighbors32.dtype), + ) + + # A sampler that was not taught the layout must raise, not read the nullptr col_idx_. + rejects = False + try: + pywrap.CPUWeightedSampler(graph32.graph_handler).sample(seeds, 2) + except RuntimeError as error: + rejects = "int32" in str(error) + guarded = _check( + "0002 untaught samplers reject int32 loudly", + rejects, + "a silent nullptr read would be the alternative", + ) + return accepted and counts_match and identical and int64_out and guarded + + +def _sysv_segments_created_by_this_process() -> int: + # /proc/sysvipc/shm is system-wide; filtering on the creator pid keeps the count immune to + # whatever else the build host is doing. + own_pid = str(os.getpid()) + count = 0 + with open("/proc/sysvipc/shm") as shm_table: + next(shm_table) + for line in shm_table: + fields = line.split() + if len(fields) > 4 and fields[4] == own_pid: + count += 1 + return count + + +def verify_0003_queue_teardown() -> bool: + """The teardown-unpin patch must be COMPILED IN, and teardown must stay leak-free. + + The runtime payload -- the deleter calling cudaHostUnregister for a mapping this process + pinned -- cannot be exercised without a CUDA device, and image builds have none. Worse, + ``pin_memory()`` on a driverless host does not raise: GLT's ``CUDACheckError`` calls + ``exit(EXIT_FAILURE)``, which would kill this verifier and the build with it. So this + check NEVER pins unless a device is actually available. Presence of the patched BINARY is + proven by a compiled capability marker instead (``supports_unpin_on_teardown``, added by + the same patch to the same .so as the deleter); the unregister BEHAVIOUR is proven by the + GPU canary probe (pin -> destroy -> repin cycles), which fails on an unpatched wheel. + + What runs everywhere: the deleter still detaches and removes the SysV segment across + repeated cycles, and the zero-copy ShmData path still holds the mapping open while a + dequeued message is in flight. + """ + import pickle + + from graphlearn_torch import py_graphlearn_torch as pywrap + + marker = getattr(pywrap.SampleQueue, "supports_unpin_on_teardown", None) + compiled_in = _check( + "0003 teardown-unpin capability is compiled in", + marker is not None and bool(marker()), + "an unpatched wheel lacks this binding", + ) + if not compiled_in: + return False + + cuda_usable = torch.cuda.is_available() + + def one_cycle() -> None: + queue = pywrap.SampleQueue(8, 1 << 20) + if cuda_usable: + queue.pin_memory() + queue.pin_memory() # second call must be a no-op, not a re-register error + queue.send({"ids": torch.arange(64, dtype=torch.int64)}) + consumer = pickle.loads(pickle.dumps(queue)) + message = consumer.receive(5000) + assert torch.equal(message["ids"], torch.arange(64, dtype=torch.int64)) + del message, consumer, queue + + # The leak check is only meaningful if the pid-filtered /proc view can SEE this + # process's segments at all (an unusual /proc mount from another pid namespace would + # count zero forever and report any leak as clean). Measure the DELTA the sentinel + # causes, not an absolute count -- a pre-existing segment attributed to the same + # numeric pid must not vouch for an invisible sentinel. + count_before_sentinel = _sysv_segments_created_by_this_process() + sentinel = pywrap.SampleQueue(8, 1 << 20) + count_with_sentinel = _sysv_segments_created_by_this_process() + del sentinel + count_after_sentinel = _sysv_segments_created_by_this_process() + observable = _check( + "0003 pid filter observes a live segment", + count_with_sentinel == count_before_sentinel + 1 + and count_after_sentinel == count_before_sentinel, + f"{count_before_sentinel} -> {count_with_sentinel} -> {count_after_sentinel}", + ) + if not observable: + return False + + one_cycle() # torch lazily allocates a process-level shm singleton on first send + baseline = _sysv_segments_created_by_this_process() + for _ in range(20): + one_cycle() + leaked = _sysv_segments_created_by_this_process() - baseline + no_leak = _check( + "0003 teardown cycles leak no shm segments", + leaked <= 0, + f"{leaked} leaked over 20 cycles" + if leaked > 0 + else f"20 cycles clean ({'pinned' if cuda_usable else 'unpinned; no CUDA device'})", + ) + + # Dequeue is zero-copy: the received tensors view the segment through ShmData. Destroying + # BOTH queue objects while the message is still held must leave the data readable -- this + # ordering is exactly why the unpin lives in the shared_ptr deleter, not in ~ShmQueue. + producer = pywrap.SampleQueue(8, 1 << 20) + producer.send({"ids": torch.arange(64, dtype=torch.int64)}) + consumer = pickle.loads(pickle.dumps(producer)) + in_flight = consumer.receive(5000) + del producer, consumer + survives = _check( + "0003 in-flight message survives queue teardown", + bool(torch.equal(in_flight["ids"], torch.arange(64, dtype=torch.int64))), + ) + del in_flight + return no_leak and survives + + +def main() -> int: + print(f"Verifying GLT patches against {Graph.__module__}") + results = { + "0001-glt-cpu-graph-col-count-bitmap": verify_0001_bitmap_col_count(), + "0002-glt-int32-csr-indices": verify_0002_int32_indices(), + "0003-glt-unpin-shm-queue-on-teardown": verify_0003_queue_teardown(), + } + failed = [name for name, passed in results.items() if not passed] + if failed: + print( + f"\nFATAL: {failed} did not take effect in the installed graphlearn_torch. The patch " + f"files applied to the source tree, so the wheel that got INSTALLED is not the one " + f"that was built from it -- check for a stale build directory or a second wheel on " + f"the path. Shipping this image would OOM (0001) or reject the int32 topology the " + f"trainer builds (0002), hours into a multi-GPU job." + ) + return 1 + print("\nAll GLT patches verified live in the installed graphlearn_torch.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/unit/scripts_test/post_install_test.py b/tests/unit/scripts_test/post_install_test.py new file mode 100644 index 000000000..00375a1c5 --- /dev/null +++ b/tests/unit/scripts_test/post_install_test.py @@ -0,0 +1,53 @@ +"""``post_install.py`` must propagate ``install_glt.sh``'s exit status as its own. + +Image builds invoke the file directly (``requirements/install_py_deps.sh`` runs +``uv run python .../post_install.py``), so a swallowed child failure produces a +SUCCESSFUL Docker layer whose GLT wheel never built, installed, or verified. That is +the exact silent failure the patch-verification gate in ``install_glt.sh`` exists to +prevent, which makes the wrapper's exit code part of the gate. + +These tests run the real file in a subprocess with ``bash`` shimmed ahead on ``PATH``, +so the entire path under test -- argument handling, ``__main__`` guard, exit-code +plumbing -- is the shipped one, with only the child's exit status under our control. +""" + +import os +import stat +import subprocess +import sys +import tempfile +from pathlib import Path + +from absl.testing import absltest + +from tests.test_assets.test_case import TestCase + +_POST_INSTALL = Path(__file__).parents[3] / "gigl" / "scripts" / "post_install.py" + + +class PostInstallExitCodeTest(TestCase): + def _run_with_shimmed_bash(self, bash_exit_code: int) -> int: + """Run the real post_install.py with a fake ``bash`` that exits as told.""" + with tempfile.TemporaryDirectory() as shim_dir: + shim = Path(shim_dir) / "bash" + shim.write_text(f"#!/bin/sh\nexit {bash_exit_code}\n") + shim.chmod(shim.stat().st_mode | stat.S_IXUSR) + environment = os.environ.copy() + environment["PATH"] = f"{shim_dir}:{environment['PATH']}" + completed = subprocess.run( + [sys.executable, str(_POST_INSTALL)], + env=environment, + capture_output=True, + text=True, + ) + return completed.returncode + + def test_a_failing_install_script_fails_the_wrapper_process(self) -> None: + self.assertEqual(self._run_with_shimmed_bash(7), 7) + + def test_a_succeeding_install_script_exits_zero(self) -> None: + self.assertEqual(self._run_with_shimmed_bash(0), 0) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/unit/utils/glt_int32_indices_test.py b/tests/unit/utils/glt_int32_indices_test.py new file mode 100644 index 000000000..9ef5f75b4 --- /dev/null +++ b/tests/unit/utils/glt_int32_indices_test.py @@ -0,0 +1,284 @@ +"""make unit_test_py PY_TEST_FILES="glt_int32_indices_test.py" + +GLT patch 0002 (``gigl/scripts/patches/0002-glt-int32-csr-indices.patch``): the COMPILED CPU +samplers must read int32 column ids and produce byte-identical results to the int64 graph. + +Why this file exists rather than a probe script: the code that can corrupt training silently +lives in a C++ extension that ``gigl/scripts/install_glt.sh`` rebuilds from source in every image +build. A wrong dispatch there does not raise -- it samples the wrong neighbors -- so the +guarantee has to be re-established against each newly built wheel, which means it belongs in the +suite and not in a scratch directory. + +These tests SKIP when the installed graphlearn_torch has no int32 support, so they pass on an +unpatched wheel instead of failing confusingly; ``test_the_patch_is_present_in_this_wheel`` +records which mode ran, so a silently-unpatched image shows up as a skip and not as green. + +The determinism trick that makes element-wise comparison possible: with ``req_num`` above the +maximum degree, ``UniformSample`` copies every neighbor instead of drawing, so full-fanout +sampling is exact and comparable id-by-id. Sub-degree draws are then checked as a distribution +(every drawn id must be a true neighbor, exact counts). +""" + +import unittest + +import torch +from graphlearn_torch import py_graphlearn_torch as pywrap +from graphlearn_torch.data import Graph, Topology + +from tests.test_assets.test_case import TestCase + + +def _build_cpu_graph(indptr: torch.Tensor, indices: torch.Tensor) -> Graph: + """A CPU ``Graph`` over a ready-made CSR, bypassing ``Topology.__init__``. + + Populating the attributes directly keeps the fixture free of the ``arange(num_edges)`` edge + ids ``Topology.__init__`` would attach, so the CSR reaches the compiled extension exactly as + given. + """ + topology = Topology.__new__(Topology) + topology._layout = "CSR" + topology._indptr = indptr + topology._indices = indices + topology._edge_ids = None + topology._edge_weights = None + graph = Graph.__new__(Graph) + graph.topo = topology + graph.mode = "CPU" + graph.device = None + graph._graph = None + graph.lazy_init() + return graph + + +def _random_csr( + num_rows: int, num_edges: int, seed: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """A valid CSR (indptr, int64 indices, degrees) with rows in ascending order.""" + generator = torch.Generator().manual_seed(seed) + row = torch.randint( + 0, num_rows, (num_edges,), generator=generator, dtype=torch.int64 + ) + col = torch.randint( + 0, num_rows, (num_edges,), generator=generator, dtype=torch.int64 + ) + order = torch.argsort(row) + row, col = row[order], col[order] + degrees = torch.bincount(row, minlength=num_rows) + indptr = torch.zeros(num_rows + 1, dtype=torch.int64) + torch.cumsum(degrees, dim=0, out=indptr[1:]) + return indptr, col.contiguous(), degrees + + +def _wheel_supports_int32_indices() -> bool: + """Whether the installed compiled extension accepts int32 CSR indices.""" + try: + _build_cpu_graph( + torch.tensor([0, 1], dtype=torch.int64), + torch.tensor([0], dtype=torch.int32), + ) + except RuntimeError: + return False + return True + + +_HAS_INT32 = _wheel_supports_int32_indices() +_SKIP_REASON = "installed graphlearn_torch has no int32 CSR support: patch 0002 is not in this wheel" + + +class Int32IndicesSupportTest(TestCase): + def test_which_mode_this_wheel_runs_in(self) -> None: + """Records whether the patch is present, so an all-skip run is self-explaining. + + This deliberately does NOT fail on an unpatched wheel. GiGL's own CI and a plain + ``make unit_test_py`` run against the released graphlearn_torch, which has no int32 + support and never will until this patch is upstreamed -- failing there would be a broken + build reporting a correct state. + + The consequence has to be stated plainly, because it is the "suite ran nothing" trap: + **on an unpatched wheel every test in this file skips, so green here proves nothing about + int32 sampling.** The signal to read is the skip COUNT. Inside an image whose wheel + ``install_glt.sh`` built, the parity tests must actually RUN; + ``gigl/scripts/verify_glt_patches.py`` is what checks that during the image build, and it + is not optional before shipping an int32 topology to a cluster. + """ + if not _HAS_INT32: + self.skipTest( + "patch 0002 absent: int32 CSR indices rejected by the compiled extension, so " + "every parity test in this file skipped. Expected on a wheel that predates the " + "patches; NOT expected inside an image whose wheel install_glt.sh built with " + "gigl/scripts/patches/0002-glt-int32-csr-indices.patch." + ) + self.assertTrue(_HAS_INT32) + + +@unittest.skipUnless(_HAS_INT32, _SKIP_REASON) +class Int32SamplingParityTest(TestCase): + """The int32 graph must sample IDENTICALLY to the int64 graph over the same CSR.""" + + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + cls.num_rows = 20_000 + cls.indptr, cls.indices64, cls.degrees = _random_csr( + cls.num_rows, 200_000, seed=7 + ) + cls.indices32 = cls.indices64.to(torch.int32) + cls.graph64 = _build_cpu_graph(cls.indptr, cls.indices64) + cls.graph32 = _build_cpu_graph(cls.indptr, cls.indices32) + generator = torch.Generator().manual_seed(11) + cls.seeds = torch.randint( + 0, cls.num_rows, (4_000,), generator=generator, dtype=torch.int64 + ) + cls.max_degree = int(cls.degrees.max().item()) + + def test_full_fanout_sampling_is_identical_element_wise(self) -> None: + """req_num > max degree makes UniformSample copy, not draw -- so this is exact.""" + sampler64 = pywrap.CPURandomSampler(self.graph64.graph_handler) + sampler32 = pywrap.CPURandomSampler(self.graph32.graph_handler) + + neighbors64, counts64 = sampler64.sample(self.seeds, self.max_degree + 1) + neighbors32, counts32 = sampler32.sample(self.seeds, self.max_degree + 1) + + self.assertTrue(torch.equal(counts64, counts32), "degree counts diverged") + self.assertTrue( + torch.equal(neighbors64, neighbors32), + "sampled neighbor ids diverged between the int64 and int32 graphs", + ) + self.assertGreater(neighbors32.numel(), 0, "fixture sampled nothing") + + def test_sampled_outputs_stay_int64(self) -> None: + """The narrowing is at REST only. + + Sampled ids flow into ``Data``/``HeteroData`` and are indexed against int64 feature + tables and label tensors, so an int32 output would propagate the dtype into the whole + batch path. The patch casts on read for exactly this reason. + """ + sampler32 = pywrap.CPURandomSampler(self.graph32.graph_handler) + neighbors, counts = sampler32.sample(self.seeds, 5) + self.assertEqual(neighbors.dtype, torch.int64) + self.assertEqual(counts.dtype, torch.int64) + + def test_sub_degree_draws_are_true_neighbors_with_exact_counts(self) -> None: + """A draw cannot be compared element-wise, so check membership and cardinality.""" + sampler32 = pywrap.CPURandomSampler(self.graph32.graph_handler) + requested = 5 + neighbors, counts = sampler32.sample(self.seeds, requested) + offsets = torch.zeros(self.seeds.numel() + 1, dtype=torch.int64) + torch.cumsum(counts, dim=0, out=offsets[1:]) + + for index in range(0, self.seeds.numel(), 173): + node = int(self.seeds[index]) + true_neighbors = set( + self.indices64[self.indptr[node] : self.indptr[node + 1]].tolist() + ) + drawn = neighbors[offsets[index] : offsets[index + 1]].tolist() + self.assertEqual( + len(drawn), + min(requested, int(self.degrees[node])), + f"wrong number of neighbors drawn for node {node}", + ) + for neighbor in drawn: + self.assertIn( + neighbor, true_neighbors, f"node {node} drew a non-neighbor" + ) + + def test_col_count_is_identical(self) -> None: + """col_count comes from the patched distinct-count, which is now dtype-generic.""" + self.assertEqual(self.graph64.col_count, self.graph32.col_count) + self.assertGreater(self.graph32.col_count, 0) + + def test_the_int32_graph_holds_half_the_bytes(self) -> None: + """The entire point: the CSC column array halves at billion-edge scale.""" + bytes64 = self.indices64.numel() * self.indices64.element_size() + bytes32 = self.indices32.numel() * self.indices32.element_size() + self.assertEqual(bytes64, 2 * bytes32) + + +@unittest.skipUnless(_HAS_INT32, _SKIP_REASON) +class Int32RejectionTest(TestCase): + """Paths NOT taught the int32 layout must fail loudly, never sample wrongly. + + ``col_idx_`` is nullptr on an int32 graph, so an untaught consumer that kept reading it + would dereference null (a crash at best, garbage at worst). Each of these asserts the + ``TORCH_CHECK`` fires instead. + """ + + def setUp(self) -> None: + super().setUp() + self.indptr, indices64, _ = _random_csr(500, 4_000, seed=13) + self.graph32 = _build_cpu_graph(self.indptr, indices64.to(torch.int32)) + self.seeds = torch.arange(10, dtype=torch.int64) + + def test_the_weighted_sampler_rejects_an_int32_graph(self) -> None: + with self.assertRaises(RuntimeError) as caught: + pywrap.CPUWeightedSampler(self.graph32.graph_handler).sample(self.seeds, 3) + self.assertIn("int32", str(caught.exception)) + + def test_an_unsupported_indices_dtype_is_rejected(self) -> None: + """Only int32 and int64 are accepted; anything else must not be reinterpreted.""" + with self.assertRaises(RuntimeError): + _build_cpu_graph( + torch.tensor([0, 1], dtype=torch.int64), + torch.tensor([0], dtype=torch.int16), + ) + + def test_a_non_contiguous_int32_column_array_is_rejected(self) -> None: + """Samplers read the column array as flat storage, so a strided view must not be taken.""" + indices = torch.arange(8, dtype=torch.int32)[::2] + self.assertFalse(indices.is_contiguous()) + with self.assertRaises(RuntimeError): + _build_cpu_graph(torch.tensor([0, 2, 4], dtype=torch.int64), indices) + + @unittest.skipUnless( + torch.cuda.is_available() and hasattr(pywrap.Graph(), "init_cuda_from_csr"), + "needs a GPU and a WITH_CUDA=ON graphlearn_torch build", + ) + def test_cuda_init_rejects_an_int32_topology(self) -> None: + """The CUDA path was deliberately NOT taught int32. + + It reads ``data_ptr()``, which throws on a dtype mismatch -- so the failure is + loud for free. Asserted rather than assumed, because "it happens to throw" is a property + of torch's accessor that a future refactor could remove. + + Skipped on a CPU-only wheel, where the entry point does not exist at all: that is the + local build, NOT the artifact that ships. This assertion only means something inside the + base image, which is where the in-image re-verification runs it. + """ + indptr, indices64, _ = _random_csr(100, 500, seed=17) + topology = Topology.__new__(Topology) + topology._layout = "CSR" + topology._indptr = indptr + topology._indices = indices64.to(torch.int32) + topology._edge_ids = None + topology._edge_weights = None + graph = Graph.__new__(Graph) + graph.topo = topology + graph.mode = "CUDA" + graph.device = 0 + graph._graph = None + with self.assertRaises(RuntimeError): + graph.lazy_init() + + +@unittest.skipUnless(_HAS_INT32, _SKIP_REASON) +class Int32EmptyGraphTest(TestCase): + """An edgeless rank is normal under range partitioning, and it still narrows to int32.""" + + def test_an_empty_int32_graph_initializes_and_samples_nothing(self) -> None: + graph = _build_cpu_graph( + torch.zeros(11, dtype=torch.int64), torch.empty(0, dtype=torch.int32) + ) + self.assertEqual(graph.col_count, 0) + + sampler = pywrap.CPURandomSampler(graph.graph_handler) + neighbors, counts = sampler.sample(torch.arange(5, dtype=torch.int64), 3) + + self.assertEqual(neighbors.numel(), 0) + self.assertEqual(int(counts.sum()), 0) + self.assertEqual(neighbors.dtype, torch.int64) + + +if __name__ == "__main__": + from absl.testing import absltest + + absltest.main()