From 26de44feb1f562310877a04782d76320f0df449a Mon Sep 17 00:00:00 2001 From: "B.D. Murphy" <161863146+bdmurphy73@users.noreply.github.com> Date: Fri, 6 Nov 2026 17:03:16 -0800 Subject: [PATCH] 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. --- .env | 32 ++ .gitignore | 6 + pyproject.toml | 2 + scripts/backup_system.sh | 134 ++++++++ scripts/cluster_worker.py | 160 +++++++++ scripts/verify_backup.sh | 115 +++++++ umbraaudio/cluster/__init__.py | 33 ++ umbraaudio/cluster/artifacts.py | 240 +++++++++++++ umbraaudio/cluster/config.py | 177 ++++++++++ umbraaudio/cluster/database.py | 581 ++++++++++++++++++++++++++++++++ umbraaudio/cluster/jobs.py | 177 ++++++++++ umbraaudio/cluster/nfs.py | 140 ++++++++ umbraaudio/cluster/node.py | 192 +++++++++++ 13 files changed, 1989 insertions(+) create mode 100755 scripts/backup_system.sh create mode 100644 scripts/cluster_worker.py create mode 100755 scripts/verify_backup.sh create mode 100644 umbraaudio/cluster/__init__.py create mode 100644 umbraaudio/cluster/artifacts.py create mode 100644 umbraaudio/cluster/config.py create mode 100644 umbraaudio/cluster/database.py create mode 100644 umbraaudio/cluster/jobs.py create mode 100644 umbraaudio/cluster/nfs.py create mode 100644 umbraaudio/cluster/node.py diff --git a/.env b/.env index 8ad0ab4..46e1881 100644 --- a/.env +++ b/.env @@ -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 diff --git a/.gitignore b/.gitignore index 67a380f..7f3803e 100644 --- a/.gitignore +++ b/.gitignore @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 5be0a9e..75b973a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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] diff --git a/scripts/backup_system.sh b/scripts/backup_system.sh new file mode 100755 index 0000000..b6fe641 --- /dev/null +++ b/scripts/backup_system.sh @@ -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 diff --git a/scripts/cluster_worker.py b/scripts/cluster_worker.py new file mode 100644 index 0000000..e25f5aa --- /dev/null +++ b/scripts/cluster_worker.py @@ -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() diff --git a/scripts/verify_backup.sh b/scripts/verify_backup.sh new file mode 100755 index 0000000..26aae08 --- /dev/null +++ b/scripts/verify_backup.sh @@ -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 "=========================================" diff --git a/umbraaudio/cluster/__init__.py b/umbraaudio/cluster/__init__.py new file mode 100644 index 0000000..663eede --- /dev/null +++ b/umbraaudio/cluster/__init__.py @@ -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", +] diff --git a/umbraaudio/cluster/artifacts.py b/umbraaudio/cluster/artifacts.py new file mode 100644 index 0000000..7178b10 --- /dev/null +++ b/umbraaudio/cluster/artifacts.py @@ -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) diff --git a/umbraaudio/cluster/config.py b/umbraaudio/cluster/config.py new file mode 100644 index 0000000..996ccf2 --- /dev/null +++ b/umbraaudio/cluster/config.py @@ -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, + ) diff --git a/umbraaudio/cluster/database.py b/umbraaudio/cluster/database.py new file mode 100644 index 0000000..c1daad5 --- /dev/null +++ b/umbraaudio/cluster/database.py @@ -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 diff --git a/umbraaudio/cluster/jobs.py b/umbraaudio/cluster/jobs.py new file mode 100644 index 0000000..55d801e --- /dev/null +++ b/umbraaudio/cluster/jobs.py @@ -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, + ) diff --git a/umbraaudio/cluster/nfs.py b/umbraaudio/cluster/nfs.py new file mode 100644 index 0000000..b5ec4bb --- /dev/null +++ b/umbraaudio/cluster/nfs.py @@ -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 diff --git a/umbraaudio/cluster/node.py b/umbraaudio/cluster/node.py new file mode 100644 index 0000000..fdb0eca --- /dev/null +++ b/umbraaudio/cluster/node.py @@ -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