hud-python 0.4.36__py3-none-any.whl → 0.4.37__py3-none-any.whl

hud/cli/flows/tasks.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import logging
 import re
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -8,10 +9,9 @@ from typing import TYPE_CHECKING, Any
 import typer
 import yaml
 
-from hud.cli.build import build_environment
 from hud.cli.push import push_environment
 from hud.cli.utils.docker import require_docker_running
-from hud.cli.utils.environment import is_environment_directory
+from hud.cli.utils.env_check import ensure_built, find_environment_dir
 from hud.cli.utils.registry import extract_name_and_tag
 from hud.utils.hud_console import hud_console
 from hud.utils.tasks import load_tasks
@@ -20,6 +20,9 @@ if TYPE_CHECKING:
     from hud.types import Task
 
 
+logger = logging.getLogger(__name__)
+
+
 def _is_remote_url(url: str) -> bool:
     """Match the remote url."""
     # See if a url is a remote url
@@ -53,62 +56,6 @@ def _validate_tasks(tasks: list[Task]) -> bool:
     return True
 
 
-def _find_environment_dir(tasks_path: Path) -> Path | None:
-    """Find the environment directory related to a tasks file.
-
-    Strategy:
-    - Prefer a directory containing hud.lock.yaml
-    - Fallback to a directory that looks like an environment (Dockerfile + pyproject.toml)
-    - Search the tasks file directory, CWD, and a couple of parents
-    """
-    candidates: list[Path] = []
-    cwd = Path.cwd()
-    candidates.extend([tasks_path.parent, cwd])
-
-    # Add parents (up to 2 levels for each)
-    for base in list(candidates):
-        p = base
-        for _ in range(2):
-            p = p.parent
-            if p not in candidates:
-                candidates.append(p)
-
-    # Prefer those with hud.lock.yaml
-    for d in candidates:
-        if (d / "hud.lock.yaml").exists():
-            return d
-
-    # Otherwise, find a plausible environment dir
-    for d in candidates:
-        try:
-            if is_environment_directory(d):
-                return d
-        except Exception as e:
-            hud_console.debug(f"Skipping path {d}: {e}")
-            continue
-
-    return None
-
-
-def _ensure_built(env_dir: Path) -> dict[str, Any]:
-    """Ensure the environment is built and a lock file exists; return lock data."""
-    lock_path = env_dir / "hud.lock.yaml"
-    if not lock_path.exists():
-        hud_console.warning("No hud.lock.yaml found. The environment hasn't been built.")
-        if not hud_console.confirm("Build the environment now (runs 'hud build')?", default=True):
-            raise typer.Exit(1)
-        # Check Docker availability before attempting a build
-        require_docker_running()
-        # Run build (non-interactive). If Docker isn't running, this will raise and stop the flow.
-        # Force linux/amd64 platform to ensure compatibility during RL flows.
-        build_environment(str(env_dir), platform="linux/amd64")
-
-    # Load lock file
-    with open(lock_path) as f:
-        lock_data = yaml.safe_load(f) or {}
-    return lock_data
-
-
 def _ensure_pushed(env_dir: Path, lock_data: dict[str, Any]) -> dict[str, Any]:
     """Ensure the environment is pushed to a registry; return updated lock data."""
     pushed = bool(lock_data.get("push"))
@@ -153,18 +100,106 @@ def _derive_remote_image(lock_data: dict[str, Any]) -> str:
     return f"{name}:{tag}"
 
 
+def _extract_existing_images(tasks: list[Task]) -> set[str]:
+    """Extract all Mcp-Image references from tasks."""
+    images = set()
+
+    def _extract_from_obj(obj: Any) -> None:
+        if isinstance(obj, dict):
+            # Check for Mcp-Image in headers
+            if "headers" in obj and isinstance(obj["headers"], dict):
+                mcp_image = obj["headers"].get("Mcp-Image")
+                if mcp_image:
+                    images.add(mcp_image)
+            # Recursively check nested objects
+            for v in obj.values():
+                _extract_from_obj(v)
+        elif isinstance(obj, list):
+            for item in obj:
+                _extract_from_obj(item)
+
+    for task in tasks:
+        if task.mcp_config:
+            _extract_from_obj(task.mcp_config)
+
+    return images
+
+
+def _env_var_to_header_key(var_name: str) -> str:
+    """Convert ENV_VAR style to Env-Env-Var header style.
+
+    Example: OPENAI_API_KEY -> Env-Openai-Api-Key
+    """
+    parts = str(var_name).split("_")
+    return f"Env-{'-'.join(part.capitalize() for part in parts)}"
+
+
+def _extract_api_key_vars(lock_data: dict[str, Any]) -> set[str]:
+    """Extract env var names from lock file's provided section (authoritative source).
+
+    We only use keys listed under environment.variables.provided, and exclude HUD_API_KEY
+    because Authorization already carries it.
+    """
+    provided_keys: set[str] = set()
+    if not isinstance(lock_data, dict):
+        return provided_keys
+    try:
+        env_section = (lock_data.get("environment") or {}).get("variables") or {}
+        provided = env_section.get("provided") or {}
+        for name in provided:
+            provided_keys.add(str(name))
+    except Exception as e:
+        logger.debug("Failed to parse provided env vars from lock data: %s", e)
+    provided_keys.discard("HUD_API_KEY")
+    return provided_keys
+
+
+def _extract_dotenv_api_key_vars(env_dir: Path) -> set[str]:
+    """Parse .env for API-like variables to suggest as headers.
+
+    We intentionally include only keys that look like secrets to avoid noise:
+    any key containing one of: api, key, token, secret, password (case-insensitive).
+    """
+    dotenv_path = env_dir / ".env"
+    detected: set[str] = set()
+    if not dotenv_path.exists():
+        return detected
+    try:
+        for line in dotenv_path.read_text(encoding="utf-8").splitlines():
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            if "=" not in line:
+                continue
+            name, _ = line.split("=", 1)
+            name = name.strip()
+            lowered = name.lower()
+            if any(s in lowered for s in ("api", "key", "token", "secret", "password")):
+                detected.add(name)
+    except Exception:
+        # Best-effort only
+        return detected
+    detected.discard("HUD_API_KEY")
+    return detected
+
+
 def convert_tasks_to_remote(tasks_file: str) -> str:
     """Convert a local tasks file to remote MCP tasks and return new filename.
 
     Steps:
     1) Find env dir; ensure built (hud.lock.yaml), otherwise build
     2) Ensure pushed to registry, otherwise push
-    3) Create remote_[tasks].json with mcp_config pointing to mcp.hud.so and Mcp-Image
-    4) Return the new tasks file path
+    3) Check for outdated images in existing task configurations
+    4) Create remote_[tasks].json with mcp_config pointing to mcp.hud.so and Mcp-Image
+    5) Return the new tasks file path
     """
     tasks_path = Path(tasks_file).resolve()
 
-    tasks = load_tasks(str(tasks_path))
+    # Load validated tasks for decision-making (may resolve env vars)
+    tasks: list[Task] = load_tasks(str(tasks_path))  # type: ignore[assignment]
+
+    # Load raw tasks to preserve placeholders when writing back to disk
+    raw_tasks: list[dict[str, Any]] = load_tasks(str(tasks_path), raw=True)  # type: ignore[assignment]
 
     # Ensure HUD_API_KEY is available: prefer process env, else load from env_dir/.env
     from hud.settings import settings
@@ -174,24 +209,119 @@ def convert_tasks_to_remote(tasks_file: str) -> str:
         hud_console.info("Set it in your environment or run: hud set HUD_API_KEY=your-key-here")
         raise typer.Exit(1)
 
+    # Check if tasks already have remote URLs
+    already_remote = _validate_tasks(tasks)
+
+    # Extract existing images from tasks
+    existing_images = _extract_existing_images(tasks)
+
     # Load tasks (supports .json and .jsonl)
-    if _validate_tasks(tasks):
+    if already_remote and not existing_images:
+        # Tasks are remote but have no image references - just return as-is
         return str(tasks_path)
 
     # Locate environment
-    env_dir = _find_environment_dir(tasks_path)
+    env_dir = find_environment_dir(tasks_path)
     if not env_dir:
         hud_console.error("Could not locate an environment directory (Dockerfile + pyproject.toml)")
         hud_console.hint("Ensure you're in or near your environment folder before running 'hud rl'")
         raise typer.Exit(1)
 
     # Ensure built and pushed
-    lock_data = _ensure_built(env_dir)
+    lock_data = ensure_built(env_dir, interactive=True)
     lock_data = _ensure_pushed(env_dir, lock_data)
 
     # Derive remote image name org/name:tag
     remote_image = _derive_remote_image(lock_data)
 
+    # Check if existing images are outdated
+    needs_update = False
+    should_update_image = False
+    if existing_images:
+        # Check if any existing image differs from the latest
+        for existing_img in existing_images:
+            if existing_img != remote_image:
+                hud_console.warning(f"Detected outdated image reference: {existing_img}")
+                hud_console.info(f"Latest pushed image: {remote_image}")
+                needs_update = True
+                break
+
+        if needs_update:
+            confirm_msg = "Update task configuration with the latest image?"
+            if hud_console.confirm(confirm_msg, default=True):
+                hud_console.info("Updating task configuration with latest image...")
+                should_update_image = True
+            else:
+                # If user doesn't want to update, just return the original file
+                if already_remote:
+                    return str(tasks_path)
+                # Otherwise, continue with conversion but keep old images
+                remote_image = next(iter(existing_images))  # Use the first existing image
+
+    # If tasks are already remote and up-to-date (no update needed), return original file
+    if already_remote and not needs_update:
+        return str(tasks_path)
+
+    # If tasks are already remote and we just need to update the image
+    if already_remote and should_update_image:
+        # Update image references in-place on RAW tasks (preserve placeholders)
+        def _update_image_refs_raw(obj: Any) -> Any:
+            if isinstance(obj, dict):
+                new_obj = {}
+                for k, v in obj.items():
+                    if k == "Mcp-Image" and isinstance(v, str) and v in existing_images:
+                        new_obj[k] = remote_image
+                    else:
+                        new_obj[k] = _update_image_refs_raw(v)
+                return new_obj
+            elif isinstance(obj, list):
+                return [_update_image_refs_raw(item) for item in obj]
+            else:
+                return obj
+
+        updated_raw_tasks: list[dict[str, Any]] = []
+        for t in raw_tasks:
+            td = dict(t)
+            if "mcp_config" in td:
+                td["mcp_config"] = _update_image_refs_raw(td["mcp_config"])
+            updated_raw_tasks.append(td)
+
+        # Write updated file (preserve original format - check if it's .jsonl)
+        if tasks_path.suffix == ".jsonl":
+            with open(tasks_path, "w", encoding="utf-8") as f:
+                for task in updated_raw_tasks:
+                    json.dump(task, f, ensure_ascii=False)
+                    f.write("\n")
+        else:
+            with open(tasks_path, "w", encoding="utf-8") as f:
+                json.dump(updated_raw_tasks, f, ensure_ascii=False, indent=2)
+                f.write("\n")
+
+        hud_console.success(f"Updated {tasks_path.name} with latest image: {remote_image}")
+        return str(tasks_path)
+
+    # Extract additional API key headers from lock and suggest from .env
+    provided_keys = _extract_api_key_vars(lock_data)
+    dotenv_keys = _extract_dotenv_api_key_vars(env_dir)
+
+    # If .env contains API-like vars not in lock, offer to include them
+    missing = sorted(dotenv_keys - provided_keys)
+    if missing:
+        names_preview = ", ".join(missing)
+        prompt = (
+            f"Detected env vars in .env that look like API keys: {names_preview}.\n"
+            "Include them as remote headers (values will be ${VAR} placeholders)?"
+        )
+        if hud_console.confirm(prompt, default=True):
+            provided_keys.update(missing)
+
+    extra_api_key_headers: dict[str, str] = {}
+    for var_name in provided_keys:
+        if str(var_name).upper() == "HUD_API_KEY":
+            continue
+        header_key = _env_var_to_header_key(var_name)
+        extra_api_key_headers[header_key] = f"${{{var_name}}}"
+
     # Helper to strip extra fields from tool calls
     def _simplify_tool_call(tool: Any) -> Any:
         def _one(x: Any) -> dict[str, Any]:
@@ -229,6 +359,9 @@ def convert_tasks_to_remote(tasks_file: str) -> str:
             },
         }
 
+        # Merge additional API key headers
+        item["mcp_config"]["hud"]["headers"].update(extra_api_key_headers)
+
         # Optional fields, omit Nones
         if t.setup_tool is not None:
             item["setup_tool"] = _simplify_tool_call(t.setup_tool)
@@ -243,7 +376,6 @@ def convert_tasks_to_remote(tasks_file: str) -> str:
 
         tasks_payload.append(item)
 
-    # Write new file: remote_<name>.json (always JSON array)
     remote_name = f"remote_{tasks_path.stem}.json"
     remote_path = tasks_path.parent / remote_name
     with open(remote_path, "w", encoding="utf-8") as f:

hud/cli/push.py

@@ -11,6 +11,7 @@ import requests
 import typer
 import yaml
 
+from hud.cli.utils.env_check import ensure_built
 from hud.utils.hud_console import HUDConsole
 
 
@@ -131,6 +132,14 @@ def push_environment(
 
     # Find hud.lock.yaml in specified directory
     env_dir = Path(directory)
+
+    # Ensure environment is built and up-to-date (hash-based); interactive prompt
+    try:
+        ensure_built(env_dir, interactive=True)
+    except typer.Exit:
+        raise
+    except Exception as e:
+        HUDConsole().debug(f"Skipping pre-push build check: {e}")
     lock_path = env_dir / "hud.lock.yaml"
 
     if not lock_path.exists():

hud/cli/rl/__init__.py

@@ -25,7 +25,7 @@ def rl_command(
     ),
     model: str | None = typer.Argument(
         None,
-        help="Model to train (default: interactive selection)",
+        help="Model to train from https://hud.so/models (default: interactive selection)",
     ),
     config_file: Path | None = typer.Option(  # noqa: B008
         None,
@@ -72,6 +72,12 @@ def rl_command(
         "--local",
         help="Run training locally instead of using remote API server",
     ),
+    yes: bool = typer.Option(
+        False,
+        "--yes",
+        "-y",
+        help="Auto-accept all prompts and use defaults (lazy mode)",
+    ),
     # Internal flag
     skip_vllm_startup: bool = typer.Option(
         False,
@@ -122,8 +128,7 @@ def rl_command(
         try:
             from hud.cli.flows.tasks import convert_tasks_to_remote
 
-            console.print("\n[cyan]Preparing remote training tasks...[/cyan]")
-            console.print("[cyan](build/push if needed)[/cyan]")
+            console.print("[cyan]Preparing remote training tasks...[/cyan]")
             tasks_file = convert_tasks_to_remote(tasks_file)
         except typer.Exit:
             raise
@@ -137,7 +142,11 @@ def rl_command(
             from .remote_runner import run_remote_training
 
             run_remote_training(
-                tasks_file=tasks_file, model=model, config_file=config_file, output_dir=output_dir
+                tasks_file=tasks_file,
+                model=model,
+                config_file=config_file,
+                output_dir=output_dir,
+                yes=yes,
             )
             return
         except Exception as e:
@@ -152,6 +161,7 @@ def rl_command(
         model=model,
         config_file=config_file,
         output_dir=output_dir,
+        yes=yes,
         restart=restart,
         verbose=verbose,
         no_ddp=no_ddp,

hud/cli/rl/celebrate.py

@@ -0,0 +1,187 @@
+# ruff: noqa: S311
+from __future__ import annotations
+
+import random
+import time
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar
+
+from rich.live import Live
+from rich.text import Text
+
+from hud.utils.hud_console import hud_console
+
+if TYPE_CHECKING:
+    from rich.console import Console
+
+
+@dataclass
+class Particle:
+    """A confetti particle with physics."""
+
+    x: float
+    y: float
+    vx: float  # velocity x
+    vy: float  # velocity y
+    char: str
+    color: str
+
+    def update(self, gravity: float = 0.5, fps: float = 30.0) -> None:
+        """Update particle position and velocity."""
+        dt = 1.0 / fps
+        self.x += self.vx * dt
+        self.vy += gravity  # Apply gravity
+        self.y += self.vy * dt
+
+
+class ConfettiSystem:
+    """Minimal confetti system inspired by confetty."""
+
+    # Confetty-style colors
+    COLORS: ClassVar[list[str]] = ["#a864fd", "#29cdff", "#78ff44", "#ff718d", "#fdff6a"]
+    # Confetty-style characters
+    CHARS: ClassVar[list[str]] = ["█", "▓", "▒", "░", "▄", "▀"]
+
+    def __init__(self, width: int, height: int) -> None:
+        self.width = width
+        self.height = height
+        self.particles: list[Particle] = []
+
+    def spawn_burst(self, num_particles: int = 75) -> None:
+        """Spawn a burst of confetti particles from the top center."""
+        center_x = self.width / 2
+
+        for _ in range(num_particles):
+            # Start from top center with some horizontal spread
+            x = center_x + (self.width / 4) * (random.random() - 0.5)
+            y = 0
+
+            # Random velocities - horizontal spread and upward/slight downward initial velocity
+            vx = (random.random() - 0.5) * 100
+            vy = random.random() * 50 - 25  # Some go up first
+
+            particle = Particle(
+                x=x,
+                y=y,
+                vx=vx,
+                vy=vy,
+                char=random.choice(self.CHARS),
+                color=random.choice(self.COLORS),
+            )
+            self.particles.append(particle)
+
+    def update(self) -> None:
+        """Update all particles and remove off-screen ones."""
+        # Update physics
+        for particle in self.particles:
+            particle.update()
+
+        # Remove particles that are off-screen
+        self.particles = [p for p in self.particles if 0 <= p.x < self.width and p.y < self.height]
+
+    def render(self) -> str:
+        """Render the particle system to a string."""
+        # Create empty grid
+        grid = [[" " for _ in range(self.width)] for _ in range(self.height)]
+
+        # Place particles
+        for particle in self.particles:
+            x, y = int(particle.x), int(particle.y)
+            if 0 <= x < self.width and 0 <= y < self.height:
+                grid[y][x] = particle.char
+
+        # Convert to string
+        return "\n".join("".join(row) for row in grid)
+
+    def render_with_colors(self) -> Text:
+        """Render the particle system with colors for Rich."""
+        text = Text()
+
+        # Create empty grid with color info
+        grid: list[list[tuple[str, str] | None]] = [
+            [None for _ in range(self.width)] for _ in range(self.height)
+        ]
+
+        # Place particles with their colors
+        for particle in self.particles:
+            x, y = int(particle.x), int(particle.y)
+            if 0 <= x < self.width and 0 <= y < self.height:
+                grid[y][x] = (particle.char, particle.color)
+
+        # Build colored text
+        for row in grid:
+            for cell in row:
+                if cell:
+                    char, color = cell
+                    text.append(char, style=color)
+                else:
+                    text.append(" ")
+            text.append("\n")
+
+        return text
+
+
+def show_confetti(console: Console, seconds: float = 2.5) -> None:
+    """Display celebratory confetti animation inspired by confetty.
+
+    Shows "Starting training!" message first, then creates two bursts of
+    falling confetti particles that fall away completely.
+
+    Args:
+        console: Rich console instance
+        seconds: Duration to show confetti
+    """
+    # Show celebratory message first
+    console.print(
+        "[bold green]🎉 Starting training! See your model on https://hud.so/models[/bold green]"
+    )
+    time.sleep(0.3)  # Brief pause to see the message
+
+    width = min(console.size.width, 120)  # Cap width for performance
+    height = min(console.size.height - 2, 30)  # Leave room for message
+
+    # Create confetti system
+    system = ConfettiSystem(width, height)
+
+    fps = 30
+    frame_time = 1.0 / fps
+
+    # First burst at the beginning
+    system.spawn_burst(num_particles=60)
+
+    # Track when to spawn second burst
+    second_burst_frame = int(fps * 0.4)  # Second burst after 0.4 seconds
+
+    with Live("", refresh_per_second=fps, console=console, transient=True) as live:
+        frame = 0
+        # Keep running until all particles have fallen off screen
+        while frame < seconds * fps or len(system.particles) > 0:
+            # Spawn second burst
+            if frame == second_burst_frame:
+                system.spawn_burst(num_particles=60)
+
+            system.update()
+            live.update(system.render_with_colors())
+            time.sleep(frame_time)
+            frame += 1
+
+
+def show_confetti_async(console: Console, seconds: float = 2.5) -> None:
+    """Non-blocking confetti animation that runs in a background thread.
+
+    The animation will run independently while training starts immediately.
+    """
+    import threading
+
+    def _run_confetti() -> None:
+        try:
+            show_confetti(console, seconds)
+        except Exception:
+            hud_console.info("Launching training...")
+
+    thread = threading.Thread(target=_run_confetti, daemon=True)
+    thread.start()
+    # Don't wait - let training start immediately while confetti plays
+
+
+__all__ = ["show_confetti", "show_confetti_async"]

hud/cli/rl/config.py

@@ -21,22 +21,29 @@ console = Console()
 def generate_config_interactive(
     model_name: str,
     presets: list[dict[str, Any]],
+    yes: bool = False,
 ) -> 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
+    if not yes:
+        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
+    if yes:
+        # Use default preset (Balanced if available, otherwise first)
+        preset_choice = 1 if len(presets) > 1 else 0
+        selected_preset = presets[preset_choice]
+        hud_console.info(f"Auto-selecting preset: {selected_preset['name']} (--yes mode)")
+    else:
+        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"]