diff --git a/.github/workflows/disk-benchmarks.yml b/.github/workflows/disk-benchmarks.yml index bd3586761..29c0a7af8 100644 --- a/.github/workflows/disk-benchmarks.yml +++ b/.github/workflows/disk-benchmarks.yml @@ -28,7 +28,6 @@ on: - 'diskann-vector/**' - 'diskann-wide/**' - 'diskann-utils/**' - - 'diskann-platform/**' - 'diskann-label-filter/**' - 'diskann-benchmark/**' - 'diskann-benchmark-runner/**' diff --git a/Cargo.lock b/Cargo.lock index aaae873b7..c623535e8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -763,7 +763,6 @@ dependencies = [ "criterion", "diskann", "diskann-linalg", - "diskann-platform", "diskann-providers", "diskann-quantization", "diskann-utils", @@ -785,6 +784,7 @@ dependencies = [ "tokio", "tracing", "vfs", + "windows-sys 0.61.2", ] [[package]] @@ -842,16 +842,6 @@ dependencies = [ "thiserror 2.0.17", ] -[[package]] -name = "diskann-platform" -version = "0.54.0" -dependencies = [ - "io-uring", - "libc", - "tracing", - "windows-sys 0.59.0", -] - [[package]] name = "diskann-providers" version = "0.54.0" @@ -864,7 +854,6 @@ dependencies = [ "cfg-if", "criterion", "diskann", - "diskann-platform", "diskann-quantization", "diskann-utils", "diskann-vector", @@ -875,7 +864,6 @@ dependencies = [ "iai-callgrind", "itertools 0.13.0", "num-traits", - "opentelemetry", "proptest", "prost", "rand", diff --git a/Cargo.toml b/Cargo.toml index 85f29ca71..3f873895b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,7 +9,6 @@ members = [ "diskann-linalg", "diskann-utils", "diskann-quantization", - "diskann-platform", # Algorithm "diskann", # Providers @@ -55,7 +54,6 @@ diskann-vector = { path = "diskann-vector", version = "0.54.0" } diskann-linalg = { path = "diskann-linalg", version = "0.54.0" } diskann-utils = { path = "diskann-utils", default-features = false, version = "0.54.0" } diskann-quantization = { path = "diskann-quantization", default-features = false, version = "0.54.0" } -diskann-platform = { path = "diskann-platform", version = "0.54.0" } # Algorithm diskann = { path = "diskann", version = "0.54.0" } # Providers diff --git a/agents.md b/agents.md index a15e29e9d..e757159f0 100644 --- a/agents.md +++ b/agents.md @@ -44,7 +44,6 @@ The repository uses a Cargo workspace with crates organized into functional tier **Tier 1: Foundation** - `diskann-wide/` - Low-level SIMD, bit manipulation, type width abstractions - `diskann-vector/` - Vector primitives and operations -- `diskann-platform/` - Platform-specific utilities **Tier 2: Core Libraries** - `diskann-linalg/` - Linear algebra operations diff --git a/diskann-benchmark/Cargo.toml b/diskann-benchmark/Cargo.toml index ce5018aad..b1759b489 100644 --- a/diskann-benchmark/Cargo.toml +++ b/diskann-benchmark/Cargo.toml @@ -73,7 +73,6 @@ bftree = ["dep:diskann-bftree"] # Enable Disk Index benchmarks disk-index = [ "diskann-disk/perf_test", - "diskann-providers/perf_test", "diskann-tools/perf_test", "dep:diskann-disk", "dep:opentelemetry", diff --git a/diskann-benchmark/src/disk_index/search.rs b/diskann-benchmark/src/disk_index/search.rs index db7bccbdc..f035d2fb4 100644 --- a/diskann-benchmark/src/disk_index/search.rs +++ b/diskann-benchmark/src/disk_index/search.rs @@ -18,7 +18,7 @@ use diskann_disk::{ disk_vertex_provider_factory::DiskVertexProviderFactory, }, storage::disk_index_reader::DiskIndexReader, - utils::{instrumentation::PerfLogger, statistics, AlignedFileReaderFactory, QueryStatistics}, + utils::{instrumentation::PerfLogger, statistics, QueryStatistics}, }; use diskann_providers::storage::StorageReadProvider; use diskann_providers::{ @@ -218,8 +218,8 @@ where CachingStrategy::None }; - let reader_factory = AlignedFileReaderFactory::new(disk_index_path); - let vertex_provider_factory = DiskVertexProviderFactory::new(reader_factory, caching_strategy)?; + let vertex_provider_factory = + DiskVertexProviderFactory::from_disk_index_path(disk_index_path, caching_strategy)?; let searcher = &DiskIndexSearcher::, _>::new( search_params.num_threads, diff --git a/diskann-disk/Cargo.toml b/diskann-disk/Cargo.toml index d49fafa17..cc8b7a7cb 100644 --- a/diskann-disk/Cargo.toml +++ b/diskann-disk/Cargo.toml @@ -33,7 +33,6 @@ cfg-if.workspace = true futures-util = { workspace = true, default-features = false } hashbrown.workspace = true half = { workspace = true, features = ["bytemuck", "num-traits"] } -diskann-platform = { workspace = true } rand.workspace = true rayon.workspace = true serde = { workspace = true, features = ["derive"] } @@ -47,8 +46,24 @@ opentelemetry = { workspace = true, optional = true } [target.'cfg(target_os = "linux")'.dependencies] io-uring = "0.6.4" + +[target.'cfg(unix)'.dependencies] libc = "0.2.148" +[target.'cfg(target_os = "windows")'.dependencies.windows-sys] +version = "0.61.2" +features = [ + "Win32_Foundation", + "Win32_Security", + "Win32_Storage", + "Win32_Storage_FileSystem", + "Win32_System_IO", + "Win32_System_Threading", + "Win32_System_WindowsProgramming", + "Win32_System_ProcessStatus", + "Win32_System_SystemInformation" +] + [dev-dependencies] rstest.workspace = true tempfile.workspace = true diff --git a/diskann-disk/benches/benchmarks/aligned_file_reader_bench.rs b/diskann-disk/benches/benchmarks/aligned_file_reader_bench.rs index 0e166bc4f..f5504307a 100644 --- a/diskann-disk/benches/benchmarks/aligned_file_reader_bench.rs +++ b/diskann-disk/benches/benchmarks/aligned_file_reader_bench.rs @@ -5,7 +5,7 @@ use std::time::Duration; use criterion::Criterion; -use diskann_disk::utils::aligned_file_reader::{ +use diskann_disk::search::provider::aligned_file_reader::{ traits::{AlignedFileReader, AlignedReaderFactory}, AlignedFileReaderFactory, AlignedRead, }; diff --git a/diskann-disk/benches/benchmarks_iai/aligned_file_reader_bench_iai.rs b/diskann-disk/benches/benchmarks_iai/aligned_file_reader_bench_iai.rs index 215597684..7da3f88b3 100644 --- a/diskann-disk/benches/benchmarks_iai/aligned_file_reader_bench_iai.rs +++ b/diskann-disk/benches/benchmarks_iai/aligned_file_reader_bench_iai.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use diskann_disk::utils::aligned_file_reader::{ +use diskann_disk::search::provider::aligned_file_reader::{ traits::{AlignedFileReader, AlignedReaderFactory}, AlignedFileReaderFactory, AlignedRead, }; diff --git a/diskann-disk/src/build/builder/core.rs b/diskann-disk/src/build/builder/core.rs index 2fb24c07a..b01528a03 100644 --- a/diskann-disk/src/build/builder/core.rs +++ b/diskann-disk/src/build/builder/core.rs @@ -626,17 +626,20 @@ impl<'a> MergedVamanaIndexWorkflow<'a> { pub(crate) mod disk_index_builder_tests { use std::{io::Read, sync::Arc}; - use crate::test_utils::{GraphDataF32VectorU32Data, GraphDataF32VectorUnitData}; + use crate::{ + test_utils::{GraphDataF32VectorU32Data, GraphDataF32VectorUnitData}, + utils::instrumentation::Timer, + }; use diskann::{ graph::config, utils::{IntoUsize, VectorRepr, ONE}, ANNResult, }; use diskann_providers::storage::VirtualStorageProvider; - use diskann_providers::{ - storage::{get_compressed_pq_file, get_disk_index_file, get_pq_pivot_file}, - utils::Timer, + use diskann_providers::storage::{ + get_compressed_pq_file, get_disk_index_file, get_pq_pivot_file, }; + use diskann_utils::test_data_root; use diskann_vector::{ distance::Metric::{self, L2}, @@ -651,11 +654,11 @@ pub(crate) mod disk_index_builder_tests { data_model::{CachingStrategy, GraphHeader}, disk_index_build_parameter::{DiskIndexBuildParameters, MemoryBudget, NumPQChunks}, search::provider::{ - disk_provider::DiskIndexSearcher, + aligned_file_reader::VirtualAlignedReaderFactory, disk_provider::DiskIndexSearcher, disk_vertex_provider_factory::DiskVertexProviderFactory, }, storage::disk_index_reader::DiskIndexReader, - utils::{QueryStatistics, VirtualAlignedReaderFactory}, + utils::QueryStatistics, }; const DEFAULT_DISK_SECTOR_LEN: usize = 4096; pub const TEST_DATA_FILE: &str = "/sift/siftsmall_learn_256pts.fbin"; diff --git a/diskann-disk/src/utils/aligned_file_reader/aligned_read.rs b/diskann-disk/src/search/provider/aligned_file_reader/aligned_read.rs similarity index 100% rename from diskann-disk/src/utils/aligned_file_reader/aligned_read.rs rename to diskann-disk/src/search/provider/aligned_file_reader/aligned_read.rs diff --git a/diskann-disk/src/utils/aligned_file_reader/aligned_file_reader_factory.rs b/diskann-disk/src/search/provider/aligned_file_reader/factory/file.rs similarity index 91% rename from diskann-disk/src/utils/aligned_file_reader/aligned_file_reader_factory.rs rename to diskann-disk/src/search/provider/aligned_file_reader/factory/file.rs index 3674765df..2474ea3db 100644 --- a/diskann-disk/src/utils/aligned_file_reader/aligned_file_reader_factory.rs +++ b/diskann-disk/src/search/provider/aligned_file_reader/factory/file.rs @@ -6,12 +6,12 @@ use diskann::ANNResult; #[cfg(all(not(miri), target_os = "linux"))] -use super::LinuxAlignedFileReader; +use crate::search::provider::aligned_file_reader::reader::LinuxAlignedFileReader; #[cfg(any(miri, target_os = "macos"))] -use super::StorageProviderAlignedFileReader; +use crate::search::provider::aligned_file_reader::reader::StorageProviderAlignedFileReader; #[cfg(all(not(miri), target_os = "windows"))] -use super::WindowsAlignedFileReader; -use crate::utils::aligned_file_reader::traits::AlignedReaderFactory; +use crate::search::provider::aligned_file_reader::reader::WindowsAlignedFileReader; +use crate::search::provider::aligned_file_reader::traits::AlignedReaderFactory; #[cfg(any(miri, target_os = "macos"))] use diskann_providers::storage::FileStorageProvider; diff --git a/diskann-disk/src/search/provider/aligned_file_reader/factory/mod.rs b/diskann-disk/src/search/provider/aligned_file_reader/factory/mod.rs new file mode 100644 index 000000000..d928957ed --- /dev/null +++ b/diskann-disk/src/search/provider/aligned_file_reader/factory/mod.rs @@ -0,0 +1,14 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Concrete [`AlignedReaderFactory`](super::traits::AlignedReaderFactory) implementations. + +mod file; +pub use file::AlignedFileReaderFactory; + +#[cfg(test)] +mod virtual_storage; +#[cfg(test)] +pub(crate) use virtual_storage::VirtualAlignedReaderFactory; diff --git a/diskann-disk/src/utils/aligned_file_reader/virtual_aligned_reader_factory.rs b/diskann-disk/src/search/provider/aligned_file_reader/factory/virtual_storage.rs similarity index 91% rename from diskann-disk/src/utils/aligned_file_reader/virtual_aligned_reader_factory.rs rename to diskann-disk/src/search/provider/aligned_file_reader/factory/virtual_storage.rs index 53c049348..12cd8436c 100644 --- a/diskann-disk/src/utils/aligned_file_reader/virtual_aligned_reader_factory.rs +++ b/diskann-disk/src/search/provider/aligned_file_reader/factory/virtual_storage.rs @@ -9,7 +9,9 @@ use diskann::ANNResult; use diskann_providers::storage::VirtualStorageProvider; use vfs::{FileSystem, MemoryFS}; -use super::{traits::AlignedReaderFactory, StorageProviderAlignedFileReader}; +use crate::search::provider::aligned_file_reader::{ + reader::StorageProviderAlignedFileReader, traits::AlignedReaderFactory, +}; pub struct VirtualAlignedReaderFactory { pub file_path: String, diff --git a/diskann-disk/src/search/provider/aligned_file_reader/mod.rs b/diskann-disk/src/search/provider/aligned_file_reader/mod.rs new file mode 100644 index 000000000..2dec15175 --- /dev/null +++ b/diskann-disk/src/search/provider/aligned_file_reader/mod.rs @@ -0,0 +1,24 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +pub mod traits; + +mod aligned_read; +pub use aligned_read::{AlignedRead, Alignment, A1, A512}; + +#[cfg(all(not(miri), any(target_os = "linux", target_os = "windows")))] +mod platform; + +mod reader; +#[cfg(all(not(miri), target_os = "linux"))] +pub use reader::LinuxAlignedFileReader; +pub use reader::StorageProviderAlignedFileReader; +#[cfg(all(not(miri), target_os = "windows"))] +pub use reader::WindowsAlignedFileReader; + +mod factory; +pub use factory::AlignedFileReaderFactory; +#[cfg(test)] +pub(crate) use factory::VirtualAlignedReaderFactory; diff --git a/diskann-platform/src/linux/ssd_io_context.rs b/diskann-disk/src/search/provider/aligned_file_reader/platform/linux.rs similarity index 66% rename from diskann-platform/src/linux/ssd_io_context.rs rename to diskann-disk/src/search/provider/aligned_file_reader/platform/linux.rs index 2ef6b52b6..fce8313ea 100644 --- a/diskann-platform/src/linux/ssd_io_context.rs +++ b/diskann-disk/src/search/provider/aligned_file_reader/platform/linux.rs @@ -9,6 +9,9 @@ use io_uring::IoUring; // The IOContext struct for disk I/O. One for each thread. pub struct IOContext { + // Kept alive so the file descriptor registered with the io_uring ring stays valid; + // the fd is used via `register_files`, so this field is never read directly. + #[allow(dead_code)] pub file_handle: File, pub ring: IoUring, } diff --git a/diskann-disk/src/search/provider/aligned_file_reader/platform/mod.rs b/diskann-disk/src/search/provider/aligned_file_reader/platform/mod.rs new file mode 100644 index 000000000..cd8a1ce9c --- /dev/null +++ b/diskann-disk/src/search/provider/aligned_file_reader/platform/mod.rs @@ -0,0 +1,19 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Platform-specific I/O primitives backing the native aligned file readers. + +#[cfg(target_os = "linux")] +mod linux; +#[cfg(target_os = "linux")] +pub use linux::IOContext; + +#[cfg(target_os = "windows")] +mod windows; +#[cfg(target_os = "windows")] +pub use windows::{ + get_queued_completion_status, read_file_to_slice, FileHandle, IOCompletionPort, IOContext, + DWORD, OVERLAPPED, ULONG_PTR, +}; diff --git a/diskann-platform/src/win/file_handle.rs b/diskann-disk/src/search/provider/aligned_file_reader/platform/windows/file_handle.rs similarity index 63% rename from diskann-platform/src/win/file_handle.rs rename to diskann-disk/src/search/provider/aligned_file_reader/platform/windows/file_handle.rs index 99467e38d..7b33191f5 100644 --- a/diskann-platform/src/win/file_handle.rs +++ b/diskann-disk/src/search/provider/aligned_file_reader/platform/windows/file_handle.rs @@ -5,12 +5,10 @@ use std::{ffi::CString, io, ptr}; use windows_sys::Win32::{ - Foundation::{ - CloseHandle, GetLastError, GENERIC_READ, GENERIC_WRITE, HANDLE, INVALID_HANDLE_VALUE, - }, + Foundation::{CloseHandle, GetLastError, GENERIC_READ, HANDLE, INVALID_HANDLE_VALUE}, Storage::FileSystem::{ CreateFileA, FILE_FLAG_NO_BUFFERING, FILE_FLAG_OVERLAPPED, FILE_FLAG_RANDOM_ACCESS, - FILE_SHARE_DELETE, FILE_SHARE_READ, FILE_SHARE_WRITE, OPEN_EXISTING, + FILE_SHARE_READ, OPEN_EXISTING, }, }; @@ -18,43 +16,6 @@ use super::DWORD; pub const FILE_ATTRIBUTE_READONLY: DWORD = 0x00000001; -/// `AccessMode` determines how a file can be accessed. -/// These modes are used when creating or opening a file to decide what operations are allowed -/// to be performed on the file. -/// -/// # Variants -/// -/// - `Read`: The file is opened in read-only mode. -/// -/// - `Write`: The file is opened in write-only mode. -/// -/// - `ReadWrite`: The file is opened for both reading and writing. -pub enum AccessMode { - Read, - Write, - ReadWrite, -} - -/// `ShareMode` determines how a file can be shared. -/// -/// These modes are used when creating or opening a file to decide what operations other -/// opening instances of the file can perform on it. -/// # Variants -/// - `None`: Prevents other processes from opening a file if they request delete, -/// read, or write access. -/// -/// - `Read`: Allows subsequent open operations on the same file to request read access. -/// -/// - `Write`: Allows subsequent open operations on the same file file to request write access. -/// -/// - `Delete`: Allows subsequent open operations on the same file file to request delete access. -pub enum ShareMode { - None, - Read, - Write, - Delete, -} - /// # Windows File Handle Wrapper /// /// Introduces a Rust-friendly wrapper around the native Windows `HANDLE` object, `FileHandle`. @@ -86,7 +47,7 @@ unsafe impl Send for FileHandle {} unsafe impl Sync for FileHandle {} impl FileHandle { - /// Creates a new `FileHandle` by opening an existing file with the given access and shared mode. + /// Creates a new `FileHandle` by opening an existing file read-only with read sharing. /// /// This function is marked unsafe because it creates a raw pointer to the filename and try to create /// a Windows `HANDLE` object without checking if you have sufficient permissions. @@ -94,41 +55,25 @@ impl FileHandle { /// # Safety /// /// Ensure that the file specified by `file_name` is valid and the calling process has - /// sufficient permissions to perform the specified `access_mode` and `share_mode` operations. + /// sufficient permissions to open it for reading. /// /// # Parameters /// /// - `file_name`: The name of the file. - /// - `access_mode`: The access mode to be used for the file. - /// - `share_mode`: The share mode to be used for the file /// /// # Errors /// This function will return an error if the `file_name` is invalid or if the file cannot - /// be opened with the specified `access_mode` and `share_mode`. - pub unsafe fn new( - file_name: &str, - access_mode: AccessMode, - share_mode: ShareMode, - ) -> io::Result { + /// be opened for reading. + pub unsafe fn new(file_name: &str) -> io::Result { let file_name_c = CString::new(file_name).map_err(|_| { io::Error::new( io::ErrorKind::InvalidData, - format!("Invalid file name. {}", file_name), + format!("Invalid file name. {file_name}"), ) })?; - let dw_desired_access = match access_mode { - AccessMode::Read => GENERIC_READ, - AccessMode::Write => GENERIC_WRITE, - AccessMode::ReadWrite => GENERIC_READ | GENERIC_WRITE, - }; - - let dw_share_mode = match share_mode { - ShareMode::None => 0, - ShareMode::Read => FILE_SHARE_READ, - ShareMode::Write => FILE_SHARE_WRITE, - ShareMode::Delete => FILE_SHARE_DELETE, - }; + let dw_desired_access = GENERIC_READ; + let dw_share_mode = FILE_SHARE_READ; let dw_flags_and_attributes = FILE_ATTRIBUTE_READONLY | FILE_FLAG_NO_BUFFERING @@ -166,7 +111,7 @@ impl Drop for FileHandle { if result == 0 { let error_code = unsafe { GetLastError() }; let error = io::Error::from_raw_os_error(error_code as i32); - tracing::warn!("Error when dropping FileHandle: {:?}", error); + tracing::warn!("Error when dropping FileHandle: {error:?}"); } } } @@ -196,9 +141,7 @@ mod tests { let path = Path::new(dummy_file_path); { - let file_handle = unsafe { - FileHandle::new(path.to_str().unwrap(), AccessMode::Read, ShareMode::Read) - }; + let file_handle = unsafe { FileHandle::new(path.to_str().unwrap()) }; // Check that the file handle is valid assert!(file_handle.is_ok()); @@ -207,15 +150,14 @@ mod tests { // Try to delete the file. If the handle was correctly dropped, this should succeed. match std::fs::remove_file(dummy_file_path) { Ok(()) => (), // File was deleted successfully, which means the handle was closed. - Err(e) => panic!("Failed to delete file: {}", e), // Failed to delete the file, likely because the handle is still open. + Err(e) => panic!("Failed to delete file: {e}"), // Failed to delete the file, likely because the handle is still open. } } #[test] fn test_file_not_found() { let path = Path::new("non_existent_file.txt"); - let file_handle = - unsafe { FileHandle::new(path.to_str().unwrap(), AccessMode::Read, ShareMode::Read) }; + let file_handle = unsafe { FileHandle::new(path.to_str().unwrap()) }; // Check that opening a non-existent file returns an error assert!(file_handle.is_err()); diff --git a/diskann-platform/src/win/file_io.rs b/diskann-disk/src/search/provider/aligned_file_reader/platform/windows/file_io.rs similarity index 95% rename from diskann-platform/src/win/file_io.rs rename to diskann-disk/src/search/provider/aligned_file_reader/platform/windows/file_io.rs index b309be974..cf0deb7a6 100644 --- a/diskann-platform/src/win/file_io.rs +++ b/diskann-disk/src/search/provider/aligned_file_reader/platform/windows/file_io.rs @@ -15,8 +15,9 @@ use windows_sys::Win32::{ System::IO::{GetQueuedCompletionStatus, OVERLAPPED}, }; +use super::file_handle::FileHandle; +use super::io_completion_port::IOCompletionPort; use super::{DWORD, ULONG_PTR}; -use crate::{FileHandle, IOCompletionPort}; /// Asynchronously queue a read request from a file into a buffer slice. /// @@ -133,7 +134,6 @@ mod tests { use std::{fs::File, io::Write, path::Path}; use super::*; - use crate::win::file_handle::{AccessMode, ShareMode}; #[test] fn test_read_file_to_slice() { @@ -147,10 +147,7 @@ mod tests { let mut buffer: [u8; 512] = [0; 512]; let mut overlapped = unsafe { std::mem::zeroed::() }; { - let file_handle = unsafe { - FileHandle::new(path.to_str().unwrap(), AccessMode::Read, ShareMode::Read) - } - .unwrap(); + let file_handle = unsafe { FileHandle::new(path.to_str().unwrap()) }.unwrap(); // Call the function under test let result = diff --git a/diskann-platform/src/win/io_completion_port.rs b/diskann-disk/src/search/provider/aligned_file_reader/platform/windows/io_completion_port.rs similarity index 93% rename from diskann-platform/src/win/io_completion_port.rs rename to diskann-disk/src/search/provider/aligned_file_reader/platform/windows/io_completion_port.rs index 6fce4afa9..6f188eeba 100644 --- a/diskann-platform/src/win/io_completion_port.rs +++ b/diskann-disk/src/search/provider/aligned_file_reader/platform/windows/io_completion_port.rs @@ -12,8 +12,8 @@ use windows_sys::Win32::{ System::IO::CreateIoCompletionPort, }; +use super::file_handle::FileHandle; use super::{DWORD, ULONG_PTR}; -use crate::FileHandle; /// This module provides a safe and idiomatic Rust interface over the IOCompletionPort handle and associated Windows API functions. /// This struct represents an I/O completion port, which is an object used in asynchronous I/O operations on Windows. @@ -97,7 +97,7 @@ impl Drop for IOCompletionPort { let error_code = unsafe { GetLastError() }; let error = io::Error::from_raw_os_error(error_code as i32); - tracing::warn!("Error when dropping IOCompletionPort: {:?}", error); + tracing::warn!("Error when dropping IOCompletionPort: {error:?}"); } } } @@ -115,13 +115,12 @@ impl Default for IOCompletionPort { #[cfg(test)] mod tests { use super::*; - use crate::win::file_handle::{AccessMode, ShareMode}; #[test] fn create_io_completion_port() { let file_name = "../test_data/delete_set_50pts.bin"; - let file_handle = unsafe { FileHandle::new(file_name, AccessMode::Read, ShareMode::Read) } - .expect("Failed to create file handle."); + let file_handle = + unsafe { FileHandle::new(file_name) }.expect("Failed to create file handle."); let io_completion_port = IOCompletionPort::new(&file_handle, None, 0, 0); @@ -134,8 +133,8 @@ mod tests { #[test] fn drop_io_completion_port() { let file_name = "../test_data/delete_set_50pts.bin"; - let file_handle = unsafe { FileHandle::new(file_name, AccessMode::Read, ShareMode::Read) } - .expect("Failed to create file handle."); + let file_handle = + unsafe { FileHandle::new(file_name) }.expect("Failed to create file handle."); let io_completion_port = IOCompletionPort::new(&file_handle, None, 0, 0) .expect("Failed to create IOCompletionPort."); diff --git a/diskann-platform/src/win/mod.rs b/diskann-disk/src/search/provider/aligned_file_reader/platform/windows/mod.rs similarity index 50% rename from diskann-platform/src/win/mod.rs rename to diskann-disk/src/search/provider/aligned_file_reader/platform/windows/mod.rs index e7ed38675..edd6af05f 100644 --- a/diskann-platform/src/win/mod.rs +++ b/diskann-disk/src/search/provider/aligned_file_reader/platform/windows/mod.rs @@ -2,28 +2,20 @@ * Copyright (c) Microsoft Corporation. * Licensed under the MIT license. */ -#![cfg_attr( - not(test), - warn(clippy::panic, clippy::unwrap_used, clippy::expect_used) -)] -pub use file_handle::{AccessMode, FileHandle, ShareMode}; +pub use file_handle::FileHandle; pub use file_io::{get_queued_completion_status, read_file_to_slice}; pub use io_completion_port::IOCompletionPort; -pub use perf::{ - get_number_of_processors, get_peak_workingset_size, get_process_cycle_time, get_process_time, - get_system_time, -}; -pub use thread_safe_handle::ThreadSafeHandle; +pub use ssd_io_context::IOContext; mod file_handle; mod file_io; mod io_completion_port; -mod perf; -pub mod ssd_io_context; -mod thread_safe_handle; +mod ssd_io_context; #[allow(non_camel_case_types)] pub type ULONG_PTR = usize; +#[allow(clippy::upper_case_acronyms)] pub type DWORD = u32; +#[allow(clippy::upper_case_acronyms)] pub type OVERLAPPED = windows_sys::Win32::System::IO::OVERLAPPED; diff --git a/diskann-disk/src/search/provider/aligned_file_reader/platform/windows/ssd_io_context.rs b/diskann-disk/src/search/provider/aligned_file_reader/platform/windows/ssd_io_context.rs new file mode 100644 index 000000000..d45ed51ba --- /dev/null +++ b/diskann-disk/src/search/provider/aligned_file_reader/platform/windows/ssd_io_context.rs @@ -0,0 +1,19 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::{FileHandle, IOCompletionPort}; + +// The IOContext struct for disk I/O. One for each thread. +#[derive(Default)] +pub struct IOContext { + pub file_handle: FileHandle, + pub io_completion_port: IOCompletionPort, +} + +impl IOContext { + pub fn new() -> Self { + Self::default() + } +} diff --git a/diskann-disk/src/utils/aligned_file_reader/linux_aligned_file_reader.rs b/diskann-disk/src/search/provider/aligned_file_reader/reader/linux.rs similarity index 99% rename from diskann-disk/src/utils/aligned_file_reader/linux_aligned_file_reader.rs rename to diskann-disk/src/search/provider/aligned_file_reader/reader/linux.rs index 9ebfb85dc..a8cf8184c 100644 --- a/diskann-disk/src/utils/aligned_file_reader/linux_aligned_file_reader.rs +++ b/diskann-disk/src/search/provider/aligned_file_reader/reader/linux.rs @@ -8,11 +8,12 @@ use std::{ }; use diskann::{ANNError, ANNResult}; -use diskann_platform::ssd_io_context::IOContext; use io_uring::IoUring; use libc; -use crate::utils::aligned_file_reader::{traits::AlignedFileReader, AlignedRead, A512}; +use crate::search::provider::aligned_file_reader::{ + platform::IOContext, traits::AlignedFileReader, AlignedRead, A512, +}; pub const MAX_IO_CONCURRENCY: usize = 128; diff --git a/diskann-disk/src/search/provider/aligned_file_reader/reader/mod.rs b/diskann-disk/src/search/provider/aligned_file_reader/reader/mod.rs new file mode 100644 index 000000000..c7dd054e8 --- /dev/null +++ b/diskann-disk/src/search/provider/aligned_file_reader/reader/mod.rs @@ -0,0 +1,19 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Concrete [`AlignedFileReader`](super::traits::AlignedFileReader) implementations. + +cfg_if::cfg_if! { + if #[cfg(all(not(miri), target_os = "linux"))] { + mod linux; + pub use linux::LinuxAlignedFileReader; + } else if #[cfg(all(not(miri), target_os = "windows"))] { + mod windows; + pub use windows::WindowsAlignedFileReader; + } +} + +mod storage_provider; +pub use storage_provider::StorageProviderAlignedFileReader; diff --git a/diskann-disk/src/utils/aligned_file_reader/storage_provider_aligned_file_reader.rs b/diskann-disk/src/search/provider/aligned_file_reader/reader/storage_provider.rs similarity index 97% rename from diskann-disk/src/utils/aligned_file_reader/storage_provider_aligned_file_reader.rs rename to diskann-disk/src/search/provider/aligned_file_reader/reader/storage_provider.rs index 631934a3f..33ad51e7d 100644 --- a/diskann-disk/src/utils/aligned_file_reader/storage_provider_aligned_file_reader.rs +++ b/diskann-disk/src/search/provider/aligned_file_reader/reader/storage_provider.rs @@ -9,8 +9,7 @@ use diskann::ANNResult; use diskann_providers::storage::StorageReadProvider; use tracing::info; -use super::traits::AlignedFileReader; -use crate::utils::aligned_file_reader::{AlignedRead, A1}; +use crate::search::provider::aligned_file_reader::{traits::AlignedFileReader, AlignedRead, A1}; pub struct StorageProviderAlignedFileReader { data: Vec, diff --git a/diskann-disk/src/utils/aligned_file_reader/windows_aligned_file_reader.rs b/diskann-disk/src/search/provider/aligned_file_reader/reader/windows.rs similarity index 96% rename from diskann-disk/src/utils/aligned_file_reader/windows_aligned_file_reader.rs rename to diskann-disk/src/search/provider/aligned_file_reader/reader/windows.rs index 543a06d48..c54fb036a 100644 --- a/diskann-disk/src/utils/aligned_file_reader/windows_aligned_file_reader.rs +++ b/diskann-disk/src/search/provider/aligned_file_reader/reader/windows.rs @@ -5,13 +5,15 @@ use std::{ptr, thread, time::Duration}; use diskann::{ANNError, ANNResult}; -use diskann_platform::{ - get_queued_completion_status, read_file_to_slice, ssd_io_context::IOContext, AccessMode, - FileHandle, IOCompletionPort, ShareMode, DWORD, OVERLAPPED, ULONG_PTR, -}; -use super::traits::AlignedFileReader; -use crate::utils::aligned_file_reader::{AlignedRead, A512}; +use crate::search::provider::aligned_file_reader::{ + platform::{ + get_queued_completion_status, read_file_to_slice, FileHandle, IOCompletionPort, IOContext, + DWORD, OVERLAPPED, ULONG_PTR, + }, + traits::AlignedFileReader, + AlignedRead, A512, +}; pub const MAX_IO_CONCURRENCY: usize = 128; pub const IO_COMPLETION_TIMEOUT: DWORD = u32::MAX; // Infinite timeout. @@ -33,7 +35,7 @@ impl WindowsAlignedFileReader { pub fn new(fname: &str) -> ANNResult { let mut io_context = IOContext::new(); tracing::debug!("Creating file handle for {}", fname); - match unsafe { FileHandle::new(fname, AccessMode::Read, ShareMode::Read) } { + match unsafe { FileHandle::new(fname) } { Ok(file_handle) => io_context.file_handle = file_handle, Err(err) => { return Err(ANNError::log_io_error(err)); @@ -127,7 +129,7 @@ mod tests { use serde::{Deserialize, Serialize}; use super::*; - use crate::utils::aligned_file_reader::{AlignedRead, A512}; + use crate::search::provider::aligned_file_reader::{AlignedRead, A512}; use diskann_quantization::alloc::{AlignedAllocator, Poly}; fn test_index_path() -> String { diff --git a/diskann-disk/src/utils/aligned_file_reader/traits/aligned_file_reader.rs b/diskann-disk/src/search/provider/aligned_file_reader/traits/aligned_file_reader.rs similarity index 91% rename from diskann-disk/src/utils/aligned_file_reader/traits/aligned_file_reader.rs rename to diskann-disk/src/search/provider/aligned_file_reader/traits/aligned_file_reader.rs index f8715e509..24acc2cbf 100644 --- a/diskann-disk/src/utils/aligned_file_reader/traits/aligned_file_reader.rs +++ b/diskann-disk/src/search/provider/aligned_file_reader/traits/aligned_file_reader.rs @@ -5,7 +5,7 @@ use diskann::ANNResult; -use crate::utils::aligned_file_reader::{aligned_read::Alignment, AlignedRead}; +use crate::search::provider::aligned_file_reader::{aligned_read::Alignment, AlignedRead}; pub trait AlignedFileReader: Send + Sync { /// Alignment requirement applied to every `AlignedRead` diff --git a/diskann-disk/src/utils/aligned_file_reader/traits/aligned_reader_factory.rs b/diskann-disk/src/search/provider/aligned_file_reader/traits/aligned_reader_factory.rs similarity index 100% rename from diskann-disk/src/utils/aligned_file_reader/traits/aligned_reader_factory.rs rename to diskann-disk/src/search/provider/aligned_file_reader/traits/aligned_reader_factory.rs diff --git a/diskann-disk/src/utils/aligned_file_reader/traits/mod.rs b/diskann-disk/src/search/provider/aligned_file_reader/traits/mod.rs similarity index 75% rename from diskann-disk/src/utils/aligned_file_reader/traits/mod.rs rename to diskann-disk/src/search/provider/aligned_file_reader/traits/mod.rs index ed21d66f0..6c8854af7 100644 --- a/diskann-disk/src/utils/aligned_file_reader/traits/mod.rs +++ b/diskann-disk/src/search/provider/aligned_file_reader/traits/mod.rs @@ -3,9 +3,8 @@ * Licensed under the MIT license. */ -#[allow(clippy::module_inception)] mod aligned_file_reader; pub use aligned_file_reader::AlignedFileReader; -pub mod aligned_reader_factory; +mod aligned_reader_factory; pub use aligned_reader_factory::AlignedReaderFactory; diff --git a/diskann-disk/src/search/provider/cached_disk_vertex_provider.rs b/diskann-disk/src/search/provider/cached_disk_vertex_provider.rs index 3df2edae3..df96119d1 100644 --- a/diskann-disk/src/search/provider/cached_disk_vertex_provider.rs +++ b/diskann-disk/src/search/provider/cached_disk_vertex_provider.rs @@ -8,12 +8,17 @@ use std::sync::Arc; use crate::data_model::GraphDataType; use diskann::{ANNError, ANNResult}; -use crate::utils::aligned_file_reader::traits::AlignedFileReader; use hashbrown::HashMap; use crate::{ data_model::{Cache, GraphHeader}, - search::{provider::disk_vertex_provider::DiskVertexProvider, traits::VertexProvider}, + search::{ + provider::{ + aligned_file_reader::traits::AlignedFileReader, + disk_vertex_provider::DiskVertexProvider, + }, + traits::VertexProvider, + }, }; pub struct CachedDiskVertexProvider diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 8e95dfa4b..64504da5f 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -49,11 +49,13 @@ use crate::{ data_model::{CachingStrategy, GraphHeader}, filter_parameter::{default_vector_filter, VectorFilter}, search::{ - provider::disk_vertex_provider_factory::DiskVertexProviderFactory, + provider::{ + aligned_file_reader::AlignedFileReaderFactory, + disk_vertex_provider_factory::DiskVertexProviderFactory, + }, traits::{VertexProvider, VertexProviderFactory}, }, storage::{api::AsyncDiskLoadContext, disk_index_reader::DiskIndexReader}, - utils::AlignedFileReaderFactory, utils::QueryStatistics, }; @@ -138,20 +140,17 @@ where ); let graph_header = { - let aligned_reader_factory = AlignedFileReaderFactory::new(get_disk_index_file( - ctx.quant_load_context.metadata.prefix(), - )); - let caching_strategy = if ctx.num_nodes_to_cache > 0 { CachingStrategy::StaticCacheWithBfsNodes(ctx.num_nodes_to_cache) } else { CachingStrategy::None }; - let vertex_provider_factory = DiskVertexProviderFactory::::new( - aligned_reader_factory, - caching_strategy, - )?; + let vertex_provider_factory = + DiskVertexProviderFactory::::from_disk_index_path( + get_disk_index_file(ctx.quant_load_context.metadata.prefix()), + caching_strategy, + )?; VertexProviderFactory::get_header(&vertex_provider_factory)? }; @@ -1155,7 +1154,8 @@ mod disk_provider_tests { use super::*; use crate::{ build::builder::core::disk_index_builder_tests::{IndexBuildFixture, TestParams}, - utils::{QueryStatistics, VirtualAlignedReaderFactory}, + search::provider::aligned_file_reader::VirtualAlignedReaderFactory, + utils::QueryStatistics, }; const TEST_INDEX_PREFIX_128DIM: &str = diff --git a/diskann-disk/src/search/provider/disk_sector_graph.rs b/diskann-disk/src/search/provider/disk_sector_graph.rs index 7d44c435d..574e3e716 100644 --- a/diskann-disk/src/search/provider/disk_sector_graph.rs +++ b/diskann-disk/src/search/provider/disk_sector_graph.rs @@ -12,7 +12,7 @@ use diskann_quantization::alloc::{AlignedAllocator, Poly}; use crate::{ data_model::GraphHeader, - utils::aligned_file_reader::{traits::AlignedFileReader, AlignedRead, Alignment}, + search::provider::aligned_file_reader::{traits::AlignedFileReader, AlignedRead, Alignment}, }; const DEFAULT_DISK_SECTOR_LEN: usize = 4096; @@ -198,13 +198,15 @@ impl Deref for DiskSectorGraph String { test_data_root() diff --git a/diskann-disk/src/search/provider/disk_vertex_provider.rs b/diskann-disk/src/search/provider/disk_vertex_provider.rs index f83759533..80b701752 100644 --- a/diskann-disk/src/search/provider/disk_vertex_provider.rs +++ b/diskann-disk/src/search/provider/disk_vertex_provider.rs @@ -12,8 +12,12 @@ use hashbrown::HashMap; use crate::{ data_model::GraphHeader, - search::{provider::disk_sector_graph::DiskSectorGraph, traits::VertexProvider}, - utils::aligned_file_reader::traits::AlignedFileReader, + search::{ + provider::{ + aligned_file_reader::traits::AlignedFileReader, disk_sector_graph::DiskSectorGraph, + }, + traits::VertexProvider, + }, }; struct Offsets { @@ -282,12 +286,12 @@ mod disk_vertex_provider_tests { disk_index_build_parameter::{ DiskIndexBuildParameters, MemoryBudget, NumPQChunks, DISK_SECTOR_LEN, }, + search::provider::aligned_file_reader::VirtualAlignedReaderFactory, search::{ provider::disk_vertex_provider_factory::DiskVertexProviderFactory, traits::{VertexProvider, VertexProviderFactory}, }, storage::DiskIndexWriter, - utils::VirtualAlignedReaderFactory, QuantizationType, }; diff --git a/diskann-disk/src/search/provider/disk_vertex_provider_factory.rs b/diskann-disk/src/search/provider/disk_vertex_provider_factory.rs index 13d66e443..f57dc12c8 100644 --- a/diskann-disk/src/search/provider/disk_vertex_provider_factory.rs +++ b/diskann-disk/src/search/provider/disk_vertex_provider_factory.rs @@ -15,6 +15,10 @@ use tracing::info; use crate::{ data_model::{Cache, CachingStrategy, GraphHeader}, + search::provider::aligned_file_reader::{ + traits::{AlignedFileReader, AlignedReaderFactory}, + AlignedFileReaderFactory, AlignedRead, + }, search::{ provider::{ cached_disk_vertex_provider::CachedDiskVertexProvider, @@ -22,10 +26,6 @@ use crate::{ }, traits::{VertexProvider, VertexProviderFactory}, }, - utils::aligned_file_reader::{ - traits::{AlignedFileReader, AlignedReaderFactory}, - AlignedRead, - }, }; const DEFAULT_DISK_SECTOR_LEN: usize = 4096; @@ -98,6 +98,22 @@ where } } +impl> + DiskVertexProviderFactory +{ + /// Creates a production `DiskVertexProviderFactory` that reads the on-disk index at + /// `disk_index_path` using the platform's native aligned file reader. + pub fn from_disk_index_path( + disk_index_path: String, + caching_strategy: CachingStrategy, + ) -> ANNResult { + Self::new( + AlignedFileReaderFactory::new(disk_index_path), + caching_strategy, + ) + } +} + impl, ReaderFactory: AlignedReaderFactory> DiskVertexProviderFactory { @@ -239,8 +255,10 @@ impl, ReaderFactory: AlignedReaderFactor #[cfg(test)] pub(crate) mod tests { use super::*; - use crate::test_utils::GraphDataF32VectorUnitData; - use crate::utils::VirtualAlignedReaderFactory; + use crate::{ + search::provider::aligned_file_reader::VirtualAlignedReaderFactory, + test_utils::GraphDataF32VectorUnitData, + }; use diskann_providers::storage::VirtualStorageProvider; use diskann_utils::test_data_root; use vfs::OverlayFS; diff --git a/diskann-disk/src/search/provider/mod.rs b/diskann-disk/src/search/provider/mod.rs index 69f697d84..eead8bd00 100644 --- a/diskann-disk/src/search/provider/mod.rs +++ b/diskann-disk/src/search/provider/mod.rs @@ -13,3 +13,5 @@ pub mod disk_provider; pub mod disk_sector_graph; pub mod disk_vertex_provider; pub mod disk_vertex_provider_factory; + +pub mod aligned_file_reader; diff --git a/diskann-disk/src/storage/quant/generator.rs b/diskann-disk/src/storage/quant/generator.rs index c14bf998a..514acbedc 100644 --- a/diskann-disk/src/storage/quant/generator.rs +++ b/diskann-disk/src/storage/quant/generator.rs @@ -9,9 +9,9 @@ use std::{ }; use diskann::{error::IntoANNResult, utils::VectorRepr, ANNError, ANNResult}; -use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; -use diskann_providers::utils::{ - load_metadata_from_file, BridgeErr, ParallelIteratorInPool, RayonThreadPoolRef, Timer, +use diskann_providers::{ + storage::{StorageReadProvider, StorageWriteProvider}, + utils::{load_metadata_from_file, BridgeErr, ParallelIteratorInPool, RayonThreadPoolRef}, }; use diskann_utils::{io::Metadata, views}; use rayon::iter::IndexedParallelIterator; @@ -23,6 +23,7 @@ use crate::{ continuation::{process_while_resource_is_available, ChunkingConfig}, }, storage::quant::compressor::{CompressionStage, QuantCompressor}, + utils::instrumentation::Timer, }; /// [`GeneratorContext`] defines parameters for vector quantization checkpoint state diff --git a/diskann-disk/src/storage/quant/pq/pq_generation.rs b/diskann-disk/src/storage/quant/pq/pq_generation.rs index 0e7b07fd4..5b915df02 100644 --- a/diskann-disk/src/storage/quant/pq/pq_generation.rs +++ b/diskann-disk/src/storage/quant/pq/pq_generation.rs @@ -13,14 +13,17 @@ use diskann_providers::{ GeneratePivotArguments, }, storage::PQStorage, - utils::{BridgeErr, RayonThreadPoolRef, Timer}, + utils::{BridgeErr, RayonThreadPoolRef}, }; use diskann_quantization::{product::TransposedTable, CompressInto}; use diskann_utils::views::MatrixBase; use diskann_vector::distance::Metric; use tracing::info; -use crate::storage::quant::compressor::{CompressionStage, QuantCompressor}; +use crate::{ + storage::quant::compressor::{CompressionStage, QuantCompressor}, + utils::instrumentation::Timer, +}; pub struct PQGenerationContext<'a, Storage> where diff --git a/diskann-disk/src/utils/aligned_file_reader/mod.rs b/diskann-disk/src/utils/aligned_file_reader/mod.rs deleted file mode 100644 index dfaa5e44e..000000000 --- a/diskann-disk/src/utils/aligned_file_reader/mod.rs +++ /dev/null @@ -1,34 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -pub mod traits; - -mod aligned_read; -pub use aligned_read::{AlignedRead, Alignment, A1, A512}; - -cfg_if::cfg_if! { - if #[cfg(all(not(miri), target_os = "linux"))] { - pub mod linux_aligned_file_reader; - pub use linux_aligned_file_reader::LinuxAlignedFileReader; - } else if #[cfg(all(not(miri), target_os = "windows"))] { - #[allow(clippy::module_inception)] - pub mod windows_aligned_file_reader; - pub use windows_aligned_file_reader::WindowsAlignedFileReader; - } - -} - -cfg_if::cfg_if! { - if #[cfg(any(feature = "virtual_storage", test))] { - pub mod virtual_aligned_reader_factory; - pub use virtual_aligned_reader_factory::VirtualAlignedReaderFactory; - } -} - -pub mod storage_provider_aligned_file_reader; -pub use storage_provider_aligned_file_reader::StorageProviderAlignedFileReader; - -pub mod aligned_file_reader_factory; -pub use aligned_file_reader_factory::AlignedFileReaderFactory; diff --git a/diskann-disk/src/utils/instrumentation/mod.rs b/diskann-disk/src/utils/instrumentation/mod.rs index 635b84955..98d44ac58 100644 --- a/diskann-disk/src/utils/instrumentation/mod.rs +++ b/diskann-disk/src/utils/instrumentation/mod.rs @@ -4,3 +4,6 @@ */ mod perf_logger; pub use perf_logger::{BuildMergedVamanaIndexCheckpoint, DiskIndexBuildCheckpoint, PerfLogger}; + +mod timer; +pub use timer::Timer; diff --git a/diskann-disk/src/utils/instrumentation/perf_logger.rs b/diskann-disk/src/utils/instrumentation/perf_logger.rs index 439ce25ae..bdd19a9a5 100644 --- a/diskann-disk/src/utils/instrumentation/perf_logger.rs +++ b/diskann-disk/src/utils/instrumentation/perf_logger.rs @@ -12,7 +12,7 @@ use opentelemetry::{ }; use tracing::info; -use diskann_providers::utils::Timer; +use super::Timer; /// Target for logging latency events. pub const LATENCY_LOG_TARGET: &str = "latency_event"; diff --git a/diskann-providers/src/utils/timer.rs b/diskann-disk/src/utils/instrumentation/timer.rs similarity index 90% rename from diskann-providers/src/utils/timer.rs rename to diskann-disk/src/utils/instrumentation/timer.rs index b81f7c9db..aa55d13a0 100644 --- a/diskann-providers/src/utils/timer.rs +++ b/diskann-disk/src/utils/instrumentation/timer.rs @@ -4,11 +4,32 @@ */ use std::time::{Duration, Instant}; -use diskann_platform::*; +#[cfg(target_os = "linux")] +mod linux; +#[cfg(target_os = "linux")] +use linux::{ + get_number_of_processors, get_peak_workingset_size, get_process_cycle_time, get_process_time, + get_system_time, +}; +#[cfg(target_os = "macos")] +mod macos; +#[cfg(target_os = "macos")] +use macos::{ + get_number_of_processors, get_peak_workingset_size, get_process_cycle_time, get_process_time, + get_system_time, +}; +#[cfg(target_os = "windows")] +mod windows; #[cfg(feature = "perf_test")] use opentelemetry::{ - KeyValue, global, - trace::{Tracer, get_active_span}, + global, + trace::{get_active_span, Tracer}, + KeyValue, +}; +#[cfg(target_os = "windows")] +use windows::{ + get_number_of_processors, get_peak_workingset_size, get_process_cycle_time, get_process_time, + get_system_time, }; #[derive(Clone)] diff --git a/diskann-platform/src/linux/perf.rs b/diskann-disk/src/utils/instrumentation/timer/linux.rs similarity index 100% rename from diskann-platform/src/linux/perf.rs rename to diskann-disk/src/utils/instrumentation/timer/linux.rs diff --git a/diskann-platform/src/macos/perf.rs b/diskann-disk/src/utils/instrumentation/timer/macos.rs similarity index 100% rename from diskann-platform/src/macos/perf.rs rename to diskann-disk/src/utils/instrumentation/timer/macos.rs diff --git a/diskann-platform/src/win/perf.rs b/diskann-disk/src/utils/instrumentation/timer/windows.rs similarity index 100% rename from diskann-platform/src/win/perf.rs rename to diskann-disk/src/utils/instrumentation/timer/windows.rs diff --git a/diskann-disk/src/utils/mod.rs b/diskann-disk/src/utils/mod.rs index 968ab2c61..0d67f9ed7 100644 --- a/diskann-disk/src/utils/mod.rs +++ b/diskann-disk/src/utils/mod.rs @@ -19,10 +19,5 @@ pub use kmeans::{k_means_clustering, k_meanspp_selecting_pivots, run_lloyds}; pub mod instrumentation; -pub mod aligned_file_reader; -pub use aligned_file_reader::AlignedFileReaderFactory; -#[cfg(any(feature = "virtual_storage", test))] -pub use aligned_file_reader::VirtualAlignedReaderFactory; - pub mod statistics; pub use statistics::{get_mean_stats, get_percentile_stats, get_sum_stats, QueryStatistics}; diff --git a/diskann-platform/Cargo.toml b/diskann-platform/Cargo.toml deleted file mode 100644 index f4459722d..000000000 --- a/diskann-platform/Cargo.toml +++ /dev/null @@ -1,44 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. -[package] -name = "diskann-platform" -edition.workspace = true -license.workspace = true -version.workspace = true -authors.workspace = true -description.workspace = true -repository.workspace = true - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -tracing.workspace = true - -[target.'cfg(target_os = "linux")'.dependencies] -io-uring = "0.6.4" - -[target.'cfg(unix)'.dependencies] -libc = "0.2.148" - -[target.'cfg(target_os = "windows")'.dependencies.windows-sys] -version = "0.59" -features = [ - "Win32_Foundation", - "Win32_Security", - "Win32_Storage", - "Win32_Storage_FileSystem", - "Win32_System_IO", - "Win32_System_Threading", - "Win32_System_WindowsProgramming", - "Win32_System_ProcessStatus", - "Win32_System_SystemInformation" -] - -[package.metadata.docs.rs] -default-target = "x86_64-pc-windows-msvc" -targets = ["x86_64-unknown-linux-gnu", "aarch64-pc-windows-msvc", "x86_64-pc-windows-msvc"] - -[lints] -# Commenting out for now. It would give warnings. -# clippy.undocumented_unsafe_blocks = "warn" -clippy.uninlined_format_args = "allow" diff --git a/diskann-platform/src/lib.rs b/diskann-platform/src/lib.rs deleted file mode 100644 index 91f878352..000000000 --- a/diskann-platform/src/lib.rs +++ /dev/null @@ -1,22 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -#[cfg(target_os = "windows")] -pub mod win; - -#[cfg(target_os = "windows")] -pub use win::*; - -#[cfg(target_os = "linux")] -pub mod linux; - -#[cfg(target_os = "linux")] -pub use linux::*; - -#[cfg(target_os = "macos")] -pub mod macos; - -#[cfg(target_os = "macos")] -pub use macos::*; diff --git a/diskann-platform/src/linux/mod.rs b/diskann-platform/src/linux/mod.rs deleted file mode 100644 index 38ccf7e08..000000000 --- a/diskann-platform/src/linux/mod.rs +++ /dev/null @@ -1,16 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ -#![cfg_attr( - not(test), - warn(clippy::panic, clippy::unwrap_used, clippy::expect_used) -)] - -mod perf; -pub use perf::{ - get_number_of_processors, get_peak_workingset_size, get_process_cycle_time, get_process_time, - get_system_time, -}; - -pub mod ssd_io_context; diff --git a/diskann-platform/src/macos/mod.rs b/diskann-platform/src/macos/mod.rs deleted file mode 100644 index 38ccf7e08..000000000 --- a/diskann-platform/src/macos/mod.rs +++ /dev/null @@ -1,16 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ -#![cfg_attr( - not(test), - warn(clippy::panic, clippy::unwrap_used, clippy::expect_used) -)] - -mod perf; -pub use perf::{ - get_number_of_processors, get_peak_workingset_size, get_process_cycle_time, get_process_time, - get_system_time, -}; - -pub mod ssd_io_context; diff --git a/diskann-platform/src/macos/ssd_io_context.rs b/diskann-platform/src/macos/ssd_io_context.rs deleted file mode 100644 index b8302eecd..000000000 --- a/diskann-platform/src/macos/ssd_io_context.rs +++ /dev/null @@ -1,22 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::fs::File; - -/// The IOContext struct for disk I/O on macOS. -/// -/// macOS does not have io_uring (Linux-only) or IOCompletionPort (Windows-only). -/// This is a simple implementation that just holds a file handle. -/// For actual async I/O operations on macOS, use the StorageProviderAlignedFileReader -/// which provides a cross-platform fallback implementation. -pub struct IOContext { - pub file_handle: File, -} - -impl IOContext { - pub fn new(file_handle: File) -> Self { - IOContext { file_handle } - } -} diff --git a/diskann-platform/src/win/ssd_io_context.rs b/diskann-platform/src/win/ssd_io_context.rs deleted file mode 100644 index eafdf1ef8..000000000 --- a/diskann-platform/src/win/ssd_io_context.rs +++ /dev/null @@ -1,35 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use crate::{FileHandle, IOCompletionPort}; - -// The IOContext struct for disk I/O. One for each thread. -pub struct IOContext { - pub status: Status, - pub file_handle: FileHandle, - pub io_completion_port: IOCompletionPort, -} - -impl Default for IOContext { - fn default() -> Self { - IOContext { - status: Status::ReadWait, - file_handle: FileHandle::default(), - io_completion_port: IOCompletionPort::default(), - } - } -} - -impl IOContext { - pub fn new() -> Self { - Self::default() - } -} - -pub enum Status { - ReadWait, - ReadSuccess, - ProcessComplete, -} diff --git a/diskann-platform/src/win/thread_safe_handle.rs b/diskann-platform/src/win/thread_safe_handle.rs deleted file mode 100644 index e71cca8f7..000000000 --- a/diskann-platform/src/win/thread_safe_handle.rs +++ /dev/null @@ -1,35 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::{ - io, - sync::{Mutex, MutexGuard}, -}; - -use windows_sys::Win32::Foundation::HANDLE; - -/// `ThreadSafeHandle` struct that wraps a native Windows `HANDLE` object with a mutex to ensure thread safety. -pub struct ThreadSafeHandle(Mutex); - -/// Implement `Send` and `Sync` for `ThreadSafeHandle` to allow it to be shared across threads. -unsafe impl Send for ThreadSafeHandle {} -unsafe impl Sync for ThreadSafeHandle {} - -impl ThreadSafeHandle { - /// Lock the mutex and return a guard to the handle. - pub fn lock(&self) -> io::Result> { - self.0.lock().map_err(|_| { - io::Error::new( - io::ErrorKind::WouldBlock, - "Unable to acquire lock on ThreadSafeHandle.", - ) - }) - } - - /// Create a new `ThreadSafeHandle` from a native Windows `HANDLE`. - pub fn new(handle: HANDLE) -> Self { - Self(Mutex::new(handle)) - } -} diff --git a/diskann-providers/Cargo.toml b/diskann-providers/Cargo.toml index 900b38cdb..6fb7544ea 100644 --- a/diskann-providers/Cargo.toml +++ b/diskann-providers/Cargo.toml @@ -21,8 +21,6 @@ rand.workspace = true rayon.workspace = true serde = { workspace = true, features = ["derive"] } thiserror.workspace = true -opentelemetry = { workspace = true, optional = true } -diskann-platform = { workspace = true } diskann-vector = { workspace = true } rand_distr.workspace = true cfg-if.workspace = true @@ -71,7 +69,6 @@ targets = [ [features] default = [] -perf_test = ["dep:opentelemetry"] testing = ["dep:tempfile"] experimental_diversity_search = ["diskann/experimental_diversity_search"] virtual_storage = ["dep:vfs"] diff --git a/diskann-providers/src/model/pq/pq_construction.rs b/diskann-providers/src/model/pq/pq_construction.rs index 4e0dd769e..8ca3f11a4 100644 --- a/diskann-providers/src/model/pq/pq_construction.rs +++ b/diskann-providers/src/model/pq/pq_construction.rs @@ -8,6 +8,7 @@ use std::{ mem::size_of, num::NonZeroUsize, sync::atomic::AtomicBool, + time::Instant, vec, }; @@ -34,7 +35,7 @@ use crate::{ model::GeneratePivotArguments, storage::PQStorage, utils::{ - BridgeErr, ParallelIteratorInPool, RandomProvider, RayonThreadPoolRef, Timer, + BridgeErr, ParallelIteratorInPool, RandomProvider, RayonThreadPoolRef, create_rnd_provider_from_seed, }, }; @@ -324,7 +325,7 @@ where T: Copy + VectorRepr, Storage: StorageWriteProvider + StorageReadProvider, { - let timer = Timer::new(); + let timer = Instant::now(); info!("Generating PQ data starting from offset {}", offset); diff --git a/diskann-providers/src/utils/mod.rs b/diskann-providers/src/utils/mod.rs index 20743f5d6..8fa464f48 100644 --- a/diskann-providers/src/utils/mod.rs +++ b/diskann-providers/src/utils/mod.rs @@ -14,9 +14,6 @@ pub use rayon_util::{ create_thread_pool_for_bench, create_thread_pool_for_test, }; -mod timer; -pub use timer::Timer; - /// Read/write block size (64 MB) for cached I/O operations such as CachedReader and CachedWriter pub const READ_WRITE_BLOCK_SIZE: u64 = 64 * 1024 * 1024; diff --git a/diskann-tools/Cargo.toml b/diskann-tools/Cargo.toml index 573b0f05c..7df37e276 100644 --- a/diskann-tools/Cargo.toml +++ b/diskann-tools/Cargo.toml @@ -47,7 +47,7 @@ diskann-utils = { workspace = true, features = ["testing"] } [features] default = [] # When a feature is enabled in tools, enable it in dependencies -perf_test = ["diskann-providers/perf_test", "diskann-disk/perf_test"] +perf_test = ["diskann-disk/perf_test"] [lints] clippy.undocumented_unsafe_blocks = "warn" diff --git a/diskann-tools/src/bin/compute_groundtruth.rs b/diskann-tools/src/bin/compute_groundtruth.rs index 75c876f39..b562cd184 100644 --- a/diskann-tools/src/bin/compute_groundtruth.rs +++ b/diskann-tools/src/bin/compute_groundtruth.rs @@ -3,7 +3,8 @@ * Licensed under the MIT license. */ use clap::Parser; -use diskann_providers::{storage::FileStorageProvider, utils::Timer}; +use diskann_disk::utils::instrumentation::Timer; +use diskann_providers::storage::FileStorageProvider; use diskann_tools::utils::{ compute_ground_truth_from_datafiles, init_subscriber, CMDResult, DataType, GraphDataF32Vector, GraphDataHalfVector, GraphDataInt8Vector, GraphDataU8Vector, diff --git a/diskann-tools/src/bin/compute_multivec_groundtruth.rs b/diskann-tools/src/bin/compute_multivec_groundtruth.rs index 452187375..3826f860e 100644 --- a/diskann-tools/src/bin/compute_multivec_groundtruth.rs +++ b/diskann-tools/src/bin/compute_multivec_groundtruth.rs @@ -3,7 +3,8 @@ * Licensed under the MIT license. */ use clap::Parser; -use diskann_providers::{storage::FileStorageProvider, utils::Timer}; +use diskann_disk::utils::instrumentation::Timer; +use diskann_providers::storage::FileStorageProvider; use diskann_tools::utils::{ compute_multivec_ground_truth_from_datafiles, init_subscriber, CMDResult, DataType, GraphDataF32Vector, GraphDataHalfVector, GraphDataInt8Vector, GraphDataU8Vector, diff --git a/diskann-tools/src/utils/build_disk_index.rs b/diskann-tools/src/utils/build_disk_index.rs index d55a1f566..116ed227a 100644 --- a/diskann-tools/src/utils/build_disk_index.rs +++ b/diskann-tools/src/utils/build_disk_index.rs @@ -18,12 +18,13 @@ use diskann_disk::{ DiskIndexBuildParameters, MemoryBudget, NumPQChunks, DISK_SECTOR_LEN, }, storage::DiskIndexWriter, + utils::instrumentation::Timer, QuantizationType, }; -use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; use diskann_providers::{ model::IndexConfiguration, - utils::{load_metadata_from_file, Timer}, + storage::{StorageReadProvider, StorageWriteProvider}, + utils::load_metadata_from_file, }; use diskann_vector::distance::Metric; use opentelemetry::global::BoxedSpan; diff --git a/diskann-tools/src/utils/build_pq.rs b/diskann-tools/src/utils/build_pq.rs index 15321cfca..8966745c4 100644 --- a/diskann-tools/src/utils/build_pq.rs +++ b/diskann-tools/src/utils/build_pq.rs @@ -4,17 +4,16 @@ */ use diskann::ANNResult; -use diskann_disk::data_model::GraphDataType; -use diskann_providers::storage::StorageReadProvider; +use diskann_disk::{data_model::GraphDataType, utils::instrumentation::Timer}; use diskann_providers::{ model::{ GeneratePivotArguments, MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, NUM_PQ_CENTROIDS, }, storage::{ get_disk_index_compressed_pq_file, get_disk_index_pq_pivot_file, FileStorageProvider, - PQStorage, + PQStorage, StorageReadProvider, }, - utils::{create_thread_pool, load_metadata_from_file, Timer}, + utils::{create_thread_pool, load_metadata_from_file}, }; use diskann_vector::distance::Metric; use tracing::info; diff --git a/diskann-tools/src/utils/search_disk_index.rs b/diskann-tools/src/utils/search_disk_index.rs index 3825d2631..577bae80f 100644 --- a/diskann-tools/src/utils/search_disk_index.rs +++ b/diskann-tools/src/utils/search_disk_index.rs @@ -9,14 +9,12 @@ use diskann::utils::IntoUsize; use diskann_disk::{ data_model::{CachingStrategy, GraphDataType}, search::provider::{ + aligned_file_reader::traits::AlignedReaderFactory, disk_provider::{DiskIndexSearcher, SearchPostProcessorKind}, disk_vertex_provider_factory::DiskVertexProviderFactory, }, storage::disk_index_reader::DiskIndexReader, - utils::{ - aligned_file_reader::traits::AlignedReaderFactory, instrumentation::PerfLogger, statistics, - QueryStatistics, - }, + utils::{instrumentation::PerfLogger, statistics, QueryStatistics}, }; use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; use diskann_providers::{