Merged in feature/remove-bcbs-main (pull request #854)
Remove main * Remove main Approved-by: Siddhant Medar
This commit is contained in:
Katon Minhas
@@ -1,298 +0,0 @@
|
||||
import concurrent.futures
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
|
||||
import pandas as pd
|
||||
|
||||
random.seed(42)
|
||||
|
||||
import src.pipelines.saas.file_processing as file_processing
|
||||
import src.prompts.prompt_templates as prompt_templates
|
||||
import src.utils.io_utils as io_utils
|
||||
import src.utils.llm_utils as llm_utils
|
||||
import src.utils.logging_utils as logging_utils
|
||||
import src.utils.usage_tracking as usage_tracking
|
||||
import src.utils.timing_utils as timing_utils
|
||||
from src.constants.constants import Constants
|
||||
from src import config
|
||||
from src.parent_child import main as parent_child_main
|
||||
from src.document_classification import main as dtc_main
|
||||
from src.qc_qa.pipeline import (
|
||||
generate_statistics,
|
||||
run_qc_qa_pipeline,
|
||||
save_qc_qa_outputs,
|
||||
)
|
||||
|
||||
|
||||
def safe_process_file(item, constants, run_timestamp):
|
||||
"""Process a single file with comprehensive error handling.
|
||||
|
||||
This function wraps the main file processing logic to ensure that:
|
||||
- Individual file failures don't crash the entire batch
|
||||
- Detailed error information is captured for debugging
|
||||
|
||||
Args:
|
||||
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: 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.
|
||||
With FAISS migration, no special resource cleanup is needed.
|
||||
"""
|
||||
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:
|
||||
result = file_processing.process_file(item, constants, run_timestamp)
|
||||
return result
|
||||
except (
|
||||
Exception
|
||||
) as e: # When there's an issue with the processing inside the future
|
||||
error_type = type(e).__name__
|
||||
error_message = str(e)
|
||||
full_traceback = traceback.format_exc()
|
||||
|
||||
logging.error(f"Error processing file {file_id}")
|
||||
logging.error(f"Error type: {error_type}")
|
||||
logging.error(f"Error Message: {error_message}")
|
||||
|
||||
logging.error(f"Full traceback:\n{full_traceback}")
|
||||
|
||||
return pd.DataFrame(
|
||||
[
|
||||
{
|
||||
"FILE_NAME": file_id,
|
||||
"error": error_message.replace(",", ";").replace(
|
||||
"\n", " "
|
||||
), # Replace problematic characters for CSV compatibility
|
||||
"error_type": error_type,
|
||||
"traceback": full_traceback.replace(",", ";").replace(
|
||||
"\n", " | "
|
||||
), # keep line breaks as separators
|
||||
}
|
||||
]
|
||||
) # Return a single-row dataframe so we can still concat it later
|
||||
|
||||
|
||||
def main(testing=False, test_params={}):
|
||||
# Set up per-file logging (creates separate log files for each document)
|
||||
logging_utils.setup_per_file_logging()
|
||||
|
||||
# Windows paths cannot contain ':'; use '-' in time component
|
||||
run_timestamp = datetime.now().strftime(f"run_%Y%m%d_%H-%M_{config.BATCH_ID}")
|
||||
|
||||
# Track overall pipeline time
|
||||
pipeline_start = time.time()
|
||||
|
||||
# Read Constants
|
||||
with timing_utils.timed_block("read_constants", log_level="INFO"):
|
||||
constants = Constants()
|
||||
|
||||
# Read and process input
|
||||
with timing_utils.timed_block("read_input", log_level="INFO"):
|
||||
if not testing:
|
||||
input_dict = io_utils.read_input()
|
||||
max_workers = config.MAX_WORKERS
|
||||
else:
|
||||
input_dict = io_utils.read_input(test_params["local_input_dir"])
|
||||
max_workers = test_params["max_workers"]
|
||||
if "input_files" in test_params:
|
||||
input_dict = {
|
||||
k: v
|
||||
for k, v in input_dict.items()
|
||||
if k in test_params["input_files"]
|
||||
}
|
||||
|
||||
logging.info(f"Total Input Files : {len(input_dict)}")
|
||||
|
||||
# Run document type classification if enabled - DTC mode only, then exit
|
||||
if config.PERFORM_DTC:
|
||||
logging.info("=" * 80)
|
||||
logging.info("DOCUMENT TYPE CLASSIFICATION MODE - Running DTC and exiting")
|
||||
logging.info("=" * 80)
|
||||
|
||||
# run DTC classification
|
||||
dtc_main.main(input_dict, run_timestamp)
|
||||
|
||||
logging.info(
|
||||
"DTC classification complete. Exiting (investment pipeline not run)."
|
||||
)
|
||||
logging.info("=" * 80)
|
||||
return
|
||||
|
||||
# Warm prompt caches before parallel processing
|
||||
# This ensures the cache is registered before multiple threads try to use it
|
||||
logging.info("Warming prompt caches before parallel processing...")
|
||||
with timing_utils.timed_block("warm_prompt_caches", log_level="INFO"):
|
||||
try:
|
||||
cacheable_instructions = prompt_templates.get_cacheable_instructions()
|
||||
for cache_key, instruction in cacheable_instructions.items():
|
||||
llm_utils.warm_prompt_cache(
|
||||
instruction=instruction,
|
||||
cache_key=cache_key,
|
||||
model_id="sonnet_latest",
|
||||
)
|
||||
logging.info(
|
||||
f"Successfully warmed {len(cacheable_instructions)} prompt caches"
|
||||
)
|
||||
except Exception as e:
|
||||
logging.warning(f"Error warming caches (will proceed anyway): {str(e)}")
|
||||
|
||||
# Process files concurrently - each thread gets its own per-file logging context
|
||||
successful_results = []
|
||||
error_results = []
|
||||
|
||||
logging.info(f"Starting parallel file processing with {max_workers} workers...")
|
||||
with timing_utils.timed_block("parallel_file_processing", log_level="INFO"):
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
# Each submitted task will set its own logging context in file_processing.process_file()
|
||||
# We use an executor to process files concurrently; use submit() instead of map()
|
||||
# so that we can catch exceptions and continue processing other files
|
||||
futures = [
|
||||
executor.submit(safe_process_file, item, constants, run_timestamp)
|
||||
for item in input_dict.items()
|
||||
]
|
||||
|
||||
# collect results as they complete
|
||||
completed_count = 0
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
try:
|
||||
results = future.result()
|
||||
if "error" in results.columns:
|
||||
error_results.append(results)
|
||||
else:
|
||||
successful_results.append(results)
|
||||
|
||||
completed_count += 1
|
||||
if completed_count % 5 == 0: # Log progress every 5 files
|
||||
logging.info(
|
||||
f"Processed {completed_count}/{len(futures)} files..."
|
||||
)
|
||||
except (
|
||||
Exception
|
||||
) as e: # When there's an issue with the future itself (not the processing inside the future)
|
||||
logging.error(f"Error processing future: {str(e)}")
|
||||
# Coerce this to the format expected by pd.concat (so a single-row dataframe)
|
||||
error_results.append(pd.DataFrame([{"error": str(e)}]))
|
||||
|
||||
# only concat if we have results
|
||||
with timing_utils.timed_block("concat_results", log_level="INFO"):
|
||||
if len(successful_results) > 0:
|
||||
FINAL_RESULT_DF = pd.concat(successful_results, ignore_index=True)
|
||||
else:
|
||||
FINAL_RESULT_DF = pd.DataFrame()
|
||||
|
||||
if len(error_results) > 0:
|
||||
ERROR_RESULT_DF = pd.concat(error_results, ignore_index=True)
|
||||
else:
|
||||
ERROR_RESULT_DF = pd.DataFrame()
|
||||
|
||||
# Run QC/QA validation if enabled (creates separate validated copy)
|
||||
if config.ENABLE_QC_QA and not FINAL_RESULT_DF.empty:
|
||||
logging.info("=" * 80)
|
||||
logging.info("Running QC/QA Validation Pipeline...")
|
||||
logging.info("=" * 80)
|
||||
with timing_utils.timed_block("qc_qa_validation", log_level="INFO"):
|
||||
try:
|
||||
# Ensure unique index before QC/QA (safety check)
|
||||
if FINAL_RESULT_DF.index.has_duplicates:
|
||||
logging.warning(
|
||||
"Detected duplicate indices in FINAL_RESULT_DF, resetting index..."
|
||||
)
|
||||
FINAL_RESULT_DF = FINAL_RESULT_DF.reset_index(drop=True)
|
||||
|
||||
# Create a copy for validation - preserve original FINAL_RESULT_DF
|
||||
validated_df = run_qc_qa_pipeline(FINAL_RESULT_DF.copy(), stats_df=None)
|
||||
logging.info("QC/QA validation completed successfully")
|
||||
|
||||
# Generate validation statistics
|
||||
logging.info("Generating QC/QA statistics...")
|
||||
qc_qa_stats_df = generate_statistics(validated_df)
|
||||
logging.info("QC/QA statistics generated successfully")
|
||||
|
||||
# Save QC/QA outputs (local + S3) - separate from main results
|
||||
save_qc_qa_outputs(
|
||||
validated_df=validated_df,
|
||||
stats_df=qc_qa_stats_df,
|
||||
run_timestamp=run_timestamp,
|
||||
write_to_s3=config.WRITE_TO_S3,
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"QC/QA validation failed: {str(e)}")
|
||||
logging.error(f"Full traceback:\n{traceback.format_exc()}")
|
||||
logging.warning("Continuing with original unvalidated results...")
|
||||
logging.info("=" * 80)
|
||||
|
||||
if not testing:
|
||||
# Get usage tracking dataframes
|
||||
per_file_usage_df, batch_summary_df = usage_tracking.get_usage_dataframes(
|
||||
config.BATCH_ID, run_timestamp
|
||||
)
|
||||
|
||||
with timing_utils.timed_block("write_outputs", log_level="INFO"):
|
||||
if config.WRITE_TO_S3:
|
||||
doczy_output_for_pc = io_utils.write_s3(
|
||||
FINAL_RESULT_DF, "", run_timestamp, "final"
|
||||
)
|
||||
if not ERROR_RESULT_DF.empty:
|
||||
io_utils.write_s3(ERROR_RESULT_DF, "", run_timestamp, "error")
|
||||
logging.warning(
|
||||
f"{len(ERROR_RESULT_DF)} files failed processing - error results written to S3"
|
||||
)
|
||||
|
||||
# Export usage and cost tracking data to S3 only if both dataframes have data
|
||||
if not per_file_usage_df.empty and not batch_summary_df.empty:
|
||||
logging.info("Exporting usage and cost tracking data to S3...")
|
||||
io_utils.write_s3(per_file_usage_df, "", run_timestamp, "usage")
|
||||
io_utils.write_s3(
|
||||
batch_summary_df, "", run_timestamp, "usage_summary"
|
||||
)
|
||||
else:
|
||||
doczy_output_for_pc = io_utils.write_local(
|
||||
FINAL_RESULT_DF, "", run_timestamp, "final"
|
||||
)
|
||||
if not ERROR_RESULT_DF.empty:
|
||||
io_utils.write_local(ERROR_RESULT_DF, "", run_timestamp, "error")
|
||||
logging.warning(
|
||||
f"{len(ERROR_RESULT_DF)} files failed processing - error results written locally"
|
||||
)
|
||||
|
||||
if config.PERFORM_PARENT_CHILD_MAPPING:
|
||||
if FINAL_RESULT_DF.empty:
|
||||
logging.warning(
|
||||
"Skipping Parent-Child mapping: No successful results to process (FINAL_RESULT_DF is empty)"
|
||||
)
|
||||
else:
|
||||
logging.info(
|
||||
f"Starting Parent-Child mapping with {len(FINAL_RESULT_DF)} records from: {doczy_output_for_pc}"
|
||||
)
|
||||
with timing_utils.timed_block("parent_child_mapping", log_level="INFO"):
|
||||
parent_child_main.main(doczy_output_for_pc)
|
||||
|
||||
# Export usage and cost tracking data locally only if both dataframes have data
|
||||
if not per_file_usage_df.empty and not batch_summary_df.empty:
|
||||
logging.info("Exporting usage and cost tracking data locally...")
|
||||
io_utils.write_local(per_file_usage_df, "", run_timestamp, "usage")
|
||||
io_utils.write_local(batch_summary_df, "", run_timestamp, "usage_summary")
|
||||
|
||||
# Print timing summary
|
||||
pipeline_duration = time.time() - pipeline_start
|
||||
logging.info("=" * 80)
|
||||
logging.info(f"PIPELINE COMPLETE - Total time: {pipeline_duration:.2f}s")
|
||||
logging.info("=" * 80)
|
||||
timing_utils.print_hierarchical_timing_summary()
|
||||
else:
|
||||
return FINAL_RESULT_DF, ERROR_RESULT_DF
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user