platform-codebase/tests/gpu_integration/conftest.py

"""Shared fixtures for GPU integration tests proving model-boss v3 migration.

This module provides pytest fixtures for testing GPU coordination across
lilith-platform ML services using model-boss v3.

Run with: pytest -m "gpu and modelboss" --real-model -v
"""
from __future__ import annotations

import asyncio
import os
from typing import TYPE_CHECKING, AsyncGenerator, Callable

import pytest
import pytest_asyncio

if TYPE_CHECKING:
    from model_boss import GPUBoss
    from model_boss_loaders import ManagedModelLoader


def pytest_addoption(parser: pytest.Parser) -> None:
    """Add CLI options for GPU tests."""
    parser.addoption(
        "--real-model",
        action="store_true",
        default=False,
        help="Run real GPU tests with actual model loading",
    )
    parser.addoption(
        "--redis-url",
        default=os.environ.get("REDIS_URL", "redis://localhost:6379"),
        help="Redis URL for GPU coordination",
    )


def pytest_configure(config: pytest.Config) -> None:
    """Register custom markers."""
    config.addinivalue_line("markers", "gpu: Requires GPU hardware")
    config.addinivalue_line("markers", "modelboss: Tests model-boss v3 integration")
    config.addinivalue_line("markers", "slow: Slow tests (model loading)")


def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
    """Skip GPU tests if --real-model not specified."""
    if not config.getoption("--real-model"):
        skip_gpu = pytest.mark.skip(reason="Use --real-model to run GPU tests")
        for item in items:
            if "gpu" in item.keywords:
                item.add_marker(skip_gpu)


@pytest.fixture(scope="session")
def redis_url(request: pytest.FixtureRequest) -> str:
    """Get Redis URL from CLI or environment."""
    return request.config.getoption("--redis-url")


@pytest.fixture(scope="session")
def gpu_available() -> bool:
    """Check if CUDA GPU is available."""
    try:
        import torch
        return torch.cuda.is_available()
    except ImportError:
        return False


@pytest.fixture(scope="session")
def gpu_vram_mb() -> int:
    """Get total GPU VRAM in MB."""
    try:
        import torch
        if not torch.cuda.is_available():
            return 0
        return torch.cuda.get_device_properties(0).total_memory // (1024 * 1024)
    except ImportError:
        return 0


@pytest.fixture(scope="session")
def gpu_name() -> str:
    """Get GPU device name."""
    try:
        import torch
        if not torch.cuda.is_available():
            return "No GPU"
        return torch.cuda.get_device_properties(0).name
    except ImportError:
        return "Unknown"


@pytest_asyncio.fixture(scope="session")
async def real_gpu_boss(
    request: pytest.FixtureRequest,
    redis_url: str,
    gpu_available: bool,
    gpu_vram_mb: int,
    gpu_name: str,
) -> AsyncGenerator["GPUBoss", None]:
    """Real GPUBoss connected to Redis with GPU initialized.

    This fixture:
    1. Connects to Redis (auto-starts if not running via model-boss daemon)
    2. Initializes the GPU with detected VRAM
    3. Yields the boss for tests
    4. Cleans up on exit

    Note: model-boss auto_start_services=True by default, so Redis
    will be started automatically if not already running.
    """
    if not request.config.getoption("--real-model"):
        pytest.skip("Use --real-model for GPU tests")

    if not gpu_available:
        pytest.skip("No GPU available")

    from model_boss import GPUBoss

    # auto_start_services=True by default - Redis starts if not running
    boss = GPUBoss(redis_url=redis_url)
    await boss.connect()

    # Initialize GPU 0 with detected VRAM
    await boss.initialize_gpu(
        gpu_index=0,
        vram_total_mb=gpu_vram_mb,
        gpu_name=gpu_name,
    )

    yield boss

    # Cleanup: release any remaining leases
    try:
        status = await boss.get_status()
        for gpu in status.gpus:
            for lease in gpu.leases:
                await boss.force_release(lease.lease_id)
    except Exception:
        pass

    await boss.close()


@pytest_asyncio.fixture
async def managed_loader_factory(
    real_gpu_boss: "GPUBoss",
) -> AsyncGenerator[Callable[[str], "ManagedModelLoader"], None]:
    """Factory for creating ManagedModelLoader instances with cleanup.

    Usage:
        loader = managed_loader_factory("my-service")
        model = await loader.load(model_id="my-model")
        # ... use model ...
        # Automatically cleaned up after test
    """
    from model_boss_loaders import ManagedModelLoader

    loaders: list[ManagedModelLoader] = []

    def _create(service_name: str = "test") -> ManagedModelLoader:
        loader = ManagedModelLoader(boss=real_gpu_boss)
        loaders.append(loader)
        return loader

    yield _create

    # Cleanup all loaders
    for loader in loaders:
        try:
            await loader.unload_all()
        except Exception:
            pass


# Helper functions are in helpers.py for direct import by test files