feat: versioned DatasetStorage for reading/writing datasets across protocol versions

.github/workflows/unit-tests.yml

@@ -111,11 +111,30 @@
       - name: Run job tests
         run: uv run pytest packages/syft-job/tests/ -v
 
+  dataset-tests:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - name: Install the project
+        run: uv sync --all-extras
+
+      - name: Run dataset tests
+        run: uv run pytest packages/syft-datasets/tests/ -v
+
   enclave-tests:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
 
       - name: Install uv
         uses: astral-sh/setup-uv@v5
 

.pre-commit-config.yaml

@@ -47,4 +47,4 @@ repos:
     rev: 'v3.0.0-alpha.9-for-vscode'
     hooks:
       - id: prettier
-        exclude: (\.min\.(js|css)$|^ci/|^test_outputs/)
+        exclude: (\.min\.(js|css)$|^ci/|^test_outputs/|^packages/syft-datasets/tests/migrations/(p2p|unit)/fixtures/)

Justfile

@@ -30,6 +30,10 @@ test-unit-job:
     #!/bin/bash
     uv run pytest -v ./packages/syft-job/tests
 
+test-unit-datasets:
+    #!/bin/bash
+    uv run pytest -n auto ./packages/syft-datasets/tests
+
 test-unit-migration:
     #!/bin/bash
     uv run pytest -n auto ./packages/syft-migration/tests

packages/enclave-model-api-example/src/enclave_model_api/logs_dataset.py

@@ -10,7 +10,10 @@
 import tempfile
 from pathlib import Path
 
+from syft_datasets.dataset_ref import DatasetNotFoundError
+
 from enclave_model_api.log_writer import LOG_FILE_NAME
+from enclave_model_api.paths import resolve_private_dataset_dir
 
 LOGS_DATASET_SUMMARY = (
     "Inference request logs (prompt, completion, stats) collected by the "
@@ -41,9 +44,12 @@ def ensure_logs_dataset(client, name: str) -> Path:
 
     Idempotent so reboots with fresh_state=false keep the existing logs.
     """
-    config = client.datasets.syftbox_config
-    if (config.get_my_mock_dataset_dir(name) / "dataset.yaml").exists():
-        return config.private_dir_for_my_dataset(name)
+    storage = client.datasets.storage
+    email = storage.config.email
+    try:
+        return storage.private_dataset_dir(storage.find_dataset_ref(email, name))
+    except DatasetNotFoundError:
+        pass
 
     mock_dir, private_dir = _seed_dirs()
     client.create_dataset(
@@ -56,4 +62,4 @@ def ensure_logs_dataset(client, name: str) -> Path:
         sync=False,
     )
     client.sync()
-    return config.private_dir_for_my_dataset(name)
+    return resolve_private_dataset_dir(storage, email, name)

packages/enclave-model-api-example/src/enclave_model_api/paths.py

@@ -9,19 +9,39 @@
 
 from syft_client.sync.syftbox_manager import get_jupyter_default_syftbox_folder
 from syft_datasets.config import SyftBoxConfig
+from syft_datasets.dataset_ref import DatasetNotFoundError, DatasetRef
+from syft_datasets.dataset_storage import DatasetStorage
 
 
 def default_syftbox_folder(email: str) -> Path:
     """The folder SyftboxManagerConfig.for_jupyter (used by for_enclave) picks."""
     return get_jupyter_default_syftbox_folder(email)
 
 
+def resolve_private_dataset_dir(storage: DatasetStorage, owner: str, name: str) -> Path:
+    """Private dir at the dataset's actual on-disk protocol layout.
+
+    A dataset may live at protocol 0 (flat) or under a ``v<n>`` segment; the
+    written layout depends on what the audience can read, not the current
+    default. Datasets not yet on disk (e.g. weights still syncing) fall back to
+    the widest-compatible protocol — where a peer running any current release
+    writes them for us.
+    """
+    try:
+        ref = storage.find_dataset_ref(owner, name)
+    except DatasetNotFoundError:
+        (widest,) = storage.target_protocol_versions_for_peers(None)
+        ref = DatasetRef(owner=owner, name=name, protocol_version=widest)
+    return storage.private_dataset_dir(ref)
+
+
 def private_dataset_dir(
     syftbox_folder: Path | str, datasite: str, dataset_name: str
 ) -> Path:
     """Private dir of *dataset_name* on *datasite* inside *syftbox_folder*."""
     config = SyftBoxConfig(syftbox_folder=Path(syftbox_folder), email=datasite)
-    return config.private_dir_for_my_dataset(dataset_name)
+    storage = DatasetStorage(config=config)
+    return resolve_private_dataset_dir(storage, datasite, dataset_name)
 
 
 def find_checkpoint_dir(weights_dir: Path | str) -> Path | None:

packages/syft-datasets/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "syft-dataset"
-version = "0.1.20"
+version = "0.1.21"
 description = "Syft Datasets"
 authors = [{ name = "OpenMined", email = "info@openmined.org" }]
 license = { text = "Apache-2.0" }
@@ -10,6 +10,7 @@ dependencies = [
     "pyyaml>=6.0.3",
     "syft-notebook-ui",
     "syft-perms==0.1.14",
+    "syft-migration",
 ]
 
 

packages/syft-datasets/scripts/export_release_artifact.py

@@ -0,0 +1,41 @@
+"""Export the release artifacts for the current syft-dataset version.
+
+Run on EVERY release (uv run python scripts/export_release_artifact.py):
+always writes the package release info; additionally writes the protocol
+artifact when this release introduces a new protocol version.
+"""
+
+import sys
+
+from syft_datasets.migrations.history import PACKAGE_ARTIFACTS_DIR, PROTOCOLS_DIR
+from syft_datasets.migrations.registry import DATASET_PROTOCOL_VERSION, dataset_registry
+from syft_datasets.version import __version__
+
+
+def main() -> None:
+    # Import the models so every versioned object is registered.
+    import syft_datasets  # noqa: F401
+
+    if dataset_registry.protocol_changed_without_bump():
+        sys.exit(
+            "The dataset protocol changed compared to the released "
+            f"protocol-{DATASET_PROTOCOL_VERSION}.json; bump "
+            "DATASET_PROTOCOL_VERSION in syft_datasets/migrations/registry.py "
+            "before releasing."
+        )
+
+    PACKAGE_ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
+    PROTOCOLS_DIR.mkdir(parents=True, exist_ok=True)
+
+    info_path = PACKAGE_ARTIFACTS_DIR / f"syft-dataset-{__version__}.json"
+    dataset_registry.compute_released_package_protocol_info().save(info_path)
+    print(f"Wrote {info_path}")
+
+    protocol_path = PROTOCOLS_DIR / f"protocol-{DATASET_PROTOCOL_VERSION}.json"
+    if not protocol_path.exists():
+        dataset_registry.compute_released_protocol().save(protocol_path)
+        print(f"Wrote {protocol_path} (new protocol version)")
+
+
+if __name__ == "__main__":
+    main()

packages/syft-datasets/scripts/generate_release_fixture.py

@@ -0,0 +1,126 @@
+"""Generate a p2p backward-compatibility fixture for a syft-dataset release.
+
+Run on EVERY release, after bumping the version:
+
+    uv run python scripts/generate_release_fixture.py
+
+Writes a SyftBox tree exactly as this release serializes datasets to disk, into
+
+    tests/migrations/p2p/fixtures/syft_dataset-<version>-protocol<p>_syftbox/
+
+Future releases loop over these fixtures (test_older_protocol_compatibility.py)
+to prove they can still read and round-trip older on-disk data.
+
+Protocol 0 / release 0.1.20 predates this script; its fixture
+(syft_dataset-0.1.20-protocol0_syftbox) is generated by passing
+``--protocol 0 --version 0.1.20`` and is treated as hand-authored, like
+protocol-0.json.
+"""
+
+import argparse
+import shutil
+import tempfile
+from pathlib import Path
+
+from syft_datasets.dataset_manager import SyftDatasetManager
+from syft_datasets.migrations import dataset_registry
+from syft_datasets.migrations.registry import DATASET_PROTOCOL_VERSION
+from syft_datasets.version import __version__
+
+DO_EMAIL = "do@test.org"
+DS_EMAIL = "ds@test.org"
+
+FIXTURES_DIR = (
+    Path(__file__).resolve().parents[1] / "tests" / "migrations" / "p2p" / "fixtures"
+)
+
+
+def _seed_syftbox(syftbox: Path, src_dir: Path, protocol_version: str) -> None:
+    """One dataset written in the given protocol layout (mock + private + readme)."""
+    mgr = SyftDatasetManager(syftbox_folder_path=syftbox, email=DO_EMAIL)
+    # Force the target protocol by advertising a peer at that version.
+    mgr.storage.peer_schemas = {
+        DS_EMAIL: dataset_registry.schema_for_protocol_version(protocol_version)
+    }
+
+    mock_path = src_dir / "mock.csv"
+    mock_path.write_text("id,value\n1,10\n2,20\n")
+    private_path = src_dir / "private.csv"
+    private_path.write_text("id,secret\n1,alpha\n2,beta\n")
+    readme_path = src_dir / "readme.md"
+    readme_path.write_text("# demo dataset\n")
+
+    mgr.create(
+        name="demo",
+        mock_path=mock_path,
+        private_path=private_path,
+        readme_path=readme_path,
+        summary="demo dataset",
+        tags=["demo"],
+        users=[DS_EMAIL],
+    )
+
+
+def _normalize(target: Path) -> None:
+    """Strip machine-specific values so the committed fixture is portable.
+
+    Drop the executable bit (fixtures are data, not scripts) and blank the
+    private config's absolute ``data_dir`` (not part of the on-disk format
+    peers see).
+    """
+    import yaml
+
+    for path in target.rglob("*"):
+        if path.is_file():
+            path.chmod(0o644)
+
+    for config in target.rglob("private_metadata.yaml"):
+        data = yaml.safe_load(config.read_text()) or {}
+        if data.get("data_dir"):
+            data["data_dir"] = ""
+        config.write_text(yaml.safe_dump(data, indent=2, sort_keys=False))
+
+
+def _build_fixture(target: Path, protocol_version: str) -> None:
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp_path = Path(tmp)
+        syftbox = tmp_path / "SyftBox"
+        syftbox.mkdir()
+        src_dir = tmp_path / "src"
+        src_dir.mkdir()
+        _seed_syftbox(syftbox, src_dir, protocol_version)
+        shutil.copytree(syftbox, target)
+        _normalize(target)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--force", action="store_true", help="overwrite an existing fixture"
+    )
+    parser.add_argument(
+        "--protocol",
+        default=DATASET_PROTOCOL_VERSION,
+        help="protocol version to write (default: current)",
+    )
+    parser.add_argument(
+        "--version",
+        default=__version__,
+        help="release version to stamp in the fixture name (default: current)",
+    )
+    args = parser.parse_args()
+
+    target = (
+        FIXTURES_DIR / f"syft_dataset-{args.version}-protocol{args.protocol}_syftbox"
+    )
+    if target.exists():
+        if not args.force:
+            raise SystemExit(f"{target} already exists; pass --force to regenerate.")
+        shutil.rmtree(target)
+
+    _build_fixture(target, args.protocol)
+    print(f"Wrote {target}")
+
+
+if __name__ == "__main__":
+    main()

packages/syft-datasets/src/syft_datasets/__init__.py

@@ -1,4 +1,26 @@
+# __version__ comes from the installed distribution metadata (see version.py).
+from .version import __version__
+
 from .config import SyftBoxConfig
+from .dataset_ref import DatasetRef
+from .dataset_storage import DatasetStorage
+from .migrations import dataset_registry
+from .migrations.history import register_historic_schemas
+from .models import Dataset, DatasetV1, PrivateDatasetConfig, PrivateDatasetConfigV1
+
+# Historic schemas list object versions that must already be registered, which
+# happens when the models above are imported.
+register_historic_schemas()
 
-__version__ = "0.1.7"
-__all__ = ["SyftBoxConfig"]
+__all__ = [
+    "__version__",
+    "SyftBoxConfig",
+    "DatasetRef",
+    "DatasetStorage",
+    "dataset_registry",
+    # Models
+    "Dataset",
+    "DatasetV1",
+    "PrivateDatasetConfig",
+    "PrivateDatasetConfigV1",
+]