""" Pod monitoring with S3 log tailing. Provides status polling, log streaming, completion detection, and auto-termination. """ import re import time from rich.console import Console from rich.live import Live from rich.text import Text from .client import RunPodClient from .config import RunPodConfig, get_config from .s3_client import S3Client from .errors import ( PodNotFoundError, PodTimeoutError, S3ObjectNotFoundError, RunPodError, ) console = Console() class PodMonitor: """ Monitor RunPod pods with S3 log tailing. NO SSH required - uses S3 byte-range requests to tail logs. """ def __init__( self, pod_id: str, config: RunPodConfig | None = None, client: RunPodClient | None = None, s3_client: S3Client | None = None, ): """ Initialize pod monitor. Args: pod_id: Pod ID to monitor config: RunPodConfig instance (uses global if None) client: RunPodClient instance (creates if None) s3_client: S3Client instance (creates if None) """ self.pod_id = pod_id self.config = config or get_config() self.client = client or RunPodClient( api_key=self.config.runpod_api_key, volume_id=self.config.runpod_volume_id, registry_auth_id=self.config.runpod_container_registry_auth_id ) self.s3_client = s3_client or S3Client(self.config) # Log tailing state self.log_position = 0 self.last_log_check = 0 # Status self.training_complete = False self.error_detected = False def wait_until_running( self, timeout: int = 300, poll_interval: int | None = None ) -> bool: """ Wait until pod is in RUNNING state. Args: timeout: Maximum wait time in seconds poll_interval: Polling interval (uses config default if None) Returns: True if pod is running Raises: PodTimeoutError: If pod doesn't start within timeout PodNotFoundError: If pod doesn't exist """ poll_interval = poll_interval or self.config.pod_status_poll_interval start_time = time.time() console.print(f"Waiting for pod {self.pod_id} to start...") while True: elapsed = time.time() - start_time if elapsed > timeout: raise PodTimeoutError(self.pod_id, "start", timeout) try: status_data = self.client.get_pod_status(self.pod_id) runtime_status = status_data.get("runtime", {}).get("status", "UNKNOWN") console.print( f" Status: {runtime_status} (elapsed: {int(elapsed)}s)", end="\r" ) if runtime_status == "RUNNING": console.print( f"\n[green]✅ Pod is running! (took {int(elapsed)}s)[/green]" ) return True if runtime_status == "FAILED": console.print(f"\n[red]❌ Pod failed to start[/red]") return False except PodNotFoundError: console.print(f"\n[red]❌ Pod not found: {self.pod_id}[/red]") raise time.sleep(poll_interval) def stream_s3_logs( self, follow: bool = True, poll_interval: int | None = None, max_lines: int | None = None, ) -> None: """ Stream logs from S3 (byte-range tailing). Args: follow: Continue streaming until completion detected poll_interval: Polling interval (uses config default if None) max_lines: Maximum lines to display (None for unlimited) Raises: RunPodError: On errors """ poll_interval = poll_interval or self.config.log_poll_interval log_key = self.config.get_log_s3_key() console.print(f"\nStreaming logs from s3://{self.s3_client.bucket}/{log_key}") console.print("(Polling every {poll_interval}s, Ctrl+C to stop)\n") console.print("-" * 70) lines_shown = 0 try: while True: # Wait for log file to appear if not self.s3_client.object_exists(log_key): if not follow: console.print( "[yellow]⚠ Log file not found (training may not have started)[/yellow]" ) return # Still waiting for log file time.sleep(poll_interval) continue # Tail new content try: content, new_position = self.s3_client.tail_log_file( log_key, start_byte=self.log_position ) if content: # Decode and print new lines text = content.decode("utf-8", errors="ignore") lines = text.splitlines() for line in lines: console.print(line) lines_shown += 1 # Check for completion/error patterns self._check_line_patterns(line) # Max lines limit if max_lines and lines_shown >= max_lines: console.print( f"\n[yellow]Reached max lines ({max_lines})[/yellow]" ) return self.log_position = new_position except S3ObjectNotFoundError: # Log file was deleted or doesn't exist yet pass # Check completion if self.training_complete or self.error_detected: console.print("\n" + "-" * 70) if self.training_complete: console.print("[green]✅ Training completed![/green]") if self.error_detected: console.print("[red]❌ Error detected in logs[/red]") return if not follow: return time.sleep(poll_interval) except KeyboardInterrupt: console.print("\n\n[yellow]Streaming stopped by user[/yellow]") def _check_line_patterns(self, line: str) -> None: """Check log line for completion/error patterns.""" # Check completion patterns for pattern in self.config.completion_check_patterns: if pattern.lower() in line.lower(): self.training_complete = True return # Check error patterns for pattern in self.config.error_patterns: if pattern in line: self.error_detected = True return def check_completion(self) -> tuple[bool, bool]: """ Check if training has completed or errored. Returns: Tuple of (completed, error_detected) """ return self.training_complete, self.error_detected def auto_terminate(self, wait_for_completion: bool = True) -> bool: """ Automatically terminate pod when training completes. Args: wait_for_completion: Wait for completion before terminating Returns: True if pod was terminated Raises: RunPodError: On errors """ if wait_for_completion: console.print("\nWaiting for training to complete...") # Stream logs until completion self.stream_s3_logs(follow=True) # Check status completed, error = self.check_completion() if not completed and not error: console.print( "[yellow]⚠ Training not complete, skipping termination[/yellow]" ) return False # Terminate pod console.print(f"\nTerminating pod {self.pod_id}...") try: self.client.terminate_pod(self.pod_id) return True except Exception as e: console.print(f"[yellow]⚠ Failed to terminate: {e}[/yellow]") return False def get_pod_info(self) -> dict: """ Get current pod information. Returns: Pod status dict """ try: return self.client.get_pod_status(self.pod_id) except PodNotFoundError: return {} def display_pod_info(self) -> None: """Display current pod information in formatted output.""" info = self.get_pod_info() if not info: console.print(f"[yellow]Pod {self.pod_id} not found[/yellow]") return console.print("\n" + "=" * 70) console.print(f"POD INFO: {self.pod_id}") console.print("=" * 70) console.print(f"Status: {info.get('desiredStatus', 'UNKNOWN')}") runtime = info.get("runtime", {}) runtime_status = runtime.get("status", "UNKNOWN") console.print(f"Runtime: {runtime_status}") machine = info.get("machine", {}) if machine: gpu_info = machine.get("gpuType", {}) console.print(f"GPU: {gpu_info.get('displayName', 'N/A')}") console.print(f"Cost/hr: ${info.get('costPerHr', 'N/A')}") console.print("=" * 70)