feat(resources): add shared VRAMLease, GpuInfo, NodeInfo models
This commit is contained in:
parent
56042dffba
commit
0888f0f16b
4 changed files with 103 additions and 0 deletions
0
circuitforge_core/resources/__init__.py
Normal file
0
circuitforge_core/resources/__init__.py
Normal file
56
circuitforge_core/resources/models.py
Normal file
56
circuitforge_core/resources/models.py
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VRAMLease:
|
||||
lease_id: str
|
||||
gpu_id: int
|
||||
node_id: str
|
||||
mb_granted: int
|
||||
holder_service: str
|
||||
priority: int
|
||||
expires_at: float
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
gpu_id: int,
|
||||
node_id: str,
|
||||
mb: int,
|
||||
service: str,
|
||||
priority: int,
|
||||
ttl_s: float = 0.0,
|
||||
) -> VRAMLease:
|
||||
return cls(
|
||||
lease_id=str(uuid.uuid4()),
|
||||
gpu_id=gpu_id,
|
||||
node_id=node_id,
|
||||
mb_granted=mb,
|
||||
holder_service=service,
|
||||
priority=priority,
|
||||
expires_at=time.time() + ttl_s if ttl_s > 0.0 else 0.0,
|
||||
)
|
||||
|
||||
def is_expired(self) -> bool:
|
||||
return self.expires_at > 0.0 and time.time() > self.expires_at
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GpuInfo:
|
||||
gpu_id: int
|
||||
name: str
|
||||
vram_total_mb: int
|
||||
vram_used_mb: int
|
||||
vram_free_mb: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class NodeInfo:
|
||||
node_id: str
|
||||
agent_url: str
|
||||
gpus: list[GpuInfo]
|
||||
last_heartbeat: float = field(default_factory=time.time)
|
||||
0
tests/test_resources/__init__.py
Normal file
0
tests/test_resources/__init__.py
Normal file
47
tests/test_resources/test_models.py
Normal file
47
tests/test_resources/test_models.py
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
import time
|
||||
from circuitforge_core.resources.models import VRAMLease, GpuInfo, NodeInfo
|
||||
|
||||
|
||||
def test_vram_lease_create_assigns_unique_ids():
|
||||
lease_a = VRAMLease.create(gpu_id=0, node_id="heimdall", mb=4096,
|
||||
service="peregrine", priority=1)
|
||||
lease_b = VRAMLease.create(gpu_id=0, node_id="heimdall", mb=4096,
|
||||
service="peregrine", priority=1)
|
||||
assert lease_a.lease_id != lease_b.lease_id
|
||||
|
||||
|
||||
def test_vram_lease_create_with_ttl_sets_expiry():
|
||||
before = time.time()
|
||||
lease = VRAMLease.create(gpu_id=0, node_id="heimdall", mb=2048,
|
||||
service="kiwi", priority=2, ttl_s=60.0)
|
||||
after = time.time()
|
||||
assert before + 60.0 <= lease.expires_at <= after + 60.0
|
||||
|
||||
|
||||
def test_vram_lease_create_no_ttl_has_zero_expiry():
|
||||
lease = VRAMLease.create(gpu_id=0, node_id="heimdall", mb=1024,
|
||||
service="snipe", priority=2)
|
||||
assert lease.expires_at == 0.0
|
||||
|
||||
|
||||
def test_vram_lease_is_immutable():
|
||||
lease = VRAMLease.create(gpu_id=0, node_id="heimdall", mb=1024,
|
||||
service="snipe", priority=2)
|
||||
import pytest
|
||||
with pytest.raises((AttributeError, TypeError)):
|
||||
lease.mb_granted = 999 # type: ignore
|
||||
|
||||
|
||||
def test_gpu_info_fields():
|
||||
info = GpuInfo(gpu_id=0, name="RTX 4000", vram_total_mb=8192,
|
||||
vram_used_mb=2048, vram_free_mb=6144)
|
||||
assert info.vram_free_mb == 6144
|
||||
|
||||
|
||||
def test_node_info_fields():
|
||||
gpu = GpuInfo(gpu_id=0, name="RTX 4000", vram_total_mb=8192,
|
||||
vram_used_mb=0, vram_free_mb=8192)
|
||||
node = NodeInfo(node_id="heimdall", agent_url="http://localhost:7701",
|
||||
gpus=[gpu], last_heartbeat=time.time())
|
||||
assert node.node_id == "heimdall"
|
||||
assert len(node.gpus) == 1
|
||||
Loading…
Reference in a new issue