perf: improve performance on parallel processing

Author: Parth220

hud/otel/exporters.py

@@ -21,6 +21,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from datetime import UTC, datetime
 from typing import TYPE_CHECKING, Any
+import httpx
 
 from mcp.types import ClientRequest, ServerResult
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
@@ -323,6 +324,11 @@ def __init__(self, *, telemetry_url: str, api_key: str) -> None:
         self._api_key = api_key
         # Track pending export futures so we can force-flush on shutdown
         self._pending_futures: list[cf.Future[SpanExportResult]] = []
+        # Persistent HTTP client to reuse connections
+        self._client = httpx.Client(
+            timeout=30.0,
+            limits=httpx.Limits(max_connections=2000, max_keepalive_connections=512, keepalive_expiry=15.0),
+        )
 
     # ------------------------------------------------------------------
     # Core API
@@ -378,6 +384,7 @@ def _sync_export():
                             url=url,
                             json=payload,
                             api_key=self._api_key,
+                            client=self._client,
                         )
                     except Exception as exc:
                         logger.exception("HUD exporter failed to send spans for task %s: %s", run_id, exc)
@@ -435,6 +442,7 @@ def _cleanup_done(f: cf.Future[SpanExportResult]) -> None:
                         url=url,
                         json=payload,
                         api_key=self._api_key,
+                        client=self._client,
                     )
                 except Exception as exc:
                     logger.exception("HUD exporter failed to send spans for task %s: %s", run_id, exc)
@@ -452,6 +460,11 @@ def shutdown(self) -> None:  # type: ignore[override]
             pass
         finally:
             self._pending_futures.clear()
+        # Close persistent client
+        try:
+            self._client.close()
+        except Exception:
+            pass
 
     def force_flush(self, timeout_millis: int | None = None) -> bool:  # type: ignore[override]
         # Wait for pending export futures

hud/telemetry/async_context.py

@@ -23,7 +23,6 @@
 from hud.telemetry.job import Job, _print_job_url, _print_job_complete_url, get_current_job
 from hud.telemetry.trace import Trace
 from hud.utils.task_tracking import track_task
-from hud.utils.task_tracking import wait_all_tasks
 
 logger = logging.getLogger(__name__)
 
@@ -119,11 +118,6 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
         if self._otel_trace:
             self._otel_trace.__exit__(exc_type, exc_val, exc_tb)
 
-        # Ensure any pending telemetry tasks complete quickly
-        try:
-            await wait_all_tasks(timeout=5.0)
-        except Exception:
-            pass
 
         logger.debug(f"Ended async trace for task_run_id={self.task_run_id}")
 
@@ -211,11 +205,6 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
         # Restore old job
         _current_job = self._old_job
 
-        # Ensure any pending telemetry tasks complete quickly
-        try:
-            await wait_all_tasks(timeout=5.0)
-        except Exception:
-            pass
 
         logger.debug(f"Ended async job {self.job.id}")