Bug Report: AttributeError when closing the httpx.AsyncClient

[giulio-utf]

Description: When trying to close the httpx.AsyncClient in an asynchronous load balancing setup, an AttributeError is raised:

AttributeError: 'AsyncLoadBalancer' object has no attribute 'aclose'

This happens because the custom AsyncLoadBalancer object does not implement the aclose() method, which is required when closing the transport in httpx.AsyncClient.

Steps to Reproduce:

1. Setup an AsyncLoadBalancer as the transport for httpx.AsyncClient.
2. Make an async request using the client.chat.completions.create() method.
3. Try to close the httpx.AsyncClient with await async_client.aclose().
4. Observe the traceback showing an AttributeError.

Expected Behavior: The async_client.aclose() should close without raising an exception, properly releasing resources used by the transport (in this case, the AsyncLoadBalancer).

Actual Behavior: An exception is raised when trying to close the httpx.AsyncClient because the custom load balancer (AsyncLoadBalancer) lacks the aclose() method.

Traceback:

Traceback (most recent call last):
  File "/workspaces/NeuralExtractor/tests/load_balancer.py", line 50, in <module>
    asyncio.run(main())
  File "/usr/local/lib/python3.11/asyncio/runners.py", line 190, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/asyncio/base_events.py", line 654, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/workspaces/NeuralExtractor/tests/load_balancer.py", line 46, in main
    await async_client.aclose()
  File "/workspaces/NeuralExtractor/src/quartapp/.venv/lib/python3.11/site-packages/httpx/_client.py", line 2018, in aclose
    await self._transport.aclose()
          ^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'AsyncLoadBalancer' object has no attribute 'aclose'

Possible Cause: The AsyncLoadBalancer class does not implement the aclose() method required by httpx.AsyncClient to close the transport gracefully.

Suggested Fix: Implement the aclose() method in the AsyncLoadBalancer class to close any resources properly when shutting down the client. Here's an example of how this might be added:

class AsyncLoadBalancer:
    # Your current methods and attributes

    async def aclose(self):
        # Close any active connections or resources here
        pass

Additional Context: This issue is observed when working with OpenAI's AsyncAzureOpenAI client and httpx in combination with custom load balancing logic. Properly closing the client is necessary to prevent resource leaks or unhandled exceptions during shutdown.

[simonkurtz-MSFT]

Hi @giulio-utf, thank you very much for the detailed issue! I'll take a look at it. You are also very welcome to submit a PR for it, of course.

[giulio-utf]

Hey @simonkurtz-MSFT, thanks for the quick response!

I have limited firepower to dig into the issue at the moment, hope that this code example help reproducing the error:

from openai import AsyncAzureOpenAI
import httpx
import asyncio
from openai_priority_loadbalancer import AsyncLoadBalancer, Backend

# Define the backend
backends = [
    Backend("XXXXXXXXXXXXXXXX.openai.azure.com", 1, None, 'XXXXXXXXXXXXXXXX')
]

# Create the async load balancer
lb = AsyncLoadBalancer(backends)

# Create the Azure OpenAI client
async_client = httpx.AsyncClient(transport=lb)
client = AsyncAzureOpenAI(
    azure_endpoint=f"https://{backends[0].host}",  # Using the first backend's host
    api_key="obtain_from_load_balancer",  # This value isn't used but must be set
    api_version="2024-09-01-preview",
    http_client=async_client
)

# Example async chat completion request
async def get_completion(prompt):
    try:
        response = await client.chat.completions.create(
            model="gpt-4o-2",  # Your deployed model name
            messages=[
                {"role": "user", "content": prompt}
            ]
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error: {str(e)}"

# Main async function to run everything
async def main():
    try:
        prompt = "What is the capital of France?"
        print('Starting request...')
        result = await get_completion(prompt)
        print(f"Prompt: {prompt}")
        print(f"Response: {result}")
    finally:
        # Close the httpx client correctly
        await async_client.aclose()

# Test the setup
if __name__ == "__main__":
    asyncio.run(main())