hud-python 0.4.28__py3-none-any.whl → 0.4.29__py3-none-any.whl

hud/cli/rl/config.py

@@ -0,0 +1,94 @@
+"""Configuration generation and management for RL training."""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any
+
+from rich.console import Console
+
+from hud.rl.config import Config, validate_vl_model
+from hud.utils.hud_console import hud_console
+
+from .display import display_preset_table
+from .presets import estimate_memory_usage
+
+if TYPE_CHECKING:
+    from pathlib import Path
+console = Console()
+
+
+def generate_config_interactive(
+    model_name: str,
+    presets: list[dict[str, Any]],
+) -> tuple[Config, float]:
+    """Generate RL training configuration interactively."""
+    # Validate model is a VL model
+    validate_vl_model(model_name)
+
+    # Display preset options
+    display_preset_table(presets, 80.0)  # Assuming A100 80GB
+
+    # Let user select preset
+    preset_choice = hud_console.select(
+        "Select a training configuration preset:",
+        choices=[{"name": p["name"], "value": i} for i, p in enumerate(presets)],
+        default=1 if len(presets) > 1 else 0,  # Default to "Balanced" if available
+    )
+
+    selected_preset = presets[preset_choice]  # type: ignore
+
+    # Use preset values directly
+    max_steps_per_episode = selected_preset["max_steps_per_episode"]
+
+    # Calculate memory estimate
+    max_pixels = 256 * 28 * 28
+    estimated_memory = estimate_memory_usage(
+        selected_preset["mini_batch_size"],
+        max_steps_per_episode,
+        selected_preset["max_new_tokens"],
+        max_pixels,
+    )
+
+    config_adds = {
+        "actor": {
+            "max_new_tokens": selected_preset["max_new_tokens"],
+            "max_parallel_episodes": selected_preset["batch_size"],
+            "max_steps_per_episode": selected_preset["max_steps_per_episode"],
+            "force_tool_choice": True,
+        },
+        "training": {
+            "mini_batch_size": selected_preset["mini_batch_size"],
+            "group_size": selected_preset["group_size"],
+            "batch_size": selected_preset["batch_size"],
+            "lr": selected_preset["lr"],
+            "epochs": selected_preset["epochs"],
+        },
+        "verbose": True,
+    }
+
+    # Create config
+    config = Config.from_dict(config_adds)
+
+    return config, estimated_memory
+
+
+def save_config(config: Config, path: Path) -> None:
+    """Save configuration to a JSON file."""
+    config_dict = config.to_dict()
+
+    with open(path, "w") as f:
+        json.dump(config_dict, f, indent=2)
+        f.write("\n")  # Add newline at end of file
+
+    if not path.name.startswith("."):  # Don't show message for temp files
+        console.print(f"[green]✅ Configuration saved to {path}[/green]")
+
+
+def load_config(path: Path) -> Config:
+    """Load configuration from a JSON file."""
+    with open(path) as f:
+        data = json.load(f)
+
+    # Use Config.from_dict which handles missing fields gracefully
+    return Config.from_dict(data)

hud/cli/rl/display.py

@@ -0,0 +1,133 @@
+"""Display utilities for RL training configuration."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from rich.console import Console
+from rich.table import Table
+
+if TYPE_CHECKING:
+    from hud.rl.config import Config
+
+console = Console()
+
+
+def display_gpu_info(gpu_info: dict[str, Any]) -> None:
+    """Display GPU information in a table."""
+    if not gpu_info["available"]:
+        console.print(f"[red]❌ CUDA not available: {gpu_info.get('error', 'Unknown error')}[/red]")
+        return
+
+    gpu_table = Table(title="🖥️  Available GPUs", title_style="bold cyan")
+    gpu_table.add_column("Index", style="yellow")
+    gpu_table.add_column("Name", style="cyan")
+    gpu_table.add_column("Memory", style="green")
+
+    for device in gpu_info["devices"]:
+        gpu_table.add_row(f"GPU {device['index']}", device["name"], f"{device['memory_gb']:.1f} GB")
+
+    console.print(gpu_table)
+
+
+def display_preset_table(presets: list[dict[str, Any]], gpu_memory_gb: float) -> None:
+    """Display training configuration presets in a table."""
+    preset_table = Table(title="📊 Training Configuration Presets", title_style="bold cyan")
+    preset_table.add_column("Option", style="yellow")
+    preset_table.add_column("Steps", style="cyan")
+    preset_table.add_column("Mini-batch", style="cyan")
+    preset_table.add_column("Group", style="cyan")
+    preset_table.add_column("Episodes/batch", style="cyan")
+
+    # Add time columns for A100
+    if gpu_memory_gb >= 40:
+        preset_table.add_column("Tasks/hour", style="green")
+        preset_table.add_column("Steps/hour", style="green")
+
+    for i, preset in enumerate(presets):
+        row = [
+            f"{i + 1}. {preset['name']}",
+            str(preset["max_steps_per_episode"]),
+            str(preset["mini_batch_size"]),
+            str(preset["group_size"]),
+            str(preset["batch_size"]),
+        ]
+        if "tasks_per_hour" in preset:
+            row.extend(
+                [
+                    str(preset["tasks_per_hour"]),
+                    str(preset["steps_per_hour"]),
+                ]
+            )
+        preset_table.add_row(*row)
+
+    console.print("\n")
+    console.print(preset_table)
+    console.print("\n")
+
+
+def display_config_summary(
+    config: Config, tasks_count: int, gpu_info: dict[str, Any], estimated_memory: float
+) -> None:
+    """Display comprehensive configuration summary for review."""
+    console.print("\n[bold cyan]📋 RL Training Configuration Summary[/bold cyan]\n")
+
+    # GPU Information
+    if gpu_info["available"]:
+        gpu_table = Table(title="🖥️  GPU Information", title_style="bold yellow")
+        gpu_table.add_column("Property", style="cyan")
+        gpu_table.add_column("Value", style="green")
+
+        device = gpu_info["devices"][0]  # Primary GPU
+        gpu_table.add_row("GPU 0", device["name"])
+        gpu_table.add_row("Memory", f"{device['memory_gb']:.1f} GB")
+        gpu_table.add_row("Compute Capability", "8.0")  # Assuming A100
+
+        console.print(gpu_table)
+
+    # Model Configuration
+    model_table = Table(title="🤖 Model Configuration", title_style="bold yellow")
+    model_table.add_column("Parameter", style="cyan")
+    model_table.add_column("Value", style="green")
+
+    model_table.add_row("Base Model", config.model.base_model)
+    model_table.add_row("LoRA Rank (r)", str(config.model.lora_r))
+    model_table.add_row("LoRA Alpha", str(config.model.lora_alpha))
+    model_table.add_row("LoRA Dropout", str(config.model.lora_dropout))
+
+    console.print(model_table)
+
+    # Training Configuration
+    training_table = Table(title="🎯 Training Configuration", title_style="bold yellow")
+    training_table.add_column("Parameter", style="cyan")
+    training_table.add_column("Value", style="green")
+
+    training_table.add_row("Tasks Count", str(tasks_count))
+    training_table.add_row("Learning Rate", f"{config.training.lr:.1e}")
+    training_table.add_row("Epochs", str(config.training.epochs))
+    training_table.add_row("Mini Batch Size", str(config.training.mini_batch_size))
+    training_table.add_row("Batch Size", str(config.training.batch_size))
+    training_table.add_row("Group Size", str(config.training.group_size))
+    training_table.add_row("Training Steps", str(config.training.training_steps))
+    training_table.add_row("Max Parallel Episodes", str(config.actor.max_parallel_episodes))
+
+    console.print(training_table)
+
+    # Memory Estimation
+    memory_table = Table(title="💾 Memory Estimation", title_style="bold yellow")
+    memory_table.add_column("Metric", style="cyan")
+    memory_table.add_column("Value", style="green")
+
+    memory_table.add_row("Estimated GPU Memory", f"{estimated_memory:.1f} GB")
+    if gpu_info["available"]:
+        available_memory = gpu_info["devices"][0]["memory_gb"]
+        memory_table.add_row("Available GPU Memory", f"{available_memory:.1f} GB")
+
+        if estimated_memory > available_memory:
+            status = "[red]⚠️  May exceed available memory[/red]"
+        else:
+            status = "[green]✅ Within memory limits[/green]"
+        memory_table.add_row("Status", status)
+
+    console.print(memory_table)
+    console.print("\n")

hud/cli/rl/gpu.py

@@ -0,0 +1,63 @@
+"""GPU detection and validation utilities for RL training."""
+
+from __future__ import annotations
+
+import subprocess
+from typing import Any
+
+
+def detect_cuda_devices() -> dict[str, Any]:
+    """Detect available CUDA devices and their properties."""
+    try:
+        # Check if CUDA is available
+        result = subprocess.run(
+            ["nvidia-smi", "--query-gpu=index,name,memory.total", "--format=csv,noheader,nounits"],  # noqa: S607
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+
+        if result.returncode != 0:
+            return {"available": False, "error": "nvidia-smi command failed"}
+
+        devices = []
+        for line in result.stdout.strip().split("\n"):
+            parts = line.split(", ")
+            if len(parts) >= 3:
+                devices.append(
+                    {
+                        "index": int(parts[0]),
+                        "name": parts[1],
+                        "memory_gb": float(parts[2]) / 1024,  # Convert MB to GB
+                    }
+                )
+
+        return {"available": True, "devices": devices}
+
+    except FileNotFoundError:
+        return {
+            "available": False,
+            "error": "nvidia-smi not found - CUDA drivers may not be installed",
+        }
+    except Exception as e:
+        return {"available": False, "error": str(e)}
+
+
+def select_gpu_for_vllm(devices: list[dict[str, Any]]) -> int:
+    """Select the best GPU for vLLM server (typically GPU 1 if available)."""
+    if len(devices) > 1:
+        # Prefer GPU 1 for vLLM to leave GPU 0 for other processes
+        return 1
+    return 0
+
+
+def validate_gpu_memory(gpu_memory_gb: float, model_size: str = "3B") -> bool:
+    """Validate if GPU has sufficient memory for the model."""
+    min_memory_requirements = {
+        "3B": 12.0,  # Minimum for Qwen 2.5 VL 3B
+        "7B": 24.0,
+        "14B": 40.0,
+    }
+
+    min_required = min_memory_requirements.get(model_size, 12.0)
+    return gpu_memory_gb >= min_required

hud/cli/rl/gpu_utils.py

@@ -0,0 +1,318 @@
+"""GPU utilities for DDP training."""
+
+from __future__ import annotations
+
+import logging
+import subprocess
+import time
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from hud.utils.hud_console import HUDConsole
+
+if TYPE_CHECKING:
+    from hud.rl.config import Config
+hud_console = HUDConsole(logging.getLogger(__name__))
+
+
+def get_gpu_memory_info() -> dict[int, dict[str, Any]]:
+    """Get memory usage information for all GPUs."""
+
+    gpu_memory = {}
+    try:
+        # Get memory info for all GPUs
+        cmd = [
+            "nvidia-smi",
+            "--query-gpu=index,memory.used,memory.total,memory.free",
+            "--format=csv,noheader,nounits",
+        ]
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)  # noqa: S603
+
+        for line in result.stdout.strip().split("\n"):
+            if not line:
+                continue
+            parts = line.split(", ")
+            if len(parts) >= 4:
+                gpu_idx = int(parts[0])
+                memory_used = float(parts[1])
+                memory_total = float(parts[2])
+                memory_free = float(parts[3])
+                gpu_memory[gpu_idx] = {
+                    "used_mb": memory_used,
+                    "total_mb": memory_total,
+                    "free_mb": memory_free,
+                    "used_pct": (memory_used / memory_total) * 100,
+                }
+
+        # Get process information per GPU
+        for gpu_idx in gpu_memory:  # noqa: PLC0206
+            cmd = [
+                "nvidia-smi",
+                "-i",
+                str(gpu_idx),
+                "--query-compute-apps=pid,used_memory",
+                "--format=csv,noheader,nounits",
+            ]
+            try:
+                result = subprocess.run(cmd, capture_output=True, text=True, check=True)  # noqa: S603
+                processes = []
+                for line in result.stdout.strip().split("\n"):
+                    if not line:
+                        continue
+                    parts = line.split(", ")
+                    if len(parts) >= 2:
+                        pid = int(parts[0])
+                        memory_mb = float(parts[1])
+                        processes.append({"pid": pid, "memory_mb": memory_mb})
+                gpu_memory[gpu_idx]["processes"] = processes
+            except Exception as e:
+                hud_console.error(f"Failed to get process info for GPU {gpu_idx}: {e}")
+                gpu_memory[gpu_idx]["processes"] = []
+
+    except Exception as e:
+        hud_console.error(f"Failed to get GPU memory info {e}")
+        return {}
+
+    return gpu_memory
+
+
+def health_check_gpus(gpu_indices: list[int]) -> dict[str, Any]:
+    """Perform health check on specified GPUs including memory status.
+
+    Returns:
+        Dict with:
+        - healthy_gpus: List of healthy GPU indices
+        - unhealthy_gpus: Dict of unhealthy GPU index -> error message
+        - all_healthy: Boolean indicating if all GPUs are healthy
+        - memory_issues: Boolean indicating if there are memory issues
+    """
+    from rich.console import Console
+    from rich.table import Table
+
+    console = Console()
+
+    console.print("\n[bold cyan]🏥 GPU Health Check[/bold cyan]")
+
+    # First get memory info
+    memory_info = get_gpu_memory_info()
+
+    healthy_gpus = []
+    unhealthy_gpus = {}
+    memory_issues = []
+
+    # Create a table for results
+    table = Table(title="GPU Health Status")
+    table.add_column("GPU", style="cyan")
+    table.add_column("Memory Usage", style="yellow")
+    table.add_column("Status", style="green")
+    table.add_column("Details", style="yellow")
+
+    for gpu_idx in gpu_indices:
+        # Memory info
+        mem_str = "Unknown"
+        if gpu_idx in memory_info:
+            mem = memory_info[gpu_idx]
+            used_gb = mem["used_mb"] / 1024
+            total_gb = mem["total_mb"] / 1024
+            mem_str = f"{used_gb:.1f}/{total_gb:.1f} GB ({mem['used_pct']:.0f}%)"
+
+            # Check for high memory usage
+            if mem["used_pct"] > 70:
+                memory_issues.append(gpu_idx)
+                proc_info = f" ({len(mem['processes'])} processes)" if mem["processes"] else ""
+                unhealthy_gpus[gpu_idx] = f"High memory usage{proc_info}"
+                table.add_row(
+                    f"GPU {gpu_idx}", mem_str, "❌ Unhealthy", f"High memory usage{proc_info}"
+                )
+                continue
+
+        # If no severe memory issue, do accessibility test
+        try:
+            # Try to allocate a small tensor on the GPU
+            torch.cuda.set_device(gpu_idx)
+            device = torch.device(f"cuda:{gpu_idx}")
+
+            # Test basic allocation
+            test_tensor = torch.zeros(100, 100, device=device)
+
+            # Test computation
+            result = torch.matmul(test_tensor, test_tensor)
+
+            # Force synchronization
+            torch.cuda.synchronize(device)
+
+            # Clean up
+            del test_tensor, result
+            torch.cuda.empty_cache()
+
+            healthy_gpus.append(gpu_idx)
+            table.add_row(f"GPU {gpu_idx}", mem_str, "✅ Healthy", "Passed all tests")
+
+        except Exception as e:
+            error_msg = str(e)
+            if "busy or unavailable" in error_msg:
+                short_msg = "Device busy or unavailable"
+            elif "out of memory" in error_msg:
+                short_msg = "Insufficient memory"
+            else:
+                short_msg = error_msg[:50] + "..." if len(error_msg) > 50 else error_msg
+
+            unhealthy_gpus[gpu_idx] = short_msg
+            table.add_row(f"GPU {gpu_idx}", mem_str, "❌ Unhealthy", short_msg)
+
+        # Small delay between GPU checks
+        time.sleep(0.1)
+
+    console.print(table)
+
+    return {
+        "healthy_gpus": healthy_gpus,
+        "unhealthy_gpus": unhealthy_gpus,
+        "all_healthy": len(unhealthy_gpus) == 0,
+        "memory_issues": memory_issues,
+    }
+
+
+def calculate_optimal_gpu_allocation(gpu_info: dict[str, Any], config: Config) -> dict[str, Any]:
+    """Calculate optimal GPU allocation for DDP GRPO training.
+
+    Key insight: In GRPO, we want to process groups in parallel.
+    Optimal case: num_gpus = num_groups (each GPU processes 1 group).
+    """
+    devices = gpu_info["devices"]
+    available_gpus = [device["index"] for device in devices]
+
+    # Need at least 2 GPUs (1 for training, 1 for vLLM)
+    if len(available_gpus) < 2:
+        return {"use_ddp": False, "reason": "Need at least 2 GPUs"}
+
+    # Reserve last GPU for vLLM
+    vllm_gpu = available_gpus[-1]
+    training_gpus = available_gpus[:-1]
+
+    # Calculate number of groups
+    batch_size = config.training.batch_size
+    group_size = config.training.group_size
+    num_groups = batch_size // group_size
+
+    if num_groups == 0:
+        num_groups = 1
+
+    # Optimal: Use exactly num_groups GPUs (each processes 1 group in parallel)
+    # But cap at available training GPUs
+    optimal_gpu_count = min(len(training_gpus), num_groups)
+
+    # Only use DDP if we have more than 1 group and more than 1 GPU
+    use_ddp = optimal_gpu_count > 1 and num_groups > 1
+
+    if not use_ddp:
+        # Single GPU training
+        return {
+            "use_ddp": False,
+            "reason": f"Single GPU sufficient for {num_groups} group(s)",
+            "training_gpus": [training_gpus[0]],
+            "vllm_gpu": vllm_gpu,
+            "num_groups": num_groups,
+        }
+
+    # Use optimal number of GPUs for DDP
+    training_gpus = training_gpus[:optimal_gpu_count]
+
+    return {
+        "use_ddp": True,
+        "training_gpus": training_gpus,
+        "vllm_gpu": vllm_gpu,
+        "num_groups": num_groups,
+        "groups_per_gpu": num_groups / len(training_gpus),
+        "parallel_efficiency": min(
+            1.0, num_groups / len(training_gpus)
+        ),  # 1.0 = perfect load balance
+    }
+
+
+def adjust_config_for_ddp(config: Config, num_gpus: int) -> Config:
+    """Adjust configuration for optimal DDP performance.
+
+    Scaling rule:
+    - For 1 GPU: batch_size = 2 * group_size
+    - For N GPUs (N > 1): batch_size = N * group_size
+
+    This ensures each GPU processes exactly 1 group in parallel for optimal performance.
+    """
+    group_size = config.training.group_size
+
+    # Apply scaling rule
+    if num_gpus == 1:
+        # Special case: 2 groups for single GPU
+        config.training.batch_size = 2 * group_size
+    else:
+        # Multi-GPU: each GPU processes 1 group
+        config.training.batch_size = num_gpus * group_size
+
+    # Update max_parallel_episodes to match
+    config.actor.max_parallel_episodes = config.training.batch_size
+
+    # Log the adjustment
+    from rich.console import Console
+
+    console = Console()
+    console.print(
+        f"\n[cyan]📊 Adjusted batch_size to {config.training.batch_size} ({config.training.batch_size // group_size} groups)[/cyan]"  # noqa: E501
+    )
+    console.print(
+        f"[cyan]   Each of the {num_gpus} GPU(s) will process {config.training.batch_size // group_size // num_gpus} group(s) in parallel[/cyan]"  # noqa: E501
+    )
+
+    return config
+
+
+def kill_high_memory_processes(memory_threshold: float = 70.0) -> int:
+    """Kill all GPU processes using more than threshold% memory.
+
+    Returns:
+        Number of processes killed
+    """
+    from rich.console import Console
+
+    console = Console()
+
+    memory_info = get_gpu_memory_info()
+    killed_count = 0
+
+    for gpu_idx, info in memory_info.items():
+        if info["used_pct"] > memory_threshold:
+            for proc in info.get("processes", []):
+                pid = proc["pid"]
+                try:
+                    # Try graceful termination first
+                    subprocess.run(["kill", "-TERM", str(pid)], check=False, capture_output=True)  # noqa: S603, S607
+                    killed_count += 1
+                    console.print(
+                        f"[yellow]Terminating PID {pid} on GPU {gpu_idx} ({proc['memory_mb'] / 1024:.1f} GB)[/yellow]"  # noqa: E501
+                    )
+                except Exception as e:
+                    console.print(f"[red]Failed to kill PID {pid}: {e}[/red]")
+
+    if killed_count > 0:
+        console.print(f"\n[yellow]Sent termination signal to {killed_count} processes...[/yellow]")
+        time.sleep(3)
+
+        # Force kill any remaining
+        for info in memory_info.values():
+            for proc in info.get("processes", []):
+                pid = proc["pid"]
+                try:
+                    # Check if still running
+                    subprocess.run(  # noqa: S603
+                        ["kill", "-0", str(pid)],  # noqa: S607
+                        check=True,
+                        capture_output=True,
+                    )
+                    # If no error, process is still running, force kill
+                    subprocess.run(["kill", "-KILL", str(pid)], check=False)  # noqa: S603, S607
+                    console.print(f"[red]Force killed PID {pid}[/red]")
+                except Exception:
+                    hud_console.error(f"Failed to kill PID {pid}")
+
+    return killed_count