foxhunt/runpod/s3_monitor.py

"""
S3 Log Monitor - Stream training logs from RunPod S3
"""

import boto3
import time
import re
from typing import Optional, Callable
from datetime import datetime, timedelta


class S3LogMonitor:
    """Monitor and stream training logs from RunPod S3."""

    def __init__(
        self,
        bucket_name: str,
        aws_access_key: str,
        aws_secret_key: str,
        endpoint_url: str = "https://s3api-eur-is-1.runpod.io"
    ):
        """
        Initialize S3 log monitor.

        Args:
            bucket_name: S3 bucket name
            aws_access_key: AWS access key ID
            aws_secret_key: AWS secret access key
            endpoint_url: S3 endpoint URL
        """
        self.bucket_name = bucket_name

        self.s3 = boto3.client(
            's3',
            aws_access_key_id=aws_access_key,
            aws_secret_access_key=aws_secret_key,
            endpoint_url=endpoint_url
        )

    def stream_logs(
        self,
        pod_id: str,
        interval: int = 10,
        timeout: Optional[int] = None,
        completion_callback: Optional[Callable[[str], bool]] = None
    ):
        """
        Stream logs from S3 in real-time.

        Args:
            pod_id: Pod ID to monitor
            interval: Polling interval in seconds
            timeout: Maximum monitoring time in seconds (None = infinite)
            completion_callback: Function to check if training is complete
                                Returns True if complete, False otherwise
        """
        log_key = f"logs/{pod_id}/training.log"
        last_position = 0
        start_time = time.time()

        print(f"\n{'='*70}")
        print(f"MONITORING POD: {pod_id}")
        print(f"{'='*70}")
        print(f"Log Key:    {log_key}")
        print(f"Interval:   {interval}s")
        if timeout:
            print(f"Timeout:    {timeout}s ({timeout // 60}m)")
        print(f"{'='*70}\n")

        try:
            while True:
                # Check timeout
                if timeout and (time.time() - start_time) > timeout:
                    print(f"\n⏱️  Monitoring timeout reached ({timeout}s)")
                    break

                # Try to fetch log file
                try:
                    response = self.s3.get_object(
                        Bucket=self.bucket_name,
                        Key=log_key,
                        Range=f"bytes={last_position}-"
                    )

                    # Read new content
                    new_content = response['Body'].read().decode('utf-8', errors='ignore')

                    if new_content:
                        print(new_content, end='')
                        last_position += len(new_content.encode('utf-8'))

                        # Check completion
                        if completion_callback and completion_callback(new_content):
                            print(f"\n✅ Training complete detected!")
                            break

                except self.s3.exceptions.NoSuchKey:
                    # Log file doesn't exist yet
                    print(f"  [Waiting for logs...]", end='\r')

                except Exception as e:
                    error_str = str(e)
                    if "InvalidRange" not in error_str:
                        print(f"\n  ⚠️  Error reading logs: {error_str[:100]}")

                # Wait before next poll
                time.sleep(interval)

        except KeyboardInterrupt:
            print(f"\n\n⏸️  Monitoring interrupted by user")

    def check_completion(self, log_content: str) -> bool:
        """
        Check if training is complete based on log content.

        Args:
            log_content: Recent log content

        Returns:
            True if training complete, False otherwise
        """
        # Common completion patterns
        completion_patterns = [
            r"Training complete",
            r"Model saved successfully",
            r"Checkpoint saved.*final",
            r"All epochs completed",
            r"✅.*complete",
        ]

        for pattern in completion_patterns:
            if re.search(pattern, log_content, re.IGNORECASE):
                return True

        return False

    def get_recent_logs(self, pod_id: str, lines: int = 50) -> Optional[str]:
        """
        Get recent log lines.

        Args:
            pod_id: Pod ID
            lines: Number of lines to fetch (approximate)

        Returns:
            Recent log content or None if not available
        """
        log_key = f"logs/{pod_id}/training.log"

        try:
            # Get object metadata to find size
            response = self.s3.head_object(
                Bucket=self.bucket_name,
                Key=log_key
            )

            size = response['ContentLength']

            # Estimate bytes to fetch (assuming ~100 bytes per line)
            fetch_bytes = lines * 100
            start = max(0, size - fetch_bytes)

            # Fetch tail of file
            response = self.s3.get_object(
                Bucket=self.bucket_name,
                Key=log_key,
                Range=f"bytes={start}-"
            )

            content = response['Body'].read().decode('utf-8', errors='ignore')

            # Return last N lines
            all_lines = content.split('\n')
            return '\n'.join(all_lines[-lines:])

        except Exception as e:
            print(f"  ⚠️  Error fetching logs: {str(e)[:100]}")
            return None