NexQuant/scripts/nexquant_parallel.py

"""
NexQuant Parallel Runner - Run multiple factor experiments concurrently.

Spawns N subprocesses, each running `nexquant.py quant` with isolated config:
- Separate log files (fin_quant_run1.log, fin_quant_run2.log, etc.)
- Separate result directories (results/runs/run1/, results/runs/run2/, etc.)
- Separate workspace directories
- API key distribution across multiple keys (round-robin)

Usage:
    python nexquant_parallel.py --runs 5 --api-keys 2
    python nexquant_parallel.py --runs 3 --model openrouter
    python nexquant_parallel.py --runs 5 --model local --api-keys 1
"""
import os
import signal
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path

from dotenv import load_dotenv
from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown
from rich.panel import Panel
from rich.table import Table

# Load environment variables from .env file
load_dotenv(Path(__file__).parent.parent / ".env")

console = Console()


class RunState:
    """Tracks the state of a single parallel run."""

    def __init__(self, run_id: int, api_key_idx: int, model: str):
        self.run_id = run_id
        self.api_key_idx = api_key_idx
        self.model = model
        self.process: subprocess.Popen | None = None
        self.status: str = "pending"  # pending, running, success, failed, stopped
        self.start_time: datetime | None = None
        self.end_time: datetime | None = None
        self.exit_code: int | None = None
        self.error_message: str | None = None
        self.log_file: str = f"fin_quant_run{run_id}.log"

    @property
    def elapsed(self) -> str:
        """Get elapsed time as human-readable string."""
        if self.start_time is None:
            return "--:--:--"
        end = self.end_time or datetime.now()
        delta = end - self.start_time
        total_seconds = int(delta.total_seconds())
        hours, remainder = divmod(total_seconds, 3600)
        minutes, seconds = divmod(remainder, 60)
        return f"{hours:02d}:{minutes:02d}:{seconds:02d}"

    @property
    def status_icon(self) -> str:
        """Get icon for current status."""
        icons = {
            "pending": "⏳",
            "running": "🔄",
            "success": "✅",
            "failed": "❌",
            "stopped": "⏹️",
        }
        return icons.get(self.status, "❓")


class ParallelRunner:
    """
    Manages multiple concurrent factor experiment runs.

    Spawns subprocesses with isolated configurations, monitors progress,
    and handles graceful shutdown.
    """

    def __init__(
        self,
        num_runs: int = 5,
        num_api_keys: int = 2,
        model: str = "openrouter",
    ):
        """
        Initialize parallel runner.

        Parameters
        ----------
        num_runs : int
            Number of concurrent runs to spawn
        num_api_keys : int
            Number of API keys to distribute across (1 or 2)
        model : str
            LLM backend: 'local' (llama.cpp) or 'openrouter' (cloud)
        """
        self.num_runs = num_runs
        self.num_api_keys = num_api_keys
        self.model = model
        self.runs: list[RunState] = []
        self.project_root = Path(__file__).parent.parent
        self._shutdown_requested = False

        # Read API keys from environment
        self.api_keys = self._load_api_keys()

        # Validate we have enough API keys
        if self.model == "openrouter" and len(self.api_keys) < num_api_keys:
            console.print(
                f"[yellow]⚠️  Requested {num_api_keys} API keys, but only {len(self.api_keys)} found in .env[/yellow]",
            )
            console.print(
                f"[dim]Distributing across {len(self.api_keys)} available key(s)[/dim]",
            )
            self.num_api_keys = len(self.api_keys)

        # Initialize run states
        for i in range(1, num_runs + 1):
            # Round-robin API key assignment
            api_key_idx = (i - 1) % max(self.num_api_keys, 1)
            run_state = RunState(run_id=i, api_key_idx=api_key_idx, model=model)
            self.runs.append(run_state)

    def _load_api_keys(self) -> list[str]:
        """Load API keys from environment variables."""
        keys = []

        if self.model == "openrouter":
            key1 = os.getenv("OPENROUTER_API_KEY", "")
            key2 = os.getenv("OPENROUTER_API_KEY_2", "")
            if key1:
                keys.append(key1)
            if key2:
                keys.append(key2)
        else:
            # For local mode, we just need the llama.cpp endpoint
            keys.append("local")

        if not keys or (len(keys) == 1 and keys[0] == "local"):
            keys = ["local"]

        return keys

    def _build_env(self, run_state: RunState) -> dict[str, str]:
        """
        Build isolated environment for a subprocess.

        Parameters
        ----------
        run_state : RunState
            The run state object containing run configuration

        Returns
        -------
        dict
            Environment variables dict for subprocess
        """
        # Start with a copy of current environment
        env = os.environ.copy()

        # Set parallel run ID for isolation
        env["PARALLEL_RUN_ID"] = str(run_state.run_id)

        # Set workspace isolation
        workspace_dir = self.project_root / f"RD-Agent_workspace_run{run_state.run_id}"
        env["RD_AGENT_WORKSPACE"] = str(workspace_dir)

        # Configure API key for this run
        if self.model == "openrouter":
            if self.num_api_keys >= 2 and len(self.api_keys) >= 2:
                env["OPENAI_API_KEY"] = f"{self.api_keys[0]},{self.api_keys[1]}"
                env["LITELLM_PARALLEL_CALLS"] = "2"
            elif run_state.api_key_idx < len(self.api_keys):
                env["OPENAI_API_KEY"] = self.api_keys[run_state.api_key_idx]
            env["OPENAI_API_BASE"] = "https://openrouter.ai/api/v1"
            env["CHAT_MODEL"] = os.getenv("OPENROUTER_MODEL", "openrouter/google/gemma-4-26b-a4b-it:free")
        elif self.model == "local":
            env["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "local")
            env["OPENAI_API_BASE"] = os.getenv("OPENAI_API_BASE", "http://localhost:8081/v1")
            env["CHAT_MODEL"] = os.getenv("CHAT_MODEL", "openai/qwen3.5-35b")

        return env

    def _build_command(self, run_state: RunState) -> list[str]:
        """
        Build the subprocess command to run nexquant quant.

        Parameters
        ----------
        run_state : RunState
            The run state object containing run configuration

        Returns
        -------
        list
            Command list for subprocess.Popen
        """
        cmd = [
            sys.executable,  # Use same Python interpreter
            str(self.project_root / "nexquant.py"),
            "quant",
            "--model", run_state.model,
            "--run-id", str(run_state.run_id),
            "--log-file", run_state.log_file,
        ]
        return cmd

    def _start_run(self, run_state: RunState) -> None:
        """
        Start a single run as a subprocess.

        Parameters
        ----------
        run_state : RunState
            The run state to start
        """
        env = self._build_env(run_state)
        cmd = self._build_command(run_state)

        # Ensure results directory exists
        results_dir = self.project_root / "results" / "runs" / f"run{run_state.run_id}"
        results_dir.mkdir(parents=True, exist_ok=True)

        # Open log file for appending
        log_path = self.project_root / run_state.log_file
        log_f = open(log_path, "a", encoding="utf-8")

        try:
            # Start subprocess
            run_state.process = subprocess.Popen(
                cmd,
                env=env,
                cwd=str(self.project_root),
                stdout=log_f,
                stderr=subprocess.STDOUT,
            )
            run_state.status = "running"
            run_state.start_time = datetime.now()
        except Exception:
            log_f.close()
            raise

        console.print(
            f"[dim]  ▶️  Run {run_state.run_id} started (PID: {run_state.process.pid}, "
            f"API Key: {run_state.api_key_idx + 1}, Model: {run_state.model})[/dim]",
        )

    def _check_run(self, run_state: RunState) -> None:
        """
        Check if a run is still running and update status.

        Parameters
        ----------
        run_state : RunState
            The run state to check
        """
        if run_state.status != "running" or run_state.process is None:
            return

        poll_result = run_state.process.poll()
        if poll_result is not None:
            # Process has finished
            run_state.exit_code = poll_result
            run_state.end_time = datetime.now()

            if poll_result == 0:
                run_state.status = "success"
                console.print(
                    f"[bold green]  ✅ Run {run_state.run_id} completed "
                    f"({run_state.elapsed})[/bold green]",
                )
            else:
                run_state.status = "failed"
                run_state.error_message = f"Exit code: {poll_result}"
                console.print(
                    f"[bold red]  ❌ Run {run_state.run_id} failed "
                    f"({run_state.elapsed}, exit code: {poll_result})[/bold red]",
                )

    def _stop_run(self, run_state: RunState) -> None:
        """
        Gracefully stop a running subprocess.

        Parameters
        ----------
        run_state : RunState
            The run state to stop
        """
        if run_state.process is None or run_state.status != "running":
            return

        try:
            # Try graceful termination first
            run_state.process.terminate()
            try:
                run_state.process.wait(timeout=10)
            except subprocess.TimeoutExpired:
                # Force kill if not responding
                run_state.process.kill()
                run_state.process.wait()
        except Exception as e:
            console.print(f"[yellow]  ⚠️  Error stopping run {run_state.run_id}: {e}[/yellow]")

        run_state.status = "stopped"
        run_state.end_time = datetime.now()

    def _render_dashboard(self) -> Panel:
        """
        Render the live dashboard panel showing all run states.

        Returns
        -------
        Panel
            Rich Panel object with dashboard content
        """
        # Summary stats
        pending = sum(1 for r in self.runs if r.status == "pending")
        running = sum(1 for r in self.runs if r.status == "running")
        success = sum(1 for r in self.runs if r.status == "success")
        failed = sum(1 for r in self.runs if r.status == "failed")
        stopped = sum(1 for r in self.runs if r.status == "stopped")

        # Build summary table
        table = Table(
            title="🔀 NexQuant Parallel Run Dashboard",
            show_header=True,
            header_style="bold cyan",
            expand=True,
        )
        table.add_column("Run", justify="center", width=6)
        table.add_column("Status", justify="center", width=10)
        table.add_column("Elapsed", justify="center", width=10)
        table.add_column("API Key", justify="center", width=8)
        table.add_column("Model", justify="center", width=12)
        table.add_column("Exit", justify="center", width=6)
        table.add_column("Log File", justify="left")

        for run in self.runs:
            table.add_row(
                f"#{run.run_id}",
                f"{run.status_icon} {run.status}",
                run.elapsed,
                str(run.api_key_idx + 1),
                run.model,
                str(run.exit_code) if run.exit_code is not None else "--",
                run.log_file,
            )

        # Summary panel
        total = len(self.runs)
        summary_text = (
            f"**Summary:** {total} total | "
            f"{success} done | "
            f"{running} running | "
            f"{pending} pending | "
            f"{failed} failed"
        )

        if self._shutdown_requested:
            summary_text += "\n⚠️  **Shutdown requested - stopping all runs...**"

        from rich.console import Group
        return Group(table, Panel(Markdown(summary_text), border_style="blue"))

    def _signal_handler(self, signum, frame) -> None:
        """Handle SIGINT/SIGTERM for graceful shutdown."""
        if self._shutdown_requested:
            # Second Ctrl+C - force kill everything
            console.print("\n[bold red]🛑 Force killing all runs![/bold red]")
            for run in self.runs:
                if run.process and run.status == "running":
                    run.process.kill()
            sys.exit(1)

        self._shutdown_requested = True
        console.print("\n[yellow]⏹️  Shutdown requested - gracefully stopping all runs...[/yellow]")
        console.print("[dim]Press Ctrl+C again to force kill[/dim]")

        for run in self.runs:
            if run.status == "running":
                self._stop_run(run)

    def run(self) -> dict[str, int]:
        """
        Execute all parallel runs and show live dashboard.

        Returns
        -------
        dict
            Summary with keys: total, success, failed, stopped
        """
        # Register signal handlers
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)

        console.print(f"\n[bold cyan]{'=' * 60}[/bold cyan]")
        console.print("[bold cyan]🔀 NexQuant Parallel Runner[/bold cyan]")
        console.print(f"[bold cyan]{'=' * 60}[/bold cyan]")
        console.print(f"  Runs: {self.num_runs}")
        console.print(f"  API Keys: {self.num_api_keys} ({len(self.api_keys)} available)")
        console.print(f"  Model: {self.model}")
        console.print(f"  Log pattern: fin_quant_run{{1..{self.num_runs}}}.log")
        console.print(f"  Results: results/runs/run{{1..{self.num_runs}}}/")
        console.print()

        # Start all runs
        for run in self.runs:
            if self._shutdown_requested:
                break
            self._start_run(run)
            # Small delay to prevent overwhelming the system
            time.sleep(1)

        # Monitor loop with live dashboard
        with Live(refresh_per_second=2, screen=True) as live:
            live.update(self._render_dashboard())
            while True:
                if self._shutdown_requested:
                    # Check if all runs are stopped
                    all_stopped = all(
                        r.status in ("success", "failed", "stopped", "pending")
                        for r in self.runs
                    )
                    if all_stopped:
                        break

                # Update all run statuses
                for run in self.runs:
                    self._check_run(run)

                # Check if all runs are complete
                all_done = all(
                    r.status in ("success", "failed", "stopped")
                    for r in self.runs
                )
                if all_done:
                    break

                live.update(self._render_dashboard())
                time.sleep(0.5)

        # Final summary
        success_count = sum(1 for r in self.runs if r.status == "success")
        failed_count = sum(1 for r in self.runs if r.status == "failed")
        stopped_count = sum(1 for r in self.runs if r.status == "stopped")

        console.print(f"\n[bold cyan]{'=' * 60}[/bold cyan]")
        console.print("[bold cyan]📊 Parallel Run Summary[/bold cyan]")
        console.print(f"[bold cyan]{'=' * 60}[/bold cyan]")
        console.print(f"  ✅ Success: {success_count}/{self.num_runs}")
        console.print(f"  ❌ Failed: {failed_count}/{self.num_runs}")
        if stopped_count > 0:
            console.print(f"  ⏹️  Stopped: {stopped_count}/{self.num_runs}")

        total_time = None
        for run in self.runs:
            if run.start_time and run.end_time:
                delta = run.end_time - run.start_time
                console.print(
                    f"    Run #{run.run_id}: {run.status} ({delta.total_seconds():.0f}s)",
                )

        return {
            "total": self.num_runs,
            "success": success_count,
            "failed": failed_count,
            "stopped": stopped_count,
        }


def main(
    runs: int = 5,
    api_keys: int = 2,
    model: str = "openrouter",
) -> dict[str, int]:
    """
    Run multiple factor experiments in parallel.

    Parameters
    ----------
    runs : int
        Number of concurrent runs to spawn
    api_keys : int
        Number of API keys to distribute across
    model : str
        LLM backend: 'local' (llama.cpp) or 'openrouter' (cloud)

    Returns
    -------
    dict
        Summary with keys: total, success, failed, stopped
    """
    runner = ParallelRunner(num_runs=runs, num_api_keys=api_keys, model=model)
    return runner.run()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="NexQuant Parallel Runner - Run multiple factor experiments concurrently",
    )
    parser.add_argument(
        "--runs", "-n",
        type=int,
        default=5,
        help="Number of concurrent runs (default: 5, max recommended: 25)",
    )
    parser.add_argument(
        "--api-keys", "-k",
        type=int,
        default=2,
        help="Number of API keys to distribute across (default: 2)",
    )
    parser.add_argument(
        "--model", "-m",
        type=str,
        default="openrouter",
        choices=["local", "openrouter"],
        help="LLM backend: 'local' (llama.cpp) or 'openrouter' (cloud)",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Skip resource warnings (allow >25 runs)",
    )

    args = parser.parse_args()

    # Resource warnings for high run counts
    if args.runs > 50 and not args.force:
        console.print(f"\n[bold red]⚠️  {args.runs} runs exceeds safe limit (50)[/bold red]")
        console.print("[yellow]This will likely cause memory exhaustion and API throttling.[/yellow]")
        console.print("[yellow]Use --force to override.[/yellow]")
        sys.exit(1)
    elif args.runs > 25:
        console.print(f"\n[yellow]⚠️  {args.runs} runs - high resource usage expected[/yellow]")
        console.print(f"   Estimated RAM: ~{args.runs * 0.65:.0f} GB")
        console.print("   Use --force to confirm.\n")
        import time
        time.sleep(2)

    result = main(runs=args.runs, api_keys=args.api_keys, model=args.model)

    # Exit with appropriate code
    if result["failed"] > 0:
        sys.exit(1)
    sys.exit(0)