"""Versioned checkpoint provenance, lineage, and resume-state truth tests.""" from __future__ import annotations import json from copy import deepcopy from hashlib import sha256 from pathlib import Path import pytest from jsonschema import Draft202012Validator from obliteratus.checkpoint_provenance import ( AdapterIdentity, ArtifactIdentity, DatasetIdentity, LineageEvent, ProvenanceRecord, ToolIdentity, TrainingIdentity, build_provenance, classify_resume_state, migrate_legacy_metadata, sanitize_command, verify_provenance_record, ) from obliteratus.run_archive import RunArchive ROOT = Path(__file__).resolve().parents[1] SCHEMA = json.loads( (ROOT / "docs/checkpoints/schemas/artifact-provenance-v1.schema.json").read_text() ) DIGEST_A = "sha256:" + "a" * 64 DIGEST_B = "sha256:" + "b" * 64 COMMIT = "c" * 40 def _record(*, output_digests=(DIGEST_B,), command=("checkpoint", "convert")): base = ArtifactIdentity("hub", "org/base", "0123456789abcdef", DIGEST_A) return build_provenance( sources=(ArtifactIdentity("local", "content-addressed-source", None, DIGEST_A),), converter=ToolIdentity("obliteratus-neutral-writer", "1.0.0", COMMIT), obliteratus_commit=COMMIT, configuration_digest=DIGEST_A, tokenizer=ArtifactIdentity("hub", "org/base", "0123456789abcdef", DIGEST_B), base_model=base, command=command, environment={"python": "3.12.7", "platform": "linux", "packages": {"torch": "2.5"}}, source_topology={"world_size": 4, "tp": 2, "pp": 2}, lineage=( LineageEvent( event_id="event-consolidation", event_type="consolidation", parent_artifact_ids=("artifact-sha256:" + "d" * 64,), tool="obliteratus-neutral-writer@1.0.0", transformations=("deduplicate_declared_replicas",), ), ), input_digests=(DIGEST_A,), output_digests=output_digests, transformations=("consolidation", "canonical_safetensors"), observed_scopes=("model_weights",), lost_state=("optimizer_state", "scheduler_state"), adapter=AdapterIdentity( adapter_type="lora", base_model=base, config_digest=DIGEST_B, key_map_digest=DIGEST_A, ), dataset=DatasetIdentity( identifier="dataset/name", revision="rev-1", digest=DIGEST_A, split="train", subset=None, record_count=42, ), training=TrainingIdentity( method="adapter_train", framework="transformers", framework_version="4.60.0", hyperparameters_digest=DIGEST_B, ), unknowns=("optimizer_producer_version",), ) def test_provenance_is_strict_versioned_canonical_and_content_addressed(): record = _record() payload = record.to_dict() Draft202012Validator(SCHEMA).validate(payload) assert payload["schema_id"] == "obliteratus.artifact-provenance" assert payload["schema_version"] == "1.0.0" assert payload["artifact_id"].startswith("artifact-sha256:") assert payload["record_digest"].startswith("sha256:") assert json.loads(record.to_json()) == payload assert record.artifact_id == _record().artifact_id assert record.to_json() == _record().to_json() def test_canonical_sets_and_maps_do_not_depend_on_caller_order(): first = _record(output_digests=(DIGEST_A, DIGEST_B)) second = _record(output_digests=(DIGEST_B, DIGEST_A)) assert first.artifact_id == second.artifact_id assert first.to_json() == second.to_json() @pytest.mark.parametrize( ("scopes", "expected"), [ (("model_weights",), "weights_only"), (("model_weights", "optimizer_state"), "model_and_optimizer"), ( ( "model_weights", "optimizer_state", "scheduler_state", "rng_state", "dataloader_state", "framework_state", ), "exact_resume", ), (("optimizer_state",), "unknown"), ((), "unknown"), ], ) def test_resume_classification_is_derived_only_from_observed_state(scopes, expected): assert classify_resume_state(scopes) == expected def test_lineage_vocabulary_keeps_surgery_distinct_from_finetuning(): allowed = { "consolidation", "reshard", "pretrain", "full_finetune", "adapter_train", "adapter_merge", "quantization", "dequantization", "surgery", } for event_type in allowed: assert LineageEvent("event", event_type, (), "tool@1", ()).event_type == event_type with pytest.raises(ValueError, match="lineage event type"): LineageEvent("event", "finetuning_surgery", (), "tool@1", ()) def test_command_redacts_secrets_prompts_and_private_paths(): sanitized = sanitize_command( ( "checkpoint", "convert", "/home/alice/private/model", "--token", "hf_abcdefghijklmnopqrstuvwxyz", "--prompt=raw private prompt", ) ) rendered = json.dumps(sanitized) assert "/home/alice" not in rendered assert "hf_" not in rendered assert "raw private prompt" not in rendered assert "[REDACTED]" in sanitized assert any(value.startswith("[LOCAL_PATH:sha256:") for value in sanitized) def test_legacy_migration_preserves_declared_facts_and_never_invents_digests(): facts = migrate_legacy_metadata( { "model": "org/base", "model_revision": "rev-1", "tokenizer_revision": None, "seed": "42", "dataset_inputs": [{"identifier": "builtin", "sha256": "e" * 64}], "unmapped_private_field": "must not leak", } ) payload = facts.to_dict() assert payload["base_model"] == { "identity": "org/base", "revision": "rev-1", "digest": None, } assert payload["tokenizer"] == {"revision": None, "digest": None} assert payload["seed"] == "42" assert payload["datasets"] == [ {"identifier": "builtin", "digest": "sha256:" + "e" * 64} ] assert "unmapped_private_field" not in json.dumps(payload) assert "base_model_digest" in payload["unknowns"] assert "tokenizer_digest" in payload["unknowns"] def test_run_archive_attaches_same_artifact_identity_without_raw_sensitive_data(tmp_path): archive = RunArchive(tmp_path) run_id = archive.begin(["org/base"]) record = _record( command=("convert", "/home/alice/private/model", "--token", "hf_secretsecretsecret") ) manifest = archive.attach_checkpoint_provenance(run_id, record) provenance_path = tmp_path / run_id / "checkpoint-provenance.json" raw = provenance_path.read_text(encoding="utf-8") assert manifest["artifact_id"] == record.artifact_id assert manifest["checkpoint_provenance"]["artifact_id"] == record.artifact_id assert json.loads(raw)["artifact_id"] == record.artifact_id assert "/home/alice" not in raw assert "hf_secret" not in raw def test_provenance_rejects_secret_bearing_environment_keys(): with pytest.raises(ValueError, match="sensitive key"): build_provenance( sources=(ArtifactIdentity("local", "source", None, DIGEST_A),), converter=ToolIdentity("writer", "1", COMMIT), obliteratus_commit=COMMIT, configuration_digest=None, tokenizer=None, base_model=None, command=("convert",), environment={"API_TOKEN": "secret"}, source_topology={}, lineage=(), input_digests=(DIGEST_A,), output_digests=(DIGEST_B,), transformations=(), observed_scopes=("model_weights",), lost_state=(), ) @pytest.mark.parametrize( ("factory", "message"), [ (lambda: ArtifactIdentity("other", "source", None, DIGEST_A), "kind is invalid"), (lambda: ArtifactIdentity("local", "", None, DIGEST_A), "non-empty bounded"), ( lambda: ArtifactIdentity("local", "hf_abcdefghijklmnopqrstuvwxyz", None, DIGEST_A), "contains a secret", ), (lambda: ArtifactIdentity("local", "/private/source", None, DIGEST_A), "private local path"), (lambda: ArtifactIdentity("local", "source", None, "bad"), "sha256 digest"), (lambda: ToolIdentity("tool", "1", "A" * 40), "40-character lowercase commit"), ( lambda: LineageEvent("event", "surgery", ("bad-parent",), "tool@1", ()), "parent artifact ID", ), ( lambda: DatasetIdentity("dataset", None, DIGEST_A, None, None, -1), "record count", ), (lambda: TrainingIdentity("invalid", None, None, None), "training method"), (lambda: TrainingIdentity("unknown", None, None, "bad"), "sha256 digest"), ], ) def test_identity_records_reject_unverifiable_or_sensitive_fields(factory, message): with pytest.raises(ValueError, match=message): factory() def _minimal_provenance(**overrides): values = { "sources": (ArtifactIdentity("local", "source", None, DIGEST_A),), "converter": ToolIdentity("writer", "1", COMMIT), "obliteratus_commit": COMMIT, "configuration_digest": None, "tokenizer": None, "base_model": None, "command": ("convert",), "environment": {"python": "3.12", "platform": "linux", "packages": {}}, "source_topology": {}, "lineage": (), "input_digests": (DIGEST_A,), "output_digests": (DIGEST_B,), "transformations": (), "observed_scopes": ("model_weights",), "lost_state": (), } values.update(overrides) return build_provenance(**values) @pytest.mark.parametrize( ("overrides", "message"), [ ({"sources": ()}, "at least one source"), ({"configuration_digest": "bad"}, "sha256 digest"), ({"environment": {"python": "3", "platform": "linux", "packages": {}, "extra": 1}}, "unsupported fields"), ({"source_topology": {"rank": 1 << 64}}, "outside int64"), ({"source_topology": {"path": "/private/checkpoint"}}, "private local path"), ({"source_topology": {"note": "hf_abcdefghijklmnopqrstuvwxyz"}}, "contains a secret"), ({"source_topology": {1: "value"}}, "non-string key"), ({"source_topology": {"prompt_text": "value"}}, "sensitive key"), ({"source_topology": {"opaque": object()}}, "non-JSON value"), ({"input_digests": ("bad",)}, "sha256 digest"), ], ) def test_provenance_builder_rejects_incomplete_or_unsafe_evidence(overrides, message): with pytest.raises(ValueError, match=message): _minimal_provenance(**overrides) def test_normalization_accepts_explicit_json_scalars_and_sequences(): record = _minimal_provenance( source_topology={"active": True, "optional": None, "ranks": [0, 1]}, ).to_dict() assert record["source_topology"] == { "active": True, "optional": None, "ranks": [0, 1], } def test_command_redaction_covers_equals_and_bare_sensitive_values(): sanitized = sanitize_command( ( "convert", "--output=/private/output", "--note=ok", "--credential=secret-value", "bearer abcdefghijklmnop", ) ) assert sanitized[1].startswith("--output=[LOCAL_PATH:sha256:") assert sanitized[2] == "--note=ok" assert sanitized[3] == "--credential=[REDACTED]" assert sanitized[4] == "[REDACTED]" def test_legacy_migration_marks_invalid_identity_and_ignores_invalid_dataset_rows(): payload = migrate_legacy_metadata( { "model": "/private/model", "model_revision": "hf_abcdefghijklmnopqrstuvwxyz", "tokenizer_revision": 42, "dataset_inputs": ["invalid", {"identifier": "/private", "sha256": "a" * 64}], } ).to_dict() assert payload["base_model"] == {"identity": None, "revision": None, "digest": None} assert payload["tokenizer"]["revision"] is None assert payload["datasets"] == [] assert "base_model_identity" in payload["unknowns"] def test_provenance_verifier_rejects_tampering_and_constructor_identity_disagreement(): record = _record() tampered = record.to_dict() tampered["unknowns"] = ["changed"] with pytest.raises(ValueError, match="record digest mismatch"): verify_provenance_record(tampered) with pytest.raises(ValueError, match="identity fields disagree"): ProvenanceRecord( "artifact-sha256:" + "f" * 64, record.record_digest, record.to_json(), ) with pytest.raises(ValueError, match="JSON is not canonical"): ProvenanceRecord( record.artifact_id, record.record_digest, json.dumps(record.to_dict()), ) def test_public_metadata_bounds_mixed_keys_and_legacy_sensitive_values_fail_closed(): with pytest.raises(ValueError, match="non-string key"): _minimal_provenance(source_topology={"rank": 0, 1: "invalid"}) with pytest.raises(ValueError, match="invalid key"): _minimal_provenance(source_topology={"x" * 513: "invalid"}) deeply_nested = {} cursor = deeply_nested for _ in range(18): cursor["next"] = {} cursor = cursor["next"] with pytest.raises(ValueError, match="nesting is too deep"): _minimal_provenance(source_topology=deeply_nested) with pytest.raises(ValueError, match="command argument"): _minimal_provenance(command=(object(),)) legacy = migrate_legacy_metadata( { "model": "org/base", "model_revision": "/private/revision", "tokenizer_revision": "hf_abcdefghijklmnopqrstuvwxyz", "seed": "/private/seed", "dataset_inputs": None, } ).to_dict() rendered = json.dumps(legacy) assert "/private" not in rendered assert "hf_" not in rendered assert legacy["seed"] is None assert "seed" in legacy["unknowns"] def test_public_provenance_parsers_reject_boundedness_and_structure_attacks(): record = _record() with pytest.raises(ValueError, match="provenance JSON is invalid"): ProvenanceRecord(record.artifact_id, record.record_digest, '{"a": 1, "a": 2}') with pytest.raises(ValueError, match="text is too large"): _minimal_provenance(source_topology={"note": "x" * 1025}) with pytest.raises(ValueError, match="too many fields"): _minimal_provenance( source_topology={f"field-{index}": index for index in range(4097)}, ) with pytest.raises(ValueError, match="too many items"): _minimal_provenance(source_topology={"items": [None] * 4097}) with pytest.raises(ValueError, match="must be an object"): verify_provenance_record([]) with pytest.raises(ValueError, match="bounded string collection"): classify_resume_state("model_weights") with pytest.raises(ValueError, match="collection is invalid"): _minimal_provenance(input_digests=DIGEST_A) with pytest.raises(ValueError, match="must contain strings"): _minimal_provenance(input_digests=(1,)) with pytest.raises(ValueError, match="sources collection"): _minimal_provenance(sources="source") with pytest.raises(ValueError, match="artifact identities"): _minimal_provenance(sources=(object(),)) with pytest.raises(ValueError, match="lineage events"): _minimal_provenance(lineage=(object(),)) with pytest.raises(ValueError, match="bounded mapping"): migrate_legacy_metadata([]) def test_provenance_verifier_rejects_each_identity_and_state_layer(): payload = _record().to_dict() malformed = {**payload, "sources": []} with pytest.raises(ValueError, match="structure is invalid"): verify_provenance_record(malformed) malformed = {**payload, "artifact_id": "invalid"} with pytest.raises(ValueError, match="artifact ID is invalid"): verify_provenance_record(malformed) malformed = {**payload, "state": {}} with pytest.raises(ValueError, match="state is invalid"): verify_provenance_record(malformed) malformed = json.loads(json.dumps(payload)) malformed["state"]["classification"] = "unknown" with pytest.raises(ValueError, match="classification is not evidence-derived"): verify_provenance_record(malformed) malformed = {**payload, "artifact_id": "artifact-sha256:" + "f" * 64} digest_input = {key: value for key, value in malformed.items() if key != "record_digest"} encoded = json.dumps( digest_input, sort_keys=True, separators=(",", ":"), ensure_ascii=True, allow_nan=False, ).encode("utf-8") malformed["record_digest"] = f"sha256:{sha256(encoded).hexdigest()}" with pytest.raises(ValueError, match="artifact ID mismatch"): verify_provenance_record(malformed) def test_command_and_legacy_scalar_paths_remain_public_and_canonical(): assert sanitize_command(("convert", f"--note=hf_{'a' * 20}")) == ( "convert", "--note=[REDACTED]", ) assert migrate_legacy_metadata({"seed": 7}).to_dict()["seed"] == "7" @pytest.mark.parametrize( ("mutation", "message"), [ (lambda value: value["sources"][0].update(kind="invalid"), "source identity"), (lambda value: value["converter"].update(commit="invalid"), "converter"), (lambda value: value["command"].append("/private/source"), "command"), (lambda value: value["environment"].update(packages=[]), "packages"), (lambda value: value.update(source_topology=[]), "source topology"), (lambda value: value["lineage"][0].update(extra=True), "lineage event fields"), (lambda value: value.update(input_digests=[]), "input digests"), ( lambda value: value.update(output_digests=[DIGEST_B, DIGEST_B]), "output digests", ), (lambda value: value["transformations"].reverse(), "transformations"), ( lambda value: value["state"].update( observed_scopes=["model_weights", "model_weights"] ), "observed scopes", ), (lambda value: value["adapter"].update(config_digest="invalid"), "adapter identity"), (lambda value: value["dataset"].update(record_count=True), "dataset identity"), (lambda value: value["training"].update(framework=[]), "training framework"), (lambda value: value.update(unknowns=["duplicate", "duplicate"]), "unknowns"), ], ) def test_provenance_verifier_rejects_noncanonical_nested_records(mutation, message): payload = deepcopy(_record().to_dict()) mutation(payload) with pytest.raises(ValueError, match=message): verify_provenance_record(payload) def test_builder_cannot_emit_schema_invalid_environment_or_duplicate_lineage(): with pytest.raises(ValueError, match="environment python"): _minimal_provenance( environment={"python": [], "platform": "linux", "packages": {}}, ) event = LineageEvent("event", "consolidation", (), "writer@1", ()) with pytest.raises(ValueError, match="sorted and unique"): _minimal_provenance(lineage=(event, event))