|
| 1 | +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | +""" |
| 16 | +Batch processing utilities for the ingestor server. |
| 17 | +
|
| 18 | +This module provides utilities for calculating optimal batch parameters |
| 19 | +based on file characteristics and system resources. |
| 20 | +""" |
| 21 | + |
| 22 | +import logging |
| 23 | +from pathlib import Path |
| 24 | + |
| 25 | +from nvidia_rag.utils.configuration import NvidiaRAGConfig |
| 26 | + |
| 27 | +logger = logging.getLogger(__name__) |
| 28 | + |
| 29 | + |
| 30 | +# Text-like file extensions that process quickly |
| 31 | +# These are mapped to TXT or HTML in EXTENSION_TO_DOCUMENT_TYPE |
| 32 | +TEXT_LIKE_EXTENSIONS = frozenset({ |
| 33 | + "txt", |
| 34 | + "md", |
| 35 | + "json", |
| 36 | + "sh", |
| 37 | + "html", |
| 38 | +}) |
| 39 | + |
| 40 | +# Optimal batch parameters for text-like files |
| 41 | +# Text files process quickly, so we use larger batches with sequential processing |
| 42 | +TEXT_FILE_BATCH_SIZE = 200 |
| 43 | +TEXT_FILE_CONCURRENT_BATCHES = 1 |
| 44 | + |
| 45 | +# Threshold percentage to determine if workload is text-heavy |
| 46 | +TEXT_FILE_PERCENTAGE_THRESHOLD = 50.0 |
| 47 | + |
| 48 | + |
| 49 | +def calculate_dynamic_batch_parameters( |
| 50 | + filepaths: list[str], |
| 51 | + config: NvidiaRAGConfig, |
| 52 | +) -> tuple[int, int]: |
| 53 | + """ |
| 54 | + Calculate optimal batch parameters dynamically based on file characteristics. |
| 55 | +
|
| 56 | + Analyzes file extensions to determine optimal batching strategy: |
| 57 | + - Text-like files (html, json, md, sh, txt): Faster processing, larger batches with less concurrency |
| 58 | + Uses files_per_batch={TEXT_FILE_BATCH_SIZE}, concurrent_batches={TEXT_FILE_CONCURRENT_BATCHES} |
| 59 | + - Complex files (pdf, images, docx, pptx, media): Smaller batches with higher concurrency |
| 60 | + Uses default configured values |
| 61 | +
|
| 62 | + The function can be enhanced in the future with additional factors: |
| 63 | + - File sizes and complexity (number of pages, resolution, etc.) |
| 64 | + - Available system resources (memory, CPU) |
| 65 | + - Historical performance metrics |
| 66 | + - Target latency and throughput requirements |
| 67 | +
|
| 68 | + Args: |
| 69 | + filepaths: List of file paths to be processed |
| 70 | + config: NvidiaRAGConfig instance containing batch configuration |
| 71 | +
|
| 72 | + Returns: |
| 73 | + Tuple of (files_per_batch, concurrent_batches) |
| 74 | + - files_per_batch: Number of files to include in each batch |
| 75 | + - concurrent_batches: Number of batches to process concurrently |
| 76 | + """ |
| 77 | + # Check if dynamic batching is enabled |
| 78 | + if not config.nv_ingest.enable_dynamic_batching: |
| 79 | + # Return default configured values without analysis |
| 80 | + logger.info( |
| 81 | + f"Dynamic batching is disabled. Using default configuration parameters for files processing: " |
| 82 | + f"files_per_batch={config.nv_ingest.files_per_batch}, concurrent_batches={config.nv_ingest.concurrent_batches}" |
| 83 | + ) |
| 84 | + return config.nv_ingest.files_per_batch, config.nv_ingest.concurrent_batches |
| 85 | + |
| 86 | + # Analyze file extensions |
| 87 | + if not filepaths: |
| 88 | + logger.warning("Empty filepaths list provided to dynamic batch calculator") |
| 89 | + return config.nv_ingest.files_per_batch, config.nv_ingest.concurrent_batches |
| 90 | + |
| 91 | + # Extract extensions and count text-like files |
| 92 | + text_file_count = 0 |
| 93 | + total_files = len(filepaths) |
| 94 | + extension_counts = {} |
| 95 | + |
| 96 | + for filepath in filepaths: |
| 97 | + # Extract extension (lowercase, without dot) |
| 98 | + ext = Path(filepath).suffix.lstrip(".").lower() |
| 99 | + |
| 100 | + # Track extension distribution |
| 101 | + extension_counts[ext] = extension_counts.get(ext, 0) + 1 |
| 102 | + |
| 103 | + # Check if it's a text-like file |
| 104 | + if ext in TEXT_LIKE_EXTENSIONS: |
| 105 | + text_file_count += 1 |
| 106 | + |
| 107 | + # Calculate percentage of text-like files |
| 108 | + text_file_percentage = (text_file_count / total_files) * 100 if total_files > 0 else 0 |
| 109 | + |
| 110 | + # Log file distribution for debugging |
| 111 | + logger.debug( |
| 112 | + f"File distribution analysis: Total={total_files}, " |
| 113 | + f"Text-like={text_file_count} ({text_file_percentage:.1f}%), " |
| 114 | + f"Extensions={dict(extension_counts)}" |
| 115 | + ) |
| 116 | + |
| 117 | + # Decision logic: If majority (>TEXT_FILE_PERCENTAGE_THRESHOLD%) are text-like files, optimize for them |
| 118 | + if text_file_percentage > TEXT_FILE_PERCENTAGE_THRESHOLD: |
| 119 | + files_per_batch = TEXT_FILE_BATCH_SIZE |
| 120 | + concurrent_batches = TEXT_FILE_CONCURRENT_BATCHES |
| 121 | + logger.info( |
| 122 | + f"Dynamic batching: Detected {text_file_percentage:.1f}% text-like files. " |
| 123 | + f"Using optimized parameters for text processing: " |
| 124 | + f"files_per_batch={files_per_batch}, concurrent_batches={concurrent_batches}" |
| 125 | + ) |
| 126 | + else: |
| 127 | + # Use default configuration for other files (PDFs, images, media, etc.) |
| 128 | + files_per_batch = config.nv_ingest.files_per_batch |
| 129 | + concurrent_batches = config.nv_ingest.concurrent_batches |
| 130 | + logger.info( |
| 131 | + f"Dynamic batching: Detected {text_file_percentage:.1f}% text-like files. " |
| 132 | + f"Using default configuration parameters for files processing: " |
| 133 | + f"files_per_batch={files_per_batch}, concurrent_batches={concurrent_batches}" |
| 134 | + ) |
| 135 | + |
| 136 | + return files_per_batch, concurrent_batches |
0 commit comments