ktransformers/kt-kernel/python/cli/commands/bench.py

"""
Bench commands for kt-cli.

Runs benchmarks for performance testing.
"""

import subprocess
import sys
from enum import Enum
from pathlib import Path
from typing import Optional

import typer

from kt_kernel.cli.i18n import t
from kt_kernel.cli.utils.console import (
    console,
    print_error,
    print_info,
    print_step,
    print_success,
)


class BenchType(str, Enum):
    """Benchmark type."""

    INFERENCE = "inference"
    MLA = "mla"
    MOE = "moe"
    LINEAR = "linear"
    ATTENTION = "attention"
    ALL = "all"


def bench(
    type: BenchType = typer.Option(
        BenchType.ALL,
        "--type",
        "-t",
        help="Benchmark type",
    ),
    model: Optional[str] = typer.Option(
        None,
        "--model",
        "-m",
        help="Model to benchmark",
    ),
    output: Optional[Path] = typer.Option(
        None,
        "--output",
        "-o",
        help="Output file for results (JSON)",
    ),
    iterations: int = typer.Option(
        10,
        "--iterations",
        "-n",
        help="Number of iterations",
    ),
) -> None:
    """Run full benchmark suite."""
    console.print()
    print_step(t("bench_starting"))
    print_info(t("bench_type", type=type.value))
    console.print()

    if type == BenchType.ALL:
        _run_all_benchmarks(model, output, iterations)
    elif type == BenchType.INFERENCE:
        _run_inference_benchmark(model, output, iterations)
    elif type == BenchType.MLA:
        _run_component_benchmark("mla", output, iterations)
    elif type == BenchType.MOE:
        _run_component_benchmark("moe", output, iterations)
    elif type == BenchType.LINEAR:
        _run_component_benchmark("linear", output, iterations)
    elif type == BenchType.ATTENTION:
        _run_component_benchmark("attention", output, iterations)

    console.print()
    print_success(t("bench_complete"))
    if output:
        console.print(f"  Results saved to: {output}")
    console.print()


def microbench(
    component: str = typer.Argument(
        "moe",
        help="Component to benchmark (moe, mla, linear, attention)",
    ),
    batch_size: int = typer.Option(
        1,
        "--batch-size",
        "-b",
        help="Batch size",
    ),
    seq_len: int = typer.Option(
        1,
        "--seq-len",
        "-s",
        help="Sequence length",
    ),
    iterations: int = typer.Option(
        100,
        "--iterations",
        "-n",
        help="Number of iterations",
    ),
    warmup: int = typer.Option(
        10,
        "--warmup",
        "-w",
        help="Warmup iterations",
    ),
    output: Optional[Path] = typer.Option(
        None,
        "--output",
        "-o",
        help="Output file for results (JSON)",
    ),
) -> None:
    """Run micro-benchmark for specific components."""
    console.print()
    console.print(f"[yellow]{t('feature_coming_soon')}[/yellow]")
    console.print()
    raise typer.Exit(0)

    # Try to find the benchmark script
    kt_kernel_path = _find_kt_kernel_path()

    if kt_kernel_path is None:
        print_error("kt-kernel not found. Install with: kt install inference")
        raise typer.Exit(1)

    bench_dir = kt_kernel_path / "bench"

    # Map component to script
    component_scripts = {
        "moe": "bench_moe.py",
        "mla": "bench_mla.py",
        "linear": "bench_linear.py",
        "attention": "bench_attention.py",
        "mlp": "bench_mlp.py",
    }

    script_name = component_scripts.get(component.lower())
    if script_name is None:
        print_error(f"Unknown component: {component}")
        console.print(f"Available: {', '.join(component_scripts.keys())}")
        raise typer.Exit(1)

    script_path = bench_dir / script_name
    if not script_path.exists():
        print_error(f"Benchmark script not found: {script_path}")
        raise typer.Exit(1)

    # Run benchmark
    cmd = [
        sys.executable,
        str(script_path),
        "--batch-size",
        str(batch_size),
        "--seq-len",
        str(seq_len),
        "--iterations",
        str(iterations),
        "--warmup",
        str(warmup),
    ]

    if output:
        cmd.extend(["--output", str(output)])

    console.print(f"[dim]$ {' '.join(cmd)}[/dim]")
    console.print()

    try:
        process = subprocess.run(cmd)

        if process.returncode == 0:
            console.print()
            print_success(t("bench_complete"))
            if output:
                console.print(f"  Results saved to: {output}")
        else:
            print_error(f"Benchmark failed with exit code {process.returncode}")
            raise typer.Exit(process.returncode)

    except FileNotFoundError as e:
        print_error(f"Failed to run benchmark: {e}")
        raise typer.Exit(1)


def _find_kt_kernel_path() -> Optional[Path]:
    """Find the kt-kernel installation path."""
    try:
        import kt_kernel

        return Path(kt_kernel.__file__).parent.parent
    except ImportError:
        pass

    # Check common locations
    possible_paths = [
        Path.home() / "Projects" / "ktransformers" / "kt-kernel",
        Path("/opt/ktransformers/kt-kernel"),
        Path.cwd() / "kt-kernel",
    ]

    for path in possible_paths:
        if path.exists() and (path / "bench").exists():
            return path

    return None


def _run_all_benchmarks(model: Optional[str], output: Optional[Path], iterations: int) -> None:
    """Run all benchmarks."""
    components = ["moe", "mla", "linear", "attention"]

    for component in components:
        console.print(f"\n[bold]Running {component} benchmark...[/bold]")
        _run_component_benchmark(component, None, iterations)


def _run_inference_benchmark(model: Optional[str], output: Optional[Path], iterations: int) -> None:
    """Run inference benchmark."""
    if model is None:
        print_error("Model required for inference benchmark. Use --model flag.")
        raise typer.Exit(1)

    print_info(f"Running inference benchmark on {model}...")
    console.print()
    console.print("[dim]This will start the server and run test requests.[/dim]")
    console.print()

    # TODO: Implement actual inference benchmarking
    print_error("Inference benchmarking not yet implemented.")


def _run_component_benchmark(component: str, output: Optional[Path], iterations: int) -> None:
    """Run a component benchmark."""
    kt_kernel_path = _find_kt_kernel_path()

    if kt_kernel_path is None:
        print_error("kt-kernel not found.")
        return

    bench_dir = kt_kernel_path / "bench"
    script_map = {
        "moe": "bench_moe.py",
        "mla": "bench_mla.py",
        "linear": "bench_linear.py",
        "attention": "bench_attention.py",
    }

    script_name = script_map.get(component)
    if script_name is None:
        print_error(f"Unknown component: {component}")
        return

    script_path = bench_dir / script_name
    if not script_path.exists():
        print_error(f"Script not found: {script_path}")
        return

    cmd = [sys.executable, str(script_path), "--iterations", str(iterations)]

    try:
        subprocess.run(cmd)
    except Exception as e:
        print_error(f"Benchmark failed: {e}")