Repository: UKGovernmentBEIS/inspect_ai
Issue: CUDA memory leak after closing APIs for local models
Root cause: Resource Mishandling

It appears that the CUDA memory associated with local model providers is not getting fully freed after closing them (I tested this with the vLLM provider, but it seems that the issue might also affect HuggingFace models). Although the model and the tokenizer are being set to None when `close` is called, the `batch_thread` keeps running in the background while maintaining a reference to the `generate` method of the model, which prevents it from being properly garbage collected.

## Reproduction steps
1. Run `eval` with a local model, such as:
```python
from inspect_ai import Task, eval, task
from inspect_ai.dataset import Sample
from inspect_ai.scorer import exact
from inspect_ai.solver import generate


@task
def hello_world():
    return Task(
        dataset=[
            Sample(
                input="Just reply with Hello World",
                target="Hello World",
            )
        ],
        solver=[generate()],
        scorer=exact(),
    )


eval(
    model="vllm/meta-llama/Llama-3.1-8B-Instruct",
    tasks=hello_world(),
)
```
2. Close the model API:
```python
from inspect_ai.model import get_model

model = get_model(
    "vllm/meta-llama/Llama-3.1-8B-Instruct",
)
await model.api.close()
```
3. Check VRAM usage:
```python
import subprocess

print(subprocess.run(["nvidia-smi"], capture_output=True, text=True).stdout)
```

I could be wrong, but I don't think that there are any references left over in the batch thread after it processes things (its input queues have functions it calls and the queues are emptied). It also appears to me as if you need to take lots of special measures to cause vllm to release GPU memory should you want that, e.g. https://github.com/vllm-project/vllm/issues/1908

It's not clear to me though whether you can then call vLLM again after that w/o entanglement in the side effects of the shutdown. 

I am also experiencing this issue and have tried all the methods in https://github.com/vllm-project/vllm/issues/1908 . None seem to work with inspect

It looks like lots of people on that thread have struggled with getting things to work. This comment towards the end seems to have outlined something that worked for at least one user: https://github.com/vllm-project/vllm/issues/1908#issuecomment-2461174904

We are going to need a PR (perhaps from @kaifronsdal who is the author of the vLLM provider?) in order to do anything different here (we don't have a vLLM setup and are not conversant in all of possibilities visa-vi multi-GPU setups). It would probably need to run in the `close()` method as it's certainly possible that all of this deallocation, etc. could make it so that you can't subsequently re-initialize vLLM.


Yes, cleaning up local models in vllm is notoriously difficult and also very dependant on which version of vllm installed (various internal functions to clean up resources have often been moved around as far as I can tell). I did my best when I implemented it, and it worked on my machine with an older version of vllm, but it was never going to be a robust solution.

As a temporary work around you can run their openai compatible server using `subprocess.Popen` to run `vllm serve ...` in the background and then kill the process when you want to use a different model. To then use with inspect, specify the model using the `openai/` prefix and set the base_url and api_key to be those used by vllm (documented [here](https://inspect.aisi.org.uk/providers.html#vllm-server)).

When I was originally implementing vllm models in inspect the openai compatible server was not in as robust a state as it is now. However, over the last year it has improved massively. I have actually been working on an updated version of the vllm provider the last couple of weeks that uses the openai compatible server rather than vllm's offline serving. This gives us a few benefits:
- Using the openai compatible server gives us native support for a lot of the nice features vllm provides such as guided decoding
- Their internal async routing system is much better than the batched routing system I built around their offline serving
- Killing the process properly cleans up all of vllm's resources
- It is possible to load multiple models in parallel on different GPUs, i.e. for llm-as-a-judge
- The code is much cleaner and easier to maintain
- Gives the option to load models on demand like the current implementation, and also to connect to an existing vllm server so you don't need to wait for the model to get loaded onto the GPU every time

The only remaining road block is that there are currently some issues with zombie processes when inspect exits since as far as I can tell, inspect does not call the model provider's close method (except when using `async with get_model` which would result in starting up and stopping the server for every request). I've currently just added a few lines to close the model providers at the end of `eval_async`, but perhaps there's a better place to ensure things are closed correctly?

Thank you for your replies! Yes, the ability to spawn and use a separate vLLM server would be very useful as a long-term solution, especially if Inspect could stop the server and free the associated GPU memory once done with the evaluation of the given model.

Regarding the current version of the provider, I had some success using the following workaround to reclaim the occupied VRAM after running `eval`:
```python
# Cleanup background processes
main_id = os.getpid()
parent = psutil.Process(main_id)
children = parent.children(recursive=True)
for child in children:
    try:
        child.terminate()
    except psutil.NoSuchProcess:
        pass
gone, still_alive = psutil.wait_procs(children, timeout=5)

if still_alive:
    warnings.warn("Failed to cleanup background processes.")
```
While this solution is somewhat hacky and potentially brittle, it seems to be working reasonably well with the latest version of vLLM, and doesn't appear to interfere with running `eval` again on a different vLLM model afterwards.


> The only remaining road block is that there are currently some issues with zombie processes when inspect exits since as far as I can tell, inspect does not call the model provider's close method (except when using `async with get_model` which would result in starting up and stopping the server for every request). I've currently just added a few lines to close the model providers at the end of `eval_async`, but perhaps there's a better place to ensure things are closed correctly?

If you register a standard Python `atexit` handler (https://docs.python.org/3/library/atexit.html) it will be called in all normal shutdown scenarios (including Ctrl+C interruptions). Just call `atexit.register` with the function that does the cleanup.


> If you register a standard Python `atexit` handler (https://docs.python.org/3/library/atexit.html) it will be called in all normal shutdown scenarios (including Ctrl+C interruptions). Just call `atexit.register` with the function that does the cleanup.

This would work in most cases. However in the scenario posed above, where we run `eval` within a python script rather than the cli `inspect eval`, the vllm model would not be unloaded when `eval` finishes. Perhaps this is the behavior we want as loading a model onto the GPU can take awhile depending on the size of the model -- i.e. we only want to unload a model when explicitly told to. Thoughts?

Another approach would be that you use the model context manager and pass the `model` explicitly to `eval()` (that way it would be released only when eval is done). 

I like this solution, but there seems to be some implementation issues with anyio/asyncio.

In order for this to work we would need to modify the`Model` class to include synchronous `__enter__` and `__exit__` methods, since right now `async with get_model` needs to be run within an async function (which is also awkward to use):

```python
async def f():
    async with get_model("vllm/meta-llama/Llama-3.1-8B-Instruct") as model:
        await eval_async(
            model=model,
            tasks=hello_world(),
        )

import asyncio
asyncio.run(f())
```
Note the use of `eval_async` since eval calls `anyio.run` internally.

This seems to work within jupyter notebooks, but not scripts:

```python
async with get_model("vllm/Qwen/Qwen2.5-7B-Instruct") as model:
    await eval_async(
        model=model,
        tasks=hello_world(),
    )
```

On the flip side, if we implement something like this in `Model`:

```python
def __enter__(self: "Model") -> "Model":
    return anyio.run(self.__aenter__)

def __exit__(
    self,
    exc_type: type[BaseException] | None,
    exc: BaseException | None,
    exc_tb: TracebackType | None,
) -> None:
    anyio.run(self.__aexit__, exc_type, exc, exc_tb)
```
and run
```python
with get_model("vllm/meta-llama/Llama-3.1-8B-Instruct") as model:
    eval(
        model=model,
        tasks=hello_world(),
    )
```

we get `RuntimeError: Event loop is closed` when the openai provider tries to close its client.

I'm not the most familiar with the anyio/asyncio packages, so perhaps there is an easy solution?

Here's an approach to making this work:

https://github.com/UKGovernmentBEIS/inspect_ai/pull/1563/files

We basically also implement a sync context manager, but with the constraint that some models require an async context manager (and `raise` if they are used from a sync one). With this it should be possible to do:

```python
with get_model("vllm/Qwen/Qwen2.5-7B-Instruct") as model:
    eval(model=model, tasks=hello_world())
```

Could you give this a try and LMK if it works well with your `close()` implementation?

@adamdejl Pinging you back to see if you've had a chance to try out this branch?

> @adamdejl Pinging you back to see if you've had a chance to try out this branch?

I think you probably meant to ping @kaifronsdal, as it doesn't seem there is much I can try on my end without the code for the updated OpenAI-compatible vLLM server provider. But happy to help with testing something if I misunderstood.

Sorry yes I meant @kaifronsdal (apologise for the misdirection!)

This should be resolved in the context manager work as well as the forthcoming updates to the vLLM provider.