"""manifest.py — Pydantic models for run manifest, milestone tracking, and pipeline context."""
from __future__ import annotations
import json
import os
import shutil
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Literal, Optional
from pydantic import BaseModel, Field
from .config import SimConfig
from .utils import init_logger
OutputKind = Literal[
"visibility",
"image_product",
"plot",
"log",
"manifest",
"weblog",
"sky_model",
]
[docs]
class Milestone(BaseModel):
"""single checkpoint in a simulation run."""
name: str
status: Literal["started", "completed", "failed"]
timestamp_utc: datetime
elapsed_s: Optional[float] = None
details: dict = Field(default_factory=dict)
[docs]
class OutputRecord(BaseModel):
"""One output produced by a run."""
kind: OutputKind
path: str
image_product_id: Optional[str] = None
imager: Optional[str] = None
role: Optional[str] = None
metadata: dict = Field(default_factory=dict)
[docs]
class RunManifest(BaseModel):
"""canonical machine-readable record of one simulation run."""
run_id: str
status: Literal["running", "completed", "failed"] = "running"
started_at: datetime
completed_at: Optional[datetime] = None
config: SimConfig
invocation: Optional[list[str]] = None
milestones: list[Milestone] = Field(default_factory=list)
outputs: list[OutputRecord] = Field(default_factory=list)
errors: list[str] = Field(default_factory=list)
[docs]
def add_milestone(
self,
name: str,
status: Literal["started", "completed", "failed"],
elapsed_s: Optional[float] = None,
details: Optional[dict] = None,
) -> Milestone:
"""append a milestone and return it."""
ms = Milestone(
name=name,
status=status,
timestamp_utc=datetime.now(timezone.utc),
elapsed_s=elapsed_s,
details=details or {},
)
self.milestones.append(ms)
return ms
[docs]
def add_output(
self,
kind: OutputKind,
path: str,
image_product_id: Optional[str] = None,
imager: Optional[str] = None,
role: Optional[str] = None,
metadata: Optional[dict] = None,
) -> OutputRecord:
"""append a structured output record and return it."""
output = OutputRecord(
kind=kind,
path=path,
image_product_id=image_product_id,
imager=imager,
role=role,
metadata=metadata or {},
)
self.outputs.append(output)
return output
[docs]
def mark_completed(self) -> None:
"""mark the run as completed."""
self.status = "completed"
self.completed_at = datetime.now(timezone.utc)
[docs]
def mark_failed(self, error: str) -> None:
"""mark the run as failed and record the error."""
self.status = "failed"
self.completed_at = datetime.now(timezone.utc)
self.errors.append(error)
[docs]
def model_dump_json(self, **kwargs) -> str:
"""serialize to pretty-printed JSON string."""
return json.dumps(self.model_dump(mode="json"), indent=2, default=str)
[docs]
class RunContext(BaseModel):
"""passed through all pipeline functions; bundles config, paths, and manifest."""
config: SimConfig
work_dir: Path
manifest: RunManifest
visibility_path: Path
log_path: Path
manifest_path: Path
weblog_path: Path
sky_file_resolved: Optional[Path] = None
[docs]
def save_manifest(self) -> None:
"""write the current manifest state to disk (overwrites)."""
self.manifest_path.write_text(self.manifest.model_dump_json(), encoding="utf-8")
[docs]
def add_milestone(self, *args, **kwargs) -> Milestone:
"""convenience: add milestone to manifest and persist to disk."""
ms = self.manifest.add_milestone(*args, **kwargs)
self.save_manifest()
return ms
[docs]
def create_run_context(config: SimConfig) -> RunContext:
"""create work_dir, init logger, build RunContext with empty manifest."""
if config.output_dir is not None:
work_dir = Path(config.output_dir).resolve()
run_id = work_dir.name
else:
run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
run_id = f"{run_id}_{config.telescope.replace('-', '_')}"
work_dir = Path(run_id).resolve()
if config.overwrite and work_dir.exists():
shutil.rmtree(work_dir)
work_dir.mkdir(parents=True, exist_ok=True)
log_file = str(work_dir / f"{work_dir.name}.log")
init_logger(log_file)
manifest = RunManifest(
run_id=run_id,
started_at=datetime.now(timezone.utc),
config=config,
invocation=list(sys.argv) if sys.argv else None,
)
ctx = RunContext(
config=config,
work_dir=work_dir,
manifest=manifest,
visibility_path=work_dir / "visibilities.MS",
log_path=Path(log_file),
manifest_path=work_dir / "run_manifest.json",
weblog_path=work_dir / "weblog.html",
sky_file_resolved=None,
)
ctx.manifest.add_output("log", ctx.log_path.name)
ctx.manifest.add_output("manifest", ctx.manifest_path.name)
if config.sky_file is not None:
fpath = config.sky_file
if not os.path.isabs(fpath):
fpath = os.path.join(os.getcwd(), fpath)
ctx.sky_file_resolved = Path(fpath).resolve()
ctx.save_manifest()
return ctx