memory leak, consumes entire system after just days of usage

[pseudotensor]

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

Progressive memory leak despite use of close(). Leads to eventually OOM on even large memory systems after just a few days.

Related: #820

To Reproduce

I'm using 1.12.0 and still hit this issue.

Using text completion. I'm also using close() in a try-finally, so close() does not help. A global client connection doesn't make sense to enforce. The old client pre 1.x had global attributes that made the API poor. I presume there are some legacy parts still in place.

In h2oGPT, we use OpenAI client for OpenAI or vLLM connections, and I see a 5GB memory leak for every 6000 connections. This happens whether I yield the generator for streaming or just exit after creating the completion.

Normally connections are not as intense, but this was easily reproducible by bisecting the OpenAI creation/generation parts of the code. For typical workloads this leads to OOM on a 256GB system after just few days of usage.

Here is repro. Please choose the to be some endpoint that you have setup like vLLM or TGI or gpt3.5 turbo so not expensive. Choose api_key and model accordingly.

import os
import psutil
from openai import OpenAI

for i in range(6000):

    client_args = dict(base_url='<choose>', api_key="EMPTY")
    client = OpenAI(**client_args)

    responses = client.completions.create(
        model='h2oai/h2ogpt-4096-llama2-13b-chat',
        prompt="Say exactly one word.",
        stream=True,
    )
    client.close()
    p = psutil.Process(os.getpid())
    print(p.memory_full_info())

The memory consumed is not increasing every step in loop, but it does monotonically increase from pss=48523264 to pss=107862016 within a few minutes (i.e. doubled) and continues this indefinitely.

The problem seems to be even stronger when doing concurrent requests in multi-threaded setup, as if the clean-up is not thread safe. I'm trying to put together a repro that would showcase the 5GB after 6000 connections that only takes half hour to run. But perhaps the above is sufficient.

Code snippets

No response

OS

ubuntu 22

Python version

Python v3.10

Library version

openai v1.12.0

[pseudotensor]

Note that the fix in 1.3.9 helped this issue to be a leak of 5GB to about 1GB, but it's still a large leak.

#820 (comment)

[rattrayalex]

You need to call response.close() if you stream a request without reading its body.

Please try this:

import os
import psutil
from openai import OpenAI

for i in range(6000):

    client_args = dict(base_url='<choose>', api_key="EMPTY")
    client = OpenAI(**client_args)

    response = client.completions.create(
        model='h2oai/h2ogpt-4096-llama2-13b-chat',
        prompt="Say exactly one word.",
        stream=True,
    )
    response.close()
    client.close()
    p = psutil.Process(os.getpid())
    print(p.memory_full_info())

[pseudotensor]

Ok thanks, will try that.

[pseudotensor]

I see same thing with the added close. Note that this is just trying to mimic a leak seen in similar circumstances with much larger code base.

base_url = <fill me>
model = <fill me>
api_key = <fill me>


import os
import psutil
from openai import OpenAI

for i in range(6000):

    client_args = dict(base_url=base_url, api_key=api_key)
    client = OpenAI(**client_args)

    responses = client.completions.create(
        model=model,
        prompt="Say exactly one word.",
        stream=True,
    )
    responses.close()
    client.close()
    p = psutil.Process(os.getpid())
    print(p.memory_full_info())

i.e. I still see it go from pss=52678656 to pss=110225408 within few minutes, about 50MB, and it keeps growing indefinitely.

[rattrayalex]

Got it, thanks for the report & repro script. We'll look into it.

[pseudotensor]

One can do a more attack style with multi-threading to get things to happen faster:

base_url = <fill me>
model = <fill me>
api_key = <fill me>

import concurrent.futures
import os
import psutil
from openai import OpenAI


def get_mem():
    p = psutil.Process(os.getpid())
    print(p.memory_full_info())


def do_something(i):
    client_args = dict(base_url=base_url, api_key=api_key)
    client = OpenAI(**client_args)

    responses = client.completions.create(
        model=model,
        prompt='\n'.join(["Say exactly one word."] * 300),
        stream=True,
    )
    responses.close()
    client.close()
    get_mem()


# Create a thread pool with 4 worker threads.
tasks = 6000
#tasks = 600
futures = []
with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
    for i in range(tasks):
        futures.append(executor.submit(do_something, i))

        # Wait for all tasks to complete.
        for f1 in futures:
            if f1.done():
                result = f1.result()
                print(result)
get_mem()

I'll open new issue if I'm able to repro the larger code base more accurately with its larger leak for this block.

[rattrayalex]

This issue is fine for this leak, no need to open another new one :)

[shivenmian]

@rattrayalex facing the same issue at 01. Repro-d on an M2 Macbook Pro with openai version 1.12.0. Things were working fine yesterday, but now program execution runs into out-of-memory errors every time I run it.

I did a memray memory profile for 01 here. Looks like the use of Pydantic is taking up an excessive amount of memory. @pseudotensor perhaps you can also do a memprof and see if you see the same thing.

cc @KillianLucas