diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index a6438c1a8d..8e4dbb47cd 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -11,450 +11,19 @@ jobs: pr-builder: needs: - build-details - - check-nightly-ci - - changed-files - - checks - - check-c-abi - - security-suite - conda-cpp-build - - conda-cpp-tests - - conda-cpp-checks - - conda-python-build - - conda-python-tests - - rocky8-clib-standalone-build - - rocky8-clib-tests - - conda-java-build - - conda-java-tests - - conda-lucene-build - - conda-lucene-tests - - rust-build - - go-build - - docs-build - - docs-build-matrix - - wheel-build-libcuvs - - wheel-build-cuvs - - wheel-tests-cuvs - - devcontainer - - telemetry-setup + - rust-old-crate + - c-host-alloc uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@main permissions: contents: read if: always() with: needs: ${{ toJSON(needs) }} - telemetry-setup: - continue-on-error: true - runs-on: ubuntu-latest - permissions: - contents: read - env: - OTEL_SERVICE_NAME: 'pr-cuvs' - steps: - - name: Telemetry setup - if: ${{ vars.TELEMETRY_ENABLED == 'true' }} - uses: rapidsai/shared-actions/telemetry-dispatch-stash-base-env-vars@main build-details: uses: rapidsai/shared-workflows/.github/workflows/compute-build-details.yaml@main - check-nightly-ci: - runs-on: ubuntu-latest - permissions: - actions: read - contents: read - id-token: write - pull-requests: read - env: - GH_TOKEN: ${{ github.token }} - steps: - - name: Get PR Info - id: get-pr-info - uses: nv-gha-runners/get-pr-info@main - - name: Check if nightly CI is passing - uses: rapidsai/shared-actions/check_nightly_success/dispatch@main - with: - repo: ${{ github.repository }} - target-branch: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).base.ref }} - max-days-without-success: 30 - changed-files: - needs: telemetry-setup - permissions: - actions: read - contents: read - packages: read - pull-requests: read - uses: rapidsai/shared-workflows/.github/workflows/changed-files.yaml@main - with: - files_yaml: | - build_docs: - - '**' - - '!**/*/REVIEW_GUIDELINES.md' - - '!.coderabbit.yaml' - - '!.devcontainer/**' - - '!.github/CODEOWNERS' - - '!.github/ISSUE_TEMPLATE/**' - - '!.github/PULL_REQUEST_TEMPLATE.md' - - '!.github/copy-pr-bot.yaml' - - '!.github/labeler.yml' - - '!.github/ops-bot.yaml' - - '!.github/release.yml' - - '!.github/workflows/build.yaml' - - '!.github/workflows/check-c-abi.yaml' - - '!.github/workflows/labeler.yml' - - '!.github/workflows/publish-rust.yaml' - - '!.github/workflows/security-suite.yml' - - '!.github/workflows/store-c-abi-baseline.yaml' - - '!.github/workflows/test.yaml' - - '!.github/workflows/trigger-breaking-change-alert.yaml' - - '!.github/workflows/update-c-abi-baseline.yaml' - - '!.pre-commit-config.yaml' - - '!.shellcheckrc' - - '!.yamllint.yaml' - - '!Dockerfile' - - '!SECURITY.md' - - '!ci/build_wheel*.sh' - - '!ci/check_style.sh' - - '!ci/release/update-version.sh' - - '!ci/run*.sh' - - '!ci/test*.sh' - - '!ci/validate_wheel.sh' - - '!thirdparty/LICENSES/**' - test_cpp: - - '.github/workflows/pr.yaml' - - 'RAPIDS_BRANCH' - - 'VERSION' - - 'build.sh' - - 'c/**' - - 'ci/build_cpp.sh' - - 'ci/get_test_data.sh' - - 'ci/test_cpp.sh' - - 'cmake/**' - - 'conda/recipes/libcuvs/**' - - 'cpp/**' - - 'dependencies.yaml' - - '!**/*/REVIEW_GUIDELINES.md' - - '!cpp/doxygen/**' - test_standalone_c: - - '.github/workflows/pr.yaml' - - 'RAPIDS_BRANCH' - - 'VERSION' - - 'build.sh' - - 'c/**' - - 'ci/standalone_c/**' - - 'ci/test_standalone_c.sh' - - 'cmake/**' - - 'cpp/**' - - 'dependencies.yaml' - - '!**/*/REVIEW_GUIDELINES.md' - - '!cpp/doxygen/**' - test_java: - - '**' - - '!**/*/REVIEW_GUIDELINES.md' - - '!.coderabbit.yaml' - - '!.devcontainer/**' - - '!.github/CODEOWNERS' - - '!.github/ISSUE_TEMPLATE/**' - - '!.github/PULL_REQUEST_TEMPLATE.md' - - '!.github/copy-pr-bot.yaml' - - '!.github/labeler.yml' - - '!.github/ops-bot.yaml' - - '!.github/release.yml' - - '!.github/workflows/build.yaml' - - '!.github/workflows/check-c-abi.yaml' - - '!.github/workflows/labeler.yml' - - '!.github/workflows/publish-rust.yaml' - - '!.github/workflows/security-suite.yml' - - '!.github/workflows/store-c-abi-baseline.yaml' - - '!.github/workflows/test.yaml' - - '!.github/workflows/trigger-breaking-change-alert.yaml' - - '!.github/workflows/update-c-abi-baseline.yaml' - - '!.pre-commit-config.yaml' - - '!.shellcheckrc' - - '!.yamllint.yaml' - - '!CODEOWNERS' - - '!Dockerfile' - - '!README.md' - - '!SECURITY.md' - - '!ci/build_docs.sh' - - '!ci/build_go.sh' - - '!ci/build_python.sh' - - '!ci/build_rust.sh' - - '!ci/build_wheel*.sh' - - '!ci/checkstyle.sh' - - '!ci/run_cuvs_pytests.sh' - - '!ci/test_cpp.sh' - - '!ci/test_python.sh' - - '!ci/test_wheel_cuvs.sh' - - '!ci/validate_wheel.sh' - - '!conda/environments/**' - - '!conda/recipes/cuvs/**' - - '!conda/recipes/cuvs-bench/**' - - '!conda/recipes/cuvs-bench-cpu/**' - - '!docs/**' - - '!examples/.clang-format' - - '!examples/.clang-tidy' - - '!examples/.clangd' - - '!examples/README.md' - - '!examples/build.sh' - - '!examples/c/**' - - '!examples/cmake/**' - - '!examples/cpp/**' - - '!examples/go/**' - - '!fern/**' - - '!go/**' - - '!img/**' - - '!notebooks/**' - - '!pyproject.toml' - - '!python/**' - - '!rust/**' - - '!thirdparty/LICENSES/**' - test_python_conda: - - '**' - - '!**/*/REVIEW_GUIDELINES.md' - - '!.coderabbit.yaml' - - '!.devcontainer/**' - - '!.github/CODEOWNERS' - - '!.github/ISSUE_TEMPLATE/**' - - '!.github/PULL_REQUEST_TEMPLATE.md' - - '!.github/copy-pr-bot.yaml' - - '!.github/labeler.yml' - - '!.github/ops-bot.yaml' - - '!.github/release.yml' - - '!.github/workflows/build.yaml' - - '!.github/workflows/check-c-abi.yaml' - - '!.github/workflows/labeler.yml' - - '!.github/workflows/publish-rust.yaml' - - '!.github/workflows/security-suite.yml' - - '!.github/workflows/store-c-abi-baseline.yaml' - - '!.github/workflows/test.yaml' - - '!.github/workflows/trigger-breaking-change-alert.yaml' - - '!.github/workflows/update-c-abi-baseline.yaml' - - '!.pre-commit-config.yaml' - - '!.shellcheckrc' - - '!.yamllint.yaml' - - '!ci/build_go.sh' - - '!ci/build_java.sh' - - '!ci/build_lucene.sh' - - '!ci/build_rust.sh' - - '!ci/check_style.sh' - - '!ci/release/update-version.sh' - - '!ci/run_ctests.sh' - - '!ci/test_cpp.sh' - - '!ci/test_java.sh' - - '!ci/test_lucene.sh' - - '!ci/test_standalone_c.sh' - - '!ci/test_wheel_cuvs.sh' - - '!ci/validate_wheel.sh' - - '!cpp/.clang-format' - - '!cpp/.clang-tidy' - - '!CODEOWNERS' - - '!Dockerfile' - - '!README.md' - - '!SECURITY.md' - - '!docs/**' - - '!examples/**' - - '!fern/**' - - '!go/**' - - '!img/**' - - '!java/**' - - '!notebooks/**' - - '!rust/**' - - '!thirdparty/LICENSES/**' - test_python_wheels: - - '**' - - '!**/*/REVIEW_GUIDELINES.md' - - '!.coderabbit.yaml' - - '!.devcontainer/**' - - '!.github/CODEOWNERS' - - '!.github/ISSUE_TEMPLATE/**' - - '!.github/PULL_REQUEST_TEMPLATE.md' - - '!.github/copy-pr-bot.yaml' - - '!.github/labeler.yml' - - '!.github/ops-bot.yaml' - - '!.github/release.yml' - - '!.github/workflows/build.yaml' - - '!.github/workflows/check-c-abi.yaml' - - '!.github/workflows/labeler.yml' - - '!.github/workflows/publish-rust.yaml' - - '!.github/workflows/security-suite.yml' - - '!.github/workflows/store-c-abi-baseline.yaml' - - '!.github/workflows/test.yaml' - - '!.github/workflows/trigger-breaking-change-alert.yaml' - - '!.github/workflows/update-c-abi-baseline.yaml' - - '!.pre-commit-config.yaml' - - '!.shellcheckrc' - - '!.yamllint.yaml' - - '!ci/build_cpp.sh' - - '!ci/build_docs.sh' - - '!ci/build_go.sh' - - '!ci/build_java.sh' - - '!ci/build_lucene.sh' - - '!ci/build_python.sh' - - '!ci/build_rust.sh' - - '!ci/build_standalone_c.sh' - - '!ci/check_style.sh' - - '!ci/release/update-version.sh' - - '!ci/run_ctests.sh' - - '!ci/test_cpp.sh' - - '!ci/test_java.sh' - - '!ci/test_lucene.sh' - - '!ci/test_python.sh' - - '!ci/test_standalone_c.sh' - - '!ci/validate_wheel.sh' - - '!CODEOWNERS' - - '!Dockerfile' - - '!README.md' - - '!SECURITY.md' - - '!conda/**' - - '!cpp/.clang-format' - - '!cpp/.clang-tidy' - - '!docs/**' - - '!examples/**' - - '!fern/**' - - '!java/**' - - '!go/**' - - '!img/**' - - '!notebooks/**' - - '!rust/**' - - '!thirdparty/LICENSES/**' - test_rust: - - '**' - - '!**/*/REVIEW_GUIDELINES.md' - - '!.coderabbit.yaml' - - '!.devcontainer/**' - - '!.github/CODEOWNERS' - - '!.github/ISSUE_TEMPLATE/**' - - '!.github/PULL_REQUEST_TEMPLATE.md' - - '!.github/copy-pr-bot.yaml' - - '!.github/labeler.yml' - - '!.github/ops-bot.yaml' - - '!.github/release.yml' - - '!.github/workflows/build.yaml' - - '!.github/workflows/check-c-abi.yaml' - - '!.github/workflows/labeler.yml' - - '!.github/workflows/publish-rust.yaml' - - '!.github/workflows/security-suite.yml' - - '!.github/workflows/store-c-abi-baseline.yaml' - - '!.github/workflows/test.yaml' - - '!.github/workflows/trigger-breaking-change-alert.yaml' - - '!.github/workflows/update-c-abi-baseline.yaml' - - '!.pre-commit-config.yaml' - - '!.shellcheckrc' - - '!.yamllint.yaml' - - '!CODEOWNERS' - - '!Dockerfile' - - '!README.md' - - '!SECURITY.md' - - '!ci/build_docs.sh' - - '!ci/build_go.sh' - - '!ci/build_java.sh' - - '!ci/build_lucene.sh' - - '!ci/build_python.sh' - - '!ci/build_wheel*.sh' - - '!ci/check_style.sh' - - '!ci/release/update-version.sh' - - '!ci/run_ctests.sh' - - '!ci/run_cuvs_pytests.sh' - - '!ci/test_cpp.sh' - - '!ci/test_java.sh' - - '!ci/test_lucene.sh' - - '!ci/test_python.sh' - - '!ci/test_wheel_cuvs.sh' - - '!ci/validate_wheel.sh' - - '!docs/**' - - '!examples/**' - - '!fern/**' - - '!go/**' - - '!img/**' - - '!java/**' - - '!notebooks/**' - - '!python/**' - - '!thirdparty/LICENSES/**' - test_go: - - '**' - - '!**/*/REVIEW_GUIDELINES.md' - - '!.coderabbit.yaml' - - '!.devcontainer/**' - - '!.github/CODEOWNERS' - - '!.github/ISSUE_TEMPLATE/**' - - '!.github/PULL_REQUEST_TEMPLATE.md' - - '!.github/copy-pr-bot.yaml' - - '!.github/labeler.yml' - - '!.github/ops-bot.yaml' - - '!.github/release.yml' - - '!.github/workflows/build.yaml' - - '!.github/workflows/check-c-abi.yaml' - - '!.github/workflows/labeler.yml' - - '!.github/workflows/publish-rust.yaml' - - '!.github/workflows/security-suite.yml' - - '!.github/workflows/store-c-abi-baseline.yaml' - - '!.github/workflows/test.yaml' - - '!.github/workflows/trigger-breaking-change-alert.yaml' - - '!.github/workflows/update-c-abi-baseline.yaml' - - '!.pre-commit-config.yaml' - - '!.shellcheckrc' - - '!.yamllint.yaml' - - '!CODEOWNERS' - - '!Dockerfile' - - '!README.md' - - '!SECURITY.md' - - '!ci/build_docs.sh' - - '!ci/build_java.sh' - - '!ci/build_lucene.sh' - - '!ci/build_python.sh' - - '!ci/build_rust.sh' - - '!ci/build_wheel*.sh' - - '!ci/check_style.sh' - - '!ci/release/update-version.sh' - - '!ci/run_ctests.sh' - - '!ci/run_cuvs_pytests.sh' - - '!ci/test_cpp.sh' - - '!ci/test_java.sh' - - '!ci/test_lucene.sh' - - '!ci/test_python.sh' - - '!ci/test_standalone_c.sh' - - '!ci/test_wheel_cuvs.sh' - - '!ci/validate_wheel.sh' - - '!docs/**' - - '!examples/.clang-format' - - '!examples/.clang-tidy' - - '!examples/.clangd' - - '!examples/README.md' - - '!examples/build.sh' - - '!examples/c/**' - - '!examples/cmake/**' - - '!examples/cpp/**' - - '!examples/java/**' - - '!fern/**' - - '!img/**' - - '!java/**' - - '!notebooks/**' - - '!python/**' - - '!rust/**' - - '!thirdparty/LICENSES/**' - checks: - needs: telemetry-setup - uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@main - permissions: - contents: read - with: - enable_check_generated_files: false - ignored_pr_jobs: "telemetry-summarize" - check-c-abi: - needs: telemetry-setup - permissions: - actions: read - contents: read - issues: write - pull-requests: read - uses: ./.github/workflows/check-c-abi.yaml - security-suite: - permissions: - actions: read - contents: read - id-token: write # OIDC -> Vault -> nvcr.io image pull - security-events: write # publish redacted SARIF to code scanning - uses: ./.github/workflows/security-suite.yml conda-cpp-build: - needs: [build-details, checks] + needs: [build-details] permissions: actions: read contents: read @@ -468,188 +37,10 @@ jobs: build-datetime: ${{ needs.build-details.outputs.build-datetime }} node_type: cpu16 script: ci/build_cpp.sh - conda-cpp-tests: - needs: [conda-cpp-build, changed-files] - permissions: - actions: read - contents: read - id-token: write - packages: read - 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 - matrix: - shard: [1, 2, 3, 4] - num_shards: [4] - with: - build_type: pull-request - script: "ci/test_cpp.sh ${{ matrix.shard }} ${{ matrix.num_shards }}" - 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 - # save newly compiled entries; restore a compatible earlier archive first. - cache-key-files: | - cpp/CMakeLists.txt - cpp/cmake/** - cpp/include/cuvs/detail/jit_lto/** - cpp/src/detail/jit_lto/** - cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/** - cpp/src/neighbors/detail/cagra/jit_lto_kernels/** - cpp/src/neighbors/detail/jit_lto_kernels/** - cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/** - cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/** - cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/** - cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/** - cpp/tests/** - dependencies.yaml - # The CUDA driver's JIT cache is independent of the Python interpreter used to run C++ tests. - cache-key-matrix-fields: | - CUDA_VER - ARCH - LINUX_VER - GPU - DRIVER - DEPENDENCIES - # PRs restore a compatible main cache but do not create merge-ref-scoped entries. - cache-read-only: true - cache-environment: | - CUDA_CACHE_PATH=.cache/cuda-jit - conda-cpp-checks: - needs: conda-cpp-build - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@main - with: - build_type: pull-request - symbol_exclusions: (void (thrust::|cub::)) - package_name: libcuvs - conda-python-build: - needs: [build-details, conda-cpp-build] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@main - with: - build_type: pull-request - build-datetime: ${{ needs.build-details.outputs.build-datetime }} - script: ci/build_python.sh - # 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] - permissions: - actions: read - contents: read - id-token: write - packages: read - 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 - rocky8-clib-standalone-build: - needs: checks - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: ./.github/workflows/clib-build.yaml - with: - build_type: pull-request - rocky8-clib-tests: - needs: [rocky8-clib-standalone-build, changed-files] - 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" - script: "ci/test_standalone_c.sh" - conda-java-build: - needs: [conda-cpp-build, changed-files] - 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_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: - build_type: pull-request - node_type: cpu8 - matrix_filter: map(select(.ARCH == "amd64")) - script: "ci/build_java.sh" - artifact-name-prefix: "cuvs-java" - file_to_upload: "java/cuvs-java/target/" - conda-java-tests: - needs: conda-java-build - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: ./.github/workflows/matrix-job.yaml - with: - build_type: pull-request - script: "ci/test_java.sh" - conda-lucene-build: - # Depends on the Java job for the cuvs-java artifact that cuvs-lucene builds against. - needs: conda-java-build - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: ./.github/workflows/matrix-job.yaml - # 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: - build_type: pull-request - node_type: cpu8 - matrix_filter: map(select(.ARCH == "amd64")) - script: "ci/build_lucene.sh" - artifact-name-prefix: "cuvs-lucene" - file_to_upload: "java/cuvs-lucene/target/" - conda-lucene-tests: - needs: conda-lucene-build - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: ./.github/workflows/matrix-job.yaml - with: - build_type: pull-request - script: "ci/test_lucene.sh" + matrix_filter: 'map(select(.ARCH == "amd64")) | sort_by(.CUDA_VER | split(".") | map(tonumber)) | [last]' rust-build: - needs: [conda-cpp-build, changed-files] + if: false + needs: [conda-cpp-build] permissions: actions: read contents: read @@ -657,15 +48,12 @@ jobs: packages: read pull-requests: read uses: ./.github/workflows/matrix-job.yaml - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_rust || 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: build_type: pull-request - matrix_filter: map(select(.ARCH == "amd64")) + matrix_filter: 'map(select(.ARCH == "amd64")) | sort_by(.CUDA_VER | split(".") | map(tonumber)) | [last] | [range(4) as $i | .[] | . + {COPY: $i, GPU: (if $i < 2 then "rtxpro6000" else "l4" end)}]' script: "ci/build_rust.sh" - go-build: - needs: [conda-cpp-build, changed-files] + rust-old-crate: + needs: [build-details] permissions: actions: read contents: read @@ -673,124 +61,20 @@ jobs: packages: read pull-requests: read uses: ./.github/workflows/matrix-job.yaml - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_go || 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: - build_type: pull-request - matrix_filter: map(select(.ARCH == "amd64")) - script: "ci/build_go.sh" - docs-build-matrix: - permissions: - contents: read - uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@main - with: - build_type: pull-request - matrix_name: docs-build - docs-build: - needs: [docs-build-matrix] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@main - if: fromJSON(needs.changed-files.outputs.changed_file_groups).build_docs - strategy: - fail-fast: false - matrix: ${{ fromJSON(needs.docs-build-matrix.outputs.matrix) }} - with: - build_type: pull-request - node_type: "gpu-${{ matrix.GPU }}-${{ matrix.DRIVER }}-1" - arch: "${{ matrix.arch }}" - container_image: "rapidsai/ci-conda:26.12-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" - script: "ci/build_docs.sh" - wheel-build-libcuvs: - needs: [build-details, checks] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@main with: build_type: pull-request - build-datetime: ${{ needs.build-details.outputs.build-datetime }} - node_type: cpu16 - script: ci/build_wheel_libcuvs.sh - # build for every combination of arch and CUDA version, but only for the latest Python - matrix_filter: group_by([.ARCH, (.CUDA_VER|split(".")|map(tonumber)|.[0])]) | map(max_by(.PY_VER|split(".")|map(tonumber))) - package-name: libcuvs - package-type: cpp - wheel-build-cuvs: - needs: [build-details, wheel-build-libcuvs] + matrix_filter: 'map(select(.ARCH == "amd64")) | sort_by(.CUDA_VER | split(".") | map(tonumber)) | [last] | [range(2) as $i | .[] | . + {COPY: $i, GPU: (if $i < 1 then "rtxpro6000" else "l4" end)}]' + script: "ci/stress_rust_old.sh" + c-host-alloc: + needs: [conda-cpp-build] permissions: actions: read contents: read id-token: write packages: read pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@main - with: - build_type: pull-request - build-datetime: ${{ needs.build-details.outputs.build-datetime }} - node_type: cpu8 - script: ci/build_wheel_cuvs.sh - package-name: cuvs - package-type: python - # 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] - permissions: - actions: read - contents: read - id-token: write - packages: read - 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 + uses: ./.github/workflows/matrix-job.yaml with: build_type: pull-request - script: ci/test_wheel_cuvs.sh - devcontainer: - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - needs: telemetry-setup - uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@main - with: - arch: '["amd64", "arm64"]' - cuda: '["13.3"]' - node_type: "cpu8" - env: | - SCCACHE_DIST_MAX_RETRIES=inf - SCCACHE_SERVER_LOG=sccache=debug - SCCACHE_DIST_FALLBACK_TO_LOCAL_COMPILE=false - build_command: | - sccache --zero-stats; - build-all -j0 --verbose 2>&1 | tee telemetry-artifacts/build.log; - sccache --show-adv-stats | tee telemetry-artifacts/sccache-stats.txt; - telemetry-summarize: - # This job must use a self-hosted runner to record telemetry traces. - runs-on: linux-amd64-cpu4 - needs: pr-builder - permissions: - actions: write - if: ${{ vars.TELEMETRY_ENABLED == 'true' && !cancelled() }} - continue-on-error: true - steps: - - name: Telemetry summarize - uses: rapidsai/shared-actions/telemetry-dispatch-summarize@main - env: - GH_TOKEN: ${{ github.token }} + matrix_filter: 'map(select(.ARCH == "amd64")) | sort_by(.CUDA_VER | split(".") | map(tonumber)) | [last]' + script: "ci/test_c_host_alloc.sh" diff --git a/c/src/core/c_api.cpp b/c/src/core/c_api.cpp index 2e215ba414..398b983663 100644 --- a/c/src/core/c_api.cpp +++ b/c/src/core/c_api.cpp @@ -347,19 +347,24 @@ extern "C" cuvsError_t cuvsRMMMemoryResourceReset() }); } -thread_local std::unique_ptr pinned_mr; +namespace { +// One process-wide resource: pinned_host_memory_resource is stateless, so a buffer allocated on +// one thread can be freed on another (a thread_local resource left it null on other threads). +rmm::mr::pinned_host_memory_resource& pinned_mr() +{ + static rmm::mr::pinned_host_memory_resource mr; + return mr; +} +} // namespace extern "C" cuvsError_t cuvsRMMHostAlloc(void** ptr, size_t bytes) { - return cuvs::core::translate_exceptions([=] { - if (pinned_mr == nullptr) { pinned_mr = std::make_unique(); } - *ptr = pinned_mr->allocate_sync(bytes); - }); + return cuvs::core::translate_exceptions([=] { *ptr = pinned_mr().allocate_sync(bytes); }); } extern "C" cuvsError_t cuvsRMMHostFree(void* ptr, size_t bytes) { - return cuvs::core::translate_exceptions([=] { pinned_mr->deallocate_sync(ptr, bytes); }); + return cuvs::core::translate_exceptions([=] { pinned_mr().deallocate_sync(ptr, bytes); }); } thread_local std::string last_error_text = ""; diff --git a/c/tests/CMakeLists.txt b/c/tests/CMakeLists.txt index cfceb58f33..b53b2d75b4 100644 --- a/c/tests/CMakeLists.txt +++ b/c/tests/CMakeLists.txt @@ -109,6 +109,7 @@ ConfigureTest(NAME cuvs_c_headers PATH core/headers.c) ConfigureTest(NAME cuvs_c_test PATH core/c_api.c) ConfigureTest(NAME cuvs_c_resource_pool_test PATH core/resource_pool.cpp) ConfigureTest(NAME cuvs_c_resource_stream_pool_test PATH core/resource_stream_pool.cpp) +ConfigureTest(NAME cuvs_c_host_alloc_test PATH core/host_alloc.cpp) target_link_libraries(cuvs_c_test PRIVATE CUDA::cudart_static) target_link_libraries(cuvs_c_resource_pool_test PRIVATE CUDA::cudart_static) target_link_libraries(cuvs_c_resource_stream_pool_test PRIVATE CUDA::cudart_static) diff --git a/c/tests/core/host_alloc.cpp b/c/tests/core/host_alloc.cpp new file mode 100644 index 0000000000..507e11fd34 --- /dev/null +++ b/c/tests/core/host_alloc.cpp @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include +#include +#include +#include + +// Pinned host buffers may be freed on a different thread than the one that allocated them, e.g. +// when a binding moves the owning object across threads. +TEST(RMMHostAlloc, FreeOnAnotherThread) +{ + constexpr std::size_t kBytes = 1024; + void* ptr = nullptr; + cuvsError_t alloc_status = CUVS_ERROR; + std::thread([&] { alloc_status = cuvsRMMHostAlloc(&ptr, kBytes); }).join(); + ASSERT_EQ(alloc_status, CUVS_SUCCESS); + ASSERT_NE(ptr, nullptr); + + cuvsError_t free_status = CUVS_ERROR; + std::thread([&] { free_status = cuvsRMMHostFree(ptr, kBytes); }).join(); + EXPECT_EQ(free_status, CUVS_SUCCESS); +} + +TEST(RMMHostAlloc, ConcurrentAllocAndFree) +{ + constexpr int kThreads = 8; + constexpr int kIters = 50; + constexpr std::size_t kBytes = 4096; + std::vector ptrs(kThreads * kIters, nullptr); + std::atomic failures{0}; + + std::vector threads; + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t] { + for (int i = 0; i < kIters; ++i) { + if (cuvsRMMHostAlloc(&ptrs[t * kIters + i], kBytes) != CUVS_SUCCESS) { failures++; } + } + }); + } + for (auto& th : threads) { + th.join(); + } + threads.clear(); + ASSERT_EQ(failures.load(), 0); + + // Free every buffer from a different thread than the one that allocated it. + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t] { + int const owner = (t + 1) % kThreads; + for (int i = 0; i < kIters; ++i) { + if (cuvsRMMHostFree(ptrs[owner * kIters + i], kBytes) != CUVS_SUCCESS) { failures++; } + } + }); + } + for (auto& th : threads) { + th.join(); + } + EXPECT_EQ(failures.load(), 0); +} diff --git a/ci/build_rust.sh b/ci/build_rust.sh index e9218a8adc..9530058872 100755 --- a/ci/build_rust.sh +++ b/ci/build_rust.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail @@ -42,12 +42,11 @@ rapids-logger "Begin Rust build" sccache --stop-server 2>/dev/null || true -bash ./build.sh rust +INITIAL_RC=0 +bash ./build.sh rust || INITIAL_RC=$? +echo "INITIAL build.sh rust (cargo test, default parallelism) rc=${INITIAL_RC}" -sccache --show-adv-stats +sccache --show-adv-stats || true sccache --stop-server >/dev/null 2>&1 || true -# Also test out that we can publish cuvs-sys via a dry-run -pushd ./rust/cuvs-sys -cargo publish --dry-run -popd +bash ci/stress_rust.sh diff --git a/ci/stress_rust.sh b/ci/stress_rust.sh new file mode 100755 index 0000000000..8f9449ce9a --- /dev/null +++ b/ci/stress_rust.sh @@ -0,0 +1,113 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# CI-only: repeat the Rust test suite with parallel test threads (issue #2233). +set -uo pipefail + +STRESS_MINUTES=${STRESS_MINUTES:-110} +ENV_NAME=${ENV_NAME:-rust} +RUST_DIR=${RUST_DIR:-rust} + +. /opt/conda/etc/profile.d/conda.sh +set +u +conda activate "${ENV_NAME}" +set -u + +rapids-logger "Installing gdb and compute-sanitizer" +rapids-mamba-retry install -y -n "${ENV_NAME}" gdb "cuda-sanitizer-api=${RAPIDS_CUDA_VERSION%.*}" || \ + rapids-mamba-retry install -y -n "${ENV_NAME}" gdb || true +which gdb compute-sanitizer || true + +nvidia-smi || true +nproc +cat /proc/sys/kernel/core_pattern || true +ulimit -c unlimited || true +ulimit -c + +cd "${RUST_DIR}" || exit 1 +export RUST_BACKTRACE=full +cargo test -p cuvs --lib --no-run 2>&1 | tee norun.log +TEST_BIN=$(sed -n 's/.*Executable unittests src\/lib.rs (\(.*\))/\1/p' norun.log | tail -1) +TEST_BIN=$(realpath "${TEST_BIN}" 2>/dev/null || true) +echo "TEST_BIN=${TEST_BIN}" +if [ ! -x "${TEST_BIN}" ]; then echo "Failed to build the test binary"; exit 1; fi +"${TEST_BIN}" --list 2>/dev/null | tail -3 +HAVE_STRESS=0 +if "${TEST_BIN}" --list --ignored 2>/dev/null | grep -q "stress_"; then HAVE_STRESS=1; fi + +LOGDIR=$(mktemp -d) +mkdir -p "${LOGDIR}/cores" +GDB_CMDS=(-q -batch -ex "set pagination off" -ex "set print thread-events off" + -ex "handle SIGPIPE nostop noprint pass" -ex "handle SIG32 nostop noprint pass" + -ex "handle SIG33 nostop noprint pass" -ex "handle SIG34 nostop noprint pass" + -ex "handle SIGUSR1 nostop noprint pass" -ex "handle SIGUSR2 nostop noprint pass" + -ex run -ex "info threads" -ex "thread apply all bt" -ex "info sharedlibrary") + +THREADS_LIST=("" 8 16 32 64) +deadline=$(( $(date +%s) + STRESS_MINUTES * 60 )) +iter=0 +fails=0 +crashes=0 +declare -A per_mode_runs per_mode_fails +while [ "$(date +%s)" -lt "${deadline}" ]; do + iter=$((iter + 1)) + t=${THREADS_LIST[$((iter % ${#THREADS_LIST[@]}))]} + targs=() + [ -n "${t}" ] && targs=(--test-threads="${t}") + mode="native" + if (( iter % 4 == 0 )); then mode="gdb"; fi + if (( iter % 6 == 0 )) && (( HAVE_STRESS )); then mode="stress"; fi + key="${mode}-t${t:-default}" + per_mode_runs[$key]=$(( ${per_mode_runs[$key]:-0} + 1 )) + log="${LOGDIR}/iter${iter}.log" + start=$(date +%s) + case "${mode}" in + native) + (cd "${LOGDIR}/cores" && timeout 900 "${TEST_BIN}" "${targs[@]}" > "${log}" 2>&1); rc=$? ;; + gdb) + timeout 1800 gdb "${GDB_CMDS[@]}" --args "${TEST_BIN}" "${targs[@]}" > "${log}" 2>&1; rc=$? + # gdb returns 0 even when the inferior crashed; detect the signal in the log. + if grep -q "received signal SIG" "${log}"; then rc=139; fi ;; + stress) + CUVS_STRESS_THREADS=16 CUVS_STRESS_ITERS=20 timeout 1200 "${TEST_BIN}" --ignored --test-threads=2 stress_ > "${log}" 2>&1; rc=$? ;; + esac + dur=$(( $(date +%s) - start )) + summary=$(grep -E "^test result:" "${log}" | tail -1) + echo "ITER ${iter} mode=${mode} threads=${t:-default} rc=${rc} dur=${dur}s ${summary}" + if [ "${rc}" -ne 0 ]; then + fails=$((fails + 1)) + per_mode_fails[$key]=$(( ${per_mode_fails[$key]:-0} + 1 )) + if [ "${rc}" -ge 128 ] && [ "${rc}" -ne 124 ]; then crashes=$((crashes + 1)); fi + echo "::group::FAILURE iter ${iter} (${key}) rc=${rc}" + if [ "${mode}" = "gdb" ]; then + sed -n '/received signal SIG/,$p' "${log}" | head -2000 + grep -E "^test .*(FAILED|ok)$" "${log}" | tail -40 + else + grep -vE "^test .* \.\.\. ok$" "${log}" | tail -300 + fi + echo "::endgroup::" + ls -la "${LOGDIR}/cores" + for core in "${LOGDIR}"/cores/core*; do + [ -e "${core}" ] || continue + gdb -q -batch -ex "info threads" -ex "thread apply all bt" "${TEST_BIN}" "${core}" 2>&1 | head -3000 + rm -f "${core}" + done + fi +done + +echo "STRESS SUMMARY: iterations=${iter} failures=${fails} crashes=${crashes}" +for k in "${!per_mode_runs[@]}"; do + echo " ${k}: runs=${per_mode_runs[$k]} fails=${per_mode_fails[$k]:-0}" +done + +if command -v compute-sanitizer > /dev/null; then + rapids-logger "compute-sanitizer memcheck, --test-threads=8" + timeout 2400 compute-sanitizer --tool memcheck --error-exitcode 1 "${TEST_BIN}" --test-threads=8 2>&1 | tail -200 + echo "memcheck rc=${PIPESTATUS[0]}" + rapids-logger "compute-sanitizer racecheck-free: synccheck skipped; initcheck on ivf_flat" + timeout 1200 compute-sanitizer --tool memcheck --error-exitcode 1 "${TEST_BIN}" --test-threads=16 ivf_flat 2>&1 | tail -100 + echo "memcheck ivf_flat rc=${PIPESTATUS[0]}" +fi + +if [ "${fails}" -ne 0 ]; then exit 1; fi diff --git a/ci/stress_rust_old.sh b/ci/stress_rust_old.sh new file mode 100755 index 0000000000..1019a52805 --- /dev/null +++ b/ci/stress_rust_old.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# CI-only: run the stress loop on the Rust crate as it was when #2233 was filed (before the +# DLPack/lifetime rewrite in #2245), against released libcuvs 26.06 (the reporter's setup). +set -euo pipefail + +OLD_SHA=78135be4af3c7b458fd18b7b570c9b6978796daa + +. /opt/conda/etc/profile.d/conda.sh +conda config --set channel_priority flexible + +rapids-dependency-file-generator \ + --output conda \ + --file-key rust \ + --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \ + | sed -e 's/==26\.12\.\*,>=0\.0\.0a0/==26.06.*/' -e 's/==26\.12\.\*/==26.06.*/' \ + -e 's/^- rapidsai-nightly$/- rapidsai/' -e 's/^ - rapidsai-nightly$/ - rapidsai/' \ + | tee env_old.yaml + +rapids-mamba-retry env create --yes -f env_old.yaml -n rustold +set +u +conda activate rustold +set -u +conda list | grep -E "libcuvs|libraft|librmm|cuda-version|^rust " || true + +git fetch --depth 1 https://github.com/NVIDIA/cuvs.git "${OLD_SHA}" +mkdir -p old +git archive "${OLD_SHA}" rust | tar -x -C old +ls old/rust +# The old crate is versioned 26.8.0 and requires a same-minor libcuvs; the reporter used 26.06. +sed -i 's/26\.8\.0/26.6.0/g' old/rust/Cargo.toml +grep -n version old/rust/Cargo.toml + +ENV_NAME=rustold RUST_DIR=old/rust bash ci/stress_rust.sh diff --git a/ci/test_c_host_alloc.sh b/ci/test_c_host_alloc.sh new file mode 100755 index 0000000000..6b3904d602 --- /dev/null +++ b/ci/test_c_host_alloc.sh @@ -0,0 +1,50 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# CI-only: run c/tests/core/host_alloc.cpp (pinned host buffer freed on another thread) against +# the nightly libcuvs (without the fix) and against this PR's libcuvs (with the fix). +set -uo pipefail + +. /opt/conda/etc/profile.d/conda.sh +CUDA_MM=${RAPIDS_CUDA_VERSION%.*} +CPP_CHANNEL=$(rapids-download-from-github "$(rapids-artifact-name conda_cpp libcuvs cuvs --cuda "$RAPIDS_CUDA_VERSION")") +nvidia-smi || true + +build_and_run() { + local label=$1 + local cxx + cxx=$(ls "${CONDA_PREFIX}"/bin/*-linux-gnu-g++ | head -1) + "${cxx}" -std=c++20 -O1 -g c/tests/core/host_alloc.cpp -o "host_alloc_${label}" \ + -I"${CONDA_PREFIX}/include" -I"${CONDA_PREFIX}/targets/x86_64-linux/include" \ + -L"${CONDA_PREFIX}/lib" -Wl,-rpath,"${CONDA_PREFIX}/lib" \ + -lcuvs_c -lgtest -lgtest_main -pthread || { echo "BUILD FAILED (${label})"; return; } + local crashes=0 fails=0 + for i in $(seq 1 10); do + "./host_alloc_${label}" > "run_${label}.log" 2>&1 + rc=$? + [ "${i}" -eq 1 ] && tail -25 "run_${label}.log" + if [ "${rc}" -ge 128 ]; then crashes=$((crashes + 1)); elif [ "${rc}" -ne 0 ]; then fails=$((fails + 1)); fi + echo "${label} run ${i}: rc=${rc}" + done + echo "RESULT ${label}: 10 runs, ${crashes} crashes, ${fails} other failures" +} + +PKGS=("cuda-version=${CUDA_MM}" cuda-cudart-dev dlpack gtest "gxx_linux-64=14.*" "sysroot_linux-64=2.28") + +rapids-logger "Baseline: nightly libcuvs without the fix" +rapids-mamba-retry create -y -n cbase -c rapidsai-nightly -c conda-forge "libcuvs=26.12" "${PKGS[@]}" +set +u; conda activate cbase; set -u +conda list | grep -E "^libcuvs |^librmm " +build_and_run baseline +set +u; conda deactivate; set -u + +rapids-logger "Fixed: this PR's libcuvs" +rapids-mamba-retry create -y -n cfix -c "${CPP_CHANNEL}" -c rapidsai-nightly -c conda-forge \ + "libcuvs=26.12" "libcuvs-tests=26.12" "${PKGS[@]}" +set +u; conda activate cfix; set -u +conda list | grep -E "^libcuvs |^libcuvs-tests |^librmm " +build_and_run fixed +GT="${CONDA_PREFIX}/bin/gtests/libcuvs/cuvs_c_host_alloc_test" +ls -la "${GT}" || true +"${GT}" --gtest_repeat=100 2>&1 | tail -15 +echo "installed cuvs_c_host_alloc_test rc=${PIPESTATUS[0]}" diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index bb96c9115b..6e4cfe163d 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -42,6 +42,7 @@ #include #include +#include #include #include #include @@ -802,7 +803,8 @@ 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; + // Shared by all threads that build balanced k-means concurrently, so it must be atomic. + static std::atomic i_primes{0}; auto average = static_cast(n_rows) / static_cast(n_clusters); auto lower_threshold = average * balance_lower_tolerance; @@ -839,8 +841,7 @@ auto adjust_centers(const raft::resources& handle, IdxT ofst; do { - i_primes = (i_primes + 1) % kPrimes.size(); - ofst = kPrimes[i_primes]; + ofst = kPrimes[(i_primes.fetch_add(1, std::memory_order_relaxed) + 1) % kPrimes.size()]; } while (n_rows % ofst == 0); rmm::device_uvector receiver_clusters(n_pairs, stream, device_memory); diff --git a/rust/cuvs/src/lib.rs b/rust/cuvs/src/lib.rs index e5748e83fe..69dbd84e86 100644 --- a/rust/cuvs/src/lib.rs +++ b/rust/cuvs/src/lib.rs @@ -22,6 +22,9 @@ pub mod version; #[cfg(test)] pub(crate) mod test_utils; +#[cfg(test)] +mod thread_stress; + pub use dlpack::{AsDlTensor, AsDlTensorMut, DLPackError, DLTensorView, DLTensorViewMut, DType}; pub use error::LibraryError; pub use resources::Resources; diff --git a/rust/cuvs/src/thread_stress.rs b/rust/cuvs/src/thread_stress.rs new file mode 100644 index 0000000000..58cf14af69 --- /dev/null +++ b/rust/cuvs/src/thread_stress.rs @@ -0,0 +1,101 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! CI-only stress tests: many host threads, each with its own `Resources`, +//! building and searching indexes concurrently (issue #2233). + +use crate::distance::DistanceType; +use crate::neighbors::{brute_force, cagra, ivf_flat, ivf_pq}; +use crate::resources::Resources; +use crate::test_utils::DeviceTensor; +use ndarray::s; +use ndarray_rand::RandomExt; +use ndarray_rand::rand_distr::Uniform; + +fn env_usize(name: &str, default: usize) -> usize { + std::env::var(name).ok().and_then(|v| v.parse().ok()).unwrap_or(default) +} + +fn one_round(kind: usize) { + let res = Resources::new().unwrap(); + let n_datapoints = 1024; + let n_features = 16; + let n_queries = 4; + let k = 10; + let dataset = ndarray::Array::::random( + (n_datapoints, n_features), + Uniform::new(0., 1.0).unwrap(), + ); + let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); + let queries = dataset.slice(s![0..n_queries, ..]).to_owned(); + let queries = DeviceTensor::from_host(&res, &queries).unwrap(); + let mut distances = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); + match kind % 4 { + 0 => { + let bp = ivf_flat::IndexParams::builder().n_lists(64).build().unwrap(); + let index = ivf_flat::Index::build(&res, &bp, &dataset_device).unwrap(); + let sp = ivf_flat::SearchParams::builder().build().unwrap(); + let mut neighbors = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); + index.search(&res, &sp, &queries, &mut neighbors, &mut distances).unwrap(); + let mut nh = ndarray::Array::::zeros((n_queries, k)); + neighbors.copy_to_host(&res, &mut nh).unwrap(); + assert_eq!(nh[[0, 0]], 0); + } + 1 => { + let bp = ivf_pq::IndexParams::builder().n_lists(64).build().unwrap(); + let index = ivf_pq::Index::build(&res, &bp, &dataset_device).unwrap(); + let sp = ivf_pq::SearchParams::builder().build().unwrap(); + let mut neighbors = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); + index.search(&res, &sp, &queries, &mut neighbors, &mut distances).unwrap(); + let mut nh = ndarray::Array::::zeros((n_queries, k)); + neighbors.copy_to_host(&res, &mut nh).unwrap(); + } + 2 => { + let bp = cagra::IndexParams::builder().build().unwrap(); + let index = cagra::Index::build(&res, &bp, &dataset_device).unwrap(); + let sp = cagra::SearchParams::builder().build().unwrap(); + let mut neighbors = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); + index.search(&res, &sp, &queries, &mut neighbors, &mut distances).unwrap(); + let mut nh = ndarray::Array::::zeros((n_queries, k)); + neighbors.copy_to_host(&res, &mut nh).unwrap(); + assert_eq!(nh[[0, 0]], 0); + } + _ => { + let index = + brute_force::Index::build(&res, DistanceType::L2Expanded, &dataset_device).unwrap(); + let mut neighbors = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); + index.search(&res, &queries, &mut neighbors, &mut distances).unwrap(); + let mut nh = ndarray::Array::::zeros((n_queries, k)); + neighbors.copy_to_host(&res, &mut nh).unwrap(); + assert_eq!(nh[[0, 0]], 0); + } + } +} + +fn stress(kinds: &[usize]) { + let threads = env_usize("CUVS_STRESS_THREADS", 16); + let iters = env_usize("CUVS_STRESS_ITERS", 20); + std::thread::scope(|scope| { + for t in 0..threads { + scope.spawn(move || { + for i in 0..iters { + one_round(kinds[(t + i) % kinds.len()]); + } + }); + } + }); +} + +#[test] +#[ignore] +fn stress_ivf_flat_threads() { + stress(&[0]); +} + +#[test] +#[ignore] +fn stress_mixed_threads() { + stress(&[0, 1, 2, 3]); +}