26de44feb1
The UmbraAudio side of the cross-machine contract: this node registers with the UmbraTrans controller, polls for jobs, claims them under a lease, and writes artifacts back over NFS. * `umbraaudio/cluster/` - configuration, PostgreSQL access, node registration and heartbeat, the job queue, artifact storage and NFS mount management. * `scripts/cluster_worker.py` - the CLI: `worker`, `check-cluster`, `register-node`. * `scripts/backup_system.sh` / `scripts/verify_backup.sh` - whole-system backup and integrity verification. * `.env` - the `UMBRA_CLUSTER__*` settings. `UMBRA_CLUSTER__PG_PASSWORD` is deliberately empty; the password comes from the environment, not from a committed file. `k3s-agent.yaml` is NOT committed: it contains a K3s server join token. It is now gitignored, and the token should be rotated if that file has ever left this machine. ruff clean on the whole package.
582 lines
19 KiB
Python
582 lines
19 KiB
Python
"""PostgreSQL database operations for UmbraTrans cluster integration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import asyncpg
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator
|
|
from datetime import datetime
|
|
|
|
from .config import ClusterConfig
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class JobRecord:
|
|
"""Job record from the queue."""
|
|
id: uuid.UUID
|
|
organization_id: uuid.UUID
|
|
project_id: uuid.UUID | None
|
|
kind: str
|
|
payload: dict[str, Any]
|
|
status: str
|
|
priority: int
|
|
attempt: int
|
|
max_attempts: int
|
|
not_before: datetime
|
|
lease_owner: str | None
|
|
lease_expires_at: datetime | None
|
|
started_at: datetime | None
|
|
completed_at: datetime | None
|
|
result: dict[str, Any] | None
|
|
error: str | None
|
|
created_at: datetime
|
|
updated_at: datetime
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ArtifactRecord:
|
|
"""Artifact record."""
|
|
id: uuid.UUID
|
|
project_id: uuid.UUID
|
|
organization_id: uuid.UUID
|
|
lang: str
|
|
kind: str
|
|
status: str
|
|
content_hash: str
|
|
byte_size: int
|
|
media_type: str
|
|
storage_key: str
|
|
storage_uri: str
|
|
integrity: str
|
|
derivation: dict[str, Any] | None
|
|
created_at: datetime
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ComputeNodeRecord:
|
|
"""Compute node record."""
|
|
id: str
|
|
organization_id: str
|
|
name: str
|
|
is_controller: bool
|
|
trust_domain: str
|
|
locality: str
|
|
transport: str
|
|
status: str
|
|
agent_version: str
|
|
agent_started_at: datetime
|
|
created_at: datetime
|
|
updated_at: datetime
|
|
|
|
|
|
class ClusterDB:
|
|
"""Async PostgreSQL connection pool for cluster operations."""
|
|
|
|
def __init__(self, config: ClusterConfig) -> None:
|
|
self.config = config
|
|
self._pool: asyncpg.Pool | None = None
|
|
|
|
async def connect(self) -> None:
|
|
"""Create connection pool."""
|
|
if self._pool is not None:
|
|
return
|
|
|
|
self._pool = await asyncpg.create_pool(
|
|
self.config.pg_dsn,
|
|
min_size=2,
|
|
max_size=10,
|
|
command_timeout=30,
|
|
)
|
|
logger.info("Cluster DB pool created: %s", self.config.pg_host)
|
|
|
|
async def close(self) -> None:
|
|
"""Close connection pool."""
|
|
if self._pool is not None:
|
|
await self._pool.close()
|
|
self._pool = None
|
|
logger.info("Cluster DB pool closed")
|
|
|
|
@asynccontextmanager
|
|
async def acquire(self) -> AsyncIterator[asyncpg.Connection]:
|
|
"""Acquire a connection from the pool."""
|
|
if self._pool is None:
|
|
await self.connect()
|
|
assert self._pool is not None
|
|
async with self._pool.acquire() as conn:
|
|
yield conn
|
|
|
|
# --- Job Queue Operations ---
|
|
|
|
async def claim_job(
|
|
self,
|
|
node_id: str,
|
|
allowed_kinds: list[str],
|
|
organization_id: str,
|
|
) -> JobRecord | None:
|
|
"""Atomically claim a queued job."""
|
|
async with self.acquire() as conn:
|
|
row = await conn.fetchrow(
|
|
"""
|
|
UPDATE public.jobs
|
|
SET status = 'running',
|
|
lease_owner = $1,
|
|
lease_expires_at = now() + interval '30 minutes',
|
|
started_at = now(),
|
|
attempt = attempt + 1,
|
|
updated_at = now()
|
|
WHERE id = (
|
|
SELECT id FROM public.jobs
|
|
WHERE status = 'queued'
|
|
AND not_before <= now()
|
|
AND (lease_expires_at IS NULL OR lease_expires_at <= now())
|
|
AND kind = ANY($2)
|
|
AND organization_id = $3
|
|
ORDER BY priority, not_before, created_at
|
|
FOR UPDATE SKIP LOCKED
|
|
LIMIT 1
|
|
)
|
|
RETURNING *;
|
|
""",
|
|
node_id,
|
|
allowed_kinds,
|
|
uuid.UUID(organization_id),
|
|
)
|
|
if row is None:
|
|
return None
|
|
return self._row_to_job(row)
|
|
|
|
async def complete_job(
|
|
self,
|
|
job_id: uuid.UUID,
|
|
node_id: str,
|
|
status: str,
|
|
result: dict[str, Any] | None = None,
|
|
error: str | None = None,
|
|
artifacts: list[dict[str, Any]] | None = None,
|
|
) -> bool:
|
|
"""Mark job as complete/failed and link artifacts."""
|
|
async with self.acquire() as conn, conn.transaction():
|
|
# Update job
|
|
updated = await conn.execute(
|
|
"""
|
|
UPDATE public.jobs
|
|
SET status = $1,
|
|
lease_owner = NULL,
|
|
lease_expires_at = NULL,
|
|
completed_at = now(),
|
|
result = $2,
|
|
error = $3,
|
|
updated_at = now()
|
|
WHERE id = $4 AND lease_owner = $5
|
|
""",
|
|
status,
|
|
json.dumps(result) if result else None,
|
|
error,
|
|
job_id,
|
|
node_id,
|
|
)
|
|
if updated == "UPDATE 0":
|
|
return False
|
|
|
|
# Link artifacts if provided
|
|
if artifacts:
|
|
for art in artifacts:
|
|
await conn.execute(
|
|
"""
|
|
INSERT INTO public.job_artifacts (job_id, artifact_id, role, created_at)
|
|
VALUES ($1, $2, $3, now())
|
|
ON CONFLICT DO NOTHING
|
|
""",
|
|
job_id,
|
|
uuid.UUID(art["artifact_id"]),
|
|
art.get("role", "produced"),
|
|
)
|
|
return True
|
|
|
|
async def requeue_job(self, job_id: uuid.UUID, node_id: str, error: str) -> bool:
|
|
"""Release job lease and requeue with error."""
|
|
async with self.acquire() as conn:
|
|
updated = await conn.execute(
|
|
"""
|
|
UPDATE public.jobs
|
|
SET status = 'queued',
|
|
lease_owner = NULL,
|
|
lease_expires_at = NULL,
|
|
error = $1,
|
|
updated_at = now()
|
|
WHERE id = $2 AND lease_owner = $3
|
|
""",
|
|
error,
|
|
job_id,
|
|
node_id,
|
|
)
|
|
return updated != "UPDATE 0"
|
|
|
|
async def get_job(self, job_id: uuid.UUID) -> JobRecord | None:
|
|
"""Get job by ID."""
|
|
async with self.acquire() as conn:
|
|
row = await conn.fetchrow("SELECT * FROM public.jobs WHERE id = $1", job_id)
|
|
return self._row_to_job(row) if row else None
|
|
|
|
# --- Artifact Operations ---
|
|
|
|
async def insert_artifact(
|
|
self,
|
|
project_id: uuid.UUID,
|
|
organization_id: uuid.UUID,
|
|
lang: str,
|
|
kind: str,
|
|
content_hash: str,
|
|
byte_size: int,
|
|
media_type: str,
|
|
storage_key: str,
|
|
storage_uri: str,
|
|
derivation: dict[str, Any] | None = None,
|
|
) -> ArtifactRecord:
|
|
"""Insert artifact (idempotent on content_hash)."""
|
|
async with self.acquire() as conn:
|
|
row = await conn.fetchrow(
|
|
"""
|
|
INSERT INTO public.artifacts (
|
|
id, project_id, organization_id, lang, kind, status,
|
|
content_hash, byte_size, media_type, storage_key,
|
|
storage_uri, integrity, derivation, created_at
|
|
) VALUES (
|
|
gen_random_uuid(), $1, $2, $3, $4, 'final',
|
|
$5, $6, $7, $8, $9, 'unverified', $10, now()
|
|
)
|
|
ON CONFLICT (project_id, content_hash, kind, lang) DO UPDATE SET
|
|
status = EXCLUDED.status,
|
|
byte_size = EXCLUDED.byte_size,
|
|
media_type = EXCLUDED.media_type,
|
|
storage_key = EXCLUDED.storage_key,
|
|
storage_uri = EXCLUDED.storage_uri,
|
|
derivation = EXCLUDED.derivation,
|
|
updated_at = now()
|
|
RETURNING *;
|
|
""",
|
|
project_id,
|
|
organization_id,
|
|
lang,
|
|
kind,
|
|
content_hash,
|
|
byte_size,
|
|
media_type,
|
|
storage_key,
|
|
storage_uri,
|
|
json.dumps(derivation) if derivation else None,
|
|
)
|
|
return self._row_to_artifact(row)
|
|
|
|
async def get_artifact(self, artifact_id: uuid.UUID) -> ArtifactRecord | None:
|
|
"""Get artifact by ID."""
|
|
async with self.acquire() as conn:
|
|
row = await conn.fetchrow("SELECT * FROM public.artifacts WHERE id = $1", artifact_id)
|
|
return self._row_to_artifact(row) if row else None
|
|
|
|
async def get_artifacts_by_job(self, job_id: uuid.UUID) -> list[ArtifactRecord]:
|
|
"""Get all artifacts linked to a job."""
|
|
async with self.acquire() as conn:
|
|
rows = await conn.fetch(
|
|
"""
|
|
SELECT a.* FROM public.artifacts a
|
|
JOIN public.job_artifacts ja ON ja.artifact_id = a.id
|
|
WHERE ja.job_id = $1
|
|
ORDER BY a.created_at
|
|
""",
|
|
job_id,
|
|
)
|
|
return [self._row_to_artifact(r) for r in rows]
|
|
|
|
# --- Compute Node Operations ---
|
|
|
|
async def register_node(
|
|
self,
|
|
node_id: str,
|
|
organization_id: str,
|
|
name: str,
|
|
is_controller: bool,
|
|
trust_domain: str,
|
|
locality: str,
|
|
transport: str,
|
|
agent_version: str,
|
|
) -> ComputeNodeRecord:
|
|
"""Register or update compute node."""
|
|
async with self.acquire() as conn:
|
|
row = await conn.fetchrow(
|
|
"""
|
|
INSERT INTO public.compute_nodes (
|
|
id, organization_id, name, is_controller, trust_domain,
|
|
locality, transport, status, agent_version, agent_started_at, created_at
|
|
) VALUES (
|
|
$1, $2, $3, $4, $5, $6, $7, 'provisioning', $8, now(), now()
|
|
)
|
|
ON CONFLICT (organization_id, name) DO UPDATE SET
|
|
status = 'provisioning',
|
|
agent_version = EXCLUDED.agent_version,
|
|
agent_started_at = EXCLUDED.agent_started_at,
|
|
updated_at = now()
|
|
RETURNING *;
|
|
""",
|
|
node_id,
|
|
organization_id,
|
|
name,
|
|
is_controller,
|
|
trust_domain,
|
|
locality,
|
|
transport,
|
|
agent_version,
|
|
)
|
|
return self._row_to_node(row)
|
|
|
|
async def set_node_online(self, node_id: str) -> None:
|
|
"""Mark node as online (accepting jobs)."""
|
|
async with self.acquire() as conn:
|
|
await conn.execute(
|
|
"UPDATE public.compute_nodes SET status = 'online', "
|
|
"updated_at = now() WHERE id = $1",
|
|
node_id,
|
|
)
|
|
|
|
async def set_node_draining(self, node_id: str) -> None:
|
|
"""Mark node as draining (finishing jobs, not accepting new)."""
|
|
async with self.acquire() as conn:
|
|
await conn.execute(
|
|
"UPDATE public.compute_nodes SET status = 'draining', "
|
|
"updated_at = now() WHERE id = $1",
|
|
node_id,
|
|
)
|
|
|
|
async def heartbeat(
|
|
self,
|
|
node_id: str,
|
|
cpu_busy_pct: float,
|
|
ram_free_mb: int,
|
|
vram_free_mb: int,
|
|
vram_total_mb: int,
|
|
disk_free_mb: int,
|
|
load_avg_1m: float,
|
|
resident_models: list[dict[str, Any]],
|
|
active_jobs: int,
|
|
agent_version: str,
|
|
raw: dict[str, Any] | None = None,
|
|
) -> None:
|
|
"""Update node heartbeat."""
|
|
async with self.acquire() as conn:
|
|
await conn.execute(
|
|
"""
|
|
INSERT INTO public.node_heartbeats (
|
|
compute_node_id, reported_at, cpu_busy_pct, ram_free_mb,
|
|
vram_free_mb, vram_total_mb, disk_free_mb, load_avg_1m,
|
|
resident_models, active_jobs, agent_version, raw
|
|
) VALUES (
|
|
$1, now(), $2, $3, $4, $5, $6, $7, $8, $9, $10, $11
|
|
)
|
|
ON CONFLICT (compute_node_id) DO UPDATE SET
|
|
reported_at = EXCLUDED.reported_at,
|
|
cpu_busy_pct = EXCLUDED.cpu_busy_pct,
|
|
ram_free_mb = EXCLUDED.ram_free_mb,
|
|
vram_free_mb = EXCLUDED.vram_free_mb,
|
|
vram_total_mb = EXCLUDED.vram_total_mb,
|
|
disk_free_mb = EXCLUDED.disk_free_mb,
|
|
load_avg_1m = EXCLUDED.load_avg_1m,
|
|
resident_models = EXCLUDED.resident_models,
|
|
active_jobs = EXCLUDED.active_jobs,
|
|
agent_version = EXCLUDED.agent_version,
|
|
raw = EXCLUDED.raw;
|
|
""",
|
|
node_id,
|
|
cpu_busy_pct,
|
|
ram_free_mb,
|
|
vram_free_mb,
|
|
vram_total_mb,
|
|
disk_free_mb,
|
|
load_avg_1m,
|
|
json.dumps(resident_models),
|
|
active_jobs,
|
|
agent_version,
|
|
json.dumps(raw) if raw else None,
|
|
)
|
|
|
|
async def register_device(
|
|
self,
|
|
node_id: str,
|
|
kind: str,
|
|
label: str,
|
|
vram_total_mb: int,
|
|
vram_free_mb: int,
|
|
ram_total_mb: int,
|
|
ram_free_mb: int,
|
|
compute_units: int,
|
|
throughput_factor: float,
|
|
supports: dict[str, Any],
|
|
) -> None:
|
|
"""Register compute device (GPU/CPU)."""
|
|
async with self.acquire() as conn:
|
|
await conn.execute(
|
|
"""
|
|
INSERT INTO public.compute_devices (
|
|
id, compute_node_id, kind, label, vram_total_mb, vram_free_mb,
|
|
ram_total_mb, ram_free_mb, compute_units, throughput_factor, supports
|
|
) VALUES (
|
|
gen_random_uuid(), $1, $2, $3, $4, $5, $6, $7, $8, $9, $10
|
|
)
|
|
ON CONFLICT (compute_node_id, label) DO UPDATE SET
|
|
vram_total_mb = EXCLUDED.vram_total_mb,
|
|
vram_free_mb = EXCLUDED.vram_free_mb,
|
|
ram_total_mb = EXCLUDED.ram_total_mb,
|
|
ram_free_mb = EXCLUDED.ram_free_mb,
|
|
updated_at = now();
|
|
""",
|
|
node_id,
|
|
kind,
|
|
label,
|
|
vram_total_mb,
|
|
vram_free_mb,
|
|
ram_total_mb,
|
|
ram_free_mb,
|
|
compute_units,
|
|
throughput_factor,
|
|
json.dumps(supports),
|
|
)
|
|
|
|
async def report_model_loaded(
|
|
self,
|
|
node_id: str,
|
|
model_id: uuid.UUID,
|
|
device_id: uuid.UUID,
|
|
vram_held_mb: int,
|
|
) -> None:
|
|
"""Report model loaded on node."""
|
|
async with self.acquire() as conn:
|
|
await conn.execute(
|
|
"""
|
|
INSERT INTO public.node_models (
|
|
compute_node_id, model_id, device_id, state, loaded_at, vram_held_mb
|
|
) VALUES ($1, $2, $3, 'loaded', now(), $4)
|
|
ON CONFLICT (compute_node_id, model_id) DO UPDATE SET
|
|
state = EXCLUDED.state,
|
|
loaded_at = EXCLUDED.loaded_at,
|
|
vram_held_mb = EXCLUDED.vram_held_mb,
|
|
updated_at = now();
|
|
""",
|
|
node_id,
|
|
model_id,
|
|
device_id,
|
|
vram_held_mb,
|
|
)
|
|
|
|
async def report_model_unloaded(self, node_id: str, model_id: uuid.UUID) -> None:
|
|
"""Report model unloaded from node."""
|
|
async with self.acquire() as conn:
|
|
await conn.execute(
|
|
"""
|
|
UPDATE public.node_models
|
|
SET state = 'available', vram_held_mb = 0, updated_at = now()
|
|
WHERE compute_node_id = $1 AND model_id = $2
|
|
""",
|
|
node_id,
|
|
model_id,
|
|
)
|
|
|
|
# --- Helpers ---
|
|
|
|
def _row_to_job(self, row: asyncpg.Record) -> JobRecord:
|
|
return JobRecord(
|
|
id=row["id"],
|
|
organization_id=row["organization_id"],
|
|
project_id=row["project_id"],
|
|
kind=row["kind"],
|
|
payload=(
|
|
row["payload"]
|
|
if isinstance(row["payload"], dict)
|
|
else json.loads(row["payload"])
|
|
),
|
|
status=row["status"],
|
|
priority=row["priority"],
|
|
attempt=row["attempt"],
|
|
max_attempts=row["max_attempts"],
|
|
not_before=row["not_before"],
|
|
lease_owner=row["lease_owner"],
|
|
lease_expires_at=row["lease_expires_at"],
|
|
started_at=row["started_at"],
|
|
completed_at=row["completed_at"],
|
|
result=(
|
|
row["result"]
|
|
if isinstance(row["result"], dict)
|
|
else (json.loads(row["result"]) if row["result"] else None)
|
|
),
|
|
error=row["error"],
|
|
created_at=row["created_at"],
|
|
updated_at=row["updated_at"],
|
|
)
|
|
|
|
def _row_to_artifact(self, row: asyncpg.Record) -> ArtifactRecord:
|
|
return ArtifactRecord(
|
|
id=row["id"],
|
|
project_id=row["project_id"],
|
|
organization_id=row["organization_id"],
|
|
lang=row["lang"],
|
|
kind=row["kind"],
|
|
status=row["status"],
|
|
content_hash=row["content_hash"],
|
|
byte_size=row["byte_size"],
|
|
media_type=row["media_type"],
|
|
storage_key=row["storage_key"],
|
|
storage_uri=row["storage_uri"],
|
|
integrity=row["integrity"],
|
|
derivation=(
|
|
row["derivation"]
|
|
if isinstance(row["derivation"], dict)
|
|
else (json.loads(row["derivation"]) if row["derivation"] else None)
|
|
),
|
|
created_at=row["created_at"],
|
|
)
|
|
|
|
def _row_to_node(self, row: asyncpg.Record) -> ComputeNodeRecord:
|
|
return ComputeNodeRecord(
|
|
id=row["id"],
|
|
organization_id=row["organization_id"],
|
|
name=row["name"],
|
|
is_controller=row["is_controller"],
|
|
trust_domain=row["trust_domain"],
|
|
locality=row["locality"],
|
|
transport=row["transport"],
|
|
status=row["status"],
|
|
agent_version=row["agent_version"],
|
|
agent_started_at=row["agent_started_at"],
|
|
created_at=row["created_at"],
|
|
updated_at=row["updated_at"],
|
|
)
|
|
|
|
|
|
# Global instance
|
|
_db_instance: ClusterDB | None = None
|
|
|
|
|
|
async def get_cluster_db(config: ClusterConfig) -> ClusterDB:
|
|
"""Get or create global cluster DB instance."""
|
|
global _db_instance
|
|
if _db_instance is None:
|
|
_db_instance = ClusterDB(config)
|
|
await _db_instance.connect()
|
|
return _db_instance
|
|
|
|
|
|
async def close_cluster_db() -> None:
|
|
"""Close global cluster DB instance."""
|
|
global _db_instance
|
|
if _db_instance is not None:
|
|
await _db_instance.close()
|
|
_db_instance = None
|