Merge remote-tracking branch 'upstream/branch-23.12' into nightly_ver…

…sions
rapidsai · Oct 30, 2023 · e055bf8 · e055bf8
2 parents 5a54f71 + 9ad76fa
commit e055bf8
Show file tree

Hide file tree

Showing 24 changed files with 327 additions and 176 deletions.
diff --git a/ci/test_wheel_raft_dask.sh b/ci/test_wheel_raft_dask.sh
@@ -12,7 +12,7 @@ RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels
 python -m pip install --no-deps ./local-pylibraft-dep/pylibraft*.whl
 
 # Always install latest dask for testing
-python -m pip install git+https://github.com/dask/dask.git@2023.9.2 git+https://github.com/dask/distributed.git@2023.9.2 git+https://github.com/rapidsai/[email protected]
+python -m pip install git+https://github.com/dask/dask.git@main git+https://github.com/dask/distributed.git@main git+https://github.com/rapidsai/[email protected]
 
 # echo to expand wildcard before adding `[extra]` requires for pip
 python -m pip install $(echo ./dist/raft_dask*.whl)[test]

diff --git a/conda/environments/all_cuda-118_arch-aarch64.yaml b/conda/environments/all_cuda-118_arch-aarch64.yaml
@@ -19,10 +19,10 @@ dependencies:
 - cupy>=12.0.0
 - cxx-compiler
 - cython>=3.0.0
-- dask-core==2023.9.2
+- dask-core>=2023.9.2
 - dask-cuda==23.12.*
-- dask==2023.9.2
-- distributed==2023.9.2
+- dask>=2023.9.2
+- distributed>=2023.9.2
 - doxygen>=1.8.20
 - gcc_linux-aarch64=11.*
 - gmock>=1.13.0

diff --git a/conda/environments/all_cuda-118_arch-x86_64.yaml b/conda/environments/all_cuda-118_arch-x86_64.yaml
@@ -19,10 +19,10 @@ dependencies:
 - cupy>=12.0.0
 - cxx-compiler
 - cython>=3.0.0
-- dask-core==2023.9.2
+- dask-core>=2023.9.2
 - dask-cuda==23.12.*
-- dask==2023.9.2
-- distributed==2023.9.2
+- dask>=2023.9.2
+- distributed>=2023.9.2
 - doxygen>=1.8.20
 - gcc_linux-64=11.*
 - gmock>=1.13.0

diff --git a/conda/environments/all_cuda-120_arch-aarch64.yaml b/conda/environments/all_cuda-120_arch-aarch64.yaml
@@ -20,10 +20,10 @@ dependencies:
 - cupy>=12.0.0
 - cxx-compiler
 - cython>=3.0.0
-- dask-core==2023.9.2
+- dask-core>=2023.9.2
 - dask-cuda==23.12.*
-- dask==2023.9.2
-- distributed==2023.9.2
+- dask>=2023.9.2
+- distributed>=2023.9.2
 - doxygen>=1.8.20
 - gcc_linux-aarch64=11.*
 - gmock>=1.13.0

diff --git a/conda/environments/all_cuda-120_arch-x86_64.yaml b/conda/environments/all_cuda-120_arch-x86_64.yaml
@@ -20,10 +20,10 @@ dependencies:
 - cupy>=12.0.0
 - cxx-compiler
 - cython>=3.0.0
-- dask-core==2023.9.2
+- dask-core>=2023.9.2
 - dask-cuda==23.12.*
-- dask==2023.9.2
-- distributed==2023.9.2
+- dask>=2023.9.2
+- distributed>=2023.9.2
 - doxygen>=1.8.20
 - gcc_linux-64=11.*
 - gmock>=1.13.0

diff --git a/conda/recipes/raft-dask/meta.yaml b/conda/recipes/raft-dask/meta.yaml
@@ -60,10 +60,10 @@ requirements:
     - cudatoolkit
     {% endif %}
     - {{ pin_compatible('cuda-version', max_pin='x', min_pin='x') }}
-    - dask ==2023.9.2
-    - dask-core ==2023.9.2
+    - dask >=2023.9.2
+    - dask-core >=2023.9.2
     - dask-cuda ={{ minor_version }}
-    - distributed ==2023.9.2
+    - distributed >=2023.9.2
     - joblib >=0.11
     - nccl >=2.9.9
     - pylibraft {{ version }}

diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt
@@ -106,10 +106,8 @@ if(RAFT_ANN_BENCH_USE_GGNN)
 endif()
 
 if(RAFT_ANN_BENCH_USE_FAISS)
-  # We need to ensure that faiss has all the conda
-  # information. So we currently use the very ugly
-  # hammer of `link_libraries` to ensure that all
-  # targets in this directory and the faiss directory
+  # We need to ensure that faiss has all the conda information. So we currently use the very ugly
+  # hammer of `link_libraries` to ensure that all targets in this directory and the faiss directory
   # will have the conda includes/link dirs
   link_libraries($<TARGET_NAME_IF_EXISTS:conda_env>)
   include(cmake/thirdparty/get_faiss.cmake)

diff --git a/cpp/bench/ann/src/common/ann_types.hpp b/cpp/bench/ann/src/common/ann_types.hpp
@@ -24,6 +24,11 @@
 
 namespace raft::bench::ann {
 
+enum Objective {
+  THROUGHPUT,  // See how many vectors we can push through
+  LATENCY      // See how fast we can push a vector through
+};
+
 enum class MemoryType {
   Host,
   HostMmap,
@@ -59,10 +64,17 @@ inline auto parse_memory_type(const std::string& memory_type) -> MemoryType
   }
 }
 
-struct AlgoProperty {
+class AlgoProperty {
+ public:
+  inline AlgoProperty() {}
+  inline AlgoProperty(MemoryType dataset_memory_type_, MemoryType query_memory_type_)
+    : dataset_memory_type(dataset_memory_type_), query_memory_type(query_memory_type_)
+  {
+  }
   MemoryType dataset_memory_type;
   // neighbors/distances should have same memory type as queries
   MemoryType query_memory_type;
+  virtual ~AlgoProperty() = default;
 };
 
 class AnnBase {
@@ -79,7 +91,8 @@ template <typename T>
 class ANN : public AnnBase {
  public:
   struct AnnSearchParam {
-    virtual ~AnnSearchParam() = default;
+    Objective metric_objective = Objective::LATENCY;
+    virtual ~AnnSearchParam()  = default;
     [[nodiscard]] virtual auto needs_dataset() const -> bool { return false; };
   };