Test for openai memory leaks

openai_test.py

"""

1. install openai and memory_profiler

2. Add following to class Completions(SyncAPIResource):

    def forceClose(self):
        self.with_raw_response = None

3. Add following to class class OpenAI(SyncAPIClient):

    def forceClose(self):
        if self.completions:
            self.completions.forceClose()
            self.completions = None


4. Comment out all of the OpenAI(SyncAPIClient) class initializations that have self reference, except for the `self.completions = resources.Completions(self)`

5. Now run `python3 openai_test.py` - the output will show increasing memory usage. It will stay the same for a while, and then increase again and never really be cleaned up.
6. Now uncomment `aiClient.forceClose()` in `openai_test.py` and run it again. The memory usage will stay the same.


"""


import asyncio
import signal
from openai import OpenAI
from memory_profiler import profile


API_KEY = "XXXX"
ORGANIZATION = "XXX"
shutdownEvent: asyncio.Event


##########
### Ai ###
##########

@profile(precision=4)
def initOpenAi() -> OpenAI:
    aiClient = OpenAI(api_key=API_KEY, organization=ORGANIZATION)

    # Cleanup
    # aiClient.forceClose()


##############
### Runner ###
##############


def shutdownHandler(loop: asyncio.AbstractEventLoop):
    shutdownEvent.set()


async def runner():
    global shutdownEvent

    shutdownEvent = asyncio.Event()

    loop = asyncio.get_running_loop()
    signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT, signal.SIGQUIT)
    for s in signals:
        loop.add_signal_handler(s, lambda s=s: shutdownHandler(loop))

    while not shutdownEvent.is_set():
        # Init ai
        initOpenAi()

        # Sleep
        await asyncio.sleep(1)


def main():
    try:
        asyncio.run(runner())
    except asyncio.CancelledError:
        pass


if __name__ == "__main__":
    main()