Memory leak in the chat completion `create` method

[CodePalAI]

Confirm this is an issue with the Python library and not an underlying OpenAI API

• This is an issue with the Python library

Describe the bug

Calling the create method on completions introduces a memory leak, according to tracemalloc.

Example:

client = OpenAI(MY_KEY)
client.chat.completions.create(
    messages=[
        {
            "role": "user",
            "content": "Can you write a poem?",
        }
    ],
    model="gpt-3.5-turbo"
)

How to determine it's a memory leak?

I use tracemalloc with my flask application:

@blueprint.route("/admin/sys/stats")
def admin_sys_stats():
    snapshot = tracemalloc.take_snapshot()
    top_stats = snapshot.statistics('lineno')

    from openai import OpenAI

    client = OpenAI(KEY)
    client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": "Can you write a poem?",
            }
        ],
        model="gpt-3.5-turbo"
    )

    stats = ""
    for stat in top_stats[:1000]:
        if grep in str(stat):
            stats += str(stat) + "\n"

    return f"<pre>{stats}</pre>", 200

When running this endpoint multiple times, one line is at the very top (which means it's the most expensive one):

\venv\Lib\site-packages\openai\_response.py:227: size=103 KiB, count=1050, average=100 B

When I refresh, the size increases. Of course, in a production environment, the numbers get high a lot quicker.

To Reproduce

There's no one way to prove there's a memory leak.
But what I did was:

1. Setup a flask application
2. Create the route provided in the bug description
3. Hit the route multiple times, you'll see an increase in the size of the object

Code snippets

No response

OS

Linux, macOS

Python version

Python v3.11.2

Library version

openai v1.2.4

[RobertCraigie]

You shouldn't be instantiating a new client for every request, if you move the client outside of the request handler you shouldn't see any memory leaks.

I did manage to replicate the memory leak you're seeing with your example and moving the client instantiation outside of the function resulted in stable memory usage.

[qxcv]

For what it's worth, I still see this issue when wrapping @CodePalAI's OpenAI client in a with block, which suggests there is some deeper problem here. I've been trying to track down a leak in my app (which existed with the pre-1.0 API, and has gotten much more aggressive after upgrading to 1.3) and I also ran into this issue.

[Filimoa]

Wouldn't "moving the client outside of the request handler" be the same thing as a module level client?

From the README, I figured the network resources are getting cleaned up when the OpenAI client gets garbage collected. Sounds like this isn't necessarily true?

This might be outside of the scope of this library but I'm also trying to understand how this fits in your typical backend app.

It seems like using a module scoped client is heavily discouraged but so is instantiating a new client with every request (ie a FastAPI dependency). What is the other option here?

[brian-goo]

@RobertCraigie

How can we implement an async client with fastapi correctly?

Is it safe to implement like this (singleton approach) ?
fastapi/fastapi#8301 (comment)

Another problem is that we have multiple api keys like below.
How can we handle this nicely if we have to instantiate a client once as mentioned in this issue #874?

    client = AsyncAzureOpenAI(
        api_key=api_keys_list[random_index],
        ...

[RobertCraigie]

@brian-goo you can create a single client and use FastAPI app events to close it when the server is stopped and then use client.with_options(api_key='foo') to use different API keys for different requests.

[RobertCraigie]

@Filimoa what are you trying to achieve that requires you to instantiate a new client for every request?

If you really want to do so you currently need to use client.with_options() or create a custom http client and pass that to the constructor every time to avoid resource leaks.

Instantiating a new http client for each request is very bad for performance regardless of whether or not it has any resource leaks as the client can not reuse connections.

[abi]

@RobertCraigie one use case I have is with https://github.com/abi/screenshot-to-code where users put in their Open AI key on the hosted version (that's how the app is free to use). I understand that performance will be worse with instantiating a new client on each request but a memory leak is even worse.

Here's a fun memory leak graph from my server :) Not 100% sure it's because of this issue but very likely.

[RobertCraigie]

@abi are you explicitly closing the client using client.close()? If you are then I would consider that a bug.

Additionally, is there a reason that client.with_options(api_key=user_api_key) doesn't work for you? That should give you all of the benefit of allowing users to specify their own API keys without creating a ton of unnecessary connections.

[abi]

@RobertCraigie I'm not closing with client.close(): https://github.com/abi/screenshot-to-code/blob/main/backend/llm.py#L26 Will try that now and report back. Thanks for the help.

The main reason client.with_options(api_key=user_api_key) is non-ideal even though it would work is that it's a bit harder to debug and error-prone when you're not starting from a clean, empty state. Just now, I even tried doing an empty AsyncOpenAI() on the module level, and setting api_key using client.with_options(api_key=user_api_key). For some reason, even before client.with_options(api_key=user_api_key), client.api_key had a value and then, I realized it was pulling an initial value from my .env. Much more straightforward to do AsyncOpenAI(api_key) and know for sure that which api_key the client is associated with. It's a bit scary to potentially serve 1 user's request with another user's api key, even if that can only happen due to a bug in my code. The ergonomics of AsyncOpenAI(api_key=) are just nicer but not a big deal if .close() fixes my memory leak.

[antont]

@abi AsyncOpenAI() on the module level, and setting api_key using client.with_options(api_key=user_api_key). For some reason, even before client.with_options(api_key=user_api_key), client.api_key had a value and then, I realized it was pulling an initial value from my .env. Much more straightforward to do AsyncOpenAI(api_key) and know for sure that which api_key the client is associated with. It's a bit scary to potentially serve 1 user's request with another user's api key, even if that can only happen due to a bug in my code.

Interesting point. Maybe there could be an option to disable the use of of that default option (what about others?), so that the client would err if you forget to provide the key using with_options in some place. Perhaps a bit awkward but would work to safeguard against such mistakes quite ok, no?

[RobertCraigie]

It's a bit scary to potentially serve 1 user's request with another user's api key, even if that can only happen due to a bug in my code

This is a fair point, I do think there are ways you could structure your code to make this practically impossible to happen though. Even for anyone making changes to your codebase for the first time.

For example, instead of using .with_options() yourself, you could define a helper function that you use whenever you need to make an OpenAI request, e.g.

# utils/openai_client.py
from openai import AsyncOpenAI

_client = AsyncOpenAI(
  # provide a dummy API key so that requests made directly will always fail
  api_key='<this client should never be used directly!>',
)

def get_openai(user: User) -> AsyncOpenAI:
  return _client.with_options(api_key=user.openai_api_key)

# backend.py

async def handler():
  user = get_user()

  completion = await get_openai(user).chat.completions.create(...)
  # ...

It's also worth noting that you could use your existing pattern and avoid memory leaks entirely by sharing the http client instance if you really don't want to use with_options, the only risk here is that you could instantiate a client and forget to pass in your http client which then gives you memory leaks again.

import httpx
from openai import AsyncOpenAI

http_client = httpx.AsyncClient()

async def handler(api_key):
  openai_client = AsyncOpenAI(api_key=api_key, http_client=http_client)

---

I would argue though that the first snippet in this comment is less likely to result in accidentally using the incorrect API key because you could accidentally just write AsyncOpenAI() and then it will use your API key from your environment.

However with a helper function similar to get_openai() you're encoding the setup of your system directly into it so instead of having to know to pass api_key=user_api_key you're forced to by your own function definition.

The only way you'll accidentally use the wrong API key is by forgetting to use this helper function, which it may be possible to define lint rules for (I don't know of an easy way to do this off the top of my head).

[abi]

@RobertCraigie that's super helpful, both of those patterns. Didn't realize you could pass in the http_client for the client.

I pushed client.close() a couple of hours ago BTW, and that solved my memory leak issues. Thanks for your help on this!