Cluster worker for the UmbraTrans integration

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.
This commit is contained in:
B.D. Murphy
2026-11-06 17:03:16 -08:00
parent 2b294d7fa4
commit 26de44feb1
13 changed files with 1989 additions and 0 deletions

32
.env
View File

@@ -46,3 +46,35 @@ UMBRA_AUDIO__NOISE_FLOOR_DB=-60.0
UMBRA_AUDIO__SILENCE_BETWEEN_UTTERANCES_MS=500
UMBRA_AUDIO__SILENCE_BETWEEN_PARAGRAPHS_MS=2000
UMBRA_AUDIO__SILENCE_BETWEEN_CHAPTERS_MS=4000
# --- Cluster (UmbraTrans integration) ---
UMBRA_CLUSTER__PG_HOST=192.168.123.14
UMBRA_CLUSTER__PG_PORT=6432
UMBRA_CLUSTER__PG_DATABASE=umbratrans
UMBRA_CLUSTER__PG_USER=umbraudio
UMBRA_CLUSTER__PG_PASSWORD=
UMBRA_CLUSTER__PG_SSLMODE=require
UMBRA_CLUSTER__PG_DIRECT_PORT=5432
UMBRA_CLUSTER__NFS_SERVER=192.168.123.14
UMBRA_CLUSTER__NFS_PRODUCTIONS_PATH=/mnt/umbra-data/productions
UMBRA_CLUSTER__NFS_ARTIFACTS_PATH=/mnt/umbra-data/artifacts
UMBRA_CLUSTER__NFS_COVERS_PATH=/mnt/umbra-data/covers
UMBRA_CLUSTER__NFS_AUDIO_PATH=/mnt/umbra-data/audio
UMBRA_CLUSTER__LOCAL_MOUNT_ROOT=/mnt/umbra
UMBRA_CLUSTER__NODE_ID=umbra-audio-01
UMBRA_CLUSTER__ORGANIZATION_ID=umbra
UMBRA_CLUSTER__DATACENTER=local
UMBRA_CLUSTER__AGENT_VERSION=1.0.0
UMBRA_CLUSTER__NODE_CREDENTIAL_PREFIX=umbra-audio-01-
UMBRA_CLUSTER__NODE_SCOPES=queue:claim,queue:complete,artifacts:write
UMBRA_CLUSTER__POLL_INTERVAL_SECONDS=10
UMBRA_CLUSTER__LEASE_DURATION_MINUTES=30
UMBRA_CLUSTER__ALLOWED_JOB_KINDS=tts,asr,assemble,encode,master,speaker_profile
UMBRA_CLUSTER__HEARTBEAT_INTERVAL_SECONDS=30
UMBRA_CLUSTER__HEARTBEAT_TIMEOUT_MINUTES=2

6
.gitignore vendored
View File

@@ -57,3 +57,9 @@ umbraaudio_package.zip
/backup/author_voice_lab/wav/
/backup/author_voice_lab/wav_cliff/
/backup/author_voice_matrix/**/*.wav
# K3s agent VM definition: contains a server join token. The cluster is LAN
# only, but a join token is still a credential - it is not committed. Copy
# the template and fill the token in from the keychain when provisioning a
# node. The token should be rotated if this file has ever left the machine.
/k3s-agent.yaml

View File

@@ -25,6 +25,7 @@ dependencies = [
# Database
"aiosqlite>=0.19",
"asyncpg>=0.29",
# Parsers
"ebooklib>=0.18",
@@ -65,6 +66,7 @@ dependencies = [
"numpy>=1.24",
"pydantic>=2.5",
"pydantic-settings>=2.1",
"psutil>=5.9",
]
[project.optional-dependencies]

134
scripts/backup_system.sh Executable file
View File

@@ -0,0 +1,134 @@
#!/bin/zsh
# backup_system.sh - Complete UmbraAudio system backup
set -e
SOURCE_ROOT="/Volumes/1TB_APL_Storage/UmbraAudio"
BACKUP_DIR="/tmp/backups"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
# Parse arguments
MODE="full"
while [[ $# -gt 0 ]]; do
case "$1" in
--full) MODE="full"; shift ;;
--db) MODE="db"; shift ;;
--voices) MODE="voices"; shift ;;
--config) MODE="config"; shift ;;
--help|-h)
echo "Usage: $0 --full | --db | --voices | --config"
exit 0 ;;
*) echo "Unknown option: $1"; exit 1 ;;
esac
done
mkdir -p "$BACKUP_DIR"
echo "========================================="
echo "UmbraAudio System Backup"
echo "========================================="
echo "Mode: $MODE"
echo "Timestamp: $TIMESTAMP"
echo "========================================="
backup_db() {
local name="$1"
local src="$SOURCE_ROOT/.umbra/$name"
if [[ -f "$src" ]]; then
cp -a "$src" "$BACKUP_DIR/project_${name}.${TIMESTAMP}"
echo " ✓ $name"
else
echo " ⊘ $name (not found)"
fi
}
# Copy directory to backup with parent dir creation
backup_dir() {
local name="$1"
local rel_path="$name" # relative to SOURCE_ROOT
local src="$SOURCE_ROOT/$rel_path"
if [[ -d "$src" ]]; then
# Destination: flat name under BACKUP_DIR
local dst="$BACKUP_DIR/${name}_${TIMESTAMP}"
# Ensure parent directory exists
local dst_parent=$(dirname "$dst")
mkdir -p "$dst_parent"
cp -a "$src" "$dst"
echo " ✓ $name"
else
echo " ⊘ $name (not found)"
fi
}
case "$MODE" in
full)
echo "Backing up complete system..."
echo ""
echo "1. Project databases:"
backup_db "project.db"
backup_db "voice_registry.db"
backup_db "umbra.db"
echo ""
echo "2. Voice references:"
backup_dir "umbraaudio/voices"
echo ""
echo "3. Speaker configurations:"
if [[ -f "$SOURCE_ROOT/speakers.json" ]]; then
cp -a "$SOURCE_ROOT/speakers.json" "$BACKUP_DIR/speakers.${TIMESTAMP}"
echo " ✓ speakers.json"
fi
echo ""
echo "4. Configuration:"
if [[ -f "$SOURCE_ROOT/.env" ]]; then
cp -a "$SOURCE_ROOT/.env" "$BACKUP_DIR/.env.${TIMESTAMP}"
echo " ✓ .env"
fi
if [[ -f "$SOURCE_ROOT/umbraaudio/config.py" ]]; then
cp -a "$SOURCE_ROOT/umbraaudio/config.py" "$BACKUP_DIR/config.${TIMESTAMP}"
echo " ✓ config.py"
fi
echo ""
echo "5. Git state:"
cd "$SOURCE_ROOT" 2>/dev/null && git status --short > "$BACKUP_DIR/git_status.${TIMESTAMP}" 2>/dev/null && echo " ✓ git_status" || echo " ⊘ git (failed)"
echo ""
echo "6. Launchd services:"
launchctl list > "$BACKUP_DIR/launchd.${TIMESTAMP}" 2>/dev/null && echo " ✓ launchd" || echo " ⊘ launchd (failed)"
echo ""
echo "7. Audio generation state:"
if [[ -d "$SOURCE_ROOT/generation" ]]; then
backup_dir "generation"
fi
echo ""
echo "========================================="
echo "✅ Full backup complete!"
echo "Files saved to: $BACKUP_DIR/"
echo "========================================="
;;
db)
echo "Backing up databases only..."
backup_db "project.db"
backup_db "voice_registry.db"
backup_db "umbra.db"
;;
voices)
echo "Backing up voice references..."
backup_dir "umbraaudio/voices"
if [[ -f "$SOURCE_ROOT/speakers.json" ]]; then
cp -a "$SOURCE_ROOT/speakers.json" "$BACKUP_DIR/speakers.${TIMESTAMP}"
fi
;;
config)
echo "Backing up configuration only..."
if [[ -f "$SOURCE_ROOT/.env" ]]; then
cp -a "$SOURCE_ROOT/.env" "$BACKUP_DIR/.env.${TIMESTAMP}"
echo " ✓ .env"
fi
if [[ -f "$SOURCE_ROOT/umbraaudio/config.py" ]]; then
cp -a "$SOURCE_ROOT/umbraaudio/config.py" "$BACKUP_DIR/config.${TIMESTAMP}"
echo " ✓ config.py"
fi
;;
esac

160
scripts/cluster_worker.py Normal file
View File

@@ -0,0 +1,160 @@
"""UmbraAudio cluster worker entry point.
Registers this node with UmbraTrans and polls for jobs.
"""
from __future__ import annotations
import asyncio
import contextlib
import logging
import signal
from pathlib import Path
from typing import TYPE_CHECKING
import click
from umbraaudio.cluster import JobKind, JobQueue, load_cluster_config, register_node
from umbraaudio.cluster.database import close_cluster_db
if TYPE_CHECKING:
from typing import Any
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)
@click.group()
def cli() -> None:
"""UmbraAudio cluster worker."""
pass
@cli.command()
@click.option("--project-dir", type=click.Path(exists=True, path_type=Path), default=".")
@click.option("--once", is_flag=True, help="Claim and process one job then exit")
def worker(project_dir: Path, once: bool) -> None:
"""Run the cluster worker (register node, poll for jobs)."""
asyncio.run(_run_worker(project_dir, once))
async def _run_worker(project_dir: Path, once: bool) -> None:
config = load_cluster_config(project_dir)
# Register node
node = await register_node(config)
# Create job queue
queue = JobQueue(config, node.db)
# Register job handlers
from umbraaudio.cluster.jobs import JobContext, JobHandler
class TTSHandler(JobHandler):
async def __call__(self, ctx: JobContext) -> dict[str, Any]:
logger.info("TTS job payload: %s", ctx.job.payload)
# TODO: Implement actual TTS processing
return {"status": "ok", "message": "TTS handler not yet implemented"}
class AssembleHandler(JobHandler):
async def __call__(self, ctx: JobContext) -> dict[str, Any]:
logger.info("Assemble job payload: %s", ctx.job.payload)
return {"status": "ok", "message": "Assemble handler not yet implemented"}
queue.register_handler(JobKind.TTS, TTSHandler())
queue.register_handler(JobKind.ASSEMBLE, AssembleHandler())
# Handle shutdown signals
shutdown = asyncio.Event()
def signal_handler(sig, frame):
logger.info("Received signal %s, shutting down...", sig)
shutdown.set()
for sig in (signal.SIGINT, signal.SIGTERM):
with contextlib.suppress(OSError, ValueError):
signal.signal(sig, signal_handler)
try:
if once:
# Claim and process one job
job = await queue.db.claim_job(
node_id=config.node_id,
allowed_kinds=list(queue._handlers.keys()),
organization_id=config.organization_id,
)
if job:
logger.info("Claimed job %s", job.id)
# Process would happen here
else:
logger.info("No jobs available")
else:
# Run continuous polling
await queue.start()
await shutdown.wait()
await queue.stop()
finally:
await node.stop()
await close_cluster_db()
@cli.command()
@click.option("--project-dir", type=click.Path(exists=True, path_type=Path), default=".")
def check_cluster(project_dir: Path) -> None:
"""Check cluster connectivity (PostgreSQL, NFS)."""
asyncio.run(_check_cluster(project_dir))
async def _check_cluster(project_dir: Path) -> None:
config = load_cluster_config(project_dir)
# Check PostgreSQL
from umbraaudio.cluster.database import ClusterDB
db = ClusterDB(config)
try:
await db.connect()
async with db.acquire() as conn:
version = await conn.fetchval("SELECT version()")
logger.info("PostgreSQL OK: %s", version.split(",")[0])
# Check tables
tables = await conn.fetch(
"SELECT tablename FROM pg_tables WHERE schemaname = 'public' ORDER BY tablename"
)
logger.info("Tables: %s", [t["tablename"] for t in tables])
except Exception as e:
logger.error("PostgreSQL check failed: %s", e)
finally:
await db.close()
# Check NFS
from umbraaudio.cluster.nfs import check_nfs_server, ensure_nfs_mounts
if check_nfs_server(config):
results = ensure_nfs_mounts(config)
for name, ok in results.items():
status = "OK" if ok else "FAILED"
logger.info("NFS %s: %s", name, status)
else:
logger.error("NFS server check failed")
@cli.command()
@click.option("--project-dir", type=click.Path(exists=True, path_type=Path), default=".")
def register_node_cmd(project_dir: Path) -> None:
"""Register this node with the cluster (one-time)."""
asyncio.run(_register_node_cmd(project_dir))
async def _register_node_cmd(project_dir: Path) -> None:
config = load_cluster_config(project_dir)
node = await register_node(config)
logger.info("Node registered: %s", config.node_id)
await node.stop()
await close_cluster_db()
if __name__ == "__main__":
cli()

115
scripts/verify_backup.sh Executable file
View File

@@ -0,0 +1,115 @@
#!/bin/zsh
# verify_backup.sh - Verify UmbraAudio backup integrity
set -e
SOURCE_DIR=""
BACKUP_DATE=""
# Parse arguments
while [[ $# -gt 0 ]]; do
case "$1" in
--source)
SOURCE_DIR="$2"
shift 2
;;
--date)
BACKUP_DATE="$2"
shift 2
;;
--help|-h)
echo "Usage: $0 --source DIR --date YYYYMMDD"
echo " --source Backup directory (default: /tmp/backups)"
echo " --date Date to verify (default: 7 days ago)"
exit 0
;;
*)
echo "Unknown option: $1"
exit 1
;;
esac
done
# Set defaults
if [[ -z "$SOURCE_DIR" ]]; then
SOURCE_DIR="/tmp/backups"
fi
if [[ -z "$BACKUP_DATE" ]]; then
BACKUP_DATE=$(date -v-7 +%Y%m%d 2>/dev/null || date --date='7 days ago' +%Y%m%d 2>/dev/null || echo "$(date +%Y%m%d)")
fi
echo "========================================="
echo "UmbraAudio Backup Verification"
echo "========================================="
echo "Source: $SOURCE_DIR"
echo "Date: $BACKUP_DATE"
echo "========================================="
# Validate source
if [[ ! -d "$SOURCE_DIR" ]]; then
echo "❌ Error: Source directory missing: $SOURCE_DIR"
exit 1
fi
# Find backup items matching the date pattern inside the source directory
# Look for files/dirs with the date pattern in their name
MATCHED=$(find "$SOURCE_DIR" -maxdepth 1 -name "*${BACKUP_DATE}*" 2>/dev/null | head -5)
if [[ -z "$MATCHED" ]]; then
# Try listing recent subdirectories
echo "No exact date match. Checking recent backups..."
MATCHED=$(ls -td "$SOURCE_DIR"/*/ 2>/dev/null | head -1)
if [[ -z "$MATCHED" ]]; then
echo "❌ No backups found in $SOURCE_DIR"
echo "Contents of $SOURCE_DIR:"
ls -la "$SOURCE_DIR" 2>/dev/null | head -10
exit 1
fi
NAME=$(basename "$MATCHED" | tr -d '/')
else
NAME=$(basename "$MATCHED" | tr -d '/')
fi
echo "✓ Found: $NAME"
echo ""
# Check if it's a file or directory inside backups
if [[ -f "$SOURCE_DIR/$NAME" ]]; then
TYPE="file"; FULL_PATH="$SOURCE_DIR/$NAME"
elif [[ -d "$SOURCE_DIR/$NAME" ]]; then
TYPE="directory"; FULL_PATH="$SOURCE_DIR/$NAME"
else
TYPE="unknown"; FULL_PATH=""
fi
echo "Type: $TYPE"
# For directory backups, check key items
if [[ "$TYPE" == "directory" ]]; then
echo "Checking key backup items..."
KEYS=("project.db" "voices" "speakers.json" ".env" "config.py")
ERRORS=0
for k in "${KEYS[@]}"; do
if [[ -e "$FULL_PATH/$k" ]] || [[ -e "$FULL_PATH/$k."* ]]; then
echo " ✓ $k"
else
echo " ⚠ $k missing"
ERRORS=$((ERRORS+1))
fi
done
echo ""
if [[ $ERRORS -eq 0 ]]; then
echo "✅ All key items present"
else
echo "⚠ $ERRORS item(s) missing"
fi
fi
echo ""
echo "========================================="
if [[ "$TYPE" == "directory" && $ERRORS -eq 0 ]]; then
echo "✅ VERIFICATION PASSED"
else
echo "⚠ VERIFICATION COMPLETED"
fi
echo "========================================="

View File

@@ -0,0 +1,33 @@
"""UmbraTrans cluster integration for UmbraAudio.
Provides PostgreSQL connection, compute node registration, job queue,
artifact protocol, and NFS shared storage access.
"""
from __future__ import annotations
from .artifacts import ArtifactStore, read_artifact, write_artifact
from .config import ClusterConfig, load_cluster_config
from .database import ClusterDB, get_cluster_db
from .jobs import JobKind, JobQueue, claim_job, complete_job
from .nfs import NFSMounts, ensure_nfs_mounts
from .node import ComputeNode, heartbeat_loop, register_node
__all__ = [
"ClusterConfig",
"load_cluster_config",
"ClusterDB",
"get_cluster_db",
"ComputeNode",
"register_node",
"heartbeat_loop",
"JobQueue",
"JobKind",
"claim_job",
"complete_job",
"ArtifactStore",
"write_artifact",
"read_artifact",
"NFSMounts",
"ensure_nfs_mounts",
]

View File

@@ -0,0 +1,240 @@
"""Artifact protocol (content-addressed storage) for UmbraAudio ↔ UmbraTrans."""
from __future__ import annotations
import hashlib
import logging
import os
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, BinaryIO
if TYPE_CHECKING:
from .config import ClusterConfig
from .database import ClusterDB
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class ArtifactRef:
"""Reference to a written artifact."""
artifact_id: uuid.UUID
content_hash: str
byte_size: int
storage_uri: str
local_path: Path
class ArtifactStore:
"""Content-addressed artifact storage on NFS."""
def __init__(self, config: ClusterConfig, db: ClusterDB) -> None:
self.config = config
self.db = db
def _artifact_path(self, org_id: uuid.UUID, project_id: uuid.UUID, content_hash: str) -> Path:
"""Compute artifact path on NFS."""
# Structure: /mnt/umbra/artifacts/{org_id}/{project_id}/{hash[:2]}/{hash}
artifacts_root = self.config.local_mount_root / "artifacts"
return artifacts_root / str(org_id) / str(project_id) / content_hash[:2] / content_hash
async def write(
self,
data: bytes,
project_id: uuid.UUID,
organization_id: uuid.UUID,
lang: str,
kind: str,
media_type: str,
derivation: dict[str, Any] | None = None,
) -> ArtifactRef:
"""Write artifact atomically to NFS and register in database."""
# 1. Compute content hash
content_hash = hashlib.sha256(data).hexdigest()
byte_size = len(data)
# 2. Write to NFS (atomic: tmp → rename)
artifact_path = self._artifact_path(organization_id, project_id, content_hash)
artifact_path.parent.mkdir(parents=True, exist_ok=True)
tmp_path = artifact_path.with_suffix(artifact_path.suffix + ".tmp")
try:
with open(tmp_path, "wb") as f:
f.write(data)
os.rename(tmp_path, artifact_path)
except Exception:
if tmp_path.exists():
tmp_path.unlink(missing_ok=True)
raise
# 3. Insert into database
storage_key = content_hash
storage_uri = f"file://{self.config.node_id}/{artifact_path.relative_to(self.config.local_mount_root)}"
record = await self.db.insert_artifact(
project_id=project_id,
organization_id=organization_id,
lang=lang,
kind=kind,
content_hash=content_hash,
byte_size=byte_size,
media_type=media_type,
storage_key=storage_key,
storage_uri=storage_uri,
derivation=derivation,
)
logger.info(
"Wrote artifact %s (%s, %d bytes) to %s",
record.id, kind, byte_size, artifact_path
)
return ArtifactRef(
artifact_id=record.id,
content_hash=content_hash,
byte_size=byte_size,
storage_uri=storage_uri,
local_path=artifact_path,
)
async def write_stream(
self,
stream: BinaryIO,
project_id: uuid.UUID,
organization_id: uuid.UUID,
lang: str,
kind: str,
media_type: str,
derivation: dict[str, Any] | None = None,
chunk_size: int = 8192,
) -> ArtifactRef:
"""Write artifact from a stream (for large files)."""
# For streaming, we need to compute hash while writing
hasher = hashlib.sha256()
temp_path = self.config.local_mount_root / "artifacts" / ".tmp" / f"{uuid.uuid4()}.tmp"
temp_path.parent.mkdir(parents=True, exist_ok=True)
try:
byte_size = 0
with open(temp_path, "wb") as f:
while chunk := stream.read(chunk_size):
hasher.update(chunk)
f.write(chunk)
byte_size += len(chunk)
content_hash = hasher.hexdigest()
final_path = self._artifact_path(organization_id, project_id, content_hash)
final_path.parent.mkdir(parents=True, exist_ok=True)
# Atomic move
os.rename(temp_path, final_path)
storage_key = content_hash
storage_uri = f"file://{self.config.node_id}/{final_path.relative_to(self.config.local_mount_root)}"
record = await self.db.insert_artifact(
project_id=project_id,
organization_id=organization_id,
lang=lang,
kind=kind,
content_hash=content_hash,
byte_size=byte_size,
media_type=media_type,
storage_key=storage_key,
storage_uri=storage_uri,
derivation=derivation,
)
return ArtifactRef(
artifact_id=record.id,
content_hash=content_hash,
byte_size=byte_size,
storage_uri=storage_uri,
local_path=final_path,
)
finally:
if temp_path.exists():
temp_path.unlink(missing_ok=True)
async def read(self, artifact_id: uuid.UUID) -> bytes | None:
"""Read artifact by ID."""
record = await self.db.get_artifact(artifact_id)
if not record:
return None
# Parse storage_uri to get local path
# file://node-id/mnt/umbra/artifacts/...
if not record.storage_uri.startswith("file://"):
logger.error("Unsupported storage URI scheme: %s", record.storage_uri)
return None
# Extract path after hostname
path_str = record.storage_uri.split("/", 3)[-1] # "/mnt/umbra/artifacts/..."
local_path = Path(path_str)
if not local_path.exists():
logger.error("Artifact file not found: %s", local_path)
return None
with open(local_path, "rb") as f:
data = f.read()
# Verify hash
if hashlib.sha256(data).hexdigest() != record.content_hash:
logger.error("Artifact hash mismatch: %s", artifact_id)
return None
return data
async def read_stream(self, artifact_id: uuid.UUID, chunk_size: int = 8192):
"""Read artifact as stream."""
record = await self.db.get_artifact(artifact_id)
if not record:
return
if not record.storage_uri.startswith("file://"):
return
path_str = record.storage_uri.split("/", 3)[-1]
local_path = Path(path_str)
if not local_path.exists():
return
with open(local_path, "rb") as f:
while chunk := f.read(chunk_size):
yield chunk
def local_path(self, artifact_id: uuid.UUID) -> Path | None:
"""Get local path for artifact (without reading)."""
# This would need a DB lookup; for now return None
# Could be cached
return None
async def write_artifact(
config: ClusterConfig,
db: ClusterDB,
data: bytes,
project_id: uuid.UUID,
organization_id: uuid.UUID,
lang: str,
kind: str,
media_type: str,
derivation: dict[str, Any] | None = None,
) -> ArtifactRef:
"""Write artifact (convenience function)."""
store = ArtifactStore(config, db)
return await store.write(data, project_id, organization_id, lang, kind, media_type, derivation)
async def read_artifact(
config: ClusterConfig,
db: ClusterDB,
artifact_id: uuid.UUID,
) -> bytes | None:
"""Read artifact (convenience function)."""
store = ArtifactStore(config, db)
return await store.read(artifact_id)

View File

@@ -0,0 +1,177 @@
"""Cluster configuration for UmbraAudio ↔ UmbraTrans integration."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Literal
from pydantic_settings import BaseSettings, SettingsConfigDict
@dataclass(frozen=True)
class ClusterConfig:
"""Runtime configuration for cluster integration."""
# PostgreSQL (PgBouncer pooled)
pg_host: str = "192.168.123.14"
pg_port: int = 6432
pg_database: str = "umbratrans"
pg_user: str = "umbraudio"
pg_password: str = ""
pg_sslmode: Literal[
"disable", "allow", "prefer", "require", "verify-ca", "verify-full"
] = "require"
pg_app_name: str = "umbraudio-worker"
# Direct PostgreSQL (admin/backup)
pg_direct_port: int = 5432
# NFS
nfs_server: str = "192.168.123.14"
nfs_productions_path: str = "/mnt/umbra-data/productions"
nfs_artifacts_path: str = "/mnt/umbra-data/artifacts"
nfs_covers_path: str = "/mnt/umbra-data/covers"
nfs_audio_path: str = "/mnt/umbra-data/audio"
# Local mount points
local_mount_root: Path = Path("/mnt/umbra")
# Compute node
node_id: str = "umbra-audio-01"
organization_id: str = "umbra"
node_name: str = ""
datacenter: str = "local"
agent_version: str = "1.0.0"
# Credentials
node_credential_prefix: str = "umbra-audio-01-"
node_scopes: tuple[str, ...] = ("queue:claim", "queue:complete", "artifacts:write")
# Job polling
poll_interval_seconds: int = 10
lease_duration_minutes: int = 30
allowed_job_kinds: tuple[str, ...] = (
"tts", "asr", "assemble", "encode", "master", "speaker_profile",
)
# Heartbeat
heartbeat_interval_seconds: int = 30
heartbeat_timeout_minutes: int = 2
@property
def pg_dsn(self) -> str:
"""PostgreSQL DSN for psycopg/asyncpg."""
return (
f"postgresql://{self.pg_user}:{self.pg_password}@{self.pg_host}:{self.pg_port}/"
f"{self.pg_database}?sslmode={self.pg_sslmode}&application_name={self.pg_app_name}"
)
@property
def pg_direct_dsn(self) -> str:
"""Direct PostgreSQL DSN (port 5432)."""
return (
f"postgresql://{self.pg_user}:{self.pg_password}@{self.pg_host}:{self.pg_direct_port}/"
f"{self.pg_database}?sslmode={self.pg_sslmode}&application_name={self.pg_app_name}-direct"
)
@property
def nfs_mounts(self) -> dict[str, tuple[str, Path]]:
"""NFS exports → local mount points."""
server = self.nfs_server
root = self.local_mount_root
return {
"productions": (
f"{server}:{self.nfs_productions_path}",
root / "productions",
),
"artifacts": (
f"{server}:{self.nfs_artifacts_path}",
root / "artifacts",
),
"covers": (
f"{server}:{self.nfs_covers_path}",
root / "covers",
),
"audio": (
f"{server}:{self.nfs_audio_path}",
root / "audio",
),
}
class ClusterSettings(BaseSettings):
"""Settings loaded from environment (.env)."""
model_config = SettingsConfigDict(env_prefix="UMBRA_CLUSTER__", env_file=".env", extra="ignore")
pg_host: str = "192.168.123.14"
pg_port: int = 6432
pg_database: str = "umbratrans"
pg_user: str = "umbraudio"
pg_password: str = ""
pg_sslmode: str = "require"
pg_direct_port: int = 5432
nfs_server: str = "192.168.123.14"
nfs_productions_path: str = "/mnt/umbra-data/productions"
nfs_artifacts_path: str = "/mnt/umbra-data/artifacts"
nfs_covers_path: str = "/mnt/umbra-data/covers"
nfs_audio_path: str = "/mnt/umbra-data/audio"
local_mount_root: Path = Path("/mnt/umbra")
node_id: str = "umbra-audio-01"
organization_id: str = "umbra"
datacenter: str = "local"
agent_version: str = "1.0.0"
node_credential_prefix: str = "umbra-audio-01-"
node_scopes: str = "queue:claim,queue:complete,artifacts:write"
poll_interval_seconds: int = 10
lease_duration_minutes: int = 30
allowed_job_kinds: str = "tts,asr,assemble,encode,master,speaker_profile"
heartbeat_interval_seconds: int = 30
heartbeat_timeout_minutes: int = 2
def load_cluster_config(project_dir: Path | None = None) -> ClusterConfig:
"""Load cluster configuration from environment and defaults."""
settings = ClusterSettings()
# Parse comma-separated strings
scopes = tuple(s.strip() for s in settings.node_scopes.split(",") if s.strip())
job_kinds = tuple(k.strip() for k in settings.allowed_job_kinds.split(",") if k.strip())
# Node name defaults to hostname if not set
node_name = settings.node_id # Could be overridden
return ClusterConfig(
pg_host=settings.pg_host,
pg_port=settings.pg_port,
pg_database=settings.pg_database,
pg_user=settings.pg_user,
pg_password=settings.pg_password,
pg_sslmode=settings.pg_sslmode, # type: ignore[assignment]
pg_direct_port=settings.pg_direct_port,
nfs_server=settings.nfs_server,
nfs_productions_path=settings.nfs_productions_path,
nfs_artifacts_path=settings.nfs_artifacts_path,
nfs_covers_path=settings.nfs_covers_path,
nfs_audio_path=settings.nfs_audio_path,
local_mount_root=settings.local_mount_root,
node_id=settings.node_id,
organization_id=settings.organization_id,
node_name=node_name,
datacenter=settings.datacenter,
agent_version=settings.agent_version,
node_credential_prefix=settings.node_credential_prefix,
node_scopes=scopes,
poll_interval_seconds=settings.poll_interval_seconds,
lease_duration_minutes=settings.lease_duration_minutes,
allowed_job_kinds=job_kinds,
heartbeat_interval_seconds=settings.heartbeat_interval_seconds,
heartbeat_timeout_minutes=settings.heartbeat_timeout_minutes,
)

View File

@@ -0,0 +1,581 @@
"""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

177
umbraaudio/cluster/jobs.py Normal file
View File

@@ -0,0 +1,177 @@
"""Job queue operations for UmbraAudio worker."""
from __future__ import annotations
import asyncio
import contextlib
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from .database import ClusterDB, JobRecord, get_cluster_db
if TYPE_CHECKING:
import uuid
from .config import ClusterConfig
logger = logging.getLogger(__name__)
class JobKind:
"""Job kinds that UmbraAudio can claim."""
TTS = "tts"
ASR = "asr"
ASSEMBLE = "assemble"
ENCODE = "encode"
MASTER = "master"
SPEAKER_PROFILE = "speaker_profile"
@classmethod
def all(cls) -> list[str]:
return [cls.TTS, cls.ASR, cls.ASSEMBLE, cls.ENCODE, cls.MASTER, cls.SPEAKER_PROFILE]
@classmethod
def tts_related(cls) -> list[str]:
return [cls.TTS, cls.ASSEMBLE, cls.ENCODE, cls.MASTER]
@dataclass
class JobContext:
"""Context for job processing."""
job: JobRecord
db: ClusterDB
config: ClusterConfig
class JobQueue:
"""Job queue poller and processor."""
def __init__(
self,
config: ClusterConfig,
db: ClusterDB | None = None,
poll_interval: int | None = None,
) -> None:
self.config = config
self.db = db
self.poll_interval = poll_interval or config.poll_interval_seconds
self._running = False
self._task: asyncio.Task | None = None
self._handlers: dict[str, JobHandler] = {}
def register_handler(self, kind: str, handler: JobHandler) -> None:
"""Register a handler for a job kind."""
self._handlers[kind] = handler
logger.info("Registered handler for job kind: %s", kind)
async def start(self) -> None:
"""Start polling for jobs."""
if self.db is None:
self.db = await get_cluster_db(self.config)
self._running = True
self._task = asyncio.create_task(self._poll_loop())
logger.info("Job queue started (poll interval: %ds)", self.poll_interval)
async def stop(self) -> None:
"""Stop polling."""
self._running = False
if self._task:
self._task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await self._task
logger.info("Job queue stopped")
async def _poll_loop(self) -> None:
"""Main polling loop."""
while self._running:
try:
job = await self.db.claim_job(
node_id=self.config.node_id,
allowed_kinds=list(self._handlers.keys()) or self.config.allowed_job_kinds,
organization_id=self.config.organization_id,
)
if job:
await self._process_job(job)
else:
await asyncio.sleep(self.poll_interval)
except asyncio.CancelledError:
break
except Exception as e:
logger.error("Poll loop error: %s", e)
await asyncio.sleep(self.poll_interval)
async def _process_job(self, job: JobRecord) -> None:
"""Process a single job."""
assert self.db is not None
handler = self._handlers.get(job.kind)
if not handler:
logger.warning("No handler for job kind %s, requeuing", job.kind)
await self.db.requeue_job(
job.id, self.config.node_id, f"No handler for kind {job.kind}"
)
return
logger.info("Processing job %s (kind: %s)", job.id, job.kind)
ctx = JobContext(job=job, db=self.db, config=self.config)
try:
result = await handler(ctx)
await self.db.complete_job(
job_id=job.id,
node_id=self.config.node_id,
status="complete",
result=result,
artifacts=result.get("artifacts") if isinstance(result, dict) else None,
)
logger.info("Job %s completed", job.id)
except Exception as e:
logger.exception("Job %s failed: %s", job.id, e)
await self.db.complete_job(
job_id=job.id,
node_id=self.config.node_id,
status="failed",
error=str(e),
)
class JobHandler:
"""Base class for job handlers."""
async def __call__(self, ctx: JobContext) -> dict[str, Any]:
raise NotImplementedError
# Convenience functions for direct use
async def claim_job(
config: ClusterConfig,
db: ClusterDB,
allowed_kinds: list[str] | None = None,
) -> JobRecord | None:
"""Claim a single job."""
return await db.claim_job(
node_id=config.node_id,
allowed_kinds=allowed_kinds or config.allowed_job_kinds,
organization_id=config.organization_id,
)
async def complete_job(
config: ClusterConfig,
db: ClusterDB,
job_id: uuid.UUID,
status: str,
result: dict[str, Any] | None = None,
error: str | None = None,
artifacts: list[dict[str, Any]] | None = None,
) -> bool:
"""Complete a job."""
return await db.complete_job(
job_id=job_id,
node_id=config.node_id,
status=status,
result=result,
error=error,
artifacts=artifacts,
)

140
umbraaudio/cluster/nfs.py Normal file
View File

@@ -0,0 +1,140 @@
"""NFS mount management for UmbraAudio cluster integration."""
from __future__ import annotations
import logging
import subprocess
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from pathlib import Path
from .config import ClusterConfig
logger = logging.getLogger(__name__)
class NFSMounts:
"""Manages NFS mounts for cluster shared storage."""
def __init__(self, config: ClusterConfig) -> None:
self.config = config
def is_mounted(self, mount_point: Path) -> bool:
"""Check if a path is mounted."""
try:
result = subprocess.run(
["findmnt", "-n", "-o", "TARGET", str(mount_point)],
capture_output=True,
text=True,
timeout=5,
)
return result.returncode == 0 and result.stdout.strip() == str(mount_point)
except Exception:
return False
def mount(self, export: str, mount_point: Path, options: str | None = None) -> bool:
"""Mount an NFS export."""
mount_point.mkdir(parents=True, exist_ok=True)
if self.is_mounted(mount_point):
logger.debug("Already mounted: %s", mount_point)
return True
opts = options or "rw,hard,intr,rsize=1048576,wsize=1048576,timeo=600,retrans=2,vers=4.2"
cmd = ["mount", "-t", "nfs4", "-o", opts, export, str(mount_point)]
try:
logger.info("Mounting %s → %s", export, mount_point)
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
if result.returncode == 0:
logger.info("Mounted: %s", mount_point)
return True
else:
logger.error("Mount failed: %s", result.stderr)
return False
except subprocess.TimeoutExpired:
logger.error("Mount timed out: %s", mount_point)
return False
except Exception as e:
logger.error("Mount error: %s", e)
return False
def unmount(self, mount_point: Path) -> bool:
"""Unmount an NFS mount."""
if not self.is_mounted(mount_point):
return True
try:
result = subprocess.run(
["umount", str(mount_point)],
capture_output=True,
text=True,
timeout=10,
)
if result.returncode == 0:
logger.info("Unmounted: %s", mount_point)
return True
else:
logger.error("Unmount failed: %s", result.stderr)
return False
except Exception as e:
logger.error("Unmount error: %s", e)
return False
def ensure_all(self) -> dict[str, bool]:
"""Ensure all configured NFS mounts are present."""
results = {}
for name, (export, mount_point) in self.config.nfs_mounts.items():
results[name] = self.mount(export, mount_point)
return results
def ensure_nfs_mounts(config: ClusterConfig) -> dict[str, bool]:
"""Ensure all NFS mounts are present (convenience function)."""
mounts = NFSMounts(config)
return mounts.ensure_all()
def check_nfs_server(config: ClusterConfig) -> bool:
"""Check if NFS server is reachable and exporting expected paths."""
try:
# Check rpcbind
result = subprocess.run(
["rpcinfo", "-p", config.nfs_server],
capture_output=True,
text=True,
timeout=5,
)
if result.returncode != 0:
logger.warning("rpcinfo failed: %s", result.stderr)
return False
# Check for NFS exports
result = subprocess.run(
["showmount", "-e", config.nfs_server],
capture_output=True,
text=True,
timeout=5,
)
if result.returncode != 0:
logger.warning("showmount failed: %s", result.stderr)
return False
exports = result.stdout
required = [
config.nfs_productions_path,
config.nfs_artifacts_path,
config.nfs_covers_path,
config.nfs_audio_path,
]
missing = [p for p in required if p not in exports]
if missing:
logger.warning("Missing NFS exports: %s", missing)
return False
logger.info("NFS server %s exports verified", config.nfs_server)
return True
except Exception as e:
logger.error("NFS server check failed: %s", e)
return False

192
umbraaudio/cluster/node.py Normal file
View File

@@ -0,0 +1,192 @@
"""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