58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
from gpu_rent.inventory import (
|
|
gpu_boot_image_score,
|
|
gpu_quota_from_compute,
|
|
match_label,
|
|
pick_boot_image,
|
|
rank_flavors,
|
|
)
|
|
|
|
|
|
class FakeFlavor:
|
|
def __init__(self, flavor_id: str, name: str, extra=None, disabled: bool = False):
|
|
self.id = flavor_id
|
|
self.name = name
|
|
self.extra_specs = extra or {}
|
|
self.is_disabled = disabled
|
|
self.vcpus = 8
|
|
self.ram = 32768
|
|
|
|
|
|
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_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)
|