Leonid Pershin
615cf81493
- Added package data configuration for the 'gpu_rent' package in pyproject.toml. - Updated README.md to include usage instructions for Windows and Unix launchers. - Enhanced CLI documentation in cli.md to reflect new commands and their functionalities. - Revised setup.md to clarify installation steps and environment setup. - Improved error handling and command descriptions in the CLI implementation. - Added new functions for model version handling and flavor resolution in the codebase. - Updated state management to include additional properties for better tracking.
91 lines
2.9 KiB
Python
91 lines
2.9 KiB
Python
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):
|
|
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_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_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)
|