update http client apis

Author: viraatc

src/inference_endpoint/commands/probe.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Probe command implementation for endpoint health checking."""
+"""Probe command for endpoint health checking."""
 
 import argparse
 import asyncio
@@ -24,12 +24,12 @@
 from urllib.parse import urljoin
 
 from inference_endpoint.core.types import Query, QueryResult
-from inference_endpoint.endpoint_client.configs import (
+from inference_endpoint.endpoint_client import (
     AioHttpConfig,
     HTTPClientConfig,
+    HTTPEndpointClient,
     ZMQConfig,
 )
-from inference_endpoint.endpoint_client.http_client import HTTPEndpointClient
 from inference_endpoint.exceptions import (
     ExecutionError,
     InputValidationError,
@@ -40,201 +40,123 @@
 
 
 async def run_probe_command(args: argparse.Namespace) -> None:
-    """Run endpoint probe to validate connectivity and basic functionality.
-
-    Actions:
-    1. Send test requests using HTTP client with futures
-    2. Measure basic latency
-    3. Report validation status
-    """
-    # Extract arguments
+    """Probe endpoint to validate connectivity and measure latency."""
     endpoint = args.endpoint
     num_requests = args.requests
-    test_prompt = args.prompt
+    prompt = args.prompt
+    model = getattr(args, "model", None)
 
-    # Model: use provided or default to valid OpenAI model name
-    model_name = getattr(args, "model", None)
-    if not model_name:
-        logger.error("Model required: --model or specify in YAML config")
+    if not model:
         raise InputValidationError("Model required: --model NAME")
-    # Note: API key handling would go in HTTP client config if needed
 
     logger.info(f"Probing: {endpoint}")
 
-    # Create temp directory for ZMQ
     tmp_dir = tempfile.mkdtemp(prefix="probe_")
     client = None
 
-    # TODO (Rashid): Add a health check with a separate timeout.
     try:
-        # Setup HTTP client with futures support
-        http_config = HTTPClientConfig(
-            endpoint_url=urljoin(endpoint, "/v1/chat/completions"),
-            num_workers=1,
-            max_concurrency=num_requests,
+        client = HTTPEndpointClient(
+            HTTPClientConfig(
+                endpoint_url=urljoin(endpoint, "/v1/chat/completions"),
+                num_workers=1,
+                max_concurrency=num_requests,
+            ),
+            AioHttpConfig(),
+            ZMQConfig(
+                zmq_request_queue_prefix=f"ipc://{tmp_dir}/req",
+                zmq_response_queue_addr=f"ipc://{tmp_dir}/resp",
+                zmq_readiness_queue_addr=f"ipc://{tmp_dir}/ready",
+            ),
         )
-        aiohttp_config = AioHttpConfig()
-        zmq_config = ZMQConfig(
-            zmq_request_queue_prefix=f"ipc://{tmp_dir}/req",
-            zmq_response_queue_addr=f"ipc://{tmp_dir}/resp",
-            zmq_readiness_queue_addr=f"ipc://{tmp_dir}/ready",
-        )
-
-        client = HTTPEndpointClient(http_config, aiohttp_config, zmq_config)
 
+        # Issue requests
         logger.info(f"Sending {num_requests} requests...")
-
-        # Send test requests
         start_times: dict[str, float] = {}
-        sent_query_ids: list[str] = []
-        issue_errors: list[str] = []
 
-        # TODO: this might not work with a real vLLM/SGLang endpoint, fix this.
         for i in range(num_requests):
             query_id = f"probe-{i}"
-            query = Query(
-                id=query_id,
-                data={
-                    "prompt": test_prompt,
-                    "model": model_name,
-                    "max_tokens": 50,
-                    "stream": False,
-                },
+            start_times[query_id] = time.time()
+            client.issue_query(
+                Query(
+                    id=query_id,
+                    data={
+                        "prompt": prompt,
+                        "model": model,
+                        "max_tokens": 50,
+                        "stream": False,
+                    },
+                )
             )
 
-            try:
-                start_times[query_id] = time.time()
-                await client.issue_query_async(query)
-                # Only track successfully issued queries
-                sent_query_ids.append(query_id)
-            except Exception as e:
-                issue_errors.append(f"{query_id}: Failed to issue - {str(e)[:50]}")
-                logger.warning(f"Failed to issue request {i}: {str(e)[:50]}")
-                continue
-
-            # Simple progress indicator
-            if (i + 1) % max(1, num_requests // 10) == 0 or i == num_requests - 1:
-                logger.info(f"  Issued {i + 1}/{num_requests} requests")
-
-        # Wait for all responses
+        # Collect responses
+        logger.info(f"Waiting for {num_requests} responses...")
         latencies: list[float] = []
-        errors: list[str] = issue_errors  # Include any issue errors
+        errors: list[str] = []
         responses: list[tuple[str, str]] = []
+        received: set[str] = set()
 
-        # Only count successfully issued queries
-        num_expected = len(sent_query_ids)
-        if num_expected == 0:
-            logger.error("✗ No queries were successfully issued")
-            raise ExecutionError("Probe failed: no queries could be issued")
-
-        # Wait for all responses with generous timeout (probe queries can be slow)
-        probe_timeout = 60.0  # 60 seconds total
+        timeout = 60.0
         start_wait = time.time()
 
-        logger.info(f"Waiting for {num_expected} responses...")
-
-        received_ids: set[str] = set()
-
-        while (
-            len(received_ids) < num_expected
-            and (time.time() - start_wait) < probe_timeout
-        ):
-            try:
-                result = await client.try_recv_response_async()
-
-                if result is None:
-                    await asyncio.sleep(0.01)
-                    continue
-
-                # Skip non-final streaming chunks
-                if not isinstance(result, QueryResult):
-                    continue
-
-                query_id = result.id
-
-                if query_id in received_ids:
-                    logger.warning(f"Received duplicate response for {query_id}")
-                    continue
-
-                received_ids.add(query_id)
-
-                # Calculate latency - should always be in start_times for issued queries
-                if query_id not in start_times:
-                    logger.warning(
-                        f"Received response for unknown query_id: {query_id}, skipping"
-                    )
-                    continue
-                latency_ms = (time.time() - start_times[query_id]) * 1000
-
-                if result.error:
-                    errors.append(f"{query_id}: {result.error}")
-                else:
-                    latencies.append(latency_ms)
-                    responses.append((query_id, result.response_output))
-
-                # Simple progress indicator
-                if (
-                    len(received_ids) % max(1, num_expected // 10) == 0
-                    or len(received_ids) == num_expected
-                ):
-                    logger.info(
-                        f"  Processed {len(received_ids)}/{num_expected} responses"
-                    )
-
-            except Exception as e:
-                logger.warning(f"Error receiving response: {str(e)[:50]}")
+        while len(received) < num_requests and (time.time() - start_wait) < timeout:
+            result = await client.try_recv_response()
+
+            if result is None:
                 await asyncio.sleep(0.01)
+                continue
+
+            if not isinstance(result, QueryResult):
+                continue
+
+            if result.id in received:
+                continue
+
+            received.add(result.id)
+            latency_ms = (time.time() - start_times.get(result.id, time.time())) * 1000
 
-        # Mark any issued but not received as timeout
-        for query_id in sent_query_ids:
-            if query_id not in received_ids:
-                errors.append(f"{query_id}: Timeout (>{probe_timeout}s)")
+            if result.error:
+                errors.append(f"{result.id}: {result.error}")
+            else:
+                latencies.append(latency_ms)
+                responses.append((result.id, result.response_output))
 
-        # Report results
+        # Report timeouts
+        for query_id in start_times:
+            if query_id not in received:
+                errors.append(f"{query_id}: Timeout")
+
+        # Results
         success_count = len(latencies)
-        logger.info(f"✓ Completed: {success_count}/{num_expected} successful")
+        logger.info(f"✓ Completed: {success_count}/{num_requests} successful")
 
         if latencies:
-            avg_latency = sum(latencies) / len(latencies)
-            logger.info(f"✓ Avg latency: {avg_latency:.0f}ms")
+            logger.info(f"✓ Avg latency: {sum(latencies) / len(latencies):.0f}ms")
             logger.info(f"✓ Range: {min(latencies):.0f}ms - {max(latencies):.0f}ms")
 
-        # Show sample responses for sanity check
         if responses:
             logger.info(f"✓ Sample responses ({len(responses)} collected):")
-            # Show first 10 responses
             for query_id, response in responses[:10]:
-                # Truncate long responses
-                response_preview = (
-                    response[:100] + "..." if len(response) > 100 else response
-                )
-                logger.info(f"  [{query_id}] {response_preview}")
+                preview = response[:100] + "..." if len(response) > 100 else response
+                logger.info(f"  [{query_id}] {preview}")
 
         if errors:
             logger.warning(f"⚠ Errors: {len(errors)}")
-            if args.verbose:
+            if getattr(args, "verbose", 0):
                 for error in errors[:3]:
                     logger.warning(f"  {error}")
-                if len(errors) > 3:
-                    logger.warning(f"  ... +{len(errors) - 3} more")
 
-        # Check if probe was successful
         if success_count < num_requests * 0.5:
-            logger.error("✗ Probe failed: Too many errors")
             raise ExecutionError(
-                f"Probe failed: only {success_count}/{num_requests} requests successful"
+                f"Probe failed: {success_count}/{num_requests} successful"
             )
 
         logger.info("✓ Probe successful")
 
     except ExecutionError:
-        # Re-raise our own exceptions
         raise
     except Exception as e:
-        logger.error("✗ Probe failed")
-        raise SetupError(f"Probe setup failed: {e}") from e
+        raise SetupError(f"Probe failed: {e}") from e
     finally:
-        # Cleanup
-        if client is not None:
-            await client.shutdown_async()
+        if client:
+            client.shutdown()
         shutil.rmtree(tmp_dir, ignore_errors=True)