loopkit 0.0.1a1__py3-none-any.whl

loopkit/cli/sweep.py

@@ -0,0 +1,313 @@
+import argparse
+import copy
+import itertools
+import json
+import re
+import subprocess
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import yaml
+
+
+def generate_sweep_configs(
+    base_config: Dict[str, Any],
+    sweep_params: Dict[str, List[Any]],
+) -> List[Dict[str, Any]]:
+    """Generate all configurations for a grid search.
+
+    Args:
+        base_config: Base configuration dictionary
+        sweep_params: Dictionary mapping parameter paths to lists of values.
+                     Example: {"training.lr": [0.001, 0.01],
+                               "model.hidden_dim": [128, 256]}
+
+    Returns:
+        List of configuration dictionaries
+    """
+    # Get all parameter names and their values
+    param_names = list(sweep_params.keys())
+    param_values = [sweep_params[name] for name in param_names]
+
+    # Generate all combinations
+    configs = []
+    for values in itertools.product(*param_values):
+        # Deep copy the base config to avoid mutations
+        config = copy.deepcopy(base_config)
+
+        # Apply sweep parameters
+        for param_name, value in zip(param_names, values):
+            # Convert dot notation to nested dict
+            keys = param_name.split(".")
+            current = config
+            for key in keys[:-1]:
+                if key not in current:
+                    current[key] = {}
+                current = current[key]
+            current[keys[-1]] = value
+
+        configs.append(config)
+
+    return configs
+
+
+def run_sweep(
+    command_template: str,
+    sweep_params: Dict[str, List[Any]],
+    base_config_file: Optional[str] = None,
+    runs_dir: str = "runs",
+    sweep_name: Optional[str] = None,
+    dry_run: bool = False,
+    parallel: bool = False,
+    max_parallel: int = 4,
+):
+    """Run a hyperparameter sweep.
+
+    Args:
+        command_template: Command to run, e.g.:
+                         "python train.py {config}"
+                         "torchrun --nproc_per_node=2 train.py {config}"
+                         "bash train.sh 2 {config}"
+        sweep_params: Parameters to sweep over
+        base_config_file: Base config file (optional)
+        runs_dir: Directory to save runs
+        sweep_name: Name for this sweep
+        dry_run: If True, only print commands without running
+        parallel: Run experiments in parallel
+        max_parallel: Maximum number of parallel runs
+
+    Example:
+        >>> run_sweep(
+        ...     command_template="python train.py {config}",
+        ...     sweep_params={
+        ...         "training.lr": [0.001, 0.01, 0.1],
+        ...         "model.hidden_dim": [128, 256],
+        ...     },
+        ...     base_config_file="conf/default.yaml",
+        ...     sweep_name="lr_hidden_sweep"
+        ... )
+    """
+    # Create sweep directory
+    if sweep_name is None:
+        sweep_name = datetime.now().strftime("sweep_%Y%m%d_%H%M%S")
+
+    sweep_dir = Path(runs_dir) / sweep_name
+    sweep_dir.mkdir(parents=True, exist_ok=True)
+
+    # Load base config if provided
+    base_config = {}
+    if base_config_file:
+        with open(base_config_file) as f:
+            base_config = yaml.safe_load(f)
+
+    # Generate all configurations
+    configs = generate_sweep_configs(base_config, sweep_params)
+
+    print(f"\n{'=' * 80}")
+    print(f"Hyperparameter Sweep: {sweep_name}")
+    print(f"{'=' * 80}")
+    print(f"Total experiments: {len(configs)}")
+    print(f"Sweep directory: {sweep_dir}")
+    print("\nSweep parameters:")
+    for param, values in sweep_params.items():
+        print(f"  {param}: {values}")
+    print(f"{'=' * 80}\n")
+
+    # Save sweep configuration
+    sweep_info = {
+        "sweep_name": sweep_name,
+        "command_template": command_template,
+        "sweep_params": sweep_params,
+        "base_config_file": base_config_file,
+        "num_experiments": len(configs),
+        "timestamp": datetime.now().isoformat(),
+    }
+
+    with open(sweep_dir / "sweep_info.json", "w") as f:
+        json.dump(sweep_info, f, indent=2)
+
+    # Run each configuration
+    results = []
+
+    for idx, config in enumerate(configs):
+        exp_name = f"exp_{idx:03d}"
+        exp_dir = sweep_dir / exp_name
+        exp_dir.mkdir(exist_ok=True)
+
+        # Save config
+        config_file = exp_dir / "config.yaml"
+        with open(config_file, "w") as f:
+            yaml.dump(config, f)
+
+        # Build command
+        command = command_template.format(config=str(config_file))
+
+        # Add run_dir override if not in command
+        if "run_dir" not in command:
+            command += f" run_dir={exp_dir}"
+
+        print(f"\n[{idx + 1}/{len(configs)}] Running: {exp_name}")
+        print(f"Command: {command}")
+
+        # Print sweep parameters for this run
+        print("Parameters:")
+        for param in sweep_params:
+            keys = param.split(".")
+            current = config
+            for key in keys:
+                current = current.get(key, {})
+            print(f"  {param}: {current}")
+
+        if dry_run:
+            results.append({"exp_name": exp_name, "status": "dry_run"})
+            continue
+
+        # Run command
+        try:
+            result = subprocess.run(
+                command,
+                shell=True,
+                capture_output=True,
+                text=True,
+                cwd=Path.cwd(),
+            )
+
+            success = result.returncode == 0
+            results.append(
+                {
+                    "exp_name": exp_name,
+                    "config": config,
+                    "command": command,
+                    "returncode": result.returncode,
+                    "status": "success" if success else "failed",
+                }
+            )
+
+            if success:
+                print("✅ Completed successfully")
+            else:
+                print(f"❌ Failed with return code {result.returncode}")
+                if result.stderr:
+                    print(f"Error: {result.stderr[:500]}")
+
+        except Exception as e:
+            print(f"❌ Exception: {e}")
+            results.append(
+                {
+                    "exp_name": exp_name,
+                    "config": config,
+                    "command": command,
+                    "status": "exception",
+                    "error": str(e),
+                }
+            )
+
+    # Save results
+    results_file = sweep_dir / "sweep_results.json"
+    with open(results_file, "w") as f:
+        json.dump(results, f, indent=2)
+
+    # Print summary
+    print(f"\n{'=' * 80}")
+    print("Sweep Summary")
+    print(f"{'=' * 80}")
+
+    successful = sum(1 for r in results if r["status"] == "success")
+    failed = sum(1 for r in results if r["status"] in ["failed", "exception"])
+
+    print(f"Total experiments: {len(results)}")
+    print(f"Successful: {successful}")
+    print(f"Failed: {failed}")
+    print(f"\nResults saved to: {sweep_dir}")
+    print(f"{'=' * 80}\n")
+
+    return results
+
+
+def setup_sweep_parser(subparsers):
+    """Setup the sweep subcommand parser."""
+    sweep_parser = subparsers.add_parser(
+        "sweep",
+        help="Run hyperparameter sweeps",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Basic sweep
+  loopkit sweep "python train.py {config}" \\
+    --config config.yaml \\
+    --sweep "training.lr=[0.001,0.01]" "model.hidden_dim=[128,256]"
+
+  # With torchrun
+  loopkit sweep "torchrun --nproc_per_node=2 train.py {config}" \\
+    --config config.yaml \\
+    --sweep "training.lr=[0.001,0.01]"
+
+  # Dry run to see commands
+  loopkit sweep "python train.py {config}" \\
+    --config config.yaml \\
+    --sweep "training.lr=[0.001,0.01]" \\
+    --dry-run
+        """,
+    )
+    sweep_parser.add_argument(
+        "command", help="Command template with {config} placeholder"
+    )
+    sweep_parser.add_argument(
+        "--config", "-c", required=True, help="Base configuration file"
+    )
+    sweep_parser.add_argument(
+        "--sweep",
+        "-s",
+        nargs="+",
+        required=True,
+        help="Sweep parameters in format: param=[val1,val2,...]",
+    )
+    sweep_parser.add_argument("--name", "-n", help="Sweep name")
+    sweep_parser.add_argument(
+        "--runs-dir", "-d", default="runs", help="Directory for runs"
+    )
+    sweep_parser.add_argument(
+        "--dry-run", action="store_true", help="Print commands without running"
+    )
+
+
+def run_sweep_command(args):
+    """Run the sweep command."""
+    # Parse sweep parameters
+    sweep_params = {}
+
+    for spec in args.sweep:
+        match = re.match(r"(\S+)=\[([^\]]+)\]", spec)
+        if not match:
+            print(f"❌ Invalid sweep spec: {spec}")
+            print("   Expected format: param=[val1,val2,...]")
+            continue
+
+        key, values_str = match.groups()
+        values = []
+        for v in values_str.split(","):
+            v = v.strip()
+            try:
+                values.append(int(v))
+            except ValueError:
+                try:
+                    values.append(float(v))
+                except ValueError:
+                    values.append(v)
+
+        sweep_params[key] = values
+
+    if not sweep_params:
+        print("❌ No valid sweep parameters provided")
+        sys.exit(1)
+
+    run_sweep(
+        command_template=args.command,
+        sweep_params=sweep_params,
+        base_config_file=args.config,
+        sweep_name=args.name,
+        runs_dir=args.runs_dir,
+        dry_run=args.dry_run,
+    )

loopkit/cli/visualize.py

@@ -0,0 +1,285 @@
+import argparse
+import json
+from pathlib import Path
+from typing import List, Optional, Union
+
+import matplotlib.pyplot as plt
+import polars as pl
+
+
+def plot_metrics(
+    run_dirs: Union[str, Path, List[Union[str, Path]]],
+    metrics: Optional[List[str]] = None,
+    output_file: Optional[str] = None,
+    show: bool = True,
+    style: str = "default",
+):
+    """Plot metrics from one or more runs.
+
+    Args:
+        run_dirs: Single run directory or list of run directories
+        metrics: List of metrics to plot (None = all)
+        output_file: Save plot to file (e.g., 'plot.png')
+        show: Whether to display the plot
+        style: Plot style ('default', 'seaborn', 'ggplot')
+    """
+    # Handle single run_dir
+    if isinstance(run_dirs, (str, Path)):
+        run_dirs = [run_dirs]
+
+    # Set style
+    if style != "default":
+        plt.style.use(style)
+
+    # Load data from all runs
+    all_data = []
+    for run_dir in run_dirs:
+        run_dir = Path(run_dir)
+        metrics_file = run_dir / "metrics.csv"
+
+        if not metrics_file.exists():
+            print(f"⚠️  No metrics found in {run_dir}")
+            continue
+
+        df = pl.read_csv(metrics_file)
+        df = df.with_columns(run_id=pl.lit(run_dir.name))
+        all_data.append(df)
+
+    if not all_data:
+        print("❌ No metrics data found")
+        return
+
+    combined_df = pl.concat(all_data)
+
+    # Filter metrics if specified
+    if metrics:
+        combined_df = combined_df.filter(pl.col("name").is_in(metrics))
+
+    # Get unique metrics and splits
+    unique_metrics = combined_df["name"].unique()
+    unique_splits = combined_df["split"].unique()
+
+    # Create subplots
+    n_metrics = len(unique_metrics)
+    n_cols = min(2, n_metrics)
+    n_rows = (n_metrics + n_cols - 1) // n_cols
+
+    fig, axes = plt.subplots(n_rows, n_cols, figsize=(6 * n_cols, 4 * n_rows))
+    if n_metrics == 1:
+        axes = [axes]
+    elif n_rows > 1 and n_cols > 1:
+        axes = axes.flatten()
+    else:
+        axes = list(axes) if n_rows > 1 or n_cols > 1 else [axes]
+
+    # Plot each metric
+    for idx, metric_name in enumerate(sorted(unique_metrics)):
+        ax = axes[idx]
+
+        for run_id in combined_df["run_id"].unique():
+            for split in unique_splits:
+                data = combined_df.filter(
+                    (pl.col("name") == metric_name)
+                    & (pl.col("run_id") == run_id)
+                    & (pl.col("split") == split)
+                )
+
+                if data.height > 0:
+                    label = f"{run_id}_{split}" if len(run_dirs) > 1 else split
+                    ax.plot(
+                        data["step"],
+                        data["value"],
+                        label=label,
+                        marker="o",
+                        markersize=3,
+                    )
+
+        ax.set_xlabel("Step")
+        ax.set_ylabel("Value")
+        ax.set_title(metric_name)
+        ax.legend()
+        ax.grid(True, alpha=0.3)
+
+    # Hide unused subplots
+    for idx in range(n_metrics, len(axes)):
+        axes[idx].axis("off")
+
+    plt.tight_layout()
+
+    if output_file:
+        plt.savefig(output_file, dpi=150, bbox_inches="tight")
+        print(f"📊 Plot saved to {output_file}")
+
+    if show:
+        plt.show()
+    else:
+        plt.close()
+
+
+def plot_resource_usage(
+    run_dir: Union[str, Path],
+    output_file: Optional[str] = None,
+    show: bool = True,
+):
+    """Plot resource usage over time.
+
+    Args:
+        run_dir: Run directory
+        output_file: Save plot to file
+        show: Whether to display the plot
+    """
+    run_dir = Path(run_dir)
+    resource_file = run_dir / "resource_usage.jsonl"
+
+    if not resource_file.exists():
+        print(f"❌ No resource monitoring data found in {run_dir}")
+        return
+
+    # Load data
+    data = []
+    with open(resource_file) as f:
+        for line in f:
+            data.append(json.loads(line))
+
+    if not data:
+        print("❌ No resource data found")
+        return
+
+    # Extract time series
+    timestamps = [d["timestamp"] for d in data]
+    start_time = timestamps[0]
+    elapsed_times = [(t - start_time) / 60 for t in timestamps]  # Convert to minutes
+
+    # Prepare subplots
+    n_plots = 0
+    has_gpu = "gpu" in data[0]
+    has_cpu = "cpu" in data[0]
+    has_memory = "memory" in data[0]
+
+    n_plots = sum([has_gpu, has_cpu, has_memory])
+
+    if n_plots == 0:
+        print("❌ No resource data to plot")
+        return
+
+    fig, axes = plt.subplots(n_plots, 1, figsize=(10, 4 * n_plots))
+    if n_plots == 1:
+        axes = [axes]
+
+    plot_idx = 0
+
+    # Plot GPU memory
+    if has_gpu:
+        ax = axes[plot_idx]
+        plot_idx += 1
+
+        for i in range(len(data[0]["gpu"])):
+            gpu_mem = [d["gpu"][i]["allocated_mb"] for d in data if "gpu" in d]
+            ax.plot(elapsed_times, gpu_mem, label=f"GPU {i}", marker="o", markersize=2)
+
+        ax.set_xlabel("Time (minutes)")
+        ax.set_ylabel("Memory (MB)")
+        ax.set_title("GPU Memory Usage")
+        ax.legend()
+        ax.grid(True, alpha=0.3)
+
+    # Plot CPU usage
+    if has_cpu:
+        ax = axes[plot_idx]
+        plot_idx += 1
+
+        cpu_percent = [d["cpu"]["percent"] for d in data if "cpu" in d]
+        ax.plot(
+            elapsed_times[: len(cpu_percent)], cpu_percent, marker="o", markersize=2
+        )
+
+        ax.set_xlabel("Time (minutes)")
+        ax.set_ylabel("CPU %")
+        ax.set_title("CPU Usage")
+        ax.grid(True, alpha=0.3)
+
+    # Plot memory usage
+    if has_memory:
+        ax = axes[plot_idx]
+        plot_idx += 1
+
+        mem_percent = [d["memory"]["percent"] for d in data if "memory" in d]
+        ax.plot(
+            elapsed_times[: len(mem_percent)], mem_percent, marker="o", markersize=2
+        )
+
+        ax.set_xlabel("Time (minutes)")
+        ax.set_ylabel("Memory %")
+        ax.set_title("System Memory Usage")
+        ax.grid(True, alpha=0.3)
+
+    plt.tight_layout()
+
+    if output_file:
+        plt.savefig(output_file, dpi=150, bbox_inches="tight")
+        print(f"📊 Resource plot saved to {output_file}")
+
+    if show:
+        plt.show()
+    else:
+        plt.close()
+
+
+def setup_visualize_parser(subparsers):
+    """Setup the visualize subcommand parser."""
+    viz_parser = subparsers.add_parser(
+        "viz",
+        help="Visualize experiment metrics",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Plot all metrics from multiple runs
+  loopkit viz runs/exp1 runs/exp2 runs/exp3
+
+  # Plot specific metrics
+  loopkit viz runs/exp_* --metrics loss accuracy
+
+  # Save plot to file
+  loopkit viz runs/exp1 --output plot.png
+
+  # Plot resource usage
+  loopkit viz runs/exp1 --resources
+        """,
+    )
+    viz_parser.add_argument("run_dirs", nargs="+", help="Run directories to plot")
+    viz_parser.add_argument(
+        "--metrics", "-m", nargs="+", help="Specific metrics to plot"
+    )
+    viz_parser.add_argument("--output", "-o", help="Save plot to file (e.g., plot.png)")
+    viz_parser.add_argument("--no-show", action="store_true", help="Don't display plot")
+    viz_parser.add_argument(
+        "--style",
+        default="default",
+        choices=["default", "seaborn", "ggplot"],
+        help="Plot style",
+    )
+    viz_parser.add_argument(
+        "--resources",
+        "-r",
+        action="store_true",
+        help="Plot resource usage instead of metrics",
+    )
+
+
+def run_visualize_command(args):
+    """Run the visualize command."""
+    if args.resources:
+        if len(args.run_dirs) > 1:
+            print("⚠️  Resource plotting only supports single run directory")
+            print(f"Plotting resources for: {args.run_dirs[0]}")
+        plot_resource_usage(
+            args.run_dirs[0], output_file=args.output, show=not args.no_show
+        )
+    else:
+        plot_metrics(
+            args.run_dirs,
+            metrics=args.metrics,
+            output_file=args.output,
+            show=not args.no_show,
+            style=args.style,
+        )