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

hud/telemetry/tests/test_trace.py

@@ -0,0 +1,270 @@
+from __future__ import annotations
+
+import asyncio
+import uuid
+from unittest.mock import MagicMock
+
+import pytest
+
+from hud.telemetry._trace import (
+    init_telemetry,
+    register_trace,
+    trace,
+)
+from hud.telemetry.context import get_current_task_run_id as actual_get_current_task_run_id
+from hud.telemetry.context import is_root_trace as actual_is_root_trace
+from hud.telemetry.context import reset_context
+from hud.telemetry.context import set_current_task_run_id as actual_set_current_task_run_id
+
+
+@pytest.fixture(autouse=True)
+def reset_telemetry_context_fixture():
+    """Ensures telemetry context is reset before and after each test in this file."""
+    reset_context()
+    yield
+    reset_context()
+
+
+class TestInitTelemetry:
+    """Test telemetry initialization."""
+
+    def test_init_telemetry(self, mocker):
+        """Test telemetry initialization calls registry.install_all."""
+        mock_registry = mocker.patch("hud.telemetry._trace.registry", autospec=True)
+        init_telemetry()
+        mock_registry.install_all.assert_called_once()
+
+
+class TestTrace:
+    """Test the trace context manager."""
+
+    def test_trace_basic(self, mocker):
+        """Test basic trace functionality and context setting."""
+        mock_flush = mocker.patch(
+            "hud.telemetry._trace.flush_buffer", return_value=[], autospec=True
+        )
+        mock_submit_loop = mocker.patch(
+            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+        )
+
+        initial_root_state = actual_is_root_trace.get()
+
+        with trace() as task_run_id:
+            assert isinstance(task_run_id, str)
+            uuid.UUID(task_run_id)
+            assert actual_get_current_task_run_id() == task_run_id
+            assert actual_is_root_trace.get() is True
+
+        assert actual_get_current_task_run_id() is None
+        assert actual_is_root_trace.get() == initial_root_state
+        mock_flush.assert_called_once()
+        mock_submit_loop.assert_not_called()
+
+    def test_trace_with_name_and_attributes(self, mocker):
+        """Test trace with name and attributes, checking they are passed on."""
+        mock_mcp_calls = [MagicMock()]
+        mock_flush = mocker.patch(
+            "hud.telemetry._trace.flush_buffer", return_value=mock_mcp_calls, autospec=True
+        )
+        mock_submit_loop = mocker.patch(
+            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+        )
+
+        trace_name = "test_trace_with_data"
+        attrs = {"key": "value", "number": 42}
+
+        with trace(name=trace_name, attributes=attrs) as task_run_id:
+            assert isinstance(task_run_id, str)
+
+        mock_flush.assert_called_once()
+        mock_submit_loop.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_trace_with_mcp_calls_exports(self, mocker):
+        """Test trace with MCP calls exports telemetry with correct data."""
+        mock_mcp_calls = [MagicMock(), MagicMock()]
+        mock_flush = mocker.patch(
+            "hud.telemetry._trace.flush_buffer", return_value=mock_mcp_calls, autospec=True
+        )
+        mock_submit_loop = mocker.patch(
+            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+        )
+
+        async def mock_export(*args, **kwargs):
+            return None
+
+        mock_export_actual_coro = mocker.patch(
+            "hud.telemetry._trace.exporter.export_telemetry",
+            side_effect=mock_export,
+        )
+
+        test_attrs = {"custom_attr": "test_val"}
+        test_name = "mcp_export_test"
+
+        with trace(name=test_name, attributes=test_attrs) as task_run_id:
+            pass
+
+        mock_flush.assert_called_once()
+        mock_submit_loop.assert_called_once()
+
+        mock_export_actual_coro.assert_called_once()
+        args, kwargs = mock_export_actual_coro.call_args
+        assert kwargs["task_run_id"] == task_run_id
+        assert kwargs["mcp_calls"] == mock_mcp_calls
+        assert kwargs["trace_attributes"]["trace_name"] == test_name
+        assert kwargs["trace_attributes"]["custom_attr"] == "test_val"
+        assert "start_time" in kwargs["trace_attributes"]
+        assert "end_time" in kwargs["trace_attributes"]
+        assert "duration" in kwargs["trace_attributes"]
+        assert kwargs["trace_attributes"]["is_root"] is True
+
+    def test_trace_nested(self, mocker):
+        """Test nested traces, verifying context restoration and root trace logic."""
+        actual_set_current_task_run_id(None)
+        actual_is_root_trace.set(False)
+
+        mock_flush_internal = mocker.patch(
+            "hud.telemetry._trace.flush_buffer", return_value=[], autospec=True
+        )
+        mock_submit_loop_internal = mocker.patch(
+            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+        )
+
+        assert actual_get_current_task_run_id() is None
+        assert actual_is_root_trace.get() is False
+
+        with trace(name="outer") as outer_id:
+            assert actual_get_current_task_run_id() == outer_id
+            assert actual_is_root_trace.get() is True
+            with trace(name="inner") as inner_id:
+                assert actual_get_current_task_run_id() == inner_id
+                assert actual_is_root_trace.get() is False
+                assert outer_id != inner_id
+            assert actual_get_current_task_run_id() == outer_id
+            assert actual_is_root_trace.get() is True
+
+        assert actual_get_current_task_run_id() is None
+        assert actual_is_root_trace.get() is False
+        assert mock_flush_internal.call_count == 2
+        mock_submit_loop_internal.assert_not_called()
+
+    def test_trace_exception_handling(self, mocker):
+        """Test trace handles exceptions properly and restores context."""
+        initial_task_id_before_trace = "pre_existing_id_123"
+        initial_root_state_before_trace = True
+        actual_set_current_task_run_id(initial_task_id_before_trace)
+        actual_is_root_trace.set(initial_root_state_before_trace)
+
+        mock_flush = mocker.patch(
+            "hud.telemetry._trace.flush_buffer", return_value=[], autospec=True
+        )
+        mock_submit_loop = mocker.patch(
+            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+        )
+
+        with pytest.raises(ValueError, match="Test exception"), trace(name="trace_with_exception"):
+            assert actual_get_current_task_run_id() != initial_task_id_before_trace
+            assert actual_is_root_trace.get() is False
+            raise ValueError("Test exception")
+
+        mock_flush.assert_called_once()
+        assert actual_get_current_task_run_id() == initial_task_id_before_trace
+        assert actual_is_root_trace.get() == initial_root_state_before_trace
+        mock_submit_loop.assert_not_called()
+
+
+class TestRegisterTrace:
+    """Test the register_trace decorator."""
+
+    def test_register_trace_sync_function(self, mocker):
+        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
+        mock_trace_context_manager.return_value.__enter__.return_value = "mocked_task_id"
+        mock_trace_context_manager.return_value.__exit__.return_value = None
+
+        @register_trace(name="test_func_sync")
+        def sync_function(x, y):
+            return x + y
+
+        result = sync_function(1, 2)
+        assert result == 3
+        mock_trace_context_manager.assert_called_once_with(name="test_func_sync", attributes=None)
+
+    def test_register_trace_async_function(self, mocker):
+        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
+        mock_trace_context_manager.return_value.__enter__.return_value = "mocked_task_id"
+        mock_trace_context_manager.return_value.__exit__.return_value = None
+
+        @register_trace(name="test_func_async")
+        async def async_function(x, y):
+            return x + y
+
+        async def run_test():
+            result = await async_function(1, 2)
+            assert result == 3
+            mock_trace_context_manager.assert_called_once_with(
+                name="test_func_async", attributes=None
+            )
+
+        asyncio.run(run_test())
+
+    def test_register_trace_with_attributes(self, mocker):
+        """Test register_trace with attributes."""
+        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
+
+        class _MockTraceContextManager:
+            def __enter__(self):
+                return "task_id"
+
+            def __exit__(self, exc_type, exc_value, traceback):
+                return None
+
+        mock_trace_context_manager.return_value = _MockTraceContextManager()
+
+        attrs = {"operation": "add"}
+
+        @register_trace(name="test_func", attributes=attrs)
+        def func_with_attrs(x):
+            return x * 2
+
+        result = func_with_attrs(5)
+        assert result == 10
+        mock_trace_context_manager.assert_called_once_with(name="test_func", attributes=attrs)
+
+    def test_register_trace_without_name(self, mocker):
+        """Test register_trace uses function name when name not provided."""
+        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
+        mock_trace_context_manager.return_value.__enter__.return_value = "task_id"
+        mock_trace_context_manager.return_value.__exit__.return_value = None
+
+        @register_trace()
+        def my_function():
+            return "result"
+
+        result = my_function()
+        assert result == "result"
+        mock_trace_context_manager.assert_called_once_with(name="my_function", attributes=None)
+
+    def test_register_trace_preserves_function_metadata(self):
+        """Test register_trace preserves original function metadata."""
+
+        @register_trace(name="test")
+        def original_function():
+            """Original docstring."""
+
+        assert original_function.__name__ == "original_function"
+        assert original_function.__doc__ == "Original docstring."
+
+    def test_register_trace_exception_propagation(self, mocker):
+        """Test register_trace propagates exceptions."""
+        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
+        mock_trace_context_manager.return_value.__enter__.return_value = "task_id"
+        mock_trace_context_manager.return_value.__exit__.return_value = None
+
+        @register_trace()
+        def failing_function():
+            raise RuntimeError("Test error")
+
+        with pytest.raises(RuntimeError, match="Test error"):
+            failing_function()
+
+        mock_trace_context_manager.assert_called_once()

hud/types.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import enum
 from pathlib import Path
-from typing import Any, Literal, TypeAlias
+from typing import Literal, TypeAlias
 
 from pydantic import BaseModel
 
@@ -12,38 +12,22 @@ class CustomGym(BaseModel):
     Public environment specification with a dockerfile and controller.
 
     If the location is remote, the env will be created on the server.
-    If the location is dev, the env will be created locally via docker.
+    If the location is local, the env will be created locally via docker.
 
     The dockerfile can be specified directly or automatically found in the controller_source_dir.
     If neither is provided, an error will be raised during validation.
     """
 
     type: Literal["public"] = "public"
-    dockerfile: str | None = None
     location: Literal["local", "remote"]
-    ports: list[int] | None = None
-    # If path, then it is a development environment on the local computer
-    # If none, then the controller must be installed in the environment through the dockerfile
-    # Can be provided as a string or Path object
-    controller_source_dir: str | Path | None = None
-
-    def model_post_init(self, __context: Any, /) -> None:
-        """Validate and set up dockerfile if not explicitly provided."""
-        # Convert string path to Path object if needed
-        if isinstance(self.controller_source_dir, str):
-            self.controller_source_dir = Path(self.controller_source_dir)
-
-        if self.dockerfile is None:
-            if self.controller_source_dir is None:
-                raise ValueError("Either dockerfile or controller_source_dir must be provided")
-
-            # Look for Dockerfile in the controller_source_dir
-            dockerfile_path = self.controller_source_dir / "Dockerfile"
-            if not dockerfile_path.exists():
-                raise ValueError(f"Dockerfile not found in {self.controller_source_dir}")
-
-            # Read the Dockerfile content
-            self.dockerfile = dockerfile_path.read_text()
+    # A. If path, then it is a docker build context on the local computer.
+    #    If the location is local, docker build will be used to create the image.
+    #    If the location is remote, we will build the image remotely.
+    #    The controller will be automatically installed and kept in sync with local changes
+    #    as long as a pyproject.toml is present at the root of the folder.
+    # B. If string, then it is the uri of the docker image to use.
+    #    The controller must already be installed in the image.
+    image_or_build_context: str | Path
 
 
 class EnvironmentStatus(str, enum.Enum):
@@ -64,7 +48,7 @@ class EnvironmentStatus(str, enum.Enum):
 
 
 # Available HUD gyms
-ServerGym: TypeAlias = Literal["qa", "hud-browser", "hud-ubuntu", "OSWorld-Ubuntu"]
+ServerGym: TypeAlias = Literal["qa", "hud-browser", "OSWorld-Ubuntu"]
 
 # Gyms can be either custom or server-side
 Gym: TypeAlias = CustomGym | ServerGym

hud/utils/common.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import io
 import logging
 import tarfile
+import zipfile
 from typing import TYPE_CHECKING, Any, TypedDict
 
 from pydantic import BaseModel
@@ -22,6 +23,7 @@ class FunctionConfig(BaseModel):
     args: list[Any]  # Must be json serializable
 
     id: str | None = None  # Optional id for remote execution
+    metadata: dict[str, Any] | None = None  # Optional metadata for telemetry
 
     def __len__(self) -> int:
         return len(self.args)
@@ -33,11 +35,12 @@ class FunctionConfig(BaseModel):
         return iter(self.args)
 
     def __str__(self) -> str:
-        return f"{self.function}: {', '.join(str(arg) for arg in self.args)}"
+        return f"FC: {self.function}: {', '.join(str(arg) for arg in self.args)} ({self.metadata})"
 
 
 # Type alias for the shorthand config, which just converts to function name and args
-ShorthandConfig = tuple[str | dict[str, Any] | list[str] | list[dict[str, Any]], ...]
+BasicType = str | int | float | bool | None
+ShorthandConfig = tuple[BasicType | dict[str, Any] | list[BasicType] | list[dict[str, Any]], ...]
 
 # Type alias for multiple config formats
 FunctionConfigs = (
@@ -62,6 +65,11 @@ class Observation(BaseModel):
     screenshot: str | None = None  # base64 string png
     text: str | None = None
 
+    def __str__(self) -> str:
+        return f"""Observation(screenshot={
+            self.screenshot[:100] if self.screenshot else "None"
+        }..., text={self.text}...)"""
+
 
 class ExecuteResult(TypedDict):
     """
@@ -107,6 +115,18 @@ def directory_to_tar_bytes(directory_path: Path) -> bytes:
     return output.getvalue()
 
 
+def directory_to_zip_bytes(context_dir: Path) -> bytes:
+    """Zip a directory and return the zip archive as bytes."""
+    output = io.BytesIO()
+    with zipfile.ZipFile(output, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for file_path in context_dir.rglob("*"):
+            if file_path.is_file():
+                rel_path = file_path.relative_to(context_dir)
+                logger.debug("Adding %s to zip archive", rel_path)
+                zipf.write(str(file_path), arcname=str(rel_path))
+    return output.getvalue()
+
+
 async def get_gym_id(gym_name_or_id: str) -> str:
     """
     Get the gym ID for a given gym name or ID.

hud/utils/misc.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from hud.server import make_request
+from hud.settings import settings
+
+if TYPE_CHECKING:
+    from hud.env.environment import Environment  # Import Environment for type hinting
+
+logger = logging.getLogger(__name__)
+
+
+async def upload_env_telemetry(
+    environment: Environment,
+    results: Any,
+    api_key: str | None = None,
+) -> None:
+    """
+    Sends telemetry data (results from a cloud runner) to the HUD telemetry upload endpoint.
+    """
+    environment_id = environment.client.env_id  # type: ignore
+
+    if not api_key:
+        api_key = settings.api_key
+
+    if not api_key:
+        raise ValueError("API key must be provided either as an argument or set in hud.settings.")
+
+    endpoint_url = f"{settings.base_url}/v2/environments/{environment_id}/telemetry-upload"
+
+    request_payload = {
+        "results": {
+            "steps": results,
+        }
+    }
+
+    logger.debug("Sending telemetry to %s for env_id: %s", endpoint_url, environment_id)
+
+    try:
+        await make_request(
+            method="POST",
+            url=endpoint_url,
+            json=request_payload,
+            api_key=api_key,
+        )
+        logger.info("Successfully uploaded telemetry for environment_id: %s", environment_id)
+    except Exception as e:
+        logger.error(
+            "Failed to upload telemetry for environment_id: %s. Error: %s", environment_id, e
+        )
+        raise

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.2.4"
+    assert hud.__version__ == "0.2.6"

hud/version.py

@@ -0,0 +1,7 @@
+"""
+Version information for the HUD SDK.
+"""
+
+from __future__ import annotations
+
+__version__ = "0.2.6"