hud-python 0.2.6__py3-none-any.whl → 0.2.7__py3-none-any.whl

hud/env/remote_client.py

@@ -2,7 +2,9 @@ from __future__ import annotations
 
 import logging
 from base64 import b64decode
-from typing import TYPE_CHECKING, Any
+from typing import Any
+
+from pydantic import BaseModel
 
 from hud.env.client import Client
 from hud.exceptions import HudResponseError
@@ -10,13 +12,18 @@ from hud.server import make_request
 from hud.settings import settings
 from hud.types import EnvironmentStatus
 from hud.utils import ExecuteResult
-
-if TYPE_CHECKING:
-    from hud.utils.config import FunctionConfig
+from hud.utils.config import FunctionConfig
 
 logger = logging.getLogger("hud.env.remote_env_client")
 
 
+class SetupRequest(BaseModel):
+    task_id: str | None = None
+    setup: FunctionConfig | None = None
+    config: dict[str, Any] | None = None
+    metadata: dict[str, Any] | None = None
+
+
 class RemoteClient(Client):
     """
     Remote environment client implementation.
@@ -183,6 +190,17 @@ class RemoteClient(Client):
 
         return data["result"], b64decode(data["stdout"]), b64decode(data["stderr"])
 
+    async def setup(self, setup_request: SetupRequest) -> dict[str, Any]:
+        """
+        Setup the environment.
+        """
+        return await make_request(
+            method="POST",
+            url=f"{settings.base_url}/v1/environments/{self.env_id}/reset",
+            json=setup_request.model_dump(),
+            api_key=settings.api_key,
+        )
+
     async def close(self) -> None:
         """
         Close the remote environment by making a request to the server.

hud/env/remote_docker_client.py

@@ -20,10 +20,14 @@ if TYPE_CHECKING:
 logger = logging.getLogger("hud.env.remote_env_client")
 
 
-async def upload_bytes_to_presigned_url(presigned_url: str, data_bytes: bytes) -> None:
+async def upload_bytes_to_presigned_url(
+    presigned_url: str,
+    data_bytes: bytes,
+    timeout: float = 600,
+) -> None:
     try:
         async with httpx.AsyncClient() as client:
-            response = await client.put(presigned_url, content=data_bytes)
+            response = await client.put(presigned_url, content=data_bytes, timeout=timeout)
             response.raise_for_status()
     except httpx.HTTPStatusError as e:
         logger.exception("Failed to upload to presigned URL")

hud/gym.py

@@ -9,13 +9,13 @@ from hud.env.local_docker_client import LocalDockerClient
 from hud.env.remote_client import RemoteClient
 from hud.env.remote_docker_client import RemoteDockerClient
 from hud.exceptions import GymMakeException
+from hud.task import Task
 from hud.telemetry.context import get_current_task_run_id
 from hud.types import CustomGym, Gym
 from hud.utils.common import get_gym_id
 
 if TYPE_CHECKING:
     from hud.job import Job
-    from hud.task import Task
 
 logger = logging.getLogger("hud.gym")
 
@@ -39,9 +39,11 @@ async def make(
     task = None
     if isinstance(env_src, str | CustomGym):
         gym = env_src
-    else:
+    elif isinstance(env_src, Task):
         gym = env_src.gym
         task = env_src
+    else:
+        raise GymMakeException(f"Invalid gym source: {env_src}", {})
 
     effective_job_id = None
     if job is not None:
@@ -89,9 +91,18 @@ async def make(
 
             if gym.location == "local":
                 logger.info("Creating local environment")
-                client = await LocalDockerClient.create(uri)
+                if gym.host_config:
+                    logger.info("Using host config: %s", gym.host_config)
+                    client = await LocalDockerClient.create(uri, gym.host_config)
+                else:
+                    client = await LocalDockerClient.create(uri)
+
             elif gym.location == "remote":
                 logger.info("Creating remote environment")
+
+                if gym.host_config:
+                    raise ValueError("host_config is not supported for remote environments")
+
                 client = await RemoteDockerClient.create(
                     image_uri=uri,
                     job_id=effective_job_id,
@@ -105,7 +116,7 @@ async def make(
                 logger.info("Setting source path %s", gym.image_or_build_context)
                 client.set_source_path(gym.image_or_build_context)
         elif isinstance(gym, str):
-            logger.info("Creating private environment")
+            logger.debug("Creating private environment")
             true_gym_id = await get_gym_id(gym)
             client, build_data = await RemoteClient.create(
                 gym_id=true_gym_id,

hud/job.py

@@ -18,12 +18,12 @@ from hud.settings import settings
 from hud.task import Task
 from hud.taskset import TaskSet
 from hud.trajectory import Trajectory
-from hud.utils.common import Observation
 from hud.utils.progress import StepProgressTracker
 
 if TYPE_CHECKING:
     from hud.adapters.common import Adapter
     from hud.agent.base import Agent
+    from hud.utils.common import Observation
 
 logger = logging.getLogger("hud.job")
 
@@ -275,7 +275,7 @@ async def _maybe_resample_action(
                 decision = await response_agent.determine_response(response_text)
                 if decision == "CONTINUE":
                     logger.info("ResponseAgent indicated CONTINUE. Retrying...")
-                    obs = Observation(text="Please continue.")
+                    obs.text = "Please continue."
                     return obs, False
                 elif decision == "CONTINUE":
                     logger.warning("Max continue retries reached. Stopping despite CONTINUE.")
@@ -321,6 +321,12 @@ async def _execute_task(
         if agent_instance is None:
             raise RuntimeError("Agent could not be instantiated")
 
+        agent_name = agent_instance.name
+        logger.info("Using agent: %s", agent_name)
+        if task.metadata is None or not isinstance(task.metadata, dict):
+            task.metadata = {}
+        task.metadata["agent_name"] = agent_name
+
         # Environment creation with semaphore
         if env_creation_semaphore:
             async with env_creation_semaphore:
@@ -328,6 +334,9 @@ async def _execute_task(
         else:
             env = await gym.make(task, job=job)
 
+        if not env:
+            raise ValueError(f"Environment creation failed for task {task_id}")
+
         obs_tuple = await env.reset()
         if obs_tuple is None:
             raise ValueError(f"env.reset() returned None for task {task_id}")
@@ -335,24 +344,45 @@ async def _execute_task(
 
         step_error = None
 
+        resampled_actions = 0
+
         for step in range(max_steps_per_task):
             action, done = (None, False)
             try:
                 # Agent prediction with semaphore
-                if agent_predict_semaphore:
-                    async with agent_predict_semaphore:
+                try:
+                    if agent_predict_semaphore:
+                        async with agent_predict_semaphore:
+                            action, done = await agent_instance.predict(obs)
+                    else:
                         action, done = await agent_instance.predict(obs)
-                else:
-                    action, done = await agent_instance.predict(obs)
+                except Exception as e:
+                    # if agent prediction fails, pass back the error to the agent
+                    logger.exception("[TR: %s] Agent prediction failed: %s", task_id, e)
+                    resampled_actions += 1
+                    if resampled_actions > 5:
+                        logger.warning(
+                            "[TR: %s] Resampled action %d times. Stopping.",
+                            task_id,
+                            resampled_actions,
+                        )
+                        break
+                    continue
 
                 if tracker:
                     tracker.increment_step(task_id)
 
-                if action is None and not done:
-                    done = True
-
-                if done and response_agent:
+                finish = False
+                if done and response_agent and action and len(action) > 0:
                     obs, finish = await _maybe_resample_action(obs, action[-1], response_agent)
+                    resampled_actions += 1
+                    if resampled_actions > 5:
+                        logger.warning(
+                            "[TR: %s] Resampled action %d times. Stopping.",
+                            task_id,
+                            resampled_actions,
+                        )
+                        break
                     if not finish:
                         continue
 
@@ -361,14 +391,12 @@ async def _execute_task(
                     terminated = True
                 else:
                     obs, _, terminated, _ = step_result
-                if terminated or done:
+                if terminated or done or finish:
                     break
 
             except Exception as agent_step_err:
                 logger.exception(
-                    "[Job: %s/%s, Task: %s] Step %d Error: %s",
-                    job.name,
-                    job.id,
+                    "[TR: %s] Step %d Error: %s",
                     task_id,
                     step + 1,
                     agent_step_err,
@@ -386,7 +414,7 @@ async def _execute_task(
                 )
                 continue
         else:
-            logger.warning("[Job: %s/%s, Task: %s] Max steps reached.", job.name, job.id, task_id)
+            logger.warning("[TR: %s] Max steps reached.", task_id)
 
         # --- Evaluate Task ---
         evaluation_result = None
@@ -401,9 +429,7 @@ async def _execute_task(
                 # logger.info("Evaluation result: %s", evaluation_result)
             except Exception as eval_err:
                 logger.exception(
-                    "[Job: %s/%s, Task: %s] Evaluation Error: %s",
-                    job.name,
-                    job.id,
+                    "[TR: %s] Evaluation Error: %s",
                     task_id,
                     eval_err,
                 )
@@ -420,7 +446,7 @@ async def _execute_task(
                 )
 
     except Exception as e:
-        logger.exception("[Job: %s/%s, Task: %s] Setup/Run Error: %s", job.name, job.id, task_id, e)
+        logger.exception("[TR: %s] Setup/Run Error: %s", task_id, e)
         status = "error"
         error_msg = str(e)
         # Store setup/initialization error in job
@@ -440,9 +466,7 @@ async def _execute_task(
             try:
                 await env.close()
             except Exception as close_err:
-                logger.exception(
-                    "[Job: %s/%s, Task: %s] Close Error: %s", job.name, job.id, task_id, close_err
-                )
+                logger.exception("[TR: %s] Close Error: %s", task_id, close_err)
                 # Store environment close error in job
                 job.errors.append(
                     {
@@ -455,9 +479,7 @@ async def _execute_task(
 
     log_suffix = f" Error: {error_msg}" if status == "error" else f" Eval: {evaluation_result}"
     logger.info(
-        "[Job: %s/%s, Task: %s] Finished local execution. Status: %s.%s",
-        job.name,
-        job.id,
+        "[TR: %s] Finished local execution. Status: %s.%s",
         task_id,
         status,
         log_suffix,
@@ -499,6 +521,7 @@ async def run_job(
     run_parallel: bool = True,
     job_metadata: dict[str, Any] | None = None,
     show_progress: bool = True,
+    verbose: bool = False,
     # Concurrency control with semaphores
     max_concurrent_env_creations: int | None = 30,  # Limits gym.make calls
     max_concurrent_agent_predictions: int | None = None,  # No limit on LLM calls
@@ -538,10 +561,16 @@ async def run_job(
     tasks_to_run: list[Task] = []
     created_job: Job | None = None
 
+    # Get hud logger
+    if not verbose:
+        logger = logging.getLogger("hud")
+        logger.setLevel(logging.CRITICAL)
+    logger = logging.getLogger("hud.job")
+
     evalset_id = None
     if isinstance(task_or_taskset, TaskSet):
         evalset_id = task_or_taskset.id
-        await task_or_taskset.fit(agent_cls)
+        task_or_taskset.fit(agent_cls)
 
     gym_id = None
     if isinstance(task_or_taskset, Task):
@@ -706,3 +735,39 @@ async def run_job(
         num_tasks,
     )
     return created_job
+
+
+"""
+c7f85f7d-3730-4c9a-85a3-a1dc436c3bd2
+
+
+de12c3cc-9d9c-4e90-82cc-1d71d30ede54
+59104743-0a63-4569-a8b5-1eda1a1b55ac
+ff759429-056c-4cde-8851-11e26729ff03
+
+
+7b98ea22-e243-4eeb-a6db-79f4a76da2b3
+
+7aad3f7b-d74f-470d-826d-d817f95fdd67
+
+e356ede6-074a-49ef-9fcd-69e5bcfbdec9
+
+26cd1192-3991-4d1b-b599-b2bed1bcb606
+
+31ece277-970f-4763-b0c8-bf19a56f56c7
+
+
+f9b722a0-5f33-466b-bce0-8ece101f2bc6
+33d1af33-8952-4945-b901-229bcfd88354
+
+6c3d6557-e745-44ab-bc10-300180a81c79
+6c3d6557-e745-44ab-bc10-300180a81c79
+502e02b5-9939-4e57-91af-4fcbcb90a979
+
+7aad3f7b-d74f-470d-826d-d817f95fdd67
+
+
+31ece277-970f-4763-b0c8-bf19a56f56c7
+
+
+e356ede6-074a-49ef-9fcd-69e5bcfbdec9"""

hud/settings.py

@@ -44,6 +44,12 @@ class Settings(BaseSettings):
         validation_alias="TELEMETRY_ENABLED",
     )
 
+    fancy_logging: bool = Field(
+        default=True,
+        description="Enable fancy logging for the HUD SDK",
+        validation_alias="FANCY_LOGGING",
+    )
+
 
 # Create a singleton instance
 settings = Settings()

hud/task.py

@@ -2,12 +2,12 @@ from __future__ import annotations
 
 import tempfile
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 from inspect_ai.util._sandbox import SandboxEnvironmentSpec
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
-from hud.types import CustomGym, Gym
+from hud.types import CustomGym, Gym, MetadataKeys, SensitiveData
 from hud.utils.common import FunctionConfig, FunctionConfigs
 
 if TYPE_CHECKING:
@@ -40,28 +40,78 @@ class Task(BaseModel):
     Attributes:
         id: The remote task ID (optional if local-only)
         prompt: The task prompt or instruction
+        system_prompt: The system prompt for the evalset (optional)
         setup: Environment setup configuration (optional)
         evaluate: Configuration for evaluating responses
         metadata: Additional task metadata
+        sensitive_data: Sensitive data such as API keys, passwords, etc.
         choices: Multiple choice answer list (for Inspect compatibility)
         target: Ideal target output (for Inspect compatibility)
         files: Files that go along with the task (for Inspect compatibility)
         gym: Environment specification
     """
 
-    id: str | None = None
-    prompt: str
+    id: str | None = None  # Remote task ID (optional if local-only)
+
+    prompt: str  # Task prompt or instruction
+    system_prompt: str | None = None  # System prompt for the evalset (optional)
+
+    gym: Gym | None = None  # Environment specification
+
+    # Setup and evaluate configurations for the environment (environment specific)
     setup: FunctionConfigs | None = None
     evaluate: FunctionConfigs | None = None
-    gym: Gym | None = None
+
+    # Overflow configuration for environments that don't conform to the standard
     config: dict[str, Any] | None = None
 
+    # Sensitive data such as API keys, passwords, etc.
+    sensitive_data: SensitiveData = Field(default_factory=dict)
+
+    # Metadata for the task evaluation, information about the agent (see MetadataKeys)
+    metadata: dict[MetadataKeys, Any] = Field(default_factory=dict)
+
+    # Description of the task, for extra information about its purpose and context
     description: str | None = None
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> Task:
         return cls(**data)
 
+    @classmethod
+    def from_serialized(cls, data: dict[str, Any]) -> Task:
+        gym_data = data.get("gym")
+        parsed_gym: Gym | None = gym_data
+
+        parsed_setup = [(param, entry) for param, entry in data.get("setup", [])]
+        parsed_evaluate = [(param, entry) for param, entry in data.get("evaluate", [])]
+
+        # Convert dict gym data to CustomGym if needed
+        if (
+            isinstance(gym_data, dict)
+            and gym_data.get("type") == "public"
+            and gym_data.get("location") in ("local", "remote")
+            and gym_data.get("image_or_build_context") is not None
+        ):
+            parsed_gym = CustomGym(
+                type=cast("Literal['public']", gym_data["type"]),
+                location=cast("Literal['local', 'remote']", gym_data["location"]),
+                image_or_build_context=Path(gym_data["image_or_build_context"]),
+            )
+
+        return cls(
+            id=data.get("id"),
+            prompt=data.get("prompt", ""),
+            system_prompt=data.get("system_prompt"),
+            setup=parsed_setup,
+            evaluate=parsed_evaluate,
+            gym=parsed_gym,
+            config=data.get("config"),
+            description=data.get("description"),
+            sensitive_data=data.get("sensitive_data", {}),
+            metadata=data.get("metadata", {}),
+        )
+
     @classmethod
     def from_inspect_sample(cls, sample: Sample) -> Task:
         """Create a Task from an Inspect dataset sample.
@@ -144,3 +194,31 @@ class Task(BaseModel):
         if self.gym is None:
             return
         self.gym = agent.transfer_gyms.get(self.gym, self.gym)
+
+    def serialize(self) -> dict[str, Any]:
+        if isinstance(self.setup, list):
+            parsed_setup = [[param, entry] for param, entry in self.setup]
+        else:
+            parsed_setup = self.setup
+        if isinstance(self.evaluate, list):
+            parsed_evaluate = [[param, entry] for param, entry in self.evaluate]
+        else:
+            parsed_evaluate = self.evaluate
+
+        if isinstance(self.gym, CustomGym):
+            parsed_gym = self.gym.model_dump()
+            parsed_gym["image_or_build_context"] = str(parsed_gym["image_or_build_context"])
+        else:  # is ServerGym
+            parsed_gym = self.gym
+
+        return {
+            "id": self.id,
+            "prompt": self.prompt,
+            "config": self.config,
+            "description": self.description,
+            "setup": parsed_setup,
+            "evaluate": parsed_evaluate,
+            "gym": parsed_gym,
+            "sensitive_data": self.sensitive_data,
+            "metadata": self.metadata,
+        }

hud/taskset.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from pathlib import PosixPath
+from typing import TYPE_CHECKING, Any, get_args
 from venv import logger
 
 from pydantic import BaseModel
@@ -9,6 +10,7 @@ from hud.env.environment import create_remote_config
 from hud.server import make_request
 from hud.settings import settings
 from hud.task import Task
+from hud.types import CustomGym, ServerGym
 from hud.utils.config import REMOTE_EVALUATE, REMOTE_SETUP
 
 if TYPE_CHECKING:
@@ -101,13 +103,30 @@ class TaskSet(BaseModel):
             else:
                 evaluate_config = None
 
+            if isinstance(task.gym, CustomGym):
+                if isinstance(task.gym.image_or_build_context, PosixPath):
+                    raise ValueError(
+                        "Local build contexts are not supported for "
+                        "remote tasksets, attach an image or existing "
+                        "gym id."
+                    )
+                gym_str = "docker"
+                image_uri = task.gym.image_or_build_context
+            elif isinstance(task.gym, str) and task.gym in get_args(ServerGym):
+                gym_str = task.gym
+                image_uri = None
+            else:
+                raise ValueError(f"Unknown gym type: {type(task.gym)}")
+
             processed_tasks.append(
                 {
                     "prompt": task.prompt,
-                    "gym": task.gym,
+                    "gym": gym_str,
                     "setup": setup_config,
                     "evaluate": evaluate_config,
                     "config": task.config,
+                    "image_uri": image_uri,
+                    "description": task.description,
                 }
             )
 
@@ -125,7 +144,15 @@ class TaskSet(BaseModel):
             "Taskset %s uploaded successfully, see it on app.hud.so/evalsets/%s", name, name
         )
 
-    async def fit(self, agent: Agent | type[Agent]) -> None:
+    def _apply(self, dict: dict[str, Any]) -> None:
+        """
+        Applies a parameter to all tasks in the taskset.
+        """
+        for task in self.tasks:
+            for key, value in dict.items():
+                setattr(task, key, value)
+
+    def fit(self, agent: Agent | type[Agent]) -> None:
         """
         Automatically adapts the taskset to the agent's transfer_gyms.
         """
@@ -133,19 +160,27 @@ class TaskSet(BaseModel):
             agent = agent()
 
         for task in self.tasks:
-            if task.gym is None:
+            if task.gym is None or isinstance(task.gym, CustomGym):
                 continue
             task.gym = agent.transfer_gyms.get(task.gym, task.gym)
 
 
-async def load_taskset(taskset_id: str, api_key: str | None = None) -> TaskSet:
+async def load_taskset(
+    taskset_id: str,
+    api_key: str | None = None,
+    metadata: dict[str, Any] | None = None,
+    load_custom_as_local: bool = False,
+    system_prompt: str | None = None,
+) -> TaskSet:
     """
     Loads a TaskSet by its ID.
 
     Args:
         taskset_id: The ID of the taskset to load
         api_key: Optional API key to use for the request
-
+        metadata: Optional metadata to apply to the taskset
+        load_custom_as_local: Whether to load custom gyms as local
+        system_prompt: Optional system prompt to override the default
     Returns:
         TaskSet: The loaded taskset
     """
@@ -161,13 +196,33 @@ async def load_taskset(taskset_id: str, api_key: str | None = None) -> TaskSet:
 
     logger.info(f"Taskset {taskset_id} loaded successfully")
 
-    return TaskSet.model_validate(
+    tasks = data["evalset"]
+    for task in tasks:
+        if system_prompt:
+            task["system_prompt"] = system_prompt
+        if task["gym"] == "docker":
+            if "image_uri" not in task:
+                raise ValueError(
+                    "No `image_uri` key found. This taskset may be "
+                    "incompatible with your version of HUD SDK."
+                )
+
+            task["gym"] = CustomGym(
+                location="local" if load_custom_as_local else "remote",
+                image_or_build_context=task["image_uri"],
+            )
+
+    taskset = TaskSet.model_validate(
         {
             "id": taskset_id,
-            "tasks": data["evalset"],
+            "tasks": tasks,
         }
     )
 
+    taskset._apply({"metadata": metadata})
+
+    return taskset
+
 
 def load_from_inspect(dataset: Dataset) -> TaskSet:
     """

hud/telemetry/exporter.py

@@ -298,12 +298,10 @@ async def _export_trace_payload_async(payload: dict[str, Any]) -> None:
         "telemetry": payload.get("mcp_calls", []),
     }
 
-    # Ensure mcp_calls is not empty if that's a requirement, or send as is. For now, send as is.
-    # if not data_to_send["mcp_calls"]:
-    #     logger.debug("No MCP calls in payload for task run %s, skipping specific export if "
-    #                  "desired.", task_run_id)
-    #     # Depending on backend, might not want to send empty mcp_calls list, or it's fine.
+    await send_telemetry_to_server(task_run_id, data_to_send)
 
+
+async def send_telemetry_to_server(task_run_id: str, data: dict[str, Any]) -> None:
     telemetry_url = f"{settings.base_url}/v2/task_runs/{task_run_id}/telemetry-upload"
 
     try:
@@ -320,7 +318,7 @@ async def _export_trace_payload_async(payload: dict[str, Any]) -> None:
             )
             response = await client.post(
                 telemetry_url,
-                json=data_to_send,  # Send the structured attributes and mcp_calls
+                json=data,  # Send the structured attributes and mcp_calls
                 headers=headers,
                 timeout=30.0,
             )

hud/trajectory.py

@@ -6,6 +6,8 @@ import datetime
 from IPython.display import HTML, Markdown, display
 from pydantic import BaseModel, Field
 
+from .adapters.common.types import LogType
+
 
 class TrajectoryStep(BaseModel):
     """Model representing a single task run's trajectory information."""
@@ -13,6 +15,7 @@ class TrajectoryStep(BaseModel):
     observation_url: str | None = None
     observation_text: str | None = None
     actions: list[dict]
+    logs: LogType | None = None
     start_timestamp: str | None = None
     end_timestamp: str | None = None