fabric/shared/schemas.py

"""Shared Pydantic schemas for Fabric control plane and node-agent."""
from __future__ import annotations
from enum import Enum
from pydantic import BaseModel, Field
from typing import Optional
import time


class PriorityClass(str, Enum):
    low = "low"
    normal = "normal"
    high = "high"
    critical = "critical"


class ReclaimStrategy(str, Enum):
    process_kill = "process-kill"
    vllm_sleep = "vllm-sleep"
    cuda_checkpoint = "cuda-checkpoint"
    none = "none"


class ModuleState(str, Enum):
    loaded = "loaded"
    running = "running"
    idle = "idle"
    sleeping = "sleeping"
    unloaded = "unloaded"
    error = "error"
    warming = "warming"


class RuntimeType(str, Enum):
    systemd_wrap = "systemd-wrap"
    subprocess = "subprocess"
    vllm = "vllm"
    python_class = "python-class"
    comfy_graph = "comfy-graph"
    proc_exec = "proc-exec"


class GpuResources(BaseModel):
    gpus: int = 1
    vram_mb: int | str = "auto"


class LifecycleConfig(BaseModel):
    reclaim: ReclaimStrategy = ReclaimStrategy.process_kill
    idle_evict_s: int = 300
    warm_default: bool = False
    idle_evict_s: int = 300
    warm_default: bool = False
    warmup_s: int = 10


class SchedulingConfig(BaseModel):
    priority_class: PriorityClass = PriorityClass.normal
    min_replicas: int = 0
    max_replicas: int = 8
    batch: bool = False


class SystemdWrapConfig(BaseModel):
    systemd_unit: str
    health_url: str | None = None
    health_timeout_s: int = 120
    infer_url: str
    infer_method: str = "POST"


class ModuleManifest(BaseModel):
    name: str
    description: str = ""
    version: str = "1.0"
    modality: str = "general"
    aliases: list[str] = Field(default_factory=list)
    runtime: RuntimeType = RuntimeType.systemd_wrap
    resources: GpuResources = Field(default_factory=GpuResources)
    lifecycle: LifecycleConfig = Field(default_factory=LifecycleConfig)
    scheduling: SchedulingConfig = Field(default_factory=SchedulingConfig)
    systemd: SystemdWrapConfig | None = None
    capabilities: list[str] = Field(default_factory=list)


class GpuPidUsage(BaseModel):
    gpu_index: int
    pid: int
    vram_mb: int
    process_name: str = ""


class GpuMemInfo(BaseModel):
    """Per-physical-GPU memory capacity + telemetry (additive). One entry per
    physical GPU (counted once, NOT per process). Absent fields degrade to
    0 / None so old and new components interoperate."""
    gpu_index: int = 0
    name: str = ""
    mem_total_mb: int = 0
    mem_used_mb: int = 0
    mem_free_mb: int = 0
    util_pct: Optional[int] = None
    temp_c: Optional[int] = None
    power_w: Optional[float] = None


class ModuleStatus(BaseModel):
    name: str
    state: ModuleState
    vram_mb: int = 0
    required_vram_mb: int = 0  # declared manifest requirement (0 = auto/unknown)
    last_infer_at: float | None = None
    health_ok: bool = False
    infer_url: str = ""
    modality: str = "general"
    aliases: list[str] = Field(default_factory=list)
    priority_class: PriorityClass = PriorityClass.normal
    reclaim: ReclaimStrategy = ReclaimStrategy.process_kill
    idle_evict_s: int = 300
    warm_default: bool = False


class Heartbeat(BaseModel):
    node_id: str
    node_ip: str
    agent_port: int = 7701
    hypervisor: str = "unknown"
    gpu_status: list[GpuPidUsage] = Field(default_factory=list)
    modules: list[ModuleStatus] = Field(default_factory=list)
    uptime_s: int = 0
    timestamp: float = Field(default_factory=time.time)
    # --- additive VRAM-capacity + telemetry (per physical GPU) ---
    gpus: list[GpuMemInfo] = Field(default_factory=list)
    vram_total_mb: int = 0
    vram_free_mb: int = 0
    driver_version: str = ""


class NodeInfo(BaseModel):
    node_id: str
    node_ip: str
    agent_port: int = 7701
    hypervisor: str = "unknown"
    gpu_status: list[GpuPidUsage] = Field(default_factory=list)
    modules: list[ModuleStatus] = Field(default_factory=list)
    last_heartbeat: float = 0
    alive: bool = True
    # --- additive VRAM-capacity + telemetry (per physical GPU) ---
    gpus: list[GpuMemInfo] = Field(default_factory=list)
    vram_total_mb: int = 0
    vram_free_mb: int = 0
    driver_version: str = ""


class PlaceRequest(BaseModel):
    model: str
    priority: PriorityClass = PriorityClass.normal


class PlaceResponse(BaseModel):
    node_id: str
    node_ip: str
    infer_url: str
    module_name: str
    state: ModuleState


class RegistryModelsResponse(BaseModel):
    models: list[ModuleStatus]
    nodes_count: int
    total_gpu_vram_mb: int = 0
    total_gpu_capacity_mb: int = 0
    total_gpu_free_mb: int = 0


class RegistryNodesResponse(BaseModel):
    nodes: list[NodeInfo]