diff --git a/packages/enclave-model-api-example/docker/Dockerfile b/packages/enclave-model-api-example/docker/Dockerfile index 92884788bf7..f896b36b6ce 100644 --- a/packages/enclave-model-api-example/docker/Dockerfile +++ b/packages/enclave-model-api-example/docker/Dockerfile @@ -28,6 +28,7 @@ SYFT_ENCLAVE_DATA_OWNERS,\ SYFT_ENCLAVE_REQUIRE_TEE,\ SYFT_ENCLAVE_FRESH_STATE,\ SYFT_ENCLAVE_USE_ENCRYPTION,\ +SYFT_ENCLAVE_PERSIST_OWNER_STATE,\ SYFT_DEFAULT_JOB_TIMEOUT_SECONDS,\ SYFT_BOOTSTRAP,\ SYFT_BOOTSTRAP_WIF_AUDIENCE,\ diff --git a/packages/enclave-model-api-example/src/enclave_model_api/__main__.py b/packages/enclave-model-api-example/src/enclave_model_api/__main__.py index 68ef7fd8795..220b98a1de8 100644 --- a/packages/enclave-model-api-example/src/enclave_model_api/__main__.py +++ b/packages/enclave-model-api-example/src/enclave_model_api/__main__.py @@ -48,7 +48,8 @@ def main() -> None: f"Enclave settings — email={settings.email} data_owners={settings.data_owners} " f"token_path={settings.token_path} poll_interval={settings.poll_interval}s " f"require_tee={settings.require_tee} fresh_state={settings.fresh_state} " - f"use_encryption={settings.use_encryption}" + f"use_encryption={settings.use_encryption} " + f"persist_owner_state={settings.persist_owner_state}" ) logger.info( f"Inference settings — model_owner={inference.model_owner} " @@ -62,6 +63,7 @@ def main() -> None: token_path=settings.token_path, data_owners=settings.data_owners, encryption=settings.use_encryption, + persist_owner_state=settings.persist_owner_state, ) logger.info("SyftEnclaveClient ready") diff --git a/packages/syft-enclave/docker/Dockerfile b/packages/syft-enclave/docker/Dockerfile index 8b1e5cb6027..38bbdd79302 100644 --- a/packages/syft-enclave/docker/Dockerfile +++ b/packages/syft-enclave/docker/Dockerfile @@ -54,6 +54,7 @@ SYFT_ENCLAVE_DATA_OWNERS,\ SYFT_ENCLAVE_REQUIRE_TEE,\ SYFT_ENCLAVE_FRESH_STATE,\ SYFT_ENCLAVE_USE_ENCRYPTION,\ +SYFT_ENCLAVE_PERSIST_OWNER_STATE,\ SYFT_DEFAULT_JOB_TIMEOUT_SECONDS,\ SYFT_BOOTSTRAP,\ SYFT_BOOTSTRAP_WIF_AUDIENCE,\ diff --git a/packages/syft-enclave/src/syft_enclaves/__main__.py b/packages/syft-enclave/src/syft_enclaves/__main__.py index e3848cab4a8..ef12ce1f7c4 100644 --- a/packages/syft-enclave/src/syft_enclaves/__main__.py +++ b/packages/syft-enclave/src/syft_enclaves/__main__.py @@ -43,7 +43,8 @@ def main() -> None: f"Enclave settings — email={settings.email} data_owners={settings.data_owners} " f"token_path={settings.token_path} poll_interval={settings.poll_interval}s " f"require_tee={settings.require_tee} fresh_state={settings.fresh_state} " - f"use_encryption={settings.use_encryption}" + f"use_encryption={settings.use_encryption} " + f"persist_owner_state={settings.persist_owner_state}" ) logger.info("Building SyftEnclaveClient...") @@ -52,6 +53,7 @@ def main() -> None: token_path=settings.token_path, data_owners=settings.data_owners, encryption=settings.use_encryption, + persist_owner_state=settings.persist_owner_state, ) logger.info("SyftEnclaveClient ready") diff --git a/packages/syft-enclave/src/syft_enclaves/client.py b/packages/syft-enclave/src/syft_enclaves/client.py index 8dce7d1e4ff..95f6d4398a6 100644 --- a/packages/syft-enclave/src/syft_enclaves/client.py +++ b/packages/syft-enclave/src/syft_enclaves/client.py @@ -435,6 +435,7 @@ def for_enclave( token_path: Path | str | None = None, data_owners: list[str] | None = None, encryption: bool = False, + persist_owner_state: bool = False, ) -> "SyftEnclaveClient": """Build an enclave client backed by a real Google Drive connection. Args: @@ -442,6 +443,10 @@ def for_enclave( token_path: Path to a pre-authorized Google Drive OAuth token. data_owners: Emails whose approval gates every job on this enclave. encryption: Enable end-to-end drive encryption. + persist_owner_state: Keep the owner-only state used to restore + this datasite later (event log, rolling state, checkpoints). + Off by default - an enclave gets ephemeral keys and wipes its + state on boot, so there is never anything to restore. """ config = SyftRDSClientConfig.for_jupyter( email=email, @@ -449,6 +454,7 @@ def for_enclave( has_do_role=True, token_path=Path(token_path) if token_path is not None else None, encryption=encryption, + persist_owner_state=persist_owner_state, ) # Note: We do not currently provide the ability to load encryption keys passed during creation of enclave. diff --git a/packages/syft-enclave/src/syft_enclaves/settings.py b/packages/syft-enclave/src/syft_enclaves/settings.py index d3d490c572f..e1ff1d01dda 100644 --- a/packages/syft-enclave/src/syft_enclaves/settings.py +++ b/packages/syft-enclave/src/syft_enclaves/settings.py @@ -94,3 +94,16 @@ def _split_data_owners(cls, v: object) -> object: "Enabled by default; set false to disable." ), ) + persist_owner_state: bool = Field( + default=False, + description=( + "Keep the owner-only state that exists solely to restore this " + "datasite later: the append-only event log, the rolling state and " + "the checkpoints. Off by default because an enclave can never " + "restore — it gets an ephemeral keypair each boot and wipes its " + "state under fresh_state — so those writes only spend Drive API " + "calls inside the poll loop. Peer-facing state (outbox, " + "collections, peers) is unaffected. Set true only for a stateful " + "enclave (fresh_state=false)." + ), + ) diff --git a/packages/syft-enclave/tests/test_settings.py b/packages/syft-enclave/tests/test_settings.py index 3b12260d1a1..32da55a0731 100644 --- a/packages/syft-enclave/tests/test_settings.py +++ b/packages/syft-enclave/tests/test_settings.py @@ -38,6 +38,7 @@ def test_defaults_applied_when_required_fields_set(required_env): assert settings.log_level == "INFO" assert settings.fresh_state is True # default: always start with a clean slate assert settings.use_encryption is True # default: encryption on + assert settings.persist_owner_state is False # default: an enclave never restores def test_data_owners_parsed_from_comma_separated_string(clean_env): @@ -112,3 +113,36 @@ def test_use_encryption_can_be_disabled_via_env(required_env): required_env.setenv("SYFT_ENCLAVE_USE_ENCRYPTION", "false") settings = EnclaveSettings(_env_file=None) assert settings.use_encryption is False + + +def test_persist_owner_state_can_be_enabled_via_env(required_env): + required_env.setenv("SYFT_ENCLAVE_PERSIST_OWNER_STATE", "true") + settings = EnclaveSettings(_env_file=None) + assert settings.persist_owner_state is True + + +def test_persist_owner_state_reaches_the_do_syncer_config(): + """Guard the `**kw` hole between for_enclave and the sync engine. + + ``SyftRDSClientConfig.for_jupyter`` forwards unknown kwargs straight to + ``SyftboxManagerConfig.for_jupyter``, so nothing type-checks this hop. A + rename on either side would silently give an enclave back the owner state + it is not supposed to keep. + """ + from syft_rds.config import SyftRDSClientConfig + + config = SyftRDSClientConfig.for_jupyter( + email="enclave@openmined.org", + has_do_role=True, + has_ds_role=True, + persist_owner_state=False, + ) + assert config.sync.persist_owner_state is False + assert config.sync.datasite_owner_syncer_config.persist_owner_state is False + + default = SyftRDSClientConfig.for_jupyter( + email="enclave@openmined.org", + has_do_role=True, + has_ds_role=True, + ) + assert default.sync.datasite_owner_syncer_config.persist_owner_state is True diff --git a/syft/sync/syftbox_manager.py b/syft/sync/syftbox_manager.py index 3585538682d..aea0d0c2375 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -113,6 +113,10 @@ class SyftboxManagerConfig(BaseModel): has_ds_role: bool = False has_do_role: bool = False use_in_memory_cache: bool = True + # Keep the owner-only state used to restore this datasite later: event + # log, rolling state, checkpoints. Off for datasites that can never + # restore (see syft-enclave's EnclaveSettings). + persist_owner_state: bool = True datasite_owner_syncer_config: DatasiteOwnerSyncerConfig peer_manager_config: PeerManagerConfig @@ -126,6 +130,7 @@ def for_colab( has_ds_role: bool = False, has_do_role: bool = False, encryption: bool = False, + persist_owner_state: bool = True, crypto_keys_path: Path | None = None, skip_peer_on_patch_version_diff: Optional[ bool @@ -150,6 +155,7 @@ def for_colab( connection_configs = [GdriveConnectionConfig(email=email, token_path=None)] datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, + persist_owner_state=persist_owner_state, syftbox_folder=syftbox_folder, collections_folder=collections_folder, collection_specs=collection_specs, @@ -192,6 +198,7 @@ def for_colab( has_do_role=has_do_role, connection_configs=connection_configs, use_in_memory_cache=False, + persist_owner_state=persist_owner_state, datasite_owner_syncer_config=datasite_owner_syncer_config, datasite_watcher_syncer_config=datasite_watcher_syncer_config, peer_manager_config=peer_manager_config, @@ -205,6 +212,7 @@ def for_jupyter( has_do_role: bool = False, token_path: Path | None = None, encryption: bool = False, + persist_owner_state: bool = True, crypto_keys_path: Path | None = None, skip_peer_on_patch_version_diff: Optional[ bool @@ -232,6 +240,7 @@ def for_jupyter( ] datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, + persist_owner_state=persist_owner_state, syftbox_folder=syftbox_folder, collections_folder=collections_folder, collection_specs=collection_specs, @@ -274,6 +283,7 @@ def for_jupyter( has_ds_role=has_ds_role, has_do_role=has_do_role, use_in_memory_cache=False, + persist_owner_state=persist_owner_state, datasite_owner_syncer_config=datasite_owner_syncer_config, datasite_watcher_syncer_config=datasite_watcher_syncer_config, peer_manager_config=peer_manager_config, @@ -288,6 +298,7 @@ def _base_config_for_testing( has_ds_role: bool = False, has_do_role: bool = False, use_in_memory_cache: bool = True, + persist_owner_state: bool = True, check_versions: bool = False, collection_specs: list["CollectionSyncSpec"] | None = None, ): @@ -306,6 +317,7 @@ def _base_config_for_testing( datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, + persist_owner_state=persist_owner_state, syftbox_folder=syftbox_folder, collections_folder=collections_folder, collection_specs=collection_specs, @@ -346,6 +358,7 @@ def _base_config_for_testing( has_ds_role=has_ds_role, has_do_role=has_do_role, use_in_memory_cache=use_in_memory_cache, + persist_owner_state=persist_owner_state, datasite_owner_syncer_config=datasite_owner_syncer_config, datasite_watcher_syncer_config=datasite_watcher_syncer_config, peer_manager_config=peer_manager_config, @@ -361,6 +374,7 @@ def for_google_drive_testing_connection( has_ds_role: bool = False, has_do_role: bool = False, use_in_memory_cache: bool = True, + persist_owner_state: bool = True, check_versions: bool = False, collection_specs: list["CollectionSyncSpec"] | None = None, ): @@ -381,6 +395,7 @@ def for_google_drive_testing_connection( ] datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, + persist_owner_state=persist_owner_state, syftbox_folder=syftbox_folder, collections_folder=collections_folder, collection_specs=collection_specs, @@ -418,6 +433,7 @@ def for_google_drive_testing_connection( email=email, syftbox_folder=syftbox_folder, write_files=write_files, + persist_owner_state=persist_owner_state, datasite_owner_syncer_config=datasite_owner_syncer_config, datasite_watcher_syncer_config=datasite_watcher_syncer_config, has_ds_role=has_ds_role, @@ -587,6 +603,7 @@ def for_colab( has_ds_role: bool = False, has_do_role: bool = False, encryption: bool = False, + persist_owner_state: bool = True, crypto_keys_path: Path | None = None, skip_peer_on_patch_version_diff: Optional[ bool @@ -599,6 +616,7 @@ def for_colab( has_ds_role=has_ds_role, has_do_role=has_do_role, encryption=encryption, + persist_owner_state=persist_owner_state, crypto_keys_path=crypto_keys_path, skip_peer_on_patch_version_diff=skip_peer_on_patch_version_diff, force_ignore_peer_version=force_ignore_peer_version, @@ -614,6 +632,7 @@ def for_jupyter( has_do_role: bool = False, token_path: Path | None = None, encryption: bool = False, + persist_owner_state: bool = True, crypto_keys_path: Path | None = None, skip_peer_on_patch_version_diff: Optional[ bool @@ -629,6 +648,7 @@ def for_jupyter( has_do_role=has_do_role, token_path=token_path, encryption=encryption, + persist_owner_state=persist_owner_state, crypto_keys_path=crypto_keys_path, skip_peer_on_patch_version_diff=skip_peer_on_patch_version_diff, force_ignore_peer_version=force_ignore_peer_version, @@ -648,6 +668,7 @@ def _pair_with_google_drive_testing_connection( add_peers: bool = True, load_peers: bool = False, use_in_memory_cache: bool = True, + persist_owner_state: bool = True, clear_caches: bool = True, check_versions: bool = False, collection_specs: list["CollectionSyncSpec"] | None = None, @@ -656,6 +677,7 @@ def _pair_with_google_drive_testing_connection( email=do_email, syftbox_folder=base_path1, use_in_memory_cache=use_in_memory_cache, + persist_owner_state=persist_owner_state, token_path=do_token_path, has_ds_role=False, has_do_role=True, @@ -669,6 +691,7 @@ def _pair_with_google_drive_testing_connection( email=ds_email, syftbox_folder=base_path2, use_in_memory_cache=use_in_memory_cache, + persist_owner_state=persist_owner_state, token_path=ds_token_path, has_ds_role=True, has_do_role=False, @@ -732,6 +755,7 @@ def pair_with_mock_drive_service_connection( sync_automatically: bool = False, add_peers: bool = True, use_in_memory_cache: bool = True, + persist_owner_state: bool = True, check_versions: bool = False, encryption: bool = False, collection_specs: list["CollectionSyncSpec"] | None = None, @@ -750,6 +774,7 @@ def pair_with_mock_drive_service_connection( sync_automatically: Whether to sync when DS sends changes add_peers: Whether to automatically add and approve peers use_in_memory_cache: Whether to use in-memory caches + persist_owner_state: Whether the DO keeps restorable owner state check_versions: Whether to check protocol/client versions Returns: @@ -762,6 +787,7 @@ def pair_with_mock_drive_service_connection( has_ds_role=False, has_do_role=True, use_in_memory_cache=use_in_memory_cache, + persist_owner_state=persist_owner_state, check_versions=check_versions, collection_specs=collection_specs, ) @@ -772,6 +798,7 @@ def pair_with_mock_drive_service_connection( has_ds_role=True, has_do_role=False, use_in_memory_cache=use_in_memory_cache, + persist_owner_state=persist_owner_state, check_versions=check_versions, collection_specs=collection_specs, ) @@ -933,7 +960,10 @@ def sync( Args: auto_checkpoint: If True, automatically create checkpoint when - event count exceeds threshold (DO only). + event count exceeds threshold (DO only). Has no + effect when this manager was built with + persist_owner_state=False, which keeps no + checkpoints at all. checkpoint_threshold: Create checkpoint when events >= this value. auto_compact: If True, after each DO sync, compact each peer's outbox if it holds at least `compact_threshold` diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index 6f0c2a48c25..9ba26e1188b 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -6,7 +6,7 @@ from pydantic import ConfigDict, Field, BaseModel, PrivateAttr from concurrent.futures import ThreadPoolExecutor from queue import Queue -from typing import List, Tuple +from typing import List, NamedTuple, Tuple from syft.sync.events.file_change_event import ( FileChangeEventsMessage, FileChangeEventsMessageFileName, @@ -52,10 +52,38 @@ MIN_MESSAGES_COMPACT = 20 +class _CheckpointRestore(NamedTuple): + """What restoring from checkpoints recovered. + + ``found_checkpoint`` is not the same as ``events_since_timestamp is not + None``: a full checkpoint may carry no event cursor, and in that case we + must NOT fall back to downloading every event. + """ + + events_since_timestamp: float | None + events: List[FileChangeEvent] + found_checkpoint: bool + + +def _later_timestamp(current: float | None, candidate: float | None) -> float | None: + """The later of two event timestamps, either of which may be missing.""" + if candidate is None: + return current + if current is None or candidate > current: + return candidate + return current + + class DatasiteOwnerSyncerConfig(BaseModel): email: str syftbox_folder: Path write_files: bool = True + # Whether to keep the owner-only state that exists solely to restore this + # datasite later: the append-only event log, the rolling state and the + # checkpoints. Off for a datasite that can never restore - an enclave gets + # ephemeral keys and wipes its state each boot - so those writes are pure + # cost. Peer-facing state (outbox, collections, peers) is unaffected. + persist_owner_state: bool = True # Full path to collections folder - must be provided explicitly collections_folder: Path | None = None # Collection sync specs (public prefix + local subpath). Empty for a bare @@ -75,6 +103,7 @@ class DatasiteOwnerSyncer(BaseModelCallbackMixin): default_factory=lambda: DataSiteOwnerEventCache() ) write_files: bool = True + persist_owner_state: bool = True connection_router: ConnectionRouter initial_sync_done: bool = False email: str @@ -115,6 +144,7 @@ def from_config(cls, config: DatasiteOwnerSyncerConfig): return cls( event_cache=DataSiteOwnerEventCache.from_config(config.cache_config), write_files=config.write_files, + persist_owner_state=config.persist_owner_state, connection_router=ConnectionRouter.from_configs( config.email, config.connection_configs ), @@ -155,7 +185,8 @@ def _save_rolling_state(self) -> None: tmp.rename(path) def sync(self, peer_emails: list[str], recompute_hashes: bool = True): - self._load_rolling_state() + if self.persist_owner_state: + self._load_rolling_state() try: if not self.initial_sync_done: self.pull_initial_state() @@ -181,7 +212,8 @@ def sync(self, peer_emails: list[str], recompute_hashes: bool = True): ) self.process_syftbox_events_queue() finally: - self._save_rolling_state() + if self.persist_owner_state: + self._save_rolling_state() def download_events_message_by_id_with_connection( self, events_message_id: str @@ -209,108 +241,128 @@ def pull_initial_state(self): """ Pull initial state from Google Drive. - Flow: - 1. Check for full (compacted) checkpoint → apply it - 2. Check for incremental checkpoints → apply them in order - 3. Check for rolling state → apply if valid - 4. Download any remaining events since last timestamp - 5. Ensure events_messages_connection is populated for get_cached_events() + Two halves: + 1. The owner's own state - checkpoints, rolling state, and the events + newer than those. Skipped entirely when `persist_owner_state` is + off: nothing was ever written, so there is nothing to read back. + 2. The owner's collections, which are shared with peers rather than + owner-only, so they are restored either way. """ - events_since_timestamp: float | None = None - restored_events: list[FileChangeEvent] = [] + if self.persist_owner_state: + self._restore_owner_state() + + # Restore ALL registered collections (public + private) generically. Each + # spec's flags decide the behaviour: `immutable` picks mirror vs restore-only; + # the local destination comes from the spec's local_subpath. + self._pull_collections_for_initial_sync() + + self.initial_sync_done = True + + def _restore_owner_state(self) -> None: + """Rebuild the cache from the owner's persisted state on the backend.""" + restore = self._restore_from_checkpoints() - # Step 1: Check for full (compacted) checkpoint + if restore.events_since_timestamp is not None: + self._download_events_since(restore.events_since_timestamp) + elif not restore.found_checkpoint: + print("No checkpoints found, downloading all events...") + self._download_all_events() + + # The restore above only populates file_hashes and file_connection, not + # events_messages_connection, which get_cached_events() reads from. + if restore.events and not self.event_cache.get_cached_events(): + self._write_events_to_messages_cache(restore.events) + + def _restore_from_checkpoints(self) -> _CheckpointRestore: + """Apply the full checkpoint, then the incrementals, then the rolling state. + + Each step restores cache state and moves the event cursor forward. + """ full_checkpoint = self.connection_router.get_latest_checkpoint() - if full_checkpoint is not None: - print( - f"Found full checkpoint with {len(full_checkpoint.files)} files, " - "restoring..." - ) - self.event_cache.apply_checkpoint( - full_checkpoint, write_files=self.write_files - ) - events_since_timestamp = full_checkpoint.last_event_timestamp + since = self._apply_full_checkpoint(full_checkpoint) - # Step 2: Check for incremental checkpoints incremental_cps = self.connection_router.get_all_incremental_checkpoints() - if incremental_cps: - print(f"Found {len(incremental_cps)} incremental checkpoints, applying...") - for inc_cp in incremental_cps: - self._apply_incremental_checkpoint_to_cache(inc_cp) - restored_events.extend(inc_cp.events) - # Update timestamp to the latest event in this checkpoint - for event in inc_cp.events: - if event.timestamp is not None: - if ( - events_since_timestamp is None - or event.timestamp > events_since_timestamp - ): - events_since_timestamp = event.timestamp + since, events = self._apply_incremental_checkpoints(incremental_cps, since) + + since, rolling_events = self._apply_rolling_state(since) + return _CheckpointRestore( + events_since_timestamp=since, + events=events + rolling_events, + found_checkpoint=full_checkpoint is not None or bool(incremental_cps), + ) - # Step 3: Check for rolling state + def _apply_full_checkpoint(self, checkpoint: Checkpoint | None) -> float | None: + """Restore a full snapshot. Returns the event cursor it carries.""" + if checkpoint is None: + return None + print(f"Found full checkpoint with {len(checkpoint.files)} files, restoring...") + self.event_cache.apply_checkpoint(checkpoint, write_files=self.write_files) + return checkpoint.last_event_timestamp + + def _apply_incremental_checkpoints( + self, + checkpoints: List[IncrementalCheckpoint], + since: float | None, + ) -> Tuple[float | None, List[FileChangeEvent]]: + """Apply incrementals in order, advancing the cursor past their events.""" + if not checkpoints: + return since, [] + + print(f"Found {len(checkpoints)} incremental checkpoints, applying...") + events: List[FileChangeEvent] = [] + for inc_cp in checkpoints: + self._apply_incremental_checkpoint_to_cache(inc_cp) + events.extend(inc_cp.events) + for event in inc_cp.events: + since = _later_timestamp(since, event.timestamp) + return since, events + + def _apply_rolling_state( + self, since: float | None + ) -> Tuple[float | None, List[FileChangeEvent]]: + """Apply the rolling state, or start an empty one when there is none.""" rolling_state = self.connection_router.get_rolling_state() - if rolling_state is not None and rolling_state.event_count > 0: - print( - f"Found rolling state with {rolling_state.event_count} events, " - "applying..." - ) - self._apply_rolling_state_to_cache(rolling_state) - self._rolling_state = rolling_state - restored_events.extend(rolling_state.events) - - # Update timestamp from rolling state - if rolling_state.last_event_timestamp is not None: - if ( - events_since_timestamp is None - or rolling_state.last_event_timestamp > events_since_timestamp - ): - events_since_timestamp = rolling_state.last_event_timestamp - else: - # Initialize empty rolling state - base_timestamp = events_since_timestamp or 0.0 + if rolling_state is None or rolling_state.event_count == 0: self._rolling_state = RollingState( email=self.email, - base_checkpoint_timestamp=base_timestamp, - ) - - # Step 4: Download any remaining events since last timestamp - if events_since_timestamp is not None: - events_messages = ( - self.connection_router.get_events_messages_since_timestamp( - events_since_timestamp - ) - ) - if events_messages: - print( - f"Downloading {len(events_messages)} events since " - "checkpoint/rolling state..." - ) - for events_message in events_messages: - self.event_cache.add_events_message_to_local_cache(events_message) - self._add_events_to_rolling_state(events_message) - elif full_checkpoint is None and not incremental_cps: - # No checkpoints at all - download all events (fallback) - print("No checkpoints found, downloading all events...") - since_timestamp = self.event_cache.latest_cached_timestamp - events_messages_list: list[FileChangeEventsMessage] = ( - self.get_all_accepted_events_messages(since_timestamp=since_timestamp) + base_checkpoint_timestamp=since or 0.0, ) - for events_message in events_messages_list: - self.event_cache.add_events_message_to_local_cache(events_message) + return since, [] - # Step 5: Ensure events from checkpoints/rolling state are in - # events_messages_connection. Steps 1-3 only populate file_hashes - # and file_connection but not events_messages_connection, which - # get_cached_events() reads from. - if restored_events and not self.event_cache.get_cached_events(): - self._write_events_to_messages_cache(restored_events) + print( + f"Found rolling state with {rolling_state.event_count} events, applying..." + ) + self._apply_rolling_state_to_cache(rolling_state) + self._rolling_state = rolling_state + return ( + _later_timestamp(since, rolling_state.last_event_timestamp), + rolling_state.events, + ) - # Restore ALL registered collections (public + private) generically. Each - # spec's flags decide the behaviour: `immutable` picks mirror vs restore-only; - # the local destination comes from the spec's local_subpath. - self._pull_collections_for_initial_sync() + def _download_events_since(self, timestamp: float) -> None: + """Download and cache every events message newer than `timestamp`.""" + events_messages = self.connection_router.get_events_messages_since_timestamp( + timestamp + ) + if not events_messages: + return - self.initial_sync_done = True + print( + f"Downloading {len(events_messages)} events since " + "checkpoint/rolling state..." + ) + for events_message in events_messages: + self.event_cache.add_events_message_to_local_cache(events_message) + self._add_events_to_rolling_state(events_message) + + def _download_all_events(self) -> None: + """Download every accepted events message - the no-checkpoint path.""" + since_timestamp = self.event_cache.latest_cached_timestamp + events_messages_list: List[FileChangeEventsMessage] = ( + self.get_all_accepted_events_messages(since_timestamp=since_timestamp) + ) + for events_message in events_messages_list: + self.event_cache.add_events_message_to_local_cache(events_message) def _apply_incremental_checkpoint_to_cache( self, checkpoint: IncrementalCheckpoint @@ -709,7 +761,11 @@ def handle_proposed_filechange_events_message( def queue_event_for_syftbox( self, recipients: list[str], file_change_events_message: FileChangeEventsMessage ): - self.syftbox_events_queue.put(file_change_events_message) + # The syftbox queue is the owner's own append-only log, read back only + # to restore this datasite. The outbox is how peers get the change, so + # it is queued either way. + if self.persist_owner_state: + self.syftbox_events_queue.put(file_change_events_message) for recipient in recipients: self.outbox_queue.put((recipient, file_change_events_message)) @@ -852,7 +908,7 @@ def _add_events_to_rolling_state( events_message: The events to add. upload_threshold: Upload to GDrive after this many events added. """ - if self._rolling_state is None: + if not self.persist_owner_state or self._rolling_state is None: return self._rolling_state.add_events_message(events_message) @@ -864,6 +920,8 @@ def _add_events_to_rolling_state( def _upload_rolling_state(self) -> None: """Upload the in-memory rolling state to GDrive.""" + if not self.persist_owner_state: + return if self._rolling_state is None or self._rolling_state.event_count == 0: return @@ -1109,8 +1167,12 @@ def try_create_checkpoint( compacting_threshold: Compact if >= this many incremental checkpoints. Returns: - The created checkpoint (incremental or compacted), or None. + The created checkpoint (incremental or compacted), or None. Always + None when this client keeps no owner state. """ + if not self.persist_owner_state: + return None + self._load_rolling_state() try: result = None diff --git a/tests/unit/test_checkpoints.py b/tests/unit/test_checkpoints.py index e4c4ad72911..abe1ae05a6c 100644 --- a/tests/unit/test_checkpoints.py +++ b/tests/unit/test_checkpoints.py @@ -603,3 +603,124 @@ def test_local_do_changes_end_up_in_incremental_checkpoint(): f"local_{i}.txt missing from incremental checkpoints. " f"process_local_changes did not add it to rolling state." ) + + +def test_no_owner_state_written_when_disabled(): + """With persist_owner_state=False, no owner-only state reaches the drive. + + Owner-only means read back by nobody but this datasite: the append-only + event log, the rolling state and the checkpoints. The second half is the + point of the flag - dropping that state must not cost correctness, so the + peer still has to see every change. + """ + ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( + use_in_memory_cache=True, + persist_owner_state=False, + ) + do_manager.datasite_owner_syncer.perm_context.open(".").grant_write_access( + ds_manager.email + ) + + do_email = do_manager.email + + # Sync well past both thresholds, on the same code path a real client takes. + for i in range(6): + ds_manager._send_file_change(f"{do_email}/test{i}.txt", f"Content {i}") + do_manager.sync(auto_checkpoint=True, checkpoint_threshold=1) + + router = do_manager._connection_router + assert router.get_latest_checkpoint() is None + assert router.get_all_incremental_checkpoints() == [] + assert router.get_rolling_state() is None + # The append-only event log is owner-only too, so it stays empty. + assert router.owner_get_all_accepted_event_file_ids() == [] + + # No local rolling state file either. + rolling_state_path = ( + do_manager.datasite_owner_syncer.syftbox_folder + / ".cache" + / "rolling_state.json" + ) + assert not rolling_state_path.exists() + + # The automatic checkpoint path stays silent rather than raising. + assert do_manager.try_create_checkpoint(threshold=1) is None + + # Sync still works: the DO holds every file and the DS receives them back + # through the outbox, which is peer-facing and so unaffected. + do_cache = do_manager.datasite_owner_syncer.event_cache + assert len(do_cache.file_hashes) == 6 + + ds_manager.sync() + ds_cache = ds_manager.datasite_watcher_syncer.datasite_watcher_cache + ds_paths = {str(event.path_in_datasite) for event in ds_cache.get_cached_events()} + for i in range(6): + assert f"test{i}.txt" in ds_paths + + +def test_disabled_owner_state_does_not_restore(): + """A datasite that persists no owner state comes back empty, by design. + + This is the behaviour an enclave wants: its keypair is ephemeral, so + restoring a previous boot's state does not fit the security model. The + restore path must not silently fall back to replaying the event log. + """ + ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( + use_in_memory_cache=True, + persist_owner_state=False, + ) + do_manager.datasite_owner_syncer.perm_context.open(".").grant_write_access( + ds_manager.email + ) + + do_email = do_manager.email + ds_manager._send_file_change(f"{do_email}/test1.txt", "Content 1") + ds_manager._send_file_change(f"{do_email}/test2.txt", "Content 2") + do_manager.sync() + assert len(do_manager.datasite_owner_syncer.event_cache.file_hashes) == 2 + + # Simulate a fresh boot: local cache gone, nothing durable behind it. + do_manager.datasite_owner_syncer.event_cache.clear_cache() + do_manager.datasite_owner_syncer.initial_sync_done = False + + # No event download should be attempted at all. + downloads = 0 + syncer = do_manager.datasite_owner_syncer + original = syncer.download_events_message_by_id_with_connection + + def counted(event_id): + nonlocal downloads + downloads += 1 + return original(event_id) + + syncer.download_events_message_by_id_with_connection = counted + + do_manager.sync() + + assert downloads == 0 + assert len(syncer.event_cache.file_hashes) == 0 + + +def test_owner_state_is_written_by_default(): + """The inverse of the flag: by default the owner-only state is all there. + + Asserted on the backend rather than on a restored cache, because what the + flag controls is what gets written - the testing config runs with + write_files=False, so a restored cache is not a sound signal here. + """ + ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( + use_in_memory_cache=True, + ) + do_manager.datasite_owner_syncer.perm_context.open(".").grant_write_access( + ds_manager.email + ) + + do_email = do_manager.email + ds_manager._send_file_change(f"{do_email}/test1.txt", "Content 1") + ds_manager._send_file_change(f"{do_email}/test2.txt", "Content 2") + do_manager.sync(auto_checkpoint=False) + + router = do_manager._connection_router + assert router.owner_get_all_accepted_event_file_ids() != [] + assert router.get_rolling_state() is not None + assert do_manager.datasite_owner_syncer._rolling_state is not None diff --git a/tests/unit/test_rolling_state.py b/tests/unit/test_rolling_state.py index c3afeec2000..1ef1dbce756 100644 --- a/tests/unit/test_rolling_state.py +++ b/tests/unit/test_rolling_state.py @@ -9,6 +9,7 @@ import time from syft.sync.syftbox_manager import SyftboxManager from syft.sync.checkpoints.rolling_state import RollingState +from syft.sync.sync.constants import CACHE_DIR, ROLLING_STATE_FILENAME from tests.unit.utils import get_mock_event @@ -244,3 +245,28 @@ def test_rolling_state_clear_resets_base_timestamp(): assert rs.event_count == 0 assert rs.base_checkpoint_timestamp == 2000.0 assert rs.last_event_timestamp is None + + +def test_no_rolling_state_when_owner_state_disabled(): + """Rolling state is owner-only state, so it goes off with the rest. + + This is the write that costs the most: the upload threshold is 1, so + without this gate an enclave would push a rolling state per event. + """ + ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( + persist_owner_state=False, + ) + do_manager.datasite_owner_syncer.perm_context.open(".").grant_write_access( + ds_manager.email + ) + + ds_manager._send_file_change(f"{do_manager.email}/file1.txt", "content1") + ds_manager._send_file_change(f"{do_manager.email}/file2.txt", "content2") + do_manager.sync() + + assert do_manager._connection_router.get_rolling_state() is None + + syncer = do_manager.datasite_owner_syncer + assert not (syncer.syftbox_folder / CACHE_DIR / ROLLING_STATE_FILENAME).exists() + # Nothing is tracked in memory either - there is nothing to track it for. + assert syncer._rolling_state is None