diff --git a/packages/syft-datasets/src/syft_datasets/dataset_manager.py b/packages/syft-datasets/src/syft_datasets/dataset_manager.py index a77f2499fd6..fa518050b71 100644 --- a/packages/syft-datasets/src/syft_datasets/dataset_manager.py +++ b/packages/syft-datasets/src/syft_datasets/dataset_manager.py @@ -2,7 +2,6 @@ from typing_extensions import Self -import yaml from syft_migration import ProtocolSchema from .types import PathLike, to_path @@ -362,5 +361,4 @@ def _private_config_without_data_dir(self, ref: DatasetRef) -> bytes: """ config = self.storage.read_private_config(ref) config.data_dir = Path("") - data = config.disk_dict() - return yaml.safe_dump(data, indent=2, sort_keys=False).encode() + return self.storage.wire_bytes(ref, config) diff --git a/packages/syft-datasets/src/syft_datasets/dataset_storage.py b/packages/syft-datasets/src/syft_datasets/dataset_storage.py index d87094c9bdb..b641e80bde0 100644 --- a/packages/syft-datasets/src/syft_datasets/dataset_storage.py +++ b/packages/syft-datasets/src/syft_datasets/dataset_storage.py @@ -529,11 +529,25 @@ def write_private_config( codec = self._codec_for(ref.protocol_version) return self._write(codec, codec.private_metadata_path(ref), config, ref) + def wire_bytes(self, ref: DatasetRef, obj: MigratableObject) -> bytes: + """The bytes to put on the wire for ``obj`` under this ref's protocol. + + Same shaping as the write methods: downgrade to the protocol schema, + then the codec's on-disk format. Use it when the bytes go to a peer + instead of to a path. + """ + codec = self._codec_for(ref.protocol_version) + return codec.dumps(self._downgrade(obj, ref)).encode() + # -- internals ----------------------------------------------------------- def _upgrade(self, data: dict, canonical_name: str) -> MigratableObject: obj = self.service.load(data) return self.service.migrate(obj, self.registry.latest_version(canonical_name)) + def _downgrade(self, obj: MigratableObject, ref: DatasetRef) -> MigratableObject: + schema = self.registry.schema_for_protocol_version(ref.protocol_version) + return self.service.migrate_to_schema(obj, schema) + def _write( self, codec: ProtocolCodec, @@ -542,7 +556,5 @@ def _write( ref: DatasetRef, ) -> Path: """Downgrade ``obj`` to the ref's protocol schema, then let the codec persist it.""" - schema = self.registry.schema_for_protocol_version(ref.protocol_version) - downgraded = self.service.migrate_to_schema(obj, schema) - codec.write(path, downgraded) + codec.write(path, self._downgrade(obj, ref)) return path diff --git a/packages/syft-datasets/src/syft_datasets/models/dataset/v1.py b/packages/syft-datasets/src/syft_datasets/models/dataset/v1.py index 9c1a1c50bf2..5837f912ea8 100644 --- a/packages/syft-datasets/src/syft_datasets/models/dataset/v1.py +++ b/packages/syft-datasets/src/syft_datasets/models/dataset/v1.py @@ -6,7 +6,6 @@ from typing import ClassVar from uuid import UUID, uuid4 -import yaml from pydantic import Field, PrivateAttr from syft_migration import MigratableObject from syft_notebook_ui.formatter_mixin import ( @@ -23,7 +22,7 @@ from ...dataset_ref import DatasetRef from ...migrations import dataset_registry from ...url import SyftBoxURL -from ..private_dataset_config.v1 import PrivateDatasetConfigV1 +from ..private_dataset_config import PrivateDatasetConfig def _utcnow() -> datetime: @@ -99,25 +98,30 @@ def get_readme(self) -> str | None: def mock_dir(self) -> Path: return self._url_to_path(self.mock_url) - @property - def private_config_path(self) -> Path: + def _require_owner(self, what: str) -> None: if self.syftbox_config.email != self.owner: raise ValueError( - "Cannot access private config for a dataset owned by another user." + f"Cannot access {what} for a dataset owned by another user." ) + + @property + def private_config_path(self) -> Path: + self._require_owner("private config") return self._private_metadata_dir / PRIVATE_METADATA_FILENAME @cached_property - def private_config(self) -> PrivateDatasetConfigV1: - config_path = self.private_config_path - if not config_path.exists(): - raise FileNotFoundError( - f"Private dataset config not found at {config_path}" - ) - data = yaml.safe_load(config_path.read_text()) or {} - data.setdefault("canonical_name", "PrivateDatasetConfig") - data.setdefault("version", "1") - return PrivateDatasetConfigV1(**data) + def private_config(self) -> PrivateDatasetConfig: + """This dataset's private config, upgraded to the latest version. + + Raises ValueError if another user owns the dataset, and + PrivateConfigNotFoundError if the file is absent. + """ + # Local import: dataset_storage imports this module. + from ...dataset_storage import DatasetStorage + + self._require_owner("private config") + storage = DatasetStorage(config=self.syftbox_config) + return storage.read_private_config(self._ref) @property def private_dir(self) -> Path: @@ -136,10 +140,7 @@ def private_dir(self) -> Path: @property def _private_metadata_dir(self) -> Path: - if self.syftbox_config.email != self.owner: - raise ValueError( - "Cannot access private data for a dataset owned by another user." - ) + self._require_owner("private data") return self.private_dir @property diff --git a/packages/syft-datasets/src/syft_datasets/protocolcodecs/base.py b/packages/syft-datasets/src/syft_datasets/protocolcodecs/base.py index 770d97e22c6..9d188295ea7 100644 --- a/packages/syft-datasets/src/syft_datasets/protocolcodecs/base.py +++ b/packages/syft-datasets/src/syft_datasets/protocolcodecs/base.py @@ -2,6 +2,7 @@ from pathlib import Path from typing import Iterator +import yaml from syft_migration import MigratableObject from ..config import SyftBoxConfig @@ -51,4 +52,20 @@ def iter_dataset_refs(self, datasite_email: str) -> Iterator[DatasetRef]: def read(self, path: Path, canonical_name: str) -> dict: ... @abstractmethod - def write(self, path: Path, obj: MigratableObject) -> None: ... + def _data_for_disk(self, obj: MigratableObject) -> dict: + """The fields to serialize for this layout. + + The one place a codec decides its on-disk shape. Abstract on purpose: a + codec states its shaping rather than inheriting a default that may not + match the layout it speaks. + """ + ... + + def dumps(self, obj: MigratableObject) -> str: + """YAML for ``obj`` in this codec's on-disk format.""" + return yaml.safe_dump(self._data_for_disk(obj), indent=2, sort_keys=False) + + def write(self, path: Path, obj: MigratableObject) -> None: + """Persist ``obj`` at ``path`` in this codec's on-disk format.""" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(self.dumps(obj)) diff --git a/packages/syft-datasets/src/syft_datasets/protocolcodecs/v0.py b/packages/syft-datasets/src/syft_datasets/protocolcodecs/v0.py index 36091c48d21..bf22326f697 100644 --- a/packages/syft-datasets/src/syft_datasets/protocolcodecs/v0.py +++ b/packages/syft-datasets/src/syft_datasets/protocolcodecs/v0.py @@ -74,10 +74,9 @@ def read(self, path: Path, canonical_name: str) -> dict: data.setdefault("version", "1") return data - def write(self, path: Path, obj: MigratableObject) -> None: + def _data_for_disk(self, obj: MigratableObject) -> dict: data = obj.disk_dict() # Byte-match the pre-versioning (<= 0.1.20) on-disk format. data.pop("canonical_name", None) data.pop("version", None) - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(yaml.safe_dump(data, indent=2, sort_keys=False)) + return data diff --git a/packages/syft-datasets/src/syft_datasets/protocolcodecs/v1.py b/packages/syft-datasets/src/syft_datasets/protocolcodecs/v1.py index b3ca28cf85b..d68e73ab446 100644 --- a/packages/syft-datasets/src/syft_datasets/protocolcodecs/v1.py +++ b/packages/syft-datasets/src/syft_datasets/protocolcodecs/v1.py @@ -98,7 +98,6 @@ def read(self, path: Path, canonical_name: str) -> dict: # Files already carry canonical_name/version on disk. return yaml.safe_load(path.read_text()) or {} - def write(self, path: Path, obj: MigratableObject) -> None: - data = obj.disk_dict() - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(yaml.safe_dump(data, indent=2, sort_keys=False)) + def _data_for_disk(self, obj: MigratableObject) -> dict: + # This layout keeps the identity fields on disk. + return obj.disk_dict() diff --git a/packages/syft-datasets/tests/migrations/p2p/test_current_protocol.py b/packages/syft-datasets/tests/migrations/p2p/test_current_protocol.py index 5d9d2f428b1..b149e65451a 100644 --- a/packages/syft-datasets/tests/migrations/p2p/test_current_protocol.py +++ b/packages/syft-datasets/tests/migrations/p2p/test_current_protocol.py @@ -170,6 +170,40 @@ def test_multi_version_write_for_mixed_audience(tmp_path: Path): assert mgr.get("demo")._ref.protocol_version == "1" +def _wire_private_metadata(mgr: SyftDatasetManager, name: str, protocol_version: str): + files = mgr.get_private_dataset_files(name, protocol_version=protocol_version) + metas = [ + content + for path, content in files.items() + if path.name == "private_metadata.yaml" + ] + assert len(metas) == 1 + return yaml.safe_load(metas[0]) + + +def test_private_files_wire_private_metadata_in_protocol_format(tmp_path: Path): + """get_private_dataset_files serializes private_metadata through the codec.""" + mock, private, readme = _create_dataset_files(tmp_path) + mgr = _dataset_manager(tmp_path) + + mgr.create(name="flat", mock_path=mock, private_path=private, readme_path=readme) + raw0 = _wire_private_metadata(mgr, "flat", "0") + assert "canonical_name" not in raw0 and "version" not in raw0 + assert "uid" in raw0 + assert not Path(str(raw0["data_dir"])).is_absolute() + + mgr.create( + name="nested", + mock_path=mock, + private_path=private, + readme_path=readme, + protocol_versions=["1"], + ) + raw1 = _wire_private_metadata(mgr, "nested", "1") + assert raw1["canonical_name"] == "PrivateDatasetConfig" and raw1["version"] == "1" + assert not Path(str(raw1["data_dir"])).is_absolute() + + def test_delete_removes_all_protocol_versions(tmp_path: Path): schema0 = dataset_registry.schema_for_protocol_version("0") schema1 = dataset_registry.schema_for_protocol_version("1") diff --git a/packages/syft-datasets/tests/migrations/p2p/test_protocol_codecs.py b/packages/syft-datasets/tests/migrations/p2p/test_protocol_codecs.py index 2863534e423..966e51517a9 100644 --- a/packages/syft-datasets/tests/migrations/p2p/test_protocol_codecs.py +++ b/packages/syft-datasets/tests/migrations/p2p/test_protocol_codecs.py @@ -18,7 +18,7 @@ from syft_datasets.dataset_storage import DatasetRef, DatasetStorage from syft_datasets.migrations import dataset_registry from syft_datasets.migrations.registry import DATASET_PROTOCOL_VERSION -from syft_datasets.models import Dataset +from syft_datasets.models import Dataset, PrivateDatasetConfig from syft_datasets.protocolcodecs import ProtocolCodec from syft_datasets.protocolcodecs.v1 import DatasetConfigV1 from syft_datasets.url import SyftBoxURL @@ -123,6 +123,92 @@ def test_v0_writes_flat_no_identity_v1_nests_with_identity(tmp_path: Path): assert loaded.version == dataset_registry.latest_version("Dataset") +def test_wire_bytes_match_write_and_strip_private_config_identity_on_v0(tmp_path: Path): + storage = _storage(tmp_path) + ref0 = DatasetRef(DO_EMAIL, "flat", "0") + ref1 = DatasetRef(DO_EMAIL, "nested", "1") + dataset = _mock_dataset(storage, ref0) + config = PrivateDatasetConfig(uid=dataset.uid, data_dir=Path("/secret/local/path")) + + p0 = storage.write_private_config(ref0, config) + p1 = storage.write_private_config(ref1, config) + + assert storage.wire_bytes(ref0, config) == p0.read_bytes() + assert storage.wire_bytes(ref1, config) == p1.read_bytes() + + raw0 = yaml.safe_load(p0.read_text()) + assert "canonical_name" not in raw0 and "version" not in raw0 + assert list(raw0) == ["uid", "data_dir"] + + raw1 = yaml.safe_load(p1.read_text()) + assert raw1["canonical_name"] == "PrivateDatasetConfig" and raw1["version"] == "1" + + +@pytest.fixture +def private_config_v2(): + """A throwaway PrivateDatasetConfig V2, as the next release would ship it. + + Subclassing a registered class registers it, so this lands in the global + dataset_registry; the teardown pops it back out, because the registry has no + deregister call and another test must not see a version "2". + """ + from syft_datasets.models import PrivateDatasetConfigV1 + + # The class body registers V2, so everything after it needs the teardown. + class PrivateDatasetConfigV2(PrivateDatasetConfigV1, registry=dataset_registry): + version: str = "2" + checksum: str = "" + + try: + dataset_registry.register_migration( + canonical_name="PrivateDatasetConfig", + from_version="1", + to_version="2", + fn=lambda obj: PrivateDatasetConfigV2( + **obj.model_dump(exclude={"version"}) + ), + ) + yield PrivateDatasetConfigV2 + finally: + dataset_registry.objects["PrivateDatasetConfig"].pop("2", None) + dataset_registry.migrations.get("PrivateDatasetConfig", {}).pop( + ("1", "2"), None + ) + + +def test_dataset_private_config_upgrades_a_v1_file_to_the_latest_version( + tmp_path: Path, private_config_v2 +): + # Dataset.private_config must read through storage, which upgrades. Reading the + # file into PrivateDatasetConfigV1 by name would return version 1 forever. + storage = _storage(tmp_path) + ref = DatasetRef(DO_EMAIL, "demo", "1") + dataset = _mock_dataset(storage, ref) + storage.write_dataset_metadata(ref, dataset) + + path = storage.private_metadata_path(ref) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + yaml.safe_dump( + { + "canonical_name": "PrivateDatasetConfig", + "version": "1", + "uid": str(dataset.uid), + "data_dir": "", + }, + indent=2, + sort_keys=False, + ) + ) + + config = storage.read_dataset(ref).private_config + + assert isinstance(config, private_config_v2) + assert config.version == "2" + # The uid comes from the file, so a fresh V2 that ignored it would fail here. + assert config.uid == dataset.uid + + # -- behavior 2: each codec scans only its own layout -------------------------- def test_scan_partitions_by_layout(tmp_path: Path): storage = _storage(tmp_path)