[Bug]: non-deterministic Python gc order leads to flaky tests

[youkaichao]

[Bug]: non-deterministic Python gc order leads to flaky tests

[DarkLight1337]

After some offline discussion, we have settled on the following code:

class HfRunner:
    def __init__(self, ...):
        self.model = create_model(...)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        del self.model
        cleanup()

Although HfRunner can only be used once as a context manager, it is sufficient for testing. This passes the type checker without needlessly complicating the existing code.

[robertgshaw2-neuralmagic]

Love this!

[rkooo567]

Actually the destructor of hf_runner should call cleanup() which calls gc.collect()

vllm/tests/conftest.py

Line 359 in 388596c

cleanup()

Is this not sufficient?

[DarkLight1337]

Actually the destructor of hf_runner should call cleanup() which calls gc.collect()

vllm/tests/conftest.py

Line 359 in 388596c

cleanup()

Is this not sufficient?

__del__ is only called when the object is actually GC'ed, not when the del statement is used.

In any case, I think the context manager approach is cleaner. However, this will require refactoring the runner classes to only initialize the model when the context manager is entered.

[DarkLight1337]

Btw in general, we should prefer weakref.finalize over __del__. See this documentation for more details.

[youkaichao]

@rkooo567 I updated the code and the output. __del__ is not enough here. You can see that A.__del__ calls del self.tensor. However, the tensor is destructed before A.__del__ is called. That's why Python gc is not reliable, especially when we need to enforce some destruction order.

[youkaichao]

In any case, I think the context manager approach is cleaner. However, this will require refactoring the runner classes to only initialize the model when the context manager is entered.

@DarkLight1337 not necessary. We can make the runner itself a context manager. e.g. adding the following lines:

class HfRunner:
    ...
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        del self
        cleanup()

Then the following change should be enough:

-    hf_model = hf_runner(model, dtype=dtype)
-    hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
-    del hf_model
+    with hf_runner(model, dtype=dtype) as hf_model:
+         hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)

We should avoid using weakref.finalize or __del__ . They are just not reliable.

[DarkLight1337]

class HfRunner:
    ...
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        del self
        cleanup()

I don't think this will work. In particular, del self does not really do anything since the HfRunner is still inside the scope of the test function. We should instead operate on the wrapped model (self.model) directly.

[DarkLight1337]

We should avoid using weakref.finalize or __del__. They are just not reliable.

Yeah, they are both dependant on GC so they aren't suitable in situations where GC is not even being invoked.

[youkaichao]

I don't think this will work. In particular, del self does not really do anything since the HfRunner is still inside the scope of the test function. We should instead operate on the wrapped model (self.model) directly.

Good point. How about del self.model inside __exit__ ?

[DarkLight1337]

I don't think this will work. In particular, del self does not really do anything since the HfRunner is still inside the scope of the test function. We should instead operate on the wrapped model (self.model) directly.

This is what I mean:

class HfRunner:
    def __init__(self, ...):
-       self.model = create_model(...)
+       self.model = None

    def __enter__(self):
+       self.model = create_model(...)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
-       del self
+       self.model = None
        cleanup()

There is no need to del self.model explicitly. Setting self.model = None should be sufficient to remove the reference to the HuggingFace model (or LLM object for VllmRunner).

[youkaichao]

I think they are equivalent. For the following code:

   with hf_runner(model, dtype=dtype) as hf_model:
        hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)

The function call chain is:

• hf_model = hf_runner() # essentially HfRunner. __init__
• hf_model.__enter__()
• hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
• hf_model. __exit__()

[DarkLight1337]

To avoid having to check self.model is not None each time in the other methods, it may be better to have a separate context manager function that returns the runner object. For example:

@contextlib.contextmanager
def hf_model_runner(*args, **kwargs):
    model = create_model(...)
    yield HfRunner(model, ...)
    del model
    cleanup()

class HfRunner:
    def __init__(self, model, ...):
        self.model = model
        # The rest is the same as original HfRunner without context manager. Also we can now remove `__del__` method