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

hud/otel/instrumentation.py

@@ -55,6 +55,9 @@ def _patch_mcp_instrumentation() -> None:
     try:
         from opentelemetry.instrumentation.mcp.instrumentation import McpInstrumentor
 
+        # First, patch the get_error_type function to handle invalid HTTP status codes
+        _patch_get_error_type()
+
         def patched_transport_wrapper(self: Any, tracer: Any) -> Callable[..., Any]:
             @asynccontextmanager
             async def traced_method(
@@ -98,3 +101,35 @@ def _patch_mcp_instrumentation() -> None:
 
         logger = logging.getLogger(__name__)
         logger.warning("Failed to patch MCP instrumentation: %s", e)
+
+
+def _patch_get_error_type() -> None:
+    """Patch get_error_type to handle invalid HTTP status codes gracefully."""
+    import re
+    from http import HTTPStatus
+
+    try:
+        import opentelemetry.instrumentation.mcp.instrumentation as mcp_inst
+
+        def patched_get_error_type(error_message: str) -> str | None:
+            """Extract HTTP status from error message, handling invalid codes."""
+            if not isinstance(error_message, str):
+                return None
+            match = re.search(r"\b(4\d{2}|5\d{2})\b", error_message)
+            if match:
+                num = int(match.group())
+                try:
+                    # Only return if it's a valid HTTPStatus
+                    if 400 <= num <= 599:
+                        return HTTPStatus(num).name
+                except ValueError:
+                    # Not a valid HTTP status code
+                    logger.debug("Ignoring invalid HTTP status code: %s", num)
+            return None
+
+        # Apply the patch
+        mcp_inst.get_error_type = patched_get_error_type
+        logger.debug("Patched get_error_type to handle invalid HTTP status codes")
+
+    except Exception as e:
+        logger.warning("Failed to patch get_error_type: %s", e)

hud/rl/README.md

@@ -0,0 +1,31 @@
+We suggest running hud rl (or with the --local flag) for optimal hyperparameters and native HuggingFace running.
+
+Install:
+```bash
+sudo apt-get update -y && sudo apt-get install -y cuda-toolkit-12-6
+uv pip install -e .[rl]
+uv pip install ninja
+uv pip install flash-attn --no-build-isolation
+```
+
+
+However, if you want to run the training directly, launch a vllm server with:
+```bash
+export VLLM_ALLOW_RUNTIME_LORA_UPDATING=True
+export TOKENIZERS_PARALLELISM=false
+export VLLM_LOGGING_LEVEL=INFO
+export CUDA_VISIBLE_DEVICES=7 # Set this to your last GPU
+
+uv run vllm serve Qwen/Qwen2.5-VL-3B-Instruct \
+    --api-key token-abc123 --host 0.0.0.0 --port 8000 --tensor-parallel-size 1 --trust-remote-code \
+    --max-model-len 16384 --enable-lora --max-lora-rank 64 --max-cpu-loras 4 --enable-auto-tool-choice \
+    --tool-call-parser hermes --disable-log-requests --dtype auto
+```
+
+And training with (replace 2 with your spare GPUs):
+```bash
+hud get hud-evals/2048-basic
+torchrun --nproc-per-node 2 -m hud.rl.train --tasks 2048-basic.json --verbose
+```
+
+Add a `--config path/to/config.json` flag to run a specific configuration (or change the defaults in config.py)

hud/rl/__init__.py

@@ -0,0 +1 @@
+"""RL module for HUD."""

hud/rl/actor.py

@@ -0,0 +1,174 @@
+"""Actor for episode collection using vLLM and HUD."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+
+import httpx
+from openai import AsyncOpenAI
+
+import hud
+from hud.agents.openai_chat_generic import GenericOpenAIChatAgent
+from hud.clients.utils.retry_transport import create_retry_httpx_client
+from hud.types import Task, Trace
+from hud.utils.hud_console import HUDConsole
+
+from .config import Config
+
+logger = logging.getLogger(__name__)
+hud_console = HUDConsole(logger)
+
+
+class Actor:
+    """Collects episodes using vLLM-served models via HUD agents."""
+
+    def __init__(self, config: Config) -> None:
+        self.config = config
+        self.actor_config = config.actor
+        self.current_adapter = config.model.base_model
+
+        # Setup OpenAI client for vLLM
+        base_url = self.actor_config.vllm_base_url.replace("localhost", "127.0.0.1")
+        self.openai_client = self._create_openai_client(base_url)
+
+    def _create_openai_client(self, base_url: str) -> AsyncOpenAI:
+        """Create OpenAI client with optimized settings for vLLM."""
+        # Match connection limits to parallel_episodes to avoid bottlenecks
+        # Use shorter per-request timeout and keep retries modest to avoid long blocking
+        http_client = create_retry_httpx_client(
+            timeout=httpx.Timeout(30.0),
+        )
+        return AsyncOpenAI(
+            base_url=base_url,
+            api_key=self.actor_config.vllm_api_key,
+            http_client=http_client,
+            max_retries=2,
+        )
+
+    def create_agent(self) -> GenericOpenAIChatAgent:
+        """Create an agent with the current adapter."""
+        return GenericOpenAIChatAgent(
+            openai_client=self.openai_client,
+            model_name=self.current_adapter,
+            allowed_tools=self.actor_config.allowed_tools,
+            append_setup_output=False,
+            system_prompt=self.actor_config.system_prompt,
+            verbose=self.config.verbose,
+            completion_kwargs={
+                "temperature": self.actor_config.temperature,
+                "max_tokens": self.actor_config.max_new_tokens,
+                "tool_choice": "required" if self.actor_config.force_tool_choice else "auto",
+            },
+        )
+
+    def update_adapter(self, adapter_name: str) -> None:
+        """Update the current adapter being used."""
+        self.current_adapter = adapter_name
+        hud_console.info(f"[Actor] Using adapter: {adapter_name}")
+
+    async def run_tasks(self, tasks: list[Task], job_id: str) -> list[Trace]:
+        """Run tasks and collect traces."""
+        traces = []
+
+        # Process tasks in batches respecting max_parallel_episodes limit
+        for batch_start in range(0, len(tasks), self.actor_config.max_parallel_episodes):
+            batch_end = min(batch_start + self.actor_config.max_parallel_episodes, len(tasks))
+            batch = tasks[batch_start:batch_end]
+
+            # Run batch in parallel with per-episode timeout protection
+            async def run_with_timeout(t: Task) -> Trace:
+                try:
+                    return await asyncio.wait_for(
+                        self._run_task(t, job_id),
+                        timeout=self.actor_config.episode_timeout_sec,
+                    )
+                except TimeoutError:
+                    hud_console.warning_log(f"Episode timed out for task {t.id}")
+                    return Trace(isError=True, content="Episode timeout")
+
+            results = await asyncio.gather(
+                *[run_with_timeout(t) for t in batch],
+                return_exceptions=True,
+            )
+
+            # Normalize exceptions to error traces
+            for res in results:
+                if isinstance(res, Exception):
+                    hud_console.warning_log(f"Episode error: {res}")
+                    traces.append(Trace(isError=True, content=str(res)))
+                else:
+                    traces.append(res)
+
+        return traces
+
+    async def _run_task(self, task: Task, job_id: str) -> Trace:
+        """Run a single task."""
+        agent = self.create_agent()
+
+        # Run the task
+        try:
+            with hud.trace(f"Training | {task.id}", job_id=job_id):
+                result = await agent.run(task, max_steps=self.actor_config.max_steps_per_episode)
+
+        except Exception:
+            logger.info("GOT EXCEPTION")
+            return Trace(isError=True)
+
+        result.info["tool_spec"] = agent.get_tool_schemas()
+
+        return result
+
+
+if __name__ == "__main__":
+    from hud.types import Task
+
+    async def test_actor() -> None:
+        """Test the actor with a single 2048 task using local hud-browser image."""
+        config = Config()
+        config.actor.max_parallel_episodes = 1
+        config.actor.max_steps_per_episode = 6
+        config.actor.episode_timeout_sec = 120
+        config.verbose = True
+
+        # Create test task with local hud-browser image
+        task_data = {
+            "id": "test_2048_128",
+            "prompt": "Play the browser-based 2048 game and try to reach the 128 tile. Start by taking a screenshot, then make strategic moves using arrow keys.",  # noqa: E501
+            "mcp_config": {
+                "local": {
+                    "command": "sh",
+                    "args": [
+                        "-c",
+                        "docker run --rm --platform linux/amd64 -i hud-browser:latest 2>/dev/null",
+                    ],
+                }
+            },
+            "setup_tool": {"name": "launch_app", "arguments": {"app_name": "2048"}},
+            "evaluate_tool": {
+                "name": "evaluate",
+                "arguments": {"name": "game_2048_max_number", "arguments": {"target": 128}},
+            },
+            "system_prompt": "You are an expert 2048 game player. Use arrow keys to reach the target tile. First take a screenshot, then make strategic moves.",  # noqa: E501
+        }
+
+        task = Task(**task_data)
+        actor = Actor(config)
+
+        logger.info("Testing actor with task: %s", task.id)
+        logger.info("Model: %s", config.model.base_model)
+        logger.info("VLLM: %s", config.actor.vllm_base_url)
+
+        traces = await actor.run_tasks([task], job_id="test_2048")
+
+        for trace in traces:
+            if trace.isError:
+                logger.info("Error: %s", trace.content)
+            else:
+                logger.info("Success!")
+                logger.info("Trace info: %s", trace.info if hasattr(trace, "info") else "No info")
+                # Check for evaluation in the trace info
+                if hasattr(trace, "info") and "evaluation" in trace.info:
+                    logger.info("  Evaluation: %s", trace.info["evaluation"])
+
+    asyncio.run(test_actor())

hud/rl/buffer.py

@@ -0,0 +1,371 @@
+"""Replay buffer for storing and sampling episodes."""
+
+from __future__ import annotations
+
+import logging
+import random
+from collections import deque
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+from hud.types import Task, Trace
+from hud.utils.hud_console import HUDConsole
+
+logger = logging.getLogger(__name__)
+hud_console = HUDConsole(logger=logger)
+
+T = TypeVar("T")
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from hud.rl.config import Config
+
+
+class Buffer(Generic[T]):
+    """Simple buffer for a list of tasks, traces or episodes."""
+
+    def __init__(self, max_size: int = 10000) -> None:
+        self.max_size = max_size
+        self.buffer: deque[T] = deque(maxlen=max_size)
+
+    def add(self, items: list[T] | T, shuffle: bool = False) -> None:
+        """Add items to buffer."""
+        if isinstance(items, list):
+            for item in items:
+                self.buffer.append(item)
+        else:
+            self.buffer.append(items)
+        if shuffle:
+            random.shuffle(self.buffer)
+
+    def add_fill(self, items: list[T] | T, target_size: int, shuffle: bool = False) -> None:
+        """Add items to buffer until the buffer is at least the target size."""
+        while len(self.buffer) < target_size:
+            self.add(items, shuffle)
+
+    def get(self, n: int = 0) -> list[T]:
+        """Get items from the buffer."""
+        if n == 0:
+            return list(self.buffer)
+        if n > len(self.buffer):
+            raise ValueError("Not enough items in buffer")
+        return list(self.buffer)[-n:]
+
+    def consume(self, n: int = 0) -> list[T]:
+        """Consume items from the buffer."""
+        if n == 0:
+            return list(self.buffer)
+        if n > len(self.buffer):
+            raise ValueError("Not enough items in buffer")
+
+        return [self.buffer.pop() for _ in range(n)]
+
+    def get_filtered(
+        self, n: int = 0, filter_fn: Callable[[T], bool] | None = None, consume: bool = False
+    ) -> list[T]:
+        """Filter the buffer by a filter function."""
+        filtered = (
+            [item for item in self.buffer if filter_fn(item)] if filter_fn else list(self.buffer)
+        )
+        if n == 0:
+            return filtered
+        return self.consume(n) if consume else self.get(n)
+
+    def sample(
+        self,
+        batch_size: int,
+        n: int = 0,
+        filter_fn: Callable[[T], bool] | None = None,
+        consume: bool = False,
+    ) -> list[T]:
+        """Sample a batch of items with optional filtering."""
+        items = self.get_filtered(n, filter_fn, consume)
+
+        if len(items) < batch_size:
+            hud_console.warning(f"Buffer has {len(items)} items, requested {batch_size}")
+            return items
+
+        return random.sample(items, batch_size)
+
+    def clear(self) -> None:
+        """Clear the buffer."""
+        self.buffer.clear()
+
+    def __len__(self) -> int:
+        """Use len() directly on Buffer instances."""
+        return len(self.buffer)
+
+
+class DatasetBuffer(Buffer[Task]):
+    """
+    Buffer for a dataset.
+    Loads in individual tasks that will be trained for a specified number of training steps.
+    """
+
+    def __init__(
+        self,
+        dataset: list[Task] | Task,
+        config: Config,
+    ) -> None:
+        self.config = config
+
+        self.group_size = config.training.group_size
+        self.batch_size = config.training.batch_size
+        self.training_steps = config.training.training_steps
+
+        if self.group_size > self.batch_size:
+            raise ValueError(
+                f"Group size is greater than batch size, {self.group_size} > {self.batch_size}"
+            )
+
+        if self.batch_size % self.group_size != 0:
+            raise ValueError(
+                f"A batch cannot have irregular groups, {self.group_size} % {self.batch_size} != 0"
+            )
+
+        if self.group_size % config.training.mini_batch_size != 0:
+            raise ValueError(
+                f"Group size is not a multiple of mini batch size, {self.group_size} % {config.training.mini_batch_size} != 0"  # noqa: E501
+            )
+
+        self.groups_per_batch = self.batch_size // self.group_size
+        self.number_of_tasks = self.training_steps * self.groups_per_batch
+
+        super().__init__(self.number_of_tasks)
+
+        dataset = dataset if isinstance(dataset, list) else [dataset]
+        tasks = self._validate_tasks(dataset)
+        if config.training.shuffle_dataset:
+            random.shuffle(tasks)
+        if len(tasks) > self.number_of_tasks:
+            leftovers = len(tasks) - self.number_of_tasks
+            hud_console.warning(
+                f"Training steps ({self.training_steps}) will lead to {leftovers} tasks not being trained"  # noqa: E501
+            )
+            tasks = tasks[: self.number_of_tasks]
+
+        # Check if the dataset is imbalanced
+        self.dataset_size = len(tasks)
+        if self.training_steps % self.dataset_size != 0:
+            leftovers = self.number_of_tasks % self.dataset_size
+            hud_console.warning(
+                f"Dataset imbalanced ({leftovers} tasks will be trained 1 more time)"
+            )
+            hud_console.warning(
+                f"This is because the number of training steps ({self.training_steps}) is not a multiple of the dataset size ({self.dataset_size})"  # noqa: E501
+            )
+
+        self.add_fill(tasks, self.number_of_tasks, config.training.shuffle_dataset)
+
+    def _validate_tasks(self, tasks: list[Task]) -> list[Task]:
+        """Validate that all tasks are proper HUD Task objects."""
+        if not tasks:
+            raise ValueError("No tasks provided to DatasetBuffer")
+
+        validated_tasks = []
+        for i, task in enumerate(tasks):
+            if not isinstance(task, Task):
+                raise TypeError(f"Task at index {i} is not a HUD Task object, got {type(task)}")
+            validated_tasks.append(task)
+
+        return validated_tasks
+
+    @property
+    def info(self) -> dict[str, int | float | str]:
+        """Get the info of the buffer."""
+        return {
+            "total_items": len(self),
+            "total_traces": self.number_of_tasks * self.group_size,
+            "total_batches": self.training_steps,
+            "task_repeats": self.number_of_tasks // self.dataset_size,
+            "dataset_size": self.dataset_size,
+            "group_size": self.group_size,
+            "batch_size": self.batch_size,
+        }
+
+    def get_tasks(self, consume: bool = True) -> list[Task]:
+        """Get tasks for a batch."""
+        tasks = self.consume(self.groups_per_batch) if consume else self.get(self.groups_per_batch)
+        # Create groups where each group contains group_size copies of the same task
+        result = []
+        for task in tasks:
+            result.extend([task] * self.group_size)
+        return result
+
+
+class ReplayBuffer(Buffer[Trace]):
+    """Buffer for traces."""
+
+    def __init__(self, config: Config) -> None:
+        self.config = config
+
+        self.buffer_steps = config.training.buffer_steps
+        self.select_strategy = config.training.select_strategy
+        self.group_size = config.training.group_size
+        self.batch_size = config.training.batch_size
+
+        buffer_size = self.buffer_steps * self.batch_size
+
+        super().__init__(buffer_size)
+
+    def sample_traces(self) -> list[Trace]:
+        """Sample traces for a batch."""
+        if self.select_strategy == "recent":
+            return self.get(self.batch_size)
+        elif self.select_strategy == "random":
+            return self.sample(self.batch_size)
+        elif self.select_strategy == "variance":
+            return self._sample_high_variance_traces()
+        else:
+            raise ValueError(f"Invalid select strategy: {self.select_strategy}")
+
+    def _sample_high_variance_traces(self) -> list[Trace]:
+        from collections import Counter, defaultdict, deque
+
+        # Expect recent window to already be grouped by task id
+
+        # Build recent window and earlier lookup (short form)
+        buf_list = list(self.buffer)
+        if len(buf_list) < self.batch_size:
+            hud_console.warning(
+                f"[group-sampler] Buffer has only {len(buf_list)} traces, need {self.batch_size}"
+            )
+            while len(buf_list) < self.batch_size:
+                take = min(len(buf_list) or 1, self.batch_size - len(buf_list))
+                buf_list.extend(buf_list[:take])
+        recent_traces = buf_list[-self.batch_size :]
+        hud_console.info(
+            f"[group-sampler] recent-window histogram: {Counter(getattr(t.task, 'id', 'NA') for t in recent_traces)}"  # noqa: E501
+        )
+
+        hud_console.info(
+            f"[group-sampler] Building earlier traces lookup, buffer size: {len(buf_list)}"
+        )
+        earlier_traces_by_task: dict[str, deque[Trace]] = defaultdict(deque)
+        for tr in buf_list[: -self.batch_size]:
+            earlier_traces_by_task[getattr(tr.task, "id", "NA")].append(tr)
+
+        # Chunk from the most-recent end
+        final_traces: list[Trace] = []
+        groups_per_batch = self.batch_size // self.group_size
+        hud_console.info(f"[group-sampler] Processing {groups_per_batch} groups")
+        for g_idx in range(groups_per_batch):
+            start = g_idx * self.group_size
+            end = start + self.group_size
+            group = recent_traces[start:end]
+
+            # Assert homogeneity: every trace in a group must share the same task id
+            cnt = Counter(getattr(t.task, "id", "NA") for t in group)
+            if len(cnt) != 1:
+                raise RuntimeError(f"Group {g_idx} is not homogeneous: {dict(cnt)}")
+            target_tid = next(iter(cnt.keys()))
+
+            # Build homogeneous group of target_tid, filling from earlier traces to increase spread
+            homogeneous: list[Trace] = [
+                t for t in group if getattr(t.task, "id", "NA") == target_tid
+            ]
+            needed = self.group_size - len(homogeneous)
+
+            # Greedy fill: choose earlier traces (same task-id) farthest from current mean reward
+            def current_mean(homogeneous: list[Trace]) -> float:
+                if not homogeneous:
+                    return 0.0
+                vals = [float(getattr(t, "reward", 0.0) or 0.0) for t in homogeneous]
+                return sum(vals) / len(vals)
+
+            while needed > 0:
+                pool = earlier_traces_by_task.get(target_tid, deque())
+                if pool:
+                    mu = current_mean(homogeneous)
+                    # pick element farthest from current mean
+                    best_i = None
+                    best_dist = -1.0
+                    for i, tr in enumerate(list(pool)):
+                        r = float(getattr(tr, "reward", 0.0) or 0.0)
+                        dist = abs(r - mu)
+                        if dist > best_dist:
+                            best_dist = dist
+                            best_i = i
+                    # pop selected
+                    chosen = list(pool)[best_i]  # type: ignore[index]
+                    # remove from deque efficiently by rotating
+                    left = list(pool)
+                    if best_i is not None:
+                        left.pop(best_i)  # O(n) but pool is small in practice
+                        earlier_traces_by_task[target_tid] = deque(left)
+                        homogeneous.append(chosen)
+                else:
+                    # duplicate extreme within current homogeneous set
+                    if not homogeneous:
+                        raise RuntimeError(f"Group {g_idx} has no traces for target {target_tid}")
+                    mu = current_mean(homogeneous)
+                    extreme = max(
+                        homogeneous, key=lambda t: abs(float(getattr(t, "reward", 0.0) or 0.0) - mu)
+                    )
+                    homogeneous.append(extreme)
+                needed -= 1
+
+            # Replacement step: swap in earlier traces to increase reward spread
+            pool = earlier_traces_by_task.get(target_tid, deque())
+            if pool:
+                # Log pool stats
+                pool_vals = [float(getattr(tr, "reward", 0.0) or 0.0) for tr in list(pool)]
+                if pool_vals:
+                    pool_mean = sum(pool_vals) / len(pool_vals)
+                    pool_var = sum((v - pool_mean) * (v - pool_mean) for v in pool_vals) / len(
+                        pool_vals
+                    )
+                    hud_console.info(
+                        f"[group-sampler] Group {g_idx}: earlier-pool size={len(pool_vals)} mean={pool_mean:.4f} std={(pool_var**0.5):.4f}"  # noqa: E501
+                    )
+
+                # Decide how many to replace (up to 1/4 of group, at least 1)
+                replace_k = max(1, self.group_size // 4)
+                replace_k = min(replace_k, len(pool), self.group_size)
+
+                if replace_k > 0:
+                    mu = current_mean(homogeneous)
+                    # Select replacement candidates from pool farthest from current mean
+                    pool_list = list(pool)
+                    pool_indices = list(range(len(pool_list)))
+                    pool_indices.sort(
+                        key=lambda i: abs(
+                            (float(getattr(pool_list[i], "reward", 0.0) or 0.0)) - mu
+                        ),
+                        reverse=True,
+                    )
+                    chosen_pool_idx = set(pool_indices[:replace_k])
+                    replacements = [pool_list[i] for i in pool_indices[:replace_k]]
+
+                    # Remove chosen from pool deque
+                    remaining = [tr for i, tr in enumerate(pool_list) if i not in chosen_pool_idx]
+                    earlier_traces_by_task[target_tid] = deque(remaining)
+
+                    # Select current group positions closest to mean to replace
+                    group_indices = list(range(len(homogeneous)))
+                    group_indices.sort(
+                        key=lambda i: abs(
+                            (float(getattr(homogeneous[i], "reward", 0.0) or 0.0)) - mu
+                        )
+                    )
+                    target_positions = group_indices[:replace_k]
+
+                    for pos, new_tr in zip(target_positions, replacements, strict=False):
+                        homogeneous[pos] = new_tr
+
+            # Validate homogeneity
+            if any(getattr(t.task, "id", "NA") != target_tid for t in homogeneous):
+                raise RuntimeError(f"Group {g_idx} is not homogeneous after sampling")
+            final_traces.extend(homogeneous)
+
+        for i in range(0, len(final_traces), self.group_size):
+            block = final_traces[i : i + self.group_size]
+            if len({getattr(t.task, "id", "NA") for t in block}) != 1:
+                raise RuntimeError(f"Homogeneity validation failed for block starting at index {i}")
+
+        hud_console.info(
+            f"[group-sampler] final histogram: {Counter(getattr(t.task, 'id', 'NA') for t in final_traces)}"  # noqa: E501
+        )
+        return final_traces
+
+        # --------------------------------------------------------------------