tests/test_gpu_pool.py

"""Tests for fpgm.datagen.batch.gpu_pool.GpuWorkerPool.

All tests inject a synthetic ``query_fn`` rather than shelling out to a real
``nvidia-smi`` -- the *real* nvidia-smi behaviour on this box is exercised
separately (see the task's own verification step), not as part of the
no-GPU-required unit suite.
"""

from __future__ import annotations

import threading
import time

import pytest

from fpgm.config_datagen import GpuPoolConfig
from fpgm.datagen.batch.gpu_pool import GpuAcquireTimeoutError, GpuDevice, GpuWorkerPool


def _cfg(**overrides) -> GpuPoolConfig:
    base = dict(
        device_ids=None, min_free_mb=16000, workers_per_device=1, max_workers=None,
        threads_per_worker=8, acquire_timeout_s=0.5, poll_interval_s=0.05,
    )
    base.update(overrides)
    return GpuPoolConfig(**base)


class TestDiscovery:
    def test_discovers_all_measured_devices(self):
        pool = GpuWorkerPool(_cfg(), query_fn=lambda: {0: 20000, 1: 22000, 2: 3000, 3: 12000})
        assert pool.device_ids == (0, 1, 2, 3)

    def test_device_ids_filters_to_config(self):
        pool = GpuWorkerPool(
            _cfg(device_ids=[1, 3]), query_fn=lambda: {0: 20000, 1: 22000, 2: 3000, 3: 12000}
        )
        assert pool.device_ids == (1, 3)

    def test_no_devices_raises(self):
        with pytest.raises(RuntimeError, match="no devices discovered"):
            GpuWorkerPool(_cfg(device_ids=[9]), query_fn=lambda: {0: 20000})


class TestAcquireAdmission:
    def test_admits_only_devices_at_or_above_min_free_mb(self):
        measured = {0: 19612, 1: 22600, 2: 3400, 3: 11700}
        pool = GpuWorkerPool(_cfg(min_free_mb=16000), query_fn=lambda: measured)
        claim1 = pool.acquire()
        claim2 = pool.acquire()
        admitted = {claim1.device.index, claim2.device.index}
        assert admitted == {0, 1}
        with pytest.raises(GpuAcquireTimeoutError):
            pool.acquire()

    def test_workers_per_device_caps_concurrent_claims(self):
        pool = GpuWorkerPool(
            _cfg(min_free_mb=1000, workers_per_device=2), query_fn=lambda: {0: 50000}
        )
        c1 = pool.acquire()
        c2 = pool.acquire()
        assert c1.device.index == c2.device.index == 0
        with pytest.raises(GpuAcquireTimeoutError):
            pool.acquire()

    def test_release_frees_a_slot_for_the_next_acquire(self):
        pool = GpuWorkerPool(
            _cfg(min_free_mb=1000, workers_per_device=1), query_fn=lambda: {0: 50000}
        )
        claim = pool.acquire()
        with pytest.raises(GpuAcquireTimeoutError):
            pool.acquire()
        pool.release(claim.device)
        claim2 = pool.acquire()  # must not raise now
        assert claim2.device.index == 0

    def test_context_manager_releases_on_exit(self):
        pool = GpuWorkerPool(
            _cfg(min_free_mb=1000, workers_per_device=1), query_fn=lambda: {0: 50000}
        )
        with pool.acquire() as device:
            assert isinstance(device, GpuDevice)
            assert device.index == 0
        # slot freed by __exit__
        claim2 = pool.acquire()
        assert claim2.device.index == 0

    def test_context_manager_releases_on_exception(self):
        pool = GpuWorkerPool(
            _cfg(min_free_mb=1000, workers_per_device=1), query_fn=lambda: {0: 50000}
        )
        with pytest.raises(ValueError):
            with pool.acquire():
                raise ValueError("boom")
        claim2 = pool.acquire()  # must not raise -- slot was released despite the exception
        assert claim2.device.index == 0

    def test_re_measures_on_every_poll_admitting_once_memory_frees_up(self):
        state = {"free": 5000}  # below min_free_mb initially

        def query():
            return {0: state["free"]}

        pool = GpuWorkerPool(_cfg(min_free_mb=16000, acquire_timeout_s=2.0, poll_interval_s=0.05),
                              query_fn=query)

        def _free_it_later():
            time.sleep(0.15)
            state["free"] = 20000

        threading.Thread(target=_free_it_later, daemon=True).start()
        claim = pool.acquire()  # blocks until the background thread raises free memory
        assert claim.device.index == 0

    def test_timeout_message_reports_last_measurement_and_claims(self):
        pool = GpuWorkerPool(_cfg(min_free_mb=99999, acquire_timeout_s=0.1, poll_interval_s=0.02),
                              query_fn=lambda: {0: 100})
        with pytest.raises(GpuAcquireTimeoutError, match=r"last measured free MB: \{0: 100\}"):
            pool.acquire()


class TestThreadSafety:
    def test_concurrent_acquires_never_exceed_workers_per_device(self):
        pool = GpuWorkerPool(
            _cfg(
                min_free_mb=1000, workers_per_device=2,
                acquire_timeout_s=1.0, poll_interval_s=0.02,
            ),
            query_fn=lambda: {0: 50000, 1: 50000},
        )
        claims: list = []
        lock = threading.Lock()
        errors: list = []

        def worker():
            try:
                c = pool.acquire()
                with lock:
                    claims.append(c)
            except GpuAcquireTimeoutError as exc:
                with lock:
                    errors.append(exc)

        threads = [threading.Thread(target=worker) for _ in range(4)]
        for t in threads:
            t.start()
        for t in threads:
            t.join(timeout=5)

        assert len(claims) == 4  # 2 devices x workers_per_device=2 == exactly 4 admitted
        assert len(errors) == 0
        by_device: dict[int, int] = {}
        for c in claims:
            by_device[c.device.index] = by_device.get(c.device.index, 0) + 1
        assert all(n <= 2 for n in by_device.values())