[MISC] Consolidate cleanup() and refactor offline_inference_with_prefix.py (#9510)

tests/entrypoints/llm/test_generate.py

@@ -4,8 +4,7 @@ from typing import List
 import pytest
 
 from vllm import LLM, RequestOutput, SamplingParams
-
-from ...conftest import cleanup
+from vllm.distributed import cleanup_dist_env_and_memory
 
 MODEL_NAME = "facebook/opt-125m"
 
@@ -39,7 +38,7 @@ def llm():
 
         del llm
 
-    cleanup()
+    cleanup_dist_env_and_memory()
 
 
 def assert_outputs_equal(o1: List[RequestOutput], o2: List[RequestOutput]):