diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c2063765..44bd1be9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -48,6 +48,7 @@ jobs: - name: Check benchmark tools on Python 3.9 run: | python3 -m py_compile \ + tools/benchmark_gpu_training.py \ tools/convert_ann_benchmarks.py \ tools/verify_binary_artifact.py diff --git a/README.md b/README.md index 953c53e3..7b57c49e 100644 --- a/README.md +++ b/README.md @@ -47,6 +47,8 @@ retention is bounded and the shared cold cache is sharded for concurrent hits. warm-up, Rust, C, C++, Java, Python, and metadata filter pushdown. - [Development and benchmarks](docs/development.html): workspace layout, build and test commands, ANN benchmarks, and storage compatibility checks. +- [Experimental GPU training and construction](docs/GPU_TRAINING.md): train IVF-SQ + centers, assign and encode vectors with optional CUDA adapters, and keep CPU-readable indexes. - [Storage format specification](core/STORAGE_FORMAT.md): normative v1 binary layout and compatibility policy. diff --git a/core/src/coarse.rs b/core/src/coarse.rs index 3e0dc0a5..5c3d6159 100644 --- a/core/src/coarse.rs +++ b/core/src/coarse.rs @@ -61,6 +61,10 @@ impl CoarseAssignment { self.approximate_enabled } + pub(crate) fn is_exact(&self, d: usize, nlist: usize) -> bool { + !self.approximate_enabled || !use_approximate_assignment(d, nlist) + } + #[cfg(test)] pub(crate) fn build_attempted(&self) -> bool { self.build_attempted diff --git a/core/src/index.rs b/core/src/index.rs index 6b3e491a..cebab279 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -1383,6 +1383,55 @@ impl VectorIndexTrainer { Ok(self) } + /// Freeze an IVF-SQ training sample for an external centroid trainer. + /// + /// Consumes the trainer on success or failure. The returned state owns the + /// preprocessed sample; no GPU runtime is required by the Rust library. + /// `nlist` has already been resolved from the configuration, not sample size. + pub fn prepare_training(self) -> io::Result { + let VectorIndexWriter::IvfSq(index) = self.writer else { + return Err(invalid_input( + "external centroid training currently requires IVF-SQ", + )); + }; + if self.training_vector_count == 0 { + return Err(invalid_input("no training vectors added")); + } + let calibration_data = { + let raw = self.training_data; + match index.preprocess_vectors(&raw, self.training_vector_count) { + std::borrow::Cow::Borrowed(_) => raw, + std::borrow::Cow::Owned(processed) => processed, + } + }; + validate_vectors( + &calibration_data, + self.training_vector_count, + index.d, + "preprocessed training data", + )?; + let max_n = index + .nlist + .checked_mul(self.ivf_training.max_points_per_centroid) + .ok_or_else(|| invalid_input("training sample limit overflows usize"))?; + let kmeans_data = (self.training_vector_count > max_n).then(|| { + crate::kmeans::subsample( + &calibration_data, + self.training_vector_count, + index.d, + max_n, + &mut StdRng::seed_from_u64(self.ivf_training.seed), + ) + }); + Ok(PreparedIvfSqTraining { + index, + calibration_data, + kmeans_data, + config: self.ivf_training, + vectors_seen: self.training_vectors_seen, + }) + } + pub fn finish(mut self) -> io::Result { if self.training_vector_count == 0 || self.training_data.is_empty() { return Err(invalid_input("no training vectors added")); @@ -1397,6 +1446,111 @@ impl VectorIndexTrainer { } } +/// CPU reference algorithms for externally trained IVF centers. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CpuIvfTrainingAlgorithm { + /// The existing CPU strategy (hierarchical for large nlist). + Auto, + /// Flat Lloyd iterations, optionally from supplied initial centers. + Lloyd, +} + +/// Immutable, owned IVF-SQ training input and residual-calibration state. +/// +/// Cosine samples are normalized once; L2 and IP samples are unchanged. All +/// centroid training uses squared L2, including for cosine and IP indexes. +/// Finishing installs centers verbatim and recalibrates SQ on the CPU. It does +/// not retrain the centers or change the index file format. +pub struct PreparedIvfSqTraining { + index: IVFSQIndex, + calibration_data: Vec, + kmeans_data: Option>, + config: KMeansConfig, + vectors_seen: usize, +} + +impl PreparedIvfSqTraining { + pub fn dimension(&self) -> usize { + self.index.d + } + pub fn nlist(&self) -> usize { + self.index.nlist + } + pub fn metric(&self) -> MetricType { + self.index.metric + } + pub fn vectors_seen(&self) -> usize { + self.vectors_seen + } + pub fn calibration_vector_count(&self) -> usize { + self.calibration_data.len() / self.dimension() + } + pub fn config(&self) -> &KMeansConfig { + &self.config + } + + /// Effective K-means input after both reservoir and per-centroid caps. + /// SQ calibration retains the full reservoir, as in the original trainer. + pub fn sample(&self) -> &[f32] { + self.kmeans_data + .as_deref() + .unwrap_or(&self.calibration_data) + } + + pub fn fit_centroids_cpu( + &self, + algorithm: CpuIvfTrainingAlgorithm, + initial_centroids: Option<&[f32]>, + ) -> io::Result> { + if let Some(initial) = initial_centroids { + validate_vectors(initial, self.nlist(), self.dimension(), "initial centroids")?; + if algorithm != CpuIvfTrainingAlgorithm::Lloyd { + return Err(invalid_input( + "initial centroids require the Lloyd algorithm", + )); + } + } + let centers = match algorithm { + CpuIvfTrainingAlgorithm::Auto => crate::kmeans::kmeans_train( + &self.config, + &self.calibration_data, + self.calibration_vector_count(), + self.dimension(), + self.nlist(), + ), + CpuIvfTrainingAlgorithm::Lloyd => crate::kmeans::kmeans_train_with_init( + &self.config, + self.sample(), + self.sample().len() / self.dimension(), + self.dimension(), + self.nlist(), + initial_centroids, + ), + }; + validate_vectors( + ¢ers, + self.nlist(), + self.dimension(), + "trained centroids", + )?; + Ok(centers) + } + + /// Consumes the prepared state on success or failure. + pub fn finish_with_ivf_centroids( + mut self, + centroids: Vec, + ) -> io::Result { + validate_vectors(¢roids, self.nlist(), self.dimension(), "IVF centroids")?; + self.index.set_quantizer_centroids(centroids); + self.index + .train_sq_from_processed(&self.calibration_data, self.calibration_vector_count()); + Ok(VectorIndexTraining { + inner: VectorIndexWriter::IvfSq(self.index), + }) + } +} + pub struct VectorIndexTraining { inner: VectorIndexWriter, } @@ -1411,6 +1565,20 @@ impl VectorIndexTraining { } } +/// Owned snapshot for external IVF-SQ encoders. Bounds are row-major [nlist, dimension]. +/// Encoded data must use these exact centers, bounds and metric preprocessing. +pub struct IvfSqEncodingModel { + pub dimension: usize, + pub nlist: usize, + pub metric: MetricType, + pub exact_assignment: bool, + pub centroids: Vec, + pub mins: Vec, + pub maxs: Vec, + /// Native f32 arithmetic boundary; 1 denotes entirely scalar encoding. + pub encoding_vector_width: usize, +} + pub enum VectorIndexWriter { IvfFlat(IVFFlatIndex), IvfSq(IVFSQIndex), @@ -1545,6 +1713,118 @@ impl VectorIndexWriter { Ok(()) } + fn ivf_sq(&self) -> io::Result<&IVFSQIndex> { + match self { + Self::IvfSq(index) => Ok(index), + _ => Err(invalid_input("external encoding requires IVF-SQ")), + } + } + + pub fn ivf_sq_encoding_model(&self) -> io::Result { + let index = self.ivf_sq()?; + validate_vectors( + index.quantizer_centroids(), + index.nlist, + index.d, + "IVF centroids", + )?; + let mins = (0..index.nlist) + .flat_map(|list| index.list_sq(list).mins.iter().copied()) + .collect::>(); + let maxs = (0..index.nlist) + .flat_map(|list| index.list_sq(list).maxs.iter().copied()) + .collect::>(); + validate_vectors(&mins, index.nlist, index.d, "SQ minima")?; + validate_vectors(&maxs, index.nlist, index.d, "SQ maxima")?; + if mins.iter().zip(&maxs).any(|(min, max)| min > max) { + return Err(invalid_input("SQ minima exceed maxima")); + } + Ok(IvfSqEncodingModel { + dimension: index.d, + nlist: index.nlist, + metric: index.metric, + exact_assignment: index.uses_exact_assignment(), + centroids: index.quantizer_centroids().to_vec(), + mins, + maxs, + encoding_vector_width: crate::sq::residual_encoding_vector_width(), + }) + } + + /// Apply the same preprocessing as native add, without changing the writer. + pub fn preprocess_ivf_sq_vectors<'a>( + &self, + data: &'a [f32], + n: usize, + ) -> io::Result> { + let index = self.ivf_sq()?; + validate_vectors(data, n, index.d, "vector data")?; + let processed = index.preprocess_vectors(data, n); + validate_vectors(&processed, n, index.d, "preprocessed vector data")?; + Ok(processed) + } + + fn validate_external_assignments( + &self, + ids: &[i64], + lists: &[u32], + n: usize, + ) -> io::Result<()> { + let index = self.ivf_sq()?; + validate_positive(n, "vector count")?; + if ids.len() != n || lists.len() != n { + return Err(invalid_input( + "ID and partition counts must match vector count", + )); + } + if lists.iter().any(|&list| list as usize >= index.nlist) { + return Err(invalid_input("partition ID must be smaller than nlist")); + } + Ok(()) + } + + /// Add raw vectors with external partition IDs; native preprocessing/SQ encoding is retained. + /// Validation errors leave the writer unchanged. Callers must use the writer's centers. + pub fn add_preassigned_vectors( + &mut self, + ids: &[i64], + data: &[f32], + lists: &[u32], + n: usize, + ) -> io::Result<()> { + self.validate_external_assignments(ids, lists, n)?; + validate_vectors(data, n, self.dimension(), "vector data")?; + if let Self::IvfSq(index) = self { + index.add_preassigned(data, ids, lists, n); + } + Ok(()) + } + + /// Add external row-major SQ8 codes using this writer's encoding model. + /// This validates shapes and partition IDs, not the caller's quantization algorithm. + /// Validation errors leave the writer unchanged. + pub fn add_encoded_vectors( + &mut self, + ids: &[i64], + codes: &[u8], + lists: &[u32], + n: usize, + ) -> io::Result<()> { + self.validate_external_assignments(ids, lists, n)?; + let len = n + .checked_mul(self.dimension()) + .ok_or_else(|| invalid_input("code length overflows usize"))?; + if codes.len() != len { + return Err(invalid_input( + "SQ8 code length must equal vector count * dimension", + )); + } + if let Self::IvfSq(index) = self { + index.add_encoded(ids, lists, codes); + } + Ok(()) + } + pub fn write(&mut self, out: &mut dyn SeekWrite) -> io::Result<()> { match self { Self::IvfFlat(index) => write_ivfflat_index(index, out), diff --git a/core/src/ivfsq.rs b/core/src/ivfsq.rs index 58e68be3..defbcd22 100644 --- a/core/src/ivfsq.rs +++ b/core/src/ivfsq.rs @@ -57,6 +57,10 @@ impl IVFSQIndex { &self.quantizer_centroids } + pub(crate) fn uses_exact_assignment(&self) -> bool { + self.coarse_assignment.is_exact(self.d, self.nlist) + } + /// Enables automatic Vamana coarse assignment for large centroid matrices. /// Disable it to keep vector assignment exact. pub(crate) fn set_approximate_coarse_assignment(&mut self, enabled: bool) { @@ -89,14 +93,24 @@ impl IVFSQIndex { let processed = self.preprocess_vectors(data, n); self.quantizer_centroids = kmeans::kmeans_train(config, &processed, n, self.d, self.nlist); self.coarse_assignment.reset(); + self.train_sq_from_processed(&processed, n); + } + + /// Calibrate residual SQ bounds using already installed IVF centers. + /// `processed` must follow this index's metric preprocessing contract. + pub(crate) fn train_sq_from_processed(&mut self, processed: &[f32], n: usize) { + // TODO: Allow externally computed calibration assignments (e.g. GPU) to + // avoid repeating the costly CPU search at large nlist. Preserve metric + // preprocessing, assignment policy, and residual SQ bounds; validate + // encoding/recall parity and end-to-end gains before changing this path. let list_ids = self.coarse_assignment.assign( - &processed, + processed, n, &self.quantizer_centroids, self.nlist, self.d, ); - self.train_list_sqs(&processed, &list_ids); + self.train_list_sqs(processed, &list_ids); } pub fn add(&mut self, data: &[f32], ids: &[i64], n: usize) { @@ -108,8 +122,19 @@ impl IVFSQIndex { self.nlist, self.d, ); + self.append_preassigned(&processed, ids, &list_ids, n); + } + + /// Inputs and partition IDs are validated by VectorIndexWriter. + pub(crate) fn add_preassigned(&mut self, data: &[f32], ids: &[i64], lists: &[u32], n: usize) { + let processed = self.preprocess_vectors(data, n); + let lists = lists.iter().map(|&list| list as usize).collect::>(); + self.append_preassigned(&processed, ids, &lists, n); + } + + fn append_preassigned(&mut self, processed: &[f32], ids: &[i64], list_ids: &[usize], n: usize) { let mut list_rows = vec![Vec::new(); self.nlist]; - for (row, list_id) in list_ids.into_iter().enumerate() { + for (row, &list_id) in list_ids.iter().enumerate() { list_rows[list_id].push(row); } @@ -126,7 +151,7 @@ impl IVFSQIndex { .enumerate() .for_each(|(list_id, ((list_ids, list_codes), rows))| { append_encoded_rows( - &processed, + processed, ids, &rows, d, @@ -144,7 +169,7 @@ impl IVFSQIndex { .enumerate() { append_encoded_rows( - &processed, + processed, ids, &rows, d, @@ -157,6 +182,27 @@ impl IVFSQIndex { } } + /// Append externally encoded, row-major SQ8 codes without assigning or encoding again. + pub(crate) fn add_encoded(&mut self, ids: &[i64], lists: &[u32], codes: &[u8]) { + let mut rows = vec![Vec::new(); self.nlist]; + for (row, &list) in lists.iter().enumerate() { + rows[list as usize].push(row); + } + let d = self.d; + self.ids + .par_iter_mut() + .zip(self.codes.par_iter_mut()) + .zip(rows) + .for_each(|((out_ids, out_codes), rows)| { + out_ids.reserve(rows.len()); + out_codes.reserve(rows.len() * d); + for row in rows { + out_ids.push(ids[row]); + out_codes.extend_from_slice(&codes[row * d..(row + 1) * d]); + } + }); + } + pub fn total_vectors(&self) -> usize { self.ids.iter().map(Vec::len).sum() } diff --git a/core/src/kmeans.rs b/core/src/kmeans.rs index 0eac2663..aa5791b3 100644 --- a/core/src/kmeans.rs +++ b/core/src/kmeans.rs @@ -1113,7 +1113,13 @@ fn weighted_kmeans_train( centroids } -fn subsample(data: &[f32], n: usize, d: usize, target_n: usize, rng: &mut StdRng) -> Vec { +pub(crate) fn subsample( + data: &[f32], + n: usize, + d: usize, + target_n: usize, + rng: &mut StdRng, +) -> Vec { let mut indices: Vec = (0..n).collect(); for i in 0..target_n { let j = rng.gen_range(i..n); diff --git a/core/src/sq.rs b/core/src/sq.rs index ffe58be0..34fcc0e6 100644 --- a/core/src/sq.rs +++ b/core/src/sq.rs @@ -953,6 +953,21 @@ unsafe fn update_bounds_batch_neon( } } +/// Vectorized dimensions multiply by a precomputed scale; scalar tails divide. +/// External encoders use this to preserve the native f32 rounding order. +pub(crate) fn residual_encoding_vector_width() -> usize { + #[cfg(target_arch = "aarch64")] + { + return 4; + } + #[cfg(target_arch = "x86_64")] + if is_x86_feature_detected!("avx2") { + return 8; + } + #[allow(unreachable_code)] + 1 +} + fn encode_residual( vector: &[f32], offset: &[f32], diff --git a/core/tests/external_build.rs b/core/tests/external_build.rs new file mode 100644 index 00000000..3a38d4db --- /dev/null +++ b/core/tests/external_build.rs @@ -0,0 +1,115 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use paimon_vindex_core::index::{VectorIndexConfig, VectorIndexTrainer, VectorIndexWriter}; +use paimon_vindex_core::io::PosWriter; +use std::collections::HashMap; + +fn writer(metric: &str, d: usize, data: &[f32]) -> VectorIndexWriter { + let config = VectorIndexConfig::from_options(&HashMap::from([ + ("index.type".into(), "ivf_sq".into()), + ("dimension".into(), d.to_string()), + ("nlist".into(), "4".into()), + ("metric".into(), metric.into()), + ("ivf.coarse-assignment".into(), "exact".into()), + ])) + .unwrap(); + VectorIndexWriter::new(VectorIndexTrainer::train(config, data, data.len() / d).unwrap()) +} + +fn bytes(writer: &mut VectorIndexWriter) -> Vec { + let mut out = Vec::new(); + writer.write(&mut PosWriter::new(&mut out)).unwrap(); + out +} + +#[test] +fn externally_assigned_and_encoded_batches_preserve_index_bytes() { + for metric in ["l2", "cosine", "inner_product"] { + for d in [7, 8, 9] { + let n = 1025; + let data = (0..n * d) + .map(|i| ((i * 37 % 997) as f32).sin()) + .collect::>(); + let ids = (0..n).map(|i| 5000 - i as i64 * 17).collect::>(); + let mut original = writer(metric, d, &data); + let model = original.ivf_sq_encoding_model().unwrap(); + assert!(model.exact_assignment); + assert_eq!(model.mins.len(), 4 * d); + original.add_vectors(&ids, &data, n).unwrap(); + let VectorIndexWriter::IvfSq(index) = &original else { + unreachable!() + }; + let mut lists = vec![0u32; n]; + let mut codes = vec![0u8; n * d]; + for list in 0..4 { + for (local, &id) in index.ids[list].iter().enumerate() { + let row = ((5000 - id) / 17) as usize; + lists[row] = list as u32; + codes[row * d..(row + 1) * d] + .copy_from_slice(&index.codes[list][local * d..(local + 1) * d]); + } + } + let mut assigned = writer(metric, d, &data); + let mut encoded = writer(metric, d, &data); + for start in (0..n).step_by(127) { + let end = (start + 127).min(n); + assigned + .add_preassigned_vectors( + &ids[start..end], + &data[start * d..end * d], + &lists[start..end], + end - start, + ) + .unwrap(); + encoded + .add_encoded_vectors( + &ids[start..end], + &codes[start * d..end * d], + &lists[start..end], + end - start, + ) + .unwrap(); + } + assert_eq!(bytes(&mut original), bytes(&mut assigned), "{metric}, {d}"); + assert_eq!(bytes(&mut original), bytes(&mut encoded), "{metric}, {d}"); + } + } +} + +#[test] +fn invalid_external_batches_do_not_mutate_writer() { + let data = [0., 0., 1., 1., 2., 2., 3., 3.]; + let mut w = writer("l2", 2, &data); + w.add_vectors(&[99], &data[..2], 1).unwrap(); + let before = bytes(&mut w); + assert!(w + .add_preassigned_vectors(&[1, 2], &data[..4], &[0, 4], 2) + .is_err()); + assert!(w + .add_preassigned_vectors(&[1, 2], &[0., 0., f32::NAN, 0.], &[0, 1], 2) + .is_err()); + assert!(w + .add_preassigned_vectors(&[1], &data[..4], &[0, 1], 2) + .is_err()); + assert!(w.add_encoded_vectors(&[1, 2], &[0; 3], &[0, 1], 2).is_err()); + assert!(w + .add_encoded_vectors(&[1, 2], &[0; 4], &[0, u32::MAX], 2) + .is_err()); + assert!(w.add_encoded_vectors(&[], &[], &[], 0).is_err()); + assert_eq!(before, bytes(&mut w)); +} diff --git a/core/tests/external_training.rs b/core/tests/external_training.rs new file mode 100644 index 00000000..04560940 --- /dev/null +++ b/core/tests/external_training.rs @@ -0,0 +1,136 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use paimon_vindex_core::index::{ + CpuIvfTrainingAlgorithm, VectorIndexConfig, VectorIndexTrainer, VectorIndexTraining, + VectorIndexWriter, +}; +use paimon_vindex_core::io::PosWriter; +use std::collections::HashMap; + +fn config(metric: &str, d: usize, nlist: usize, max_points: usize) -> VectorIndexConfig { + VectorIndexConfig::from_options(&HashMap::from([ + ("index.type".into(), "ivf_sq".into()), + ("dimension".into(), d.to_string()), + ("nlist".into(), nlist.to_string()), + ("metric".into(), metric.into()), + ( + "ivf.train.max-points-per-centroid".into(), + max_points.to_string(), + ), + ])) + .unwrap() +} + +fn serialize(training: VectorIndexTraining, data: &[f32], n: usize) -> Vec { + let mut writer = VectorIndexWriter::new(training); + writer + .add_vectors(&(0..n as i64).collect::>(), data, n) + .unwrap(); + let mut bytes = Vec::new(); + writer.write(&mut PosWriter::new(&mut bytes)).unwrap(); + bytes +} + +#[test] +fn prepared_cpu_auto_preserves_original_index_bytes() { + for (metric, n, d, nlist, max_points) in [ + ("l2", 128, 8, 4, 256), + ("l2", 128, 8, 4, 2), + ("cosine", 128, 8, 4, 256), + ("inner_product", 128, 8, 4, 256), + ("l2", 600, 4, 300, 2), + ] { + let data = (0..n * d) + .map(|i| ((i * 37 % 997) as f32).sin()) + .collect::>(); + let original = + VectorIndexTrainer::train(config(metric, d, nlist, max_points), &data, n).unwrap(); + let prepared = VectorIndexTrainer::new(config(metric, d, nlist, max_points)) + .unwrap() + .add_training_vectors(&data, n) + .unwrap() + .prepare_training() + .unwrap(); + let centers = prepared + .fit_centroids_cpu(CpuIvfTrainingAlgorithm::Auto, None) + .unwrap(); + let external = prepared.finish_with_ivf_centroids(centers).unwrap(); + assert_eq!( + serialize(original, &data, n), + serialize(external, &data, n), + "{metric}, nlist={nlist}" + ); + } +} + +#[test] +fn supplied_centers_are_retained_and_sq_is_recalibrated() { + let data = [-2.0, -2.0, 10.0, 10.0, 12.0, 12.0]; + let centers = vec![0.0, 0.0, 10.0, 10.0, 20.0, 20.0]; + let prepared = VectorIndexTrainer::new(config("l2", 2, 3, 256)) + .unwrap() + .add_training_vectors(&data, 3) + .unwrap() + .prepare_training() + .unwrap(); + let training = prepared.finish_with_ivf_centroids(centers.clone()).unwrap(); + let VectorIndexWriter::IvfSq(index) = VectorIndexWriter::new(training) else { + unreachable!() + }; + assert_eq!(index.quantizer_centroids(), centers); + assert_eq!(index.sq.mins, [-2.0, -2.0]); + assert_eq!(index.sq.maxs, [2.0, 2.0]); + assert!(index.ids.iter().all(Vec::is_empty)); + assert!(index + .list_sqs + .iter() + .all(|sq| sq.mins == index.sq.mins && sq.maxs == index.sq.maxs)); +} + +#[test] +fn prepared_sampling_is_bounded_and_independent_of_batch_boundaries() { + let data = (0..70_000 * 2).map(|i| i as f32).collect::>(); + let prepare = |batch_rows: usize| { + let mut trainer = VectorIndexTrainer::new(config("l2", 2, 2, 3)).unwrap(); + for chunk in data.chunks(batch_rows * 2) { + trainer + .add_training_vectors_mut(chunk, chunk.len() / 2) + .unwrap(); + } + trainer.prepare_training().unwrap() + }; + let a = prepare(70_000); + let b = prepare(511); + assert_eq!(a.sample(), b.sample()); + assert_eq!(a.sample().len(), 6 * 2); + assert_eq!(a.calibration_vector_count(), 65_536); + assert_eq!(a.vectors_seen(), 70_000); +} + +#[test] +fn external_centers_reject_wrong_shapes_and_non_finite_values() { + for centers in [vec![0.0], vec![f32::NAN; 4], vec![f32::INFINITY; 4]] { + let prepared = VectorIndexTrainer::new(config("l2", 2, 2, 256)) + .unwrap() + .add_training_vectors(&[0.0, 0.0, 1.0, 1.0], 2) + .unwrap() + .prepare_training() + .unwrap(); + assert!(prepared.finish_with_ivf_centroids(centers).is_err()); + } +} diff --git a/docs/GPU_TRAINING.md b/docs/GPU_TRAINING.md new file mode 100644 index 00000000..038ae4b3 --- /dev/null +++ b/docs/GPU_TRAINING.md @@ -0,0 +1,329 @@ + + +# Experimental GPU training and construction for IVF-SQ + +IVF-SQ can train its IVF centers with NVIDIA cuVS, calibrate SQ with Rust, +and optionally assign and encode full batches on GPU. Rust assembles and writes +the existing v1 index. Existing CPU Readers +can read these files. The Rust library and normal Python imports do not require +CUDA. The optional adapters are `paimon_vindex.gpu.CuvsKMeans` and +`paimon_vindex.gpu.CuvsIvfSqWriter`. + +This first implementation exposes preparation and external-center completion +in Rust, C and Python. The CUDA adapter is Python-only. It does not add GPU +search or a CUDA runtime to the Java/JNI or C++ wrapper APIs. + +GPU execution and speedups must be validated on a CUDA machine. CPU round-trip +tests and benchmark smoke runs do not establish GPU performance. + +## Build and run + +Build the matching native library from this checkout: + +```sh +cargo build --release -p paimon-vindex-ffi +python -m pip install -e 'python[test]' +export PAIMON_VINDEX_LIB_PATH="$PWD/target/release" +``` + +In the GPU worker environment, install matching CuPy and cuVS packages for its +CUDA version using the [NVIDIA installation guide](https://docs.nvidia.com/cuvs/installation). +The required API is `cuvs.cluster.kmeans.KMeansParams` / `fit`, with array +initialization, and `cuvs.common.Resources`. See the +[official K-means API](https://docs.nvidia.com/cuvs/api-reference/python-api-cluster-kmeans). +Package/driver installation is external to this library; no CUDA dependency is +installed by the normal Python package. The adapter reports installed versions. + +For example, install the optional dependencies in a separate virtual environment: + +```sh +python -m pip install 'cuvs-cu12==26.8.1' 'cupy-cuda12x[ctk]==14.2.0' +``` + +The `ctk` extra supplies CUDA user-space libraries. This command does not install +or update the host NVIDIA driver. + +```python +import numpy as np +from paimon_vindex import VectorIndexTrainer, VectorIndexWriter +from paimon_vindex.gpu import CuvsKMeans, CuvsIvfSqWriter + +vectors = np.load("base.npy", mmap_mode="r") # shape (N, dimension) +options = { + "index.type": "ivf_sq", + "dimension": str(vectors.shape[1]), + "nlist": "1024", + "expected-vector-count": str(len(vectors)), + "metric": "cosine", + "ivf.coarse-assignment": "exact", +} + +# Reuse one worker across successive jobs to amortize GPU initialization. +with CuvsKMeans(device=0) as gpu: + with VectorIndexTrainer.create(options) as trainer: + for start in range(0, len(vectors), 8192): + trainer.add_training_vectors(vectors[start:start + 8192]) + with trainer.prepare_training() as prepared: + print(prepared.info) + centers = gpu.fit(prepared) + print(gpu.last_run) + training = prepared.finish_with_ivf_centroids(centers) + + with VectorIndexWriter(training) as writer: + with CuvsIvfSqWriter(writer, device=0) as builder: + for start in range(0, len(vectors), 8192): + batch = vectors[start:start + 8192] + builder.add_vectors(np.arange(start, start + len(batch), dtype=np.int64), batch) + print(writer.ivf_sq_partition_sizes()) + with open("vectors.index", "wb") as output: + writer.write(output) +``` + +The prepared sample is owned by Rust. `prepared.sample` returns an independent, +read-only float32 copy, usable by another training library. External centers +must have shape `(nlist, dimension)` and finite values. The centers are installed +verbatim; Rust assigns calibration samples and recomputes pooled residual SQ +bounds without retraining them. + +`prepare_training()` consumes the Trainer on success or failure. Native +completion consumes the prepared state on success or failure; Python shape/type +checks that fail before calling native code leave it open. Both classes support +context managers. Sample copies remain valid after the native state closes. + +## Full-vector construction + +`CuvsIvfSqWriter` snapshots the writer's centers and per-list SQ bounds and +retains them on its selected device. Each `add_vectors(ids, data)` call uploads +one batch, predicts its nearest centers, computes residual SQ8 codes, and +returns partition IDs and uint8 codes. Rust validates their shapes and partition +IDs, appends them, and performs the existing ID sorting and blocked-code +serialization. No second CPU assignment or encoding is performed. + +- Set `encode=False` to use GPU assignment with the existing CPU SQ encoder. + Call `writer.add_vectors` directly to retain the complete CPU add path. +- The adapter borrows the writer: closing the adapter releases its device + references and resources, while the writer remains usable. CuPy's process-wide + allocator may cache freed buffers. Do not close the writer during a build. +- Use bounded batches; the adapter does not automatically split an oversized + call. Batch scratch is bounded by the caller's batch size. The native writer + still retains the complete encoded index in host memory until serialization. +- Cosine preprocessing uses the native Rust implementation. L2/IP inputs remain + unchanged. SQ8 encoding follows native clipping, round-half-up behavior and + f32 operation order, including scalar tail dimensions and constant bounds. + Its arithmetic preserves subnormal values independently of CuPy's default + flush-to-zero compilation mode. +- GPU assignment is exact squared L2. Configure `ivf.coarse-assignment=exact` + before training for consistent SQ calibration. If the default policy would + use approximate Vamana, the adapter rejects that model with an actionable error. +- GPU and CPU distance calculations can choose different centers near ties. + Check recall and partition distribution; bitwise identity of independently + assigned indexes is not guaranteed. +- Invalid input or backend results are rejected before append. A failed batch + leaves prior successful batches in the writer. Input validation errors can be + corrected and retried. CUDA failures may require recreating the worker or its + process. The adapter never substitutes a CPU build after a GPU failure. +- Closing releases the adapter's resource references even if CUDA synchronization + fails. Explicit `close()` reports that failure; context-manager cleanup preserves + an exception already raised by the operation. A closed adapter cannot be reused. + +`builder.last_run` reports the last successful batch's validation/preprocessing, +H2D, assignment, encoding, D2H and native append times. Timings synchronize the +device. Constructor setup is `initialization_seconds`; CUDA kernel compilation, +when needed, is included in the first encoding call. Imports of the module do +not load CUDA. GPU construction additionally requires CuPy's CUDA kernel +compilation support (NVRTC), supplied by the installation above. + +The backend-independent native interfaces are `writer.ivf_sq_encoding_model()`, +`writer.add_preassigned_vectors(ids, raw_vectors, partition_ids)` and +`writer.add_encoded_vectors(ids, uint8_codes, partition_ids)`, with corresponding +Rust and C functions. The latter validates representation and partition IDs; +external callers are responsible for using the exact model and preprocessing +that belong to that writer. Encoding model arrays are independent, read-only +copies with shape `(nlist, dimension)`. + +## Training contract and comparisons + +- `nlist` is resolved before preparation. With automatic nlist, supply the + **full** `expected-vector-count`, rather than the number of sampled rows. +- The existing deterministic reservoir keeps at most + `max(65536, 64 * nlist)` rows. The effective K-means sample also honors + `ivf.train.max-points-per-centroid`. Metadata reports `vectors_seen`, + `calibration_count` and `sample_count` separately. SQ calibration uses the + whole reservoir, including when the K-means sample is further capped. +- Cosine inputs are normalized once by Rust. L2 and IP inputs are unchanged. + All center fitting uses **squared L2**. Do not switch the external trainer + to inner-product or spherical K-means based only on the final search metric. +- The adapter uses float32 flat Lloyd iterations, at most the prepared state's + iteration budget (currently 25), a `1e-6` tolerance and one initial center set. + By default it samples initial rows with the prepared seed; callers may supply + `initial_centroids`. This is a different initialization from CPU Auto. +- CPU Auto retains the existing hierarchical strategy above 256 centers. + `prepared.fit_centroids_cpu("lloyd", initial_centroids=...)` offers a closer + algorithmic comparison. Matching initial centers and iteration budgets does + not guarantee bitwise-equal GPU results or identical convergence behavior. +- The GPU adapter requires at least `nlist` effective sample rows and raises + errors for missing CUDA dependencies, device failures or invalid results. + It never silently falls back to CPU. + +cuVS releases differ in exposed tiling parameters. By default the adapter uses +the library's tiling. On versions exposing the corresponding properties, use +`CuvsKMeans(batch_samples=16384, batch_centroids=1024)`. Unsupported explicit +options raise an error. Input, centers, initialization and temporary buffers +all consume device memory; sample bytes alone are not a peak-memory estimate. + +## Reproducible benchmark + +The harness compares CPU Auto, CPU Lloyd and cuVS. CPU Lloyd and cuVS use the +same initial rows from the same prepared sample. CPU Auto uses the existing +algorithm and initialization. Order rotates between repeats. Each run records +sample/center hashes, resolved parameters, stage timings and center artifacts. + +| Preset | Synthetic rows | Dimensions | nlist | Raw sample bytes | +|---|---:|---:|---:|---:| +| smoke | 2,048 | 16 | 8 | 128 KiB | +| baseline | 65,536 | 960 | 1,024 | 240 MiB | +| medium | 262,144 | 1,536 | 4,096 | 1.5 GiB | +| large | 1,048,576 | 1,536 | 16,384 | 6 GiB | + +Synthetic presets exercise training loads. They are not production recall +benchmarks, and they do not represent a larger full corpus. With `--base`, +actual rows and dimension come from the file; all rows enter the native +reservoir, and `nlist` comes from the preset. The large CPU Lloyd baseline can +take substantial time. Start with smoke to verify the environment. + +CPU-only smoke, including encoding and serialization: + +```sh +python tools/benchmark_gpu_training.py --preset smoke --synthetic \ + --backends cpu-auto cpu-lloyd --build-index --output-dir /tmp/ivfsq-cpu-smoke +``` + +On a CUDA host, include `cuvs` explicitly: + +```sh +python tools/benchmark_gpu_training.py --preset baseline --synthetic \ + --backends cpu-auto cpu-lloyd cuvs --threads 8 --repeats 3 \ + --output-dir /tmp/ivfsq-gpu-baseline +``` + +Repeat with `--preset medium` and `--preset large`, using a new output directory +each time. Existing output directories are rejected. The default measures +training only; `--build-index` also builds and writes every base row per run. + +Evaluate real data using `.fvecs`/`.ivecs` or float32/integer `.npy` files: + +```sh +python tools/benchmark_gpu_training.py --preset baseline \ + --base /data/gist1m/base.fvecs --queries /data/gist1m/query.fvecs \ + --neighbors /data/gist1m/ground_truth.ivecs --metric l2 \ + --backends cpu-auto cpu-lloyd cuvs --build-index --nprobe 64 \ + --output-dir /tmp/ivfsq-gist-gpu +``` + +To measure complete GPU construction, add `--build-backend cuvs-encode` and +`--coarse-assignment exact`. Use `--build-backend cuvs-assign` to isolate GPU +assignment while retaining CPU SQ encoding. Keep the same assignment policy +for the CPU reference. + +For a comparison with identical centers, reuse a saved center artifact: + +```sh +for builder in cpu cuvs-assign cuvs-encode; do + python tools/benchmark_gpu_training.py --preset baseline \ + --base /data/gist1m/base.fvecs --queries /data/gist1m/query.fvecs \ + --neighbors /data/gist1m/ground_truth.ivecs --metric l2 \ + --backends external --centroids /tmp/ivfsq-gist-gpu/cuvs-0-centroids.npy \ + --build-index --build-backend "$builder" --coarse-assignment exact \ + --threads 8 --repeats 3 --nprobe 64 --output-dir "/tmp/ivfsq-fixed-$builder" +done +``` + +`external` skips centroid fitting: its totals describe rebuilding with supplied +centers, not end-to-end model training. Use `cpu-auto` and `cuvs` for complete +training/build comparisons. The manifest records the external center file hash, +builder, device and batch size. + +Ground truth must match the supplied base, metric and zero-based row IDs. +The harness records actual empty/max/P95 partition sizes, Recall@K, first-batch +QPS/read bytes and sequential P95 after that batch. A fixed `nprobe` reveals +quality changes; tune `nprobe` to the same recall target before claiming a +query-performance advantage. No quality claim is produced without ground truth. + +`--coarse-assignment auto` preserves the CPU writer/calibrator's existing policy: +it can switch to approximate Vamana assignment when `dimension * nlist >= 1000000`. +Use `--coarse-assignment exact` to compare against exact nearest-center assignment. +This changes CPU SQ calibration and full-row assignment, not the GPU K-means +algorithm. Report this policy with the timing and recall results. Partition +diagnostics describe the written index, not the GPU trainer's cluster labels. + +Interpret timing fields as follows: + +- `prepare_seconds`: trainer creation, input ingestion/reservoir sampling and + preprocessing. Native-to-Python export for CPU Lloyd's initialization is + included in its `centroid_fit_seconds`. +- `centroid_fit_seconds`: centroid training, including required sample export + and initialization. GPU input/output copies and synchronization are included. + GPU `kmeans_seconds`, H2D and D2H are also reported separately. +- `sq_calibration_seconds`: CPU residual parameter training using the new centers. +- `training_seconds`: preparation + centroid fitting + SQ calibration. +- `gpu_setup_seconds`: imports, resource/device initialization for the first + GPU job. `training_including_setup_seconds` adds this cost. Inspect the first + run separately; a multi-run median can hide one-time startup costs. +- `add_seconds`: full-row addition using the chosen builder. Includes GPU builder + initialization/teardown, input validation, transfers, encoding and native + append. `gpu_build_stages` sums per-batch timings; `gpu_build_setup_seconds` + is already included in `add_seconds`. These are not extra costs to add again. + `gpu_build` records the device name and CuPy/cuVS versions even when centroid + training uses a different backend. +- `build_seconds`: training + full-row add + CPU serialization when requested. + Writing includes buffered flush, not an `fsync` durability guarantee. + +Source mapping/generation, diagnostic hashing, common-initialization diagnostics, +center artifact export and query evaluation are outside those timers. File +pages may be cold on the first run. GPU peak memory is explicitly **not +measured**; profile it on the target hardware before setting capacity limits. + +`manifest.json`, incremental `runs.jsonl`, `summary.json` and per-run centers +are written to the output directory. If CUDA fails, the process fails rather +than fabricating a CPU-backed GPU result. Preserve the tested package versions +and compare both training and full-build metrics when deciding whether to +enable GPU workers in production. + +## Verification + +```sh +cargo test -p paimon-vindex-core --test external_training --test external_build +python -m pytest python/tests/test_gpu_training.py python/tests/test_gpu_build.py +``` + +The CPU checks cover fixed-center SQ calibration, reservoir bounds, batch +invariance, validation, ownership, and byte-identical CPU reference files. +Actual GPU integration tests are explicit and fail if the requested CUDA +environment is unavailable: + +```sh +PAIMON_TEST_CUVS=1 python -m pytest python/tests/test_gpu_training.py python/tests/test_gpu_build.py +``` + +They train on GPU, reuse the worker, write v1 files with CPU, and check CPU +retrieval against exact neighbors for L2, cosine and inner product. +Construction tests also compare GPU and native SQ8 encoding at rounding/clipping +boundaries and subnormal ranges, exercise tail dimensions, and test held-out queries with partial +partition probing and more than 256 centers. Ordinary CPU CI covers malformed +backend labels, recovery, writer ownership, CUDA cleanup failures and atomic validation failures using +a fake CUDA interface; these checks do not replace actual GPU tests. diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs index 90b99160..6f5a72c6 100644 --- a/ffi/src/lib.rs +++ b/ffi/src/lib.rs @@ -19,10 +19,10 @@ use paimon_vindex_core::distance::MetricType; use paimon_vindex_core::index::{ - IvfPqBatchTableReuseMode, SearchWidth, VectorIndexConfig, VectorIndexMetadata, - VectorIndexReadPlan, VectorIndexReader, VectorIndexReaderOptions, VectorIndexTrainer, - VectorIndexTraining, VectorIndexWriter, VectorSearchParams, - DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES, + CpuIvfTrainingAlgorithm, IvfPqBatchTableReuseMode, PreparedIvfSqTraining, SearchWidth, + VectorIndexConfig, VectorIndexMetadata, VectorIndexReadPlan, VectorIndexReader, + VectorIndexReaderOptions, VectorIndexTrainer, VectorIndexTraining, VectorIndexWriter, + VectorSearchParams, DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES, }; use paimon_vindex_core::io::{ReadRequest, SeekRead, SeekReadCapabilities, SeekWrite}; use std::cell::RefCell; @@ -326,6 +326,34 @@ pub struct PaimonVindexTrainerHandle { inner: Option, } +pub struct PaimonVindexPreparedTrainingHandle { + inner: Option, +} + +/// Resolved, immutable IVF-SQ training parameters. Samples use squared L2 +/// clustering; `metric` identifies preprocessing and the final index metric. +#[repr(C)] +pub struct PaimonVindexPreparedTrainingInfo { + pub dimension: usize, + pub nlist: usize, + pub sample_count: usize, + pub calibration_count: usize, + pub vectors_seen: usize, + pub iterations: usize, + pub restarts: usize, + pub seed: u64, + pub metric: u32, +} + +#[repr(C)] +pub struct PaimonVindexIvfSqEncodingInfo { + pub dimension: usize, + pub nlist: usize, + pub metric: u32, + pub exact_assignment: u32, + pub encoding_vector_width: usize, +} + pub struct PaimonVindexTrainingHandle { inner: Option, } @@ -780,6 +808,166 @@ pub unsafe extern "C" fn paimon_vindex_trainer_add_training_vectors( }) } +/// Freezes an IVF-SQ sample and consumes the trainer, including on failure. +/// Call `paimon_vindex_trainer_free` afterwards. Prepared handles must be +/// serialized by the caller, just like trainer and writer handles. +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_trainer_prepare( + handle: *mut PaimonVindexTrainerHandle, +) -> *mut PaimonVindexPreparedTrainingHandle { + ffi_ptr(|| { + let trainer = unsafe { trainer_mut(handle) }? + .inner + .take() + .ok_or("trainer has already finished")?; + let prepared = trainer.prepare_training().map_err(|e| e.to_string())?; + Ok(Box::into_raw(Box::new( + PaimonVindexPreparedTrainingHandle { + inner: Some(prepared), + }, + ))) + }) +} + +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_prepared_training_free( + handle: *mut PaimonVindexPreparedTrainingHandle, +) { + if !handle.is_null() { + unsafe { + drop(Box::from_raw(handle)); + } + } +} + +unsafe fn prepared_ref<'a>( + handle: *const PaimonVindexPreparedTrainingHandle, +) -> Result<&'a PreparedIvfSqTraining, String> { + if handle.is_null() { + return Err("null prepared training handle".into()); + } + unsafe { &*handle } + .inner + .as_ref() + .ok_or_else(|| "prepared training has already finished".into()) +} + +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_prepared_training_info( + handle: *const PaimonVindexPreparedTrainingHandle, + out: *mut PaimonVindexPreparedTrainingInfo, +) -> c_int { + ffi_status(|| { + if out.is_null() { + return Err("info pointer is null".into()); + } + let p = unsafe { prepared_ref(handle) }?; + unsafe { + *out = PaimonVindexPreparedTrainingInfo { + dimension: p.dimension(), + nlist: p.nlist(), + sample_count: p.sample().len() / p.dimension(), + calibration_count: p.calibration_vector_count(), + vectors_seen: p.vectors_seen(), + iterations: p.config().niter, + restarts: p.config().nredo, + seed: p.config().seed, + metric: p.metric() as u32, + }; + } + Ok(()) + }) +} + +/// Copies preprocessed row-major f32 samples to caller-owned memory. `out_len` +/// must equal sample_count * dimension, in elements (not bytes). +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_prepared_training_copy_sample( + handle: *const PaimonVindexPreparedTrainingHandle, + out: *mut f32, + out_len: usize, +) -> c_int { + ffi_status(|| { + let p = unsafe { prepared_ref(handle) }?; + if out_len != p.sample().len() { + return Err("sample output length mismatch".into()); + } + unsafe { mut_slice(out, out_len, "sample output") }?.copy_from_slice(p.sample()); + Ok(()) + }) +} + +/// CPU reference trainer: algorithm 0 = existing Auto, 1 = flat Lloyd. +/// Optional initial centers are supported only for Lloyd; pass NULL, 0 to +/// omit. Input/output lengths are f32 elements, nlist * dimension. +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_prepared_training_fit_cpu( + handle: *const PaimonVindexPreparedTrainingHandle, + algorithm: u32, + initial: *const f32, + initial_len: usize, + out: *mut f32, + out_len: usize, +) -> c_int { + ffi_status(|| { + let p = unsafe { prepared_ref(handle) }?; + let expected = checked_len(p.nlist(), p.dimension(), "centroids")?; + if out_len != expected { + return Err("centroid output length mismatch".into()); + } + if out.is_null() { + return Err("centroid output pointer is null".into()); + } + let algorithm = match algorithm { + 0 => CpuIvfTrainingAlgorithm::Auto, + 1 => CpuIvfTrainingAlgorithm::Lloyd, + _ => return Err("unknown CPU training algorithm".into()), + }; + let initial = if initial_len == 0 { + None + } else { + if initial_len != expected { + return Err("initial centroid length mismatch".into()); + } + Some(unsafe { const_slice(initial, initial_len, "initial centroids") }?) + }; + let centers = p + .fit_centroids_cpu(algorithm, initial) + .map_err(|e| e.to_string())?; + unsafe { mut_slice(out, out_len, "centroid output") }?.copy_from_slice(¢ers); + Ok(()) + }) +} + +/// Consumes the prepared state, including on validation failure, but does not +/// free its handle. Installs centers and calibrates residual SQ on the CPU. +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_prepared_training_finish( + handle: *mut PaimonVindexPreparedTrainingHandle, + centers: *const f32, + centers_len: usize, +) -> *mut PaimonVindexTrainingHandle { + ffi_ptr(|| { + if handle.is_null() { + return Err("null prepared training handle".into()); + } + let p = unsafe { &mut *handle } + .inner + .take() + .ok_or("prepared training has already finished")?; + if centers_len != checked_len(p.nlist(), p.dimension(), "centroids")? { + return Err("IVF centroid length mismatch".into()); + } + let centers = unsafe { const_slice(centers, centers_len, "IVF centroids") }?.to_vec(); + let training = p + .finish_with_ivf_centroids(centers) + .map_err(|e| e.to_string())?; + Ok(Box::into_raw(Box::new(PaimonVindexTrainingHandle { + inner: Some(training), + }))) + }) +} + /// Finishes training and consumes the trainer's internal state, but does not free `handle`. /// Callers must still call `paimon_vindex_trainer_free(handle)` after this returns. #[no_mangle] @@ -837,6 +1025,171 @@ pub unsafe extern "C" fn paimon_vindex_writer_free(handle: *mut PaimonVindexWrit } } +/// Describe the immutable encoding model owned by an IVF-SQ writer. +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_writer_ivf_sq_encoding_info( + handle: *const PaimonVindexWriterHandle, + out: *mut PaimonVindexIvfSqEncodingInfo, +) -> c_int { + ffi_status(|| { + if out.is_null() { + return Err("encoding info pointer is null".into()); + } + let model = unsafe { writer_ref(handle) }? + .inner + .ivf_sq_encoding_model() + .map_err(|e| e.to_string())?; + unsafe { + *out = PaimonVindexIvfSqEncodingInfo { + dimension: model.dimension, + nlist: model.nlist, + metric: model.metric as u32, + exact_assignment: u32::from(model.exact_assignment), + encoding_vector_width: model.encoding_vector_width, + }; + } + Ok(()) + }) +} + +/// Copy centers and per-list SQ bounds. Each nonoverlapping output holds +/// exactly nlist * dimension f32 elements in row-major order. +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_writer_ivf_sq_copy_model( + handle: *const PaimonVindexWriterHandle, + centers: *mut f32, + mins: *mut f32, + maxs: *mut f32, + len: usize, +) -> c_int { + ffi_status(|| { + let model = unsafe { writer_ref(handle) }? + .inner + .ivf_sq_encoding_model() + .map_err(|e| e.to_string())?; + if len != model.centroids.len() { + return Err("encoding model output length mismatch".into()); + } + if centers.is_null() || mins.is_null() || maxs.is_null() { + return Err("encoding model output pointer is null".into()); + } + unsafe { + mut_slice(centers, len, "centers")?.copy_from_slice(&model.centroids); + mut_slice(mins, len, "mins")?.copy_from_slice(&model.mins); + mut_slice(maxs, len, "maxs")?.copy_from_slice(&model.maxs); + } + Ok(()) + }) +} + +/// Preprocess raw vectors without mutating the writer. Input and output must +/// not overlap. Output length is vector_count * dimension f32 elements. +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_writer_ivf_sq_preprocess( + handle: *const PaimonVindexWriterHandle, + data: *const f32, + vector_count: usize, + out: *mut f32, + out_len: usize, +) -> c_int { + ffi_status(|| { + let writer = &unsafe { writer_ref(handle) }?.inner; + let len = checked_len(vector_count, writer.dimension(), "vector data")?; + if out_len != len || out.is_null() { + return Err("preprocessing output shape or pointer is invalid".into()); + } + let input = unsafe { const_slice(data, len, "vector data") }?; + let processed = writer + .preprocess_ivf_sq_vectors(input, vector_count) + .map_err(|e| e.to_string())?; + unsafe { mut_slice(out, len, "preprocessed output") }?.copy_from_slice(&processed); + Ok(()) + }) +} + +/// Append raw vectors with caller-provided partition IDs. Uses native metric +/// preprocessing and SQ encoding. Validation failure does not append any rows. +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_writer_add_preassigned_vectors( + handle: *mut PaimonVindexWriterHandle, + ids: *const i64, + data: *const f32, + lists: *const u32, + vector_count: usize, +) -> c_int { + ffi_status(|| { + let writer = &mut unsafe { writer_mut(handle) }?.inner; + let len = checked_len(vector_count, writer.dimension(), "vector data")?; + let ids = unsafe { const_slice(ids, vector_count, "IDs") }?; + let data = unsafe { const_slice(data, len, "vector data") }?; + let lists = unsafe { const_slice(lists, vector_count, "partition IDs") }?; + writer + .add_preassigned_vectors(ids, data, lists, vector_count) + .map_err(|e| e.to_string()) + }) +} + +/// Append row-major SQ8 codes made with this writer's encoding model. +/// codes_len is vector_count * dimension bytes. Validation failure appends no rows. +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_writer_add_encoded_vectors( + handle: *mut PaimonVindexWriterHandle, + ids: *const i64, + codes: *const u8, + codes_len: usize, + lists: *const u32, + vector_count: usize, +) -> c_int { + ffi_status(|| { + let writer = &mut unsafe { writer_mut(handle) }?.inner; + if codes_len != checked_len(vector_count, writer.dimension(), "SQ8 codes")? { + return Err("SQ8 code length mismatch".into()); + } + let ids = unsafe { const_slice(ids, vector_count, "IDs") }?; + let codes = unsafe { const_slice(codes, codes_len, "SQ8 codes") }?; + let lists = unsafe { const_slice(lists, vector_count, "partition IDs") }?; + writer + .add_encoded_vectors(ids, codes, lists, vector_count) + .map_err(|e| e.to_string()) + }) +} + +/// IVF-SQ partition diagnostics after add. Query nlist with NULL, 0; otherwise +/// provide exactly nlist size_t entries. Does not alter the writer. +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_writer_ivf_sq_partition_sizes( + handle: *const PaimonVindexWriterHandle, + out: *mut usize, + out_len: usize, + out_nlist: *mut usize, +) -> c_int { + ffi_status(|| { + if out_nlist.is_null() { + return Err("nlist output pointer is null".into()); + } + let writer = unsafe { writer_ref(handle) }?; + let VectorIndexWriter::IvfSq(index) = &writer.inner else { + return Err("partition sizes currently require IVF-SQ".into()); + }; + unsafe { + *out_nlist = index.nlist; + } + if out.is_null() && out_len == 0 { + return Ok(()); + } + if out_len != index.nlist { + return Err("partition sizes output length mismatch".into()); + } + for (dst, ids) in unsafe { mut_slice(out, out_len, "partition sizes") }? + .iter_mut() + .zip(&index.ids) + { + *dst = ids.len(); + } + Ok(()) + }) +} + #[no_mangle] pub unsafe extern "C" fn paimon_vindex_writer_dimension( handle: *const PaimonVindexWriterHandle, diff --git a/python/paimon_vindex/__init__.py b/python/paimon_vindex/__init__.py index bba1ef07..3726ea1b 100644 --- a/python/paimon_vindex/__init__.py +++ b/python/paimon_vindex/__init__.py @@ -386,6 +386,123 @@ def __del__(self): pass +@dataclass(frozen=True) +class PreparedTrainingInfo: + dimension: int + nlist: int + sample_count: int + calibration_count: int + vectors_seen: int + iterations: int + restarts: int + seed: int + metric: str + + +class PreparedIvfSqTraining: + """Owned IVF-SQ sample for external centroid training. + + Samples are already preprocessed (normalized for cosine). Fit centers with + squared L2 even for IP/cosine indexes. Finishing calibrates residual SQ on + CPU and consumes the native state, including on native validation failure. + """ + + def __init__(self, handle): + self._native_handle_lock = _NativeHandleLock() + self._handle = handle + try: + info = _ffi.PaimonVindexPreparedTrainingInfo() + if lib.paimon_vindex_prepared_training_info(handle, ctypes.byref(info)) != 0: + _check_error("prepared training info failed") + fields = {name: getattr(info, name) for name, _ in info._fields_} + fields["metric"] = METRICS[info.metric] + self._info = PreparedTrainingInfo(**fields) + except Exception: + self.close() + raise + + def _require_open(self): + if not self._handle: + raise RuntimeError("PreparedIvfSqTraining is closed") + + @property + def info(self): + return self._info + + @property + def sample(self): + """Return an independent, read-only float32 copy of the effective sample.""" + with self._native_handle_lock: + self._require_open() + sample = np.empty((self.info.sample_count, self.info.dimension), dtype=np.float32) + rc = lib.paimon_vindex_prepared_training_copy_sample( + self._handle, sample.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), sample.size, + ) + if rc != 0: + _check_error("copy training sample failed") + sample.flags.writeable = False + return sample + + def fit_centroids_cpu(self, algorithm="auto", initial_centroids=None): + """Fit the existing CPU strategy or flat Lloyd for GPU comparisons.""" + if algorithm not in ("auto", "lloyd"): + raise ValueError("algorithm must be auto or lloyd") + initial = None + if initial_centroids is not None: + initial = _float32_matrix(initial_centroids, "initial_centroids") + if initial.shape != (self.info.nlist, self.info.dimension): + raise ValueError("initial centroid shape must be (nlist, dimension)") + centers = np.empty((self.info.nlist, self.info.dimension), dtype=np.float32) + with self._native_handle_lock: + self._require_open() + rc = lib.paimon_vindex_prepared_training_fit_cpu( + self._handle, int(algorithm == "lloyd"), + None if initial is None else initial.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + 0 if initial is None else initial.size, + centers.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), centers.size, + ) + if rc != 0: + _check_error("CPU centroid training failed") + return centers + + def finish_with_ivf_centroids(self, centroids): + centroids = _float32_matrix(centroids, "centroids") + if centroids.shape != (self.info.nlist, self.info.dimension): + raise ValueError("centroid shape must be (nlist, dimension)") + with self._native_handle_lock: + self._require_open() + handle = self._handle + training = lib.paimon_vindex_prepared_training_finish( + handle, centroids.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), centroids.size, + ) + lib.paimon_vindex_prepared_training_free(handle) + self._handle = None + if not training: + _check_error("finish external centroid training failed") + return VectorIndexTraining(training) + + def close(self): + with self._native_handle_lock: + if self._handle: + lib.paimon_vindex_prepared_training_free(self._handle) + self._handle = None + + def __enter__(self): + with self._native_handle_lock: + self._require_open() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + return False + + def __del__(self): + try: + self.close() + except Exception: + pass + + class VectorIndexTrainer: def __init__(self, options: Mapping[str, str]): self._native_handle_lock = _NativeHandleLock() @@ -464,6 +581,19 @@ def add_training_vectors(self, data): _check_error("add training vectors failed") return self + def prepare_training(self): + """Consume this IVF-SQ trainer and freeze its sample for external fitting.""" + with self._native_handle_lock: + self._require_open() + handle = self._handle + prepared = lib.paimon_vindex_trainer_prepare(handle) + lib.paimon_vindex_trainer_free(handle) + self._handle = None + self._closed = True + if not prepared: + _check_error("prepare training failed") + return PreparedIvfSqTraining(prepared) + def finish_training(self): with self._native_handle_lock: self._require_open() @@ -499,6 +629,28 @@ def __del__(self): pass +@dataclass(frozen=True) +class IvfSqEncodingModel: + """Independent model snapshot; arrays have shape (nlist, dimension).""" + dimension: int + nlist: int + metric: str + exact_assignment: bool + centroids: np.ndarray + mins: np.ndarray + maxs: np.ndarray + _encoding_vector_width: int + + +def _partition_ids(value): + array = np.asarray(value) + if array.ndim != 1 or not np.issubdtype(array.dtype, np.integer): + raise ValueError("partition_ids must be a one-dimensional integer array") + if np.any(array < 0) or np.any(array > np.iinfo(np.uint32).max): + raise ValueError("partition_ids must fit uint32") + return np.ascontiguousarray(array, dtype=np.uint32) + + class VectorIndexWriter: def __init__(self, training: VectorIndexTraining): if not isinstance(training, VectorIndexTraining): @@ -557,6 +709,95 @@ def add_vectors(self, ids, data): if rc != 0: _check_error("add_vectors failed") + def ivf_sq_partition_sizes(self): + """Return actual per-partition row counts for an IVF-SQ writer.""" + with self._native_handle_lock: + self._require_open() + nlist = ctypes.c_size_t() + rc = lib.paimon_vindex_writer_ivf_sq_partition_sizes( + self._handle, None, 0, ctypes.byref(nlist), + ) + if rc != 0: + _check_error("partition sizes failed") + sizes = np.empty(nlist.value, dtype=np.uintp) + rc = lib.paimon_vindex_writer_ivf_sq_partition_sizes( + self._handle, sizes.ctypes.data_as(ctypes.POINTER(ctypes.c_size_t)), + sizes.size, ctypes.byref(nlist), + ) + if rc != 0: + _check_error("partition sizes failed") + return sizes + + def ivf_sq_encoding_model(self): + """Copy centers and SQ bounds for an external encoder using this writer.""" + with self._native_handle_lock: + self._require_open() + info = _ffi.PaimonVindexIvfSqEncodingInfo() + if lib.paimon_vindex_writer_ivf_sq_encoding_info(self._handle, ctypes.byref(info)) != 0: + _check_error("encoding model info failed") + arrays = [np.empty((info.nlist, info.dimension), dtype=np.float32) for _ in range(3)] + pointers = [a.ctypes.data_as(ctypes.POINTER(ctypes.c_float)) for a in arrays] + if lib.paimon_vindex_writer_ivf_sq_copy_model(self._handle, *pointers, arrays[0].size) != 0: + _check_error("copy encoding model failed") + for array in arrays: + array.flags.writeable = False + return IvfSqEncodingModel(info.dimension, info.nlist, METRICS[info.metric], + bool(info.exact_assignment), *arrays, info.encoding_vector_width) + + def _preprocess_ivf_sq_vectors(self, data): + data = _float32_matrix(data, "data") + if data.shape[1] != self._dimension: + raise ValueError("data dimension does not match writer") + result = np.empty(data.shape, dtype=np.float32) + with self._native_handle_lock: + self._require_open() + rc = lib.paimon_vindex_writer_ivf_sq_preprocess( + self._handle, data.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), len(data), + result.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), result.size) + if rc != 0: + _check_error("preprocess vectors failed") + return result + + def add_preassigned_vectors(self, ids, data, partition_ids): + """Add raw vectors with external labels; native code preprocesses and encodes. + + Labels must refer to this writer's centers. Validation failures append no rows. + """ + data = _float32_matrix(data, "data") + ids = _int64_vector(ids, "ids") + lists = _partition_ids(partition_ids) + if data.shape != (len(ids), self._dimension) or len(lists) != len(ids): + raise ValueError("data, IDs and partition IDs must have matching shapes") + with self._native_handle_lock: + self._require_open() + rc = lib.paimon_vindex_writer_add_preassigned_vectors( + self._handle, ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), + data.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + lists.ctypes.data_as(ctypes.POINTER(ctypes.c_uint32)), len(ids)) + if rc != 0: + _check_error("add preassigned vectors failed") + + def add_encoded_vectors(self, ids, codes, partition_ids): + """Append row-major uint8 SQ8 codes made with this writer's encoding model. + + Validates shapes and labels; the caller is responsible for the encoding + algorithm and model. Validation failures append no rows. + """ + ids = _int64_vector(ids, "ids") + lists = _partition_ids(partition_ids) + codes = np.asarray(codes) + if codes.dtype != np.uint8 or codes.shape != (len(ids), self._dimension) or len(lists) != len(ids): + raise ValueError("codes must be uint8 (len(ids), dimension), with one partition ID per row") + codes = np.ascontiguousarray(codes) + with self._native_handle_lock: + self._require_open() + rc = lib.paimon_vindex_writer_add_encoded_vectors( + self._handle, ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), + codes.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8)), codes.size, + lists.ctypes.data_as(ctypes.POINTER(ctypes.c_uint32)), len(ids)) + if rc != 0: + _check_error("add encoded vectors failed") + def write(self, file): pos = 0 @@ -875,6 +1116,9 @@ def __del__(self): __all__ = [ + "IvfSqEncodingModel", + "PreparedIvfSqTraining", + "PreparedTrainingInfo", "IvfPqBatchTableReuseMode", "SearchParams", "VectorIndexMetadata", diff --git a/python/paimon_vindex/_ffi.py b/python/paimon_vindex/_ffi.py index 02cd900c..50c26f23 100644 --- a/python/paimon_vindex/_ffi.py +++ b/python/paimon_vindex/_ffi.py @@ -162,6 +162,15 @@ class PaimonVindexReaderOptions(Structure): ] +class PaimonVindexPreparedTrainingInfo(Structure): + _fields_ = [ + ("dimension", c_size_t), ("nlist", c_size_t), + ("sample_count", c_size_t), ("calibration_count", c_size_t), + ("vectors_seen", c_size_t), ("iterations", c_size_t), + ("restarts", c_size_t), ("seed", c_uint64), ("metric", c_uint32), + ] + + class PaimonVindexReadPlan(Structure): _fields_ = [ ("random_read_latency_nanos", c_uint64), @@ -202,6 +211,27 @@ class PaimonVindexReadPlan(Structure): lib.paimon_vindex_trainer_finish.argtypes = [c_void_p] lib.paimon_vindex_trainer_finish.restype = c_void_p +lib.paimon_vindex_trainer_prepare.argtypes = [c_void_p] +lib.paimon_vindex_trainer_prepare.restype = c_void_p +lib.paimon_vindex_prepared_training_free.argtypes = [c_void_p] +lib.paimon_vindex_prepared_training_free.restype = None +lib.paimon_vindex_prepared_training_info.argtypes = [ + c_void_p, POINTER(PaimonVindexPreparedTrainingInfo), +] +lib.paimon_vindex_prepared_training_info.restype = c_int +lib.paimon_vindex_prepared_training_copy_sample.argtypes = [ + c_void_p, POINTER(c_float), c_size_t, +] +lib.paimon_vindex_prepared_training_copy_sample.restype = c_int +lib.paimon_vindex_prepared_training_fit_cpu.argtypes = [ + c_void_p, c_uint32, POINTER(c_float), c_size_t, POINTER(c_float), c_size_t, +] +lib.paimon_vindex_prepared_training_fit_cpu.restype = c_int +lib.paimon_vindex_prepared_training_finish.argtypes = [ + c_void_p, POINTER(c_float), c_size_t, +] +lib.paimon_vindex_prepared_training_finish.restype = c_void_p + lib.paimon_vindex_training_free.argtypes = [c_void_p] lib.paimon_vindex_training_free.restype = None @@ -214,6 +244,27 @@ class PaimonVindexReadPlan(Structure): lib.paimon_vindex_writer_dimension.argtypes = [c_void_p, POINTER(c_size_t)] lib.paimon_vindex_writer_dimension.restype = c_int +class PaimonVindexIvfSqEncodingInfo(Structure): + _fields_ = [("dimension", c_size_t), ("nlist", c_size_t), ("metric", c_uint32), + ("exact_assignment", c_uint32), ("encoding_vector_width", c_size_t)] + + +lib.paimon_vindex_writer_ivf_sq_encoding_info.argtypes = [c_void_p, POINTER(PaimonVindexIvfSqEncodingInfo)] +lib.paimon_vindex_writer_ivf_sq_encoding_info.restype = c_int +lib.paimon_vindex_writer_ivf_sq_copy_model.argtypes = [c_void_p, POINTER(c_float), POINTER(c_float), POINTER(c_float), c_size_t] +lib.paimon_vindex_writer_ivf_sq_copy_model.restype = c_int +lib.paimon_vindex_writer_ivf_sq_preprocess.argtypes = [c_void_p, POINTER(c_float), c_size_t, POINTER(c_float), c_size_t] +lib.paimon_vindex_writer_ivf_sq_preprocess.restype = c_int +lib.paimon_vindex_writer_add_preassigned_vectors.argtypes = [c_void_p, POINTER(c_int64), POINTER(c_float), POINTER(c_uint32), c_size_t] +lib.paimon_vindex_writer_add_preassigned_vectors.restype = c_int +lib.paimon_vindex_writer_add_encoded_vectors.argtypes = [c_void_p, POINTER(c_int64), POINTER(c_uint8), c_size_t, POINTER(c_uint32), c_size_t] +lib.paimon_vindex_writer_add_encoded_vectors.restype = c_int + +lib.paimon_vindex_writer_ivf_sq_partition_sizes.argtypes = [ + c_void_p, POINTER(c_size_t), c_size_t, POINTER(c_size_t), +] +lib.paimon_vindex_writer_ivf_sq_partition_sizes.restype = c_int + lib.paimon_vindex_writer_add_vectors.argtypes = [ c_void_p, POINTER(c_int64), diff --git a/python/paimon_vindex/gpu.py b/python/paimon_vindex/gpu.py new file mode 100644 index 00000000..58de7f81 --- /dev/null +++ b/python/paimon_vindex/gpu.py @@ -0,0 +1,388 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Optional NVIDIA cuVS training and construction of CPU-readable IVF-SQ indexes. + +CuPy/cuVS are imported only when CuvsKMeans is constructed. This module never +silently falls back to CPU. Install matching CUDA, CuPy and cuVS packages as +described at https://docs.nvidia.com/cuvs/installation . +""" + +import threading +from time import perf_counter + +import numpy as np + +from . import (PreparedIvfSqTraining, VectorIndexWriter, _float32_matrix, + _int64_vector, _size_t) + + +class CuvsKMeans: + """Reusable, serialized GPU worker for squared-L2 IVF centroid training. + + The GPU runs flat Lloyd iterations. The existing CPU Auto strategy uses + hierarchical splitting for large nlist, so evaluate final index quality. + Caller-supplied initial centers allow comparison with CPU Lloyd. Otherwise + initial rows are chosen without replacement using the prepared state's + seed. This differs from the CPU Auto strategy's initialization. + """ + + def __init__(self, device=0, *, batch_samples=None, batch_centroids=None): + started = perf_counter() + self.device = _size_t(device, "device", allow_zero=True) + self._tile_options = {} + for key, value in (("batch_samples", batch_samples), ("batch_centroids", batch_centroids)): + if value is not None: + self._tile_options[key] = _size_t(value, key, allow_zero=False) + self._lock = threading.Lock() + self.last_run = None + try: + import cupy as cp + import cuvs + from cuvs.cluster.kmeans import KMeansParams, fit + from cuvs.common import Resources + except (ImportError, OSError) as exc: + raise RuntimeError( + "GPU training requires compatible CuPy and cuVS packages with " + "cuvs.cluster.kmeans.KMeansParams/fit on an NVIDIA CUDA host. " + "See https://docs.nvidia.com/cuvs/installation . " + "CPU training remains available via fit_centroids_cpu()." + ) from exc + self._cp = cp + self._params_type = KMeansParams + self._fit = fit + for key in self._tile_options: + if not hasattr(KMeansParams, key): + raise RuntimeError(f"installed cuVS does not expose {key}; upgrade cuVS or use its default tiling") + self.cuvs_version = getattr(cuvs, "__version__", "unknown") + self.cupy_version = cp.__version__ + self._resources = None + try: + with cp.cuda.Device(self.device): + self._resources = Resources() + name = cp.cuda.runtime.getDeviceProperties(self.device)["name"] + self.device_name = name.decode() if isinstance(name, bytes) else str(name) + cp.cuda.Device(self.device).synchronize() + except Exception: + try: + self.close() + except Exception: + pass # Preserve the initialization error if cleanup also fails. + raise + self.initialization_seconds = perf_counter() - started + + def fit(self, prepared, *, initial_centroids=None): + """Return host float32 centers. Timings include sample export and copies. + + The input sample stays on the device throughout the Lloyd iterations. + All timings synchronize the device, including work on cuVS streams. + `last_run` describes the last successful call; it is cleared on failure. + """ + with self._lock: + self.last_run = None + if self._resources is None: + raise RuntimeError("CuvsKMeans is closed") + if not isinstance(prepared, PreparedIvfSqTraining): + raise TypeError("prepared must be a PreparedIvfSqTraining") + started = perf_counter() + info = prepared.info + sample = prepared.sample + if info.sample_count < info.nlist: + raise ValueError("cuVS training requires at least nlist sample rows") + if initial_centroids is None: + rows = np.random.default_rng(info.seed).choice( + info.sample_count, size=info.nlist, replace=False, + ) + initial = np.ascontiguousarray(sample[rows]) + else: + initial = _float32_matrix(initial_centroids, "initial_centroids") + if initial.shape != (info.nlist, info.dimension) or not np.isfinite(initial).all(): + raise ValueError("initial_centroids must be finite with shape (nlist, dimension)") + cp = self._cp + with cp.cuda.Device(self.device): + # Array initialization fixes the initial centers across CPU/GPU + # references; one supplied initialization implies one run. + params = self._params_type( + metric="sqeuclidean", n_clusters=info.nlist, + init_method="Array", max_iter=info.iterations, n_init=1, + tol=1e-6, hierarchical=False, + **self._tile_options, + ) + cp.cuda.Device(self.device).synchronize() + transfer_started = perf_counter() + device_sample = cp.asarray(sample) + device_initial = cp.asarray(initial) + cp.cuda.Device(self.device).synchronize() + h2d_seconds = perf_counter() - transfer_started + fit_started = perf_counter() + device_centers, inertia, n_iter = self._fit( + params, device_sample, centroids=device_initial, + resources=self._resources, + ) + cp.cuda.Device(self.device).synchronize() + kmeans_seconds = perf_counter() - fit_started + transfer_started = perf_counter() + centers = np.ascontiguousarray(cp.asnumpy(cp.asarray(device_centers)), dtype=np.float32) + cp.cuda.Device(self.device).synchronize() + d2h_seconds = perf_counter() - transfer_started + if centers.shape != initial.shape or not np.isfinite(centers).all(): + raise RuntimeError("cuVS returned invalid IVF centroids") + self.last_run = { + "backend": "cuvs", "algorithm": "lloyd", "device": self.device, + "device_name": self.device_name, "cuvs_version": self.cuvs_version, + "cupy_version": self.cupy_version, "dtype": "float32", + "tile_options": dict(self._tile_options), + "initialization": "provided" if initial_centroids is not None else "sample_random", + "seed": info.seed, "sample_bytes": sample.nbytes, + "centroid_bytes": centers.nbytes, "iterations": int(n_iter), + "inertia": float(inertia), "h2d_seconds": h2d_seconds, + "kmeans_seconds": kmeans_seconds, "d2h_seconds": d2h_seconds, + "total_seconds": perf_counter() - started, + } + return centers + + def close(self): + with self._lock: + if self._resources is not None: + try: + with self._cp.cuda.Device(self.device): + self._cp.cuda.Device(self.device).synchronize() + finally: + # CUDA errors must not leave the worker open or retain its + # resources. A poisoned CUDA context may require a new process. + self._resources = None + + def __enter__(self): + with self._lock: + if self._resources is None: + raise RuntimeError("CuvsKMeans is closed") + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + try: + self.close() + except Exception: + if exc_type is None: + raise + return False + + +class CuvsIvfSqWriter: + """Batch GPU assignment and optional SQ8 encoding for an existing writer. + + The caller owns the native writer and writes/closes it as usual. Model + arrays stay on the selected device across batches. Set encode=False to + keep SQ encoding on CPU. GPU assignment is exact squared L2 for all metrics; + cosine preprocessing is delegated to Rust. For large centroid matrices, + train with ivf.coarse-assignment=exact so SQ calibration uses the same policy. + """ + + def __init__(self, writer, device=0, *, encode=True): + started = perf_counter() + if not isinstance(writer, VectorIndexWriter): + raise TypeError("writer must be a VectorIndexWriter") + if not isinstance(encode, bool): + raise TypeError("encode must be bool") + self._lock = threading.Lock() + self._worker = None + self._arrays = None + self._closed = False + self.last_run = None + self._writer = writer + self._encode = encode + self._model = writer.ivf_sq_encoding_model() + if not self._model.exact_assignment: + raise ValueError("GPU building requires SQ calibration with exact assignment; " + "train again with ivf.coarse-assignment=exact") + try: + self._worker = CuvsKMeans(device) + from cuvs.cluster.kmeans import predict + self._predict = predict + worker = self._worker + cp = worker._cp + model = self._model + with cp.cuda.Device(worker.device): + self._params = worker._params_type(metric="sqeuclidean", n_clusters=model.nlist, + hierarchical=False) + self._arrays = {"centers": cp.asarray(model.centroids)} + if encode: + scales = np.zeros_like(model.mins) + with np.errstate(over="ignore", divide="ignore", invalid="ignore"): + np.divide(np.float32(255), model.maxs - model.mins, out=scales, + where=model.mins < model.maxs) + self._arrays.update(mins=cp.asarray(model.mins), maxs=cp.asarray(model.maxs), + scales=cp.asarray(scales)) + self._kernel = cp.ElementwiseKernel( + "float32 x, raw I labels, raw float32 centers, raw float32 mins, " + "raw float32 maxs, raw float32 scales, int32 d, int32 vector_width", + "uint8 code", r''' + int dim = i % d; + long long at = (long long)labels[i / d] * d + dim; + float lo = mins[at], hi = maxs[at]; + float value = 0.0f; + if (paimon_lt(lo, hi)) { + float residual = paimon_sub_rn(x, centers[at]); + float shifted = paimon_sub_rn(residual, lo); + int cutoff = vector_width > 1 ? d / vector_width * vector_width : 0; + value = dim < cutoff ? paimon_mul_rn(shifted, scales[at]) + : paimon_div_rn(paimon_mul_rn(shifted, 255.0f), paimon_sub_rn(hi, lo)); + } + value = fminf(255.0f, fmaxf(0.0f, value)); + float lower = floorf(value); + code = (unsigned char)(lower + (value - lower >= 0.5f ? 1.0f : 0.0f)); + ''', "paimon_ivfsq_encode_v1", options=("--fmad=false",), preamble=r''' + // CuPy appends -ftz=true after user options. Explicit PTX + // without .ftz preserves subnormal inputs and results, + // including comparisons of tiny nonconstant SQ bounds. + __device__ __forceinline__ unsigned int paimon_lt(float a, float b) { + unsigned int out; + asm("set.lt.u32.f32 %0, %1, %2;" : "=r"(out) : "f"(a), "f"(b)); + return out; + } + __device__ __forceinline__ float paimon_sub_rn(float a, float b) { + float out; + asm("sub.rn.f32 %0, %1, %2;" : "=f"(out) : "f"(a), "f"(b)); + return out; + } + __device__ __forceinline__ float paimon_mul_rn(float a, float b) { + float out; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(out) : "f"(a), "f"(b)); + return out; + } + __device__ __forceinline__ float paimon_div_rn(float a, float b) { + float out; + asm("div.rn.f32 %0, %1, %2;" : "=f"(out) : "f"(a), "f"(b)); + return out; + } + ''') + cp.cuda.Device(worker.device).synchronize() + self.initialization_seconds = perf_counter() - started + except Exception: + try: + self.close() + except Exception: + pass # Preserve the build initialization error. + raise + + def add_vectors(self, ids, data): + """Process one batch; last_run includes validation, transfers and native append. + + Initialize once and call repeatedly with bounded batches. A failed batch + is not appended; prior successful batches stay in the caller-owned writer. + The first encoding call includes any CUDA kernel compilation cost. + """ + with self._lock: + self.last_run = None + if self._closed: + raise RuntimeError("CuvsIvfSqWriter is closed") + started = perf_counter() + raw = _float32_matrix(data, "data") + ids = _int64_vector(ids, "ids") + model = self._model + if raw.shape != (len(ids), model.dimension) or not len(ids): + raise ValueError("data must have shape (len(ids), dimension) with at least one row") + # Check the writer before doing device work; each native operation + # still locks/checks its handle in case another thread closes it. + self._writer.dimension + if model.metric == "cosine": + processed = self._writer._preprocess_ivf_sq_vectors(raw) + else: + if not np.isfinite(raw).all(): + raise ValueError("data must contain only finite values") + processed = raw + input_seconds = perf_counter() - started + worker = self._worker + cp = worker._cp + with cp.cuda.Device(worker.device): + cp.cuda.Device(worker.device).synchronize() + tick = perf_counter() + vectors = cp.asarray(processed) + cp.cuda.Device(worker.device).synchronize() + h2d_seconds = perf_counter() - tick + tick = perf_counter() + labels, _ = self._predict(self._params, vectors, self._arrays["centers"], + resources=worker._resources) + cp.cuda.Device(worker.device).synchronize() + labels = cp.asarray(labels).reshape(-1) + assignment_seconds = perf_counter() - tick + tick = perf_counter() + host_labels = cp.asnumpy(labels) + cp.cuda.Device(worker.device).synchronize() + d2h_seconds = perf_counter() - tick + # Validate before the quantization kernel uses labels as addresses. + if (host_labels.shape != (len(ids),) or host_labels.dtype.kind not in "iu" + or np.any(host_labels < 0) or np.any(host_labels >= model.nlist)): + raise RuntimeError("cuVS returned invalid partition IDs") + host_labels = np.ascontiguousarray(host_labels, dtype=np.uint32) + encode_seconds = 0.0 + code_bytes = 0 + if self._encode: + tick = perf_counter() + arrays = self._arrays + codes = self._kernel(vectors, labels, arrays["centers"], arrays["mins"], + arrays["maxs"], arrays["scales"], np.int32(model.dimension), + np.int32(model._encoding_vector_width)) + cp.cuda.Device(worker.device).synchronize() + encode_seconds = perf_counter() - tick + tick = perf_counter() + host_codes = cp.asnumpy(codes) + cp.cuda.Device(worker.device).synchronize() + d2h_seconds += perf_counter() - tick + code_bytes = host_codes.nbytes + tick = perf_counter() + if self._encode: + self._writer.add_encoded_vectors(ids, host_codes, host_labels) + else: + self._writer.add_preassigned_vectors(ids, raw, host_labels) + append_seconds = perf_counter() - tick + self.last_run = { + "backend": "cuvs-encode" if self._encode else "cuvs-assign", + "assignment": "exact", "rows": len(ids), "device": worker.device, + "device_name": worker.device_name, + "cuvs_version": worker.cuvs_version, "cupy_version": worker.cupy_version, + "input_seconds": input_seconds, "h2d_seconds": h2d_seconds, + "assignment_seconds": assignment_seconds, "encode_seconds": encode_seconds, + "d2h_seconds": d2h_seconds, "append_seconds": append_seconds, + "total_seconds": perf_counter() - started, + "h2d_bytes": processed.nbytes, "d2h_bytes": host_labels.nbytes + code_bytes, + } + + def close(self): + with self._lock: + if self._closed: + return + try: + if self._worker is not None: + self._worker.close() + finally: + self._arrays = None + self._worker = None + self._closed = True + + def __enter__(self): + with self._lock: + if self._closed: + raise RuntimeError("CuvsIvfSqWriter is closed") + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + try: + self.close() + except Exception: + if exc_type is None: + raise + return False diff --git a/python/tests/test_gpu_build.py b/python/tests/test_gpu_build.py new file mode 100644 index 00000000..1b8871a6 --- /dev/null +++ b/python/tests/test_gpu_build.py @@ -0,0 +1,348 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import ctypes +import importlib.util +import io +import json +import os +import sys +import subprocess +import threading +import types +from pathlib import Path + +import numpy as np +import pytest + +from paimon_vindex import SearchParams, VectorIndexReader, VectorIndexTrainer, VectorIndexWriter +from paimon_vindex._ffi import lib +from paimon_vindex.gpu import CuvsIvfSqWriter + + +class BytesInput: + def __init__(self, payload): + self.payload = payload + + def pread_many(self, ranges): + return [self.payload[p:p+n] for p,n in ranges] + + +def make_writer(data, metric="l2", nlist=4, centers=None): + options = {"index.type":"ivf_sq", "dimension":str(data.shape[1]), "nlist":str(nlist), + "metric":metric, "ivf.coarse-assignment":"exact"} + with VectorIndexTrainer.create(options) as trainer: + trainer.add_training_vectors(data) + if centers is None: + training = trainer.finish_training() + else: + with trainer.prepare_training() as prepared: + training = prepared.finish_with_ivf_centroids(centers) + return VectorIndexWriter(training) + + +def payload(writer): + result = io.BytesIO() + writer.write(result) + return result.getvalue() + + +@pytest.mark.parametrize("metric", ["l2", "cosine", "inner_product"]) +def test_preassigned_native_writer_and_model(metric): + data = np.random.default_rng(42).normal(size=(257,9)).astype(np.float32) + ids = np.arange(len(data),dtype=np.int64)*13-2000 + with make_writer(data, metric) as original, make_writer(data, metric) as assigned: + model = assigned.ivf_sq_encoding_model() + assert model.exact_assignment and model.centroids.shape == model.mins.shape == model.maxs.shape == (4,9) + assert not model.centroids.flags.writeable + processed = assigned._preprocess_ivf_sq_vectors(data) + lists = ((processed.astype(np.float64)[:,None,:]-model.centroids.astype(np.float64)[None,:,:])**2).sum(axis=2).argmin(axis=1) + for start in range(0,len(data),31): + assigned.add_preassigned_vectors(ids[start:start+31], data[start:start+31], lists[start:start+31]) + original.add_vectors(ids,data) + assert payload(original) == payload(assigned) + model.centroids.flags.writeable = True + model.centroids[:] = 123 + assert not np.all(assigned.ivf_sq_encoding_model().centroids == 123) + + +def test_invalid_batches_and_ffi_model_buffers_preserve_writer(): + data = np.arange(64,dtype=np.float32).reshape(8,8) + with make_writer(data) as writer: + writer.add_vectors([100],data[:1]) + before=payload(writer) + for labels in [[0,4], [-1,0], [2**32,0], [0.5,1.0]]: + with pytest.raises((ValueError,RuntimeError)): + writer.add_preassigned_vectors([1,2],data[:2],labels) + with pytest.raises((ValueError,RuntimeError)): + writer.add_encoded_vectors([1,2],np.zeros((2,8),dtype=np.uint8),labels) + with pytest.raises(ValueError): + writer.add_encoded_vectors([1],np.zeros((1,8),dtype=np.int32),[0]) + with pytest.raises((ValueError,RuntimeError)): + writer.add_preassigned_vectors([1],np.full((1,8),np.nan,dtype=np.float32),[0]) + out=np.full((4,8),123,dtype=np.float32) + ptr=out.ctypes.data_as(ctypes.POINTER(ctypes.c_float)) + assert lib.paimon_vindex_writer_ivf_sq_copy_model(writer._handle,ptr,None,ptr,out.size) != 0 + assert np.all(out==123) + assert lib.paimon_vindex_writer_add_encoded_vectors(writer._handle,None,None,0,None,1) != 0 + assert payload(writer)==before + writer.add_encoded_vectors([1],np.zeros((1,8),dtype=np.uint8),[0]) + assert writer.ivf_sq_partition_sizes().sum()==2 + + +def install_fake_cuda(monkeypatch, predict, *, asarray=np.asarray, synchronize=None): + from paimon_vindex import gpu + workers = [] + + class Device: + def __init__(self, device): pass + def __enter__(self): return self + def __exit__(self, *args): pass + def synchronize(self): + if synchronize is not None: + synchronize() + + class Worker(gpu.CuvsKMeans): + def __init__(self, device=0): + self._lock = threading.Lock() + self.device=device + self.device_name="fake-contract-device" + self.cuvs_version=self.cupy_version="fake-contract-test" + self._resources=object() + self._params_type=lambda **kwargs: kwargs + self._cp=types.SimpleNamespace(cuda=types.SimpleNamespace(Device=Device), + asarray=asarray,asnumpy=np.asarray) + workers.append(self) + + monkeypatch.setattr(gpu,"CuvsKMeans",Worker) + monkeypatch.setitem(sys.modules,"cuvs.cluster.kmeans",types.SimpleNamespace(predict=predict)) + return workers + + +@pytest.mark.parametrize("target", ["worker", "builder"]) +@pytest.mark.parametrize("body_error", [False, True]) +def test_cuda_cleanup_failure_releases_resources_and_preserves_original_error(monkeypatch, target, body_error): + install_fake_cuda(monkeypatch, lambda *args, **kwargs: (np.array([0]), 0.)) + data = np.arange(64, dtype=np.float32).reshape(8, 8) + with make_writer(data) as writer: + builder = CuvsIvfSqWriter(writer, encode=False) + worker = builder._worker + resource = worker if target == "worker" else builder + + def fail_sync(self): + raise RuntimeError("CUDA synchronization failed") + + monkeypatch.setattr(worker._cp.cuda.Device, "synchronize", fail_sync) + if body_error: + with pytest.raises(ValueError, match="original operation failed"): + with resource: + raise ValueError("original operation failed") + else: + with pytest.raises(RuntimeError, match="CUDA synchronization failed"): + resource.close() + assert worker._resources is None + resource.close() # Failed cleanup still makes close idempotent. + with pytest.raises(RuntimeError, match="closed"): + with resource: + pass + if target == "builder": + assert builder._arrays is None + builder.close() + writer.add_vectors([1], data[:1]) + assert writer.ivf_sq_partition_sizes().sum() == 1 + + +def test_gpu_constructor_preserves_allocation_error_when_cleanup_also_fails(monkeypatch): + def fail_copy(*args): + raise MemoryError("model allocation failed") + + def fail_sync(): + raise RuntimeError("cleanup synchronization failed") + + workers = install_fake_cuda(monkeypatch, lambda *args: None, + asarray=fail_copy, synchronize=fail_sync) + data = np.arange(64, dtype=np.float32).reshape(8, 8) + with make_writer(data) as writer: + with pytest.raises(MemoryError, match="model allocation failed"): + CuvsIvfSqWriter(writer, encode=False) + assert workers[0]._resources is None + writer.add_vectors([1], data[:1]) + + +def test_invalid_training_input_clears_previous_telemetry(monkeypatch): + workers = install_fake_cuda(monkeypatch, lambda *args: None) + data = np.arange(64, dtype=np.float32).reshape(8, 8) + with make_writer(data) as writer, CuvsIvfSqWriter(writer, encode=False): + worker = workers[0] + worker.last_run = {"rows": 123} + with pytest.raises(TypeError, match="PreparedIvfSqTraining"): + worker.fit(object()) + assert worker.last_run is None + + +def test_adapter_rejects_bad_backend_labels_and_recovers(monkeypatch): + state={"labels":np.array([-1,0],dtype=np.int32)} + install_fake_cuda(monkeypatch,lambda *args,**kwargs:(state["labels"],0.)) + data=np.arange(64,dtype=np.float32).reshape(8,8) + with make_writer(data) as writer: + with CuvsIvfSqWriter(writer,encode=False) as gpu: + before=payload(writer) + for labels in [np.array([-1,0]), np.array([0,4]), np.array([0]), np.array([0.,1.])]: + state["labels"]=labels + with pytest.raises(RuntimeError,match="invalid partition"): + gpu.add_vectors([1,2],data[:2]) + assert gpu.last_run is None and payload(writer)==before + state["labels"]=np.array([0,1],dtype=np.int32) + gpu.add_vectors([1,2],data[:2]) + assert gpu.last_run["rows"]==2 + with pytest.raises(ValueError): gpu.add_vectors([3],np.full((1,8),np.inf)) + assert gpu.last_run is None + assert writer.ivf_sq_partition_sizes().sum()==2 + with pytest.raises(RuntimeError,match="closed"): gpu.add_vectors([3],data[:1]) + writer.add_vectors([3],data[:1]) # adapter never owns the writer + assert writer.ivf_sq_partition_sizes().sum()==3 + + +def test_adapter_requires_consistent_assignment_before_importing_cuda(): + data=np.zeros((4,1024),dtype=np.float32) + options={"index.type":"ivf_sq","dimension":"1024","nlist":"1024","metric":"l2"} + with VectorIndexTrainer.create(options) as trainer: + trainer.add_training_vectors(data) + with trainer.prepare_training() as p: + training=p.finish_with_ivf_centroids(np.zeros((1024,1024),dtype=np.float32)) + with VectorIndexWriter(training) as writer: + with pytest.raises(ValueError,match="ivf.coarse-assignment=exact"): + CuvsIvfSqWriter(writer) + + +CUDA=pytest.mark.skipif(os.environ.get("PAIMON_TEST_CUVS")!="1",reason="set PAIMON_TEST_CUVS=1 on CUDA host") + + +@pytest.mark.parametrize("build_backend", ["cpu", pytest.param("cuvs-encode", marks=CUDA)]) +def test_fixed_model_benchmark_and_portable_reader(tmp_path, monkeypatch, build_backend): + data=np.random.default_rng(1234).normal(size=(256,9)).astype(np.float32) + with make_writer(data,nlist=8) as writer: + model=writer.ivf_sq_encoding_model() + np.save(tmp_path/"base.npy",data) + np.save(tmp_path/"centers.npy",model.centroids) + script=Path(__file__).resolve().parents[2]/"tools"/"benchmark_gpu_training.py" + subprocess.run([sys.executable,str(script),"--preset","smoke","--base",str(tmp_path/"base.npy"), + "--backends","external","--centroids",str(tmp_path/"centers.npy"),"--build-index", + "--build-backend",build_backend, + "--coarse-assignment","exact","--threads","2","--repeats","1", + "--output-dir",str(tmp_path/"result")],check=True,capture_output=True,text=True) + run = json.loads((tmp_path/"result"/"runs.jsonl").read_text()) + if build_backend != "cpu": + assert run["gpu_build"]["backend"] == build_backend + for key in ("device_name", "cuvs_version", "cupy_version"): + assert run["gpu_build"][key] + spec=importlib.util.spec_from_file_location("gpu_build_benchmark_test",script) + module=importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + monkeypatch.setattr(module,"os",types.SimpleNamespace()) # Windows has no os.pread + queries=data[:8] + truth=np.argsort(((queries[:,None,:]-data[None,:,:])**2).sum(axis=2),axis=1)[:,:10] + result=module.evaluate_index(tmp_path/"result"/"external-0.index",queries,truth,10,8) + assert result["recall_at_k"]>=.9 + + +@CUDA +@pytest.mark.parametrize("dimension",[1,7,8,9,129]) +def test_gpu_encoding_clipping_rounding_and_scalar_tails(dimension): + train=np.stack([np.full(dimension,-1,np.float32),np.full(dimension,1,np.float32)]) + centers=np.zeros((1,dimension),dtype=np.float32) + values=np.concatenate([np.array([-3,-1,0,1,3],dtype=np.float32), + ((np.arange(255,dtype=np.float32)+.5)/np.float32(127.5)-1)]) + data=np.resize(values,(523,dimension)).astype(np.float32) + ids=np.arange(len(data),dtype=np.int64)[::-1] + with make_writer(train,nlist=1,centers=centers) as cpu, make_writer(train,nlist=1,centers=centers) as accelerated: + cpu.add_vectors(ids,data) + with CuvsIvfSqWriter(accelerated) as gpu: + for start in range(0,len(data),127): gpu.add_vectors(ids[start:start+127],data[start:start+127]) + assert payload(cpu)==payload(accelerated) + + +@CUDA +@pytest.mark.parametrize("metric",["l2","cosine","inner_product"]) +def test_gpu_encoding_constant_bounds_and_zero_vectors(metric): + data=np.zeros((32,9),dtype=np.float32) + centers=np.zeros((1,9),dtype=np.float32) + with make_writer(data,metric,1,centers) as cpu, make_writer(data,metric,1,centers) as accelerated: + ids=np.arange(len(data),dtype=np.int64) + cpu.add_vectors(ids,data) + with CuvsIvfSqWriter(accelerated) as gpu: + gpu.add_vectors(ids,data) + assert payload(cpu)==payload(accelerated) + + +@CUDA +@pytest.mark.parametrize("dimension", [1, 8, 9]) +def test_gpu_encoding_subnormal_bounds_and_overflowing_scale(dimension): + # Finite, nonconstant bounds can have an infinite f32 encoding scale. + # Exercise NaN from 0 * inf at the minimum, saturation, and scalar tails. + tiny = np.float32(np.finfo(np.float32).tiny / 16) + train = np.stack([np.full(dimension, -tiny, np.float32), + np.full(dimension, tiny, np.float32)]) + centers = np.zeros((1, dimension), dtype=np.float32) + data = np.repeat((np.arange(-2, 3, dtype=np.float32) * tiny)[:, None], dimension, axis=1) + ids = np.arange(len(data), dtype=np.int64) + with make_writer(train, nlist=1, centers=centers) as cpu, make_writer(train, nlist=1, centers=centers) as accelerated: + cpu.add_vectors(ids, data) + with CuvsIvfSqWriter(accelerated) as gpu: + gpu.add_vectors(ids, data) + assert payload(cpu) == payload(accelerated) + + +@CUDA +@pytest.mark.parametrize("metric",["l2","cosine","inner_product"]) +def test_gpu_build_default_training_and_heldout_retrieval(metric): + from paimon_vindex.gpu import CuvsKMeans + rng=np.random.default_rng(42) + k,d,n=300,16,12000 + means=rng.normal(size=(k,d)).astype(np.float32)*5 + means *= np.float32(20) / np.linalg.norm(means,axis=1)[:,None] + data=means[np.arange(n)%k]+rng.normal(0,.2,size=(n,d)).astype(np.float32) + queries=means[rng.integers(k,size=96)]+rng.normal(0,.2,size=(96,d)).astype(np.float32) + options={"index.type":"ivf_sq","dimension":str(d),"nlist":str(k),"metric":metric, + "ivf.coarse-assignment":"exact"} + with VectorIndexTrainer.create(options) as trainer: + trainer.add_training_vectors(data) + with trainer.prepare_training() as p, CuvsKMeans() as gpu: + centers=gpu.fit(p) # exercise default initialization with nlist > 256 + if metric=="l2": + scores=((queries[:,None,:]-data[None,:,:])**2).sum(axis=2) + elif metric=="cosine": + scores=-(queries@data.T)/(np.linalg.norm(queries,axis=1)[:,None]*np.linalg.norm(data,axis=1)[None,:]) + else: + scores=-(queries@data.T) + truth=np.argsort(scores,axis=1)[:,:10] + payloads=[] + for mode in ["cpu","cuvs-assign","cuvs-encode"]: + with make_writer(data,metric,k,centers) as writer: + if mode=="cpu": writer.add_vectors(np.arange(n,dtype=np.int64),data) + else: + with CuvsIvfSqWriter(writer,encode=mode=="cuvs-encode") as gpu: + for start in range(0,n,1024): + gpu.add_vectors(np.arange(start,min(start+1024,n),dtype=np.int64),data[start:start+1024]) + assert writer.ivf_sq_partition_sizes().sum()==n + payloads.append(payload(writer)) + with VectorIndexReader(BytesInput(payloads[-1])) as reader: + ids,distances=reader.search_batch(queries,SearchParams.ivf(10,16)) + recall=sum(len(set(found)&set(expected)) for found,expected in zip(ids,truth))/(len(queries)*10) + assert recall>=.9, (metric,mode,recall) + assert np.isfinite(distances).all() + assert payloads[1]==payloads[2] # same GPU assignments, native versus GPU encoding diff --git a/python/tests/test_gpu_training.py b/python/tests/test_gpu_training.py new file mode 100644 index 00000000..39bca942 --- /dev/null +++ b/python/tests/test_gpu_training.py @@ -0,0 +1,226 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import ctypes +import io +import json +import os +from pathlib import Path +import subprocess +import sys +import types + +import numpy as np +import pytest + +from paimon_vindex import SearchParams, VectorIndexReader, VectorIndexTrainer, VectorIndexWriter +from paimon_vindex._ffi import lib + + +class BytesInput: + def __init__(self, data): + self.data = data + + def pread_many(self, ranges): + return [self.data[pos:pos+length] for pos, length in ranges] + + +def make_trainer(data, metric="l2", **extra): + options = {"index.type": "ivf_sq", "dimension": str(data.shape[1]), + "nlist": "4", "metric": metric, **extra} + return VectorIndexTrainer.create(options).add_training_vectors(data) + + +def index_bytes(training, data): + output = io.BytesIO() + with VectorIndexWriter(training) as writer: + assert writer.ivf_sq_partition_sizes().sum() == 0 + writer.add_vectors(np.arange(len(data), dtype=np.int64), data) + assert writer.ivf_sq_partition_sizes().sum() == len(data) + writer.write(output) + return output.getvalue() + + +@pytest.mark.parametrize("metric", ["l2", "cosine", "inner_product"]) +def test_prepared_cpu_roundtrip_preserves_bytes_and_cpu_results(metric): + data = np.random.default_rng(1234).normal(size=(256, 8)).astype(np.float32) + original = index_bytes(make_trainer(data, metric).finish_training(), data) + trainer = make_trainer(data, metric) + with trainer.prepare_training() as prepared: + with pytest.raises(RuntimeError, match="closed"): + trainer.add_training_vectors(data) + sample = prepared.sample + expected = data / np.linalg.norm(data, axis=1, keepdims=True) if metric == "cosine" else data + np.testing.assert_allclose(sample, expected, atol=2e-7) + assert not sample.flags.writeable + # A caller-owned copy cannot mutate native SQ calibration state. + sample.flags.writeable = True + sample[:] = 0 + np.testing.assert_allclose(prepared.sample, expected, atol=2e-7) + centers = prepared.fit_centroids_cpu() + training = prepared.finish_with_ivf_centroids(centers) + with pytest.raises(RuntimeError, match="closed"): + prepared.fit_centroids_cpu() + actual = index_bytes(training, data) + assert actual == original + with VectorIndexReader(BytesInput(actual)) as reader: + ids, distances = reader.search(data[0], SearchParams.ivf(5, 4)) + assert len(ids) == 5 + assert np.isfinite(distances).all() + if metric != "inner_product": + assert ids[0] == 0 + + +def test_prepared_respects_resolved_nlist_and_both_sample_caps(): + data = np.random.default_rng(0).normal(size=(512, 8)).astype(np.float32) + with make_trainer(data, nlist="auto", **{"expected-vector-count": "1000000"}).prepare_training() as p: + assert p.info.nlist == 1024 # Not inferred from the 512-row sample. + with make_trainer(data, **{"ivf.train.max-points-per-centroid": "2"}).prepare_training() as p: + assert p.info.sample_count == 8 + assert p.info.calibration_count == 512 + assert p.sample.shape == (8, 8) + + +def test_prepared_validation_and_consumption(): + data = np.ones((8, 2), dtype=np.float32) + with make_trainer(data).prepare_training() as p: + # Transposed shapes have the same element count but are still invalid. + with pytest.raises(ValueError, match="shape"): + p.finish_with_ivf_centroids(np.zeros((2, 4), dtype=np.float32)) + with pytest.raises(RuntimeError, match="Lloyd"): + p.fit_centroids_cpu(initial_centroids=np.ones((4, 2), dtype=np.float32)) + with pytest.raises(RuntimeError, match="finite"): + p.finish_with_ivf_centroids(np.full((4, 2), np.nan, dtype=np.float32)) + with pytest.raises(RuntimeError, match="closed"): + p.sample + with make_trainer(data, **{"index.type": "ivf_flat"}) as trainer: + with pytest.raises(RuntimeError, match="IVF-SQ"): + trainer.prepare_training() + with pytest.raises(RuntimeError, match="closed"): + trainer.finish_training() + + +def test_ffi_sample_copy_validates_buffer_before_writing(): + data = np.ones((8, 2), dtype=np.float32) + with make_trainer(data).prepare_training() as p: + output = np.full(data.size, 123.0, dtype=np.float32) + assert lib.paimon_vindex_prepared_training_copy_sample( + p._handle, output.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), output.size - 1, + ) != 0 + assert (output == 123).all() + assert lib.paimon_vindex_prepared_training_copy_sample(p._handle, None, output.size) != 0 + + +def test_gpu_module_does_not_import_cuda_for_cpu_users(): + subprocess.run([sys.executable, "-c", "import sys; import paimon_vindex.gpu; " + "assert 'cupy' not in sys.modules; assert 'cuvs' not in sys.modules"], check=True) + + +def test_missing_cuda_dependencies_raise_an_actionable_error(monkeypatch): + from paimon_vindex.gpu import CuvsKMeans + + monkeypatch.setitem(sys.modules, "cupy", None) + with pytest.raises(RuntimeError, match="requires compatible CuPy and cuVS"): + CuvsKMeans() + + +def test_worker_initialization_failure_releases_resources(monkeypatch): + from paimon_vindex.gpu import CuvsKMeans + + class Device: + def __init__(self, device): pass + def __enter__(self): return self + def __exit__(self, *args): pass + def synchronize(self): raise RuntimeError("cleanup synchronization failed") + + def fail_properties(device): + raise ValueError("device properties failed") + + monkeypatch.setitem(sys.modules, "cupy", types.SimpleNamespace( + __version__="fake", cuda=types.SimpleNamespace( + Device=Device, runtime=types.SimpleNamespace(getDeviceProperties=fail_properties)))) + monkeypatch.setitem(sys.modules, "cuvs", types.SimpleNamespace(__version__="fake")) + monkeypatch.setitem(sys.modules, "cuvs.common", types.SimpleNamespace(Resources=object)) + monkeypatch.setitem(sys.modules, "cuvs.cluster.kmeans", types.SimpleNamespace(KMeansParams=object, fit=None)) + worker = CuvsKMeans.__new__(CuvsKMeans) + with pytest.raises(ValueError, match="device properties failed"): + worker.__init__() + assert worker._resources is None + worker.close() + + +@pytest.mark.parametrize("format", ["npy", "fvecs"]) +def test_benchmark_builds_and_evaluates_real_input_formats(tmp_path, format): + rng = np.random.default_rng(42) + base = rng.normal(size=(256, 8)).astype(np.float32) + queries = rng.normal(size=(8, 8)).astype(np.float32) + truth = np.argsort(((queries[:, None, :] - base[None, :, :]) ** 2).sum(axis=2), axis=1)[:, :10].astype(np.int32) + base_path = tmp_path / f"base.{format}" + query_path = tmp_path / f"query.{format}" + truth_path = tmp_path / ("truth.npy" if format == "npy" else "truth.ivecs") + for path, array, integer in [(base_path, base, False), (query_path, queries, False), (truth_path, truth, True)]: + if format == "npy": + np.save(path, array) + else: + records = np.empty((len(array), array.shape[1] + 1), dtype="= .9 + assert run["partitions"]["mean"] * 8 == 256 + assert run["build_seconds"] >= run["training_seconds"] > 0 + manifest = (output / "manifest.json").read_bytes() + assert subprocess.run(command, capture_output=True).returncode != 0 + assert (output / "manifest.json").read_bytes() == manifest + + +@pytest.mark.skipif(os.environ.get("PAIMON_TEST_CUVS") != "1", reason="set PAIMON_TEST_CUVS=1 on a CUDA host") +@pytest.mark.parametrize("metric", ["l2", "cosine", "inner_product"]) +def test_cuvs_training_produces_cpu_readable_index(metric): + from paimon_vindex.gpu import CuvsKMeans + + rng = np.random.default_rng(1234) + means = np.eye(4, 8, dtype=np.float32) * 20 + data = means[np.arange(512) % 4] + rng.normal(0, .1, size=(512, 8)).astype(np.float32) + with CuvsKMeans() as gpu: + for _ in range(2): # Reuse the same resources across independent jobs. + with make_trainer(data, metric).prepare_training() as p: + centers = gpu.fit(p, initial_centroids=p.sample[:4].copy()) + assert centers.shape == (4, 8) + assert gpu.last_run["kmeans_seconds"] > 0 + payload = index_bytes(p.finish_with_ivf_centroids(centers), data) + with VectorIndexReader(BytesInput(payload)) as reader: + query = data[0] + scores = ((data - query) ** 2).sum(axis=1) + if metric == "inner_product": + scores = -(data @ query) + elif metric == "cosine": + scores = 1 - (data @ query) / (np.linalg.norm(data, axis=1) * np.linalg.norm(query)) + exact = set(np.argsort(scores)[:10]) + ids, distances = reader.search(query, SearchParams.ivf(10, 4)) + assert len(exact.intersection(ids.tolist())) >= 8 + assert np.isfinite(distances).all() diff --git a/tools/README.md b/tools/README.md index 17cac1cf..7efaeb24 100644 --- a/tools/README.md +++ b/tools/README.md @@ -21,6 +21,17 @@ This directory contains helper scripts used by release managers and committers. +## GPU training benchmark + +`benchmark_gpu_training.py` compares the current CPU IVF-SQ trainer, CPU Lloyd +and optional NVIDIA cuVS center training on identical prepared samples. It +supports CPU, GPU assignment, and GPU assignment plus SQ8 encoding builders; +saved centers can be reused to compare builders with an identical model. It +records training stages and can also measure full index construction, +partition distribution and recall with supplied ground truth. See the +[GPU training guide](../docs/GPU_TRAINING.md) for setup, workload presets and +timing boundaries. + ## ANN-Benchmarks dataset conversion `convert_ann_benchmarks.py` converts a dense diff --git a/tools/benchmark_gpu_training.py b/tools/benchmark_gpu_training.py new file mode 100644 index 00000000..6e23f46b --- /dev/null +++ b/tools/benchmark_gpu_training.py @@ -0,0 +1,352 @@ +#!/usr/bin/env python3 +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Compare IVF-SQ CPU Auto, CPU Lloyd, and optional cuVS training. + +Source loading, diagnostic hashes and artifact export are outside stage timers. +Training includes native sampling/preprocessing, centroid fit and CPU SQ +calibration. GPU worker setup is reported separately and included in the first +GPU job's inclusive time. No GPU timings or speedups are synthesized. +""" + +import argparse +from contextlib import nullcontext +from dataclasses import asdict +import hashlib +import json +import os +from pathlib import Path +import platform +import subprocess +import threading +from time import perf_counter + +import numpy as np + + +PRESETS = { + "smoke": (2048, 16, 8), + "baseline": (65536, 960, 1024), + "medium": (262144, 1536, 4096), + "large": (1048576, 1536, 16384), +} + + +def positive(value): + number = int(value) + if number <= 0: + raise argparse.ArgumentTypeError("must be positive") + return number + + +def load_matrix(path, *, neighbors=False): + """Memory-map .npy or ANN-Benchmarks .fvecs/.ivecs matrices.""" + if path.suffix == ".npy": + result = np.load(path, mmap_mode="r", allow_pickle=False) + else: + expected = ".ivecs" if neighbors else ".fvecs" + if path.suffix != expected: + raise ValueError(f"expected .npy or {expected}: {path}") + if path.stat().st_size < 4 or path.stat().st_size % 4: + raise ValueError(f"invalid record file length: {path}") + raw = np.memmap(path, mode="r", dtype="= 1_000_000: + parser.error("GPU building requires --coarse-assignment exact for this centroid matrix") + fixed_centers = load_matrix(args.centroids) if args.centroids else None + if fixed_centers is not None and (fixed_centers.shape != (nlist, dimension) or not np.isfinite(fixed_centers).all()): + parser.error("external centers must be finite with shape (nlist, dimension)") + nprobe = args.nprobe or min(nlist, max(1, nlist // 16)) + if nprobe > nlist: + parser.error("nprobe must not exceed nlist") + queries = truth = None + if args.queries: + queries = load_matrix(args.queries)[:args.query_limit] + truth = load_matrix(args.neighbors, neighbors=True)[:args.query_limit] + if queries.shape[1] != dimension or len(queries) != len(truth) or truth.shape[1] < args.top_k: + parser.error("query/ground-truth shapes do not match the base and top-k") + if not np.isfinite(queries).all() or np.any(truth < 0) or np.any(truth >= rows): + parser.error("queries must be finite and ground-truth IDs must refer to the supplied base") + args.output_dir.mkdir(parents=True, exist_ok=False) + options = {"index.type": "ivf_sq", "dimension": str(dimension), "nlist": str(nlist), + "metric": args.metric, "expected-vector-count": str(rows), + "ivf.coarse-assignment": args.coarse_assignment} + try: + revision = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=Path(__file__).resolve().parent, + text=True, stderr=subprocess.DEVNULL).strip() + dirty = bool(subprocess.check_output(["git", "status", "--porcelain"], + cwd=Path(__file__).resolve().parent, + stderr=subprocess.DEVNULL)) + except (OSError, subprocess.CalledProcessError): + revision = None + dirty = None + native_library = Path(_ffi.lib._name).resolve() + manifest = {"schema_version": 1, "preset": args.preset, "source": source_name, + "base_shape": list(base.shape), "options": options, "backends": args.backends, + "repeats": args.repeats, "rayon_threads": args.threads, + "platform": platform.platform(), "python": platform.python_version(), + "numpy": np.__version__, "git_revision": revision, "git_dirty": dirty, + "native_library": str(native_library), "native_library_sha256": file_hash(native_library), + "benchmark_sha256": file_hash(Path(__file__)), + "gpu_peak_memory": "not measured", "quality_evaluated": queries is not None} + manifest.update(build_backend=args.build_backend, device=args.device, batch_rows=args.batch_rows, + external_centroids=None if args.centroids is None else str(args.centroids.resolve()), + external_centroids_sha256=None if args.centroids is None else file_hash(args.centroids)) + (args.output_dir / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + gpu = None + runs = [] + try: + with (args.output_dir / "runs.jsonl").open("x") as results: + for repeat in range(args.repeats): + # Rotate execution order to avoid always favoring one backend. + offset = repeat % len(args.backends) + order = args.backends[offset:] + args.backends[:offset] + for backend in order: + print(f"{backend} repeat={repeat}: {rows}x{dimension}, nlist={nlist}", flush=True) + start = perf_counter() + with VectorIndexTrainer.create(options) as trainer: + for begin in range(0, rows, args.batch_rows): + trainer.add_training_vectors(base[begin:begin+args.batch_rows]) + prepared = trainer.prepare_training() + prepare_seconds = perf_counter() - start + with prepared: + info = asdict(prepared.info) + start = perf_counter() + sample = prepared.sample + sample_export_seconds = perf_counter() - start + sample_sha256 = array_hash(sample) + start = perf_counter() + initial = np.ascontiguousarray(sample[np.random.default_rng(prepared.info.seed).choice( + len(sample), size=nlist, replace=False)]) + initial_prepare_seconds = (sample_export_seconds + perf_counter() - start) if backend == "cpu-lloyd" else 0.0 + initial_sha256 = array_hash(initial) + del sample + setup_seconds = 0.0 + if backend == "cuvs" and gpu is None: + from paimon_vindex.gpu import CuvsKMeans + gpu = CuvsKMeans(args.device) + setup_seconds = gpu.initialization_seconds + start = perf_counter() + if backend == "cuvs": + # The adapter uses the same seed/row-selection rule + # as `initial`, and times its own export/initialization. + centers = gpu.fit(prepared) + elif backend == "external": + centers = fixed_centers + else: + centers = prepared.fit_centroids_cpu( + "auto" if backend == "cpu-auto" else "lloyd", + initial_centroids=None if backend == "cpu-auto" else initial, + ) + fit_seconds = perf_counter() - start + initial_prepare_seconds + start = perf_counter() + training = prepared.finish_with_ivf_centroids(centers) + sq_seconds = perf_counter() - start + prefix = f"{backend}-{repeat}" + np.save(args.output_dir / f"{prefix}-centroids.npy", centers, allow_pickle=False) + run = {"backend": backend, "repeat": repeat, "training_info": info, + "sample_sha256": sample_sha256, + "initial_centroids_sha256": None if backend in ("cpu-auto", "external") else initial_sha256, + "centroids_sha256": array_hash(centers), + "external_initialization_seconds": initial_prepare_seconds, + "prepare_seconds": prepare_seconds, "centroid_fit_seconds": fit_seconds, + "sq_calibration_seconds": sq_seconds, "gpu_setup_seconds": setup_seconds, + "build_backend": args.build_backend, + "training_seconds": prepare_seconds + fit_seconds + sq_seconds, + "training_including_setup_seconds": prepare_seconds + fit_seconds + sq_seconds + setup_seconds} + if backend == "cuvs": + run["gpu"] = gpu.last_run + try: + if args.build_index: + with VectorIndexWriter(training) as writer: + start = perf_counter() + if args.build_backend == "cpu": + builder_context = nullcontext(writer) + else: + from paimon_vindex.gpu import CuvsIvfSqWriter + builder_context = CuvsIvfSqWriter(writer, args.device, encode=args.build_backend == "cuvs-encode") + run["gpu_build_setup_seconds"] = builder_context.initialization_seconds + stages = {} + with builder_context as builder: + for begin in range(0, rows, args.batch_rows): + chunk = base[begin:begin+args.batch_rows] + builder.add_vectors(np.arange(begin, begin+len(chunk), dtype=np.int64), chunk) + if args.build_backend != "cpu": + if "gpu_build" not in run: + run["gpu_build"] = {key: builder.last_run[key] for key in ( + "backend", "assignment", "device", "device_name", + "cuvs_version", "cupy_version", + )} + for key, value in builder.last_run.items(): + if key.endswith("_seconds") or key.endswith("_bytes") or key == "rows": + stages[key] = stages.get(key, 0) + value + add_seconds = perf_counter() - start + if stages: + run["gpu_build_stages"] = stages + sizes = writer.ivf_sq_partition_sizes() + run["partitions"] = {"empty": int((sizes == 0).sum()), "max": int(sizes.max()), + "p95": float(np.percentile(sizes, 95)), "mean": float(sizes.mean())} + index_path = args.output_dir / f"{prefix}.index" + start = perf_counter() + with index_path.open("xb") as out: + writer.write(out) + write_seconds = perf_counter() - start + run.update(add_seconds=add_seconds, serialize_seconds=write_seconds, + index_bytes=index_path.stat().st_size, + build_seconds=run["training_seconds"] + add_seconds + write_seconds, + build_including_setup_seconds=run["training_including_setup_seconds"] + add_seconds + write_seconds) + if queries is not None: + run["query"] = evaluate_index(index_path, queries, truth, args.top_k, nprobe) + finally: + training.close() + runs.append(run) + results.write(json.dumps(run, allow_nan=False) + "\n") + results.flush() + summary = {} + for backend in args.backends: + selected = [r for r in runs if r["backend"] == backend] + summary[backend] = {key: float(np.median([r[key] for r in selected])) + for key in ("centroid_fit_seconds", "training_seconds", "training_including_setup_seconds", + "build_seconds", "build_including_setup_seconds") if key in selected[0]} + (args.output_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary, indent=2)) + finally: + if gpu is not None: + gpu.close() + + +if __name__ == "__main__": + main()