feat(resources): add shared VRAMLease, GpuInfo, NodeInfo models

This commit is contained in:
pyr0ball 2026-03-30 20:21:37 -07:00
parent 56042dffba
commit 0888f0f16b
4 changed files with 103 additions and 0 deletions

View file

View 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)

View file

View 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