Files
ModelForge/scripts/m15_record_artifact_recovery.py

117 lines
4.8 KiB
Python

"""Record observed artifact-recovery evidence against a planned recovery operation.
The rehydration itself is performed by the existing acquisition plane: an exact-revision download
plan is executed by the Node Agent, which quarantines, hashes and only then promotes. This script
reads the *observed* result out of the control-plane database and journals it on the M15 artifact
recovery operation, so the recovery record carries measured evidence rather than an assumption.
python scripts/m15_record_artifact_recovery.py \
--operation-id <uuid> --storage-root-id <uuid> --artifact-job-id <uuid> \
--download-plan-id <uuid> --database-url postgresql+psycopg://...
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "backend" / "src"))
from modelforge_api.domain.recovery import ( # noqa: E402
ArtifactRecoveryState,
RecoveryFailureCode,
)
from modelforge_api.persistence.models import ( # noqa: E402
ArtifactJob,
ArtifactLocation,
ModelArtifact,
)
from modelforge_api.services.recovery import RecoveryService # noqa: E402
from modelforge_api.settings import Settings # noqa: E402
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--operation-id", required=True)
parser.add_argument("--storage-root-id", required=True)
parser.add_argument("--artifact-job-id", default=None)
parser.add_argument("--download-plan-id", default=None)
parser.add_argument("--database-url", required=True)
parser.add_argument(
"--blocked-reason",
default=None,
help="record an upstream outage as ARTIFACT_REHYDRATION_BLOCKED instead of a recovery",
)
args = parser.parse_args()
engine = create_engine(args.database_url, pool_pre_ping=True)
try:
with Session(engine) as session:
service = RecoveryService(
session, Settings(database_url=args.database_url), "operator", "m15-rehearsal"
)
if args.blocked_reason:
response = service.record_artifact_recovery(
__import__("uuid").UUID(args.operation_id),
state=ArtifactRecoveryState.BLOCKED,
verified_files=[],
bytes_recovered=0,
duration_seconds=None,
failure_code=RecoveryFailureCode.ARTIFACT_REHYDRATION_BLOCKED.value,
failure_reason=args.blocked_reason,
)
print(json.dumps(json.loads(response.model_dump_json()), indent=2))
return 0
rows = session.execute(
select(ModelArtifact, ArtifactLocation)
.join(ArtifactLocation, ArtifactLocation.artifact_id == ModelArtifact.id)
.where(ArtifactLocation.storage_root_id == args.storage_root_id)
.order_by(ModelArtifact.filename)
).all()
verified = [
{
"filename": artifact.filename,
"sha256": location.observed_sha256,
"expected_sha256": artifact.sha256,
"size_bytes": location.size_bytes,
"relative_path": location.relative_path,
"status": location.status,
"security_status": artifact.security_status,
}
for artifact, location in rows
]
duration = None
if args.artifact_job_id:
job = session.get(ArtifactJob, __import__("uuid").UUID(args.artifact_job_id))
if job and job.completed_at and job.started_at:
duration = (job.completed_at - job.started_at).total_seconds()
response = service.record_artifact_recovery(
__import__("uuid").UUID(args.operation_id),
state=ArtifactRecoveryState.RECOVERED,
verified_files=verified,
bytes_recovered=sum(int(item["size_bytes"] or 0) for item in verified),
duration_seconds=duration,
download_plan_id=(
__import__("uuid").UUID(args.download_plan_id)
if args.download_plan_id
else None
),
artifact_job_id=(
__import__("uuid").UUID(args.artifact_job_id) if args.artifact_job_id else None
),
)
print(json.dumps(json.loads(response.model_dump_json()), indent=2))
return 0 if response.state is ArtifactRecoveryState.RECOVERED else 1
finally:
engine.dispose()
if __name__ == "__main__":
raise SystemExit(main())