bf83bb35cc
The remote had one commit this branch did not (03f8a85), which adds the
Apache-2.0 LICENSE, the NOTICE and SPDX headers to every .py and .sh
file. The merge was clean - the headers sit at the top of each file and
my changes are further down.
The new files this branch added (languages.py, translations.py, the
cluster package, the new tests and scripts) were written after the
merge base, so the licence commit never saw them. They now carry the
same header, so the convention holds across the whole tree.
Nothing else changed. 883 tests passed / 8 failed - the same 8
pre-existing failures.
585 lines
19 KiB
Python
585 lines
19 KiB
Python
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""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
|