diff --git a/mltrail_verifiable_provenance_ledger_for/__init__.py b/mltrail_verifiable_provenance_ledger_for/__init__.py index c2d4481..8d5ead0 100644 --- a/mltrail_verifiable_provenance_ledger_for/__init__.py +++ b/mltrail_verifiable_provenance_ledger_for/__init__.py @@ -20,6 +20,7 @@ from .adapters import MLFlowAdapter, WandBAdapter, Adapter from .registry import ContractRegistry from .util import environment_hash from .governance import DID, GovernanceLogEntry +from .reprobundle import ReproBundle __all__ = [ "Ledger", @@ -39,4 +40,5 @@ __all__ = [ "environment_hash", "DID", "GovernanceLogEntry", + "ReproBundle", ] diff --git a/mltrail_verifiable_provenance_ledger_for/reprobundle.py b/mltrail_verifiable_provenance_ledger_for/reprobundle.py new file mode 100644 index 0000000..0760be8 --- /dev/null +++ b/mltrail_verifiable_provenance_ledger_for/reprobundle.py @@ -0,0 +1,43 @@ +from dataclasses import dataclass, asdict +from typing import List, Dict, Any, Optional +import json +import hashlib + +from .contracts import Environment + + +@dataclass +class ReproBundle: + """A compact, content-addressed snapshot describing everything needed to + deterministically replay a run (code commit, environment, dataset refs, + model ref, and a minimal run manifest). + + This is intentionally small and deterministic so it can be cheaply + exchanged between peers and used to compute Merkle-like bundle hashes. + """ + code_commit: str + environment: Environment + dataset_refs: List[Dict[str, Any]] + model_ref: Dict[str, Any] + run_manifest: Dict[str, Any] + container_hash: Optional[str] = None + + def to_dict(self) -> Dict[str, Any]: + # Convert to a canonical dict suitable for hashing/serialization. + return { + "code_commit": self.code_commit, + "environment": self.environment.to_dict(), + "container_hash": self.container_hash, + "dataset_refs": sorted(self.dataset_refs, key=lambda d: d.get("id") or ""), + "model_ref": self.model_ref, + "run_manifest": self.run_manifest, + } + + def compute_bundle_hash(self) -> str: + """Compute a deterministic SHA-256 digest over the canonical bundle + representation. Uses JSON with sorted keys so identical bundles always + produce the same digest. + """ + payload = self.to_dict() + serialized = json.dumps(payload, sort_keys=True, default=str).encode("utf-8") + return hashlib.sha256(serialized).hexdigest() diff --git a/tests/test_reprobundle.py b/tests/test_reprobundle.py new file mode 100644 index 0000000..c014d5c --- /dev/null +++ b/tests/test_reprobundle.py @@ -0,0 +1,54 @@ +from mltrail_verifiable_provenance_ledger_for.reprobundle import ReproBundle +from mltrail_verifiable_provenance_ledger_for.contracts import Environment + + +def make_sample_environment(): + return Environment( + id="env1", + language="python", + version="3.9", + dependencies={"numpy": "1.24.0", "pandas": "1.5.3"}, + container_hash=None, + ) + + +def test_bundle_hash_deterministic(): + env = make_sample_environment() + bundle1 = ReproBundle( + code_commit="deadbeef", + environment=env, + dataset_refs=[{"id": "ds1", "fingerprint": "fp1"}], + model_ref={"id": "m1", "fingerprint": "mf1"}, + run_manifest={"seed": 42, "cmd": "python train.py"}, + ) + bundle2 = ReproBundle( + code_commit="deadbeef", + environment=make_sample_environment(), + dataset_refs=[{"id": "ds1", "fingerprint": "fp1"}], + model_ref={"id": "m1", "fingerprint": "mf1"}, + run_manifest={"seed": 42, "cmd": "python train.py"}, + ) + + h1 = bundle1.compute_bundle_hash() + h2 = bundle2.compute_bundle_hash() + assert h1 == h2 + + +def test_bundle_hash_changes_on_difference(): + env = make_sample_environment() + base = ReproBundle( + code_commit="deadbeef", + environment=env, + dataset_refs=[{"id": "ds1", "fingerprint": "fp1"}], + model_ref={"id": "m1", "fingerprint": "mf1"}, + run_manifest={"seed": 42, "cmd": "python train.py"}, + ) + modified = ReproBundle( + code_commit="cafebabe", # different commit + environment=env, + dataset_refs=[{"id": "ds1", "fingerprint": "fp1"}], + model_ref={"id": "m1", "fingerprint": "mf1"}, + run_manifest={"seed": 42, "cmd": "python train.py"}, + ) + + assert base.compute_bundle_hash() != modified.compute_bundle_hash()