import pandas as pd import boto3 import csv from datetime import datetime from pathlib import Path import subprocess import config import os class BatchTracker: def __init__(self): self.s3_client = boto3.client('s3') self.batch_tracker_bucket = config.S3_OUTPUT_BUCKET self.batch_tracker_key = 'batch-tracker/doczyai-batch-tracker.csv' self.current_batch_start = datetime.now() self.files_processed = 0 self.processed_files = set() # Track unique files self.processing_events = 0 # Track total processing events self.total_input_files = 0 self.reimbursement_files = 0 # Define the file log fields in correct order self.file_log_fields = [ 'filename', 'process_type', 'start_time', 'end_time', 'processing_time_seconds', 'delta_time_hours', 'status', 'error', 'memory_usage_mb', 'cpu_percent' ] def set_input_file_counts(self, total_files, reimbursement_files): """Set the input file counts for tracking""" self.total_input_files = total_files self.reimbursement_files = reimbursement_files def get_git_info(self): """Get current git branch, commit hash, and metadata""" try: # Get current branch branch = subprocess.check_output( ['git', 'rev-parse', '--abbrev-ref', 'HEAD'], universal_newlines=True ).strip() # Get current commit hash commit = subprocess.check_output( ['git', 'rev-parse', 'HEAD'], universal_newlines=True ).strip() # Get commit message commit_msg = subprocess.check_output( ['git', 'log', '-1', '--pretty=%B'], universal_newlines=True ).strip() # Get last commit date commit_date = subprocess.check_output( ['git', 'log', '-1', '--format=%cd'], universal_newlines=True ).strip() return { 'branch': branch, 'commit_hash': commit, 'commit_message': commit_msg, 'commit_date': commit_date } except Exception as e: return { 'branch': 'Unknown', 'commit_hash': 'Unknown', 'commit_message': 'Unknown', 'commit_date': 'Unknown' } def get_batch_tracker(self): """Retrieve or create the batch tracker DataFrame""" try: response = self.s3_client.get_object( Bucket=self.batch_tracker_bucket, Key=self.batch_tracker_key ) return pd.read_csv(response['Body']) except Exception as e: print(f"Creating new batch tracker: {str(e)}") return pd.DataFrame(columns=[ 'batch_id', 'state', 'run_mode', 'read_mode', 'input_dir', 'output_dir', 'fields', 'max_workers', 'start_time', 'end_time', 'total_input_files', 'reimbursement_files', 'files_processed', 'contracts_per_hour', 'status', 'machine_info', 'git_branch', 'git_commit', 'git_commit_message', 'git_commit_date' ]) def get_machine_info(self): """Get information about the machine running the process""" import platform import socket return { 'hostname': socket.gethostname(), 'platform': platform.platform(), 'processor': platform.processor(), 'python_version': platform.python_version() } def add_batch_run(self): """Add new batch run to the tracker""" df = self.get_batch_tracker() machine_info = self.get_machine_info() git_info = self.get_git_info() new_row = { 'batch_id': config.BATCH_ID, 'state': config.STATE, 'run_mode': config.RUN_MODE, 'read_mode': config.READ_MODE, 'input_dir': config.LOCAL_PATH, 'output_dir': config.OUTPUT_DIRECTORY, 'fields': config.FIELDS, 'max_workers': config.MAX_WORKERS, 'start_time': self.current_batch_start.strftime('%Y-%m-%d %H:%M:%S'), 'end_time': None, 'total_input_files': self.total_input_files, 'reimbursement_files': self.reimbursement_files, 'files_processed': 0, 'contracts_per_hour': 0, 'status': 'Running', 'machine_info': str(machine_info), 'git_branch': git_info['branch'], 'git_commit': git_info['commit_hash'], 'git_commit_message': git_info['commit_message'], 'git_commit_date': git_info['commit_date'] } df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True) self.update_tracker(df) # Initialize the file log with headers if it doesn't exist if not os.path.exists(config.FILE_LOG_NAME): with open(config.FILE_LOG_NAME, 'w', newline='', encoding='utf-8') as f: writer = csv.DictWriter(f, fieldnames=self.file_log_fields) writer.writeheader() def update_batch_run(self, end_time=None): """Update the current batch run stats""" df = self.get_batch_tracker() current_batch_entries = df[df['batch_id'] == config.BATCH_ID] if current_batch_entries.empty: self.add_batch_run() return latest_idx = current_batch_entries.index[-1] if end_time: df.loc[latest_idx, 'end_time'] = end_time.strftime('%Y-%m-%d %H:%M:%S') df.loc[latest_idx, 'status'] = 'Completed' df.loc[latest_idx, 'files_processed'] = len(self.processed_files) # Use unique files count try: start_time = datetime.strptime(df.loc[latest_idx, 'start_time'], '%Y-%m-%d %H:%M:%S') current_time = end_time if end_time else datetime.now() hours_elapsed = (current_time - start_time).total_seconds() / 3600 if hours_elapsed > 0: contracts_per_hour = len(self.processed_files) / hours_elapsed # Use unique files count df.loc[latest_idx, 'contracts_per_hour'] = round(contracts_per_hour, 2) except Exception as e: print(f"Error calculating processing speed: {str(e)}") self.update_tracker(df) if end_time and len(self.processed_files) > 0: try: print(f"\nBatch Performance Metrics:") print(f"Unique files processed: {len(self.processed_files)}") print(f"Total processing events: {self.processing_events}") print(f"Total time elapsed: {round(hours_elapsed, 2)} hours") print(f"Average processing speed: {round(contracts_per_hour, 2)} unique files/hour") except: print("Could not calculate final performance metrics") def update_tracker(self, df): """Save tracker to S3""" try: csv_buffer = df.to_csv(index=False).encode() self.s3_client.put_object( Bucket=self.batch_tracker_bucket, Key=self.batch_tracker_key, Body=csv_buffer ) df.to_csv('doczyai-batch-tracker-local.csv', index=False) except Exception as e: print(f"Error updating tracker in S3: {str(e)}") df.to_csv('doczyai-batch-tracker-local.csv', index=False) def update_file_progress(self, filename, process_type, start_time, end_time, status='Completed', error=None): """Update progress for individual file processing""" try: processing_time = (end_time - start_time).total_seconds() delta_time_hours = processing_time / 3600 # Convert to hours system_metrics = self.get_system_metrics() new_row = { 'filename': filename, 'process_type': process_type, 'start_time': start_time.strftime('%Y-%m-%d %H:%M:%S'), 'end_time': end_time.strftime('%Y-%m-%d %H:%M:%S'), 'processing_time_seconds': processing_time, 'delta_time_hours': round(delta_time_hours, 4), 'status': status, 'error': error, 'memory_usage_mb': system_metrics['memory_usage_mb'], 'cpu_percent': system_metrics['cpu_percent'] } # Write to file log if not os.path.exists(config.FILE_LOG_NAME): with open(config.FILE_LOG_NAME, 'w', newline='', encoding='utf-8') as f: writer = csv.DictWriter(f, fieldnames=self.file_log_fields) writer.writeheader() with open(config.FILE_LOG_NAME, 'a', newline='', encoding='utf-8') as f: writer = csv.DictWriter(f, fieldnames=self.file_log_fields) writer.writerow(new_row) # Track both unique files and processing events if status == 'Completed': self.processed_files.add(filename) # Add to set of unique files self.processing_events += 1 # Increment total events self.files_processed = len(self.processed_files) # Update files_processed to unique count self.update_batch_run() except Exception as e: print(f"Error updating file progress: {str(e)}") def get_system_metrics(self): """Get current system metrics""" import psutil try: process = psutil.Process() return { 'memory_usage_mb': process.memory_info().rss / 1024 / 1024, 'cpu_percent': process.cpu_percent() } except: return { 'memory_usage_mb': 0, 'cpu_percent': 0 }