hud-python 0.4.51__py3-none-any.whl → 0.4.53__py3-none-any.whl

hud/cli/utils/version_check.py

@@ -0,0 +1,257 @@
+"""Version checking utilities for HUD CLI.
+
+This module handles checking for updates to the hud-python package
+and prompting users to upgrade when a new version is available.
+
+Features:
+- Checks PyPI for the latest version of hud-python
+- Caches results for 6 hours to avoid excessive API calls
+- Displays a friendly prompt when an update is available
+- Can be disabled with HUD_SKIP_VERSION_CHECK=1 environment variable
+
+The version check runs automatically at the start of most CLI commands,
+but is skipped for help and version commands to keep them fast.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+import time
+from pathlib import Path
+from typing import NamedTuple
+
+import httpx
+from packaging import version
+
+from hud.utils.hud_console import HUDConsole
+
+# Cache location for version check data
+CACHE_DIR = Path.home() / ".hud" / ".cache"
+VERSION_CACHE_FILE = CACHE_DIR / "version_check.json"
+
+# Cache duration in seconds (6 hours)
+CACHE_DURATION = 6 * 60 * 60
+
+# PyPI API URL for package info
+PYPI_URL = "https://pypi.org/pypi/hud-python/json"
+
+
+class VersionInfo(NamedTuple):
+    """Version information from PyPI."""
+
+    latest: str
+    current: str
+    is_outdated: bool
+    checked_at: float
+
+
+def _get_current_version() -> str:
+    """Get the currently installed version of hud-python."""
+    try:
+        from hud import __version__
+
+        return __version__
+    except ImportError:
+        return "unknown"
+
+
+def _fetch_latest_version() -> str | None:
+    """Fetch the latest version from PyPI.
+
+    Returns:
+        The latest version string, or None if the request fails.
+    """
+    try:
+        with httpx.Client(timeout=3.0) as client:
+            response = client.get(PYPI_URL)
+            if response.status_code == 200:
+                data = response.json()
+                return data["info"]["version"]
+    except Exception:  # noqa: S110
+        # Silently fail - we don't want to disrupt the user's workflow
+        # if PyPI is down or there's a network issue
+        pass
+    return None
+
+
+def _load_cache() -> VersionInfo | None:
+    """Load cached version information.
+
+    Returns:
+        Cached VersionInfo if valid, None otherwise.
+    """
+    if not VERSION_CACHE_FILE.exists():
+        return None
+
+    try:
+        with open(VERSION_CACHE_FILE) as f:
+            data = json.load(f)
+
+        # Check if cache is still valid
+        if time.time() - data["checked_at"] > CACHE_DURATION:
+            return None
+
+        return VersionInfo(
+            latest=data["latest"],
+            current=data["current"],
+            is_outdated=data["is_outdated"],
+            checked_at=data["checked_at"],
+        )
+    except Exception:
+        # If cache is corrupted, return None
+        return None
+
+
+def _save_cache(info: VersionInfo) -> None:
+    """Save version information to cache.
+
+    Args:
+        info: Version information to cache.
+    """
+    try:
+        # Create cache directory if it doesn't exist
+        CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+        with open(VERSION_CACHE_FILE, "w") as f:
+            json.dump(
+                {
+                    "latest": info.latest,
+                    "current": info.current,
+                    "is_outdated": info.is_outdated,
+                    "checked_at": info.checked_at,
+                },
+                f,
+            )
+    except Exception:  # noqa: S110
+        # Silently fail if we can't write cache
+        pass
+
+
+def _compare_versions(current: str, latest: str) -> bool:
+    """Compare versions to determine if an update is available.
+
+    Args:
+        current: Current version string
+        latest: Latest version string
+
+    Returns:
+        True if latest is newer than current, False otherwise.
+    """
+    if current == "unknown":
+        return False
+
+    try:
+        current_v = version.parse(current)
+        latest_v = version.parse(latest)
+        return latest_v > current_v
+    except Exception:
+        # If we can't parse versions, assume no update needed
+        return False
+
+
+def check_for_updates() -> VersionInfo | None:
+    """Check for updates to hud-python.
+
+    This function checks PyPI for the latest version and caches the result
+    for 6 hours to avoid excessive API calls.
+
+    Returns:
+        VersionInfo if check succeeds, None otherwise.
+    """
+    # Check if we're in CI/testing environment
+    if os.environ.get("CI") or os.environ.get("HUD_SKIP_VERSION_CHECK"):
+        return None
+
+    # Get current version first
+    current = _get_current_version()
+    if current == "unknown":
+        return None
+
+    # Try to load from cache
+    cached_info = _load_cache()
+
+    # If cache exists but current version has changed (user upgraded), invalidate cache
+    if cached_info and cached_info.current != current:
+        cached_info = None  # Force fresh check
+
+    if cached_info:
+        # Update the current version in the cached info to reflect reality
+        # but keep the cached latest version and timestamp
+        return VersionInfo(
+            latest=cached_info.latest,
+            current=current,  # Use actual current version, not cached
+            is_outdated=_compare_versions(current, cached_info.latest),
+            checked_at=cached_info.checked_at,
+        )
+
+    # Fetch latest version from PyPI
+    latest = _fetch_latest_version()
+    if not latest:
+        return None
+
+    # Compare versions
+    is_outdated = _compare_versions(current, latest)
+
+    # Create version info
+    info = VersionInfo(
+        latest=latest,
+        current=current,
+        is_outdated=is_outdated,
+        checked_at=time.time(),
+    )
+
+    # Save to cache
+    _save_cache(info)
+
+    return info
+
+
+def display_update_prompt(console: HUDConsole | None = None) -> None:
+    """Display update prompt if a new version is available.
+
+    This function checks for updates and displays a prompt to the user
+    if their version is outdated.
+
+    Args:
+        console: HUDConsole instance for output. If None, creates a new one.
+    """
+    if console is None:
+        console = HUDConsole()
+
+    try:
+        info = check_for_updates()
+        if info and info.is_outdated:
+            # Create update message
+            update_msg = (
+                f"🆕 A new version of hud-python is available: "
+                f"[bold cyan]{info.latest}[/bold cyan] "
+                f"(current: [dim]{info.current}[/dim])\n"
+                f"   Run: [bold yellow]uv tool upgrade hud-python[/bold yellow] to update"
+            )
+
+            # Display as a subtle but noticeable panel
+            console._stdout_console.print(
+                f"\n[yellow]{update_msg}[/yellow]\n",
+                highlight=False,
+            )
+    except Exception:  # noqa: S110
+        # Never let version checking disrupt the user's workflow
+        pass
+
+
+def force_version_check() -> VersionInfo | None:
+    """Force a version check, bypassing the cache.
+
+    This is useful for explicit version checks or testing.
+
+    Returns:
+        VersionInfo if check succeeds, None otherwise.
+    """
+    # Clear the cache to force a fresh check
+    if VERSION_CACHE_FILE.exists():
+        with contextlib.suppress(Exception):
+            VERSION_CACHE_FILE.unlink()
+
+    return check_for_updates()

hud/clients/base.py

@@ -170,7 +170,7 @@ class BaseHUDClient(AgentMCPClient):
         if self._initialized:
             await self._disconnect()
             self._initialized = False
-            hud_console.info("Shutdown completed")
+            hud_console.info("Environment Shutdown completed")
         else:
             hud_console.debug("Client was not initialized, skipping disconnect")
 

hud/clients/mcp_use.py

@@ -92,7 +92,9 @@ class MCPUseHUDClient(BaseHUDClient):
         try:
             assert self._client is not None  # noqa: S101
             self._sessions = await self._client.create_all_sessions()
-            hud_console.info(f"Created {len(self._sessions)} MCP sessions")
+            session_count = len(self._sessions)
+            session_text = "session" if session_count == 1 else "sessions"
+            hud_console.info(f"Created {session_count} MCP {session_text}")
 
             # Configure validation for all sessions based on client setting
             try:

hud/datasets/parallel.py

@@ -111,13 +111,13 @@ def _process_worker(
             """Process a single task with telemetry tracking."""
             async with sem:
                 try:
-                    # Create trace for this task (linked to the job) - match original format
+                    # Create trace for this task (linked to the job)
                     task_name = task_dict.get("prompt") or f"Task {index}"
 
                     # Use the job_id to group all tasks under the same job
                     raw_task_id = task_dict.get("id")
                     safe_task_id = str(raw_task_id) if raw_task_id is not None else None
-                    with hud.trace(task_name, job_id=job_id, task_id=safe_task_id):
+                    async with hud.async_trace(task_name, job_id=job_id, task_id=safe_task_id):
                         # Convert dict to Task
                         task = Task(**task_dict)
 

hud/datasets/runner.py

@@ -28,8 +28,11 @@ async def run_dataset(
     split: str = "train",
     auto_respond: bool = False,
 ) -> list[Any]:
-    """
-    Run all tasks in a dataset with automatic job tracking.
+    """Run all tasks in a dataset with automatic job and telemetry tracking.
+
+    This function handles concurrent task execution with proper telemetry collection.
+    All tasks are executed in parallel up to `max_concurrent`, with full telemetry
+    automatically uploaded to the HUD platform.
 
     Args:
         name: Name for the job
@@ -37,23 +40,27 @@ async def run_dataset(
                 Dataset object, OR list of Task objects
         agent_class: Agent class to instantiate (e.g., ClaudeAgent)
         agent_config: Configuration/kwargs for agent (model, etc.)
-        max_concurrent: Maximum parallel task execution
+        max_concurrent: Maximum parallel task execution. Higher values improve throughput
+                       but may increase memory usage. Recommended: 30-200 depending on
+                       task complexity and available resources.
         metadata: Optional metadata for the job
         max_steps: Maximum steps per task
         split: Dataset split to use when loading from string (default: "train")
         auto_respond: Whether to use auto-response agent
 
     Returns:
-        List of results from agent.run() in dataset order
+        List of results from agent.run() in dataset order. Telemetry is automatically
+        collected and uploaded for all tasks.
 
     Example:
         >>> from hud.agents import ClaudeAgent
-        >>> # Option 1: From dataset string identifier
+        >>> # Basic usage with dataset identifier
         >>> results = await run_dataset(
         ...     "SheetBench Eval",
         ...     "hud-evals/SheetBench-50",
         ...     ClaudeAgent,
         ...     {"model": "claude-3-5-sonnet-20241022"},
+        ...     max_concurrent=100,  # Adjust based on your needs
         ... )
         >>> # Option 2: From HuggingFace dataset object
         >>> from datasets import load_dataset
@@ -62,9 +69,12 @@ async def run_dataset(
         >>> # Option 3: From list of dicts
         >>> tasks = [{"prompt": "...", "mcp_config": {...}, ...}, ...]
         >>> results = await run_dataset("browser_eval", tasks, ClaudeAgent)
+
+    Note:
+        Telemetry collection and upload is handled automatically. The function ensures
+        all telemetry is flushed before returning, even at high concurrency levels.
     """
-    # Import here to avoid circular imports
-    import hud
+    import hud  # Import here to avoid circular imports
 
     dataset_link = None
 
@@ -91,33 +101,84 @@ async def run_dataset(
         except Exception:
             logger.warning("Failed to extract dataset verification info")
 
-    with hud.job(name, metadata=job_metadata, dataset_link=dataset_link) as job_obj:
+    # Use async job context manager for high-concurrency telemetry
+    async with hud.async_job(name, metadata=job_metadata, dataset_link=dataset_link) as job_obj:
         # Run tasks with semaphore for concurrency control
         sem = asyncio.Semaphore(max_concurrent)
         results: list[Any | None] = [None] * len(dataset)
 
         async def _worker(index: int, task_dict: Any, max_steps: int = 10) -> None:
             async with sem:
-                # Create trace for this task
-                task_name = task_dict.get("prompt") or f"Task {index}"
-
-                # Ensure task_id is a string for baggage propagation
-                raw_task_id = task_dict.get("id")
-                safe_task_id = str(raw_task_id) if raw_task_id is not None else None
-                with hud.trace(task_name, job_id=job_obj.id, task_id=safe_task_id):
-                    # Convert dict to Task here, at trace level
-                    task = Task(**task_dict)
-
-                    agent = agent_class(**(agent_config or {}))
-
-                    if auto_respond:
-                        agent.response_agent = ResponseAgent()
-                    results[index] = await agent.run(task, max_steps=max_steps)
+                try:
+                    # Create trace for this task
+                    task_name = task_dict.get("prompt") or f"Task {index}"
+
+                    # Ensure task_id is a string for baggage propagation
+                    raw_task_id = task_dict.get("id")
+                    safe_task_id = str(raw_task_id) if raw_task_id is not None else None
+                    async with hud.async_trace(task_name, job_id=job_obj.id, task_id=safe_task_id):
+                        # with hud.trace(task_name, job_id=job_obj.id, task_id=safe_task_id):
+                        # Convert dict to Task here, at trace level
+                        task = Task(**task_dict)
+
+                        agent = agent_class(**(agent_config or {}))
+
+                        if auto_respond:
+                            agent.response_agent = ResponseAgent()
+                        results[index] = await agent.run(task, max_steps=max_steps)
+                except Exception as e:
+                    logger.exception("Task %s failed: %s", index, e)
+                    results[index] = None
 
         # Execute all tasks
-        await asyncio.gather(
+        worker_results = await asyncio.gather(
             *[_worker(i, task, max_steps=max_steps) for i, task in enumerate(dataset)],
             return_exceptions=True,  # Don't fail entire batch on one error
         )
 
+        # Log any exceptions that occurred
+        for i, result in enumerate(worker_results):
+            if isinstance(result, Exception):
+                logger.error("Worker %s failed with exception: %s", i, result, exc_info=result)
+
+    # Ensure all telemetry is uploaded before returning
+    await _flush_telemetry()
+
     return results
+
+
+async def _flush_telemetry() -> None:
+    """Flush all pending telemetry operations.
+
+    Ensures complete telemetry upload by:
+    1. Waiting for all async status updates to complete
+    2. Forcing OpenTelemetry span processor to export remaining spans
+
+    This prevents telemetry loss at high concurrency (200+ tasks) by ensuring
+    all operations complete before process exit.
+    """
+    from hud.otel.config import is_telemetry_configured
+    from hud.utils import hud_console
+    from hud.utils.task_tracking import wait_all_tasks
+
+    hud_console.info("Uploading telemetry...")
+
+    # Step 1: Wait for async status updates (job/trace status)
+    completed_tasks = await wait_all_tasks(timeout_seconds=20.0)
+    if completed_tasks > 0:
+        hud_console.info(f"Completed {completed_tasks} pending telemetry tasks")
+
+    # Step 2: Flush OpenTelemetry span exports
+    if is_telemetry_configured():
+        try:
+            from opentelemetry import trace
+            from opentelemetry.sdk.trace import TracerProvider
+
+            provider = trace.get_tracer_provider()
+            if isinstance(provider, TracerProvider):
+                provider.force_flush(timeout_millis=20000)
+                logger.debug("OpenTelemetry spans flushed successfully")
+        except Exception as e:
+            logger.warning("Failed to flush OpenTelemetry: %s", e)
+
+    hud_console.info("Telemetry uploaded successfully")

hud/datasets/tests/test_runner.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from hud.datasets.runner import _flush_telemetry
+
+
+@pytest.mark.asyncio
+async def test_flush_telemetry():
+    """Test _flush_telemetry function."""
+    with (
+        patch("hud.otel.config.is_telemetry_configured", return_value=True),
+        patch("hud.utils.hud_console.hud_console"),
+        patch("hud.utils.task_tracking.wait_all_tasks", new_callable=AsyncMock) as mock_wait,
+        patch("opentelemetry.trace.get_tracer_provider") as mock_get_provider,
+    ):
+        from opentelemetry.sdk.trace import TracerProvider
+
+        mock_provider = MagicMock(spec=TracerProvider)
+        mock_provider.force_flush.return_value = True
+        mock_get_provider.return_value = mock_provider
+
+        mock_wait.return_value = 5
+
+        await _flush_telemetry()
+
+        mock_wait.assert_called_once()
+        mock_provider.force_flush.assert_called_once_with(timeout_millis=20000)
+
+
+@pytest.mark.asyncio
+async def test_flush_telemetry_no_telemetry():
+    """Test _flush_telemetry when telemetry is not configured."""
+    with (
+        patch("hud.otel.config.is_telemetry_configured", return_value=False),
+        patch("hud.utils.hud_console.hud_console"),
+        patch("hud.utils.task_tracking.wait_all_tasks", new_callable=AsyncMock) as mock_wait,
+        patch("opentelemetry.trace.get_tracer_provider"),
+    ):
+        mock_wait.return_value = 0
+
+        await _flush_telemetry()
+
+        mock_wait.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_flush_telemetry_exception():
+    """Test _flush_telemetry handles exceptions gracefully."""
+    with (
+        patch("hud.otel.config.is_telemetry_configured", return_value=True),
+        patch("hud.utils.hud_console.hud_console"),
+        patch("hud.utils.task_tracking.wait_all_tasks", new_callable=AsyncMock) as mock_wait,
+        patch("opentelemetry.trace.get_tracer_provider") as mock_get_provider,
+    ):
+        from opentelemetry.sdk.trace import TracerProvider
+
+        mock_provider = MagicMock(spec=TracerProvider)
+        mock_provider.force_flush.side_effect = Exception("Flush failed")
+        mock_get_provider.return_value = mock_provider
+
+        mock_wait.return_value = 3
+
+        # Should not raise
+        await _flush_telemetry()
+
+
+@pytest.mark.asyncio
+async def test_flush_telemetry_no_completed_tasks():
+    """Test _flush_telemetry when no tasks were completed."""
+    with (
+        patch("hud.otel.config.is_telemetry_configured", return_value=True),
+        patch("hud.utils.hud_console.hud_console"),
+        patch("hud.utils.task_tracking.wait_all_tasks", new_callable=AsyncMock) as mock_wait,
+        patch("opentelemetry.trace.get_tracer_provider") as mock_get_provider,
+    ):
+        from opentelemetry.sdk.trace import TracerProvider
+
+        mock_provider = MagicMock(spec=TracerProvider)
+        mock_get_provider.return_value = mock_provider
+
+        mock_wait.return_value = 0
+
+        await _flush_telemetry()
+
+        mock_provider.force_flush.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_flush_telemetry_non_sdk_provider():
+    """Test _flush_telemetry with non-SDK TracerProvider."""
+    with (
+        patch("hud.otel.config.is_telemetry_configured", return_value=True),
+        patch("hud.utils.hud_console.hud_console"),
+        patch("hud.utils.task_tracking.wait_all_tasks", new_callable=AsyncMock) as mock_wait,
+        patch("opentelemetry.trace.get_tracer_provider") as mock_get_provider,
+    ):
+        # Return a non-TracerProvider object
+        mock_get_provider.return_value = MagicMock(spec=object)
+
+        mock_wait.return_value = 2
+
+        # Should not raise
+        await _flush_telemetry()