Approach (brief): capture minimal metadata at training runtime, persist it with the artifact (model file + manifest.json), surface metadata in logs, and add a CI/PR check that validates every produced artifact includes required fields (git-sha, data-version, hyperparams, metrics). Use existing tracking libs (MLflow/W&B) if available, but keep a simple manifest file so artifacts are self-describing and reproducible.Minimal code changes (python example):- add a small helper to collect metadata (reads git sha, dataset checksum), record hyperparams/metrics, and write manifest.json next to the saved model.python
import json, subprocess, hashlib, os
from datetime import datetime
def git_sha():
return subprocess.check_output(["git","rev-parse","HEAD"]).decode().strip()
def file_checksum(path):
h = hashlib.sha256()
with open(path,"rb") as f:
for chunk in iter(lambda: f.read(8192), b""):
h.update(chunk)
return h.hexdigest()
def save_model_with_manifest(model_bytes, out_dir, hyperparams, metrics, dataset_path, extra=None):
os.makedirs(out_dir, exist_ok=True)
model_path = os.path.join(out_dir, "model.pt")
with open(model_path,"wb") as f:
f.write(model_bytes) # e.g., torch.save bytes or file write
manifest = {
"git_sha": git_sha(),
"timestamp": datetime.utcnow().isoformat()+"Z",
"hyperparams": hyperparams,
"metrics": metrics,
"data_version": file_checksum(dataset_path),
"dataset_path": os.path.basename(dataset_path),
}
if extra: manifest.update(extra)
with open(os.path.join(out_dir,"manifest.json"), "w") as f:
json.dump(manifest, f, indent=2)
print("Saved model and manifest:", model_path)
How to surface metadata:- Print key fields to training logs (git_sha, data_version, major metrics) so CI/log ingestion can index them.- Persist manifest.json alongside model artifact in artifact store (S3, mlflow artifacts, W&B artifact).- Include environment spec (requirements.txt / conda.yaml) and training script path to guarantee runnable reproduction.Enforcement via CI / PR checks:- Add a lightweight validation step (run on artifact publish or as pre-merge on PRs that change training code/data) that verifies: - manifest.json exists and contains required keys (git_sha, data_version, hyperparams, metrics). - git_sha in manifest matches the commit that produced the artifact (if artifact produced in the pipeline, capture CURRENT_SHA env). - data_version matches a trusted dataset registry or checksum.- Example validator script (run in CI):python
# validate_manifest.py
import json, sys
m = json.load(open(sys.argv[1]))
required = ["git_sha","data_version","hyperparams","metrics"]
missing = [k for k in required if k not in m]
if missing:
print("Missing metadata fields:", missing); sys.exit(1)
# Optionally compare m["git_sha"] to env var CI_COMMIT
- Wire into CI (GitHub Actions/Gitlab) to run on PRs and on artifact publish. Fail builds if validation fails.Best practices / trade-offs:- Prefer recording git SHA and training script path; for reproducibility, build and publish a Docker image or wheel tied to that SHA.- Record dataset checksums or an immutable dataset version id from your data registry.- Use experiment trackers (MLflow/W&B) to get search/indexing and UI, but always keep local manifest.json so artifacts remain self-describing even outside tracker.- Keep manifest small, stable schema, and version the manifest format (manifest_version) to support future changes.This pattern ensures artifacts are coupled to the exact code and data used, are easy to inspect, and can be enforced automatically in CI/PR pipelines.