# GPU Memory Management Refactoring Analysis

## Executive Summary

The test suite currently has **redundant and inefficient GPU memory cleanup** patterns. LLM tests execute cleanup **3 times per test**, and cleanup logic is scattered across multiple conftest files and individual test functions. This document outlines the current state and proposed refactoring.

---

## Current State

### Cleanup Locations

#### 1. Root `tests/conftest.py` (lines 24-40)
```python
@pytest.fixture(autouse=True, scope="function")
def cleanup():
    yield
    # Cleans CUDA cache if CUDA is initialized
    if torch.cuda.is_available() and torch.cuda.is_initialized():
        torch.cuda.empty_cache()

    # Runs gc.collect() every 5 tests
    if cleanup.call_count % 5 == 0:
        gc.collect()
```
**Scope**: ALL tests in the suite
**Frequency**: After every test

#### 2. Root `tests/conftest.py` (lines 290-314)
```python
@pytest.fixture
def deepspeed_env():
    # Sets up DeepSpeed environment variables
    yield
    finally:
        # Cleanup env vars
        gc.collect()
        torch.cuda.empty_cache()
```
**Scope**: Used by 50 tests (46 in LLM tests, 4 in other tests)
**Issue**: Should be in LLM conftest, not root

#### 3. LLM `tests/test_algorithms/test_llms/conftest.py` (lines 54-66)
```python
@pytest.fixture(autouse=True)
def cleanup_after_test(request):
    yield
    if "vllm" in request.node.name:
        cleanup_dist_env_and_memory()
        # Reset DeepSpeed groups
        for attr in dir(ds_groups):
            if attr.startswith("_") and attr.endswith("_GROUP"):
                setattr(ds_groups, attr, None)
        ds_comm.cdb = None
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.synchronize()
    AcceleratorState._reset_state(True)
```
**Scope**: All LLM tests
**Issue**: Runs IN ADDITION to root cleanup (2x cleanup per test)

#### 4. LLM `tests/test_algorithms/test_llms/conftest.py` (lines 81-103)
```python
@pytest.fixture(scope="function")
def accelerator_factory():
    def generate_accelerator(...):
        gc.collect()  # <-- Cleanup at START
        torch.cuda.empty_cache()
        AcceleratorState._reset_state(True)
        # ... create accelerator
    return generate_accelerator
```
**Issue**: Cleans up at START, conflicts with fixture-based cleanup

#### 5. `tests/test_hpo/test_mutation.py` (lines 77-83)
```python
@pytest.fixture(autouse=True)
def cleanup_after_test():
    yield
    gc.collect()
    torch.cuda.synchronize()
    torch.cuda.empty_cache()
```
**Scope**: All mutation tests
**Issue**: Runs IN ADDITION to root cleanup (2x cleanup per test)

#### 6. `tests/test_modules/test_make_evolvable.py`
Multiple fixtures with teardown cleanup:
```python
@pytest.fixture
def simple_mlp():
    network = nn.Sequential(...)
    yield network
    del network
    gc.collect()
    torch.cuda.empty_cache()
```
**Count**: 6 fixtures with this pattern
**Issue**: Runs IN ADDITION to root cleanup

#### 7. Individual Test Functions
Manual cleanup in 60+ test functions across multiple files:
- `tests/test_algorithms/test_llms/test_grpo.py`: 25 occurrences
- `tests/test_algorithms/test_llms/test_dpo.py`: 2 occurrences
- `tests/test_utils/test_probe_envs.py`: 20 occurrences
- `tests/test_utils/test_probe_envs_ma.py`: 8 occurrences
- etc.

---

## Problems Identified

### 1. **Redundant Cleanup (Most Critical)**

**LLM Tests**: Run cleanup **3 times**:
1. Root `cleanup` fixture (function-scoped autouse)
2. LLM `cleanup_after_test` fixture (autouse)
3. Often manual cleanup in test body

**Result**: Excessive `gc.collect()` and `torch.cuda.empty_cache()` calls slow down test suite

**Mutation Tests**: Run cleanup **2 times**:
1. Root `cleanup` fixture
2. Local `cleanup_after_test` fixture

### 2. **Defensive Exception Handling**

In `tests/test_algorithms/test_llms/conftest.py` (lines 19-52):
```python
def cleanup_vllm_instances():
    try:
        vllm._global_llm_engine.shutdown()
    except Exception:  # <-- TOO BROAD
        pass  # <-- HIDES ERRORS
```

**Issues**:
- Catches generic `Exception`
- Silently passes, hiding real errors
- Makes debugging impossible
- Appears 3 times in the function

### 3. **Fixture Locality**

**`deepspeed_env` fixture**:
- Lives in: `tests/conftest.py` (root)
- Used by:
  - `tests/test_algorithms/test_llms/test_grpo.py` (45 times)
  - `tests/test_algorithms/test_llms/test_dpo.py` (4 times)
  - `tests/test_hpo/test_mutation.py` (1 time)

**Should be**: In `tests/test_algorithms/test_llms/conftest.py`
**Reason**: 98% of usage is LLM tests; root conftest should only contain truly universal fixtures

### 4. **Unclear Cleanup Order**

No documentation explaining:
- Why `AcceleratorState._reset_state()` happens after GPU cleanup
- Why `torch.cuda.synchronize()` is needed in some places but not others
- Order requirements for vLLM → DeepSpeed → Accelerator cleanup

### 5. **Inconsistent Patterns**

- Some fixtures clean up in teardown (`test_make_evolvable.py`)
- Some tests clean up manually in test body
- Some rely only on autouse fixtures
- `accelerator_factory` cleans at START instead of teardown

---

## Proposed Solution

### Design Principles

1. **Single Responsibility**: One authoritative cleanup location per test scope
2. **No Redundancy**: Eliminate duplicate cleanup calls
3. **Explicit Over Implicit**: Replace blanket exception handling with specific handling
4. **Locality**: Specialist fixtures belong near their usage
5. **Documentation**: Explain non-obvious cleanup ordering

### Implementation Plan

#### 1. **Root `tests/conftest.py`**

**Keep**:
- Minimal, smart GPU cleanup for general tests
- Only runs when CUDA was actually used
- Conservative gc.collect() frequency

**Remove**:
- `deepspeed_env` fixture (move to LLM conftest)

**Optimized cleanup fixture**:
```python
@pytest.fixture(autouse=True, scope="function")
def cleanup_gpu_memory():
    """
    Minimal GPU memory cleanup for non-specialist tests.

    LLM tests override this with their own cleanup fixture that handles
    distributed training framework cleanup (vLLM, DeepSpeed, Accelerator).
    """
    yield

    # Only clear CUDA cache if CUDA was actually used
    if torch.cuda.is_available() and torch.cuda.is_initialized():
        torch.cuda.empty_cache()

    # Run garbage collection periodically (not every test)
    if not hasattr(cleanup_gpu_memory, "call_count"):
        cleanup_gpu_memory.call_count = 0
    cleanup_gpu_memory.call_count += 1

    if cleanup_gpu_memory.call_count % 10 == 0:
        gc.collect()
```

#### 2. **LLM `tests/test_algorithms/test_llms/conftest.py`**

**Add**: `deepspeed_env` fixture (moved from root)

**Replace** both `cleanup_after_test` and cleanup in `accelerator_factory` with single fixture:
```python
@pytest.fixture(autouse=True)
def cleanup_distributed_training(request):
    """
    Cleanup fixture for LLM tests using distributed training frameworks.

    This fixture handles the proper teardown order for distributed systems:
    1. vLLM cleanup (if used) - must happen first to release GPU memory properly
    2. DeepSpeed state reset - clears process group state
    3. Accelerator state reset - clears HuggingFace Accelerate state
    4. GPU memory cleanup - releases cached memory
    5. Garbage collection - frees Python objects

    The order matters because:
    - vLLM holds GPU memory that must be released before other cleanup
    - DeepSpeed process groups must be cleared before Accelerator reset
    - Accelerator state must be reset before next test initialization
    """
    yield

    # Step 1: Clean up vLLM if this test used it
    if "vllm" in request.node.name or _test_used_vllm(request):
        cleanup_vllm_instances()
        cleanup_dist_env_and_memory()

        # Reset DeepSpeed global state
        # These must be cleared or subsequent tests will fail with process group errors
        for attr in dir(ds_groups):
            if attr.startswith("_") and attr.endswith("_GROUP"):
                setattr(ds_groups, attr, None)
        ds_comm.cdb = None

    # Step 2: GPU synchronization and memory cleanup
    if torch.cuda.is_available() and torch.cuda.is_initialized():
        torch.cuda.synchronize()  # Wait for all GPU operations to complete
        torch.cuda.empty_cache()   # Release cached memory

    # Step 3: Reset Accelerator state for next test
    AcceleratorState._reset_state(True)

    # Step 4: Garbage collection
    gc.collect()


def cleanup_vllm_instances():
    """
    Clean up vLLM LLM instances and engines.

    Note: We use bare except here because vLLM's shutdown process can fail
    in various ways (CUDA errors, process group errors, etc.) and we need
    to ensure all cleanup attempts are made even if some fail.
    """
    import vllm
    import vllm.engine.llm_engine

    # Clean global engine
    if hasattr(vllm, "_global_llm_engine"):
        try:
            if vllm._global_llm_engine is not None:
                vllm._global_llm_engine.shutdown()
        except Exception as e:
            # vLLM shutdown can fail due to CUDA/distributed errors during test teardown
            # We log but don't raise to ensure other cleanup continues
            import warnings
            warnings.warn(f"vLLM global engine shutdown failed: {e}")
        finally:
            del vllm._global_llm_engine

    # Clean any cached engines
    if hasattr(vllm.engine.llm_engine, "_cached_engines"):
        for engine in vllm.engine.llm_engine._cached_engines.values():
            try:
                if engine is not None:
                    engine.shutdown()
            except Exception as e:
                import warnings
                warnings.warn(f"vLLM cached engine shutdown failed: {e}")
        vllm.engine.llm_engine._cached_engines.clear()

    # Clean LLM class instances
    if hasattr(vllm, "LLM") and hasattr(vllm.LLM, "_instances"):
        for instance in vllm.LLM._instances:
            try:
                if hasattr(instance, "llm_engine") and instance.llm_engine is not None:
                    instance.llm_engine.shutdown()
            except Exception as e:
                import warnings
                warnings.warn(f"vLLM instance shutdown failed: {e}")
        vllm.LLM._instances.clear()


def _test_used_vllm(request):
    """Check if test used vLLM by inspecting fixtures."""
    return any("vllm" in str(fixture).lower() for fixture in request.fixturenames)
```

**Update `accelerator_factory`**:
```python
@pytest.fixture(scope="function")
def accelerator_factory():
    """Factory for creating Accelerator instances with DeepSpeed support."""
    def generate_accelerator(use_deepspeed_optimizer, config):
        # Reset state before creating new accelerator
        # This is safe because cleanup_distributed_training fixture
        # already ran after the previous test
        AcceleratorState._reset_state(True)

        if use_deepspeed_optimizer and (config is not None):
            config["optimizer"] = {
                "type": "AdamW",
                "params": {
                    "lr": 1e-4,
                    "betas": [0.9, 0.999],
                    "eps": 1e-8,
                    "weight_decay": 0.01,
                },
            }
        return (
            Accelerator(deepspeed_plugin=DeepSpeedPlugin(hf_ds_config=config))
            if config is not None
            else None
        )

    return generate_accelerator
```

**Move `deepspeed_env` from root conftest**:
```python
# At top of file
import socket

dist_env = dict(
    ACCELERATE_USE_DEEPSPEED="true",
    MASTER_ADDR="localhost",
    MASTER_PORT="10999",
    RANK="0",
    LOCAL_RANK="0",
    WORLD_SIZE="1",
    PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True",
    CUDA_VISIBLE_DEVICES="0",
)


def get_free_port():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("", 0))
        return s.getsockname()[1]


@pytest.fixture
def deepspeed_env():
    """
    Set up environment variables for DeepSpeed distributed training.

    This fixture is scoped to LLM tests but can be imported by other tests
    that need DeepSpeed (e.g., test_hpo/test_mutation.py).
    """
    import os

    dynamic_dist_env = dist_env.copy()
    dynamic_dist_env["MASTER_PORT"] = str(get_free_port())
    existing_vars = {}

    for key, value in dynamic_dist_env.items():
        key = key.upper()
        if key in os.environ:
            existing_vars[key] = os.environ[key]
        os.environ[key] = str(value)

    try:
        yield
    finally:
        for key in dynamic_dist_env:
            key = key.upper()
            if key in existing_vars:
                os.environ[key] = existing_vars[key]
            else:
                os.environ.pop(key, None)
```

#### 3. **Remove Manual Cleanup from Test Files**

**Files to modify**:
- `tests/test_algorithms/test_llms/test_grpo.py` (remove ~25 cleanup calls)
- `tests/test_algorithms/test_llms/test_dpo.py` (remove ~2 cleanup calls)
- `tests/test_utils/test_probe_envs.py` (remove ~20 cleanup calls)
- `tests/test_utils/test_probe_envs_ma.py` (remove ~8 cleanup calls)
- `tests/test_modules/test_cnn.py` (remove 1 cleanup call)
- `tests/test_modules/test_bert.py` (remove 1 cleanup call)

**Pattern to remove**:
```python
# DELETE THESE LINES:
gc.collect()
torch.cuda.empty_cache()
```

#### 4. **Remove Local Cleanup Fixtures**

**`tests/test_hpo/test_mutation.py`**:
- Remove `cleanup_after_test` fixture (lines 77-83)
- Rely on root cleanup fixture instead

**`tests/test_modules/test_make_evolvable.py`**:
- Remove cleanup from fixture teardown in:
  - `simple_mlp` (lines 69-70)
  - `simple_mlp_2` (lines 80-81)
  - `mlp_norm_layers` (lines 98-99)
  - `simple_cnn` (lines 122-123)
  - `simple_cnn_norm_layers` (lines 148-149)
  - `simple_cnn_2` (lines 157-158)

**Pattern to remove**:
```python
@pytest.fixture
def simple_mlp():
    network = nn.Sequential(...)
    yield network
    del network
    gc.collect()  # <-- DELETE
    torch.cuda.empty_cache()  # <-- DELETE
```

**After refactoring**:
```python
@pytest.fixture
def simple_mlp():
    network = nn.Sequential(...)
    yield network
    # Cleanup handled by root cleanup fixture
```

#### 5. **Update `tests/test_hpo/test_mutation.py` imports**

Since `deepspeed_env` is moving to LLM conftest:
```python
# Add this import at top of file
import sys
sys.path.insert(0, str(Path(__file__).parent.parent / "test_algorithms" / "test_llms"))
from conftest import deepspeed_env  # noqa: E402
```

Or better, use `pytest_plugins`:
```python
# At top of test_mutation.py
pytest_plugins = ["tests.test_algorithms.test_llms.conftest"]
```

---

## Expected Impact

### Performance
- **LLM tests**: Reduce cleanup overhead by 66% (3 cleanups → 1 cleanup per test)
- **Mutation tests**: Reduce cleanup overhead by 50% (2 cleanups → 1 cleanup per test)
- **General tests**: Minimal impact (cleanup already optimized)

### Maintainability
- Single source of truth for cleanup logic
- Clear documentation of cleanup order
- Easier to debug (no silent exception swallowing)

### Test Isolation
- No impact - cleanup still happens after every test
- Distributed state properly reset between tests
- GPU memory properly released

---

## Migration Path

1. ✅ Document current state and issues
2. ⏭️ Refactor root `tests/conftest.py`
3. ⏭️ Refactor LLM `tests/test_algorithms/test_llms/conftest.py`
4. ⏭️ Remove manual cleanup from test files (can be done in batches)
5. ⏭️ Remove cleanup from local fixtures
6. ⏭️ Update imports in `test_mutation.py`
7. ⏭️ Run full test suite to validate
8. ⏭️ Monitor for any test failures or GPU memory issues

---

## Risk Assessment

**Low Risk Changes**:
- Removing manual cleanup from test bodies (autouse fixtures will handle it)
- Documenting cleanup order (documentation only)

**Medium Risk Changes**:
- Consolidating LLM cleanup fixtures (well-tested pattern, just restructured)
- Moving `deepspeed_env` fixture (requires import update in one file)

**Higher Risk Changes**:
- Changing exception handling in vLLM cleanup (but adding warnings for visibility)
- Removing cleanup from `accelerator_factory` startup (should be handled by fixture)

**Mitigation**:
- Test incrementally (one conftest at a time)
- Keep manual cleanup in tests initially, remove after validating fixtures work
- Monitor CI for GPU OOM errors

---

## Open Questions

1. **Should `test_mutation.py` use a local `deepspeed_env` or import from LLM conftest?**
   - Recommendation: Import from LLM conftest to avoid duplication

2. **Should we add a pytest marker for vLLM tests instead of checking test name?**
   - Recommendation: Yes, add `@pytest.mark.vllm` marker for clarity

3. **Should `torch.cuda.synchronize()` be in root cleanup too?**
   - Recommendation: No, only needed for distributed training (LLM tests)

4. **What's the right gc.collect() frequency?**
   - Current: Every 5 tests (root), every test (LLM)
   - Proposed: Every 10 tests (root), every test (LLM)
   - Open to adjustment based on memory profiling
