diff --git a/fieldExtraction/src/investment/main.py b/fieldExtraction/src/investment/main.py index 92a8e51..4e37ed1 100644 --- a/fieldExtraction/src/investment/main.py +++ b/fieldExtraction/src/investment/main.py @@ -1,4 +1,5 @@ import concurrent.futures +import gc import logging import os import random @@ -11,30 +12,46 @@ random.seed(42) import traceback +from constants.constants import Constants from sentence_transformers import SentenceTransformer import src.investment.file_processing as file_processing import src.utils.io_utils as io_utils -from constants.constants import Constants from src import config def safe_process_file(item, constants, run_timestamp): - """call process_file and catch any exceptions that occur + """Process a single file with error handling and resource cleanup. + + This function wraps the main file processing logic to ensure that: + - Individual file failures don't crash the entire batch + - Resources are properly cleaned up after each file + - Detailed error information is captured for debugging + - Garbage collection is forced to help prevent file descriptor leaks Args: - item: A file object - all_dataset: Dictionary containing the crosswalk data for different levels for different codes - run_timestamp: The timestamp when the process is executed + item: Tuple of (filename, contract_text) representing the file to process + constants: Constants object containing configuration and processing rules + run_timestamp: Timestamp string for output file naming and tracking Returns: - pd.DataFrame: DataFrame containing the results or an error message + pd.DataFrame: Either: + - DataFrame with extracted field results (success case) + - Single-row DataFrame with error details (failure case) + + Note: + Always returns a DataFrame to maintain consistent output format for pd.concat. + Forces garbage collection after processing to help with resource cleanup, + particularly important for preventing file descriptor leaks in concurrent processing. """ file_id = ( item[0] if isinstance(item, (list, tuple)) and len(item) > 0 else item ) # unpack file_id from item (passed in as a tuple below) try: - return file_processing.process_file(item, constants, run_timestamp) + result = file_processing.process_file(item, constants, run_timestamp) + # Force garbage collection after each file + gc.collect() + return result except ( Exception ) as e: # When there's an issue with the processing inside the future @@ -47,6 +64,9 @@ def safe_process_file(item, constants, run_timestamp): logging.error(f"Error Message: {error_message}") logging.error(f"Full traceback:\n{full_traceback}") + + # Force garbage collection on error too + gc.collect() return pd.DataFrame( [ { diff --git a/fieldExtraction/src/investment/smart_chunking_funcs.py b/fieldExtraction/src/investment/smart_chunking_funcs.py index 8c25857..4ded0f6 100644 --- a/fieldExtraction/src/investment/smart_chunking_funcs.py +++ b/fieldExtraction/src/investment/smart_chunking_funcs.py @@ -8,6 +8,7 @@ for contract processing. It includes: - Multiple chunking strategies (hierarchical, OR, AND, regex) - Semantic and keyword-based chunk retrieval - Thread-safe operations for multi-threaded contract processing +- Monitoring and cleanup to prevent file descriptor leaks ("too many open files" errors) Key Components: - field_context(): Main function for semantic chunk retrieval with caching @@ -18,8 +19,24 @@ Performance Notes: - Uses module-level vectorstore caching to share embeddings across threads - Reference counting ensures proper cleanup of temporary vectorstores - Designed to handle concurrent processing without file descriptor leaks +- Emergency cleanup via atexit handler prevents resource leaks on process termination + +File Descriptor Management: +- Monitors open file descriptors to prevent "too many open files" errors +- Implements reference counting for shared vectorstores +- Forces Chroma client connection cleanup via _client.reset() +- Includes retry logic for directory cleanup operations +- Logs detailed statistics every 10 files for monitoring + +Troubleshooting: +- Check logs for "High number of open files" warnings +- Monitor cache hit/miss ratios for performance +- Watch for cleanup failure counts in cache statistics +- Emergency cleanup logs appear during process shutdown """ +import atexit +import gc import hashlib import json import logging @@ -27,18 +44,22 @@ import os import re import shutil import threading +import time +import weakref from collections import defaultdict from itertools import combinations from typing import Any, Dict, List, Optional, Set, Tuple, Union +import chromadb import nltk +import psutil +from constants.constants import Constants from langchain_aws.embeddings.bedrock import BedrockEmbeddings from langchain_chroma import Chroma from langchain_community.retrievers import BM25Retriever from langchain_text_splitters import RecursiveCharacterTextSplitter from nltk.tokenize import word_tokenize -from constants.constants import Constants from src import config from src.prompts.fieldset import FieldSet @@ -49,6 +70,17 @@ nltk.data.path.append( ) ) +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +# Enable Chroma reset functionality +try: + # Enable reset functionality to allow proper cleanup + chromadb.config.Settings(allow_reset=True) + logger.info("Chroma reset functionality enabled") +except Exception as e: + logger.warning(f"Could not enable Chroma reset: {e}") + # =================================================================== # VECTORSTORE CACHING @@ -57,6 +89,53 @@ nltk.data.path.append( _vectorstore_lock = threading.Lock() _vectorstore_cache: Dict[str, Dict[str, Any]] = {} +# Track active vectorstores for cleanup +_active_vectorstores = weakref.WeakSet() + + +def _cleanup_all_vectorstores(): + """Emergency cleanup function for atexit""" + logger.info("Emergency cleanup starting...") + with _vectorstore_lock: + cleanup_count = 0 + for cache_key in list(_vectorstore_cache.keys()): + try: + vectorstore_data = _vectorstore_cache[cache_key] + vectorstore = vectorstore_data["vectorstore"] + + # Force close the vectorstore + if hasattr(vectorstore, "_client") and vectorstore._client: + vectorstore._client.reset() + if hasattr(vectorstore, "delete_collection"): + vectorstore.delete_collection() + + # Clean up directory + chroma_filename = vectorstore_data.get("chroma_filename") + if os.path.exists(chroma_filename): + shutil.rmtree(chroma_filename, ignore_errors=True) + + cleanup_count += 1 + except Exception as e: + logger.warning(f"Error during emergency cleanup of {cache_key}: {e}") + + _vectorstore_cache.clear() + logger.info( + f"Emergency cleanup completed. Cleaned up {cleanup_count} vectorstores." + ) + + +# Register emergency cleanup +atexit.register(_cleanup_all_vectorstores) + +# Monitoring counters +_cache_stats = { + "cache_hits": 0, + "cache_misses": 0, + "vectorstores_created": 0, + "vectorstores_cleaned": 0, + "cleanup_failures": 0, + "open_files_peak": 0, +} # Setup bedrock embeddings_model = BedrockEmbeddings( @@ -77,6 +156,72 @@ TOP_K = 8 DATE_PATTERN = r"(?i)\bday of\b|\b(?:jan(?:uary)?|feb(?:ruary)?|mar(?:ch)?|apr(?:il)?|may|jun(?:e)?|jul(?:y)?|aug(?:ust)?|sep(?:tember)?|oct(?:ober)?|nov(?:ember)?|dec(?:ember)?)\b|\b(?:\d{1,2}[\/\.-]\d{1,2}[\/\.-](?:\d{2}|\d{4})|\d{4}[\/\.-]\d{1,2}[\/\.-]\d{1,2}|\d{8}|\d{2}\d{2}|\d{4}[\/\.-]\d{1,2}[\/\.-]\d{1,2}|\d{4}[\.]\d{1,2}[\.]\d{1,2})\b" +def get_open_file_count() -> int: + """Get current number of open file descriptors for monitoring. + + Uses Linux /proc filesystem when available, falls back to psutil. + Returns -1 if neither method works (shouldn't happen in normal operation). + + Returns: + int: Number of open file descriptors, or -1 if unable to determine + """ + + try: # Linux-specific: count files in /proc/PID/fd directory + return len(os.listdir(f"/proc/{os.getpid()}/fd")) + except: + # Fallback if /proc not available (e.g., on Windows or Mac) + try: + return len(psutil.Process().open_files()) + except: + return -1 + + +def log_cache_stats(): + """Log comprehensive cache statistics for monitoring and debugging. + + Logs include: + - Cache hit/miss ratios (higher hit ratio = better performance) + - Vectorstore creation/cleanup counts (should be roughly equal) + - Current and peak open file descriptor counts + - Active cache size (number of cached vectorstores) + - Cleanup failure count (should remain at 0) + """ + open_files = get_open_file_count() + _cache_stats["open_files_peak"] = max(_cache_stats["open_files_peak"], open_files) + + logger.info( + f"Cache Stats - Hits: {_cache_stats['cache_hits']}, " + f"Misses: {_cache_stats['cache_misses']}, " + f"Created: {_cache_stats['vectorstores_created']}, " + f"Cleaned: {_cache_stats['vectorstores_cleaned']}, " + f"Failures: {_cache_stats['cleanup_failures']}, " + f"Open Files: {open_files}, " + f"Peak Files: {_cache_stats['open_files_peak']}, " + f"Active Cache: {len(_vectorstore_cache)}" + ) + + +def monitor_file_descriptors(filename: str, operation: str) -> int: + """Monitor file descriptor usage and warn if approaching limits. + + Args: + filename: Name of file being processed (for logging context) + operation: Operation being performed (e.g., "start", "vectorstore_create") + + Returns: + int: Current number of open file descriptors + + Note: + Logs warning if file descriptor count exceeds 800 (threshold for 1024 limit) + """ + open_files = get_open_file_count() + if open_files > 800: # Warning threshold + logger.warning( + f"High number of open files ({open_files}) detected during {operation} on {filename}" + ) + return open_files + + def parse_chunk(chunk: str) -> tuple[int, str]: """Extract chunk number and content from a chunk string @@ -166,17 +311,40 @@ def field_context( one_to_one_fields: FieldSet, constants: Constants, contract_text: str, filename: str ) -> Tuple[FieldSet, List[str]]: """Thread-safe field context extraction with vectorstore caching. + This function creates chunks from whole contract text and returns relevant chunks for each field based on semantic search and keyword search. + This function is the core of the file descriptor leak prevention system. + It implements: + - Thread-safe vectorstore caching with reference counting + - Comprehensive cleanup with retries and forced garbage collection + - File descriptor monitoring throughout the process + - Emergency cleanup registration for process termination + Args: - contract_text (str): full text of the contract + one_to_one_fields: Fields to process for semantic chunking + constants: Configuration constants + contract_text: Full contract text to chunk and embed + filename: Filename for logging context Returns: - dict[Any, set[str]]: retrieved text for each field - list[str]: list of all chunks + Tuple of (updated fieldset with retrieved text, list of all chunks) + + Raises: + Various exceptions from Chroma/embedding operations + + Note: + This function is designed to be called concurrently from multiple threads. + The vectorstore cache prevents duplicate embedding work and the reference + counting ensures cleanup only happens when all threads are done. """ + # === MONITORING SETUP === + # Track processing time and file descriptor usage for this operation + start_time = time.time() + initial_files = monitor_file_descriptors(filename, "start_field_context") + # Convert contract_text to lowercase for case-insensitive processing contract_text = contract_text.lower() @@ -195,81 +363,172 @@ def field_context( for idx, split in enumerate(splits) ] - # Thread-safe vectorstore creation/retrieval - with _vectorstore_lock: - if cache_key not in _vectorstore_cache: - # Create vectorstore with persistence (for performance) - vectorstore = Chroma.from_texts( - texts=splits, - embedding=embeddings_model, - persist_directory=chroma_filename, - ) + vectorstore = None + try: + # === THREAD-SAFE CACHE ACCESS === + # Multiple threads may try to create / access the same vectorstore + # Lock ensures only one thread creates it, others wait and reuse + with _vectorstore_lock: + if cache_key not in _vectorstore_cache: + # === VECTORSTORE CREATION === + # This is the expensive operation we would like to cache + # Creates Chroma database with persistent storage for performance + logger.info(f"Creating new vectorstore for {filename} (cache miss)") + _cache_stats["cache_misses"] += 1 + _cache_stats["vectorstores_created"] += 1 - _vectorstore_cache[cache_key] = { - "vectorstore": vectorstore, - "splits": splits, - "chroma_filename": chroma_filename, - "ref_count": 0, - } + # Create vectorstore with persistence (for performance) + vectorstore = Chroma.from_texts( + texts=splits, + embedding=embeddings_model, + persist_directory=chroma_filename, + ) - # Increment reference count - _vectorstore_cache[cache_key]["ref_count"] += 1 - vectorstore_data = _vectorstore_cache[cache_key] + _vectorstore_cache[cache_key] = { + "vectorstore": vectorstore, + "splits": splits, + "chroma_filename": chroma_filename, + "ref_count": 0, + "created_at": time.time(), + } - # Use the cached/created vectorstore - vectorstore = vectorstore_data["vectorstore"] - splits = vectorstore_data["splits"] + # === EMERGENCY CLEANUP REGISTRATION === + # Add to weak set so emergency cleanup can find it if process crashes + _active_vectorstores.add(vectorstore) # Track for emergency cleanup - vectorstore_retriever = vectorstore.as_retriever(search_kwargs={"k": TOP_K}) - keyword_retriever = BM25Retriever.from_texts( - splits, k=TOP_K, preprocess_func=word_tokenize - ) - # keyword_retriever.k = 3 - # ensemble_retriever = EnsembleRetriever(retrievers=[vectorstore_retriever, keyword_retriever], weights=[0.5, 0.5]) + after_create_files = monitor_file_descriptors( + filename, "vectorstore_create" + ) + logger.info( + f"Vectorstore created. Files {initial_files} -> {after_create_files}" + ) + else: + # === CACHE HIT === + # Reusing existing vectorstore saves significant time and resources + logger.info(f"Using cached vectorstore for {filename} (cache hit)") + _cache_stats["cache_hits"] += 1 - for field in one_to_one_fields.fields: - if field.field_type == "smart_chunked": - questions = {} - questions[field.field_name] = field.get_prompt(constants) - question = json.dumps(questions) + # === REFERENCE COUNTING === + # Track how many threads are using this vectorstore + # Only cleanup when count reaches 0 + _vectorstore_cache[cache_key]["ref_count"] += 1 + vectorstore_data = _vectorstore_cache[cache_key] - # Convert field keywords to lowercase for case-insensitive processing - field_keywords = [keyword.lower() for keyword in field.keywords] + # Use the cached/created vectorstore + vectorstore = vectorstore_data["vectorstore"] + splits = vectorstore_data["splits"] - # get relevant chunks based on semantic search and keyword search and merge them - semantic_relevant_docs = vectorstore_retriever.invoke(question) - keyword_relevant_docs = keyword_retriever.invoke(" ;".join(field.keywords)) - semantic_relevant_chunks = [ - doc.page_content for doc in semantic_relevant_docs - ] - keyword_relevant_chunks = [ - doc.page_content for doc in keyword_relevant_docs - ] - if field.field_name == "CONTRACT_TITLE": - keyword_relevant_chunks = keyword_relevant_chunks + splits[:5] - field.retrieved_text = sorted( - list(set(semantic_relevant_chunks + keyword_relevant_chunks)) - ) + vectorstore_retriever = vectorstore.as_retriever(search_kwargs={"k": TOP_K}) + keyword_retriever = BM25Retriever.from_texts( + splits, k=TOP_K, preprocess_func=word_tokenize + ) - # Check if no context or full_context, if so, switch attribute - if len(field.retrieved_text) == 0 or len(field.retrieved_text) == len( - splits - ): - field.field_type = "full_context" + for field in one_to_one_fields.fields: + if field.field_type == "smart_chunked": + questions = {} + questions[field.field_name] = field.get_prompt(constants) + question = json.dumps(questions) - # Thread-safe cleanup with reference counting - with _vectorstore_lock: - _vectorstore_cache[cache_key]["ref_count"] -= 1 + # Convert field keywords to lowercase for case-insensitive processing + field_keywords = [keyword.lower() for keyword in field.keywords] - # Only cleanup when no more references - if _vectorstore_cache[cache_key]["ref_count"] == 0: - vectorstore.delete_collection() - try: - shutil.rmtree(_vectorstore_cache[cache_key]["chroma_filename"]) - except OSError as e: - logging.info(f"Error cleaning up {cache_key}: {e}") - logging.error(f"Error in field_context for {filename}: {e}") - del _vectorstore_cache[cache_key] # Remove from cache + # get relevant chunks based on semantic search and keyword search and merge them + semantic_relevant_docs = vectorstore_retriever.invoke(question) + keyword_relevant_docs = keyword_retriever.invoke( + " ;".join(field.keywords) + ) + semantic_relevant_chunks = [ + doc.page_content for doc in semantic_relevant_docs + ] + keyword_relevant_chunks = [ + doc.page_content for doc in keyword_relevant_docs + ] + if field.field_name == "CONTRACT_TITLE": + keyword_relevant_chunks = keyword_relevant_chunks + splits[:5] + field.retrieved_text = sorted( + list(set(semantic_relevant_chunks + keyword_relevant_chunks)) + ) + + # Check if no context or full_context, if so, switch attribute + if len(field.retrieved_text) == 0 or len(field.retrieved_text) == len( + splits + ): + field.field_type = "full_context" + + finally: + # === CLEANUP WITH REFERENCE COUNTING === + # This section is critical for preventing file descriptor leaks + with _vectorstore_lock: + if cache_key in _vectorstore_cache: + _vectorstore_cache[cache_key]["ref_count"] -= 1 + + # === ACTUAL CLEANUP === + # Only cleanup when no other threads are using this vectorstore + if _vectorstore_cache[cache_key]["ref_count"] == 0: + logger.info(f"Cleaning up vectorstore for {filename}") + try: + vectorstore_to_cleanup = _vectorstore_cache[cache_key][ + "vectorstore" + ] + chroma_filename_to_cleanup = _vectorstore_cache[cache_key][ + "chroma_filename" + ] + + # === FORCE CONNECTION CLEANUP === + # Directory deletion can fail due to timing issues + # Retry with small delays to handle race conditions + if ( + hasattr(vectorstore_to_cleanup, "_client") + and vectorstore_to_cleanup._client + ): + vectorstore_to_cleanup._client.reset() + + # Delete collection first + if hasattr(vectorstore_to_cleanup, "delete_collection"): + vectorstore_to_cleanup.delete_collection() + + # === RETRY LOGIC === + # Directory deletion can fail due to timing issues + # Retry with small delays to handle race conditions + # Clean up directory with retries + for attempt in range(3): + try: + if os.path.exists(chroma_filename_to_cleanup): + shutil.rmtree( + chroma_filename_to_cleanup, ignore_errors=True + ) + break # Success + except OSError as e: + if attempt == 2: # Last attempt + logger.warning( + f"Failed to cleanup {chroma_filename_to_cleanup} after 3 attempts: {e}" + ) + _cache_stats["cleanup_failures"] += 1 + else: + time.sleep(0.1) # Brief delay before retry + + # Force garbage collection + gc.collect() + _cache_stats["vectorstores_cleaned"] += 1 + + except Exception as e: + logger.warning(f"Error cleaning up {cache_key}: {e}") + _cache_stats["cleanup_failures"] += 1 + finally: + # Always remove from cache + del _vectorstore_cache[cache_key] + + end_time = time.time() + final_files = monitor_file_descriptors(filename, "completion") + processing_time = end_time - start_time + logger.info( + f"Completed {filename} in {processing_time:.2f}s. " + f"Files: {initial_files} -> {final_files}" + ) + + # Log stats every 10 files + if (_cache_stats["cache_hits"] + _cache_stats["cache_misses"]) % 10 == 0: + log_cache_stats() return one_to_one_fields, splits