UmbraAudio/umbraaudio/cluster/node.py

# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""Compute node registration and heartbeat for UmbraAudio worker."""

from __future__ import annotations

import asyncio
import contextlib
import logging
import platform
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

import psutil

from .database import ClusterDB, get_cluster_db

if TYPE_CHECKING:
    from .config import ClusterConfig

logger = logging.getLogger(__name__)


@dataclass
class ComputeNode:
    """Compute node state."""
    config: ClusterConfig
    db: ClusterDB | None = None
    _heartbeat_task: asyncio.Task | None = field(default=None, init=False)
    _running: bool = field(default=False, init=False)

    async def start(self) -> None:
        """Start node: register, register devices, start heartbeat."""
        self.db = await get_cluster_db(self.config)
        await self.register()
        await self.register_devices()
        await self.set_online()
        self._running = True
        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
        logger.info("Compute node %s started", self.config.node_id)

    async def stop(self) -> None:
        """Stop node: set draining, stop heartbeat."""
        self._running = False
        if self._heartbeat_task:
            self._heartbeat_task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._heartbeat_task
        if self.db:
            await self.set_draining()
        logger.info("Compute node %s stopped", self.config.node_id)

    async def register(self) -> None:
        """Register compute node in cluster."""
        assert self.db is not None
        await self.db.register_node(
            node_id=self.config.node_id,
            organization_id=self.config.organization_id,
            name=self.config.node_name or self.config.node_id,
            is_controller=False,
            trust_domain="remote",
            locality=self.config.datacenter,
            transport="https",
            agent_version=self.config.agent_version,
        )
        logger.info("Registered compute node: %s", self.config.node_id)

    async def register_devices(self) -> None:
        """Register GPUs and CPUs."""
        assert self.db is not None

        # Register GPUs
        try:
            import torch
            if torch.cuda.is_available():
                for i in range(torch.cuda.device_count()):
                    props = torch.cuda.get_device_properties(i)
                    vram_total = props.total_memory // (1024 * 1024)
                    # Approximate free VRAM
                    vram_free = vram_total - (torch.cuda.memory_allocated(i) // (1024 * 1024))
                    await self.db.register_device(
                        node_id=self.config.node_id,
                        kind="gpu_discrete",
                        label=f"gpu:{i}",
                        vram_total_mb=vram_total,
                        vram_free_mb=max(0, vram_free),
                        ram_total_mb=0,
                        ram_free_mb=0,
                        compute_units=props.multi_processor_count,
                        throughput_factor=1.0,
                        supports={"cuda": torch.version.cuda or "unknown"},
                    )
                    logger.info("Registered GPU %d: %s (%d MB VRAM)", i, props.name, vram_total)
        except Exception as e:
            logger.warning("Could not register GPUs: %s", e)

        # Register CPU as fallback
        ram = psutil.virtual_memory()
        await self.db.register_device(
            node_id=self.config.node_id,
            kind="cpu",
            label="cpu:0",
            vram_total_mb=0,
            vram_free_mb=0,
            ram_total_mb=ram.total // (1024 * 1024),
            ram_free_mb=ram.available // (1024 * 1024),
            compute_units=psutil.cpu_count(logical=True) or 1,
            throughput_factor=1.0,
            supports={"arch": platform.machine()},
        )
        logger.info(
            "Registered CPU: %d cores, %d MB RAM",
            psutil.cpu_count(),
            ram.total // (1024 * 1024),
        )

    async def set_online(self) -> None:
        """Mark node as online (accepting jobs)."""
        assert self.db is not None
        await self.db.set_node_online(self.config.node_id)

    async def set_draining(self) -> None:
        """Mark node as draining."""
        assert self.db is not None
        await self.db.set_node_draining(self.config.node_id)

    async def _heartbeat_loop(self) -> None:
        """Send heartbeat every interval."""
        interval = self.config.heartbeat_interval_seconds
        while self._running:
            try:
                await self._send_heartbeat()
            except Exception as e:
                logger.error("Heartbeat failed: %s", e)
            await asyncio.sleep(interval)

    async def _send_heartbeat(self) -> None:
        """Send single heartbeat with system metrics."""
        assert self.db is not None

        # System metrics
        cpu_pct = psutil.cpu_percent(interval=0.1)
        ram = psutil.virtual_memory()
        disk = psutil.disk_usage("/")
        load_avg = psutil.getloadavg()[0] if hasattr(psutil, "getloadavg") else 0.0

        # GPU metrics
        vram_free = 0
        vram_total = 0
        resident_models = []
        try:
            import torch
            if torch.cuda.is_available():
                for i in range(torch.cuda.device_count()):
                    torch.cuda.memory_allocated(i)
                    reserved = torch.cuda.memory_reserved(i)
                    free = (torch.cuda.get_device_properties(i).total_memory - reserved)
                    vram_free += free // (1024 * 1024)
                    vram_total += torch.cuda.get_device_properties(i).total_memory // (1024 * 1024)
        except Exception:
            pass

        # Active jobs - placeholder (would track from job queue)
        active_jobs = 0

        await self.db.heartbeat(
            node_id=self.config.node_id,
            cpu_busy_pct=cpu_pct,
            ram_free_mb=ram.available // (1024 * 1024),
            vram_free_mb=vram_free,
            vram_total_mb=vram_total,
            disk_free_mb=disk.free // (1024 * 1024),
            load_avg_1m=load_avg,
            resident_models=resident_models,
            active_jobs=active_jobs,
            agent_version=self.config.agent_version,
            raw={
                "platform": platform.platform(),
                "python": platform.python_version(),
            },
        )


async def register_node(config: ClusterConfig) -> ComputeNode:
    """Create and start a compute node."""
    node = ComputeNode(config)
    await node.start()
    return node


async def heartbeat_loop(node: ComputeNode) -> None:
    """Run heartbeat loop (alternative to node.start())."""
    # This is handled by node.start()/_heartbeat_loop
    pass