[DO NOT MERGE] Stress test CI with combined flaky-test fixes

.github/workflows/pr.yaml

@@ -36,6 +36,7 @@ jobs:
       - wheel-tests-cuvs
       - devcontainer
       - telemetry-setup
+      - stress-rerun-anchor
     uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@main
     permissions:
       contents: read
@@ -53,6 +54,14 @@ jobs:
       - name: Telemetry setup
         if: ${{ vars.TELEMETRY_ENABLED == 'true' }}
         uses: rapidsai/shared-actions/telemetry-dispatch-stash-base-env-vars@main
+  stress-rerun-anchor:
+    # Stress testing only: every test job needs this job, so re-running it re-runs all
+    # test jobs against the existing build artifacts without re-running any builds.
+    runs-on: ubuntu-latest
+    permissions: {}
+    steps:
+      - name: Anchor for test job re-runs
+        run: echo "Re-running the test jobs"
   build-details:
     uses: rapidsai/shared-workflows/.github/workflows/compute-build-details.yaml@main
   check-nightly-ci:
@@ -469,7 +478,7 @@ jobs:
       node_type: cpu16
       script: ci/build_cpp.sh
   conda-cpp-tests:
-    needs: [conda-cpp-build, changed-files]
+    needs: [conda-cpp-build, changed-files, stress-rerun-anchor]
     permissions:
       actions: read
       contents: read
@@ -478,7 +487,6 @@ jobs:
       pull-requests: read
     secrets: inherit # zizmor: ignore[secrets-inherit]
     uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@main
-    if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp
     # Split the C++ tests across several jobs for each CUDA/GPU configuration
     strategy:
       fail-fast: false
@@ -488,6 +496,7 @@ jobs:
     with:
       build_type: pull-request
       script: "ci/test_cpp.sh ${{ matrix.shard }} ${{ matrix.num_shards }}"
+      matrix_type: "pull-request,nightly"
       cache-paths: .cache/cuda-jit
       cache-key-prefix: libcuvs-cudajit-v1
       # Caches are immutable. Re-key when the JIT-relevant sources change so a completed run can
@@ -548,7 +557,7 @@ jobs:
       # Build a conda package for each CUDA x ARCH x minimum supported Python version
       matrix_filter: group_by({CUDA_VER, ARCH}) | map(min_by(.PY_VER | split(".") | map(tonumber)))
   conda-python-tests:
-    needs: [conda-python-build, changed-files]
+    needs: [conda-python-build, changed-files, stress-rerun-anchor]
     permissions:
       actions: read
       contents: read
@@ -557,10 +566,10 @@ jobs:
       pull-requests: read
     secrets: inherit # zizmor: ignore[secrets-inherit]
     uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@main
-    if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_conda
     with:
       build_type: pull-request
       script: ci/test_python.sh
+      matrix_type: "pull-request,nightly"
   rocky8-clib-standalone-build:
     needs: checks
     permissions:
@@ -573,15 +582,14 @@ jobs:
     with:
       build_type: pull-request
   rocky8-clib-tests:
-    needs: [rocky8-clib-standalone-build, changed-files]
+    needs: [rocky8-clib-standalone-build, changed-files, stress-rerun-anchor]
     permissions:
       actions: read
       contents: read
       id-token: write
       packages: read
       pull-requests: read
     uses: ./.github/workflows/matrix-job.yaml
-    if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_standalone_c
     with:
       build_type: pull-request
       container-image-repository: "rapidsai/citestwheel"
@@ -595,7 +603,6 @@ jobs:
       packages: read
       pull-requests: read
     uses: ./.github/workflows/matrix-job.yaml
-    if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_java || fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp
     # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions.
     # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions.
     with:
@@ -606,7 +613,7 @@ jobs:
       artifact-name-prefix: "cuvs-java"
       file_to_upload: "java/cuvs-java/target/"
   conda-java-tests:
-    needs: conda-java-build
+    needs: [conda-java-build, stress-rerun-anchor]
     permissions:
       actions: read
       contents: read
@@ -637,7 +644,7 @@ jobs:
       artifact-name-prefix: "cuvs-lucene"
       file_to_upload: "java/cuvs-lucene/target/"
   conda-lucene-tests:
-    needs: conda-lucene-build
+    needs: [conda-lucene-build, stress-rerun-anchor]
     permissions:
       actions: read
       contents: read
@@ -746,7 +753,7 @@ jobs:
       # Build a wheel for each CUDA x ARCH x minimum supported Python version
       matrix_filter: group_by({CUDA_VER, ARCH}) | map(min_by(.PY_VER | split(".") | map(tonumber)))
   wheel-tests-cuvs:
-    needs: [wheel-build-cuvs, changed-files]
+    needs: [wheel-build-cuvs, changed-files, stress-rerun-anchor]
     permissions:
       actions: read
       contents: read
@@ -755,10 +762,10 @@ jobs:
       pull-requests: read
     secrets: inherit # zizmor: ignore[secrets-inherit]
     uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main
-    if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels
     with:
       build_type: pull-request
       script: ci/test_wheel_cuvs.sh
+      matrix_type: "pull-request,nightly"
   devcontainer:
     permissions:
       actions: read

cpp/CMakeLists.txt

@@ -1472,6 +1472,7 @@ if(NOT BUILD_CPU_ONLY)
     src/neighbors/detail/cagra/topk_for_cagra/topk.cu
     ${cuvs_cagra_search_cuda_inst_files}
     src/neighbors/detail/cagra/cagra_helpers.cpp
+    src/neighbors/detail/pointer_attributes.cpp
     src/neighbors/dynamic_batching.cu
     src/neighbors/composite/index.cu
     $<$<BOOL:${BUILD_CAGRA_HNSWLIB}>:src/neighbors/cagra.cpp>

cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h

@@ -66,13 +66,15 @@ inline void maybe_log_cagra_persistent_concurrency_hint(bool persistent_search)
     threads_rec);
 }
 
-/** Whether the GPU can dereference this pointer: device, managed or mapped host memory. */
+/**
+ * Whether the GPU can dereference this pointer: device, managed or mapped host memory, or pageable
+ * memory on systems where the device reads it through the host page tables (ATS).
+ */
 inline auto is_device_accessible(const void* ptr) -> bool
 {
   cudaPointerAttributes attrs{};
   RAFT_CUDA_TRY(cudaPointerGetAttributes(&attrs, ptr));
-  return attrs.type == cudaMemoryTypeDevice || attrs.type == cudaMemoryTypeManaged ||
-         attrs.devicePointer != nullptr;
+  return cuvs::neighbors::detail::is_device_accessible(attrs);
 }
 
 /**

cpp/cmake/modules/generate_cutile_kernels.cmake

@@ -34,14 +34,23 @@ function(_cutile_detect)
     return()
   endif()
 
+  # bin2c is installed next to the real nvcc. In conda environments, CUDAToolkit_BIN_DIR is
+  # targets/<arch>/bin, which only contains an nvcc symlink, so also search the symlink target's
+  # directory.
+  set(bin2c_dirs ${CUDAToolkit_BIN_DIR})
+  if(CUDAToolkit_NVCC_EXECUTABLE)
+    get_filename_component(nvcc_real_path "${CUDAToolkit_NVCC_EXECUTABLE}" REALPATH)
+    get_filename_component(nvcc_real_dir "${nvcc_real_path}" DIRECTORY)
+    list(APPEND bin2c_dirs "${nvcc_real_dir}")
+  endif()
   find_program(
     CUVS_CUTILE_BIN2C
     NAMES bin2c
-    PATHS ${CUDAToolkit_BIN_DIR}
+    PATHS ${bin2c_dirs}
     NO_DEFAULT_PATH
   )
   if(NOT CUVS_CUTILE_BIN2C)
-    message(STATUS "cuTile disabled: bin2c not found in ${CUDAToolkit_BIN_DIR}.")
+    message(STATUS "cuTile disabled: bin2c not found in ${bin2c_dirs}.")
     return()
   endif()
 

cpp/include/cuvs/neighbors/common.hpp

@@ -1094,6 +1094,18 @@ template <typename T, typename I, typename L>
 
 namespace detail {
 
+/**
+ * Whether memory described by `attr` may be dereferenced in device code on the current device.
+ *
+ * Device and managed memory are device-accessible, and registered (pinned) host memory is when it
+ * is mapped (`devicePointer` is set). Pageable (unregistered) host memory is device-accessible only
+ * when the device reads it through the host page tables (ATS). On systems with HMM,
+ * `cudaPointerGetAttributes` can also report a non-null `devicePointer` for pageable memory (for
+ * example after the first GPU access to pageable memory in the process), but device accesses to it
+ * are then serviced through GPU page faults, so it is treated as host-only there.
+ */
+[[nodiscard]] auto is_device_accessible(const cudaPointerAttributes& attr) -> bool;
+
 template <typename SrcT>
 [[nodiscard]] inline uint32_t mdspan_row_stride_elements(SrcT const& src)
 {
@@ -1115,7 +1127,10 @@ template <typename ValueT, typename SrcT>
   if (ptr_attrs.type == cudaMemoryTypeDevice || ptr_attrs.type == cudaMemoryTypeManaged) {
     return const_cast<ValueT*>(src.data_handle());
   }
-  auto* device_ptr = reinterpret_cast<ValueT*>(ptr_attrs.devicePointer);
+  // Host memory needs a device mapping; pageable memory only counts when the device reads it
+  // through the host page tables (ATS), not through HMM page faults.
+  auto* device_ptr =
+    is_device_accessible(ptr_attrs) ? reinterpret_cast<ValueT*>(ptr_attrs.devicePointer) : nullptr;
   RAFT_EXPECTS(device_ptr != nullptr, "%s", error_msg);
   return device_ptr;
 }

cpp/src/cluster/detail/kmeans.cuh

@@ -1082,7 +1082,8 @@ void kmeans_fit(
       compute_centroid_shift<DataT, IndexT>(handle,
                                             raft::make_const_mdspan(centroids_const),
                                             raft::make_const_mdspan(new_centroids_view),
-                                            sqrdNormError.view());
+                                            sqrdNormError.view(),
+                                            ws);
 
       std::swap(cur_centroids_ptr, new_centroids_ptr);
 

cpp/src/cluster/detail/kmeans_balanced.cuh

@@ -43,6 +43,7 @@
 
 #include <algorithm>
 #include <limits>
+#include <numeric>
 #include <optional>
 #include <tuple>
 #include <type_traits>
@@ -640,7 +641,7 @@ __launch_bounds__((raft::WarpSize * BlockDimY)) RAFT_KERNEL
                                      IdxT average,
                                      MathT centroid_offset,
                                      IdxT seed,
-                                     IdxT* search_count,
+                                     IdxT stride,
                                      IdxT* update_count,
                                      MappingOpT mapping_op)
 {
@@ -652,12 +653,17 @@ __launch_bounds__((raft::WarpSize * BlockDimY)) RAFT_KERNEL
   IdxT i = n_rows;
   IdxT j = raft::laneId();
   if (j == 0) {
+    // The candidate sequence of each receiver is derived from its index (not from a shared atomic
+    // counter), so the donor points do not depend on the order in which the warps run.
+    // `stride >= n_clusters`, so different receivers use different counters (and thus, before
+    // wrapping around, different candidates); `stride` is coprime with `n_rows`, so each receiver
+    // visits every row within `n_rows` attempts and finds a donor whenever one exists.
     IdxT attempt = 0;
     do {
-      auto old = atomicAdd(search_count, IdxT{1});
-      auto candidate =
-        static_cast<IdxT>((static_cast<int64_t>(seed) * static_cast<int64_t>(old + 1)) %
-                          static_cast<int64_t>(n_rows));
+      auto counter = static_cast<int64_t>(attempt) * static_cast<int64_t>(stride) +
+                     static_cast<int64_t>(receiver_cluster);
+      auto candidate = static_cast<IdxT>((static_cast<int64_t>(seed) * (counter + 1)) %
+                                         static_cast<int64_t>(n_rows));
       if (static_cast<IdxT>(cluster_sizes[labels[candidate]]) >= average) { i = candidate; }
       ++attempt;
     } while (i >= n_rows && attempt < n_rows);
@@ -768,6 +774,8 @@ __launch_bounds__((raft::WarpSize * BlockDimY)) RAFT_KERNEL
  *                   max_cluster_size > average_size * balance_upper_tolerance
  *                   balance_upper_tolerance > 1
  * @param[in] centroid_offset offset from the donor cluster centroid towards a donor point
+ * @param[in] donor_selection how the donor points are chosen (see above)
+ * @param[in] balancing_step index of the current EM iteration; selects the pseudo-random offset
  * @param[in] mapping_op Mapping operation from T to MathT
  * @param[inout] device_memory  memory resource to use for temporary allocations
  *
@@ -791,6 +799,7 @@ auto adjust_centers(const raft::resources& handle,
                     MathT balance_upper_tolerance,
                     MathT centroid_offset,
                     cuvs::cluster::kmeans::balanced_donor_selection donor_selection,
+                    uint32_t balancing_step,
                     MappingOpT mapping_op,
                     rmm::device_async_resource_ref device_memory) -> bool
 {
@@ -802,7 +811,6 @@ auto adjust_centers(const raft::resources& handle,
                                       601,  659,  733,  809,  863,  941,  1013, 1069, 1151, 1223,
                                       1291, 1373, 1451, 1511, 1583, 1657, 1733, 1811, 1889, 1987,
                                       2053, 2129, 2213, 2287, 2357, 2423, 2531, 2617, 2687, 2741};
-  static IdxT i_primes = 0;
 
   auto average         = static_cast<MathT>(n_rows) / static_cast<MathT>(n_clusters);
   auto lower_threshold = average * balance_lower_tolerance;
@@ -837,10 +845,12 @@ auto adjust_centers(const raft::resources& handle,
   auto n_pairs = static_cast<IdxT>(host_receiver_clusters.size());
   if (n_pairs == 0) { return false; }
 
+  // Rotate the pseudo-random offset with the balancing step rather than with a process-wide
+  // counter, so that the donor points do not depend on how many fits ran before in the process.
   IdxT ofst;
+  size_t i_primes = balancing_step;
   do {
-    i_primes = (i_primes + 1) % kPrimes.size();
-    ofst     = kPrimes[i_primes];
+    ofst = kPrimes[i_primes++ % kPrimes.size()];
   } while (n_rows % ofst == 0);
 
   rmm::device_uvector<IdxT> receiver_clusters(n_pairs, stream, device_memory);
@@ -851,8 +861,11 @@ auto adjust_centers(const raft::resources& handle,
   update_count.set_value_to_zero_async(stream);
 
   if (donor_selection == cuvs::cluster::kmeans::balanced_donor_selection::Random) {
-    rmm::device_scalar<IdxT> search_count(stream, device_memory);
-    search_count.set_value_to_zero_async(stream);
+    // The smallest stride >= n_clusters that is coprime with n_rows (see the kernel).
+    auto stride = std::max<IdxT>(n_clusters, 1);
+    while (std::gcd(static_cast<uint64_t>(stride), static_cast<uint64_t>(n_rows)) != 1) {
+      ++stride;
+    }
     const dim3 grid_dim(raft::ceildiv(n_clusters, static_cast<IdxT>(kBlockDimY)), 1, 1);
     adjust_centers_random_donor_kernel<kBlockDimY>
       <<<grid_dim, block_dim, 0, stream.get()>>>(centers,
@@ -866,7 +879,7 @@ auto adjust_centers(const raft::resources& handle,
                                                  static_cast<IdxT>(n_rows / n_clusters),
                                                  centroid_offset,
                                                  ofst,
-                                                 search_count.data(),
+                                                 stride,
                                                  update_count.data(),
                                                  mapping_op);
     return update_count.value(stream) > 0;  // NB: rmm scalar performs the sync
@@ -979,6 +992,7 @@ void balancing_em_iters(const raft::resources& handle,
                                    balance_upper_tolerance,
                                    static_cast<MathT>(params.centroid_offset),
                                    params.donor_selection,
+                                   iter,
                                    mapping_op,
                                    device_memory)) {
       if (balancing_counter++ >= balancing_pullback) {