From a4a020cd49459b1d7fd5d15837e177dd57939c29 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Thu, 10 Sep 2026 17:24:31 -0300 Subject: [PATCH] feat: Implement dataset upgrade functionality - Added `DatasetUpgrade` and `UpgradeReport` classes to manage dataset upgrades. - Implemented `upgrade()` method in `SyftRDSClient` to promote datasets to the current layout. - Enhanced protocol version handling in dataset creation and sharing processes. - Updated tests to cover new upgrade functionality, including edge cases for dataset sharing and audience recovery. - Refactored existing tests to align with the new protocol versioning logic. - Ensured backward compatibility with older protocol versions during dataset upgrades. --- .../docker/inference_server.py | 12 +- .../scripts/local_server.py | 12 +- .../src/enclave_model_api/paths.py | 55 +++- .../src/enclave_model_api/service.py | 14 +- .../tests/test_inference_server.py | 168 +++++++++- .../src/syft_datasets/dataset_manager.py | 59 +++- .../src/syft_datasets/dataset_storage.py | 121 +++++-- .../migrations/p2p/test_current_protocol.py | 207 +++++++++++- .../p2p/test_protocol_sanity_checks.py | 29 +- packages/syft-rds/src/syft_rds/__init__.py | 4 +- packages/syft-rds/src/syft_rds/client.py | 285 +++++++++++++++-- .../tests/test_create_dataset_cleanup.py | 5 +- .../syft-rds/tests/test_dataset_upgrade.py | 294 ++++++++++++++++++ .../p2p/test_dataset_multicopy_delivery.py | 141 ++++++++- .../p2p/test_dataset_schema_negotiation.py | 15 +- 15 files changed, 1320 insertions(+), 101 deletions(-) create mode 100644 packages/syft-rds/tests/test_dataset_upgrade.py diff --git a/packages/enclave-model-api-example/docker/inference_server.py b/packages/enclave-model-api-example/docker/inference_server.py index 2ba08923c29..cad573a4b60 100644 --- a/packages/enclave-model-api-example/docker/inference_server.py +++ b/packages/enclave-model-api-example/docker/inference_server.py @@ -9,7 +9,11 @@ from attestation_server import app from syft_enclaves.settings import EnclaveSettings -from enclave_model_api.paths import default_syftbox_folder, private_dataset_dir +from enclave_model_api.paths import ( + default_syftbox_folder, + private_dataset_dir, + resolve_weights_dir, +) from enclave_model_api.server import build_router from enclave_model_api.service import InferenceService from enclave_model_api.settings import InferenceSettings @@ -32,9 +36,13 @@ service = InferenceService( backend=backend, model_size=inference.model_size, - weights_dir=private_dataset_dir( + # Resolved on each poll: the model owner writes the layout its own release + # decided, and this process starts before the weights arrive. + weights_dir=lambda: resolve_weights_dir( syftbox_folder, inference.model_owner, inference.model_dataset ), + # The enclave owns its logs dataset and writes it, so the current layout is + # the right answer while it is absent. logs_dir=private_dataset_dir( syftbox_folder, settings.email, inference.logs_dataset ), diff --git a/packages/enclave-model-api-example/scripts/local_server.py b/packages/enclave-model-api-example/scripts/local_server.py index 722c79a2389..b5676054772 100644 --- a/packages/enclave-model-api-example/scripts/local_server.py +++ b/packages/enclave-model-api-example/scripts/local_server.py @@ -13,7 +13,11 @@ import uvicorn from syft_enclaves.settings import EnclaveSettings -from enclave_model_api.paths import default_syftbox_folder, private_dataset_dir +from enclave_model_api.paths import ( + default_syftbox_folder, + private_dataset_dir, + resolve_weights_dir, +) from enclave_model_api.server import create_app from enclave_model_api.service import InferenceService from enclave_model_api.settings import InferenceSettings @@ -33,7 +37,11 @@ service = InferenceService( backend=backend, model_size=inf.model_size, - weights_dir=private_dataset_dir(folder, inf.model_owner, inf.model_dataset), + # Resolved on each poll: the model owner writes the layout its own release + # decided, and this process starts before the weights arrive. + weights_dir=lambda: resolve_weights_dir(folder, inf.model_owner, inf.model_dataset), + # The enclave owns its logs dataset and writes it, so the current layout is + # the right answer while it is absent. logs_dir=private_dataset_dir(folder, settings.email, inf.logs_dataset), ) service.start_polling() diff --git a/packages/enclave-model-api-example/src/enclave_model_api/paths.py b/packages/enclave-model-api-example/src/enclave_model_api/paths.py index 7f22802a112..ef03369c0c2 100644 --- a/packages/enclave-model-api-example/src/enclave_model_api/paths.py +++ b/packages/enclave-model-api-example/src/enclave_model_api/paths.py @@ -18,21 +18,56 @@ def default_syftbox_folder(email: str) -> Path: 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. +def candidate_private_dataset_dirs( + storage: DatasetStorage, owner: str, name: str +) -> list[Path]: + """The private dirs a dataset could occupy, newest layout first. - A dataset may live at protocol 0 (flat) or under a ``v`` 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. + One entry, the layout on disk, once the dataset is there. While it is + absent, one entry for each protocol layout this client reads: the owner + writes the layout its own release and audience decided, and a dataset that + has not arrived cannot tell us which that is. A reader that waits must + therefore watch them all. """ 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) + return [ + storage.private_dataset_dir( + DatasetRef(owner=owner, name=name, protocol_version=protocol_version) + ) + for protocol_version in storage.supported_protocol_versions + ] + return [storage.private_dataset_dir(ref)] + + +def resolve_private_dataset_dir(storage: DatasetStorage, owner: str, name: str) -> Path: + """Private dir at the dataset's on-disk protocol layout. + + The current layout while the dataset is absent, which is the layout this + client writes for a dataset of its own. A reader waiting for a dataset that + another datasite writes must use ``candidate_private_dataset_dirs``, because + that owner may write an older layout. + """ + return candidate_private_dataset_dirs(storage, owner, name)[0] + + +def resolve_weights_dir( + syftbox_folder: Path | str, datasite: str, dataset_name: str +) -> Path: + """The layout that holds the synced weights, or the newest candidate so far. + + Re-resolved on each poll, not fixed at startup: the weights arrive in the + layout their owner writes, and an owner on an earlier release writes the + flat one. + """ + config = SyftBoxConfig(syftbox_folder=Path(syftbox_folder), email=datasite) + storage = DatasetStorage(config=config) + candidates = candidate_private_dataset_dirs(storage, datasite, dataset_name) + for path in candidates: + if weights_ready(path): + return path + return candidates[0] def private_dataset_dir( diff --git a/packages/enclave-model-api-example/src/enclave_model_api/service.py b/packages/enclave-model-api-example/src/enclave_model_api/service.py index dcde9276c86..48af7d110c4 100644 --- a/packages/enclave-model-api-example/src/enclave_model_api/service.py +++ b/packages/enclave-model-api-example/src/enclave_model_api/service.py @@ -9,6 +9,7 @@ import threading import time from pathlib import Path +from typing import Callable from enclave_model_api.log_writer import append_log_record, build_log_record from enclave_model_api.paths import weights_ready @@ -21,16 +22,25 @@ def __init__( self, backend, model_size: str, - weights_dir: Path | str, + weights_dir: Path | str | Callable[[], Path], logs_dir: Path | str, ): self.backend = backend self.model_size = model_size - self.weights_dir = Path(weights_dir) + # A callable is re-resolved on each read. The weights arrive in the + # layout their owner writes, so a path fixed at startup can watch a + # layout the owner never writes, and the poll then never ends. + self._weights_dir = ( + weights_dir if callable(weights_dir) else (lambda: Path(weights_dir)) + ) self.logs_dir = Path(logs_dir) self._loaded = None self._lock = threading.Lock() + @property + def weights_dir(self) -> Path: + return Path(self._weights_dir()) + @property def loaded(self) -> bool: return self._loaded is not None diff --git a/packages/enclave-model-api-example/tests/test_inference_server.py b/packages/enclave-model-api-example/tests/test_inference_server.py index 196b73a10a4..5866ea319de 100644 --- a/packages/enclave-model-api-example/tests/test_inference_server.py +++ b/packages/enclave-model-api-example/tests/test_inference_server.py @@ -1,6 +1,8 @@ """Unit tests for the inference service pieces: paths, log writer, FastAPI app.""" +import ast import json +from pathlib import Path from fastapi.testclient import TestClient @@ -9,9 +11,17 @@ append_log_record, build_log_record, ) -from enclave_model_api.paths import private_dataset_dir, weights_ready +from enclave_model_api.paths import ( + candidate_private_dataset_dirs, + private_dataset_dir, + resolve_weights_dir, + weights_ready, +) from enclave_model_api.server import create_app from enclave_model_api.service import InferenceService +from syft_datasets.config import SyftBoxConfig +from syft_datasets.dataset_storage import DatasetStorage +from syft_datasets.migrations.registry import DATASET_PROTOCOL_VERSION from inference_stub import STUB_COMPLETION_PREFIX, StubBackend, make_stub_weights @@ -44,16 +54,46 @@ def test_weights_ready(tmp_path): def test_private_dataset_dir_layout(tmp_path): + # A dataset not yet on disk falls back to the current protocol, which is the + # layout this client writes for a dataset of its own (the enclave's logs). + # A dataset another datasite writes may arrive in an older layout, and + # candidate_private_dataset_dirs covers that case. path = private_dataset_dir(tmp_path, "enclave@openmined.org", "inference_logs") assert path == ( tmp_path / "enclave@openmined.org" / "private" / "syft_datasets" + / f"v{DATASET_PROTOCOL_VERSION}" / "inference_logs" ) +def test_private_dataset_dir_follows_the_layout_on_disk(tmp_path): + # A dataset already on disk decides its own layout, whatever the fallback is. + flat = ( + tmp_path + / "enclave@openmined.org" + / "private" + / "syft_datasets" + / "inference_logs" + ) + flat.mkdir(parents=True) + (flat / "private_metadata.yaml").write_text("uid: x\n") + public = ( + tmp_path + / "enclave@openmined.org" + / "public" + / "syft_datasets" + / "inference_logs" + ) + public.mkdir(parents=True) + (public / "dataset.yaml").write_text("name: inference_logs\n") + + path = private_dataset_dir(tmp_path, "enclave@openmined.org", "inference_logs") + assert path == flat + + def test_inference_server_full_lifecycle(tmp_path): """503 before weights → load after weights sync → /infer logs each request.""" weights_dir = tmp_path / "weights" @@ -96,3 +136,129 @@ def test_inference_server_full_lifecycle(tmp_path): assert len(records) == 1 assert records[0]["prompt"] == "What is the capital of NL?" assert records[0]["completion"] == body["completion"] + + +def _storage(tmp_path, email: str) -> DatasetStorage: + return DatasetStorage(config=SyftBoxConfig(syftbox_folder=tmp_path, email=email)) + + +def _write_weights(private_dir): + private_dir.mkdir(parents=True, exist_ok=True) + make_stub_weights(private_dir) + + +def test_candidate_dirs_cover_every_layout_while_the_dataset_is_absent(tmp_path): + # The owner writes the layout its own release decided, so a reader that is + # still waiting cannot know which one arrives. + storage = _storage(tmp_path, "owner@test.org") + candidates = candidate_private_dataset_dirs(storage, "owner@test.org", "weights") + + assert [c.name for c in candidates] == ["weights"] * len(candidates) + segments = [c.parent.name for c in candidates] + # Newest layout first, down to the floor. + assert segments[0] == f"v{DATASET_PROTOCOL_VERSION}" + assert "syft_datasets" in segments + + +def test_candidate_dirs_follow_the_layout_on_disk(tmp_path): + storage = _storage(tmp_path, "owner@test.org") + flat_private = tmp_path / "owner@test.org" / "private" / "syft_datasets" / "weights" + flat_private.mkdir(parents=True) + (flat_private / "private_metadata.yaml").write_text("uid: x\n") + flat_public = tmp_path / "owner@test.org" / "public" / "syft_datasets" / "weights" + flat_public.mkdir(parents=True) + (flat_public / "dataset.yaml").write_text("name: weights\n") + + assert candidate_private_dataset_dirs(storage, "owner@test.org", "weights") == [ + flat_private + ] + + +def test_resolve_weights_dir_finds_the_flat_layout_of_an_earlier_release(tmp_path): + # Every data owner in the fleet today writes protocol 0, so the weights land + # flat. A guess fixed at the current layout never sees them. + owner = "owner@test.org" + flat = tmp_path / owner / "private" / "syft_datasets" / "weights" + _write_weights(flat) + + assert resolve_weights_dir(tmp_path, owner, "weights") == flat + assert weights_ready(resolve_weights_dir(tmp_path, owner, "weights")) + + +def test_resolve_weights_dir_returns_the_newest_candidate_while_absent(tmp_path): + owner = "owner@test.org" + resolved = resolve_weights_dir(tmp_path, owner, "weights") + assert resolved.parent.name == f"v{DATASET_PROTOCOL_VERSION}" + assert not weights_ready(resolved) + + +def test_the_service_sees_weights_that_land_in_an_older_layout(tmp_path): + # The poll exists so the enclave need not restart. It re-resolves the + # layout, so weights arriving flat after startup are picked up. + owner = "owner@test.org" + service = InferenceService( + backend=StubBackend(), + model_size="270m", + weights_dir=lambda: resolve_weights_dir(tmp_path, owner, "weights"), + logs_dir=tmp_path / "logs", + ) + assert not service.weights_present + + _write_weights(tmp_path / owner / "private" / "syft_datasets" / "weights") + + assert service.weights_present + assert service.try_load() + + +# --- entrypoint wiring ---------------------------------------------------- +# +# Neither entrypoint can be imported by a test: each starts a server at import +# time. The invariant they must hold is checked on their syntax tree instead. +# A path fixed at startup defeats the poll, and both entrypoints have to pass a +# callable that re-resolves the layout. + + +def _entrypoints(): + root = Path(__file__).resolve().parents[1] + return sorted((root / "scripts").glob("*.py")) + sorted( + (root / "docker").glob("inference_server.py") + ) + + +def _weights_argument(path: Path): + """The ``weights_dir`` argument of the InferenceService call in a module.""" + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "InferenceService" + ): + for keyword in node.keywords: + if keyword.arg == "weights_dir": + return keyword.value + return None + + +def test_every_entrypoint_re_resolves_the_weights_layout(): + checked = [] + for path in _entrypoints(): + argument = _weights_argument(path) + if argument is None: + continue + checked.append(path.name) + assert isinstance(argument, ast.Lambda), ( + f"{path.name} fixes the weights path at startup; the owner may write " + "an older layout, and the poll would never see it" + ) + called = { + n.func.id + for n in ast.walk(argument) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) + } + assert "resolve_weights_dir" in called, ( + f"{path.name} must resolve the weights layout through " + "resolve_weights_dir, which probes every layout it supports" + ) + # Both entrypoints, or the glob stopped matching. + assert sorted(checked) == ["inference_server.py", "local_server.py"] diff --git a/packages/syft-datasets/src/syft_datasets/dataset_manager.py b/packages/syft-datasets/src/syft_datasets/dataset_manager.py index a77f2499fd6..ca3077d698e 100644 --- a/packages/syft-datasets/src/syft_datasets/dataset_manager.py +++ b/packages/syft-datasets/src/syft_datasets/dataset_manager.py @@ -14,6 +14,8 @@ from syft_datasets.dataset_storage import DatasetSourceFiles, DatasetStorage from .config import PRIVATE_METADATA_FILENAME, SyftBoxConfig +from syft_permissions.spec.ruleset import PERMISSION_FILE_NAME, RuleSet + from .permissions import set_mock_dataset_permissions, set_private_dataset_permissions DATASET_COLLECTION_PREFIX = "syft_datasetcollection" @@ -33,8 +35,8 @@ def __init__( ) # peer_schemas (peer email -> dataset ProtocolSchema): syft # passes PeerManager's live map here (updated in place as peer version - # files load). Peers without an entry resolve to the widest-compatible - # protocol, so datasets stay readable by unknown-version peers. + # files load). A named peer without an entry resolves to the floor, so + # the dataset stays readable by a peer of unknown version. self.storage = DatasetStorage( config=self.syftbox_config, peer_schemas=peer_schemas ) @@ -141,15 +143,62 @@ def migrate( """Rewrite an owned dataset into another protocol layout, re-applying permissions. Storage copies the files + writes metadata for the new layout; the manager - re-applies read permissions. The audience (``users``) must be supplied by - the caller as it is on create — granted readers are not recoverable from - disk via the permissions API. + re-applies read permissions. ``users`` is the audience of the new layout. + A caller that no longer holds it reads it back with + ``recover_audience_from_ruleset``, which covers the explicit grants; the + "any" flag lives on the transport, so a caller that shares collections + must check that too. """ ref = self.storage.find_dataset_ref(self.syftbox_config.email, name) migrated = self.storage.migrate_dataset(ref, to_version) self._set_new_dataset_permissions(dataset=migrated, users=users) return migrated + def recover_audience_from_ruleset( + self, name: str, protocol_version: str | None = None + ) -> list[str] | str: + """The explicit read grants recorded on a layout's mock directory. + + The audience of record for a dataset already on disk. ``"*"`` maps back + to ``SHARE_WITH_ANY``. An empty list is a real answer: a dataset kept by + the owner alone has no ruleset, and it must stay that way. + + Reads the layout ``protocol_version``, or the newest layout by default. + Covers the explicit grants only. A collection tagged for any peer + carries that flag on the transport, not here. + """ + ref = self.storage.find_dataset_ref( + self.syftbox_config.email, name, protocol_version=protocol_version + ) + ruleset_path = self.storage.public_dataset_dir(ref) / PERMISSION_FILE_NAME + if not ruleset_path.exists(): + return [] + emails: list[str] = [] + for rule in RuleSet.load(ruleset_path).rules: + for user in rule.access.read: + if user == "*": + return SHARE_WITH_ANY + if user not in emails: + emails.append(user) + return emails + + def grant_read_on_every_layout(self, name: str, users: list[str] | str) -> None: + """Add ``users`` to the ruleset of every on-disk layout of a dataset. + + A share reaches a peer through the transport, so a share that only + updates the transport leaves the rulesets naming the audience at create + time. The ruleset is the audience of record for ``upgrade()``, so every + layout records the grant. + """ + for ref in self.storage.iter_dataset_refs_all_protocols( + self.syftbox_config.email + ): + if ref.name != name: + continue + self._set_new_dataset_permissions( + dataset=self.storage.read_dataset(ref), users=users + ) + @staticmethod def _peer_emails(users: list[str] | str | None) -> list[str] | None: """Audience emails for protocol negotiation; None means no/any peers.""" diff --git a/packages/syft-datasets/src/syft_datasets/dataset_storage.py b/packages/syft-datasets/src/syft_datasets/dataset_storage.py index d87094c9bdb..15deada0dcc 100644 --- a/packages/syft-datasets/src/syft_datasets/dataset_storage.py +++ b/packages/syft-datasets/src/syft_datasets/dataset_storage.py @@ -87,9 +87,9 @@ class DatasetStorage: version for new datasets). The last codec is the current protocol. Unlike jobs (each copy targets one peer), a public dataset is a single copy - read by the whole audience, so new datasets are written at the version(s) the - audience can read (see ``target_protocol_versions_for_peers``), defaulting to - the widest-compatible (oldest) protocol when no peers are known. + read by the whole audience. A new dataset is therefore written at the current + protocol, plus one layout for each older version its audience reads (see + ``create_protocol_versions``). """ def __init__( @@ -102,9 +102,9 @@ def __init__( self.registry = registry self.service = MigrationService(registry=registry) # peer email -> dataset ProtocolSchema; syft passes PeerManager's - # live map here (updated in place as peer version files load). Peers - # without an entry cannot be assumed to read the current layout, so - # they resolve to the widest-compatible (oldest) protocol. + # live map here (updated in place as peer version files load). A named + # peer without an entry cannot be assumed to read the current layout, so + # it resolves to the floor. # `is not None`, not `or`: the live dict starts empty and `or {}` would # drop the shared reference, freezing negotiation at construction time. self.peer_schemas: dict[str, ProtocolSchema] = ( @@ -130,9 +130,30 @@ def _codec_for(self, protocol_version: str) -> ProtocolCodec: return codec @property - def _widest_protocol_version(self) -> str: - """The oldest protocol version any current client can read (widest compat).""" - return min(self._codec_by_protocol_version, key=int) + def _floor_protocol_version(self) -> str: + """The oldest dataset protocol this client still supports. + + The backfill target for a named peer whose version we have not read. It + is the floor and not ``min(codec versions)``: a codec that still reads a + layout we no longer support must not pull a new copy into that layout. + Raising the floor therefore retires a layout in one place. + """ + return self.registry.min_supported_protocol_version + + @property + def supported_protocol_versions(self) -> list[str]: + """Every protocol layout this client reads, newest first. + + The layout an owner wrote depends on that owner's release and audience, + so a reader that has not yet received a dataset cannot know which one it + will arrive in. It looks in each of these. + """ + floor = self._floor_protocol_version + return sorted( + (v for v in self._codec_by_protocol_version if int(v) >= int(floor)), + key=int, + reverse=True, + ) # -- peers / protocol ---------------------------------------------------- def negotiated_protocol_version_for_peer( @@ -169,30 +190,46 @@ def negotiated_protocol_version_for_peer( def target_protocol_versions_for_peers( self, peer_emails: Optional[list[str]] = None ) -> set[str]: - """The set of protocol versions to write so every peer can read a copy. + """The older layouts an audience needs, beyond the current protocol. - A dataset is written once per distinct version in the audience. A known - peer contributes ``min(ours, theirs)``; an unknown peer (or no audience) - contributes the widest-compatible protocol, since we cannot assume it - can read a newer layout. + A known peer contributes ``min(ours, theirs)``. A named peer with no + known schema contributes the floor, because we cannot assume it reads a + newer layout. An empty audience contributes nothing: with no peer to + serve, a create writes the current layout alone, and a peer that arrives + later gets its layout from the backfill. The two unknown-peer answers differ on purpose. This method serves an - audience. An unknown peer therefore takes the widest protocol, and every - reader can read a copy. ``negotiated_protocol_version_for_peer`` serves - one peer, so an unknown peer takes the current protocol. The caller of - that method accepts the risk when it passes ``raise_on_unknown=False``. + audience, so a named peer of unknown version takes the floor and can + read its copy. ``negotiated_protocol_version_for_peer`` serves one peer, + so an unknown peer takes the current protocol. The caller of that method + accepts the risk when it passes ``raise_on_unknown=False``. """ if not peer_emails: - return {self._widest_protocol_version} + return set() versions: set[str] = set() for email in peer_emails: schema = self.peer_schemas.get(email) if schema is not None: versions.add(min(DATASET_PROTOCOL_VERSION, schema.version, key=int)) else: - versions.add(self._widest_protocol_version) + versions.add(self._floor_protocol_version) return versions + def create_protocol_versions( + self, peer_emails: Optional[list[str]] = None + ) -> list[str]: + """The layouts a create writes: the current one, plus what the audience reads. + + The current layout is always written, so a fresh dataset never lands in a + layout the installed client has moved past, and ``upgrade()`` stays a + one-off for each release instead of recurring maintenance. + """ + return sorted( + {DATASET_PROTOCOL_VERSION} + | self.target_protocol_versions_for_peers(peer_emails), + key=int, + ) + def new_dataset_ref(self, name: str, protocol_version: str) -> DatasetRef: """A ref for a new dataset owned by the current user.""" return DatasetRef( @@ -217,9 +254,9 @@ def create_dataset( Copies the source files into each version's on-disk layout and writes the metadata/private config. Returns {protocol_version: written Dataset}. - By default the versions are inferred from ``peer_emails`` (no/unknown peers - => the widest-compatible protocol). Pass ``protocol_versions`` to write - exactly those versions instead, skipping inference. + By default the versions are the current protocol plus one layout for each + older version the audience reads (``create_protocol_versions``). Pass + ``protocol_versions`` to write exactly those versions instead. """ self.validate_dataset_name(name) if source.mock.is_dir() and (source.mock / METADATA_FILENAME).exists(): @@ -228,9 +265,7 @@ def create_dataset( f"{METADATA_FILENAME}. Please rename it and try again." ) if protocol_versions is None: - protocol_versions = list( - self.target_protocol_versions_for_peers(peer_emails) - ) + protocol_versions = self.create_protocol_versions(peer_emails) now = _utcnow() fields = _DatasetFields( uid=uid or uuid4(), @@ -253,9 +288,20 @@ def migrate_dataset(self, ref: DatasetRef, target_protocol_version: str) -> Data Copies the source layout's files into the target layout and writes the metadata/private config there, preserving identity (uid/timestamps). Does not delete the source copy. Owner-only. + + Idempotent and resumable: a complete target layout is returned as it is, + and the debris of an interrupted copy is removed before the copy runs + again. A migrate to the layout the ref already holds reads it back. """ if ref.owner != self.config.email: raise ValueError("Can only migrate datasets you own.") + if ref.protocol_version == target_protocol_version: + # Nothing to copy, and the source must never be cleared as debris. + return self.read_dataset(ref) + target_ref = self.new_dataset_ref(ref.name, target_protocol_version) + complete = self._reuse_or_clear_target(target_ref) + if complete is not None: + return complete old = self.read_dataset(ref) source = DatasetSourceFiles( mock=self.public_dataset_dir(ref), @@ -271,11 +317,32 @@ def migrate_dataset(self, ref: DatasetRef, target_protocol_version: str) -> Data location=old.location, tags=old.tags, ) - target_ref = self.new_dataset_ref(ref.name, target_protocol_version) return self._materialize_version( target_ref, fields, source, exclude_names=_NON_PAYLOAD_FILES ) + def _reuse_or_clear_target(self, target_ref: DatasetRef) -> Optional[Dataset]: + """The dataset already complete at ``target_ref``, else None once cleared. + + ``_materialize_version`` copies the payload first and writes both metadata + files last, so a target that holds them is a finished copy. A target + without them is the debris of an interrupted run, which no scan can see + (``iter_dataset_refs`` needs ``dataset.yaml``) and which would make + ``_copy_mock_data`` raise. It is removed so the copy can run again. + """ + if ( + self.metadata_path(target_ref).exists() + and self.private_metadata_path(target_ref).exists() + ): + return self.read_dataset(target_ref) + for directory in ( + self.public_dataset_dir(target_ref), + self.private_dataset_dir(target_ref), + ): + if directory.exists(): + shutil.rmtree(directory) + return None + def _materialize_version( self, ref: DatasetRef, 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..2a393735485 100644 --- a/packages/syft-datasets/tests/migrations/p2p/test_current_protocol.py +++ b/packages/syft-datasets/tests/migrations/p2p/test_current_protocol.py @@ -1,16 +1,17 @@ """End-to-end dataset flow through SyftDatasetManager. -Datasets are broadcast (one public copy read by the whole audience), so by -default they are written in the widest-compatible (oldest) protocol so every -current peer can read them; the new v layout is written only for peers that -advertise support for it. +Datasets are broadcast: one public copy is read by the whole audience. A create +therefore writes the current protocol, plus one older layout for each version +its audience reads. A peer that arrives after the create gets its layout from +the backfill, at the share. """ from pathlib import Path import yaml -from syft_datasets.dataset_manager import SyftDatasetManager +from syft_datasets.dataset_manager import SHARE_WITH_ANY, SyftDatasetManager from syft_datasets.migrations import dataset_registry +from syft_datasets.migrations.registry import DATASET_PROTOCOL_VERSION DO_EMAIL = "do@test.org" DS0 = "old@test.org" @@ -39,7 +40,8 @@ def _dataset_manager(tmp_path: Path, peer_schemas=None) -> SyftDatasetManager: def test_create_with_explicit_protocol_versions_skips_inference(tmp_path: Path): - # No peers => inference would write only protocol 0; explicit versions override. + # No peers => the default writes the current layout. The explicit list + # overrides that, so an older layout alone is what lands. mgr = _dataset_manager(tmp_path) mock, private, readme = _create_dataset_files(tmp_path) @@ -48,21 +50,29 @@ def test_create_with_explicit_protocol_versions_skips_inference(tmp_path: Path): mock_path=mock, private_path=private, readme_path=readme, - protocol_versions=["1"], + protocol_versions=["0"], ) public_root = mgr.syftbox_config.datasite_public_root(DO_EMAIL) / "syft_datasets" - # Exactly the requested version is written: v1 layout, not the flat default. - assert (public_root / "v1" / "demo" / "dataset.yaml").exists() - assert not (public_root / "demo").exists() - assert mgr.get("demo")._ref.protocol_version == "1" + # Exactly the requested version is written: the flat layout, not the + # current one the default would have chosen. + assert (public_root / "demo" / "dataset.yaml").exists() + assert not (public_root / f"v{DATASET_PROTOCOL_VERSION}").exists() + assert mgr.get("demo")._ref.protocol_version == "0" def test_migrate_dataset_v0_to_v1_preserves_identity(tmp_path: Path): - # No peers => created in the widest-compatible (flat, protocol 0) layout. + # An explicit protocol 0 create, standing in for a dataset written by an + # earlier release: the flat layout with no v segment. mgr = _dataset_manager(tmp_path) mock, private, readme = _create_dataset_files(tmp_path) - mgr.create(name="demo", mock_path=mock, private_path=private, readme_path=readme) + mgr.create( + name="demo", + mock_path=mock, + private_path=private, + readme_path=readme, + protocol_versions=["0"], + ) ref0 = mgr.storage.find_dataset_ref(DO_EMAIL, "demo") assert ref0.protocol_version == "0" @@ -99,17 +109,100 @@ def test_migrate_dataset_v0_to_v1_preserves_identity(tmp_path: Path): assert all_datasets[0]._ref.protocol_version == "1" -def test_default_create_writes_protocol_0(tmp_path: Path): +def _v0_dataset(tmp_path: Path) -> SyftDatasetManager: + """A manager holding one dataset in the flat protocol 0 layout.""" + mgr = _dataset_manager(tmp_path) + mock, private, readme = _create_dataset_files(tmp_path) + mgr.create( + name="demo", + mock_path=mock, + private_path=private, + readme_path=readme, + protocol_versions=["0"], + ) + return mgr + + +def test_a_second_migrate_returns_the_existing_copy(tmp_path: Path): + # A sweep can run twice, and a migrate must not raise the second time. + mgr = _v0_dataset(tmp_path) + ref0 = mgr.storage.find_dataset_ref(DO_EMAIL, "demo", protocol_version="0") + + first = mgr.storage.migrate_dataset(ref0, "1") + second = mgr.storage.migrate_dataset(ref0, "1") + + assert second.uid == first.uid + assert second.created_at == first.created_at + assert second._ref.protocol_version == "1" + # Still one copy for each layout, and the source survives. + versions = sorted( + r.protocol_version + for r in mgr.storage.iter_dataset_refs_all_protocols(DO_EMAIL) + if r.name == "demo" + ) + assert versions == ["0", "1"] + + +def test_a_migrate_over_an_interrupted_copy_succeeds(tmp_path: Path): + # The payload is copied before the metadata is written, so an interruption + # leaves a target that no scan can see and that a plain copy refuses. + mgr = _v0_dataset(tmp_path) + ref0 = mgr.storage.find_dataset_ref(DO_EMAIL, "demo", protocol_version="0") + target_ref = mgr.storage.new_dataset_ref("demo", "1") + + debris = mgr.storage.public_dataset_dir(target_ref) + debris.mkdir(parents=True) + (debris / "mock.csv").write_text("half a copy\n") + assert not mgr.storage.metadata_path(target_ref).exists() + + migrated = mgr.storage.migrate_dataset(ref0, "1") + + assert migrated._ref.protocol_version == "1" + assert mgr.storage.metadata_path(target_ref).exists() + # The debris is replaced by a copy of the source, not merged with it. + assert (debris / "mock.csv").read_text() == "id,value\n1,10\n" + assert mgr.storage.read_private_config(target_ref).uid == migrated.uid + + +def test_a_migrate_over_a_target_with_no_private_config_redoes_the_copy( + tmp_path: Path, +): + # The public metadata is written before the private config, so a target with + # only the first one is also incomplete. + mgr = _v0_dataset(tmp_path) + ref0 = mgr.storage.find_dataset_ref(DO_EMAIL, "demo", protocol_version="0") + target_ref = mgr.storage.new_dataset_ref("demo", "1") + mgr.storage.migrate_dataset(ref0, "1") + mgr.storage.private_metadata_path(target_ref).unlink() + + mgr.storage.migrate_dataset(ref0, "1") + + assert mgr.storage.read_private_config(target_ref).uid == ( + mgr.storage.read_dataset(target_ref).uid + ) + + +def test_default_create_writes_the_current_protocol(tmp_path: Path): + # No audience, so no older layout is needed. A peer that arrives later gets + # its layout from the backfill, at the share. mgr = _dataset_manager(tmp_path) mock, private, readme = _create_dataset_files(tmp_path) dataset = mgr.create( name="demo", mock_path=mock, private_path=private, readme_path=readme ) - # Flat layout, no v, no identity fields (byte-compatible with 0.1.20). - assert dataset.mock_dir.parent.name == "syft_datasets" + # Versioned layout, with the identity fields protocol 1 carries. + assert dataset.mock_dir.parent.name == f"v{DATASET_PROTOCOL_VERSION}" raw = yaml.safe_load((dataset.mock_dir / "dataset.yaml").read_text()) - assert "canonical_name" not in raw + assert raw["canonical_name"] == "Dataset" + + # The current layout is the only one written. + versions = sorted( + r.protocol_version + for r in mgr.storage.iter_dataset_refs_all_protocols(DO_EMAIL) + if r.name == "demo" + ) + assert versions == [DATASET_PROTOCOL_VERSION] got = mgr.get("demo") assert got.name == "demo" @@ -200,3 +293,83 @@ def test_delete_removes_all_protocol_versions(tmp_path: Path): assert not (private_root / "demo").exists() assert not (private_root / "v1" / "demo").exists() assert mgr.get_all() == [] + + +def test_the_audience_reads_back_from_the_ruleset(tmp_path: Path): + # upgrade() must know who a dataset was shared with, and the ruleset of the + # source layout records the explicit grants. + schema0 = dataset_registry.schema_for_protocol_version("0") + mgr = _dataset_manager(tmp_path, peer_schemas={DS0: schema0}) + mock, private, readme = _create_dataset_files(tmp_path) + mgr.create( + name="demo", + mock_path=mock, + private_path=private, + readme_path=readme, + users=[DS0], + ) + + assert mgr.recover_audience_from_ruleset("demo") == [DS0] + assert mgr.recover_audience_from_ruleset("demo", protocol_version="0") == [DS0] + + +def test_an_owner_only_dataset_recovers_an_empty_audience(tmp_path: Path): + # No ruleset is written when there is no audience, and that is a real + # answer: the dataset must stay owner-only. + mgr = _dataset_manager(tmp_path) + mock, private, readme = _create_dataset_files(tmp_path) + mgr.create(name="demo", mock_path=mock, private_path=private, readme_path=readme) + + assert mgr.recover_audience_from_ruleset("demo") == [] + + +def test_a_dataset_shared_with_any_recovers_the_any_marker(tmp_path: Path): + mgr = _dataset_manager(tmp_path) + mock, private, readme = _create_dataset_files(tmp_path) + mgr.create( + name="demo", + mock_path=mock, + private_path=private, + readme_path=readme, + users=SHARE_WITH_ANY, + ) + + assert mgr.recover_audience_from_ruleset("demo") == SHARE_WITH_ANY + + +def test_a_grant_reaches_the_ruleset_of_every_layout(tmp_path: Path): + # The drift fix: a share after create must record the new user on every + # layout, or upgrade() recovers the create-time audience and the promoted + # copy never reaches that user. + schema0 = dataset_registry.schema_for_protocol_version("0") + mgr = _dataset_manager(tmp_path, peer_schemas={DS0: schema0}) + mock, private, readme = _create_dataset_files(tmp_path) + mgr.create( + name="demo", + mock_path=mock, + private_path=private, + readme_path=readme, + users=[DS0], + ) + + mgr.grant_read_on_every_layout("demo", ["late@test.org"]) + + for protocol_version in ("0", "1"): + assert sorted( + mgr.recover_audience_from_ruleset("demo", protocol_version=protocol_version) + ) == sorted([DS0, "late@test.org"]) + + +def test_a_migrate_to_the_same_layout_keeps_the_source(tmp_path: Path): + # The debris check clears an incomplete target, so a migrate onto the + # source's own layout must never reach it. + mgr = _v0_dataset(tmp_path) + ref0 = mgr.storage.find_dataset_ref(DO_EMAIL, "demo", protocol_version="0") + mgr.storage.private_metadata_path(ref0).unlink() + + same = mgr.storage.migrate_dataset(ref0, "0") + + assert same.name == "demo" + assert mgr.storage.metadata_path(ref0).exists() + assert (mgr.storage.public_dataset_dir(ref0) / "mock.csv").exists() + assert (mgr.storage.private_dataset_dir(ref0) / "private.csv").exists() diff --git a/packages/syft-datasets/tests/migrations/p2p/test_protocol_sanity_checks.py b/packages/syft-datasets/tests/migrations/p2p/test_protocol_sanity_checks.py index 34b3a437ea4..2cfbf50bab7 100644 --- a/packages/syft-datasets/tests/migrations/p2p/test_protocol_sanity_checks.py +++ b/packages/syft-datasets/tests/migrations/p2p/test_protocol_sanity_checks.py @@ -53,11 +53,34 @@ def test_target_protocol_versions_for_peers(tmp_path: Path): tmp_path, peer_schemas={"old@test.org": schema0, "new@test.org": schema1} ) - # No audience -> widest-compatible (oldest) protocol. - assert storage.target_protocol_versions_for_peers() == {"0"} - # Unknown peer -> also widest-compatible. + # No audience -> nothing beyond the current layout. + assert storage.target_protocol_versions_for_peers() == set() + # A named peer of unknown version -> the floor, which it can read. assert storage.target_protocol_versions_for_peers(["stranger@test.org"]) == {"0"} # Mixed audience -> a copy per distinct version. assert storage.target_protocol_versions_for_peers( ["old@test.org", "new@test.org"] ) == {"0", "1"} + + +def test_create_protocol_versions(tmp_path: Path): + schema0 = dataset_registry.schema_for_protocol_version("0") + storage = _storage(tmp_path, peer_schemas={"old@test.org": schema0}) + + # No audience -> the current layout alone. + assert storage.create_protocol_versions() == [DATASET_PROTOCOL_VERSION] + # An older peer in the audience -> its layout as well as the current one. + assert storage.create_protocol_versions(["old@test.org"]) == [ + "0", + DATASET_PROTOCOL_VERSION, + ] + + +def test_the_floor_is_what_we_support_not_what_we_can_read(tmp_path: Path): + # A codec may still read a layout the floor has retired. The backfill target + # follows the floor, so no new copy lands in a retired layout. + storage = _storage(tmp_path) + assert storage._floor_protocol_version == ( + dataset_registry.min_supported_protocol_version + ) + assert "0" in storage._codec_by_protocol_version diff --git a/packages/syft-rds/src/syft_rds/__init__.py b/packages/syft-rds/src/syft_rds/__init__.py index d4826ca6b8a..3954622f053 100644 --- a/packages/syft-rds/src/syft_rds/__init__.py +++ b/packages/syft-rds/src/syft_rds/__init__.py @@ -1,6 +1,6 @@ """syft-rds: Remote Data Science product composed on top of syft.""" -from syft_rds.client import SyftRDSClient +from syft_rds.client import DatasetUpgrade, SyftRDSClient, UpgradeReport from syft_rds.config import SyftRDSClientConfig from syft_rds.job_auto_approval import auto_approve_and_run_jobs, job_matches_criteria from syft_rds.login import login_do, login_ds @@ -15,6 +15,8 @@ __all__ = [ "SyftRDSClient", + "UpgradeReport", + "DatasetUpgrade", "SyftRDSClientConfig", "login_do", "login_ds", diff --git a/packages/syft-rds/src/syft_rds/client.py b/packages/syft-rds/src/syft_rds/client.py index 3af8076ccd6..b803a8a5db0 100644 --- a/packages/syft-rds/src/syft_rds/client.py +++ b/packages/syft-rds/src/syft_rds/client.py @@ -19,8 +19,9 @@ from syft.sync.version.peer_manager import CompatAction from syft_job.client import JobClient from syft_job.job_runner import SyftJobRunner -from syft_datasets.dataset_manager import SyftDatasetManager +from syft_datasets.dataset_manager import SHARE_WITH_ANY, SyftDatasetManager from syft_datasets.dataset_ref import DatasetNotFoundError +from syft_datasets.migrations.registry import DATASET_PROTOCOL_VERSION from syft_rds.config import ( DATASET_COLLECTION_SPECS, MOCK_DATASET_SPEC, @@ -31,6 +32,64 @@ logger = logging.getLogger(__name__) +_ALREADY_PUBLISHED = f"already published at protocol {DATASET_PROTOCOL_VERSION}" + + +class DatasetUpgrade(BaseModel): + """What ``upgrade()`` did, or would do, for one dataset.""" + + tag: str + published_before: list[str] + published_after: list[str] + upload_bytes: int = 0 + skipped_reason: str | None = None + error: str | None = None + + @property + def promoted(self) -> bool: + return self.skipped_reason is None and self.error is None + + def __str__(self) -> str: + if self.error is not None: + return f"{self.tag}: failed ({self.error})" + if self.skipped_reason is not None: + return f"{self.tag}: skipped ({self.skipped_reason})" + before = ",".join(self.published_before) or "none" + after = ",".join(self.published_after) or "none" + return f"{self.tag}: {before} -> {after} ({self.upload_bytes} bytes)" + + +class UpgradeReport(BaseModel): + """The result of one ``upgrade()`` sweep.""" + + dry_run: bool + datasets: list[DatasetUpgrade] = [] + + @property + def promoted(self) -> list[DatasetUpgrade]: + return [d for d in self.datasets if d.promoted] + + @property + def failed(self) -> list[DatasetUpgrade]: + return [d for d in self.datasets if d.error is not None] + + @property + def upload_bytes(self) -> int: + return sum(d.upload_bytes for d in self.datasets) + + def __str__(self) -> str: + head = "upgrade (dry run)" if self.dry_run else "upgrade" + lines = [ + f"{head}: {len(self.promoted)} promoted, {len(self.failed)} failed, " + f"{self.upload_bytes} bytes" + ] + lines += [f" {d}" for d in self.datasets] + return "\n".join(lines) + + def _repr_html_(self) -> str: + rows = "".join(f"{d}" for d in self.datasets) + return f"{self}{rows}
" if rows else f"{self}" + class SyftRDSClient(BaseModel): # Holds live service objects (sync engine + RDS-owned managers), not @@ -252,24 +311,44 @@ def _share_any_datasets_with_peer(self, peer_email: str) -> None: Google Drive "anyone with link" files are not discoverable via search, so explicit user sharing is added. Reads the cache populated during ``pull_initial_state()`` in the nested DatasiteOwnerSyncer. + + An "any" dataset holds the layouts its audience read at create time, so + a peer approved later may read none of them. Its layout is materialized + first, because a grant on a collection the peer cannot list reaches + nobody. The new layout is then marked "any" like the rest of the tag, + so the next peer finds it in the cache instead of a fresh listing. """ - for ( - wire_prefix, - tag, - content_hash, - ) in self.sync_engine.datasite_owner_syncer.any_shared_collections: + cached = self.sync_engine.datasite_owner_syncer.any_shared_collections + # The cache holds one entry for each layout, so the backfill is keyed on + # the tag and runs once for each dataset. + for tag in dict.fromkeys(tag for _, tag, _ in cached): try: - self.sync_engine.share_collection( - wire_prefix, tag, content_hash, [peer_email] + self._ensure_dataset_layouts_for( + tag, + [peer_email], + {self._protocol_of(c) for c in self._mock_collections_for(tag)}, ) + # Listed again, not read from the cache: a layout materialized + # just now is not in the cache yet, and it is the one this peer + # reads. + for collection in self._mock_collections_for(tag): + wire_prefix = MOCK_DATASET_SPEC.wire_prefix(collection.variant) + if not collection.has_any_permission: + self.sync_engine.tag_collection_as_any( + wire_prefix, tag, collection.content_hash + ) + self.sync_engine.datasite_owner_syncer.register_any_shared_collection( + wire_prefix, tag, collection.content_hash + ) + self.sync_engine.share_collection( + wire_prefix, tag, collection.content_hash, [peer_email] + ) except Exception: - # One collection failing (missing folder, quota, network) must - # not stop us sharing the rest with this peer. "alreadyShared" - # is already handled in _batch_add_permissions, so anything + # One dataset failing (missing folder, quota, network) must not + # stop us sharing the rest with this peer. "alreadyShared" is + # already handled in _batch_add_permissions, so anything # reaching here is a real failure worth a traceback. - logger.exception( - "Failed to share collection %r with %s", tag, peer_email - ) + logger.exception("Failed to share dataset %r with %s", tag, peer_email) # ------------------------------------------------------------------ # # job product surface (RDS-owned) @@ -431,8 +510,10 @@ def create_dataset( private_folder_ids: list[str] = [] try: - # Create the dataset locally, in one layout for each protocol - # version the audience reads. + # Create the dataset locally: the current layout, plus one layout + # for each older protocol version the audience reads. Resolved here + # and not in storage, because "any" is an audience of the approved + # peers and only the client holds that list. created = self.dataset_manager.create_all( name=name, mock_path=mock_path, @@ -442,6 +523,9 @@ def create_dataset( location=location, tags=tags, users=users, + protocol_versions=self.dataset_manager.storage.create_protocol_versions( + self._audience_emails(users) + ), ) created_local = True # The newest copy is the one to hand back to the owner. @@ -685,6 +769,11 @@ def share_dataset(self, tag: str, users: list[str] | str, sync=True): users, ) + # A share changes the audience, so every layout's ruleset records it. + # The transport alone would leave the rulesets naming the create-time + # audience, and upgrade() recovers the audience from there. + self.dataset_manager.grant_read_on_every_layout(tag, users) + if sync: self.sync() @@ -713,6 +802,22 @@ def _protocol_of(collection) -> str: """The dataset protocol version a collection's wire variant stands for.""" return collection.variant.removeprefix("v") or "0" + def _audience_emails(self, users: list[str] | str | None) -> list[str]: + """The peers whose layouts an audience needs. + + ``SHARE_WITH_ANY`` is an audience of the approved peers, whose versions + we know, so it needs no floor copy. No audience is an empty list: only + the current layout is written, and a peer that arrives later gets its + layout from the backfill. + """ + if users == SHARE_WITH_ANY: + return [p.email for p in self.sync_engine.peer_manager.approved_peers] + if users is None: + return [] + if isinstance(users, str): + return [users] + return list(users) + def _ensure_dataset_layouts_for( self, tag: str, users: list[str] | str, existing_versions: set[str] ) -> bool: @@ -725,8 +830,9 @@ def _ensure_dataset_layouts_for( was added. """ storage = self.dataset_manager.storage - peer_emails = self.dataset_manager._peer_emails(users) - needed = storage.target_protocol_versions_for_peers(peer_emails) + needed = storage.target_protocol_versions_for_peers( + self._audience_emails(users) + ) missing = { version for version in needed @@ -769,6 +875,149 @@ def _materialize_dataset_copy( if self._private_collections_for(tag): self._upload_private_dataset_to_collection(copy) + def upgrade(self, *, dry_run: bool = False, sync: bool = True) -> UpgradeReport: + """Publish every owned dataset in the protocol layout of this client. + + A dataset's object versions follow its layout directory, so a dataset is + promoted by adding the current layout. Every older layout is kept, and a + peer that reads one keeps reading it. + + The audience of each dataset is recovered from the ruleset of its newest + layout, plus the "any" flag on its collections. A dataset kept by the + owner alone stays that way. + + One sweep for each release: a create writes the current layout, so only + the datasets of an earlier release need this. A repeat run re-shares + each layout with the recovered audience, so a sweep that failed part way + through repairs itself. ``dry_run`` reports the plan and writes nothing. + """ + if not self.has_do_role: + raise ValueError("Only dataset owners can upgrade datasets") + if self.dataset_manager is None: + raise ValueError("Dataset manager is not set") + + # Load the peers first: the audience of an "any" dataset is the approved + # peers, and a stale map would decide it. + self.load_peers() + + report = UpgradeReport(dry_run=dry_run) + storage = self.dataset_manager.storage + for ref in list(storage.iter_dataset_refs(self.email)): + report.datasets.append(self._upgrade_dataset(ref.name, dry_run=dry_run)) + + if sync and not dry_run: + # Housekeeping. The publish already happened, per dataset. + self.sync() + return report + + def _upgrade_dataset(self, tag: str, *, dry_run: bool) -> DatasetUpgrade: + """Promote one dataset to the current layout, or say why it was skipped.""" + published = sorted( + {self._protocol_of(c) for c in self._mock_collections_for(tag)}, key=int + ) + # Published only. A layout on disk with no collection of its own is the + # case _materialize_dataset_copy resumes, so folding the disk into this + # skip would drop exactly that case and never publish the copy. + already_published = DATASET_PROTOCOL_VERSION in published + + if dry_run: + return DatasetUpgrade( + tag=tag, + published_before=published, + published_after=published + if already_published + else sorted({*published, DATASET_PROTOCOL_VERSION}, key=int), + upload_bytes=0 + if already_published + else self._upgrade_upload_bytes(tag), + skipped_reason=_ALREADY_PUBLISHED if already_published else None, + ) + + audience = self._recover_audience(tag) + error: str | None = None + if not already_published: + try: + self._materialize_dataset_copy( + tag, DATASET_PROTOCOL_VERSION, users=audience + ) + except Exception as exc: + # An added layout is additive, so a partial sweep is safe to + # leave. The rest of the datasets still upgrade. + logger.exception("Failed to promote dataset %r", tag) + error = f"{type(exc).__name__}: {exc}" + + if error is None: + # Always, even for a layout that was already published: the promote + # uploads unshared, and a share that failed after the upload + # succeeded would never be retried otherwise, because the next + # sweep sees the layout published and skips it. A share ignores a + # grant that already exists, so a repeat costs nothing. + try: + for collection in self._mock_collections_for(tag): + self._share_dataset_collection( + MOCK_DATASET_SPEC.wire_prefix(collection.variant), + tag, + collection.content_hash, + audience, + ) + except Exception as exc: + logger.exception("Failed to share the layouts of dataset %r", tag) + error = f"{type(exc).__name__}: {exc}" + + return DatasetUpgrade( + tag=tag, + published_before=published, + published_after=sorted( + {self._protocol_of(c) for c in self._mock_collections_for(tag)}, + key=int, + ), + upload_bytes=0 + if (already_published or error is not None) + else self._upgrade_upload_bytes(tag), + skipped_reason=_ALREADY_PUBLISHED + if (already_published and error is None) + else None, + error=error, + ) + + def _recover_audience(self, tag: str) -> list[str] | str: + """Who a dataset on disk was shared with. + + Two sources, because a peer reads a dataset through a shared collection + and not through ``syft.pub.yaml``: the explicit grants come from the + ruleset, and the "any" flag from the collections. "any" has priority, + because it is the wider audience and a single bit cannot be recovered + from a list of emails. + """ + if any(c.has_any_permission for c in self._mock_collections_for(tag)): + return SHARE_WITH_ANY + return self.dataset_manager.recover_audience_from_ruleset(tag) + + def _upgrade_upload_bytes(self, tag: str) -> int: + """The bytes a promote uploads, measured on the local copy. + + The cost of a sweep is what goes to the transport, not a count of files + on disk. The private payload counts too when the copies of the dataset + are Drive-backed, because the promote uploads a private collection as + well as the mock one. + """ + storage = self.dataset_manager.storage + try: + ref = storage.find_dataset_ref( + self.email, tag, protocol_version=DATASET_PROTOCOL_VERSION + ) + except DatasetNotFoundError: + ref = storage.find_dataset_ref(self.email, tag) + dataset = storage.read_dataset(ref) + total = sum( + len(payload) for payload in self._collect_mock_files(dataset).values() + ) + if self._private_collections_for(tag): + total += sum( + f.stat().st_size for f in dataset.private_dir.iterdir() if f.is_file() + ) + return total + def share_private_dataset(self, tag: str, enclave_email: str): """Share private dataset files with an enclave via outbox events. diff --git a/packages/syft-rds/tests/test_create_dataset_cleanup.py b/packages/syft-rds/tests/test_create_dataset_cleanup.py index d9491413da7..94156aff17d 100644 --- a/packages/syft-rds/tests/test_create_dataset_cleanup.py +++ b/packages/syft-rds/tests/test_create_dataset_cleanup.py @@ -6,6 +6,7 @@ import pytest from syft.sync.syftbox_manager import SyftboxManager +from syft_datasets.migrations.registry import DATASET_PROTOCOL_VERSION from syft_rds import SyftRDSClient from syft_datasets.dataset_manager import DATASET_COLLECTION_PREFIX from dataset_test_utils import create_tmp_dataset_files @@ -72,9 +73,9 @@ def test_cleanup_on_private_upload_failure(self): do_manager = self._make_do_manager() # Compute expected local paths before the test so we can verify deletion. - # A dataset with no peers is written at the widest-compatible protocol. + # A dataset with no peers is written at the current protocol. storage = do_manager.dataset_manager.storage - ref = storage.new_dataset_ref("testdataset", storage._widest_protocol_version) + ref = storage.new_dataset_ref("testdataset", DATASET_PROTOCOL_VERSION) mock_dir = storage.public_dataset_dir(ref) private_metadata_dir = storage.private_dataset_dir(ref) diff --git a/packages/syft-rds/tests/test_dataset_upgrade.py b/packages/syft-rds/tests/test_dataset_upgrade.py new file mode 100644 index 00000000000..846287223f4 --- /dev/null +++ b/packages/syft-rds/tests/test_dataset_upgrade.py @@ -0,0 +1,294 @@ +"""``upgrade()`` promotes owned datasets to the layout of the installed client. + +A dataset's object versions follow its layout directory, so a promote is the +addition of the current layout. These tests drive the sweep through the mock +Drive: what it publishes, what it skips, who it reaches, and what it reports. +""" + +import pytest +from syft_datasets.dataset_manager import SHARE_WITH_ANY +from syft_datasets.migrations.registry import DATASET_PROTOCOL_VERSION +from syft_rds import SyftRDSClient +from syft_rds.config import MOCK_DATASET_SPEC, dataset_variant + +from dataset_test_utils import create_tmp_dataset_files + +OLD_PEER = "old@test.org" + + +@pytest.fixture +def pair(): + return SyftRDSClient.pair_with_mock_drive_service_connection( + use_in_memory_cache=False, + ) + + +def _published(do_manager, tag: str) -> set: + return {do_manager._protocol_of(c) for c in do_manager._mock_collections_for(tag)} + + +def _create_v0_dataset(do_manager, tag: str, users=None, upload_private: bool = False): + """A dataset of an earlier release: the flat layout, published and shared.""" + mock_path, private_path, readme_path = create_tmp_dataset_files() + created = do_manager.dataset_manager.create_all( + name=tag, + mock_path=mock_path, + private_path=private_path, + readme_path=readme_path, + users=users, + protocol_versions=["0"], + ) + do_manager._upload_dataset_to_collection(created["0"], users=users or []) + if upload_private: + do_manager._upload_private_dataset_to_collection(created["0"]) + return created["0"] + + +def test_upgrade_promotes_a_v0_dataset_and_publishes_it(pair): + ds_manager, do_manager = pair + _create_v0_dataset(do_manager, "legacy", users=[ds_manager.email]) + assert _published(do_manager, "legacy") == {"0"} + + report = do_manager.upgrade(sync=False) + + assert _published(do_manager, "legacy") == {"0", DATASET_PROTOCOL_VERSION} + assert [d.tag for d in report.promoted] == ["legacy"] + assert report.datasets[0].published_before == ["0"] + assert report.datasets[0].published_after == ["0", DATASET_PROTOCOL_VERSION] + assert report.upload_bytes > 0 + assert not report.failed + + +def test_the_promoted_copy_reaches_the_audience(pair): + ds_manager, do_manager = pair + _create_v0_dataset(do_manager, "legacy", users=[ds_manager.email]) + + do_manager.upgrade(sync=False) + ds_manager.sync() + + dataset = ds_manager.datasets.get("legacy", datasite=do_manager.email) + # The DS reads the newest layout it can, which is the promoted copy. + assert dataset.protocol_version == DATASET_PROTOCOL_VERSION + assert dataset.mock_files + for path in dataset.mock_files: + assert path.exists() + + +def test_upgrade_publishes_a_layout_that_is_on_disk_but_unpublished(pair): + # The regression test for the published-only skip: a migrate can succeed and + # the upload fail, so the layout exists locally with no collection. A skip + # that folded the disk into the published set would never publish it. + ds_manager, do_manager = pair + _create_v0_dataset(do_manager, "halfway", users=[ds_manager.email]) + do_manager.dataset_manager.migrate( + "halfway", DATASET_PROTOCOL_VERSION, users=[ds_manager.email] + ) + storage = do_manager.dataset_manager.storage + assert storage.find_dataset_ref( + do_manager.email, "halfway", protocol_version=DATASET_PROTOCOL_VERSION + ) + assert _published(do_manager, "halfway") == {"0"} + + report = do_manager.upgrade(sync=False) + + assert _published(do_manager, "halfway") == {"0", DATASET_PROTOCOL_VERSION} + assert [d.tag for d in report.promoted] == ["halfway"] + + +def test_upgrade_is_a_no_op_when_the_current_layout_is_published(pair): + ds_manager, do_manager = pair + mock_path, private_path, readme_path = create_tmp_dataset_files() + do_manager.create_dataset( + name="current", + mock_path=mock_path, + private_path=private_path, + readme_path=readme_path, + users=[ds_manager.email], + sync=False, + ) + before = _published(do_manager, "current") + + report = do_manager.upgrade(sync=False) + + assert _published(do_manager, "current") == before + assert not report.promoted + assert "already published" in report.datasets[0].skipped_reason + + +def test_dry_run_writes_nothing_and_reports_the_upload_cost(pair): + ds_manager, do_manager = pair + _create_v0_dataset(do_manager, "legacy", users=[ds_manager.email]) + storage = do_manager.dataset_manager.storage + target = storage.new_dataset_ref("legacy", DATASET_PROTOCOL_VERSION) + + report = do_manager.upgrade(dry_run=True) + + assert report.dry_run + assert _published(do_manager, "legacy") == {"0"} + assert not storage.public_dataset_dir(target).exists() + assert report.datasets[0].published_after == ["0", DATASET_PROTOCOL_VERSION] + assert report.upload_bytes > 0 + + +def test_dry_run_applies_the_same_skip_as_the_sweep(pair): + ds_manager, do_manager = pair + mock_path, private_path, readme_path = create_tmp_dataset_files() + do_manager.create_dataset( + name="current", + mock_path=mock_path, + private_path=private_path, + readme_path=readme_path, + users=[ds_manager.email], + sync=False, + ) + + report = do_manager.upgrade(dry_run=True) + + assert not report.promoted + assert report.upload_bytes == 0 + + +def test_an_owner_only_dataset_stays_owner_only(pair): + _, do_manager = pair + _create_v0_dataset(do_manager, "mine") + + do_manager.upgrade(sync=False) + + promoted = [ + c + for c in do_manager._mock_collections_for("mine") + if do_manager._protocol_of(c) == DATASET_PROTOCOL_VERSION + ] + assert len(promoted) == 1 + assert not promoted[0].has_any_permission + # No ruleset is written for an audience of nobody. + storage = do_manager.dataset_manager.storage + target = storage.new_dataset_ref("mine", DATASET_PROTOCOL_VERSION) + assert not (storage.public_dataset_dir(target) / "syft.pub.yaml").exists() + + +def test_an_any_dataset_is_still_any_after_the_upgrade(pair): + # The "any" flag lives on the collection, not in the ruleset, and it is the + # wider audience, so it has priority over the recovered emails. + _, do_manager = pair + _create_v0_dataset(do_manager, "open", users=SHARE_WITH_ANY) + assert any(c.has_any_permission for c in do_manager._mock_collections_for("open")) + + do_manager.upgrade(sync=False) + + promoted = [ + c + for c in do_manager._mock_collections_for("open") + if do_manager._protocol_of(c) == DATASET_PROTOCOL_VERSION + ] + assert len(promoted) == 1 + assert promoted[0].has_any_permission + + +def test_a_user_added_after_create_is_on_the_promoted_copy(pair): + # The drift case. A share updates the transport, so the ruleset of every + # layout must record it too, or the recovered audience is the create-time + # one and the promoted copy never reaches the later user. + ds_manager, do_manager = pair + _create_v0_dataset(do_manager, "shared", users=[ds_manager.email]) + do_manager.share_dataset("shared", [OLD_PEER], sync=False) + + do_manager.upgrade(sync=False) + + audience = do_manager.dataset_manager.recover_audience_from_ruleset( + "shared", protocol_version=DATASET_PROTOCOL_VERSION + ) + assert sorted(audience) == sorted([ds_manager.email, OLD_PEER]) + + +def test_one_failing_dataset_is_reported_and_the_rest_continue(pair): + ds_manager, do_manager = pair + _create_v0_dataset(do_manager, "aaa broken", users=[ds_manager.email]) + _create_v0_dataset(do_manager, "zzz fine", users=[ds_manager.email]) + + real_materialize = do_manager._materialize_dataset_copy + + def fail_for_one(tag, protocol_version, users): + if tag == "aaa broken": + raise RuntimeError("migrate exploded") + return real_materialize(tag, protocol_version, users) + + do_manager._materialize_dataset_copy = fail_for_one + try: + report = do_manager.upgrade(sync=False) + finally: + do_manager._materialize_dataset_copy = real_materialize + + failed = {d.tag for d in report.failed} + promoted = {d.tag for d in report.promoted} + assert failed == {"aaa broken"} + assert promoted == {"zzz fine"} + assert "migrate exploded" in report.datasets[0].error + assert _published(do_manager, "zzz fine") == {"0", DATASET_PROTOCOL_VERSION} + + +def test_upgrade_is_refused_without_the_do_role(pair): + ds_manager, _ = pair + with pytest.raises(ValueError, match="Only dataset owners"): + ds_manager.upgrade() + + +def test_a_share_that_fails_after_the_upload_is_retried_next_sweep(pair): + # The upload and the share are separate steps. If the share fails, the + # layout is published but unshared, and a skip on "already published" would + # strand the audience on the old copy with no way back. + ds_manager, do_manager = pair + _create_v0_dataset(do_manager, "flaky", users=[ds_manager.email]) + + real_share = do_manager._share_dataset_collection + calls = [] + + def fail_once(wire_prefix, tag, content_hash, users): + calls.append((tag, content_hash, users)) + raise RuntimeError("share exploded") + + do_manager._share_dataset_collection = fail_once + try: + first = do_manager.upgrade(sync=False) + finally: + do_manager._share_dataset_collection = real_share + + # The copy went up, the share did not, and the failure is reported. + assert _published(do_manager, "flaky") == {"0", DATASET_PROTOCOL_VERSION} + assert "share exploded" in first.datasets[0].error + assert not first.promoted + + shared = [] + + def record(wire_prefix, tag, content_hash, users): + shared.append((wire_prefix, users)) + return real_share(wire_prefix, tag, content_hash, users) + + do_manager._share_dataset_collection = record + try: + second = do_manager.upgrade(sync=False) + finally: + do_manager._share_dataset_collection = real_share + + # The second sweep re-shares every layout, including the promoted one. + assert not second.failed + assert {prefix for prefix, _ in shared} == { + MOCK_DATASET_SPEC.wire_prefix(dataset_variant(v)) + for v in ("0", DATASET_PROTOCOL_VERSION) + } + assert all(users == [ds_manager.email] for _, users in shared) + + +def test_dry_run_counts_the_private_payload_of_a_drive_backed_dataset(pair): + # A promote uploads a private collection too when the copies are + # Drive-backed, so the reported cost must include it. + ds_manager, do_manager = pair + _create_v0_dataset(do_manager, "mock only", users=[ds_manager.email]) + _create_v0_dataset( + do_manager, "with private", users=[ds_manager.email], upload_private=True + ) + + report = do_manager.upgrade(dry_run=True) + by_tag = {d.tag: d.upload_bytes for d in report.datasets} + + assert by_tag["with private"] > by_tag["mock only"] diff --git a/tests/migrations/p2p/test_dataset_multicopy_delivery.py b/tests/migrations/p2p/test_dataset_multicopy_delivery.py index b30fc2a96b7..56fbed743b0 100644 --- a/tests/migrations/p2p/test_dataset_multicopy_delivery.py +++ b/tests/migrations/p2p/test_dataset_multicopy_delivery.py @@ -128,23 +128,22 @@ def test_the_local_directory_of_a_collection_follows_its_protocol(pair): def test_a_dataset_for_a_protocol0_peer_arrives_flat_and_reads(pair): ds_manager, do_manager = pair - # The DS advertises dataset protocol 0, as an earlier client does. - do_manager.peer_manager.live_peer_schemas("syft-dataset")[ds_manager.email] = ( - _dataset_schema("0") - ) - + # A dataset of an earlier release: the flat layout is the only one published, + # so it is the only one a peer can take. mock_path, private_path, readme_path = create_tmp_dataset_files() - do_manager.create_dataset( + created = do_manager.dataset_manager.create_all( name="skew dataset", mock_path=mock_path, private_path=private_path, readme_path=readme_path, - users=[ds_manager.email], + protocol_versions=["0"], ) + do_manager._upload_dataset_to_collection(created["0"], users=[ds_manager.email]) ds_manager.sync() dataset = ds_manager.datasets.get("skew dataset", datasite=do_manager.email) - # The owner wrote the layout that this peer reads, not its own newest. + # The flat copy carries the whole dataset, and the metadata points at files + # the peer actually got. assert dataset.protocol_version == "0" assert ( dataset.mock_dir @@ -160,6 +159,34 @@ def test_a_dataset_for_a_protocol0_peer_arrives_flat_and_reads(pair): ) +def test_a_create_for_a_protocol0_peer_publishes_both_layouts(pair): + # The create writes the current layout always, and the layout of each peer + # in the audience. A peer that reads only the older one still finds a copy, + # and a peer that reads both prefers the newest. + ds_manager, do_manager = pair + do_manager.peer_manager.live_peer_schemas("syft-dataset")[OLD_PEER] = ( + _dataset_schema("0") + ) + + mock_path, private_path, readme_path = create_tmp_dataset_files() + do_manager.create_dataset( + name="skew dataset", + mock_path=mock_path, + private_path=private_path, + readme_path=readme_path, + users=[ds_manager.email, OLD_PEER], + ) + + assert _collections_for(do_manager, "skew dataset") == {"0", "1"} + + ds_manager.sync() + dataset = ds_manager.datasets.get("skew dataset", datasite=do_manager.email) + assert dataset.protocol_version == "1" + assert dataset.mock_files + for path in dataset.mock_files: + assert path.exists() + + def test_a_mixed_audience_gets_one_collection_for_each_protocol(pair): _, do_manager = pair mock_path, private_path, readme_path = create_tmp_dataset_files() @@ -490,3 +517,101 @@ def test_a_share_uploads_a_local_copy_that_has_no_collection(pair): do_manager.share_dataset("half done", [OLD_PEER], sync=False) assert _collections_for(do_manager, "half done") == {"0", "1"} + + +def test_upgrade_publishes_the_current_layout_and_the_peer_prefers_it(pair): + # A dataset of an earlier release: the flat layout alone, on disk and on the + # transport. A current peer already reads it, because its floor is "0", so + # this proves the promote and the publish happened -- not that a read was + # broken. + ds_manager, do_manager = pair + mock_path, private_path, readme_path = create_tmp_dataset_files() + created = do_manager.dataset_manager.create_all( + name="from an earlier release", + mock_path=mock_path, + private_path=private_path, + readme_path=readme_path, + users=[ds_manager.email], + protocol_versions=["0"], + ) + do_manager._upload_dataset_to_collection(created["0"], users=[ds_manager.email]) + assert _collections_for(do_manager, "from an earlier release") == {"0"} + + do_manager.upgrade(sync=False) + + assert _collections_for(do_manager, "from an earlier release") == {"0", "1"} + + ds_manager.sync() + dataset = ds_manager.datasets.get( + "from an earlier release", datasite=do_manager.email + ) + assert dataset.protocol_version == "1" + assert dataset.mock_dir == ( + ds_manager.syftbox_folder + / do_manager.email + / COLLECTION_SUBPATH + / "v1" + / "from an earlier release" + ) + assert dataset.mock_files + for path in dataset.mock_files: + assert path.exists(), ( + f"the metadata points to a file the peer does not get: {path}" + ) + + +def test_a_peer_approved_after_an_any_create_gets_a_layout_it_reads(pair): + # An "any" dataset holds the layouts its audience read at create time. A + # peer approved later may read none of them, so the approval materializes + # and publishes its layout before granting on the collection. + ds_manager, do_manager = pair + mock_path, private_path, readme_path = create_tmp_dataset_files() + created = do_manager.dataset_manager.create_all( + name="open dataset", + mock_path=mock_path, + private_path=private_path, + readme_path=readme_path, + users="any", + protocol_versions=["1"], + ) + do_manager._upload_dataset_to_collection(created["1"], users="any") + assert _collections_for(do_manager, "open dataset") == {"1"} + + # A protocol-0 peer arrives and is approved. + do_manager.peer_manager.live_peer_schemas("syft-dataset")[OLD_PEER] = ( + _dataset_schema("0") + ) + do_manager._share_any_datasets_with_peer(OLD_PEER) + + assert _collections_for(do_manager, "open dataset") == {"0", "1"} + + +def test_a_backfilled_any_layout_is_marked_any_too(pair): + # The "any" bit lives on the collection. A layout added at approval time + # must carry it, or a later cache rebuild from has_any_permission loses it + # and the next peer never sees that layout. + ds_manager, do_manager = pair + mock_path, private_path, readme_path = create_tmp_dataset_files() + created = do_manager.dataset_manager.create_all( + name="open dataset", + mock_path=mock_path, + private_path=private_path, + readme_path=readme_path, + users="any", + protocol_versions=["1"], + ) + do_manager._upload_dataset_to_collection(created["1"], users="any") + + do_manager.peer_manager.live_peer_schemas("syft-dataset")[OLD_PEER] = ( + _dataset_schema("0") + ) + do_manager._share_any_datasets_with_peer(OLD_PEER) + + collections = do_manager._mock_collections_for("open dataset") + assert {do_manager._protocol_of(c) for c in collections} == {"0", "1"} + assert all(c.has_any_permission for c in collections) + # And the cache names every layout, so the next approval needs no listing. + cached = do_manager.sync_engine.datasite_owner_syncer.any_shared_collections + assert {hash_ for _, tag, hash_ in cached if tag == "open dataset"} == { + c.content_hash for c in collections + } diff --git a/tests/migrations/p2p/test_dataset_schema_negotiation.py b/tests/migrations/p2p/test_dataset_schema_negotiation.py index 12c3e461608..d189072e5ae 100644 --- a/tests/migrations/p2p/test_dataset_schema_negotiation.py +++ b/tests/migrations/p2p/test_dataset_schema_negotiation.py @@ -48,17 +48,26 @@ def test_all_current_audience_drops_legacy_layout(tmp_path): } -def test_unknown_peer_gets_widest_protocol(tmp_path): +def test_a_named_peer_of_unknown_version_gets_the_floor(tmp_path): + # We cannot assume a peer we have not read reads the current layout, so it + # takes the oldest protocol we still support -- a layout it can read. storage = _storage(tmp_path, {}) versions = storage.target_protocol_versions_for_peers([UNKNOWN_PEER]) - assert versions == {storage._widest_protocol_version} + assert versions == {storage._floor_protocol_version} + + +def test_no_audience_needs_no_older_layout(tmp_path): + # With no peer to serve, a create writes the current layout alone. + storage = _storage(tmp_path, {}) + assert storage.target_protocol_versions_for_peers([]) == set() + assert storage.target_protocol_versions_for_peers() == set() def test_live_map_updates_are_seen_by_storage(tmp_path): live: dict = {} storage = _storage(tmp_path, live) assert storage.target_protocol_versions_for_peers([NEW_PEER]) == { - storage._widest_protocol_version + storage._floor_protocol_version } live[NEW_PEER] = _dataset_schema(DATASET_PROTOCOL_VERSION) assert storage.target_protocol_versions_for_peers([NEW_PEER]) == {