diff --git a/docs/input-bundles.md b/docs/input-bundles.md new file mode 100644 index 0000000..fed9f0c --- /dev/null +++ b/docs/input-bundles.md @@ -0,0 +1,60 @@ +# Shared registered research inputs + +`quant_data_kit.input_bundles.InputBundleStore` shares registered CSV, JSON, text +and arbitrary file bytes. It does not assign financial, PIT, schema or executable +trust to the content. Parquet market queries continue to use `SnapshotStore`. + +```python +from pathlib import Path +from quant_data_kit.artifacts import ArtifactStore +from quant_data_kit.input_bundles import InputBundleStore + +cas = ArtifactStore(Path("D:/research/cas")) +bundles = InputBundleStore(cas) +bundle = bundles.publish( + {"prices.csv": Path("D:/input/prices.csv"), "plan.json": Path("D:/input/plan.json")}, + owner="project/study-1", metadata={"version": 1}, +) +with bundles.open_member(bundle.artifact, "plan.json") as stream: + plan_bytes = stream.read() +print(bundles.describe(bundle.artifact)) +``` + +Use the same ArtifactRef for candidates/retries. Manifests contain logical names +and content references, not source paths. Names cannot contain path separators. +Up to 1000 files and a 1 MiB manifest are accepted. Files stream into CAS; the +complete input set is not loaded into memory during publication. + +## Integrity, concurrency and invalidation + +The existing CAS process lock and SQLite transaction cover publication, dependency +edges and owner pins. Concurrent identical publications produce one manifest and +one object per distinct byte content. Changing bytes or metadata produces a new +identity; old pinned versions remain. Duplicate imports still read/hash bytes. + +Every load/read verifies hashes, sizes and the manifest's registered dependencies. +Missing/modified members fail closed. The existing CAS rejects symlink/reparse +points and hard links, and checks for source mutation during import. Streams are +read-only and hold the store lock against GC until closed. Close them promptly: +v1 readers are serialized. No writable CAS path or hard link is handed to a trial. +Trusted host-user code still has OS privileges; this is not an untrusted sandbox. + +## Retention and crash recovery + +Owner pins protect manifests and their member closure. Pipeline research batches +add submission/job pins before execution. Failure, cancellation, retry and restart +never release inputs or delete history. There is no automatic expiry or GC. + +After auditing active consumers and historical evidence needs, an operator can +explicitly release an owner. `release` removes pins, not bytes. Separate existing +`collect(min_age_seconds=...)` reclaims only unprotected objects after the grace +period. Savings come from reuse, without requiring historic result deletion. + +A killed publisher may leave unindexed files but cannot expose a partial bundle. +Restart verifies/reuses complete bytes and publishes atomically. Explicit CAS GC +can reclaim unindexed leftovers after its grace period. + +`describe` separates logical member bytes, distinct member bytes and manifest +bytes. Repeated integrity reads still incur I/O. No zero-I/O or streaming factor +computation is claimed. Existing Notebook inputs are not automatically migrated: +handlers must explicitly consume references; unregistered reads stay outside. diff --git a/src/quant_data_kit/input_bundles.py b/src/quant_data_kit/input_bundles.py new file mode 100644 index 0000000..54acdbb --- /dev/null +++ b/src/quant_data_kit/input_bundles.py @@ -0,0 +1,161 @@ +"""Named immutable research inputs, shared across candidates and retries. + +Unlike SnapshotStore this layer accepts arbitrary registered file bytes; it makes +no market-data, PIT, or Parquet claims. Readers receive read-only verified streams. +""" + +from __future__ import annotations + +import io +import json +import re +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import BinaryIO + +from .artifacts import ArtifactRef, ArtifactStore, _text, canonical_json +from .exceptions import ValidationError + +_SCHEMA = "quant.input-bundle/v1" +_MEDIA = "application/vnd.quant.input-bundle+json" +_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}\Z") + + +@dataclass(frozen=True) +class InputBundle: + artifact: ArtifactRef + document: dict + + +class InputBundleStore: + """Use one content-addressed bundle reference instead of copying per trial. + + Pins protect the manifest and every member. No automatic release, expiration, + hard links, writable CAS paths, or deletion is performed by this API. + """ + + def __init__(self, artifacts: ArtifactStore): + self.artifacts = artifacts + + @staticmethod + def _name(value: str) -> str: + if not isinstance(value, str) or not _NAME.fullmatch(value): + raise ValidationError("Input names must be simple identifiers without paths") + return value + + def publish( + self, files: Mapping[str, Path], *, owner: str, metadata: Mapping | None = None + ) -> InputBundle: + """Stream files once into CAS and atomically publish members, edges and pin. + + Repeated publication reads and hashes bytes again but stores one copy. + Metadata is part of the immutable identity; source paths are not retained. + A crash can leave unindexed bytes, handled by ArtifactStore.collect's + explicit grace policy. It cannot expose a partially published bundle. + """ + _text(owner, "owner") + if not isinstance(files, Mapping) or not 1 <= len(files) <= 1000: + raise ValidationError("Input bundle requires 1 to 1000 files") + for name in files: + self._name(name) + if metadata is not None and not isinstance(metadata, Mapping): + raise ValidationError("Input metadata must be an object") + metadata = json.loads(canonical_json(dict(metadata or {}))) + with self.artifacts._session() as db: + members = [ + {"name": name, "artifact": self.artifacts._put_file(db, path, None).to_dict()} + for name, path in sorted(files.items()) + ] + document = {"schema_version": _SCHEMA, "metadata": metadata, "members": members} + encoded = canonical_json(document) + if len(encoded) > 1024 * 1024: + raise ValidationError("Input bundle manifest exceeds 1 MiB") + ref = self.artifacts._put_stream(db, io.BytesIO(encoded), _MEDIA) + for member in members: + db.execute( + "INSERT OR IGNORE INTO edges VALUES (?,?)", + (ref.sha256, member["artifact"]["sha256"]), + ) + self.artifacts._pin(db, owner, [ref]) + self.artifacts._event( + "input_bundle_published", sha256=ref.sha256, members=len(members), owner=owner + ) + return InputBundle(ref, document) + + def _load(self, db, ref: ArtifactRef) -> InputBundle: + if ref.size_bytes > 1024 * 1024: + raise ValidationError("Input bundle manifest exceeds 1 MiB") + path = self.artifacts._verify(db, ref) + try: + document = json.loads(path.read_bytes()) + canonical_json(document) # Reject forged nonfinite JSON metadata too. + if ( + set(document) != {"schema_version", "metadata", "members"} + or document["schema_version"] != _SCHEMA + or not isinstance(document["metadata"], dict) + or not isinstance(document["members"], list) + or not 1 <= len(document["members"]) <= 1000 + ): + raise ValidationError("Invalid input bundle manifest") + names, expected_edges = set(), set() + for member in document["members"]: + if not isinstance(member, dict) or set(member) != {"name", "artifact"}: + raise ValidationError("Invalid input bundle member") + name = self._name(member["name"]) + if name in names: + raise ValidationError("Duplicate input name") + names.add(name) + child = ArtifactRef.from_dict(member["artifact"]) + self.artifacts._verify(db, child) + expected_edges.add(child.sha256) + actual_edges = { + row[0] + for row in db.execute("SELECT child FROM edges WHERE parent=?", (ref.sha256,)) + } + if actual_edges != expected_edges: + raise ValidationError("Input bundle dependencies are not registered") + except ValidationError: + raise + except (KeyError, TypeError, ValueError) as exc: + raise ValidationError("Invalid input bundle manifest") from exc + return InputBundle(ref, document) + + def load(self, ref: ArtifactRef) -> InputBundle: + with self.artifacts._session() as db: + return self._load(db, ref) + + @contextmanager + def open_member(self, ref: ArtifactRef, name: str) -> Iterator[BinaryIO]: + """Hold the store lock against GC for the entire verified read-only stream. + + This serializes simultaneous readers in the v1 CAS. It is byte reuse, not + zero I/O or streaming factor computation. Always close readers promptly. + """ + self._name(name) + with self.artifacts._session() as db: + bundle = self._load(db, ref) + member = next( + (item for item in bundle.document["members"] if item["name"] == name), None + ) + if member is None: + raise ValidationError("Input name does not exist") + path = self.artifacts._verify(db, ArtifactRef.from_dict(member["artifact"])) + with path.open("rb") as stream: + yield stream + + def describe(self, ref: ArtifactRef) -> dict: + """Report logical and deduplicated payload bytes without changing retention.""" + bundle = self.load(ref) + refs = [ArtifactRef.from_dict(item["artifact"]) for item in bundle.document["members"]] + unique = {item.sha256: item.size_bytes for item in refs} + return { + "schema_version": "quant.input-bundle-usage/v1", + "artifact": ref.to_dict(), + "members": len(refs), + "logical_input_bytes": sum(item.size_bytes for item in refs), + "unique_input_bytes": sum(unique.values()), + "manifest_bytes": ref.size_bytes, + "retention": "explicit-owner-pins-no-automatic-release", + } diff --git a/tests/test_input_bundles.py b/tests/test_input_bundles.py new file mode 100644 index 0000000..7c5e3c1 --- /dev/null +++ b/tests/test_input_bundles.py @@ -0,0 +1,184 @@ +import io +import json +import multiprocessing as mp +import sqlite3 +from pathlib import Path + +import pytest + +from quant_data_kit.artifacts import ArtifactRef, ArtifactStore, canonical_json +from quant_data_kit.exceptions import ValidationError +from quant_data_kit.input_bundles import InputBundleStore + + +def object_path(store, ref): + return store.root / "objects" / ref.sha256[:2] / ref.sha256 + + +def concurrent_publish(root, source, start, queue): + store = InputBundleStore(ArtifactStore(Path(root))) + start.wait(20) + bundle = store.publish({"prices.csv": Path(source)}, owner=f"pid-{mp.current_process().pid}") + queue.put(bundle.artifact.to_dict()) + + +def interrupted_publish(root, source, ready): + import time + + store = ArtifactStore(Path(root)) + original = store._pin + + def stop_before_commit(*args): + original(*args) + ready.set() + time.sleep(120) + + store._pin = stop_before_commit + InputBundleStore(store).publish({"prices.csv": Path(source)}, owner="uncommitted") + + +def test_reuse_and_transitive_retention_without_per_trial_copies(tmp_path): + source = tmp_path / "prices.csv" + source.write_bytes(b"date,close\n2026-01-01,1\n") + artifacts = ArtifactStore(tmp_path / "cas") + store = InputBundleStore(artifacts) + first = store.publish({"prices.csv": source, "same.csv": source}, owner="project") + second = store.publish({"same.csv": source, "prices.csv": source}, owner="retry") + assert first.artifact == second.artifact + usage = store.describe(first.artifact) + assert usage["logical_input_bytes"] == 2 * source.stat().st_size + assert usage["unique_input_bytes"] == source.stat().st_size + assert len(list((artifacts.root / "objects").glob("*/*"))) == 2 + source.unlink() + with store.open_member(first.artifact, "prices.csv") as stream: + assert stream.read().startswith(b"date,close") + assert not stream.writable() + artifacts.release("project") + assert artifacts.collect(min_age_seconds=0) == [] + assert store.load(first.artifact).document == first.document + artifacts.release("retry") + assert len(artifacts.collect(min_age_seconds=0)) == 2 + + +def test_changed_input_and_metadata_create_new_identity_keep_old(tmp_path): + source = tmp_path / "input.txt" + source.write_text("first") + store = InputBundleStore(ArtifactStore(tmp_path / "cas")) + first = store.publish({"input": source}, owner="run", metadata={"version": 1}) + source.write_text("second") + second = store.publish({"input": source}, owner="run", metadata={"version": 1}) + third = store.publish({"input": source}, owner="run", metadata={"version": 2}) + assert len({first.artifact.sha256, second.artifact.sha256, third.artifact.sha256}) == 3 + with store.open_member(first.artifact, "input") as stream: + assert stream.read() == b"first" + + +@pytest.mark.parametrize("name", ["../x", "a/b", "a\\b", "", ".", 1]) +def test_invalid_names_rejected(tmp_path, name): + store = InputBundleStore(ArtifactStore(tmp_path / "cas")) + with pytest.raises(ValidationError): + store.publish({name: tmp_path / "missing"}, owner="invalid") + + +def test_missing_files_and_oversized_metadata_never_publish_partial_bundle(tmp_path): + store = InputBundleStore(ArtifactStore(tmp_path / "cas")) + source = tmp_path / "a" + source.write_text("data") + with pytest.raises(ValidationError): + store.publish({"a": source, "z": tmp_path / "missing"}, owner="invalid") + with pytest.raises(ValidationError, match="1 MiB"): + store.publish({"a": source}, owner="large", metadata={"x": "x" * (1024 * 1024)}) + with sqlite3.connect(store.artifacts.root / "index.sqlite3") as db: + assert db.execute("SELECT COUNT(*) FROM objects").fetchone()[0] == 0 + assert db.execute("SELECT COUNT(*) FROM pins").fetchone()[0] == 0 + + +def test_member_tamper_is_rejected_at_load_read_and_retention(tmp_path): + source = tmp_path / "input" + source.write_bytes(b"trusted") + artifacts = ArtifactStore(tmp_path / "cas") + store = InputBundleStore(artifacts) + bundle = store.publish({"input": source}, owner="first") + member = ArtifactRef.from_dict(bundle.document["members"][0]["artifact"]) + object_path(artifacts, member).write_bytes(b"changed") + with pytest.raises(ValidationError): + store.load(bundle.artifact) + with pytest.raises(ValidationError), store.open_member(bundle.artifact, "input"): + pass + with pytest.raises(ValidationError): + artifacts.retain("next", [bundle.artifact]) + + +def test_forged_manifest_or_dependency_edges_rejected(tmp_path): + source = tmp_path / "input" + source.write_bytes(b"trusted") + artifacts = ArtifactStore(tmp_path / "cas") + store = InputBundleStore(artifacts) + bundle = store.publish({"input": source}, owner="first") + forged = dict(bundle.document, metadata={"forged": True}) + ref = artifacts.put_bytes(canonical_json(forged)) + with pytest.raises(ValidationError, match="dependencies"): + store.load(ref) + with ( + pytest.raises(ValidationError, match="does not exist"), + store.open_member(bundle.artifact, "absent"), + ): + pass + for value in ([], {}, {"schema_version": "unknown"}): + bad = artifacts.put_bytes(json.dumps(value).encode()) + with pytest.raises(ValidationError): + store.load(bad) + with artifacts._session() as db: + duplicate = json.loads(json.dumps(bundle.document)) + duplicate["members"].append(duplicate["members"][0]) + bad = artifacts._put_stream(db, io.BytesIO(canonical_json(duplicate)), None) + with pytest.raises(ValidationError, match="Duplicate"): + store.load(bad) + + +def test_real_concurrent_publish_reuses_one_bundle(tmp_path): + source = tmp_path / "prices.csv" + source.write_bytes(b"row,value\n" * 1000) + root = tmp_path / "cas" + ArtifactStore(root) + context = mp.get_context("spawn") + start, queue = context.Event(), context.Queue() + children = [ + context.Process(target=concurrent_publish, args=(root, source, start, queue)) + for _ in range(3) + ] + for child in children: + child.start() + start.set() + results = [queue.get(timeout=60) for _ in children] + for child in children: + child.join(30) + assert child.exitcode == 0 + assert all(ref == results[0] for ref in results) + assert len(list((root / "objects").glob("*/*"))) == 2 + + +def test_killed_publisher_never_exposes_partial_bundle_and_restart_succeeds(tmp_path): + source = tmp_path / "prices.csv" + source.write_bytes(b"important-data") + root = tmp_path / "cas" + ArtifactStore(root) + context = mp.get_context("spawn") + ready = context.Event() + process = context.Process(target=interrupted_publish, args=(root, source, ready)) + process.start() + try: + assert ready.wait(30) + process.kill() + process.join(15) + with sqlite3.connect(root / "index.sqlite3") as db: + assert db.execute("SELECT COUNT(*) FROM pins").fetchone()[0] == 0 + assert db.execute("SELECT COUNT(*) FROM objects").fetchone()[0] == 0 + store = InputBundleStore(ArtifactStore(root)) + bundle = store.publish({"prices.csv": source}, owner="recovered") + assert store.load(bundle.artifact).artifact == bundle.artifact + assert store.artifacts.collect(min_age_seconds=0) == [] + finally: + if process.is_alive(): + process.kill() + process.join(15)