feat(preprocess): resume evidence jobs by run id
This commit is contained in:
@@ -6,6 +6,7 @@ import hashlib
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from tht.corpus.chunk import ChunkPolicy, chunk
|
||||
from tht.corpus.models import CanonicalChunk, CanonicalDocument, CorpusManifest
|
||||
@@ -14,6 +15,19 @@ from tht.corpus.store import CorpusStore
|
||||
from tht.ports.evidence import EvidenceSource, SourceObject
|
||||
from tht.ports.vector import VectorStore, VectorWriteRecord
|
||||
from tht.vectorstore.records import VectorRecord
|
||||
from tht.jobs.models import JobSpec
|
||||
from tht.jobs.runner import JobContext, run_job
|
||||
|
||||
|
||||
EVIDENCE_STAGE_IDS = (
|
||||
"discover",
|
||||
"acquire_normalize_chunk",
|
||||
"embed",
|
||||
"vector_upsert",
|
||||
"stage_validate",
|
||||
"publish",
|
||||
"retention_cleanup",
|
||||
)
|
||||
|
||||
|
||||
class PipelineError(RuntimeError):
|
||||
@@ -29,6 +43,8 @@ class PipelineResult:
|
||||
unchanged: tuple[str, ...]
|
||||
removed: tuple[str, ...]
|
||||
manifest: CorpusManifest
|
||||
run_id: str | None = None
|
||||
resumed_from: str | None = None
|
||||
|
||||
def model_dump(self, mode=None):
|
||||
value = asdict(self)
|
||||
@@ -71,6 +87,247 @@ class CorpusPipeline:
|
||||
with self.store.writer_lock():
|
||||
return self._run(dry_run=dry_run, resume=resume)
|
||||
|
||||
def run_as_job(
|
||||
self,
|
||||
*,
|
||||
workspace_id: str,
|
||||
workspace_root: Path,
|
||||
config_fingerprint: str,
|
||||
input_fingerprint: str,
|
||||
dry_run: bool = False,
|
||||
resume_run_id: str | None = None,
|
||||
) -> PipelineResult:
|
||||
"""Execute preprocessing through the durable shared job envelope."""
|
||||
discovered = self._discover()
|
||||
discovered_fingerprint = _fingerprint(
|
||||
{item.source_id: item.fingerprint for _, item in discovered}
|
||||
)
|
||||
source_by_id = {item.source_id: (source, item) for source, item in discovered}
|
||||
compatibility = _fingerprint({
|
||||
"pipeline": self.pipeline_version,
|
||||
"model": self.embedding_model,
|
||||
"dimensions": self.embedding_dimensions,
|
||||
"chunk_policy": asdict(self.chunk_policy),
|
||||
})
|
||||
spec = JobSpec(
|
||||
workspace_id=workspace_id,
|
||||
job_type="evidence",
|
||||
workspace_root=workspace_root,
|
||||
spec_version="jobs-v1",
|
||||
pipeline_version=self.pipeline_version,
|
||||
config_fingerprint=config_fingerprint,
|
||||
input_fingerprint=_fingerprint([input_fingerprint, discovered_fingerprint]),
|
||||
stage_ids=EVIDENCE_STAGE_IDS,
|
||||
dry_run=dry_run,
|
||||
resume_run_id=resume_run_id,
|
||||
)
|
||||
|
||||
def artifact(context: JobContext, name: str) -> Path:
|
||||
root = context.run_dir / "artifacts"
|
||||
root.mkdir(exist_ok=True)
|
||||
return root / name
|
||||
|
||||
def write(context: JobContext, name: str, value) -> None:
|
||||
artifact(context, name).write_text(
|
||||
json.dumps(value, sort_keys=True, separators=(",", ":")), encoding="utf-8"
|
||||
)
|
||||
|
||||
def read(context: JobContext, name: str):
|
||||
try:
|
||||
return json.loads(artifact(context, name).read_text(encoding="utf-8"))
|
||||
except (OSError, ValueError) as error:
|
||||
raise PipelineError("preprocessing checkpoint artifact is corrupt") from error
|
||||
|
||||
def discover_stage(context: JobContext) -> None:
|
||||
previous = self.store.active_manifest()
|
||||
prior = {doc.source_id: doc for doc in previous.documents} if previous else {}
|
||||
fingerprints = {item.source_id: item.fingerprint for _, item in discovered}
|
||||
rebuild = bool(previous and previous.metadata.get("compatibility_fingerprint") != compatibility)
|
||||
changed = sorted(
|
||||
item.source_id for _, item in discovered
|
||||
if rebuild or item.source_id not in prior
|
||||
or prior[item.source_id].source_fingerprint != item.fingerprint
|
||||
)
|
||||
unchanged = sorted(set(fingerprints) - set(changed))
|
||||
removed = sorted(set(prior) - set(fingerprints))
|
||||
write(context, "plan.json", {
|
||||
"generation": f"gen:{context.run_id}",
|
||||
"compatibility": compatibility,
|
||||
"fingerprints": fingerprints,
|
||||
"changed": changed,
|
||||
"unchanged": unchanged,
|
||||
"removed": removed,
|
||||
"previous": previous.model_dump(mode="json") if previous else None,
|
||||
})
|
||||
|
||||
def acquire_stage(context: JobContext) -> None:
|
||||
if context.dry_run:
|
||||
return
|
||||
plan = read(context, "plan.json")
|
||||
previous = CorpusManifest.model_validate(plan["previous"]) if plan["previous"] else None
|
||||
prior = {doc.source_id: doc for doc in previous.documents} if previous else {}
|
||||
documents = [prior[source_id] for source_id in plan["unchanged"]]
|
||||
for source_id in plan["changed"]:
|
||||
source, item = source_by_id[source_id]
|
||||
documents.append(normalize(source.acquire(item), self.pipeline_version))
|
||||
documents.sort(key=lambda value: value.source_id)
|
||||
chunks = [part for document in documents for part in chunk(document, self.chunk_policy)]
|
||||
previous_generations = dict(previous.metadata.get("document_generations", {})) if previous else {}
|
||||
changed = set(plan["changed"])
|
||||
generations = {
|
||||
document.document_id: (
|
||||
plan["generation"] if document.source_id in changed
|
||||
else previous_generations.get(document.document_id, previous.vector_generation)
|
||||
) for document in documents
|
||||
}
|
||||
manifest = CorpusManifest(
|
||||
pipeline_version=self.pipeline_version,
|
||||
embedding_model=self.embedding_model,
|
||||
embedding_dimensions=self.embedding_dimensions,
|
||||
vector_generation=plan["generation"],
|
||||
documents=tuple(documents), chunks=tuple(chunks),
|
||||
metadata={
|
||||
"compatibility_fingerprint": compatibility,
|
||||
"fingerprints": plan["fingerprints"],
|
||||
"removed": plan["removed"],
|
||||
"document_generations": generations,
|
||||
},
|
||||
)
|
||||
write(context, "manifest.json", manifest.model_dump(mode="json"))
|
||||
|
||||
def embed_stage(context: JobContext) -> None:
|
||||
if context.dry_run:
|
||||
return
|
||||
plan = read(context, "plan.json")
|
||||
manifest = CorpusManifest.model_validate(read(context, "manifest.json"))
|
||||
changed_docs = {doc.document_id for doc in manifest.documents if doc.source_id in plan["changed"]}
|
||||
parts = [part for part in manifest.chunks if part.document_id in changed_docs]
|
||||
embeddings = self.embedder.embed_documents([part.content for part in parts])
|
||||
if len(embeddings) != len(parts) or any(
|
||||
len(vector) != self.embedding_dimensions for vector in embeddings
|
||||
):
|
||||
raise PipelineError("embedding output is incompatible")
|
||||
write(context, "embeddings.json", embeddings)
|
||||
|
||||
def records(context: JobContext):
|
||||
plan = read(context, "plan.json")
|
||||
manifest = CorpusManifest.model_validate(read(context, "manifest.json"))
|
||||
changed_docs = {doc.document_id for doc in manifest.documents if doc.source_id in plan["changed"]}
|
||||
parts = [part for part in manifest.chunks if part.document_id in changed_docs]
|
||||
embeddings = read(context, "embeddings.json")
|
||||
return [self._vector_record(part, vector, plan["generation"])
|
||||
for part, vector in zip(parts, embeddings, strict=True)]
|
||||
|
||||
def compensate(context: JobContext) -> None:
|
||||
generation = read(context, "plan.json")["generation"]
|
||||
self.store.discard(generation)
|
||||
try:
|
||||
self.vector_store.delete_generation("evidence", generation)
|
||||
except Exception:
|
||||
pass
|
||||
write(context, "compensated.json", {"generation": generation})
|
||||
|
||||
def rotate_compensated_generation(context: JobContext) -> None:
|
||||
marker = artifact(context, "compensated.json")
|
||||
if not marker.exists():
|
||||
return
|
||||
plan = read(context, "plan.json")
|
||||
old = plan["generation"]
|
||||
plan["generation"] = f"gen:{uuid.uuid4().hex}"
|
||||
write(context, "plan.json", plan)
|
||||
manifest = CorpusManifest.model_validate(read(context, "manifest.json"))
|
||||
changed = set(plan["changed"])
|
||||
generations = dict(manifest.metadata["document_generations"])
|
||||
for document in manifest.documents:
|
||||
if document.source_id in changed and generations.get(document.document_id) == old:
|
||||
generations[document.document_id] = plan["generation"]
|
||||
metadata = dict(manifest.metadata)
|
||||
metadata["document_generations"] = generations
|
||||
manifest = manifest.model_copy(update={
|
||||
"vector_generation": plan["generation"], "metadata": metadata,
|
||||
})
|
||||
write(context, "manifest.json", manifest.model_dump(mode="json"))
|
||||
marker.unlink()
|
||||
|
||||
def vector_stage(context: JobContext) -> None:
|
||||
if context.dry_run:
|
||||
return
|
||||
rotate_compensated_generation(context)
|
||||
values = records(context)
|
||||
try:
|
||||
if values and self.vector_store.upsert("evidence", values) != len(values):
|
||||
raise PipelineError("vector write count mismatch")
|
||||
except Exception:
|
||||
compensate(context)
|
||||
raise
|
||||
|
||||
def stage_stage(context: JobContext) -> None:
|
||||
if context.dry_run:
|
||||
return
|
||||
plan = read(context, "plan.json")
|
||||
manifest = CorpusManifest.model_validate(read(context, "manifest.json"))
|
||||
try:
|
||||
self.store.stage(
|
||||
manifest, {doc.document_id: doc.content for doc in manifest.documents},
|
||||
generation=plan["generation"],
|
||||
)
|
||||
self.store.manifest(plan["generation"])
|
||||
except Exception:
|
||||
compensate(context)
|
||||
raise
|
||||
|
||||
def publish_stage(context: JobContext) -> None:
|
||||
if context.dry_run:
|
||||
return
|
||||
if artifact(context, "compensated.json").exists():
|
||||
rotate_compensated_generation(context)
|
||||
values = records(context)
|
||||
if values and self.vector_store.upsert("evidence", values) != len(values):
|
||||
compensate(context)
|
||||
raise PipelineError("vector write count mismatch")
|
||||
manifest = CorpusManifest.model_validate(read(context, "manifest.json"))
|
||||
generation = read(context, "plan.json")["generation"]
|
||||
self.store.stage(
|
||||
manifest, {doc.document_id: doc.content for doc in manifest.documents},
|
||||
generation=generation,
|
||||
)
|
||||
generation = read(context, "plan.json")["generation"]
|
||||
try:
|
||||
self.store.publish(generation)
|
||||
except Exception:
|
||||
compensate(context)
|
||||
raise
|
||||
|
||||
def retention_stage(context: JobContext) -> None:
|
||||
# Retention policy is intentionally a stable no-op until configured.
|
||||
return
|
||||
|
||||
report = run_job(spec, [
|
||||
discover_stage, acquire_stage, embed_stage, vector_stage,
|
||||
stage_stage, publish_stage, retention_stage,
|
||||
])
|
||||
run_dir = workspace_root / ".tht-jobs" / "evidence" / "runs" / report.run_id
|
||||
plan = json.loads((run_dir / "artifacts" / "plan.json").read_text())
|
||||
if dry_run:
|
||||
manifest = self.store.active_manifest() or CorpusManifest(pipeline_version=self.pipeline_version)
|
||||
generation = None
|
||||
published = False
|
||||
elif report.status == "succeeded":
|
||||
generation = plan["generation"]
|
||||
manifest = self.store.manifest(generation)
|
||||
published = True
|
||||
else:
|
||||
generation = plan["generation"]
|
||||
manifest_path = run_dir / "artifacts" / "manifest.json"
|
||||
manifest = (CorpusManifest.model_validate_json(manifest_path.read_text())
|
||||
if manifest_path.exists() else CorpusManifest(pipeline_version=self.pipeline_version))
|
||||
published = False
|
||||
return PipelineResult(
|
||||
report.status, generation, published, tuple(plan["changed"]),
|
||||
tuple(plan["unchanged"]), tuple(plan["removed"]), manifest,
|
||||
report.run_id, report.resumed_from,
|
||||
)
|
||||
|
||||
def _run(self, *, dry_run: bool = False, resume: str | None = None) -> PipelineResult:
|
||||
generation = None
|
||||
vector_written = False
|
||||
|
||||
Reference in New Issue
Block a user