26de44feb1
The UmbraAudio side of the cross-machine contract: this node registers with the UmbraTrans controller, polls for jobs, claims them under a lease, and writes artifacts back over NFS. * `umbraaudio/cluster/` - configuration, PostgreSQL access, node registration and heartbeat, the job queue, artifact storage and NFS mount management. * `scripts/cluster_worker.py` - the CLI: `worker`, `check-cluster`, `register-node`. * `scripts/backup_system.sh` / `scripts/verify_backup.sh` - whole-system backup and integrity verification. * `.env` - the `UMBRA_CLUSTER__*` settings. `UMBRA_CLUSTER__PG_PASSWORD` is deliberately empty; the password comes from the environment, not from a committed file. `k3s-agent.yaml` is NOT committed: it contains a K3s server join token. It is now gitignored, and the token should be rotated if that file has ever left this machine. ruff clean on the whole package.
241 lines
7.6 KiB
Python
241 lines
7.6 KiB
Python
"""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)
|