1097 lines
42 KiB
Python
1097 lines
42 KiB
Python
"""Crash-safe, resumable DWH catalog and LSH preprocessing stages."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import atexit
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import stat
|
|
import tempfile
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from tht.jobs.models import JobReport, JobSpec
|
|
from tht.jobs.runner import (
|
|
CorruptCheckpointError,
|
|
JobContext,
|
|
StageArtifacts,
|
|
run_job,
|
|
seal_stage_artifacts,
|
|
)
|
|
|
|
DWH_STAGE_IDS = ("introspect", "lsh")
|
|
_RUN_ID = re.compile(r"^[0-9a-f]{32}$")
|
|
_SAFE_FILE = re.compile(r"^[A-Za-z0-9_-]+\.(?:pkl|json)$")
|
|
GENERATION_MANIFEST = "generation-manifest.json"
|
|
OWNER_MARKER = "OWNER.json"
|
|
_SNAPSHOT_DIRS: set[Path] = set()
|
|
|
|
|
|
def _cleanup_snapshot_dirs() -> None:
|
|
for path in tuple(_SNAPSHOT_DIRS):
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
_SNAPSHOT_DIRS.discard(path)
|
|
|
|
|
|
atexit.register(_cleanup_snapshot_dirs)
|
|
|
|
|
|
def config_dwh_binding(cfg) -> dict[str, str]:
|
|
workspace_id = getattr(cfg, "_workspace_id", None)
|
|
config_source = getattr(cfg, "_config_source", None)
|
|
if not isinstance(workspace_id, str) or not isinstance(config_source, str):
|
|
raise CorruptCheckpointError("DWH workspace identity is unavailable; reload configuration")
|
|
try:
|
|
# P3: the versioned canonical effective configuration. Only DWH-affecting fields are
|
|
# included, so content-only/Evidence-only changes reuse the generation; a changed
|
|
# endpoint/transport/database/schema/identity fails closed via the OWNER.json compare.
|
|
from tht.config import (
|
|
effective_config_fingerprint,
|
|
effective_config_input_fingerprint,
|
|
)
|
|
config_fingerprint = effective_config_fingerprint(cfg)
|
|
input_fingerprint = effective_config_input_fingerprint(workspace_id, cfg)
|
|
except Exception: # noqa: BLE001
|
|
# Lightweight test doubles predating the canonical form retain the legacy seam: the
|
|
# model dump minus session persistence and runtime identity (unchanged behavior).
|
|
model_dump = getattr(cfg, "model_dump", None)
|
|
if callable(model_dump):
|
|
payload = model_dump(mode="json")
|
|
if not isinstance(payload, dict):
|
|
raise CorruptCheckpointError("DWH workspace configuration is unavailable; reload configuration")
|
|
payload.pop("session_storage", None)
|
|
payload.pop("runtime_identity", None)
|
|
config_fingerprint = fingerprint(json.dumps(payload, separators=(",", ":"), ensure_ascii=False))
|
|
input_fingerprint = fingerprint(config_source)
|
|
else:
|
|
config_fingerprint = fingerprint(cfg.model_dump_json())
|
|
input_fingerprint = fingerprint(config_source)
|
|
return {
|
|
"workspace_id": workspace_id,
|
|
"config_fingerprint": config_fingerprint,
|
|
"input_fingerprint": input_fingerprint,
|
|
}
|
|
|
|
|
|
def _binding_digest(binding: dict[str, str]) -> str:
|
|
payload = json.dumps(binding, sort_keys=True, separators=(",", ":"))
|
|
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _read_root_binding_fd(root_fd: int) -> dict[str, str]:
|
|
try:
|
|
fd = os.open(OWNER_MARKER, os.O_RDONLY | os.O_NOFOLLOW, dir_fd=root_fd)
|
|
try:
|
|
info = os.fstat(fd)
|
|
if (
|
|
not stat.S_ISREG(info.st_mode)
|
|
or info.st_uid != os.getuid()
|
|
or info.st_nlink != 1
|
|
or stat.S_IMODE(info.st_mode) != 0o400
|
|
):
|
|
raise OSError("unsafe DWH ownership marker")
|
|
chunks = []
|
|
while chunk := os.read(fd, 1024 * 1024):
|
|
chunks.append(chunk)
|
|
finally:
|
|
os.close(fd)
|
|
payload = json.loads(b"".join(chunks).decode("utf-8"))
|
|
binding = payload["binding"]
|
|
if (
|
|
payload.get("schema_version") != 1
|
|
or not isinstance(binding, dict)
|
|
or set(binding) != {"workspace_id", "config_fingerprint", "input_fingerprint"}
|
|
or payload.get("binding_sha256") != _binding_digest(binding)
|
|
):
|
|
raise ValueError
|
|
return binding
|
|
except (OSError, KeyError, TypeError, ValueError, UnicodeDecodeError) as error:
|
|
raise CorruptCheckpointError("DWH workspace ownership marker is missing or invalid") from error
|
|
|
|
|
|
def _validate_root_binding_fd(root_fd: int, expected: dict[str, str]) -> None:
|
|
if _read_root_binding_fd(root_fd) != expected:
|
|
raise CorruptCheckpointError("DWH artifacts belong to a different workspace configuration")
|
|
entries = set(os.listdir(root_fd))
|
|
active_exists = "ACTIVE" in entries
|
|
try:
|
|
generations_fd = os.open(
|
|
"generations", os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=root_fd
|
|
)
|
|
except FileNotFoundError:
|
|
generation_entries = []
|
|
except OSError as error:
|
|
raise CorruptCheckpointError("DWH generations directory is invalid") from error
|
|
else:
|
|
try:
|
|
info = os.fstat(generations_fd)
|
|
if info.st_uid != os.getuid() or stat.S_IMODE(info.st_mode) != 0o700:
|
|
raise OSError("unsafe DWH generations directory")
|
|
generation_entries = os.listdir(generations_fd)
|
|
finally:
|
|
os.close(generations_fd)
|
|
if generation_entries and not active_exists:
|
|
raise CorruptCheckpointError("DWH generations exist without a consistent ACTIVE pointer")
|
|
|
|
|
|
def _claim_or_validate_root_binding(
|
|
root_fd: int, binding: dict[str, str]
|
|
) -> None:
|
|
temporary: str | None = None
|
|
try:
|
|
entries = set(os.listdir(root_fd))
|
|
if OWNER_MARKER in entries:
|
|
_validate_root_binding_fd(root_fd, binding)
|
|
return
|
|
if not entries <= {"generation.lock", "generations"} or "generation.lock" not in entries:
|
|
raise CorruptCheckpointError(
|
|
"DWH artifacts are unbound; migrate them explicitly or use an empty root"
|
|
)
|
|
if "generations" in entries:
|
|
generations_fd = os.open(
|
|
"generations", os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=root_fd
|
|
)
|
|
try:
|
|
info = os.fstat(generations_fd)
|
|
if info.st_uid != os.getuid() or stat.S_IMODE(info.st_mode) != 0o700:
|
|
raise OSError("unsafe DWH generations directory")
|
|
if os.listdir(generations_fd):
|
|
raise CorruptCheckpointError(
|
|
"DWH artifacts are unbound; migrate them explicitly or use an empty root"
|
|
)
|
|
finally:
|
|
os.close(generations_fd)
|
|
payload = {
|
|
"schema_version": 1,
|
|
"binding": binding,
|
|
"binding_sha256": _binding_digest(binding),
|
|
}
|
|
temporary = f".{OWNER_MARKER}.{uuid.uuid4().hex}.tmp"
|
|
fd = os.open(
|
|
temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW,
|
|
0o600, dir_fd=root_fd,
|
|
)
|
|
try:
|
|
data = (json.dumps(payload, sort_keys=True, separators=(",", ":")) + "\n").encode()
|
|
offset = 0
|
|
while offset < len(data):
|
|
offset += os.write(fd, data[offset:])
|
|
os.fsync(fd)
|
|
os.fchmod(fd, 0o400)
|
|
os.fsync(fd)
|
|
finally:
|
|
os.close(fd)
|
|
os.replace(temporary, OWNER_MARKER, src_dir_fd=root_fd, dst_dir_fd=root_fd)
|
|
temporary = None
|
|
os.fsync(root_fd)
|
|
except OSError as error:
|
|
raise CorruptCheckpointError("DWH workspace root is invalid") from error
|
|
finally:
|
|
if temporary is not None:
|
|
try:
|
|
os.unlink(temporary, dir_fd=root_fd)
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DwhArtifactSnapshot:
|
|
generation: str | None
|
|
physical: Path
|
|
lsh_dir: Path
|
|
_holder: object | None = None
|
|
|
|
|
|
@dataclass
|
|
class _GenerationLease:
|
|
root_fd: int
|
|
lock_fd: int
|
|
|
|
def assert_root_identity(self, root: Path) -> None:
|
|
try:
|
|
path_info = os.stat(root, follow_symlinks=False)
|
|
opened_info = os.fstat(self.root_fd)
|
|
except OSError as error:
|
|
raise CorruptCheckpointError("DWH workspace root changed while locked") from error
|
|
if (
|
|
not stat.S_ISDIR(path_info.st_mode)
|
|
or (path_info.st_dev, path_info.st_ino)
|
|
!= (opened_info.st_dev, opened_info.st_ino)
|
|
):
|
|
raise CorruptCheckpointError("DWH workspace root changed while locked")
|
|
|
|
def close(self) -> None:
|
|
try:
|
|
fcntl.flock(self.lock_fd, fcntl.LOCK_UN)
|
|
finally:
|
|
try:
|
|
os.close(self.lock_fd)
|
|
finally:
|
|
os.close(self.root_fd)
|
|
|
|
|
|
class DwhSnapshotLease:
|
|
def __init__(self, cfg) -> None:
|
|
self.cfg = cfg
|
|
self._lease: _GenerationLease | None = None
|
|
self.snapshot: DwhArtifactSnapshot | None = None
|
|
|
|
def __enter__(self) -> DwhArtifactSnapshot:
|
|
binding = config_dwh_binding(self.cfg)
|
|
self._lease = _acquire_existing_generation_lock(
|
|
self.cfg.paths.artifacts.parent, exclusive=False
|
|
)
|
|
try:
|
|
self.snapshot = _resolve_dwh_snapshot_locked(
|
|
self.cfg, binding, self._lease.root_fd
|
|
)
|
|
return self.snapshot
|
|
except BaseException:
|
|
self.__exit__(None, None, None)
|
|
raise
|
|
|
|
def __exit__(self, *_args) -> None:
|
|
if self._lease is not None:
|
|
lease, self._lease = self._lease, None
|
|
lease.close()
|
|
if self.snapshot is not None and isinstance(self.snapshot._holder, Path):
|
|
shutil.rmtree(self.snapshot._holder, ignore_errors=True)
|
|
_SNAPSHOT_DIRS.discard(self.snapshot._holder)
|
|
|
|
|
|
def lease_dwh_snapshot(cfg) -> DwhSnapshotLease:
|
|
return DwhSnapshotLease(cfg)
|
|
|
|
|
|
def _acquire_generation_lock(workspace_root: Path, *, exclusive: bool) -> _GenerationLease:
|
|
root = workspace_root / ".tht-dwh"
|
|
DwhPreprocessPipeline._ensure_owned_dir(root)
|
|
root_fd = os.open(root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
|
|
try:
|
|
fd = os.open(
|
|
"generation.lock", os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW,
|
|
0o600, dir_fd=root_fd,
|
|
)
|
|
except BaseException:
|
|
os.close(root_fd)
|
|
raise
|
|
try:
|
|
info = os.fstat(fd)
|
|
if (
|
|
not stat.S_ISREG(info.st_mode)
|
|
or info.st_uid != os.getuid()
|
|
or info.st_nlink != 1
|
|
or stat.S_IMODE(info.st_mode) != 0o600
|
|
):
|
|
raise OSError("unsafe DWH generation lock")
|
|
fcntl.flock(fd, fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH)
|
|
return _GenerationLease(root_fd, fd)
|
|
except BaseException:
|
|
os.close(fd)
|
|
os.close(root_fd)
|
|
raise
|
|
|
|
|
|
def _acquire_existing_generation_lock(
|
|
workspace_root: Path, *, exclusive: bool
|
|
) -> _GenerationLease:
|
|
root = workspace_root / ".tht-dwh"
|
|
try:
|
|
root_fd = os.open(root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
|
|
try:
|
|
root_info = os.fstat(root_fd)
|
|
if root_info.st_uid != os.getuid() or stat.S_IMODE(root_info.st_mode) != 0o700:
|
|
raise OSError("unsafe DWH workspace root")
|
|
fd = os.open("generation.lock", os.O_RDWR | os.O_NOFOLLOW, dir_fd=root_fd)
|
|
except BaseException:
|
|
os.close(root_fd)
|
|
raise
|
|
try:
|
|
info = os.fstat(fd)
|
|
if (
|
|
not stat.S_ISREG(info.st_mode)
|
|
or info.st_uid != os.getuid()
|
|
or info.st_nlink != 1
|
|
or stat.S_IMODE(info.st_mode) != 0o600
|
|
):
|
|
raise OSError("unsafe DWH generation lock")
|
|
fcntl.flock(fd, fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH)
|
|
return _GenerationLease(root_fd, fd)
|
|
except BaseException:
|
|
os.close(fd)
|
|
os.close(root_fd)
|
|
raise
|
|
except OSError as error:
|
|
raise CorruptCheckpointError(
|
|
"DWH workspace ownership is not initialized; run preprocessing first"
|
|
) from error
|
|
|
|
|
|
def _digest(path: Path) -> str:
|
|
return hashlib.sha256(path.read_bytes()).hexdigest()
|
|
|
|
|
|
def _read_owned_at(directory_fd: int, name: str, *, readonly: bool) -> bytes:
|
|
fd = os.open(name, os.O_RDONLY | os.O_NOFOLLOW, dir_fd=directory_fd)
|
|
try:
|
|
info = os.fstat(fd)
|
|
if (
|
|
not stat.S_ISREG(info.st_mode)
|
|
or info.st_uid != os.getuid()
|
|
or info.st_nlink != 1
|
|
or (readonly and bool(info.st_mode & 0o222))
|
|
):
|
|
raise OSError("unsafe DWH generation file")
|
|
chunks = []
|
|
while chunk := os.read(fd, 1024 * 1024):
|
|
chunks.append(chunk)
|
|
return b"".join(chunks)
|
|
finally:
|
|
os.close(fd)
|
|
|
|
|
|
def _open_generations_fd(root_fd: int, *, create: bool = False) -> int:
|
|
if create:
|
|
try:
|
|
os.mkdir("generations", 0o700, dir_fd=root_fd)
|
|
os.fsync(root_fd)
|
|
except FileExistsError:
|
|
pass
|
|
fd = os.open(
|
|
"generations", os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=root_fd
|
|
)
|
|
info = os.fstat(fd)
|
|
if (
|
|
not stat.S_ISDIR(info.st_mode)
|
|
or info.st_uid != os.getuid()
|
|
or stat.S_IMODE(info.st_mode) != 0o700
|
|
):
|
|
os.close(fd)
|
|
raise CorruptCheckpointError("DWH generations directory is invalid")
|
|
return fd
|
|
|
|
|
|
def _open_generation_fd(generations_fd: int, generation: str) -> int:
|
|
if not _RUN_ID.fullmatch(generation):
|
|
raise CorruptCheckpointError("DWH generation identity is invalid")
|
|
try:
|
|
return os.open(
|
|
generation, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW,
|
|
dir_fd=generations_fd,
|
|
)
|
|
except OSError as error:
|
|
raise CorruptCheckpointError("active DWH generation is missing") from error
|
|
|
|
|
|
def _active_generation_fd(
|
|
root_fd: int, expected_binding: dict[str, str], *, validate: bool = True
|
|
) -> tuple[str, int] | None:
|
|
_validate_root_binding_fd(root_fd, expected_binding)
|
|
try:
|
|
generation = _read_owned_at(root_fd, "ACTIVE", readonly=False).decode().strip()
|
|
except FileNotFoundError:
|
|
try:
|
|
generations_fd = _open_generations_fd(root_fd)
|
|
except FileNotFoundError:
|
|
return None
|
|
try:
|
|
if os.listdir(generations_fd):
|
|
raise CorruptCheckpointError(
|
|
"DWH generations exist without a consistent ACTIVE pointer"
|
|
)
|
|
finally:
|
|
os.close(generations_fd)
|
|
return None
|
|
except (OSError, UnicodeDecodeError) as error:
|
|
raise CorruptCheckpointError("DWH ACTIVE pointer is invalid") from error
|
|
if not _RUN_ID.fullmatch(generation):
|
|
raise CorruptCheckpointError("DWH ACTIVE pointer is invalid")
|
|
generations_fd = _open_generations_fd(root_fd)
|
|
try:
|
|
generation_fd = _open_generation_fd(generations_fd, generation)
|
|
finally:
|
|
os.close(generations_fd)
|
|
try:
|
|
if validate:
|
|
validate_generation_fd(generation_fd, generation, expected_binding)
|
|
return generation, generation_fd
|
|
except BaseException:
|
|
os.close(generation_fd)
|
|
raise
|
|
|
|
|
|
def _materialize_generation_fd(
|
|
generation_fd: int, generation: str, binding: dict[str, str]
|
|
) -> tuple[Path, Path]:
|
|
holder = Path(tempfile.mkdtemp(prefix="tht-dwh-snapshot-"))
|
|
_SNAPSHOT_DIRS.add(holder)
|
|
snapshot_root = holder
|
|
try:
|
|
_, payloads = _read_validated_generation_fd(generation_fd, generation, binding)
|
|
for name, payload in payloads.items():
|
|
destination = snapshot_root / name
|
|
destination.write_bytes(payload)
|
|
destination.chmod(0o400)
|
|
return holder, snapshot_root
|
|
except BaseException:
|
|
shutil.rmtree(holder, ignore_errors=True)
|
|
_SNAPSHOT_DIRS.discard(holder)
|
|
raise
|
|
|
|
|
|
def validate_generation_fd(
|
|
directory_fd: int, generation: str, expected_binding: dict[str, str] | None = None,
|
|
) -> dict:
|
|
manifest, _ = _read_validated_generation_fd(
|
|
directory_fd, generation, expected_binding
|
|
)
|
|
return manifest
|
|
|
|
|
|
def _read_validated_generation_fd(
|
|
directory_fd: int, generation: str, expected_binding: dict[str, str] | None = None,
|
|
) -> tuple[dict, dict[str, bytes]]:
|
|
try:
|
|
directory_info = os.fstat(directory_fd)
|
|
if (
|
|
not stat.S_ISDIR(directory_info.st_mode)
|
|
or directory_info.st_uid != os.getuid()
|
|
or stat.S_IMODE(directory_info.st_mode) != 0o700
|
|
):
|
|
raise ValueError
|
|
manifest = json.loads(
|
|
_read_owned_at(directory_fd, GENERATION_MANIFEST, readonly=True).decode("utf-8")
|
|
)
|
|
files = manifest["files"]
|
|
if (
|
|
manifest["generation"] != generation
|
|
or not isinstance(files, dict)
|
|
or not re.fullmatch(r"sha256:[0-9a-f]{64}", manifest["job_spec_fingerprint"])
|
|
or not re.fullmatch(r"[0-9a-f]{64}", manifest["artifact_manifest_sha256"])
|
|
or not re.fullmatch(r"[a-z][a-z0-9_-]{0,63}", manifest["workspace_id"])
|
|
or not re.fullmatch(r"sha256:[0-9a-f]{64}", manifest["config_fingerprint"])
|
|
or not re.fullmatch(r"sha256:[0-9a-f]{64}", manifest["input_fingerprint"])
|
|
):
|
|
raise ValueError
|
|
if expected_binding is not None and any(
|
|
manifest.get(key) != value for key, value in expected_binding.items()
|
|
):
|
|
raise CorruptCheckpointError(
|
|
"DWH artifacts belong to a different workspace configuration"
|
|
)
|
|
if set(os.listdir(directory_fd)) != set(files) | {GENERATION_MANIFEST}:
|
|
raise ValueError
|
|
payloads = {}
|
|
for name, expected in files.items():
|
|
if name != "physical.yaml" and not _SAFE_FILE.fullmatch(name):
|
|
raise ValueError
|
|
payload = _read_owned_at(directory_fd, name, readonly=True)
|
|
if hashlib.sha256(payload).hexdigest() != expected:
|
|
raise ValueError
|
|
payloads[name] = payload
|
|
return manifest, payloads
|
|
except (OSError, KeyError, TypeError, ValueError, UnicodeDecodeError, json.JSONDecodeError) as error:
|
|
raise CorruptCheckpointError("published DWH generation is invalid") from error
|
|
|
|
|
|
def resolve_dwh_snapshot(cfg) -> DwhArtifactSnapshot:
|
|
binding = config_dwh_binding(cfg)
|
|
lease = _acquire_existing_generation_lock(
|
|
cfg.paths.artifacts.parent, exclusive=False
|
|
)
|
|
try:
|
|
return _resolve_dwh_snapshot_locked(cfg, binding, lease.root_fd)
|
|
finally:
|
|
lease.close()
|
|
|
|
|
|
def _resolve_dwh_snapshot_locked(
|
|
cfg, binding: dict[str, str], root_fd: int,
|
|
) -> DwhArtifactSnapshot:
|
|
active = _active_generation_fd(root_fd, binding, validate=False)
|
|
if active is None:
|
|
return DwhArtifactSnapshot(
|
|
None, cfg.paths.artifacts / "mschema" / "physical.yaml", cfg.paths.indexes / "lsh"
|
|
)
|
|
generation, generation_fd = active
|
|
try:
|
|
holder, snapshot_root = _materialize_generation_fd(
|
|
generation_fd, generation, binding
|
|
)
|
|
finally:
|
|
os.close(generation_fd)
|
|
return DwhArtifactSnapshot(
|
|
generation, snapshot_root / "physical.yaml", snapshot_root, holder
|
|
)
|
|
|
|
|
|
def active_generation_dir(
|
|
workspace_root: Path, expected_binding: dict[str, str]
|
|
) -> Path | None:
|
|
lease = _acquire_existing_generation_lock(workspace_root, exclusive=False)
|
|
try:
|
|
active = _active_generation_fd(lease.root_fd, expected_binding)
|
|
if active is None:
|
|
return None
|
|
generation, generation_fd = active
|
|
os.close(generation_fd)
|
|
lease.assert_root_identity(workspace_root / ".tht-dwh")
|
|
return workspace_root / ".tht-dwh" / "generations" / generation
|
|
finally:
|
|
lease.close()
|
|
|
|
|
|
class DwhPreprocessPipeline:
|
|
"""Stage a complete artifact bundle, then publish it through one atomic pointer."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
workspace_id: str,
|
|
workspace_root: Path,
|
|
config_fingerprint: str,
|
|
input_fingerprint: str,
|
|
introspect: Callable[[Path], object],
|
|
build_lsh: Callable[[Path, Path], object],
|
|
lsh_filenames: tuple[str, str, str] | None = None,
|
|
current_physical: Path | None = None,
|
|
current_lsh_dir: Path | None = None,
|
|
after_publish: Callable[[str], object] | None = None,
|
|
retain_generations: int = 3,
|
|
) -> None:
|
|
self.workspace_id = workspace_id
|
|
self.workspace_root = workspace_root
|
|
self.config_fingerprint = config_fingerprint
|
|
self.input_fingerprint = input_fingerprint
|
|
self.introspect = introspect
|
|
self.build_lsh = build_lsh
|
|
self.lsh_filenames = lsh_filenames or (
|
|
f"{workspace_id}_lsh.pkl", f"{workspace_id}_minhashes.pkl",
|
|
f"{workspace_id}_meta.json",
|
|
)
|
|
self.current_physical = current_physical
|
|
self.current_lsh_dir = current_lsh_dir
|
|
self._snapshot_holder = None
|
|
self.after_publish = after_publish
|
|
if isinstance(retain_generations, bool) or retain_generations < 1:
|
|
raise ValueError("retain_generations must be positive")
|
|
self.retain_generations = retain_generations
|
|
if len(set(self.lsh_filenames)) != 3 or any(
|
|
not _SAFE_FILE.fullmatch(name) for name in self.lsh_filenames
|
|
):
|
|
raise ValueError("LSH filenames must be unique flat safe names")
|
|
|
|
@property
|
|
def binding(self) -> dict[str, str]:
|
|
return {
|
|
"workspace_id": self.workspace_id,
|
|
"config_fingerprint": self.config_fingerprint,
|
|
"input_fingerprint": self.input_fingerprint,
|
|
}
|
|
|
|
def _assert_active_binding(self, root_fd: int) -> None:
|
|
_validate_root_binding_fd(root_fd, self.binding)
|
|
active = _active_generation_fd(root_fd, self.binding)
|
|
if active is not None:
|
|
_, generation_fd = active
|
|
os.close(generation_fd)
|
|
|
|
def run(
|
|
self, steps: tuple[str, ...] = DWH_STAGE_IDS, *, resume_run_id: str | None = None
|
|
) -> JobReport:
|
|
self._validate_steps(steps)
|
|
lease = _acquire_generation_lock(self.workspace_root, exclusive=True)
|
|
try:
|
|
self._assert_no_legacy_artifacts(lease.root_fd)
|
|
_claim_or_validate_root_binding(lease.root_fd, self.binding)
|
|
lease.assert_root_identity(self.workspace_root / ".tht-dwh")
|
|
self._assert_active_binding(lease.root_fd)
|
|
active = _active_generation_fd(lease.root_fd, self.binding, validate=False)
|
|
if active is not None:
|
|
generation, generation_fd = active
|
|
try:
|
|
self._snapshot_holder, snapshot_root = _materialize_generation_fd(
|
|
generation_fd, generation, self.binding
|
|
)
|
|
finally:
|
|
os.close(generation_fd)
|
|
self.current_physical = snapshot_root / "physical.yaml"
|
|
self.current_lsh_dir = snapshot_root
|
|
finally:
|
|
lease.close()
|
|
try:
|
|
if resume_run_id is not None:
|
|
self._validate_resume_publication(resume_run_id)
|
|
spec = JobSpec(
|
|
workspace_id=self.workspace_id,
|
|
job_type="dwh",
|
|
workspace_root=self.workspace_root,
|
|
spec_version="jobs-v1",
|
|
pipeline_version="dwh-v2",
|
|
config_fingerprint=self.config_fingerprint,
|
|
input_fingerprint=self.input_fingerprint,
|
|
stage_ids=steps,
|
|
resume_run_id=resume_run_id,
|
|
)
|
|
except BaseException:
|
|
self._release_pipeline_snapshot()
|
|
raise
|
|
|
|
def introspect_stage(context: JobContext):
|
|
artifacts = self._artifacts(context)
|
|
physical = artifacts / "physical.yaml"
|
|
try:
|
|
self.introspect(physical)
|
|
except BaseException:
|
|
physical.unlink(missing_ok=True)
|
|
raise
|
|
if steps[-1] == "introspect":
|
|
self._copy_current_lsh(artifacts)
|
|
return self._publish_stage(context, "introspect", spec)
|
|
return StageArtifacts(("physical.yaml",))
|
|
|
|
def lsh_stage(context: JobContext):
|
|
artifacts = self._artifacts(context)
|
|
physical = artifacts / "physical.yaml"
|
|
if not physical.exists():
|
|
source = self.current_physical
|
|
if source is None or not source.is_file():
|
|
raise FileNotFoundError("physical catalog is missing")
|
|
shutil.copyfile(source, physical)
|
|
try:
|
|
self.build_lsh(physical, artifacts)
|
|
except BaseException:
|
|
for name in self.lsh_filenames:
|
|
(artifacts / name).unlink(missing_ok=True)
|
|
raise
|
|
return self._publish_stage(context, "lsh", spec)
|
|
|
|
implementations = {"introspect": introspect_stage, "lsh": lsh_stage}
|
|
try:
|
|
return run_job(
|
|
spec, tuple(implementations[step] for step in steps),
|
|
reconcile_effects=self._reconcile_effects,
|
|
)
|
|
finally:
|
|
self._release_pipeline_snapshot()
|
|
|
|
def _release_pipeline_snapshot(self) -> None:
|
|
holder, self._snapshot_holder = self._snapshot_holder, None
|
|
if isinstance(holder, Path):
|
|
shutil.rmtree(holder, ignore_errors=True)
|
|
_SNAPSHOT_DIRS.discard(holder)
|
|
if self.current_physical is not None and holder in self.current_physical.parents:
|
|
self.current_physical = None
|
|
if self.current_lsh_dir is not None and (
|
|
self.current_lsh_dir == holder or holder in self.current_lsh_dir.parents
|
|
):
|
|
self.current_lsh_dir = None
|
|
|
|
def _reconcile_effects(self, source, run_dir: Path) -> set[str]:
|
|
running = next(
|
|
(stage for stage in source.stages if stage.status == "running" and stage.effect_state == "intent"),
|
|
None,
|
|
)
|
|
if running is None:
|
|
return set()
|
|
lease = _acquire_existing_generation_lock(self.workspace_root, exclusive=True)
|
|
try:
|
|
generations_fd = _open_generations_fd(lease.root_fd)
|
|
try:
|
|
self._validate_published_fd(
|
|
generations_fd, source.run_id,
|
|
run_dir / "artifacts", running.artifact_files,
|
|
)
|
|
active = _active_generation_fd(lease.root_fd, self.binding)
|
|
try:
|
|
if active is None or active[0] != source.run_id:
|
|
raise CorruptCheckpointError("sealed DWH publication is not ACTIVE")
|
|
finally:
|
|
if active is not None:
|
|
os.close(active[1])
|
|
finally:
|
|
os.close(generations_fd)
|
|
finally:
|
|
lease.close()
|
|
return {running.name}
|
|
|
|
def _validate_resume_publication(self, run_id: str) -> None:
|
|
run_dir = self.workspace_root / ".tht-jobs" / "dwh" / "runs" / run_id
|
|
try:
|
|
checkpoint = json.loads((run_dir / "checkpoint.json").read_text(encoding="utf-8"))
|
|
except (OSError, ValueError, TypeError) as error:
|
|
raise CorruptCheckpointError("checkpoint is invalid and cannot be resumed") from error
|
|
lease = _acquire_existing_generation_lock(self.workspace_root, exclusive=True)
|
|
try:
|
|
try:
|
|
generations_fd = _open_generations_fd(lease.root_fd)
|
|
except FileNotFoundError:
|
|
if checkpoint.get("status") == "succeeded":
|
|
raise CorruptCheckpointError("published DWH generation is missing")
|
|
return
|
|
try:
|
|
if run_id not in os.listdir(generations_fd):
|
|
if checkpoint.get("status") == "succeeded":
|
|
raise CorruptCheckpointError("published DWH generation is missing")
|
|
return
|
|
try:
|
|
manifest = json.loads(
|
|
(run_dir / "artifacts" / "artifact-manifest.json").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
)
|
|
required = tuple(
|
|
name
|
|
for stage in manifest["stages"].values()
|
|
for name in stage["required"]
|
|
)
|
|
except (OSError, KeyError, ValueError, TypeError) as error:
|
|
raise CorruptCheckpointError(
|
|
"resume artifact manifest is invalid"
|
|
) from error
|
|
self._validate_published_fd(
|
|
generations_fd, run_id, run_dir / "artifacts", required
|
|
)
|
|
finally:
|
|
os.close(generations_fd)
|
|
finally:
|
|
lease.close()
|
|
|
|
def _assert_no_legacy_artifacts(self, root_fd: int) -> None:
|
|
if OWNER_MARKER in os.listdir(root_fd):
|
|
return
|
|
legacy_physical = self.current_physical is not None and (
|
|
self.current_physical.exists() or self.current_physical.is_symlink()
|
|
)
|
|
legacy_lsh = self.current_lsh_dir is not None and any(
|
|
(self.current_lsh_dir / name).exists()
|
|
or (self.current_lsh_dir / name).is_symlink()
|
|
for name in self.lsh_filenames
|
|
)
|
|
if legacy_physical or legacy_lsh:
|
|
raise CorruptCheckpointError(
|
|
"DWH legacy artifacts are unbound; migrate them explicitly or use an empty root"
|
|
)
|
|
|
|
@staticmethod
|
|
def _artifacts(context: JobContext) -> Path:
|
|
root = context.run_dir / "artifacts"
|
|
root.mkdir(exist_ok=True)
|
|
return root
|
|
|
|
def _required(self, artifacts: Path) -> tuple[str, ...]:
|
|
required = ("physical.yaml",) + tuple(
|
|
name for name in self.lsh_filenames if (artifacts / name).is_file()
|
|
)
|
|
if not (artifacts / "physical.yaml").is_file():
|
|
raise CorruptCheckpointError("staged physical catalog is missing")
|
|
lsh_count = len(required) - 1
|
|
if lsh_count not in (0, len(self.lsh_filenames)):
|
|
raise CorruptCheckpointError("staged LSH artifact set is incomplete")
|
|
return required
|
|
|
|
def _copy_current_lsh(self, artifacts: Path) -> None:
|
|
if self.current_lsh_dir is None:
|
|
return
|
|
existing = [self.current_lsh_dir / name for name in self.lsh_filenames]
|
|
if not any(path.exists() for path in existing):
|
|
return
|
|
if not all(path.is_file() for path in existing):
|
|
raise CorruptCheckpointError("current LSH artifact set is incomplete")
|
|
for source in existing:
|
|
shutil.copyfile(source, artifacts / source.name)
|
|
|
|
def _publish_stage(self, context: JobContext, stage: str, spec: JobSpec):
|
|
artifacts = self._artifacts(context)
|
|
required = self._required(artifacts)
|
|
seal_stage_artifacts(context, stage, required, spec)
|
|
lease = _acquire_generation_lock(self.workspace_root, exclusive=True)
|
|
try:
|
|
_validate_root_binding_fd(lease.root_fd, self.binding)
|
|
self._assert_active_binding(lease.root_fd)
|
|
self._publish(context.run_id, artifacts, required, lease.root_fd)
|
|
if self.after_publish is not None:
|
|
self.after_publish(context.run_id)
|
|
self._cleanup_generations_fd(lease.root_fd)
|
|
finally:
|
|
lease.close()
|
|
return StageArtifacts(required)
|
|
|
|
def _publish(
|
|
self, generation: str, artifacts: Path, required: tuple[str, ...], root_fd: int
|
|
) -> None:
|
|
generations_fd = _open_generations_fd(root_fd, create=True)
|
|
temporary = f".{generation}.{uuid.uuid4().hex}.tmp"
|
|
try:
|
|
if generation in os.listdir(generations_fd):
|
|
self._validate_published_fd(generations_fd, generation, artifacts, required)
|
|
else:
|
|
os.mkdir(temporary, 0o700, dir_fd=generations_fd)
|
|
temporary_fd = os.open(
|
|
temporary, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW,
|
|
dir_fd=generations_fd,
|
|
)
|
|
try:
|
|
for name in required:
|
|
self._write_readonly_at(
|
|
temporary_fd, name, (artifacts / name).read_bytes()
|
|
)
|
|
artifact_manifest = artifacts / "artifact-manifest.json"
|
|
manifest = {
|
|
"schema_version": 1,
|
|
"generation": generation,
|
|
"files": {name: _digest(artifacts / name) for name in required},
|
|
"job_spec_fingerprint": json.loads(
|
|
artifact_manifest.read_text(encoding="utf-8")
|
|
)["spec_fingerprint"],
|
|
"artifact_manifest_sha256": _digest(artifact_manifest),
|
|
**self.binding,
|
|
}
|
|
self._write_readonly_at(
|
|
temporary_fd, GENERATION_MANIFEST,
|
|
(json.dumps(manifest, sort_keys=True, separators=(",", ":")) + "\n").encode(),
|
|
)
|
|
os.fsync(temporary_fd)
|
|
finally:
|
|
os.close(temporary_fd)
|
|
os.rename(
|
|
temporary, generation,
|
|
src_dir_fd=generations_fd, dst_dir_fd=generations_fd,
|
|
)
|
|
temporary = ""
|
|
os.fsync(generations_fd)
|
|
previous = None
|
|
try:
|
|
previous = _read_owned_at(root_fd, "ACTIVE", readonly=False)
|
|
except FileNotFoundError:
|
|
pass
|
|
self._replace_active_at(root_fd, (generation + "\n").encode(), previous)
|
|
finally:
|
|
if temporary:
|
|
self._safe_delete_generation(generations_fd, temporary, allow_temporary=True)
|
|
os.close(generations_fd)
|
|
|
|
@staticmethod
|
|
def _write_readonly_at(directory_fd: int, name: str, payload: bytes) -> None:
|
|
fd = os.open(
|
|
name, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW,
|
|
0o600, dir_fd=directory_fd,
|
|
)
|
|
try:
|
|
offset = 0
|
|
while offset < len(payload):
|
|
offset += os.write(fd, payload[offset:])
|
|
os.fsync(fd)
|
|
os.fchmod(fd, 0o400)
|
|
os.fsync(fd)
|
|
finally:
|
|
os.close(fd)
|
|
|
|
def _replace_active_at(self, root_fd: int, payload: bytes, previous: bytes | None) -> None:
|
|
temporary = f".ACTIVE.{uuid.uuid4().hex}.tmp"
|
|
fd = os.open(
|
|
temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW,
|
|
0o600, dir_fd=root_fd,
|
|
)
|
|
try:
|
|
offset = 0
|
|
while offset < len(payload):
|
|
offset += os.write(fd, payload[offset:])
|
|
os.fsync(fd)
|
|
finally:
|
|
os.close(fd)
|
|
try:
|
|
os.replace(temporary, "ACTIVE", src_dir_fd=root_fd, dst_dir_fd=root_fd)
|
|
temporary = ""
|
|
try:
|
|
os.fsync(root_fd)
|
|
except BaseException:
|
|
self._restore_pointer_at(root_fd, previous)
|
|
raise
|
|
finally:
|
|
if temporary:
|
|
try:
|
|
os.unlink(temporary, dir_fd=root_fd)
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
def _validate_published_fd(
|
|
self, generations_fd: int, generation: str,
|
|
artifacts: Path, required: tuple[str, ...],
|
|
) -> None:
|
|
generation_fd = _open_generation_fd(generations_fd, generation)
|
|
try:
|
|
manifest = validate_generation_fd(generation_fd, generation, self.binding)
|
|
if set(manifest["files"]) != set(required):
|
|
raise CorruptCheckpointError("published DWH generation is invalid")
|
|
artifact_digest = _digest(artifacts / "artifact-manifest.json")
|
|
if manifest["artifact_manifest_sha256"] != artifact_digest:
|
|
raise CorruptCheckpointError("published DWH job manifest digest mismatch")
|
|
for name in required:
|
|
if hashlib.sha256((artifacts / name).read_bytes()).hexdigest() != manifest["files"][name]:
|
|
raise CorruptCheckpointError("published DWH artifact digest mismatch")
|
|
finally:
|
|
os.close(generation_fd)
|
|
|
|
def _restore_pointer_at(self, root_fd: int, previous: bytes | None) -> None:
|
|
if previous is None:
|
|
try:
|
|
os.unlink("ACTIVE", dir_fd=root_fd)
|
|
except FileNotFoundError:
|
|
pass
|
|
else:
|
|
restore = f".ACTIVE.restore.{uuid.uuid4().hex}.tmp"
|
|
fd = os.open(
|
|
restore, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW,
|
|
0o600, dir_fd=root_fd,
|
|
)
|
|
try:
|
|
offset = 0
|
|
while offset < len(previous):
|
|
offset += os.write(fd, previous[offset:])
|
|
os.fsync(fd)
|
|
finally:
|
|
os.close(fd)
|
|
os.replace(restore, "ACTIVE", src_dir_fd=root_fd, dst_dir_fd=root_fd)
|
|
os.fsync(root_fd)
|
|
|
|
def _cleanup_generations(self) -> None:
|
|
lease = _acquire_existing_generation_lock(self.workspace_root, exclusive=True)
|
|
try:
|
|
self._cleanup_generations_fd(lease.root_fd)
|
|
finally:
|
|
lease.close()
|
|
|
|
def _cleanup_generations_fd(self, root_fd: int) -> None:
|
|
_validate_root_binding_fd(root_fd, self.binding)
|
|
active = _active_generation_fd(root_fd, self.binding)
|
|
active_name = active[0] if active else None
|
|
if active is not None:
|
|
os.close(active[1])
|
|
protected = {active_name} if active_name else set()
|
|
runs = self.workspace_root / ".tht-jobs" / "dwh" / "runs"
|
|
for checkpoint in runs.glob("*/checkpoint.json") if runs.exists() else ():
|
|
try:
|
|
status = json.loads(checkpoint.read_text(encoding="utf-8"))["status"]
|
|
if status in {"running", "failed"}:
|
|
protected.add(checkpoint.parent.name)
|
|
except (OSError, KeyError, ValueError):
|
|
continue
|
|
generations = []
|
|
try:
|
|
generations_fd = _open_generations_fd(root_fd)
|
|
except FileNotFoundError:
|
|
return
|
|
try:
|
|
for name in os.listdir(generations_fd):
|
|
if not _RUN_ID.fullmatch(name):
|
|
continue
|
|
try:
|
|
candidate_fd = os.open(
|
|
name, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW,
|
|
dir_fd=generations_fd,
|
|
)
|
|
except OSError:
|
|
continue
|
|
try:
|
|
info = os.fstat(candidate_fd)
|
|
validate_generation_fd(candidate_fd, name, self.binding)
|
|
except CorruptCheckpointError:
|
|
continue
|
|
finally:
|
|
os.close(candidate_fd)
|
|
generations.append((info.st_mtime_ns, name))
|
|
generations.sort(key=lambda value: (value[0], value[1]))
|
|
rollback = [value for value in generations if value[1] != active_name]
|
|
rollback_count = self.retain_generations - 1
|
|
keep_recent = {
|
|
name for _, name in (rollback[-rollback_count:] if rollback_count else ())
|
|
}
|
|
for _, name in generations:
|
|
if name in protected | keep_recent:
|
|
continue
|
|
self._safe_delete_generation(generations_fd, name)
|
|
os.fsync(generations_fd)
|
|
finally:
|
|
os.close(generations_fd)
|
|
|
|
@staticmethod
|
|
def _safe_delete_generation(
|
|
root_fd: int, name: str, *, allow_temporary: bool = False
|
|
) -> None:
|
|
if not (_RUN_ID.fullmatch(name) or (allow_temporary and name.startswith("."))):
|
|
return
|
|
try:
|
|
generation_fd = os.open(
|
|
name, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=root_fd
|
|
)
|
|
except OSError:
|
|
return
|
|
try:
|
|
info = os.fstat(generation_fd)
|
|
if not stat.S_ISDIR(info.st_mode) or info.st_uid != os.getuid():
|
|
return
|
|
entries = os.listdir(generation_fd)
|
|
opened = []
|
|
try:
|
|
for child in entries:
|
|
try:
|
|
fd = os.open(
|
|
child, os.O_RDONLY | os.O_NOFOLLOW, dir_fd=generation_fd
|
|
)
|
|
except OSError:
|
|
return
|
|
child_info = os.fstat(fd)
|
|
if (
|
|
not stat.S_ISREG(child_info.st_mode)
|
|
or child_info.st_uid != os.getuid()
|
|
or child_info.st_nlink != 1
|
|
):
|
|
os.close(fd)
|
|
return
|
|
opened.append((child, fd))
|
|
for child, fd in opened:
|
|
os.fchmod(fd, 0o600)
|
|
os.close(fd)
|
|
os.unlink(child, dir_fd=generation_fd)
|
|
opened.clear()
|
|
finally:
|
|
for _, fd in opened:
|
|
os.close(fd)
|
|
finally:
|
|
os.close(generation_fd)
|
|
try:
|
|
os.rmdir(name, dir_fd=root_fd)
|
|
except OSError:
|
|
return
|
|
|
|
@staticmethod
|
|
def _ensure_owned_dir(path: Path) -> None:
|
|
try:
|
|
path.mkdir(mode=0o700)
|
|
except FileExistsError:
|
|
pass
|
|
info = path.lstat()
|
|
if not stat.S_ISDIR(info.st_mode) or info.st_uid != os.getuid():
|
|
raise OSError("unsafe DWH publication directory")
|
|
path.chmod(0o700)
|
|
|
|
@staticmethod
|
|
def _validate_steps(steps: tuple[str, ...]) -> None:
|
|
if not steps or len(steps) != len(set(steps)) or any(
|
|
step not in DWH_STAGE_IDS for step in steps
|
|
):
|
|
raise ValueError("DWH preprocessing steps must be unique introspect/lsh stages")
|
|
if tuple(sorted(steps, key=DWH_STAGE_IDS.index)) != steps:
|
|
raise ValueError("DWH preprocessing steps must follow introspect,lsh order")
|
|
|
|
|
|
def fingerprint(value: str) -> str:
|
|
return "sha256:" + hashlib.sha256(value.encode("utf-8")).hexdigest()
|