gpu-rent/tests/test_inventory.py

from gpu_rent.inventory import (
    gpu_boot_image_score,
    gpu_quota_from_compute,
    match_label,
    pick_boot_image,
    rank_flavors,
    resolve_flavor,
)


class FakeFlavor:
    def __init__(
        self,
        flavor_id: str,
        name: str,
        extra=None,
        disabled: bool = False,
        vcpus: int = 8,
        ram: int = 32768,
        disk: int = 0,
        availability_zones=None,
    ):
        self.id = flavor_id
        self.name = name
        self.extra_specs = extra or {}
        self.is_disabled = disabled
        self.vcpus = vcpus
        self.ram = ram
        self.disk = disk
        self.availability_zones = list(availability_zones or [])


def test_match_4090_24_not_48():
    a = FakeFlavor("1", "GPU 1x RTX 4090 24GB")
    b = FakeFlavor("2", "GPU 1x RTX 4090 48GB")
    assert match_label("4090-24", a)
    assert not match_label("4090-24", b)
    assert match_label("4090-48", b)


def test_rank_skips_disabled():
    flavors = [
        FakeFlavor("d", "RTX 4090 24GB", disabled=True),
        FakeFlavor("ok", "RTX A5000 24GB"),
    ]
    ranked = rank_flavors(flavors, ("4090-24", "a5000"))
    assert [x.id for x in ranked] == ["ok"]
    assert ranked[0].label == "a5000"


def test_gpu_quota_from_compute():
    assert gpu_quota_from_compute({"cores": 10}) is None
    assert gpu_quota_from_compute({"gpu": 0}) == 0
    assert gpu_quota_from_compute({"GPU_limit": 2}) == 2


def test_resolve_flavor_explicit_and_preference():
    flavors = [
        FakeFlavor("d", "RTX 4090 24GB", disabled=True),
        FakeFlavor("a5", "RTX A5000 24GB"),
        FakeFlavor("ok", "GPU 1x RTX 4090 24GB"),
    ]
    picked = resolve_flavor(flavors, ("4090-24", "a5000"))
    assert picked.id == "ok"
    explicit = resolve_flavor(flavors, ("a5000",), explicit="a5")
    assert explicit.id == "a5"


def test_resolve_flavor_prefers_list_over_default_id():
    flavors = [
        FakeFlavor("a5", "RTX A5000 24GB"),
        FakeFlavor("ok", "GPU 1x RTX 4090 24GB"),
    ]
    picked = resolve_flavor(flavors, ("4090-24", "a5000"), default_id="a5", fallback=True)
    assert picked.id == "ok"


def test_resolve_flavor_no_fallback_requires_default():
    flavors = [FakeFlavor("ok", "GPU 1x RTX 4090 24GB")]
    try:
        resolve_flavor(flavors, ("4090-24",), fallback=False)
        raise AssertionError("expected ValueError")
    except ValueError:
        pass
    picked = resolve_flavor(flavors, ("a5000",), default_id="ok", fallback=False)
    assert picked.id == "ok"


def test_rank_cheap_prefers_1gpu_min_sku():
    extra = {"aggregate_instance_extra_specs:gpu": "RTX4090", "gpu_memory": "24"}
    flavors = [
        FakeFlavor("dual", "GL10.8-65536-0-2GPU", extra, vcpus=8, ram=65536),
        FakeFlavor("fat", "GL10.8-32768-0-1GPU", extra, vcpus=8, ram=32768),
        FakeFlavor("tiny", "GL10.4-16384-0-1GPU", extra, vcpus=4, ram=16384),
    ]
    ranked = rank_flavors(flavors, ("4090-24",), size_preset="cheap")
    assert ranked[0].id == "tiny"
    assert ranked[0].gpu_count == 1
    balanced = rank_flavors(flavors, ("4090-24",), size_preset="balanced")
    assert balanced[0].id == "fat"


def test_rank_filters_availability_zone():
    extra = {"aggregate_instance_extra_specs:gpu": "RTX4090", "gpu_memory": "24"}
    flavors = [
        FakeFlavor(
            "a",
            "GL10.4-16384-0-1GPU",
            extra,
            vcpus=4,
            ram=16384,
            availability_zones=["ru-6a"],
        ),
        FakeFlavor(
            "h200",
            "GL-H200",
            {"aggregate_instance_extra_specs:gpu": "H200"},
            vcpus=12,
            ram=120000,
            availability_zones=["ru-6b"],
        ),
    ]
    in_a = rank_flavors(flavors, ("4090-24", "a5000"), az="ru-6a")
    assert [x.id for x in in_a] == ["a"]
    in_b = rank_flavors(flavors, ("4090-24",), az="ru-6b")
    assert in_b == []


def test_pick_boot_image_prefers_24_580_without_docker():
    class Img:
        def __init__(self, name):
            self.name = name

    images = [
        Img("Ubuntu 24.04 LTS 64-bit GPU Driver 580 Docker"),
        Img("Ubuntu 24.04 LTS 64-bit GPU Driver 535"),
        Img("Ubuntu 24.04 LTS 64-bit GPU Driver 580"),
        Img("Data Science VM (Ubuntu 22.04 LTS 64-bit)"),
    ]
    picked = pick_boot_image(images)
    assert picked.name == "Ubuntu 24.04 LTS 64-bit GPU Driver 580"
    assert gpu_boot_image_score(picked.name) > gpu_boot_image_score(images[0].name)