chore(http-client): cleanup types, improve coverage, remove orjson

[viraatc]

What does this PR do?

(1) before:

$ -> python scripts/benchmark_httpclient.py --duration 10
Final Results:
  Send Duration:  10.00s
  Total Duration: 10.00s
  Total Sent:     1,510,500
  Total Recv:     1,510,500
  Errors:         0
  Outstanding:    0
  Send Rate:      151,047 req/s
  Recv Rate:      151,045 resp/s

(2) before, with asserts disabled:

$ -> python -O scripts/benchmark_httpclient.py --duration 10
Final Results:
  Send Duration:  10.00s
  Total Duration: 10.00s
  Total Sent:     1,565,600
  Total Recv:     1,565,600
  Errors:         0
  Outstanding:    0
  Send Rate:      156,554 req/s
  Recv Rate:      156,553 resp/s

after: (within run-to-run variation of (2))

$ -> python scripts/benchmark_httpclient.py --duration 10
Final Results:
  Send Duration:  10.00s
  Total Duration: 10.00s
  Total Sent:     1,605,600
  Total Recv:     1,605,600
  Errors:         0
  Outstanding:    0
  Send Rate:      160,555 req/s
  Recv Rate:      160,554 resp/s

measurements made using standalone MR: #122

Type of change

• Bug fix
• New feature
• Documentation update
• Refactor/cleanup

Related issues

Testing

• Tests added/updated
• All tests pass locally
• Manual testing completed

Checklist

• Code follows project style
• Pre-commit hooks pass
• Documentation updated (if needed)

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[gemini-code-assist]

Summary of Changes

Hello @viraatc, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request delivers a significant refactor and cleanup of the HTTP client, primarily addressing type hinting inconsistencies and improving the overall robustness and performance. The changes streamline the client's internal mechanisms, particularly around connection handling and event loop management, leading to measurable performance gains. Additionally, comprehensive new tests have been introduced to ensure the correct behavior of core client functionalities like polling, receiving, draining, and shutting down.

Highlights

• Type Cleanup and Refactoring: Refactored HTTP client types for clarity and correctness, removing Optional from fields that are guaranteed to be initialized in __post_init__ or run() methods. This includes adapter, accumulator, worker_pool_transport in config.py and various internal components in worker.py.
• Performance Improvements: Improved the performance of the HTTP client, as evidenced by benchmark results showing increased request/response rates (e.g., from 1.56M req/s to 1.60M req/s). This was achieved through various optimizations, including removing redundant assert statements and streamlining internal logic.
• Enhanced Test Coverage: Significantly enhanced test coverage for core HTTP client functionalities, including poll(), recv(), drain(), and shutdown() methods, ensuring their correct behavior and robustness.
• Streamlined Connection Handling: Streamlined internal worker logic by passing PooledConnection objects directly to response handling methods (_process_response, _handle_streaming_body, _handle_non_streaming_body), reducing reliance on req.connection and simplifying the flow.
• Robust Shutdown Mechanism: Implemented an explicit and idempotent shutdown() method for the HTTP client, ensuring graceful termination of workers and event loops, and handling requests dropped during shutdown.
• Simplified Drain Method: The drain() method was simplified using iter(self.poll, None) for more concise and potentially efficient retrieval of all available responses.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• src/inference_endpoint/endpoint_client/config.py
  • Updated type annotations for adapter, accumulator, and worker_pool_transport to reflect their guaranteed initialization in __post_init__.
• src/inference_endpoint/endpoint_client/http_client.py
  • Refactored event loop management, removing redundant _loop_thread and simplifying self.loop assignment.
  • Removed several assert statements, improving code clarity and potential runtime performance.
  • Simplified the drain() method using iter(self.poll, None).
  • Implemented a robust and idempotent shutdown() method for graceful client termination.
  • Removed the issue() method override in HTTPEndpointClient.
• src/inference_endpoint/endpoint_client/http_sample_issuer.py
  • Removed an unnecessary assert statement related to the event loop.
• src/inference_endpoint/endpoint_client/worker.py
  • Cleaned up type hints for internal components (_pool, _http_template, _loop, _requests, _responses), marking them for assignment in run().
  • Eliminated multiple assert statements for initialized attributes.
  • Refactored _fire_request to return PooledConnection and updated _process_response, _handle_streaming_body, and _handle_non_streaming_body to accept the connection object directly.
  • Added a more descriptive error message to an assert for event_logs_dir.
• tests/integration/endpoint_client/test_http_client.py
  • Added helper functions _create_client and _make_query for test setup.
  • Introduced a new http_client fixture for consistent client initialization.
  • Expanded the test suite with new classes TestPoll, TestRecv, TestDrain, and TestShutdown to cover these client functionalities comprehensively.

Activity

• The author, viraatc, initiated this pull request with the primary goal of refactoring and cleaning up the HTTP client, as well as improving test coverage.
• Benchmark results were provided in the PR description, demonstrating performance improvements after the changes.
• The PR includes new and updated tests, indicating a focus on verifying the correctness and robustness of the refactored code.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[Copilot]

Pull request overview

Refactors the HTTP client types/initialization to reduce runtime asserts, and expands integration coverage for non-blocking/buffering APIs (poll/recv/drain/shutdown).

Changes:

• Added new integration tests covering poll(), recv(), drain(), and shutdown idempotency/behavior.
• Refactored worker/client/config typing by making several fields non-optional and removing some runtime asserts.
• Simplified drain() implementation and adjusted shutdown flow/idempotency checks.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
tests/integration/endpoint_client/test_http_client.py	Adds new integration coverage for poll/recv/drain/shutdown and introduces shared client/query helpers.
src/inference_endpoint/endpoint_client/worker.py	Removes several runtime asserts and changes request firing to pass pooled connections explicitly.
src/inference_endpoint/endpoint_client/http_sample_issuer.py	Drops an assert that is no longer needed due to non-optional loop typing.
src/inference_endpoint/endpoint_client/http_client.py	Simplifies loop/thread setup, makes drain() more concise, and tweaks shutdown idempotency behavior.
src/inference_endpoint/endpoint_client/config.py	Makes adapter/accumulator/transport non-optional (defaulted in __post_init__).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[gemini-code-assist]

Code Review

This pull request introduces several cleanups to type hinting, which simplifies the code by removing numerous assert is not None checks. It also significantly improves test coverage for the HTTP client's response retrieval and shutdown mechanisms. While most changes are positive, I've identified a couple of concerns regarding thread safety and shutdown logic that could introduce bugs. My review includes specific suggestions to address these potential issues.

I am having trouble creating individual review comments. Click here to see my feedback.

src/inference_endpoint/endpoint_client/http_client.py (163-169)

This issue method override was crucial for thread safety and should not have been removed. HTTPEndpointClient is a synchronous wrapper intended for use from other threads. The underlying transport sockets are not thread-safe, so calls must be scheduled on the event loop thread. This implementation correctly used loop.call_soon_threadsafe to achieve that. Removing it introduces a race condition. Please restore this method.

src/inference_endpoint/endpoint_client/http_client.py (126-143)

Moving self._shutdown = True to the end of the method creates a race condition. During the await self.worker_manager.shutdown() call, new requests can still be accepted by issue() because self._shutdown is False. These requests will likely fail ungracefully as workers are terminated, and they won't be counted as dropped. To ensure new requests are gracefully rejected during shutdown, self._shutdown = True should be set at the beginning of the method, immediately after the idempotency check.

        if self._shutdown:  # Already shutdown, no-op
            return
        self._shutdown = True

        logger.info(f"[{self.client_id}] Shutting down...")

        # Shutdown workers
        await self.worker_manager.shutdown()

        # Stop event loop if we own it (scheduled to run after this coroutine completes)
        if self._owns_loop:
            self.loop.call_soon(self.loop.stop)

        if self._dropped_requests > 0:
            logger.info(
                f"[{self.client_id}] Dropped {self._dropped_requests} requests during shutdown"
            )
        logger.info(f"[{self.client_id}] Shutdown complete.")

[viraatc]

info:

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 6 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 15 out of 15 changed files in this pull request and generated 1 comment.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 14 out of 14 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (1)

src/inference_endpoint/endpoint_client/http_client.py:174

• HTTPEndpointClient no longer overrides issue() to schedule onto the client’s event-loop thread. AsyncHttpEndpointClient.issue() now calls self.pool.send(...) directly from the caller thread, but the underlying ZMQ sender can call loop.add_writer(...) (not thread-safe). This can cause races or runtime errors under backpressure; restore a thread-safe issue() wrapper (e.g., loop.call_soon_threadsafe) for the sync client.

class HTTPEndpointClient(AsyncHttpEndpointClient):
    """
    Sync HTTP client for LLM inference.
    Inherits from AsyncHttpEndpointClient and provides sync interface.

    Usage:
        client = HTTPEndpointClient(config)
        client.issue(query)
    """

    def shutdown(self) -> None:  # type: ignore[override]
        """Sync shutdown wrapper - blocks until base class async shutdown completes."""
        if self._shutdown:  # Already shutdown, no-op
            return
        asyncio.run_coroutine_threadsafe(super().shutdown(), self.loop).result()

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[nv-alicheng]

Changes look like they do what what the PR is about but what is the reasoning for removing orjson? Is msgpack encode faster or just that it fits into the pipeline more since we use msgpack structs throughout the project?

[viraatc]

Changes look like they do what what the PR is about but what is the reasoning for removing orjson? Is msgpack encode faster or just that it fits into the pipeline more since we use msgpack structs throughout the project?

ORjson doesnt support py3.14t (GILL-Free builds) as of yet, msgspec already does.
This allows building/testing the inference-endpoints on all py3.14t.

[Copilot]

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated 3 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (1)

src/inference_endpoint/metrics/recorder.py:400

• When JSON encoding fails, the recorded error payload uses "error_type": "JSONEncodeError", but the exception being caught is now msgspec.EncodeError. This makes logs harder to interpret; consider updating error_type to match the actual exception/library (e.g., msgspec.EncodeError).

                    msgspec.json.encode(
                        {
                            "error_type": "JSONEncodeError",
                            "error_message": str(e),
                        }

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[arekay-nv]

LGTM.

[viraatc]

continuing in #159