doczyai-pipelines/fieldExtraction/src/batch_tracking.py

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
            }