hud-python 0.4.42__py3-none-any.whl → 0.4.44__py3-none-any.whl

hud/agents/openai_chat_generic.py

@@ -205,7 +205,7 @@ class GenericOpenAIChatAgent(MCPAgent):
         try:
             response = await self._invoke_chat_completion(
                 messages=messages,
-                tools=tools, # type: ignore
+                tools=tools,  # type: ignore
                 extra=extra,
             )
         except Exception as e:

hud/cli/__init__.py

@@ -1178,6 +1178,11 @@ def rl(
         "--vllm-gpu",
         help="Specific GPU for vLLM server",
     ),
+    vllm_gpu_count: int = typer.Option(
+        1,
+        "--vllm-gpu-count",
+        help="Number of GPUs for vLLM server",
+    ),
     skip_vllm_startup: bool = typer.Option(
         False,
         "--skip-vllm-startup",
@@ -1199,6 +1204,7 @@ def rl(
         no_ddp=no_ddp,
         ddp_gpus=ddp_gpus,
         vllm_gpu=vllm_gpu,
+        vllm_gpu_count=vllm_gpu_count,
         yes=yes,
         skip_vllm_startup=skip_vllm_startup,
     )

hud/cli/dev.py

@@ -73,6 +73,24 @@ def create_proxy_server(
         "PYTHONUNBUFFERED=1",  # Ensure Python output is not buffered
     ]
 
+    # Check for .env file in the project directory and add env vars
+    env_file = project_path / ".env"
+    loaded_env_vars = {}
+    if env_file.exists():
+        try:
+            from hud.cli.utils.config import parse_env_file
+
+            env_contents = env_file.read_text(encoding="utf-8")
+            loaded_env_vars = parse_env_file(env_contents)
+            for key, value in loaded_env_vars.items():
+                docker_cmd.extend(["-e", f"{key}={value}"])
+            if verbose and loaded_env_vars:
+                hud_console.info(
+                    f"Loaded {len(loaded_env_vars)} environment variable(s) from .env file"
+                )
+        except Exception as e:
+            hud_console.warning(f"Failed to load .env file: {e}")
+
     # Add user-provided Docker arguments
     if docker_args:
         docker_cmd.extend(docker_args)
@@ -112,8 +130,12 @@ def create_proxy_server(
             hud_console.info("The container's CMD determines reload behavior")
         hud_console.command_example(f"docker logs -f {container_name}", "View container logs")
 
-        # Show the full Docker command if there are environment variables
-        if docker_args and any(arg == "-e" or arg.startswith("--env") for arg in docker_args):
+        # Show the full Docker command if there are environment variables (from .env or args)
+        has_env_from_args = docker_args and any(
+            arg == "-e" or arg.startswith("--env") for arg in docker_args
+        )
+        has_env_from_file = bool(loaded_env_vars)
+        if has_env_from_args or has_env_from_file:
             hud_console.info("")
             hud_console.info("Docker command with environment variables:")
             hud_console.info(" ".join(docker_cmd))

hud/cli/eval.py

@@ -298,16 +298,15 @@ async def run_single_task(
                 agent_config["allowed_tools"] = allowed_tools
 
         # Run with grouping
-        with hud.trace(name=f"{task_prompt} (group_size={group_size})"):
-            stats = await run_tasks_grouped(
-                tasks=[task],
-                agent_class=agent_class,
-                agent_config=agent_config,
-                group_size=group_size,
-                max_parallel_episodes=48,  # Same as RL default
-                max_steps=max_steps,
-                verbose=verbose,
-            )
+        stats = await run_tasks_grouped(
+            tasks=[task],
+            agent_class=agent_class,
+            agent_config=agent_config,
+            group_size=group_size,
+            max_parallel_episodes=48,  # Same as RL default
+            max_steps=max_steps,
+            verbose=verbose,
+        )
 
         # Display results
         display_group_statistics(stats, show_details=True)
@@ -499,7 +498,7 @@ async def run_full_dataset(
             )
 
         # Display results
-        display_group_statistics(stats, show_details=len(stats) <= 20)
+        display_group_statistics(stats, show_details=len(stats) <= 50)
 
         # Return stats for consistency with other modes
         return stats

hud/cli/flows/tasks.py

@@ -212,17 +212,14 @@ def convert_tasks_to_remote(tasks_file: str) -> str:
     # Check if tasks already have remote URLs
     already_remote = _validate_tasks(tasks)
 
-    # If tasks already reference a remote MCP URL, do not require a local environment
-    # or attempt any image updates. Use the dataset as-is.
-    if already_remote:
-        return str(tasks_path)
-
     # Extract existing images from tasks
     existing_images = _extract_existing_images(tasks)
 
     # Locate environment
     env_dir = find_environment_dir(tasks_path)
     if not env_dir:
+        if already_remote:
+            return str(tasks_path)
         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)
@@ -373,6 +370,8 @@ def convert_tasks_to_remote(tasks_file: str) -> str:
             item["system_prompt"] = t.system_prompt
         if t.metadata:
             item["metadata"] = t.metadata
+        if t.id is not None:
+            item["id"] = t.id
 
         tasks_payload.append(item)
 

hud/cli/rl/__init__.py

@@ -78,6 +78,11 @@ def rl_command(
         "-y",
         help="Auto-accept all prompts and use defaults (lazy mode)",
     ),
+    vllm_gpu_count: int = typer.Option(
+        None,
+        "--vllm-gpu-count",
+        help="Number of GPUs for vLLM server",
+    ),
     skip_vllm_startup: bool = typer.Option(
         False,
         "--skip-vllm-startup",
@@ -145,6 +150,7 @@ def rl_command(
                 model=model,
                 config_file=config_file,
                 output_dir=output_dir,
+                vllm_gpu_count=vllm_gpu_count,
                 yes=yes,
             )
             return

hud/cli/rl/config.py

@@ -84,7 +84,7 @@ 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:
+    with open(path, "w", encoding="utf-8") as f:
         json.dump(config_dict, f, indent=2)
         f.write("\n")  # Add newline at end of file
 
@@ -94,7 +94,7 @@ def save_config(config: Config, path: Path) -> None:
 
 def load_config(path: Path) -> Config:
     """Load configuration from a JSON file."""
-    with open(path) as f:
+    with open(path, encoding="utf-8") as f:
         data = json.load(f)
 
     # Use Config.from_dict which handles missing fields gracefully

hud/cli/rl/gpu_utils.py

@@ -245,10 +245,12 @@ def adjust_config_for_ddp(config: Config, num_gpus: int) -> Config:
     # Apply scaling rule
     if num_gpus == 1:
         # Special case: 2 groups for single GPU
+        groups_per_gpu = 2
         config.training.batch_size = 2 * group_size
     else:
-        # Multi-GPU: each GPU processes 1 group
-        config.training.batch_size = num_gpus * group_size
+        groups_per_gpu = config.training.batch_size // group_size
+        # Multi-GPU: each GPU processes groups_per_gpu groups
+        config.training.batch_size = num_gpus * group_size * groups_per_gpu
 
     # Update max_parallel_episodes to match
     config.actor.max_parallel_episodes = config.training.batch_size
@@ -263,7 +265,7 @@ def adjust_config_for_ddp(config: Config, num_gpus: int) -> Config:
         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
+        f"[cyan]   Each of the {num_gpus} GPU(s) will process {groups_per_gpu} group(s) in parallel[/cyan]"  # noqa: E501
     )
 
     return config

hud/cli/rl/remote_runner.py

@@ -32,7 +32,9 @@ GPU_PRICING = {
 }
 
 
-def ensure_vllm_deployed(model_name: str, gpu_type: str = "A100", timeout: int = 600) -> None:
+def ensure_vllm_deployed(
+    model_name: str, gpu_type: str = "A100", gpu_count: int = 1, timeout: int = 600
+) -> None:
     """Deploy vLLM for a model if needed and wait until it's ready.
 
     Args:
@@ -47,7 +49,7 @@ def ensure_vllm_deployed(model_name: str, gpu_type: str = "A100", timeout: int =
         return
 
     hud_console.info(f"Deploying vLLM server for {model_name}...")
-    rl_api.deploy_vllm(model_name, gpu_type=gpu_type)
+    rl_api.deploy_vllm(model_name, gpu_type=gpu_type, gpu_count=gpu_count)
     hud_console.success("vLLM deployment started")
 
     hud_console.info("Waiting for vLLM server to be ready...")
@@ -72,6 +74,7 @@ def run_remote_training(
     model: str | None,
     config_file: Path | None,
     output_dir: str,
+    vllm_gpu_count: int = 1,
     yes: bool = False,
 ) -> None:
     """Run RL training remotely via the API server following the new interactive flow."""
@@ -183,14 +186,18 @@ def run_remote_training(
 
             # Ask for model type
             if yes:
-                model_type = "Qwen/Qwen2.5-VL-3B-Instruct"  # Default model in yes mode
+                if config_file:
+                    config = load_config(config_file)
+                    model_type = config.model.base_model
+                else:
+                    model_type = "Qwen/Qwen2.5-VL-3B-Instruct"
                 hud_console.info(f"Auto-selecting base model: {model_type} (--yes mode)")
             else:
                 model_type = hud_console.select(
                     "Select base model type:",
                     choices=[
                         {"name": "Qwen2.5-VL-3B-Instruct", "value": "Qwen/Qwen2.5-VL-3B-Instruct"},
-                        # {"name": "Qwen2.5-VL-7B-Instruct", "value": "Qwen/Qwen2.5-VL-7B-Instruct"}, # noqa: E501
+                        {"name": "Qwen2.5-3B-Instruct", "value": "Qwen/Qwen2.5-3B-Instruct"},
                     ],
                     default=0,
                 )
@@ -218,7 +225,7 @@ def run_remote_training(
             try:
                 rl_api.create_model(model_name, model_type)
                 hud_console.success(f"Created model: {model_name}")
-                ensure_vllm_deployed(model_name, gpu_type="A100")
+                ensure_vllm_deployed(model_name, gpu_type="A100", gpu_count=vllm_gpu_count)
 
             except Exception as e:
                 # If the name already exists, suggest a new name and prompt once
@@ -247,7 +254,7 @@ def run_remote_training(
                         rl_api.create_model(chosen, model_type)
                         hud_console.success(f"Created model: {chosen}")
                         model_name = chosen
-                        ensure_vllm_deployed(model_name, gpu_type="A100")
+                        ensure_vllm_deployed(model_name, gpu_type="A100", gpu_count=vllm_gpu_count)
                     except Exception as e2:
                         hud_console.error(f"Failed to create model: {e2}")
                         raise
@@ -281,7 +288,7 @@ def run_remote_training(
                     return
 
             # Ensure vLLM is deployed
-            ensure_vllm_deployed(model_name, gpu_type="A100")
+            ensure_vllm_deployed(model_name, gpu_type="A100", gpu_count=vllm_gpu_count)
     except KeyboardInterrupt:
         hud_console.dim_info("Training cancelled", "")
         return
@@ -323,7 +330,7 @@ def run_remote_training(
             )
 
         if yes:
-            num_gpus = 2 # Default to 2 GPUs in yes mode
+            num_gpus = 2  # Default to 2 GPUs in yes mode
             hud_console.info(f"Auto-selecting {num_gpus} GPU(s) (--yes mode)")
         else:
             num_gpus = hud_console.select(
@@ -425,10 +432,12 @@ def run_remote_training(
         # Load provided config
         hud_console.info(f"Loading configuration from: {config_file}")
         config = load_config(config_file)
-        config_dict = config.to_dict()
         gpu_choice = config.training.gpu_type
         num_gpus = config.training.num_gpus
 
+        config = adjust_config_for_ddp(config, int(num_gpus))
+        config_dict = config.to_dict()
+
     # Launch training
     try:
         # Little celebration before launching

hud/cli/rl/rl_api.py

@@ -61,12 +61,12 @@ def list_models() -> list[RLModelInfo]:
     ]
 
 
-def deploy_vllm(model_name: str, gpu_type: str = "A100") -> dict[str, Any]:
+def deploy_vllm(model_name: str, gpu_type: str = "A100", gpu_count: int = 1) -> dict[str, Any]:
     """Deploy a vLLM server for a model."""
     return make_request_sync(
         method="POST",
         url=f"{settings.hud_rl_url}/models/{model_name}/deploy",
-        json={"gpu_type": gpu_type},
+        json={"gpu_type": gpu_type, "gpu_count": gpu_count},
         api_key=settings.api_key,
     )
 

hud/cli/utils/environment.py

@@ -127,8 +127,4 @@ def is_environment_directory(path: str | Path) -> bool:
         return False
 
     # Must have pyproject.toml
-    if not (dir_path / "pyproject.toml").exists():
-        hud_console.error("pyproject.toml not found")
-        return False
-
-    return True
+    return (dir_path / "pyproject.toml").exists()

hud/rl/config.py

@@ -13,6 +13,7 @@ SUPPORTED_MODELS = [
     "Qwen/Qwen2.5-VL-32B-Instruct",
     "Qwen/Qwen2.5-VL-72B-Instruct",
     "Qwen/Qwen2.5-7B-Instruct",
+    "Qwen/Qwen2.5-3B-Instruct",
 ]
 
 
@@ -39,9 +40,9 @@ class ModelConfig:
     """Model and LoRA configuration."""
 
     base_model: str = "Qwen/Qwen2.5-VL-3B-Instruct"
-    lora_r: int = 8
-    lora_alpha: int = 16
-    lora_dropout: float = 0.05
+    lora_r: int = 16
+    lora_alpha: int = 32
+    lora_dropout: float = 0.1
     target_modules: tuple[str, ...] = (
         "q_proj",
         "k_proj",
@@ -61,6 +62,7 @@ class ModelConfig:
 @dataclass
 class TrainingConfig:
     """Training hyperparameters."""
+
     # GPU parameters
     gpu_type: str = "A100"
     num_gpus: int = 2
@@ -71,9 +73,9 @@ class TrainingConfig:
     save_every_batches: int = 1
 
     # Batching parameters
-    epochs: int = 2
-    batch_size: int = 24
-    group_size: int = 4
+    epochs: int = 1
+    batch_size: int = 16
+    group_size: int = 8
     mini_batch_size: int = 1
     update_after_group: bool = True  # Whether to update the policy after each task group
     accumulate_over_minibatches: bool = False  # Whether to accumulate over minibatches
@@ -84,7 +86,7 @@ class TrainingConfig:
     leave_one_out: bool = True
 
     # Replay buffer parameters
-    buffer_steps: int = 4
+    buffer_steps: int = 8
     select_strategy: Literal["recent", "variance", "random"] = "variance"
 
     # Aggregation parameters
@@ -92,8 +94,8 @@ class TrainingConfig:
     token_agg: Literal["mean", "sum"] = "mean"  # noqa: S105
 
     # Regularization parameters
-    kl_beta: float = 0.0
-    entropy_beta: float = 0.0
+    kl_beta: float = 0.001
+    entropy_beta: float = 0.001
     top_eps: float = 0.2
     bottom_eps: float = 0.1
 
@@ -143,6 +145,7 @@ class Config:
     job_id: str | None = None  # Use existing job ID if provided
     stats_interval: int = 1
     verbose: bool = False
+    very_verbose: bool = False
 
     # Paths
     out_dir: str = "./checkpoints"
@@ -166,6 +169,7 @@ class Config:
             job_id=d.get("job_id"),
             stats_interval=d.get("stats_interval", 1),
             verbose=d.get("verbose", False),
+            very_verbose=d.get("very_verbose", False),
             out_dir=d.get("out_dir", "./checkpoints"),
             adapter_prefix=d.get("adapter_prefix", "cua-grpo-step"),
             seed=d.get("seed", 1234),
@@ -181,6 +185,7 @@ class Config:
             "job_id": self.job_id,
             "stats_interval": self.stats_interval,
             "verbose": self.verbose,
+            "very_verbose": self.very_verbose,
             "out_dir": self.out_dir,
             "adapter_prefix": self.adapter_prefix,
             "seed": self.seed,

hud/rl/distributed.py

@@ -66,7 +66,13 @@ def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor:
 
 
 def broadcast_object(obj: Any, src: int = 0) -> Any:
-    """Broadcast a Python object from src rank to all ranks."""
+    """Broadcast a Python object from src rank to all ranks.
+
+    Args:
+        obj: Object to broadcast (used on src rank)
+        src: Source rank
+        device: Device for temporary tensor buffer during pickling transfer
+    """
     if not dist.is_initialized():
         return obj
 
@@ -75,6 +81,33 @@ def broadcast_object(obj: Any, src: int = 0) -> Any:
     return obj_list[0]
 
 
+def scatter_object(
+    obj_list: list[Any] | None,
+    src: int = 0,
+) -> Any:
+    """Scatter a list of Python objects from src so each rank receives one object.
+
+    Usage:
+        - On src rank: pass the full list (length == world_size)
+        - On non-src ranks: pass None
+
+    Returns:
+        The object intended for this rank.
+    """
+    if not dist.is_initialized():
+        # Single-process: return first element if provided, else None
+        if obj_list is None or len(obj_list) == 0:
+            return None
+        return obj_list[0]
+
+    out: list[Any] = [None]
+    if dist.get_rank() == src:
+        dist.scatter_object_list(out, obj_list, src=src)
+    else:
+        dist.scatter_object_list(out, None, src=src)
+    return out[0]
+
+
 def gather_tensors(tensor: torch.Tensor) -> list[torch.Tensor] | None:
     """Gather tensors from all ranks to rank 0.
 

hud/rl/learner.py

@@ -240,6 +240,8 @@ class GRPOLearner:
                 if sample.inputs:
                     sample = sample.to_device(self.device)
                     sample.old_logprobs, _ = self.compute_logprobs(self.policy, sample.inputs)
+                    # Free GPU memory for this sample immediately
+                    sample.to_device(torch.device("cpu"))
 
             policy_module = self.policy.module if hasattr(self.policy, "module") else self.policy
             with policy_module.disable_adapter():
@@ -247,7 +249,10 @@ class GRPOLearner:
                     if is_main_process():
                         progress.update(f"Processing batch of traces... {i}/{len(batch)}")
                     if sample.inputs:
+                        # Move back to GPU for reference computation, then free
+                        sample = sample.to_device(self.device)
                         sample.ref_logprobs, _ = self.compute_logprobs(self.policy, sample.inputs)
+                        sample.to_device(torch.device("cpu"))
 
         hud_console.info_log("Creating mini-batches...")
         group_size = self.config.training.group_size
@@ -488,15 +493,21 @@ class GRPOLearner:
             out = model(**model_inputs)
 
             logits = out.logits / self.config.actor.temperature
-            log_probs = F.log_softmax(logits, dim=-1)
 
+            # Compute token log-probs via negative cross-entropy to avoid materializing full log_probs
             targets = inputs["input_ids"][:, 1:]
-            token_log_probs = log_probs[:, :-1].gather(-1, targets.unsqueeze(-1)).squeeze(-1)
+            logits_slice = logits[:, :-1, :]
+            loss_flat = F.cross_entropy(
+                logits_slice.reshape(-1, logits_slice.size(-1)),
+                targets.reshape(-1),
+                reduction="none",
+            )
+            token_log_probs = (-loss_flat).reshape_as(targets)
 
             # Compute entropy only for assistant tokens to save memory
             assistant_mask = inputs["assistant_mask"]
             entropy = torch.zeros_like(token_log_probs)
-            if assistant_mask.any():
+            if assistant_mask.any() and getattr(self.config.training, "entropy_beta", 0.0) != 0.0:
                 entropy[assistant_mask] = entropy_from_logits(logits[:, :-1][assistant_mask])
 
             return token_log_probs, entropy
@@ -506,8 +517,20 @@ class GRPOLearner:
             # Return dummy values that match expected shapes
             seq_len = inputs["input_ids"].shape[1] - 1 if "input_ids" in inputs else 0
             batch_size = inputs["input_ids"].shape[0] if "input_ids" in inputs else 1
-            dummy_logprobs = torch.zeros(batch_size, seq_len, device=self.device)
-            dummy_entropy = torch.zeros(batch_size, seq_len, device=self.device)
+            # Create dummy tensors that still participate in autograd so backward doesn't fail
+            try:
+                param_sum = torch.sum(
+                    next(self.policy.parameters())
+                )  # touch params to build a graph
+                base = param_sum * 0.0
+            except StopIteration:
+                base = torch.tensor(0.0, device=self.device)
+            dummy_logprobs = (
+                base + torch.zeros(batch_size, seq_len, device=self.device)
+            ).requires_grad_(True)
+            dummy_entropy = (
+                base + torch.zeros(batch_size, seq_len, device=self.device)
+            ).requires_grad_(True)
             return dummy_logprobs, dummy_entropy
 
     def save(self, path: str) -> None:

hud/rl/train.py

@@ -13,7 +13,7 @@ import json
 import logging
 from datetime import datetime
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import hud
 from hud.rl.actor import Actor
@@ -25,6 +25,7 @@ from hud.rl.distributed import (
     get_global_rank,
     get_world_size,
     is_main_process,
+    scatter_object,
     setup_distributed,
     synchronize,
 )
@@ -56,6 +57,10 @@ async def train(config: Config, tasks: list[Task]) -> None:
         logging.basicConfig(level=logging.INFO)
         # Remove httpx logger
         logging.getLogger("httpx").setLevel(logging.WARNING)
+    if config.very_verbose:
+        logging.basicConfig(level=logging.DEBUG)
+        # Remove httpx logger
+        logging.getLogger("httpx").setLevel(logging.INFO)
 
     if is_main_process():
         hud_console.header("Starting GRPO Training")
@@ -103,10 +108,9 @@ async def train(config: Config, tasks: list[Task]) -> None:
     if is_main_process():
         hud_console.info(f"Creating job with config.job_id: {config.job_id}")
         job_obj = hud.create_job(
-            job_id=config.job_id, name=config.job_name, metadata={
-                "config": config.to_dict(),
-                "agent_class": config.model.base_model
-            }
+            job_id=config.job_id,
+            name=config.job_name,
+            metadata={"config": config.to_dict(), "agent_class": config.model.base_model},
         )
         hud_console.info(f"Created job with job_obj.id: {job_obj.id}")
         job_obj.update_status_sync("running")
@@ -130,53 +134,71 @@ async def train(config: Config, tasks: list[Task]) -> None:
             global_reward_stats = None
             global_advantage_stats = None
 
-            # Only rank 0 runs tasks and collects traces
+            # Step-state gate: ensure all ranks branch coherently
+            state = {"ok": False, "err": None, "num_samples": 0}
+            rank_samples = None
+            episode_time_value = None
+
+            # Only rank 0 runs tasks and prepares distribution
             if is_main_process() and actor is not None:
                 import time
 
-                episode_start_time = time.time()
-                traces = await actor.run_tasks(tasks, job_id=job_id)
-                episode_time = time.time() - episode_start_time
-                hud_console.info(f"Sampled {len(traces)} traces in {episode_time:.1f}s")
-                trace_buffer.add(traces)
-                global_reward_stats = [trace.reward for trace in traces]
-
-                # Get all traces from buffer for distribution
-                all_traces = trace_buffer.sample_traces()
-
-                assert len(traces) == len(all_traces)  # noqa: S101
-
-                # Preprocess traces to training samples
-                preprocessed_traces = preprocess_advantages(all_traces, config)
-
-                # Store these for later use in metrics
-                global_advantage_stats = [sample.advantage for sample in preprocessed_traces]
-
-                # Distribute preprocessed samples in groups across ranks
-                gpu_batch_size = len(preprocessed_traces) // num_gpus
-                rank_samples = [
-                    preprocessed_traces[i : i + gpu_batch_size]
-                    for i in range(0, len(preprocessed_traces), gpu_batch_size)
-                ]
+                try:
+                    episode_start_time = time.time()
+                    traces = await actor.run_tasks(tasks, job_id=job_id)
+                    episode_time = time.time() - episode_start_time
+                    hud_console.info(f"Sampled {len(traces)} traces in {episode_time:.1f}s")
+                    trace_buffer.add(traces)
+                    global_reward_stats = [trace.reward for trace in traces]
+
+                    # Get all traces from buffer for distribution
+                    all_traces = trace_buffer.sample_traces()
+
+                    # Preprocess traces to training samples
+                    preprocessed_traces = preprocess_advantages(all_traces, config)
+
+                    # Store these for later use in metrics
+                    global_advantage_stats = [sample.advantage for sample in preprocessed_traces]
+
+                    # Distribute preprocessed samples in groups across ranks via scatter
+                    # Ensure list length is a multiple of num_gpus by allowing empty per-rank slices
+                    gpu_batch_size = max(1, (len(preprocessed_traces) + num_gpus - 1) // num_gpus)
+                    rank_samples = [
+                        preprocessed_traces[i : i + gpu_batch_size]
+                        for i in range(0, len(preprocessed_traces), gpu_batch_size)
+                    ]
+                    # Pad rank_samples to exactly num_gpus entries
+                    if len(rank_samples) < num_gpus:
+                        rank_samples.extend([[] for _ in range(num_gpus - len(rank_samples))])
+
+                    # Log distribution info
+                    dist_msg = (
+                        f"Distributing {len(preprocessed_traces)} samples as {gpu_batch_size} "
+                        f"sized batches across {num_gpus} GPUs"
+                    )
+                    hud_console.info(dist_msg)
+                    for rank in range(num_gpus):
+                        n_samples = len(rank_samples[rank]) if rank < len(rank_samples) else 0
+                        hud_console.info(f"  Rank {rank}: {n_samples} samples")
+
+                    hud_console.section_title(f"Training on {len(all_traces)} traces")
+                    episode_time_value = episode_time
+
+                    state.update({"ok": True, "num_samples": len(preprocessed_traces)})
+                except Exception as e:
+                    state.update({"ok": False, "err": str(e)})
+
+            # Broadcast step-state to keep ranks in lockstep
+            state = broadcast_object(state, src=0)
+            if not state.get("ok", False):
+                hud_console.warning("Step failed on rank 0; skipping this step coherently")
+                synchronize()
+                continue
 
-                # Log distribution info
-                hud_console.info(
-                    f"Distributing {len(preprocessed_traces)} samples as {gpu_batch_size} sized batches across {num_gpus} GPUs"  # noqa: E501
-                )
-                for rank in range(num_gpus):
-                    n_samples = len(rank_samples[rank])
-                    hud_console.info(f"  Rank {rank}: {n_samples} samples")
-
-                hud_console.section_title(f"Training on {len(all_traces)} traces")
-                episode_time_value = episode_time
-            else:
-                rank_samples = None
-                episode_time_value = None
-
-            # Broadcast each rank's samples and episode time
-            rank_samples = broadcast_object(rank_samples, src=0)
+            # Scatter per-rank samples; each rank receives only its slice
+            my_samples = scatter_object(rank_samples if is_main_process() else None, src=0)
+            # Broadcast the episode time (small object)
             episode_time_value = broadcast_object(episode_time_value, src=0)
-            my_samples = rank_samples[get_global_rank()] if rank_samples else []
 
             # Process only assigned samples
             last_metrics = learner.update(my_samples)
@@ -299,7 +321,7 @@ async def main() -> None:
 
     # Load config
     if args.config:
-        with open(args.config) as f:  # noqa: ASYNC230
+        with open(args.config, encoding="utf-8") as f:  # noqa: ASYNC230
             config_dict = json.load(f)
         config = Config.from_dict(config_dict)
     else:
@@ -337,7 +359,7 @@ async def main() -> None:
     # Load tasks
     if args.tasks_json:
         # Tasks provided as JSON list via command line
-        tasks = load_tasks(args.tasks_jso)
+        tasks = load_tasks(args.tasks_json)
     elif args.tasks:
         # Tasks provided as file path or HuggingFace dataset
         tasks = load_tasks(args.tasks)
@@ -353,7 +375,8 @@ async def main() -> None:
             )
 
     # Run training
-    await train(config, tasks)
+    tasks_typed = cast("list[Task]", tasks)
+    await train(config, tasks_typed)
 
 
 if __name__ == "__main__":

hud/utils/group_eval.py

@@ -189,7 +189,7 @@ def display_group_statistics(stats: list[dict[str, Any]], show_details: bool = T
     hud_console.info(f"Overall mean reward: {overall_mean:.3f} ± {overall_std:.3f}")
 
     # Detailed table
-    if show_details and len(stats) <= 20:  # Only show for reasonable dataset sizes
+    if show_details and len(stats) <= 50:  # Only show for reasonable dataset sizes
         table = Table(title="\nPer-Task Performance Distribution")
         table.add_column("Task", style="cyan", no_wrap=True)
         table.add_column("Mean±Std", justify="right", style="green")
@@ -216,7 +216,7 @@ def display_group_statistics(stats: list[dict[str, Any]], show_details: bool = T
     # High variance tasks
     high_variance_tasks = [s for s in stats if s["std_reward"] > 0.3 and s["group_size"] > 1]
     if high_variance_tasks:
-        hud_console.warning(f"\n⚠️  {len(high_variance_tasks)} tasks show high variance (std > 0.3)")
+        hud_console.warning(f"\n{len(high_variance_tasks)} tasks show high variance (std > 0.3)")
         for task in high_variance_tasks[:3]:
             hud_console.info(
                 f"  • {task['task_id']}: μ={task['mean_reward']:.3f}, σ={task['std_reward']:.3f}"  # noqa: RUF001

hud/utils/tasks.py

@@ -40,7 +40,7 @@ def load_tasks(tasks_input: str | list[dict], *, raw: bool = False) -> list[Task
         if Path(tasks_input).exists():
             file_path = Path(tasks_input)
 
-            with open(file_path) as f:
+            with open(file_path, encoding="utf-8") as f:
                 # Handle JSON files (array of tasks)
                 if file_path.suffix.lower() == ".json":
                     data = json.load(f)

hud/utils/tests/test_version.py

@@ -5,4 +5,4 @@ def test_import():
     """Test that the package can be imported."""
     import hud
 
-    assert hud.__version__ == "0.4.42"
+    assert hud.__version__ == "0.4.44"

hud/version.py

@@ -4,4 +4,4 @@ Version information for the HUD SDK.
 
 from __future__ import annotations
 
-__version__ = "0.4.42"
+__version__ = "0.4.44"

{hud_python-0.4.42.dist-info → hud_python-0.4.44.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hud-python
-Version: 0.4.42
+Version: 0.4.44
 Summary: SDK for the HUD platform.
 Project-URL: Homepage, https://github.com/hud-evals/hud-python
 Project-URL: Bug Tracker, https://github.com/hud-evals/hud-python/issues

{hud_python-0.4.42.dist-info → hud_python-0.4.44.dist-info}/RECORD

@@ -2,7 +2,7 @@ hud/__init__.py,sha256=JMDFUE1pP0J1Xl_miBdt7ERvoffZmTzSFe8yxz512A8,552
 hud/__main__.py,sha256=YR8Dq8OhINOsVfQ55PmRXXg4fEK84Rt_-rMtJ5rvhWo,145
 hud/settings.py,sha256=disObWa-DgXzoDcCDp3y1dTPaNsbR0IvoMJL9Eg4zyo,3947
 hud/types.py,sha256=pmPj_8emfMIfEY_fRS8NgIJ56kCsolWSqQjyCzXDaGY,11072
-hud/version.py,sha256=cq2TbXqf0IU0JmINUgqBhq-0yqdfvngAxHjSWlTSBns,105
+hud/version.py,sha256=j-0v9E6ZVwBdP3D1A-70Ie5rXP137HYVUJCZeIwO3_0,105
 hud/agents/__init__.py,sha256=UoIkljWdbq4bM0LD-mSaw6w826EqdEjOk7r6glNYwYQ,286
 hud/agents/base.py,sha256=_u1zR3gXzZ1RlTCUYdMcvgHqdJBC4-AB1lZt0yBx8lg,35406
 hud/agents/claude.py,sha256=TGhm5gE2ltINDAdEsDxKuT9iGMQ5G87R6kmabU3KPt8,16101
@@ -10,7 +10,7 @@ hud/agents/grounded_openai.py,sha256=U-FHjB2Nh1_o0gmlxY5F17lWJ3oHsNRIB2a7z-IKB64
 hud/agents/langchain.py,sha256=1EgCy8jfjunsWxlPC5XfvfLS6_XZVrIF1ZjtHcrvhYw,9584
 hud/agents/lite_llm.py,sha256=_3wbUiYCp7q8Vyu9rhaoJDvmb_bsyUsLYWP3iQJ2bHo,2239
 hud/agents/openai.py,sha256=O1xV1h1l-W8lmnmXqTYr5CwnmnaniMqOxAZbl2CTTng,14576
-hud/agents/openai_chat_generic.py,sha256=RUfXDZSUbOXQ2leR4_8PGpqvUzz5PJOWeR3PTticKUY,12150
+hud/agents/openai_chat_generic.py,sha256=_vAID9dZ_UxL0elYwafskRcsdrSsLsxJ4zPrP58oBiw,12151
 hud/agents/misc/__init__.py,sha256=BYi4Ytp9b_vycpZFXnr5Oyw6ncKLNNGml8Jrb7bWUb4,136
 hud/agents/misc/response_agent.py,sha256=uMuRDkz5QgaMQliNzBRepond5sb7KyqIiKm3LstjVnw,3753
 hud/agents/tests/__init__.py,sha256=W-O-_4i34d9TTyEHV-O_q1Ai1gLhzwDaaPo02_TWQIY,34
@@ -19,14 +19,14 @@ hud/agents/tests/test_claude.py,sha256=0nZnfsbGoECvsLPdmaRnc9jVmrehVvc3kxeyiCQI2
 hud/agents/tests/test_client.py,sha256=uikgh6yhjPPX2RBU4XJQMz1mNox9uXjuwsP8t93id18,13337
 hud/agents/tests/test_grounded_openai_agent.py,sha256=VK8lUvHIjWicMX00VKPE-FZyjiJqTEhb80MuRRa9fVc,5437
 hud/agents/tests/test_openai.py,sha256=Npbdr0acgLExGLbrleXze-k3w9LHfmqzQjPk9TnjN68,7620
-hud/cli/__init__.py,sha256=v4602N3FWOXD2raK1h3APJXFshsCrJV8P3JFcK--IvM,45596
+hud/cli/__init__.py,sha256=lwyaA7z7H4BOt9ksySpT0AnRERoYEiVgUdwV_5s9wIg,45768
 hud/cli/__main__.py,sha256=fDH7XITyuDITwSDIVwRso06aouADO0CzTHKqp5TOwJE,143
 hud/cli/analyze.py,sha256=4u5oYfJMquOjT9PzzRTYVcTZDxDi0ilNP_g532_hpOU,14716
 hud/cli/build.py,sha256=h-4SAoe3j8Pth3mPYf26vh7q1Do5JADlvKKwkZrf2AU,19551
 hud/cli/clone.py,sha256=AwVDIuhr8mHb1oT2Af2HrD25SiTdwATpE6zd93vzLgA,6099
 hud/cli/debug.py,sha256=jtFW8J5F_3rhq1Hf1_SkJ7aLS3wjnyIs_LsC8k5cnzc,14200
-hud/cli/dev.py,sha256=J0Q_ndHbQcXe64gMjXfqiccWYWpdiYWvTKbJhCAvlgI,30666
-hud/cli/eval.py,sha256=d1RouB3rxP3axca2sRblNWZMNvHGP1EugST5fCJ-7tc,25790
+hud/cli/dev.py,sha256=2zUeVz5S__WrV-DLSDqOlQawcJS7eYPKiDRVUaJ8mAk,31579
+hud/cli/eval.py,sha256=zoRC9ExxrsOEj3myTUz_72LVSnFF557lS1aJfhQ9kHg,25681
 hud/cli/get.py,sha256=sksKrdzBGZa7ZuSoQkc0haj-CvOGVSSikoVXeaUd3N4,6274
 hud/cli/init.py,sha256=YkWxkIDCnhnxGGpbm7IvYMcfDqWuO1X9wxDxE4k-9ew,9721
 hud/cli/list_func.py,sha256=EVi2Vc3Lb3glBNJxFx4MPnZknZ4xmuJz1OFg_dc8a_E,7177
@@ -34,17 +34,17 @@ hud/cli/pull.py,sha256=XGEZ8n60tbzLQP_8d9h7XYmzyCW0e2-Rkr3_tLG7jvw,12449
 hud/cli/push.py,sha256=DsXFrMtWBZ-HUxt6VoLihpklk8JJIe2gy-GA4AMg6Kw,18805
 hud/cli/remove.py,sha256=8vGQyXDqgtjz85_vtusoIG8zurH4RHz6z8UMevQRYM4,6861
 hud/cli/flows/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-hud/cli/flows/tasks.py,sha256=d-RsIV0rvzBuYN6WR7AOId19wzEy0z7SSW9c5RnEaqQ,14564
-hud/cli/rl/__init__.py,sha256=q0GIYRARpNkY8o1Sza5CjQq6cks_1W_SjGprybaCxq4,5193
+hud/cli/flows/tasks.py,sha256=z3qdn6MBLJRK7dvfp3IFkdCzuImKRps7jbcedmCTjDA,14485
+hud/cli/rl/__init__.py,sha256=pGx4WGaL-yGdogJNzgEE7BtjFtT4I9CTI_UdCm49h98,5376
 hud/cli/rl/celebrate.py,sha256=trGEJn3xebexlHwFVKPJKhRujVVV8sy7TQTJvRd2p9A,5947
-hud/cli/rl/config.py,sha256=VZ8fiOI22Aw6YTRk7gj1ozpF-TU7NK8QWQgWFwMbNs0,3235
+hud/cli/rl/config.py,sha256=A-4WWwAS68GRKx1cP_DJ-NZD_96cFNnGwx0P3pQT1ps,3271
 hud/cli/rl/display.py,sha256=hqJVGmO9csYinladhZwjF-GMvppYWngxDHajTyIJ_gM,5214
 hud/cli/rl/gpu.py,sha256=peXS-NdUF5RyuSs0aZoCzGLboneBUpCy8f9f99WMrG0,2009
-hud/cli/rl/gpu_utils.py,sha256=f7olAhIxLmZmvII9igwtqnjIRPenY9DP6m70hp4lP0U,11113
+hud/cli/rl/gpu_utils.py,sha256=VSdEWJDH-P9LjRZscQXPju5vB3FomP4Iy2znPcpUZc4,11199
 hud/cli/rl/local_runner.py,sha256=NFsNmRZ4nenPnb45ZtdsILeICKEq11wmpLwq9E-a8ZE,22614
 hud/cli/rl/presets.py,sha256=DzOO82xL5QyzdVtlX-Do1CODMvDz9ILMPapjU92jcZg,3051
-hud/cli/rl/remote_runner.py,sha256=AXnmFGDqo4NcCpJ2RE-Q1HaKzIYlqFP1TWU4tQC_DFk,17500
-hud/cli/rl/rl_api.py,sha256=INJobvSa50ccR037u_GPsDa_9WboWyNwqEaoh9hcXj0,4306
+hud/cli/rl/remote_runner.py,sha256=fKmOVKSBUWfakunfe9-HAllpUJDxfRNZwL00fPw-QTI,17837
+hud/cli/rl/rl_api.py,sha256=fvRMWQXhTSLM2zQaWWDas_u47RreH8erLgtXRKnQGeA,4350
 hud/cli/rl/viewer.py,sha256=ExQs1IX3T8x_9aBzc4JojZ779jmFvFTh7EjOYIHzYsU,4441
 hud/cli/rl/vllm.py,sha256=Gq_M6KsQArGz7FNIdemuM5mk16mu3xe8abpO2GCCuOE,6093
 hud/cli/rl/wait_utils.py,sha256=FyIvqYWLOydANTetukoE5Rp2AOQi67qkiAlIJp4HpL8,2577
@@ -69,7 +69,7 @@ hud/cli/utils/config.py,sha256=AnsN6FEa8V3jg3EWaqUJN38-UuYC6tVZxPfBb_5LFBs,2652
 hud/cli/utils/cursor.py,sha256=fy850p0rVp5k_1wwOCI7rK1SggbselJrywFInSQ2gio,3009
 hud/cli/utils/docker.py,sha256=oGVzPfp0Rn89o9d6tgSEziKy9GXFrYaWn_mjBmGRHe4,6326
 hud/cli/utils/env_check.py,sha256=TqsmwgTfMDzfP0Ii50YxDkOP4_T5nqks9JMTxIq60-s,7095
-hud/cli/utils/environment.py,sha256=EfATQyAz8Jybj4N9QNaaADUrpiZ2JMp2elQYnAG9gU8,4371
+hud/cli/utils/environment.py,sha256=cxsNwCfwX2PtCHht9xH_Yo5jpcqANf7h0wa3gfiy5tY,4278
 hud/cli/utils/interactive.py,sha256=sHhTjaImxlwlZ5_DTXb23Jwrjy5oJ7diB-8duhHbImU,16647
 hud/cli/utils/local_runner.py,sha256=jnPFoJu3sCq65LSUapKCkakdlEuz__96oJU_FfOYtEg,6542
 hud/cli/utils/logging.py,sha256=DyOWuzZUg6HeKCqfs6ufb703XS3bW4G2pzaXVAvDqvA,9018
@@ -120,10 +120,10 @@ hud/rl/__init__.py,sha256=yYL7U1WV6L3mr3Hig48-4lhnryTaWj4nCXm4lG5vrYI,25
 hud/rl/actor.py,sha256=H6gwRGRY1YpkOyiaJ9yai8yQwcI-Gx0dFxd18jpLx_Q,6950
 hud/rl/buffer.py,sha256=z47HOjOBJx3umUzzUfdtq_N4ZoJ8FMBPkX8YQKBtd3A,15457
 hud/rl/chat_template.jinja,sha256=XTdzI8oFGEcSA-exKxyHaprwRDmX5Am1KEb0VxvUc6U,4965
-hud/rl/config.py,sha256=0ejjbIMXqxjhz3VnIHX0QFtZ_0KmXNynTuHdg-i0M5o,5454
-hud/rl/distributed.py,sha256=8avhrb0lHYkhW22Z7MfkqSnlczWj5jMrUMEtkcoCf74,2473
-hud/rl/learner.py,sha256=FKIgIIghsNiDr_g090xokOO_BxNmTSj1O-TSJzIq_Uw,24703
-hud/rl/train.py,sha256=i8_HocC92yC-wP_ZeQYUXLhJKNvsfX8mdnhO8gOf-8c,13641
+hud/rl/config.py,sha256=akQ2a53NX3Dh1UWgMyw7mTxq33eiQbZcBpmKTzd79Xk,5624
+hud/rl/distributed.py,sha256=ZIh5GTMuRl_tHV_62iWsYgrV--AylBelp_TZQnhwfy4,3391
+hud/rl/learner.py,sha256=GowGqhWyCMPfrxD9V3KyOdqF0FDeUMUSCA0QPnE1RWE,25855
+hud/rl/train.py,sha256=zO5TVvGWQdYfdhSCOSMaahfBVwcWb0Fxa80LiInx01c,15005
 hud/rl/types.py,sha256=lrLKo7iaqodYth2EyeuOQfLiuzXfYM2eJjPmpObrD7c,3965
 hud/rl/utils.py,sha256=IsgVUUibxnUzb32a4mu1sYrgJC1CwoG9E-Dd5y5VDOA,19115
 hud/rl/vllm_adapter.py,sha256=2wnTfoXPI4C9EzhVxk0GU-ArLjX7hgXS0BndMwN8Ppg,4751
@@ -204,12 +204,12 @@ hud/tools/tests/test_utils.py,sha256=qaujM1uyTMaKqWIeEgxty5GOFyfSUtrYCEHhmIazoy4
 hud/utils/__init__.py,sha256=nk9Re6ls2RudAWnAHDWYbLG28AwNF4qMFYf5xQIJhQA,181
 hud/utils/agent_factories.py,sha256=cvfXByqG6gOYHtm1VGeJjCpxoLxM4aJez8rH-AerP_A,3186
 hud/utils/async_utils.py,sha256=5cKrJcnaHV2eJNxeyx0r7fPcdPTDBK7kM9-nLaF51X4,2409
-hud/utils/group_eval.py,sha256=oaoBqlQN6g5gRQmuY_JmqM5bpuf2sFIgu4uDZ7X-3a0,8360
+hud/utils/group_eval.py,sha256=6yXEH8ZRKkR4bBy9-QWGmjlm2IbCnTUZppEFbjTvndY,8352
 hud/utils/hud_console.py,sha256=ywTrzyNhWFoQN2PpzpDDKp_32b-ACDvfKQuWxDoF8iE,21898
 hud/utils/mcp.py,sha256=pMadd7A0DH6Y_aWywKU8jVYu2pRHGPEndV2ZQFrrj60,2888
 hud/utils/pretty_errors.py,sha256=WGeL4CTHtlA6KgPuV_JSX5l6H4-xbuTp6Y6tw1bkiFg,2430
 hud/utils/progress.py,sha256=suikwFM8sdSfkV10nAOEaInDhG4XKgOSvFePg4jSj1A,5927
-hud/utils/tasks.py,sha256=4v8FvVhXXefbWFldf564XPCQIActYI3tmA510-SA4LE,4967
+hud/utils/tasks.py,sha256=7i36ck84gz1GZxhn9jryMBvKgMmcvLVu1YH5n3Y23-c,4985
 hud/utils/telemetry.py,sha256=hrVIx2rUjSGyy9IVxTZ_3Jii83PiHjyFRd5ls2whimM,1863
 hud/utils/tool_shorthand.py,sha256=_haLgK3yazLR2Y0jlEHUUQjw9uZCxi9yTipAwdOAJ70,2148
 hud/utils/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -218,10 +218,10 @@ hud/utils/tests/test_init.py,sha256=2QLQSGgyP9wJhOvPCusm_zjJad0qApOZi1BXpxcdHXQ,
 hud/utils/tests/test_mcp.py,sha256=0pUa16mL-bqbZDXp5NHBnt1gO5o10BOg7zTMHZ1DNPM,4023
 hud/utils/tests/test_progress.py,sha256=QSF7Kpi03Ff_l3mAeqW9qs1nhK50j9vBiSobZq7T4f4,7394
 hud/utils/tests/test_telemetry.py,sha256=5jl7bEx8C8b-FfFUko5pf4UY-mPOR-9HaeL98dGtVHM,2781
-hud/utils/tests/test_version.py,sha256=nPdE5Z40E9Q4XglmFTXKz9ffcgtlH7iV_yqI267-7ic,160
+hud/utils/tests/test_version.py,sha256=B9UhswFSFbHf544swTgKJdq6TMat27bGIzFb8Sy-bKc,160
 hud/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-hud_python-0.4.42.dist-info/METADATA,sha256=cN0LvsjgX38My66ATDP5l_KqaokO9oRS8vWetmAL_So,22275
-hud_python-0.4.42.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hud_python-0.4.42.dist-info/entry_points.txt,sha256=jJbodNFg1m0-CDofe5AHvB4zKBq7sSdP97-ohaQ3ae4,63
-hud_python-0.4.42.dist-info/licenses/LICENSE,sha256=yIzBheVUf86FC1bztAcr7RYWWNxyd3B-UJQ3uddg1HA,1078
-hud_python-0.4.42.dist-info/RECORD,,
+hud_python-0.4.44.dist-info/METADATA,sha256=bjz1T1aLq3yUaoW_Ih9ZQjGD8X-nKRTYmgeggS568LM,22275
+hud_python-0.4.44.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hud_python-0.4.44.dist-info/entry_points.txt,sha256=jJbodNFg1m0-CDofe5AHvB4zKBq7sSdP97-ohaQ3ae4,63
+hud_python-0.4.44.dist-info/licenses/LICENSE,sha256=yIzBheVUf86FC1bztAcr7RYWWNxyd3B-UJQ3uddg1HA,1078
+hud_python-0.4.44.dist-info/RECORD,,