Add support for concurrent file processing in ingestion module

BryanChasko · BryanChasko · commit 71a216402466 · 2025-07-20T20:10:18.000-06:00
- Introduced `max_concurrent_files` option in `IngestConfig` to control the number of files processed simultaneously.
- Updated `process_directory_tree` to utilize a semaphore for managing concurrent tasks.
- Refactored file processing functions to support parallel execution.
- Enhanced command-line interface to accept `max_concurrent` argument.
- Improved metadata loading with modern file reading methods.
diff --git a/rust_ingest/src/ingest.rs b/rust_ingest/src/ingest.rs
@@ -65,6 +65,9 @@ pub struct IngestConfig {
     pub max_chars: usize,
     /// Maximum tokens for embedding requests.
     pub max_tokens: usize,
+    /// Maximum number of files to process concurrently.
+    /// If None, defaults to the number of CPU cores.
+    pub max_concurrent_files: Option<usize>,
 }
 
 impl Default for IngestConfig {
@@ -73,6 +76,7 @@ impl Default for IngestConfig {
             root_dir: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
             max_chars: MAX_FILE_CHARS,
             max_tokens: MAX_EMBEDDING_TOKENS,
+            max_concurrent_files: None, // Use CPU core count by default
         }
     }
 }
@@ -168,14 +172,44 @@ fn create_hnsw_index() -> Hnsw<'static, f32, DistCosine> {
     )
 }
 
-/// Processes all files in the directory tree.
+/// Processes all files in the directory tree with parallel execution.
 async fn process_directory_tree(
     config: &IngestConfig,
     client: &reqwest::Client,
     index: &Hnsw<'_, f32, DistCosine>,
     file_metadata: &mut Vec<PathBuf>,
     stats: &mut IngestStats,
 ) -> Result<()> {
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::sync::{Arc, Mutex};
+    use tokio::sync::Semaphore;
+    use tokio::task::JoinSet;
+
+    // Create a semaphore to limit concurrent operations
+    let max_concurrent = config.max_concurrent_files.unwrap_or_else(|| {
+        // Default to number of CPUs if not specified
+        std::thread::available_parallelism()
+            .map(|n| n.get())
+            .unwrap_or(4)
+    });
+
+    // Use a semaphore to limit concurrent embedding operations
+    let semaphore = Arc::new(Semaphore::new(max_concurrent));
+
+    // Use JoinSet to manage async tasks
+    let mut tasks = JoinSet::new();
+
+    // Use atomics for thread-safe counters
+    let processed_count = Arc::new(AtomicUsize::new(0));
+    let skipped_count = Arc::new(AtomicUsize::new(0));
+
+    // Use a mutex to protect the file metadata vector
+    let file_paths = Arc::new(Mutex::new(Vec::new()));
+
+    // Collect candidate files first
+    let mut candidate_files = Vec::new();
+
+    // First pass: find all valid files to process
     for entry in WalkDir::new(&config.root_dir)
         .into_iter()
         .filter_map(|e| e.ok())
@@ -192,29 +226,82 @@ async fn process_directory_tree(
             continue;
         }
 
-        match process_single_file(
-            path,
-            config,
-            client,
-            index,
-            file_metadata,
-            stats.files_processed,
-        )
-        .await
-        {
-            Ok(()) => {
-                stats.files_processed += 1;
-                if stats.files_processed % PROGRESS_INTERVAL == 0 {
-                    println!("Processed {} files…", stats.files_processed);
+        // Add to candidates
+        candidate_files.push(path.to_path_buf());
+    }
+
+    println!("Found {} files to process", candidate_files.len());
+
+    // Second pass: process files concurrently
+    for (file_id, path) in candidate_files.into_iter().enumerate() {
+        // Clone references for the async task
+        let semaphore_clone = semaphore.clone();
+        let client_clone = client.clone();
+        let config_clone = config.clone();
+        let processed_count_clone = processed_count.clone();
+        let skipped_count_clone = skipped_count.clone();
+        let file_paths_clone = file_paths.clone();
+        let path_clone = path.clone();
+
+        // Spawn a task for each file
+        tasks.spawn(async move {
+            // Acquire a permit from the semaphore
+            let _permit = semaphore_clone.acquire().await.unwrap();
+
+            // Process the file
+            match process_single_file_for_embedding(&path_clone, &config_clone, &client_clone).await
+            {
+                Ok(embedding) => {
+                    // Successfully processed
+                    let count = processed_count_clone.fetch_add(1, Ordering::SeqCst) + 1;
+
+                    // Store result
+                    let mut metadata = file_paths_clone.lock().unwrap();
+                    metadata.push((file_id, path_clone, embedding));
+
+                    // Show progress periodically
+                    if count % PROGRESS_INTERVAL == 0 {
+                        println!("Processed {} files…", count);
+                    }
+
+                    Ok(())
+                }
+                Err(e) => {
+                    // Log error and count as skipped
+                    eprintln!(
+                        "Warning: Failed to process file {}: {}",
+                        path_clone.display(),
+                        e
+                    );
+                    skipped_count_clone.fetch_add(1, Ordering::SeqCst);
+                    Err(e)
                 }
             }
-            Err(e) => {
-                eprintln!("Warning: Failed to process file {}: {}", path.display(), e);
-                stats.files_skipped += 1;
-            }
-        }
+        });
     }
 
+    // Wait for all tasks to complete
+    while let Some(result) = tasks.join_next().await {
+        // Just check for panics, errors are already handled in the task
+        result?;
+    }
+
+    // Update stats from atomic counters
+    stats.files_processed += processed_count.load(Ordering::SeqCst);
+    stats.files_skipped += skipped_count.load(Ordering::SeqCst);
+
+    // Sort results by file_id and insert into the index
+    let mut results = file_paths.lock().unwrap();
+    results.sort_by_key(|(id, _, _)| *id);
+
+    // Now populate the index and metadata
+    for (_, path, embedding) in results.iter() {
+        let file_id = file_metadata.len();
+        index.insert((embedding.as_slice(), file_id));
+        file_metadata.push(path.clone());
+    }
+
+    println!("Successfully indexed {} files", file_metadata.len());
     Ok(())
 }
 
@@ -232,29 +319,26 @@ fn is_supported_file(path: &std::path::Path) -> bool {
         .is_some_and(|ext| SUPPORTED_EXTENSIONS.contains(&ext))
 }
 
-/// Processes a single file and adds it to the index.
+/// Processes a single file for embedding without modifying the index.
+///
+/// This function handles just the embedding part, making it suitable for
+/// parallel processing in our async pipeline.
 ///
 /// # Arguments
 /// * `path` - Path to the file being processed
 /// * `config` - Configuration settings for ingestion
 /// * `client` - HTTP client for embedding API requests
-/// * `index` - HNSW index to insert embeddings into
-/// * `file_metadata` - Collection of file paths to track processed files
-/// * `file_id` - Unique identifier for this file in the index
 ///
 /// # Returns
-/// Success if the file was processed and added to the index.
+/// The embedding vector on success
 ///
 /// # Errors
 /// Returns error if file reading or embedding generation fails.
-async fn process_single_file(
+async fn process_single_file_for_embedding(
     path: &std::path::Path,
     config: &IngestConfig,
     client: &reqwest::Client,
-    index: &Hnsw<'_, f32, DistCosine>,
-    file_metadata: &mut Vec<PathBuf>,
-    file_id: usize,
-) -> Result<()> {
+) -> Result<Vec<f32>> {
     // Read and truncate file content
     let content = std::fs::read_to_string(path)
         .with_context(|| format!("Failed to read file: {}", path.display()))?;
@@ -266,6 +350,37 @@ async fn process_single_file(
         .await
         .with_context(|| format!("Failed to generate embedding for file: {}", path.display()))?;
 
+    Ok(embedding)
+}
+
+/// Processes a single file and adds it to the index.
+///
+/// # Arguments
+/// * `path` - Path to the file being processed
+/// * `config` - Configuration settings for ingestion
+/// * `client` - HTTP client for embedding API requests
+/// * `index` - HNSW index to insert embeddings into
+/// * `file_metadata` - Collection of file paths to track processed files
+/// * `file_id` - Unique identifier for this file in the index
+///
+/// # Returns
+/// Success if the file was processed and added to the index.
+///
+/// # Errors
+/// Returns error if file reading or embedding generation fails.
+///
+/// @deprecated Use the parallel processing pipeline instead
+async fn process_single_file(
+    path: &std::path::Path,
+    config: &IngestConfig,
+    client: &reqwest::Client,
+    index: &Hnsw<'_, f32, DistCosine>,
+    file_metadata: &mut Vec<PathBuf>,
+    file_id: usize,
+) -> Result<()> {
+    // Generate embedding
+    let embedding = process_single_file_for_embedding(path, config, client).await?;
+
     // Insert into index
     index.insert((embedding.as_slice(), file_id));
 
diff --git a/rust_ingest/src/main.rs b/rust_ingest/src/main.rs
@@ -48,6 +48,11 @@ enum Commands {
         /// Maximum tokens for embedding requests
         #[arg(long, default_value = "600")]
         max_tokens: usize,
+
+        /// Maximum number of files to process concurrently
+        /// If not specified, defaults to number of CPU cores
+        #[arg(long)]
+        max_concurrent: Option<usize>,
     },
 
     /// Ask a question using the pre-built index.
@@ -88,6 +93,7 @@ async fn main() -> Result<()> {
             root,
             max_chars,
             max_tokens,
+            max_concurrent,
         } => {
             // Create ingest configuration
             let mut config = rust_ingest::ingest::IngestConfig::default();
@@ -97,6 +103,7 @@ async fn main() -> Result<()> {
             }
             config.max_chars = max_chars;
             config.max_tokens = max_tokens;
+            config.max_concurrent_files = max_concurrent;
 
             // Run ingestion process
             let stats = ingest::run_with_config(config).await?;
diff --git a/rust_ingest/src/query.rs b/rust_ingest/src/query.rs
@@ -159,21 +159,24 @@ fn load_index_and_metadata(
 ) -> Result<(Hnsw<'static, f32, DistCosine>, Vec<PathBuf>)> {
     let data_dir = config.root_dir.join("data");
 
-    // Load the HNSW index using the correct API with HnswIo
-    let mut hnsw_loader = HnswIo::new(&data_dir, "index");
-    let loaded_index: Hnsw<'_, f32, DistCosine> = hnsw_loader
-        .load_hnsw()
+    // Load the HNSW index using the modern file_load method
+    let index_path = data_dir.join("index");
+    let loaded_index: Hnsw<'_, f32, DistCosine> = HnswIo::file_load(&index_path)
         .context("Failed to load HNSW index - ensure ingestion has been run")?;
 
     // Convert the loaded index to an owned index with 'static lifetime
     // SAFETY: This transmute extends the lifetime of the index, which is safe
     // because we're taking ownership of the index and ensuring it outlives
     // its original borrow
-    let index: Hnsw<'static, f32, DistCosine> = unsafe { std::mem::transmute(loaded_index) }; // Load file metadata
-    let metadata_file =
-        fs::File::open(data_dir.join("meta.json")).context("Failed to open metadata file")?;
+    let index: Hnsw<'static, f32, DistCosine> = unsafe { std::mem::transmute(loaded_index) };
+
+    // Load file metadata
+    let metadata_path = data_dir.join("meta.json");
+    let metadata_content = fs::read_to_string(&metadata_path)
+        .with_context(|| format!("Failed to read metadata file: {}", metadata_path.display()))?;
+
     let metadata: Vec<PathBuf> =
-        serde_json::from_reader(metadata_file).context("Failed to parse metadata JSON")?;
+        serde_json::from_str(&metadata_content).context("Failed to parse metadata JSON")?;
 
     Ok((index, metadata))
 }
