feat(jobs): add resumable preprocessing envelope
This commit is contained in:
@@ -0,0 +1,57 @@
|
|||||||
|
# Evidence Task 4 — shared job envelope
|
||||||
|
|
||||||
|
Status: complete
|
||||||
|
|
||||||
|
## Delivered
|
||||||
|
|
||||||
|
- Immutable `JobSpec`, `JobRun`, `JobReport`, per-stage state, sanitized error, and UTC
|
||||||
|
timestamp records.
|
||||||
|
- `run_job(spec, stages)` with a durable checkpoint at job start, before and after every stage,
|
||||||
|
and at terminal state. Successful stages are skipped when a prior run is resumed.
|
||||||
|
- Atomic JSON checkpoint/report replacement using a unique same-directory temporary file,
|
||||||
|
file `fsync`, atomic `os.replace`, and parent-directory `fsync`.
|
||||||
|
- Public reports contain fixed operational fields only. Workspace paths, stage return values,
|
||||||
|
exception messages, source content, credentials, and arbitrary metadata are not serialized.
|
||||||
|
- `WorkspaceJobLock` uses non-blocking kernel `flock` on a stable workspace/job-specific inode.
|
||||||
|
Locks are released by the kernel on process exit; lock files are never removed based on PID,
|
||||||
|
avoiding stale-lock and PID-reuse deletion races. Evidence and DWH use distinct lock files.
|
||||||
|
- Dry-run intent is immutable in the spec/report and exposed to every stage through `JobContext`.
|
||||||
|
|
||||||
|
## TDD evidence
|
||||||
|
|
||||||
|
Initial focused collection failed because `tht.jobs` did not exist. Tests then drove:
|
||||||
|
|
||||||
|
- failure, sanitized reporting, resume, and idempotent successful-stage skipping;
|
||||||
|
- corrupt-checkpoint refusal before stage execution;
|
||||||
|
- JSON schema and path/secret/PII exclusion;
|
||||||
|
- dry-run propagation and ordered aware timestamps;
|
||||||
|
- multiprocessing exclusion, distinct Evidence/DWH jobs, traversal rejection, and recovery after
|
||||||
|
a lock-owning process crashes.
|
||||||
|
|
||||||
|
Final focused result:
|
||||||
|
|
||||||
|
```text
|
||||||
|
11 passed in 0.42s
|
||||||
|
```
|
||||||
|
|
||||||
|
## Verification
|
||||||
|
|
||||||
|
```text
|
||||||
|
cd harness && .venv/bin/pytest -q
|
||||||
|
597 passed, 5 deselected, 17 warnings in 28.45s
|
||||||
|
|
||||||
|
cd harness && .venv/bin/ruff check tht/jobs tests/test_job_runner.py tests/test_job_locking.py
|
||||||
|
All checks passed!
|
||||||
|
```
|
||||||
|
|
||||||
|
The full Ruff invocation was also run. It reports 34 pre-existing violations in unrelated legacy
|
||||||
|
tests; no Task 4 file is among them. L2 tests remain deselected by the repository configuration.
|
||||||
|
|
||||||
|
## Operational notes
|
||||||
|
|
||||||
|
- `fcntl.flock` intentionally targets the supported Linux/macOS deployment environments; it is not
|
||||||
|
a Windows locking implementation.
|
||||||
|
- The envelope does not publish or mutate an active corpus. Later pipeline stages must use
|
||||||
|
`JobContext.run_dir` for staging and perform their own final atomic publish only after validation.
|
||||||
|
- A dry run is an execution mode foundation: the runner exposes and records it; individual stages
|
||||||
|
remain responsible for suppressing external mutations.
|
||||||
@@ -0,0 +1,61 @@
|
|||||||
|
import multiprocessing
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from tht.jobs.locking import JobAlreadyRunningError, WorkspaceJobLock
|
||||||
|
|
||||||
|
|
||||||
|
def _hold_lock(root: str, ready, release):
|
||||||
|
with WorkspaceJobLock(Path(root), "demo", "evidence"):
|
||||||
|
ready.set()
|
||||||
|
release.wait(10)
|
||||||
|
|
||||||
|
|
||||||
|
def _crash_with_lock(root: str, ready):
|
||||||
|
lock = WorkspaceJobLock(Path(root), "demo", "evidence")
|
||||||
|
lock.acquire()
|
||||||
|
ready.set()
|
||||||
|
raise SystemExit(7)
|
||||||
|
|
||||||
|
|
||||||
|
def test_same_workspace_and_job_are_exclusive_across_processes(tmp_path):
|
||||||
|
context = multiprocessing.get_context("spawn")
|
||||||
|
ready = context.Event()
|
||||||
|
release = context.Event()
|
||||||
|
process = context.Process(target=_hold_lock, args=(str(tmp_path), ready, release))
|
||||||
|
process.start()
|
||||||
|
assert ready.wait(10)
|
||||||
|
try:
|
||||||
|
with pytest.raises(JobAlreadyRunningError):
|
||||||
|
WorkspaceJobLock(tmp_path, "demo", "evidence").acquire()
|
||||||
|
finally:
|
||||||
|
release.set()
|
||||||
|
process.join(10)
|
||||||
|
assert process.exitcode == 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_evidence_and_dwh_jobs_have_distinct_locks(tmp_path):
|
||||||
|
with WorkspaceJobLock(tmp_path, "demo", "evidence"):
|
||||||
|
with WorkspaceJobLock(tmp_path, "demo", "dwh"):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def test_lock_keys_cannot_escape_lock_directory(tmp_path):
|
||||||
|
with pytest.raises(ValueError, match="filesystem-safe"):
|
||||||
|
WorkspaceJobLock(tmp_path, "demo", "../evidence")
|
||||||
|
|
||||||
|
|
||||||
|
def test_lock_is_recoverable_after_process_crash_without_stale_deletion(tmp_path):
|
||||||
|
context = multiprocessing.get_context("spawn")
|
||||||
|
ready = context.Event()
|
||||||
|
process = context.Process(target=_crash_with_lock, args=(str(tmp_path), ready))
|
||||||
|
process.start()
|
||||||
|
assert ready.wait(10)
|
||||||
|
process.join(10)
|
||||||
|
assert process.exitcode == 7
|
||||||
|
|
||||||
|
lock_path = WorkspaceJobLock(tmp_path, "demo", "evidence").path
|
||||||
|
assert lock_path.exists()
|
||||||
|
with WorkspaceJobLock(tmp_path, "demo", "evidence"):
|
||||||
|
assert lock_path.exists()
|
||||||
@@ -0,0 +1,127 @@
|
|||||||
|
import json
|
||||||
|
import pytest
|
||||||
|
from pydantic import ValidationError
|
||||||
|
|
||||||
|
from tht.jobs.models import JobSpec
|
||||||
|
from tht.jobs.runner import CorruptCheckpointError, run_job
|
||||||
|
|
||||||
|
|
||||||
|
def _spec(tmp_path, **updates):
|
||||||
|
values = {
|
||||||
|
"workspace_id": "demo",
|
||||||
|
"job_type": "evidence",
|
||||||
|
"workspace_root": tmp_path,
|
||||||
|
}
|
||||||
|
values.update(updates)
|
||||||
|
return JobSpec(**values)
|
||||||
|
|
||||||
|
|
||||||
|
def test_job_models_are_immutable(tmp_path):
|
||||||
|
spec = _spec(tmp_path)
|
||||||
|
with pytest.raises(ValidationError, match="Instance is frozen"):
|
||||||
|
spec.job_type = "dwh"
|
||||||
|
|
||||||
|
resumed = spec.with_resume("a" * 32)
|
||||||
|
assert resumed.workspace_root == tmp_path
|
||||||
|
assert resumed.resume_run_id == "a" * 32
|
||||||
|
|
||||||
|
|
||||||
|
def test_failed_stage_is_resumable_and_skips_completed_stage(tmp_path):
|
||||||
|
calls = []
|
||||||
|
|
||||||
|
def discover(context):
|
||||||
|
calls.append(("discover", context.dry_run))
|
||||||
|
|
||||||
|
def acquire(_context):
|
||||||
|
calls.append(("acquire", False))
|
||||||
|
raise RuntimeError("source /customer/alice token=secret unavailable")
|
||||||
|
|
||||||
|
first = run_job(_spec(tmp_path), [discover, acquire])
|
||||||
|
assert first.status == "failed"
|
||||||
|
assert [stage.status for stage in first.stages] == ["succeeded", "failed"]
|
||||||
|
assert first.stages[1].error.model_dump() == {
|
||||||
|
"category": "RuntimeError",
|
||||||
|
"message": "stage execution failed",
|
||||||
|
}
|
||||||
|
|
||||||
|
def acquire(_context):
|
||||||
|
calls.append(("recovered", False))
|
||||||
|
|
||||||
|
second = run_job(
|
||||||
|
_spec(tmp_path, resume_run_id=first.run_id),
|
||||||
|
[discover, acquire],
|
||||||
|
)
|
||||||
|
assert second.resumed_from == first.run_id
|
||||||
|
assert second.status == "succeeded"
|
||||||
|
assert calls == [("discover", False), ("acquire", False), ("recovered", False)]
|
||||||
|
|
||||||
|
|
||||||
|
def test_successful_job_is_idempotently_resumable(tmp_path):
|
||||||
|
calls = []
|
||||||
|
|
||||||
|
def normalize(_context):
|
||||||
|
calls.append("normalize")
|
||||||
|
|
||||||
|
first = run_job(_spec(tmp_path), [normalize])
|
||||||
|
second = run_job(_spec(tmp_path, resume_run_id=first.run_id), [normalize])
|
||||||
|
assert first.status == second.status == "succeeded"
|
||||||
|
assert second.resumed_from == first.run_id
|
||||||
|
assert calls == ["normalize"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_checkpoints_and_report_are_json_safe_and_do_not_disclose_workspace_path(tmp_path):
|
||||||
|
def publish(_context):
|
||||||
|
return {"ignored": "/customer/alice", "password": "secret"}
|
||||||
|
|
||||||
|
report = run_job(_spec(tmp_path), [publish])
|
||||||
|
run_dir = tmp_path / ".tht-jobs" / "evidence" / "runs" / report.run_id
|
||||||
|
checkpoint = json.loads((run_dir / "checkpoint.json").read_text())
|
||||||
|
payload = (run_dir / "report.json").read_text()
|
||||||
|
parsed = json.loads(payload)
|
||||||
|
|
||||||
|
assert checkpoint["status"] == "succeeded"
|
||||||
|
assert parsed["schema_version"] == 1
|
||||||
|
assert parsed["run_id"] == report.run_id
|
||||||
|
assert str(tmp_path) not in payload
|
||||||
|
assert "alice" not in payload
|
||||||
|
assert "secret" not in payload
|
||||||
|
assert parsed["started_at"].endswith("Z")
|
||||||
|
assert parsed["finished_at"].endswith("Z")
|
||||||
|
|
||||||
|
|
||||||
|
def test_dry_run_is_exposed_to_stages_and_report(tmp_path):
|
||||||
|
observed = []
|
||||||
|
|
||||||
|
def plan(context):
|
||||||
|
observed.append(context.dry_run)
|
||||||
|
|
||||||
|
report = run_job(_spec(tmp_path, dry_run=True), [plan])
|
||||||
|
assert observed == [True]
|
||||||
|
assert report.dry_run is True
|
||||||
|
assert report.status == "succeeded"
|
||||||
|
|
||||||
|
|
||||||
|
def test_corrupt_checkpoint_is_rejected_without_running_stages(tmp_path):
|
||||||
|
first = run_job(_spec(tmp_path), [lambda _context: None])
|
||||||
|
checkpoint = (
|
||||||
|
tmp_path / ".tht-jobs" / "evidence" / "runs" / first.run_id / "checkpoint.json"
|
||||||
|
)
|
||||||
|
checkpoint.write_text("{not-json")
|
||||||
|
called = False
|
||||||
|
|
||||||
|
def stage(_context):
|
||||||
|
nonlocal called
|
||||||
|
called = True
|
||||||
|
|
||||||
|
with pytest.raises(CorruptCheckpointError, match="checkpoint is invalid"):
|
||||||
|
run_job(_spec(tmp_path, resume_run_id=first.run_id), [stage])
|
||||||
|
assert called is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_stage_timestamps_are_aware_and_ordered(tmp_path):
|
||||||
|
report = run_job(_spec(tmp_path), [lambda _context: None])
|
||||||
|
stage = report.stages[0]
|
||||||
|
assert stage.started_at.tzinfo is not None
|
||||||
|
assert stage.finished_at.tzinfo is not None
|
||||||
|
assert stage.started_at <= stage.finished_at
|
||||||
|
assert report.started_at <= stage.started_at <= report.finished_at
|
||||||
@@ -0,0 +1,15 @@
|
|||||||
|
"""Shared execution envelope for resumable preprocessing jobs."""
|
||||||
|
|
||||||
|
from tht.jobs.locking import JobAlreadyRunningError, WorkspaceJobLock
|
||||||
|
from tht.jobs.models import JobReport, JobRun, JobSpec
|
||||||
|
from tht.jobs.runner import JobContext, run_job
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"JobAlreadyRunningError",
|
||||||
|
"JobContext",
|
||||||
|
"JobReport",
|
||||||
|
"JobRun",
|
||||||
|
"JobSpec",
|
||||||
|
"WorkspaceJobLock",
|
||||||
|
"run_job",
|
||||||
|
]
|
||||||
@@ -0,0 +1,67 @@
|
|||||||
|
"""Crash-safe interprocess locking scoped by workspace and job type."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import fcntl
|
||||||
|
import hashlib
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
from pathlib import Path
|
||||||
|
from types import TracebackType
|
||||||
|
|
||||||
|
|
||||||
|
class JobAlreadyRunningError(RuntimeError):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
_JOB_KEY = re.compile(r"^[a-z][a-z0-9_-]{0,63}$")
|
||||||
|
|
||||||
|
|
||||||
|
def _lock_name(workspace_id: str, job_type: str) -> str:
|
||||||
|
if not _JOB_KEY.fullmatch(workspace_id) or not _JOB_KEY.fullmatch(job_type):
|
||||||
|
raise ValueError("lock identifiers must be lowercase filesystem-safe keys")
|
||||||
|
workspace_key = hashlib.sha256(workspace_id.encode("utf-8")).hexdigest()[:16]
|
||||||
|
return f"{workspace_key}-{job_type}.lock"
|
||||||
|
|
||||||
|
|
||||||
|
class WorkspaceJobLock:
|
||||||
|
"""Advisory kernel lock; the inode remains stable and is never deleted by PID."""
|
||||||
|
|
||||||
|
def __init__(self, workspace_root: Path, workspace_id: str, job_type: str) -> None:
|
||||||
|
self.path = workspace_root / ".tht-jobs" / ".locks" / _lock_name(
|
||||||
|
workspace_id, job_type
|
||||||
|
)
|
||||||
|
self._fd: int | None = None
|
||||||
|
|
||||||
|
def acquire(self) -> "WorkspaceJobLock":
|
||||||
|
if self._fd is not None:
|
||||||
|
raise RuntimeError("job lock is already held by this object")
|
||||||
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
fd = os.open(self.path, os.O_RDWR | os.O_CREAT, 0o600)
|
||||||
|
try:
|
||||||
|
fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||||
|
except BlockingIOError as error:
|
||||||
|
os.close(fd)
|
||||||
|
raise JobAlreadyRunningError("this workspace job is already running") from error
|
||||||
|
self._fd = fd
|
||||||
|
return self
|
||||||
|
|
||||||
|
def release(self) -> None:
|
||||||
|
if self._fd is None:
|
||||||
|
return
|
||||||
|
fd, self._fd = self._fd, None
|
||||||
|
try:
|
||||||
|
fcntl.flock(fd, fcntl.LOCK_UN)
|
||||||
|
finally:
|
||||||
|
os.close(fd)
|
||||||
|
|
||||||
|
def __enter__(self) -> "WorkspaceJobLock":
|
||||||
|
return self.acquire()
|
||||||
|
|
||||||
|
def __exit__(
|
||||||
|
self,
|
||||||
|
exc_type: type[BaseException] | None,
|
||||||
|
exc: BaseException | None,
|
||||||
|
traceback: TracebackType | None,
|
||||||
|
) -> None:
|
||||||
|
self.release()
|
||||||
@@ -0,0 +1,114 @@
|
|||||||
|
"""Immutable, secret-free records for preprocessing execution."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Literal, Self
|
||||||
|
|
||||||
|
from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator
|
||||||
|
|
||||||
|
|
||||||
|
_JOB_KEY = re.compile(r"^[a-z][a-z0-9_-]{0,63}$")
|
||||||
|
_RUN_ID = re.compile(r"^[0-9a-f]{32}$")
|
||||||
|
JobStatus = Literal["pending", "running", "succeeded", "failed"]
|
||||||
|
StageStatus = Literal["pending", "running", "succeeded", "failed"]
|
||||||
|
|
||||||
|
|
||||||
|
def utc_now() -> datetime:
|
||||||
|
return datetime.now(UTC)
|
||||||
|
|
||||||
|
|
||||||
|
def _validate_job_key(value: str) -> str:
|
||||||
|
if not _JOB_KEY.fullmatch(value):
|
||||||
|
raise ValueError("job identifiers must be lowercase filesystem-safe keys")
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
def _validate_run_id(value: str | None) -> str | None:
|
||||||
|
if value is not None and not _RUN_ID.fullmatch(value):
|
||||||
|
raise ValueError("run id must contain 32 lowercase hexadecimal characters")
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
class _FrozenModel(BaseModel):
|
||||||
|
model_config = ConfigDict(frozen=True, extra="forbid", validate_default=True)
|
||||||
|
|
||||||
|
def model_copy(self, *, update=None, deep: bool = False) -> Self:
|
||||||
|
data = self.model_dump(round_trip=True)
|
||||||
|
if update:
|
||||||
|
data.update(update)
|
||||||
|
return type(self).model_validate(data)
|
||||||
|
|
||||||
|
@field_serializer("*", when_used="json", check_fields=False)
|
||||||
|
def serialize_utc(self, value):
|
||||||
|
if isinstance(value, datetime):
|
||||||
|
return value.astimezone(UTC).isoformat().replace("+00:00", "Z")
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
class JobSpec(_FrozenModel):
|
||||||
|
"""Execution input. The local root is deliberately excluded from serialization."""
|
||||||
|
|
||||||
|
workspace_id: str
|
||||||
|
job_type: str
|
||||||
|
workspace_root: Path = Field(exclude=True)
|
||||||
|
dry_run: bool = False
|
||||||
|
resume_run_id: str | None = None
|
||||||
|
|
||||||
|
_workspace_key = field_validator("workspace_id")(_validate_job_key)
|
||||||
|
_job_type_key = field_validator("job_type")(_validate_job_key)
|
||||||
|
_resume_id = field_validator("resume_run_id")(_validate_run_id)
|
||||||
|
|
||||||
|
def model_copy(self, *, update=None, deep: bool = False) -> Self:
|
||||||
|
data = {
|
||||||
|
"workspace_id": self.workspace_id,
|
||||||
|
"job_type": self.job_type,
|
||||||
|
"workspace_root": self.workspace_root,
|
||||||
|
"dry_run": self.dry_run,
|
||||||
|
"resume_run_id": self.resume_run_id,
|
||||||
|
}
|
||||||
|
if update:
|
||||||
|
data.update(update)
|
||||||
|
return type(self).model_validate(data)
|
||||||
|
|
||||||
|
def with_resume(self, run_id: str) -> "JobSpec":
|
||||||
|
return self.model_copy(update={"resume_run_id": run_id})
|
||||||
|
|
||||||
|
|
||||||
|
class StageError(_FrozenModel):
|
||||||
|
category: str = Field(pattern=r"^[A-Za-z][A-Za-z0-9_]{0,127}$")
|
||||||
|
message: Literal["stage execution failed"] = "stage execution failed"
|
||||||
|
|
||||||
|
|
||||||
|
class StageRun(_FrozenModel):
|
||||||
|
name: str
|
||||||
|
status: StageStatus = "pending"
|
||||||
|
started_at: datetime | None = None
|
||||||
|
finished_at: datetime | None = None
|
||||||
|
error: StageError | None = None
|
||||||
|
|
||||||
|
_name_key = field_validator("name")(_validate_job_key)
|
||||||
|
|
||||||
|
|
||||||
|
class JobRun(_FrozenModel):
|
||||||
|
"""Durable checkpoint, persisted after every state transition."""
|
||||||
|
|
||||||
|
schema_version: Literal[1] = 1
|
||||||
|
run_id: str
|
||||||
|
job_type: str
|
||||||
|
dry_run: bool
|
||||||
|
status: JobStatus
|
||||||
|
started_at: datetime
|
||||||
|
finished_at: datetime | None = None
|
||||||
|
resumed_from: str | None = None
|
||||||
|
stages: tuple[StageRun, ...] = ()
|
||||||
|
|
||||||
|
_run_id = field_validator("run_id")(_validate_run_id)
|
||||||
|
_job_type = field_validator("job_type")(_validate_job_key)
|
||||||
|
_resumed_from = field_validator("resumed_from")(_validate_run_id)
|
||||||
|
|
||||||
|
|
||||||
|
class JobReport(JobRun):
|
||||||
|
"""Public machine-readable terminal report (contains no paths or stage outputs)."""
|
||||||
@@ -0,0 +1,171 @@
|
|||||||
|
"""Resumable stage runner with durable atomic checkpoints and reports."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import uuid
|
||||||
|
from collections.abc import Callable, Sequence
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from pydantic import ValidationError
|
||||||
|
|
||||||
|
from tht.jobs.locking import WorkspaceJobLock
|
||||||
|
from tht.jobs.models import JobReport, JobRun, JobSpec, StageError, StageRun, utc_now
|
||||||
|
|
||||||
|
|
||||||
|
class CorruptCheckpointError(RuntimeError):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class JobContext:
|
||||||
|
run_id: str
|
||||||
|
job_type: str
|
||||||
|
dry_run: bool
|
||||||
|
workspace_root: Path
|
||||||
|
run_dir: Path
|
||||||
|
|
||||||
|
|
||||||
|
Stage = Callable[[JobContext], Any]
|
||||||
|
|
||||||
|
|
||||||
|
def _atomic_write(path: Path, payload: str) -> None:
|
||||||
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
|
||||||
|
fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
||||||
|
try:
|
||||||
|
with os.fdopen(fd, "w", encoding="utf-8") as stream:
|
||||||
|
stream.write(payload)
|
||||||
|
stream.flush()
|
||||||
|
os.fsync(stream.fileno())
|
||||||
|
os.replace(temporary, path)
|
||||||
|
directory_fd = os.open(path.parent, os.O_RDONLY)
|
||||||
|
try:
|
||||||
|
os.fsync(directory_fd)
|
||||||
|
finally:
|
||||||
|
os.close(directory_fd)
|
||||||
|
except BaseException:
|
||||||
|
try:
|
||||||
|
temporary.unlink()
|
||||||
|
except FileNotFoundError:
|
||||||
|
pass
|
||||||
|
raise
|
||||||
|
|
||||||
|
|
||||||
|
def _persist(path: Path, run: JobRun) -> None:
|
||||||
|
_atomic_write(path, run.model_dump_json(indent=2) + "\n")
|
||||||
|
|
||||||
|
|
||||||
|
def _load_checkpoint(path: Path) -> JobRun:
|
||||||
|
try:
|
||||||
|
return JobRun.model_validate_json(path.read_text(encoding="utf-8"))
|
||||||
|
except (OSError, ValidationError, ValueError, json.JSONDecodeError) as error:
|
||||||
|
raise CorruptCheckpointError("checkpoint is invalid and cannot be resumed") from error
|
||||||
|
|
||||||
|
|
||||||
|
def _stage_name(stage: Stage) -> str:
|
||||||
|
name = getattr(stage, "__name__", "")
|
||||||
|
if name == "<lambda>":
|
||||||
|
name = "stage"
|
||||||
|
return name.replace("_", "-")
|
||||||
|
|
||||||
|
|
||||||
|
def _new_run(spec: JobSpec, run_id: str, stages: Sequence[Stage]) -> JobRun:
|
||||||
|
names = [_stage_name(stage) for stage in stages]
|
||||||
|
if len(names) != len(set(names)):
|
||||||
|
raise ValueError("stage names must be unique")
|
||||||
|
return JobRun(
|
||||||
|
run_id=run_id,
|
||||||
|
job_type=spec.job_type,
|
||||||
|
dry_run=spec.dry_run,
|
||||||
|
status="running",
|
||||||
|
started_at=utc_now(),
|
||||||
|
resumed_from=spec.resume_run_id,
|
||||||
|
stages=tuple(StageRun(name=name) for name in names),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _resume_run(spec: JobSpec, run_id: str, stages: Sequence[Stage], source: JobRun) -> JobRun:
|
||||||
|
if source.job_type != spec.job_type:
|
||||||
|
raise CorruptCheckpointError("checkpoint job type does not match resume request")
|
||||||
|
requested_names = [_stage_name(stage) for stage in stages]
|
||||||
|
source_by_name = {stage.name: stage for stage in source.stages}
|
||||||
|
resumed_stages = []
|
||||||
|
for name in requested_names:
|
||||||
|
previous = source_by_name.get(name)
|
||||||
|
resumed_stages.append(
|
||||||
|
previous
|
||||||
|
if previous is not None and previous.status == "succeeded"
|
||||||
|
else StageRun(name=name)
|
||||||
|
)
|
||||||
|
return JobRun(
|
||||||
|
run_id=run_id,
|
||||||
|
job_type=spec.job_type,
|
||||||
|
dry_run=spec.dry_run,
|
||||||
|
status="running",
|
||||||
|
started_at=utc_now(),
|
||||||
|
resumed_from=source.run_id,
|
||||||
|
stages=tuple(resumed_stages),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def run_job(spec: JobSpec, stages: Sequence[Stage]) -> JobReport:
|
||||||
|
"""Run stages once, returning a terminal report instead of leaking stage exceptions."""
|
||||||
|
with WorkspaceJobLock(spec.workspace_root, spec.workspace_id, spec.job_type):
|
||||||
|
jobs_root = spec.workspace_root / ".tht-jobs" / spec.job_type / "runs"
|
||||||
|
run_id = uuid.uuid4().hex
|
||||||
|
run_dir = jobs_root / run_id
|
||||||
|
checkpoint_path = run_dir / "checkpoint.json"
|
||||||
|
if spec.resume_run_id is None:
|
||||||
|
run = _new_run(spec, run_id, stages)
|
||||||
|
else:
|
||||||
|
source_path = jobs_root / spec.resume_run_id / "checkpoint.json"
|
||||||
|
source = _load_checkpoint(source_path)
|
||||||
|
run = _resume_run(spec, run_id, stages, source)
|
||||||
|
_persist(checkpoint_path, run)
|
||||||
|
context = JobContext(run_id, spec.job_type, spec.dry_run, spec.workspace_root, run_dir)
|
||||||
|
|
||||||
|
for index, stage_callable in enumerate(stages):
|
||||||
|
if run.stages[index].status == "succeeded":
|
||||||
|
continue
|
||||||
|
stage = run.stages[index].model_copy(
|
||||||
|
update={"status": "running", "started_at": utc_now()}
|
||||||
|
)
|
||||||
|
run = run.model_copy(
|
||||||
|
update={"stages": run.stages[:index] + (stage,) + run.stages[index + 1 :]}
|
||||||
|
)
|
||||||
|
_persist(checkpoint_path, run)
|
||||||
|
try:
|
||||||
|
stage_callable(context)
|
||||||
|
except Exception as error:
|
||||||
|
failed = stage.model_copy(
|
||||||
|
update={
|
||||||
|
"status": "failed",
|
||||||
|
"finished_at": utc_now(),
|
||||||
|
"error": StageError(category=type(error).__name__),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
run = run.model_copy(
|
||||||
|
update={
|
||||||
|
"status": "failed",
|
||||||
|
"finished_at": utc_now(),
|
||||||
|
"stages": run.stages[:index] + (failed,) + run.stages[index + 1 :],
|
||||||
|
}
|
||||||
|
)
|
||||||
|
_persist(checkpoint_path, run)
|
||||||
|
break
|
||||||
|
succeeded = stage.model_copy(update={"status": "succeeded", "finished_at": utc_now()})
|
||||||
|
run = run.model_copy(
|
||||||
|
update={"stages": run.stages[:index] + (succeeded,) + run.stages[index + 1 :]}
|
||||||
|
)
|
||||||
|
_persist(checkpoint_path, run)
|
||||||
|
else:
|
||||||
|
run = run.model_copy(update={"status": "succeeded", "finished_at": utc_now()})
|
||||||
|
_persist(checkpoint_path, run)
|
||||||
|
|
||||||
|
report = JobReport.model_validate(run.model_dump())
|
||||||
|
_atomic_write(run_dir / "report.json", report.model_dump_json(indent=2) + "\n")
|
||||||
|
return report
|
||||||
Reference in New Issue
Block a user