mptyl/ThothII
1
0
Fork 0
Code Issues 25 Pull Requests Actions Packages Projects Releases 1 Wiki Activity

fix(preprocess): claim DWH roots atomically

Browse Source
This commit is contained in:
marcopan
2026-07-12 06:51:46 +02:00
parent 8110793f61
commit c4dd6c8900
2 changed files with 199 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
harness/tht/jobs/dwh_pipeline.py
+124 -3
View File
@@ -28,6 +28,7 @@ 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"
def config_dwh_binding(cfg) -> dict[str, str]:
@@ -42,6 +43,96 @@ def config_dwh_binding(cfg) -> dict[str, str]:
}
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(workspace_root: Path) -> dict[str, str]:
marker = workspace_root / ".tht-dwh" / OWNER_MARKER
try:
payload = json.loads(_read_owned(marker, readonly=True).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(workspace_root: Path, expected: dict[str, str]) -> None:
if _read_root_binding(workspace_root) != expected:
raise CorruptCheckpointError("DWH artifacts belong to a different workspace configuration")
root = workspace_root / ".tht-dwh"
generations = root / "generations"
active_exists = (root / "ACTIVE").exists()
try:
generations_fd = os.open(
generations, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
)
except FileNotFoundError:
generation_entries = []
except OSError as error:
raise CorruptCheckpointError("DWH generations directory is invalid") from error
else:
try:
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(workspace_root: Path, binding: dict[str, str]) -> None:
root = workspace_root / ".tht-dwh"
marker = root / OWNER_MARKER
if marker.exists() or marker.is_symlink():
_validate_root_binding(workspace_root, binding)
return
generations = root / "generations"
generations_nonempty = False
try:
generations_fd = os.open(
generations, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
)
except FileNotFoundError:
pass
except OSError as error:
raise CorruptCheckpointError("DWH generations directory is invalid") from error
else:
try:
generations_nonempty = bool(os.listdir(generations_fd))
finally:
os.close(generations_fd)
if (root / "ACTIVE").exists() or generations_nonempty:
raise CorruptCheckpointError(
"DWH artifacts are unbound; migrate them explicitly or use an empty root"
)
payload = {
"schema_version": 1,
"binding": binding,
"binding_sha256": _binding_digest(binding),
}
temporary = root / f".{OWNER_MARKER}.{uuid.uuid4().hex}.tmp"
fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
try:
with os.fdopen(fd, "w", encoding="utf-8") as stream:
stream.write(json.dumps(payload, sort_keys=True, separators=(",", ":")) + "\n")
stream.flush()
os.fsync(stream.fileno())
temporary.chmod(0o400)
os.replace(temporary, marker)
DwhPreprocessPipeline._fsync(root)
except BaseException:
temporary.unlink(missing_ok=True)
raise
@dataclass(frozen=True)
class DwhArtifactSnapshot:
generation: str | None
@@ -56,9 +147,16 @@ class DwhSnapshotLease:
self.snapshot: DwhArtifactSnapshot | None = None
def __enter__(self) -> DwhArtifactSnapshot:
binding = config_dwh_binding(self.cfg)
claim_fd = _acquire_generation_lock(self.cfg.paths.artifacts.parent, exclusive=True)
try:
_claim_or_validate_root_binding(self.cfg.paths.artifacts.parent, binding)
finally:
fcntl.flock(claim_fd, fcntl.LOCK_UN)
os.close(claim_fd)
self._fd = _acquire_generation_lock(self.cfg.paths.artifacts.parent, exclusive=False)
try:
self.snapshot = resolve_dwh_snapshot(self.cfg)
self.snapshot = _resolve_dwh_snapshot_locked(self.cfg, binding)
return self.snapshot
except BaseException:
self.__exit__(None, None, None)
@@ -190,12 +288,31 @@ def validate_generation(
def resolve_dwh_snapshot(cfg) -> DwhArtifactSnapshot:
binding = config_dwh_binding(cfg)
claim_fd = _acquire_generation_lock(cfg.paths.artifacts.parent, exclusive=True)
try:
_claim_or_validate_root_binding(cfg.paths.artifacts.parent, binding)
finally:
fcntl.flock(claim_fd, fcntl.LOCK_UN)
os.close(claim_fd)
lease_fd = _acquire_generation_lock(cfg.paths.artifacts.parent, exclusive=False)
try:
return _resolve_dwh_snapshot_locked(cfg, binding)
finally:
fcntl.flock(lease_fd, fcntl.LOCK_UN)
os.close(lease_fd)
def _resolve_dwh_snapshot_locked(
cfg, binding: dict[str, str],
) -> DwhArtifactSnapshot:
_validate_root_binding(cfg.paths.artifacts.parent, binding)
target = active_generation_dir(cfg.paths.artifacts.parent)
if target is None:
return DwhArtifactSnapshot(
None, cfg.paths.artifacts / "mschema" / "physical.yaml", cfg.paths.indexes / "lsh"
)
validate_generation(target, config_dwh_binding(cfg))
validate_generation(target, binding)
return DwhArtifactSnapshot(target.name, target / "physical.yaml", target)
@@ -263,6 +380,7 @@ class DwhPreprocessPipeline:
}
def _assert_active_binding(self) -> None:
_validate_root_binding(self.workspace_root, self.binding)
active = active_generation_dir(self.workspace_root)
if active is not None:
validate_generation(active, self.binding)
@@ -271,8 +389,9 @@ class DwhPreprocessPipeline:
self, steps: tuple[str, ...] = DWH_STAGE_IDS, *, resume_run_id: str | None = None
) -> JobReport:
self._validate_steps(steps)
lease_fd = _acquire_generation_lock(self.workspace_root, exclusive=False)
lease_fd = _acquire_generation_lock(self.workspace_root, exclusive=True)
try:
_claim_or_validate_root_binding(self.workspace_root, self.binding)
self._assert_active_binding()
finally:
fcntl.flock(lease_fd, fcntl.LOCK_UN)
@@ -399,6 +518,7 @@ class DwhPreprocessPipeline:
seal_stage_artifacts(context, stage, required, spec)
lease_fd = _acquire_generation_lock(self.workspace_root, exclusive=True)
try:
_validate_root_binding(self.workspace_root, self.binding)
self._assert_active_binding()
self._publish(context.run_id, artifacts, required)
if self.after_publish is not None:
@@ -509,6 +629,7 @@ class DwhPreprocessPipeline:
root = self.workspace_root / ".tht-dwh" / "generations"
if not root.exists():
return
_validate_root_binding(self.workspace_root, self.binding)
active = active_generation_dir(self.workspace_root)
if active is not None:
validate_generation(active, self.binding)