ci: memory management in tests (#721)

* fix(vllm): implement shared backend to prevent GPU OOM errors

- Add session-scoped shared_vllm_backend fixture using Granite 4 Micro
- Update test_vllm.py and test_vllm_tools.py to use shared backend
- Fall back to module-scoped backends when --isolate-heavy flag is set
- Both modules now use consistent Granite 4 Micro model
- Enhance CUDA OOM error message with actionable solutions
- Maintains backward compatibility with existing isolation mechanism

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

* reduce vllm GPU allocation for tests

* implement backend test grouping via reordering

* add gpu cleanup between backend groups

* delay vllm backend creation until after openai vllm group

* adding explicit served model name for vllm openai test

* fix: rag intrinscis are not for the hybrid model (I think)

* testing a fix in tests for all the gpu issues

* more gpu cleaning

* adding docs tooling to mypy exclude

* removing kv cache also from GPU in cleanup for tests

* moving test order around and also fixing a fixture bug

* rolling back some changes from exclusive process

* some changes to the error message in vllm and also conftest cleaning

* adding an end-to-end script for tests with ollama

* adding a port finder (just in case)

* adding direct download of ollama binary from github

* warm starting ollama

* warm starting ollama

* adding cuda paths for ollama

* some extra checks for vllm and teardown

* making group by backend default

* making the script executable

* test: remove heavy ram pytest marks added in #623

Signed-off-by: Alex Bozarth <ajbozart@us.ibm.com>

* ruff formatting

* small changes to script and adding cleaup to guardian and core

* making log dir more easy to set

* increasing ollama startup to 2 mins

* adding pytest-json-report

---------

Signed-off-by: Alex Bozarth <ajbozart@us.ibm.com>
Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
Co-authored-by: avinash.bala@us.ibm.com;0J8455897;AVINASH BALAKRISHNAN <avinashbala@p5-r03-n2.bluevela.rmf.ibm.com>
Co-authored-by: Alex Bozarth <ajbozart@us.ibm.com>

Author: 4 people

mellea/backends/vllm.py

@@ -184,11 +184,15 @@ def __init__(
                     or engine_args["gpu_memory_utilization"] < 0.1
                 ):
                     raise RuntimeError(
-                        "no matter how I reduced max_model_len and max_num_seqs, there is not enough memory! \n"
-                        "final values:\n"
-                        f"gpu_memory_utilization: {engine_args['gpu_memory_utilization']}\n"
-                        f"max_model_len: {engine_args['max_model_len']}\n"
-                        f"max_num_seqs: {engine_args['max_num_seqs']}\n"
+                        f"Failed to initialize vLLM engine (last error: {e}).\n"
+                        "\n"
+                        "Exhausted all retry attempts:\n"
+                        f"  • gpu_memory_utilization: {engine_args['gpu_memory_utilization']:.2f}\n"
+                        f"  • max_model_len: {engine_args['max_model_len']}\n"
+                        f"  • max_num_seqs: {engine_args['max_num_seqs']}\n"
+                        "\n"
+                        "Please check the full stack trace above for details.\n"
+                        "Common issues: OOM, CUDA fork errors, GPU exclusive_process mode.\n"
                     )
                 logger.info(
                     f"Reducing vllm model parameters to make it fit in the GPU memory.\n"

pyproject.toml

@@ -128,6 +128,7 @@ test = [
     "pytest-timeout",  # For test suite timeout configuration
     "pytest-recording",  # For blocking network access and using canned server responses
     "pytest-xdist>=3.8.0", # For concurrent test execution (uv run pytest -n auto test/)
+    "pytest-json-report",  # For JSON test reports in nightly CI
     "psutil",  # For test infrastructure: RAM detection in conftest.py
     "nbmake>=1.5.5",  # Notebook testing
     "sentencepiece==0.2.1",  # Necessary for test_huggingface_tools test because of Mistral model
@@ -313,7 +314,7 @@ disable_error_code = [
 # -----------------------------
 
 [tool.codespell]
-ignore-words-list = 'mellea,hashi,noo,Asai,asai,nd,mot,rouge,Rouge,Strat,Wight,Aer,aer'
+ignore-words-list = 'mellea,hashi,noo,Asai,asai,nd,mot,rouge,Rouge,Strat,Wight,Aer,aer,preemptable'
 check-filenames = true
 check-hidden = false
 regex = "(?<![a-z])[a-z'`]+|[A-Z][a-z'`]*|[a-z]+'[a-z]*|[a-z]+(?=[_-])|[a-z]+(?=[A-Z])|\\d+"

test/backends/test_huggingface.py

@@ -64,7 +64,11 @@ def backend():
     backend.add_adapter(
         IntrinsicAdapter("answerability", base_model_name=backend.base_model_name)
     )
-    return backend
+    yield backend
+
+    from test.conftest import cleanup_gpu_backend
+
+    cleanup_gpu_backend(backend, "huggingface")
 
 
 @pytest.fixture(scope="function")

test/backends/test_huggingface_tools.py

@@ -32,7 +32,11 @@ def backend():
         model_id=model_ids.MISTRALAI_MISTRAL_0_3_7B, cache=SimpleLRUCache(5)
     )
     # add_granite_aloras(backend)
-    return backend
+    yield backend
+
+    from test.conftest import cleanup_gpu_backend
+
+    cleanup_gpu_backend(backend, "huggingface-tools")
 
 
 @pytest.fixture(scope="function")

test/backends/test_openai_vllm.py

@@ -45,6 +45,8 @@ def vllm_process():
                 "vllm",
                 "serve",
                 IBM_GRANITE_4_MICRO_3B.hf_model_name,
+                "--served-model-name",
+                IBM_GRANITE_4_MICRO_3B.hf_model_name,
                 "--enable-lora",
                 "--dtype",
                 "bfloat16",
@@ -55,6 +57,9 @@ def vllm_process():
             # the process will have a new session id, so
             # entire process tree is killable at once
             start_new_session=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,  # merge stderr into stdout
+            text=True,
         )
         url = "http://127.0.0.1:8000/ping"
         timeout = 600  # vllm initialization takes quite a while
@@ -63,7 +68,11 @@ def vllm_process():
         # Wait for readiness message
         while True:
             if process.poll() is not None:
-                raise RuntimeError("vLLM server exited before startup.")
+                output = process.stdout.read() if process.stdout else ""
+                raise RuntimeError(
+                    f"vLLM server exited before startup (code {process.returncode}).\n"
+                    f"--- vLLM output ---\n{output}\n--- end ---"
+                )
 
             try:
                 response = requests.get(url, timeout=2)
@@ -73,17 +82,36 @@ def vllm_process():
                 pass
 
             if time.time() - start_time > timeout:
+                output = ""
+                if process.stdout:
+                    try:
+                        # Read whatever is available without blocking
+                        import select
+
+                        if select.select([process.stdout], [], [], 0)[0]:
+                            output = process.stdout.read()
+                    except Exception:
+                        pass
                 raise TimeoutError(
-                    f"Timed out waiting for server health check at {url}"
+                    f"Timed out waiting for server health check at {url}\n"
+                    f"--- vLLM output (last lines) ---\n{output[-2000:]}\n--- end ---"
                 )
 
         yield process
 
     except Exception as e:
-        pytest.skip(
-            f"vLLM process not available: {e}. May need to install with: pip install mellea[vllm]",
-            allow_module_level=True,
+        output = ""
+        if process is not None and process.stdout:
+            try:
+                output = process.stdout.read()
+            except Exception:
+                pass
+        skip_msg = (
+            f"vLLM process not available: {e}\n"
+            f"--- vLLM output ---\n{output}\n--- end ---"
         )
+        print(skip_msg)  # visible with -s flag
+        pytest.skip(skip_msg, allow_module_level=True)
 
     # --- Teardown (always runs) ---
     finally:

test/backends/test_vllm.py

@@ -34,26 +34,34 @@
     )
 
 
+# vLLM tests use hybrid backend strategy (see conftest.py):
+# - Default: Shared session-scoped backend (fast, no fragmentation)
+# - --isolate-heavy: Module-scoped backends in separate processes
 @pytest.fixture(scope="module")
-def backend():
-    """Shared vllm backend for all tests in this module."""
+def backend(shared_vllm_backend):
+    """Use shared session-scoped backend, or create module-scoped if isolated.
+
+    Without --isolate-heavy: Uses shared backend (fast, no fragmentation)
+    With --isolate-heavy: Creates module-scoped backend (process isolation)
+    """
+    if shared_vllm_backend is not None:
+        yield shared_vllm_backend
+        return  # skip cleanup — shared backend cleaned up by conftest
+
+    # Isolation mode - create module-scoped backend
     backend = LocalVLLMBackend(
-        model_id=model_ids.QWEN3_0_6B,
-        # formatter=TemplateFormatter(model_id="ibm-granite/granite-4.0-tiny-preview"),
+        model_id=model_ids.IBM_GRANITE_4_MICRO_3B,
         model_options={
-            # made smaller for a testing environment with smaller gpus.
-            # such an environment could possibly be running other gpu applications, including slack
-            "gpu_memory_utilization": 0.8,
-            "max_model_len": 8192,
-            "max_num_seqs": 8,
+            "gpu_memory_utilization": 0.6,
+            "max_model_len": 4096,
+            "max_num_seqs": 4,
         },
     )
     yield backend
 
-    # Cleanup using shared function (best-effort within module)
-    from test.conftest import cleanup_vllm_backend
+    from test.conftest import cleanup_gpu_backend
 
-    cleanup_vllm_backend(backend)
+    cleanup_gpu_backend(backend, "vllm")
 
 
 @pytest.fixture(scope="function")

test/backends/test_vllm_tools.py

@@ -31,26 +31,36 @@
     )
 
 
+# vLLM tests use hybrid backend strategy (see conftest.py):
+# - Default: Shared session-scoped backend (fast, no fragmentation)
+# - --isolate-heavy: Module-scoped backends in separate processes
+# Note: Originally used Mistral-7B, now uses Granite 4 Micro for consistency.
+# Granite 4 Micro supports tool calling and is sufficient for testing.
 @pytest.fixture(scope="module")
-def backend():
-    """Shared vllm backend for all tests in this module."""
+def backend(shared_vllm_backend):
+    """Use shared session-scoped backend, or create module-scoped if isolated.
 
+    Without --isolate-heavy: Uses shared backend (fast, no fragmentation)
+    With --isolate-heavy: Creates module-scoped backend (process isolation)
+    """
+    if shared_vllm_backend is not None:
+        yield shared_vllm_backend
+        return
+
+    # Isolation mode - create module-scoped backend
     backend = LocalVLLMBackend(
-        model_id=model_ids.MISTRALAI_MISTRAL_0_3_7B,
+        model_id=model_ids.IBM_GRANITE_4_MICRO_3B,
         model_options={
-            # made smaller for a testing environment with smaller gpus.
-            # such an environment could possibly be running other gpu applications, including slack
-            "gpu_memory_utilization": 0.8,
-            "max_model_len": 8192,
-            "max_num_seqs": 8,
+            "gpu_memory_utilization": 0.6,
+            "max_model_len": 4096,
+            "max_num_seqs": 4,
         },
     )
     yield backend
 
-    # Cleanup using shared function (best-effort within module)
-    from test.conftest import cleanup_vllm_backend
+    from test.conftest import cleanup_gpu_backend
 
-    cleanup_vllm_backend(backend)
+    cleanup_gpu_backend(backend, "vllm-tools")
 
 
 @pytest.fixture(scope="function")

test/cli/test_alora_train_integration.py

@@ -287,6 +287,16 @@ def test_alora_training_integration():
             "✅ Verified adapter activation: both with/without invocation tokens generate successfully"
         )
 
+        # Cleanup GPU memory
+        base_model.cpu()
+        del model_with_adapter
+        del base_model
+        import gc
+
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
 
 def test_lora_training_integration():
     """Integration test: Train a tiny standard LoRA adapter and verify it works.