doczyai-pipelines/src/tracking/batch_tracking.py

import csv
import os
import subprocess
from datetime import datetime
from pathlib import Path

import boto3
import pandas as pd
from src import config


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}