"""Crash-safe, resumable DWH catalog and LSH preprocessing stages.""" from __future__ import annotations import hashlib import json import fcntl import atexit 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") return { "workspace_id": workspace_id, "config_fingerprint": fingerprint(cfg.model_dump_json()), "input_fingerprint": fingerprint(config_source), } 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()