From 0888f0f16b5d87aba54d427481c9cec14b860ab5 Mon Sep 17 00:00:00 2001
From: pyr0ball <pyroballpcs@gmail.com>
Date: Mon, 30 Mar 2026 20:21:37 -0700
Subject: [PATCH] feat(resources): add shared VRAMLease, GpuInfo, NodeInfo
 models

---
 circuitforge_core/resources/__init__.py |  0
 circuitforge_core/resources/models.py   | 56 +++++++++++++++++++++++++
 tests/test_resources/__init__.py        |  0
 tests/test_resources/test_models.py     | 47 +++++++++++++++++++++
 4 files changed, 103 insertions(+)
 create mode 100644 circuitforge_core/resources/__init__.py
 create mode 100644 circuitforge_core/resources/models.py
 create mode 100644 tests/test_resources/__init__.py
 create mode 100644 tests/test_resources/test_models.py

diff --git a/circuitforge_core/resources/__init__.py b/circuitforge_core/resources/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/circuitforge_core/resources/models.py b/circuitforge_core/resources/models.py
new file mode 100644
index 0000000..45c6829
--- /dev/null
+++ b/circuitforge_core/resources/models.py
@@ -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)
diff --git a/tests/test_resources/__init__.py b/tests/test_resources/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/test_resources/test_models.py b/tests/test_resources/test_models.py
new file mode 100644
index 0000000..cdef834
--- /dev/null
+++ b/tests/test_resources/test_models.py
@@ -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