From 7bac9664f4f0ac30cfda3d468edfcc39efbe1462 Mon Sep 17 00:00:00 2001 From: Koen van der Veen Date: Tue, 8 Sep 2026 14:38:20 +0200 Subject: [PATCH 1/2] feat: make checkpoints configurable and turn them off for enclaves Adds a use_checkpoints flag to the sync engine, defaulting to True so notebook and Colab clients are unaffected. When off, the DO syncer skips checkpoint uploads, rolling-state uploads (one Drive write per event) and the checkpoint reads on cold start, falling through to the existing download-all-events path. Enclaves default to off via SYFT_ENCLAVE_USE_CHECKPOINTS: an enclave boots with fresh_state, so there is never a snapshot to restore and every checkpoint write just spends Drive API calls inside the poll loop. sync(auto_checkpoint=...) is unchanged but subordinate to the config, so no existing sync() call site needed touching. --- .../docker/Dockerfile | 1 + .../src/enclave_model_api/__main__.py | 4 +- packages/syft-enclave/docker/Dockerfile | 1 + .../src/syft_enclaves/__main__.py | 4 +- .../syft-enclave/src/syft_enclaves/client.py | 5 + .../src/syft_enclaves/settings.py | 11 + packages/syft-enclave/tests/test_settings.py | 33 +++ syft/sync/syftbox_manager.py | 30 +- syft/sync/sync/datasite_owner_syncer.py | 268 ++++++++++++------ tests/unit/test_checkpoints.py | 86 ++++++ tests/unit/test_rolling_state.py | 26 ++ 11 files changed, 373 insertions(+), 96 deletions(-) diff --git a/packages/enclave-model-api-example/docker/Dockerfile b/packages/enclave-model-api-example/docker/Dockerfile index 92884788bf7..917318ec337 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_USE_CHECKPOINTS,\ 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..88b42cd5574 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"use_checkpoints={settings.use_checkpoints}" ) 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, + use_checkpoints=settings.use_checkpoints, ) logger.info("SyftEnclaveClient ready") diff --git a/packages/syft-enclave/docker/Dockerfile b/packages/syft-enclave/docker/Dockerfile index 8b1e5cb6027..2f66f9f0ecb 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_USE_CHECKPOINTS,\ 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..5534141bd87 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"use_checkpoints={settings.use_checkpoints}" ) 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, + use_checkpoints=settings.use_checkpoints, ) 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..d5a34125511 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, + use_checkpoints: bool = False, ) -> "SyftEnclaveClient": """Build an enclave client backed by a real Google Drive connection. Args: @@ -442,6 +443,9 @@ 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. + use_checkpoints: Write and restore sync checkpoints. Off by + default - an enclave wipes its state on boot, so there is + nothing for a checkpoint to restore. """ config = SyftRDSClientConfig.for_jupyter( email=email, @@ -449,6 +453,7 @@ def for_enclave( has_do_role=True, token_path=Path(token_path) if token_path is not None else None, encryption=encryption, + use_checkpoints=use_checkpoints, ) # 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..2decd7e4835 100644 --- a/packages/syft-enclave/src/syft_enclaves/settings.py +++ b/packages/syft-enclave/src/syft_enclaves/settings.py @@ -94,3 +94,14 @@ def _split_data_owners(cls, v: object) -> object: "Enabled by default; set false to disable." ), ) + use_checkpoints: bool = Field( + default=False, + description=( + "Write and restore sync checkpoints. Off by default: checkpoints " + "only pay off on a cold start, and an enclave boots with " + "fresh_state, so there is never a snapshot to restore — every " + "checkpoint write would just spend Drive API calls inside the " + "poll loop. 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..8409fb37d80 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.use_checkpoints is False # default: no cold start to speed up def test_data_owners_parsed_from_comma_separated_string(clean_env): @@ -112,3 +113,35 @@ 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_use_checkpoints_can_be_enabled_via_env(required_env): + required_env.setenv("SYFT_ENCLAVE_USE_CHECKPOINTS", "true") + settings = EnclaveSettings(_env_file=None) + assert settings.use_checkpoints is True + + +def test_use_checkpoints_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 checkpoints back. + """ + from syft_rds.config import SyftRDSClientConfig + + config = SyftRDSClientConfig.for_jupyter( + email="enclave@openmined.org", + has_do_role=True, + has_ds_role=True, + use_checkpoints=False, + ) + assert config.sync.use_checkpoints is False + assert config.sync.datasite_owner_syncer_config.use_checkpoints 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.use_checkpoints is True diff --git a/syft/sync/syftbox_manager.py b/syft/sync/syftbox_manager.py index 3585538682d..b376e82db46 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -113,6 +113,9 @@ class SyftboxManagerConfig(BaseModel): has_ds_role: bool = False has_do_role: bool = False use_in_memory_cache: bool = True + # Checkpoints speed up a cold start at the cost of Drive writes. Off for + # datasites that never restore (see syft-enclave's EnclaveSettings). + use_checkpoints: bool = True datasite_owner_syncer_config: DatasiteOwnerSyncerConfig peer_manager_config: PeerManagerConfig @@ -126,6 +129,7 @@ def for_colab( has_ds_role: bool = False, has_do_role: bool = False, encryption: bool = False, + use_checkpoints: bool = True, crypto_keys_path: Path | None = None, skip_peer_on_patch_version_diff: Optional[ bool @@ -150,6 +154,7 @@ def for_colab( connection_configs = [GdriveConnectionConfig(email=email, token_path=None)] datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, + use_checkpoints=use_checkpoints, syftbox_folder=syftbox_folder, collections_folder=collections_folder, collection_specs=collection_specs, @@ -192,6 +197,7 @@ def for_colab( has_do_role=has_do_role, connection_configs=connection_configs, use_in_memory_cache=False, + use_checkpoints=use_checkpoints, datasite_owner_syncer_config=datasite_owner_syncer_config, datasite_watcher_syncer_config=datasite_watcher_syncer_config, peer_manager_config=peer_manager_config, @@ -205,6 +211,7 @@ def for_jupyter( has_do_role: bool = False, token_path: Path | None = None, encryption: bool = False, + use_checkpoints: bool = True, crypto_keys_path: Path | None = None, skip_peer_on_patch_version_diff: Optional[ bool @@ -232,6 +239,7 @@ def for_jupyter( ] datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, + use_checkpoints=use_checkpoints, syftbox_folder=syftbox_folder, collections_folder=collections_folder, collection_specs=collection_specs, @@ -274,6 +282,7 @@ def for_jupyter( has_ds_role=has_ds_role, has_do_role=has_do_role, use_in_memory_cache=False, + use_checkpoints=use_checkpoints, datasite_owner_syncer_config=datasite_owner_syncer_config, datasite_watcher_syncer_config=datasite_watcher_syncer_config, peer_manager_config=peer_manager_config, @@ -288,6 +297,7 @@ def _base_config_for_testing( has_ds_role: bool = False, has_do_role: bool = False, use_in_memory_cache: bool = True, + use_checkpoints: bool = True, check_versions: bool = False, collection_specs: list["CollectionSyncSpec"] | None = None, ): @@ -306,6 +316,7 @@ def _base_config_for_testing( datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, + use_checkpoints=use_checkpoints, syftbox_folder=syftbox_folder, collections_folder=collections_folder, collection_specs=collection_specs, @@ -346,6 +357,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, + use_checkpoints=use_checkpoints, datasite_owner_syncer_config=datasite_owner_syncer_config, datasite_watcher_syncer_config=datasite_watcher_syncer_config, peer_manager_config=peer_manager_config, @@ -361,6 +373,7 @@ def for_google_drive_testing_connection( has_ds_role: bool = False, has_do_role: bool = False, use_in_memory_cache: bool = True, + use_checkpoints: bool = True, check_versions: bool = False, collection_specs: list["CollectionSyncSpec"] | None = None, ): @@ -381,6 +394,7 @@ def for_google_drive_testing_connection( ] datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, + use_checkpoints=use_checkpoints, syftbox_folder=syftbox_folder, collections_folder=collections_folder, collection_specs=collection_specs, @@ -418,6 +432,7 @@ def for_google_drive_testing_connection( email=email, syftbox_folder=syftbox_folder, write_files=write_files, + use_checkpoints=use_checkpoints, datasite_owner_syncer_config=datasite_owner_syncer_config, datasite_watcher_syncer_config=datasite_watcher_syncer_config, has_ds_role=has_ds_role, @@ -587,6 +602,7 @@ def for_colab( has_ds_role: bool = False, has_do_role: bool = False, encryption: bool = False, + use_checkpoints: bool = True, crypto_keys_path: Path | None = None, skip_peer_on_patch_version_diff: Optional[ bool @@ -599,6 +615,7 @@ def for_colab( has_ds_role=has_ds_role, has_do_role=has_do_role, encryption=encryption, + use_checkpoints=use_checkpoints, 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 +631,7 @@ def for_jupyter( has_do_role: bool = False, token_path: Path | None = None, encryption: bool = False, + use_checkpoints: bool = True, crypto_keys_path: Path | None = None, skip_peer_on_patch_version_diff: Optional[ bool @@ -629,6 +647,7 @@ def for_jupyter( has_do_role=has_do_role, token_path=token_path, encryption=encryption, + use_checkpoints=use_checkpoints, 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 +667,7 @@ def _pair_with_google_drive_testing_connection( add_peers: bool = True, load_peers: bool = False, use_in_memory_cache: bool = True, + use_checkpoints: bool = True, clear_caches: bool = True, check_versions: bool = False, collection_specs: list["CollectionSyncSpec"] | None = None, @@ -656,6 +676,7 @@ def _pair_with_google_drive_testing_connection( email=do_email, syftbox_folder=base_path1, use_in_memory_cache=use_in_memory_cache, + use_checkpoints=use_checkpoints, token_path=do_token_path, has_ds_role=False, has_do_role=True, @@ -669,6 +690,7 @@ def _pair_with_google_drive_testing_connection( email=ds_email, syftbox_folder=base_path2, use_in_memory_cache=use_in_memory_cache, + use_checkpoints=use_checkpoints, token_path=ds_token_path, has_ds_role=True, has_do_role=False, @@ -732,6 +754,7 @@ def pair_with_mock_drive_service_connection( sync_automatically: bool = False, add_peers: bool = True, use_in_memory_cache: bool = True, + use_checkpoints: bool = True, check_versions: bool = False, encryption: bool = False, collection_specs: list["CollectionSyncSpec"] | None = None, @@ -750,6 +773,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 + use_checkpoints: Whether the DO writes and restores checkpoints check_versions: Whether to check protocol/client versions Returns: @@ -762,6 +786,7 @@ def pair_with_mock_drive_service_connection( has_ds_role=False, has_do_role=True, use_in_memory_cache=use_in_memory_cache, + use_checkpoints=use_checkpoints, check_versions=check_versions, collection_specs=collection_specs, ) @@ -772,6 +797,7 @@ def pair_with_mock_drive_service_connection( has_ds_role=True, has_do_role=False, use_in_memory_cache=use_in_memory_cache, + use_checkpoints=use_checkpoints, check_versions=check_versions, collection_specs=collection_specs, ) @@ -933,7 +959,9 @@ 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 + use_checkpoints=False. 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..983fb950b59 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,36 @@ 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 + # Checkpoints only speed up a cold start. Turn them off for a datasite that + # never restores (e.g. an enclave booting with fresh_state) to save the + # Drive writes they cost. + use_checkpoints: 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 +101,7 @@ class DatasiteOwnerSyncer(BaseModelCallbackMixin): default_factory=lambda: DataSiteOwnerEventCache() ) write_files: bool = True + use_checkpoints: bool = True connection_router: ConnectionRouter initial_sync_done: bool = False email: str @@ -115,6 +142,7 @@ def from_config(cls, config: DatasiteOwnerSyncerConfig): return cls( event_cache=DataSiteOwnerEventCache.from_config(config.cache_config), write_files=config.write_files, + use_checkpoints=config.use_checkpoints, connection_router=ConnectionRouter.from_configs( config.email, config.connection_configs ), @@ -155,7 +183,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.use_checkpoints: + self._load_rolling_state() try: if not self.initial_sync_done: self.pull_initial_state() @@ -181,7 +210,8 @@ def sync(self, peer_emails: list[str], recompute_hashes: bool = True): ) self.process_syftbox_events_queue() finally: - self._save_rolling_state() + if self.use_checkpoints: + self._save_rolling_state() def download_events_message_by_id_with_connection( self, events_message_id: str @@ -210,107 +240,142 @@ 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() + 1. Restore from the full checkpoint, the incrementals and the rolling + state. Skipped entirely when checkpoints are off, which leaves the + download-all-events fallback below. + 2. Download any events newer than the restored cursor, or all events + when there is no cursor and no checkpoint to anchor one. + 3. Ensure events_messages_connection is populated for get_cached_events() """ - events_since_timestamp: float | None = None - restored_events: list[FileChangeEvent] = [] + restore = self._restore_checkpoints_or_start_empty() + + 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) + + # 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 - # Step 1: Check for full (compacted) checkpoint + def _restore_checkpoints_or_start_empty(self) -> _CheckpointRestore: + """Restore from checkpoints, or begin from an empty rolling state. + + With checkpoints off there is nothing on Drive to restore, so we skip + the three downloads entirely and let the caller fall back to events. + """ + if self.use_checkpoints: + return self._restore_from_checkpoints() + + self._rolling_state = RollingState( + email=self.email, + base_checkpoint_timestamp=0.0, + ) + return _CheckpointRestore( + events_since_timestamp=None, + events=[], + found_checkpoint=False, + ) + + 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), + ) + + 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 - # Step 3: Check for rolling state + 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, + base_checkpoint_timestamp=since or 0.0, ) + return since, [] - # 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) - ) - for events_message in events_messages_list: - self.event_cache.add_events_message_to_local_cache(events_message) - - # 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 @@ -852,7 +917,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.use_checkpoints or self._rolling_state is None: return self._rolling_state.add_events_message(events_message) @@ -864,6 +929,8 @@ def _add_events_to_rolling_state( def _upload_rolling_state(self) -> None: """Upload the in-memory rolling state to GDrive.""" + if not self.use_checkpoints: + return if self._rolling_state is None or self._rolling_state.event_count == 0: return @@ -874,6 +941,14 @@ def _upload_rolling_state(self) -> None: # CHECKPOINT METHODS # ========================================================================= + def _raise_if_checkpoints_disabled(self) -> None: + """Refuse an explicit checkpoint request when checkpoints are off.""" + if not self.use_checkpoints: + raise ValueError( + "This client was created with use_checkpoints=False, so it " + "cannot create checkpoints." + ) + def create_incremental_checkpoint(self) -> IncrementalCheckpoint: """ Create an incremental checkpoint from the current rolling state. @@ -885,6 +960,7 @@ def create_incremental_checkpoint(self) -> IncrementalCheckpoint: Returns: The created IncrementalCheckpoint object. """ + self._raise_if_checkpoints_disabled() if self._rolling_state is None or self._rolling_state.event_count == 0: raise ValueError("No rolling state to create checkpoint from") @@ -941,6 +1017,7 @@ def create_checkpoint(self) -> Checkpoint: Returns: The created Checkpoint object. """ + self._raise_if_checkpoints_disabled() self._load_rolling_state() try: last_event_timestamp = self.event_cache.get_latest_event_timestamp() @@ -1010,6 +1087,7 @@ def compact_checkpoints(self) -> Checkpoint: Returns: The compacted Checkpoint object. """ + self._raise_if_checkpoints_disabled() print("Compacting incremental checkpoints...") # Download existing full checkpoint (if exists) @@ -1109,8 +1187,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 has checkpoints turned off. """ + if not self.use_checkpoints: + 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..cc12585dc8f 100644 --- a/tests/unit/test_checkpoints.py +++ b/tests/unit/test_checkpoints.py @@ -10,6 +10,8 @@ from pathlib import Path from uuid import uuid4 +import pytest + from syft.sync.checkpoints.checkpoint import ( Checkpoint, IncrementalCheckpoint, @@ -603,3 +605,87 @@ 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_checkpoints_written_when_disabled(): + """With use_checkpoints=False, nothing checkpoint-shaped reaches the drive. + + The second half is the point of the flag: turning checkpoints off must not + cost correctness, so the DS still has to see every DO change. + """ + ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( + use_in_memory_cache=True, + use_checkpoints=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 + + # 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() + + # Sync still works: the DO holds every file and the DS receives them back. + 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_checkpoints_restore_via_all_events(): + """A cold start with checkpoints off rebuilds state from the events instead.""" + ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( + use_in_memory_cache=True, + use_checkpoints=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 login: no checkpoint exists to restore from. + do_manager.datasite_owner_syncer.event_cache.clear_cache() + do_manager.datasite_owner_syncer.initial_sync_done = False + + do_manager.sync() + + assert len(do_manager.datasite_owner_syncer.event_cache.file_hashes) == 2 + + +def test_explicit_checkpoint_calls_raise_when_disabled(): + """An explicit request fails loudly rather than silently doing nothing.""" + _, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( + use_in_memory_cache=True, + use_checkpoints=False, + ) + + with pytest.raises(ValueError, match="use_checkpoints=False"): + do_manager.create_checkpoint() + + # The automatic path stays silent instead - it is not a caller error. + assert do_manager.try_create_checkpoint(threshold=1) is None diff --git a/tests/unit/test_rolling_state.py b/tests/unit/test_rolling_state.py index c3afeec2000..5499d4f728d 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_checkpoints_disabled(): + """Rolling state is part of the checkpoint machinery, so it goes off too. + + 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( + use_checkpoints=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() + # The in-memory state is initialised but never accumulates events. + assert syncer._rolling_state.event_count == 0 From b1eece9364613c77081a7b5fc55c06bb4fb84e16 Mon Sep 17 00:00:00 2001 From: Koen van der Veen Date: Tue, 8 Sep 2026 15:01:03 +0200 Subject: [PATCH 2/2] refactor: widen the flag to all owner-only state and rename to persist_owner_state The first pass only covered checkpoints and the rolling state, so the DO still wrote its append-only event log to Drive and could still restore from it. That misses the point: an enclave gets an ephemeral keypair each boot, so restoring a previous boot's state does not fit the security model and none of those writes are needed. Renames use_checkpoints to persist_owner_state and extends it to the event log, gated at the single enqueue site in queue_event_for_syftbox. The peer outbox is a separate queue and is unaffected, so peers still receive every change. pull_initial_state now skips the whole owner-state restore rather than falling back to replaying the log; collections stay restored either way since they are peer-facing. Also drops the ValueError guards on the manual checkpoint methods - callers simply do not call them. --- .../docker/Dockerfile | 2 +- .../src/enclave_model_api/__main__.py | 4 +- packages/syft-enclave/docker/Dockerfile | 2 +- .../src/syft_enclaves/__main__.py | 4 +- .../syft-enclave/src/syft_enclaves/client.py | 11 ++- .../src/syft_enclaves/settings.py | 16 ++-- packages/syft-enclave/tests/test_settings.py | 21 ++-- syft/sync/syftbox_manager.py | 56 +++++------ syft/sync/sync/datasite_owner_syncer.py | 96 ++++++++----------- tests/unit/test_checkpoints.py | 77 +++++++++++---- tests/unit/test_rolling_state.py | 10 +- 11 files changed, 160 insertions(+), 139 deletions(-) diff --git a/packages/enclave-model-api-example/docker/Dockerfile b/packages/enclave-model-api-example/docker/Dockerfile index 917318ec337..f896b36b6ce 100644 --- a/packages/enclave-model-api-example/docker/Dockerfile +++ b/packages/enclave-model-api-example/docker/Dockerfile @@ -28,7 +28,7 @@ SYFT_ENCLAVE_DATA_OWNERS,\ SYFT_ENCLAVE_REQUIRE_TEE,\ SYFT_ENCLAVE_FRESH_STATE,\ SYFT_ENCLAVE_USE_ENCRYPTION,\ -SYFT_ENCLAVE_USE_CHECKPOINTS,\ +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 88b42cd5574..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 @@ -49,7 +49,7 @@ def main() -> None: 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_checkpoints={settings.use_checkpoints}" + f"persist_owner_state={settings.persist_owner_state}" ) logger.info( f"Inference settings — model_owner={inference.model_owner} " @@ -63,7 +63,7 @@ def main() -> None: token_path=settings.token_path, data_owners=settings.data_owners, encryption=settings.use_encryption, - use_checkpoints=settings.use_checkpoints, + 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 2f66f9f0ecb..38bbdd79302 100644 --- a/packages/syft-enclave/docker/Dockerfile +++ b/packages/syft-enclave/docker/Dockerfile @@ -54,7 +54,7 @@ SYFT_ENCLAVE_DATA_OWNERS,\ SYFT_ENCLAVE_REQUIRE_TEE,\ SYFT_ENCLAVE_FRESH_STATE,\ SYFT_ENCLAVE_USE_ENCRYPTION,\ -SYFT_ENCLAVE_USE_CHECKPOINTS,\ +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 5534141bd87..ef12ce1f7c4 100644 --- a/packages/syft-enclave/src/syft_enclaves/__main__.py +++ b/packages/syft-enclave/src/syft_enclaves/__main__.py @@ -44,7 +44,7 @@ def main() -> None: 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_checkpoints={settings.use_checkpoints}" + f"persist_owner_state={settings.persist_owner_state}" ) logger.info("Building SyftEnclaveClient...") @@ -53,7 +53,7 @@ def main() -> None: token_path=settings.token_path, data_owners=settings.data_owners, encryption=settings.use_encryption, - use_checkpoints=settings.use_checkpoints, + 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 d5a34125511..95f6d4398a6 100644 --- a/packages/syft-enclave/src/syft_enclaves/client.py +++ b/packages/syft-enclave/src/syft_enclaves/client.py @@ -435,7 +435,7 @@ def for_enclave( token_path: Path | str | None = None, data_owners: list[str] | None = None, encryption: bool = False, - use_checkpoints: bool = False, + persist_owner_state: bool = False, ) -> "SyftEnclaveClient": """Build an enclave client backed by a real Google Drive connection. Args: @@ -443,9 +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. - use_checkpoints: Write and restore sync checkpoints. Off by - default - an enclave wipes its state on boot, so there is - nothing for a checkpoint to restore. + 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, @@ -453,7 +454,7 @@ def for_enclave( has_do_role=True, token_path=Path(token_path) if token_path is not None else None, encryption=encryption, - use_checkpoints=use_checkpoints, + 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 2decd7e4835..e1ff1d01dda 100644 --- a/packages/syft-enclave/src/syft_enclaves/settings.py +++ b/packages/syft-enclave/src/syft_enclaves/settings.py @@ -94,14 +94,16 @@ def _split_data_owners(cls, v: object) -> object: "Enabled by default; set false to disable." ), ) - use_checkpoints: bool = Field( + persist_owner_state: bool = Field( default=False, description=( - "Write and restore sync checkpoints. Off by default: checkpoints " - "only pay off on a cold start, and an enclave boots with " - "fresh_state, so there is never a snapshot to restore — every " - "checkpoint write would just spend Drive API calls inside the " - "poll loop. Set true only for a stateful enclave " - "(fresh_state=false)." + "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 8409fb37d80..32da55a0731 100644 --- a/packages/syft-enclave/tests/test_settings.py +++ b/packages/syft-enclave/tests/test_settings.py @@ -38,7 +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.use_checkpoints is False # default: no cold start to speed up + assert settings.persist_owner_state is False # default: an enclave never restores def test_data_owners_parsed_from_comma_separated_string(clean_env): @@ -115,18 +115,19 @@ def test_use_encryption_can_be_disabled_via_env(required_env): assert settings.use_encryption is False -def test_use_checkpoints_can_be_enabled_via_env(required_env): - required_env.setenv("SYFT_ENCLAVE_USE_CHECKPOINTS", "true") +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.use_checkpoints is True + assert settings.persist_owner_state is True -def test_use_checkpoints_reaches_the_do_syncer_config(): +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 checkpoints back. + 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 @@ -134,14 +135,14 @@ def test_use_checkpoints_reaches_the_do_syncer_config(): email="enclave@openmined.org", has_do_role=True, has_ds_role=True, - use_checkpoints=False, + persist_owner_state=False, ) - assert config.sync.use_checkpoints is False - assert config.sync.datasite_owner_syncer_config.use_checkpoints is 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.use_checkpoints is 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 b376e82db46..aea0d0c2375 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -113,9 +113,10 @@ class SyftboxManagerConfig(BaseModel): has_ds_role: bool = False has_do_role: bool = False use_in_memory_cache: bool = True - # Checkpoints speed up a cold start at the cost of Drive writes. Off for - # datasites that never restore (see syft-enclave's EnclaveSettings). - use_checkpoints: 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 @@ -129,7 +130,7 @@ def for_colab( has_ds_role: bool = False, has_do_role: bool = False, encryption: bool = False, - use_checkpoints: bool = True, + persist_owner_state: bool = True, crypto_keys_path: Path | None = None, skip_peer_on_patch_version_diff: Optional[ bool @@ -154,7 +155,7 @@ def for_colab( connection_configs = [GdriveConnectionConfig(email=email, token_path=None)] datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, - use_checkpoints=use_checkpoints, + persist_owner_state=persist_owner_state, syftbox_folder=syftbox_folder, collections_folder=collections_folder, collection_specs=collection_specs, @@ -197,7 +198,7 @@ def for_colab( has_do_role=has_do_role, connection_configs=connection_configs, use_in_memory_cache=False, - use_checkpoints=use_checkpoints, + 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, @@ -211,7 +212,7 @@ def for_jupyter( has_do_role: bool = False, token_path: Path | None = None, encryption: bool = False, - use_checkpoints: bool = True, + persist_owner_state: bool = True, crypto_keys_path: Path | None = None, skip_peer_on_patch_version_diff: Optional[ bool @@ -239,7 +240,7 @@ def for_jupyter( ] datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, - use_checkpoints=use_checkpoints, + persist_owner_state=persist_owner_state, syftbox_folder=syftbox_folder, collections_folder=collections_folder, collection_specs=collection_specs, @@ -282,7 +283,7 @@ def for_jupyter( has_ds_role=has_ds_role, has_do_role=has_do_role, use_in_memory_cache=False, - use_checkpoints=use_checkpoints, + 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, @@ -297,7 +298,7 @@ def _base_config_for_testing( has_ds_role: bool = False, has_do_role: bool = False, use_in_memory_cache: bool = True, - use_checkpoints: bool = True, + persist_owner_state: bool = True, check_versions: bool = False, collection_specs: list["CollectionSyncSpec"] | None = None, ): @@ -316,7 +317,7 @@ def _base_config_for_testing( datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, - use_checkpoints=use_checkpoints, + persist_owner_state=persist_owner_state, syftbox_folder=syftbox_folder, collections_folder=collections_folder, collection_specs=collection_specs, @@ -357,7 +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, - use_checkpoints=use_checkpoints, + 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, @@ -373,7 +374,7 @@ def for_google_drive_testing_connection( has_ds_role: bool = False, has_do_role: bool = False, use_in_memory_cache: bool = True, - use_checkpoints: bool = True, + persist_owner_state: bool = True, check_versions: bool = False, collection_specs: list["CollectionSyncSpec"] | None = None, ): @@ -394,7 +395,7 @@ def for_google_drive_testing_connection( ] datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, - use_checkpoints=use_checkpoints, + persist_owner_state=persist_owner_state, syftbox_folder=syftbox_folder, collections_folder=collections_folder, collection_specs=collection_specs, @@ -432,7 +433,7 @@ def for_google_drive_testing_connection( email=email, syftbox_folder=syftbox_folder, write_files=write_files, - use_checkpoints=use_checkpoints, + 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, @@ -602,7 +603,7 @@ def for_colab( has_ds_role: bool = False, has_do_role: bool = False, encryption: bool = False, - use_checkpoints: bool = True, + persist_owner_state: bool = True, crypto_keys_path: Path | None = None, skip_peer_on_patch_version_diff: Optional[ bool @@ -615,7 +616,7 @@ def for_colab( has_ds_role=has_ds_role, has_do_role=has_do_role, encryption=encryption, - use_checkpoints=use_checkpoints, + 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, @@ -631,7 +632,7 @@ def for_jupyter( has_do_role: bool = False, token_path: Path | None = None, encryption: bool = False, - use_checkpoints: bool = True, + persist_owner_state: bool = True, crypto_keys_path: Path | None = None, skip_peer_on_patch_version_diff: Optional[ bool @@ -647,7 +648,7 @@ def for_jupyter( has_do_role=has_do_role, token_path=token_path, encryption=encryption, - use_checkpoints=use_checkpoints, + 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, @@ -667,7 +668,7 @@ def _pair_with_google_drive_testing_connection( add_peers: bool = True, load_peers: bool = False, use_in_memory_cache: bool = True, - use_checkpoints: bool = True, + persist_owner_state: bool = True, clear_caches: bool = True, check_versions: bool = False, collection_specs: list["CollectionSyncSpec"] | None = None, @@ -676,7 +677,7 @@ def _pair_with_google_drive_testing_connection( email=do_email, syftbox_folder=base_path1, use_in_memory_cache=use_in_memory_cache, - use_checkpoints=use_checkpoints, + persist_owner_state=persist_owner_state, token_path=do_token_path, has_ds_role=False, has_do_role=True, @@ -690,7 +691,7 @@ def _pair_with_google_drive_testing_connection( email=ds_email, syftbox_folder=base_path2, use_in_memory_cache=use_in_memory_cache, - use_checkpoints=use_checkpoints, + persist_owner_state=persist_owner_state, token_path=ds_token_path, has_ds_role=True, has_do_role=False, @@ -754,7 +755,7 @@ def pair_with_mock_drive_service_connection( sync_automatically: bool = False, add_peers: bool = True, use_in_memory_cache: bool = True, - use_checkpoints: bool = True, + persist_owner_state: bool = True, check_versions: bool = False, encryption: bool = False, collection_specs: list["CollectionSyncSpec"] | None = None, @@ -773,7 +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 - use_checkpoints: Whether the DO writes and restores checkpoints + persist_owner_state: Whether the DO keeps restorable owner state check_versions: Whether to check protocol/client versions Returns: @@ -786,7 +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, - use_checkpoints=use_checkpoints, + persist_owner_state=persist_owner_state, check_versions=check_versions, collection_specs=collection_specs, ) @@ -797,7 +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, - use_checkpoints=use_checkpoints, + persist_owner_state=persist_owner_state, check_versions=check_versions, collection_specs=collection_specs, ) @@ -961,7 +962,8 @@ def sync( auto_checkpoint: If True, automatically create checkpoint when event count exceeds threshold (DO only). Has no effect when this manager was built with - use_checkpoints=False. + 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 983fb950b59..9ba26e1188b 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -78,10 +78,12 @@ class DatasiteOwnerSyncerConfig(BaseModel): email: str syftbox_folder: Path write_files: bool = True - # Checkpoints only speed up a cold start. Turn them off for a datasite that - # never restores (e.g. an enclave booting with fresh_state) to save the - # Drive writes they cost. - use_checkpoints: 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 @@ -101,7 +103,7 @@ class DatasiteOwnerSyncer(BaseModelCallbackMixin): default_factory=lambda: DataSiteOwnerEventCache() ) write_files: bool = True - use_checkpoints: bool = True + persist_owner_state: bool = True connection_router: ConnectionRouter initial_sync_done: bool = False email: str @@ -142,7 +144,7 @@ def from_config(cls, config: DatasiteOwnerSyncerConfig): return cls( event_cache=DataSiteOwnerEventCache.from_config(config.cache_config), write_files=config.write_files, - use_checkpoints=config.use_checkpoints, + persist_owner_state=config.persist_owner_state, connection_router=ConnectionRouter.from_configs( config.email, config.connection_configs ), @@ -183,7 +185,7 @@ def _save_rolling_state(self) -> None: tmp.rename(path) def sync(self, peer_emails: list[str], recompute_hashes: bool = True): - if self.use_checkpoints: + if self.persist_owner_state: self._load_rolling_state() try: if not self.initial_sync_done: @@ -210,7 +212,7 @@ def sync(self, peer_emails: list[str], recompute_hashes: bool = True): ) self.process_syftbox_events_queue() finally: - if self.use_checkpoints: + if self.persist_owner_state: self._save_rolling_state() def download_events_message_by_id_with_connection( @@ -239,15 +241,26 @@ def pull_initial_state(self): """ Pull initial state from Google Drive. - Flow: - 1. Restore from the full checkpoint, the incrementals and the rolling - state. Skipped entirely when checkpoints are off, which leaves the - download-all-events fallback below. - 2. Download any events newer than the restored cursor, or all events - when there is no cursor and no checkpoint to anchor one. - 3. 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. """ - restore = self._restore_checkpoints_or_start_empty() + 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() if restore.events_since_timestamp is not None: self._download_events_since(restore.events_since_timestamp) @@ -260,32 +273,6 @@ def pull_initial_state(self): if restore.events and not self.event_cache.get_cached_events(): self._write_events_to_messages_cache(restore.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() - - self.initial_sync_done = True - - def _restore_checkpoints_or_start_empty(self) -> _CheckpointRestore: - """Restore from checkpoints, or begin from an empty rolling state. - - With checkpoints off there is nothing on Drive to restore, so we skip - the three downloads entirely and let the caller fall back to events. - """ - if self.use_checkpoints: - return self._restore_from_checkpoints() - - self._rolling_state = RollingState( - email=self.email, - base_checkpoint_timestamp=0.0, - ) - return _CheckpointRestore( - events_since_timestamp=None, - events=[], - found_checkpoint=False, - ) - def _restore_from_checkpoints(self) -> _CheckpointRestore: """Apply the full checkpoint, then the incrementals, then the rolling state. @@ -774,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)) @@ -917,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 not self.use_checkpoints or 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) @@ -929,7 +920,7 @@ def _add_events_to_rolling_state( def _upload_rolling_state(self) -> None: """Upload the in-memory rolling state to GDrive.""" - if not self.use_checkpoints: + if not self.persist_owner_state: return if self._rolling_state is None or self._rolling_state.event_count == 0: return @@ -941,14 +932,6 @@ def _upload_rolling_state(self) -> None: # CHECKPOINT METHODS # ========================================================================= - def _raise_if_checkpoints_disabled(self) -> None: - """Refuse an explicit checkpoint request when checkpoints are off.""" - if not self.use_checkpoints: - raise ValueError( - "This client was created with use_checkpoints=False, so it " - "cannot create checkpoints." - ) - def create_incremental_checkpoint(self) -> IncrementalCheckpoint: """ Create an incremental checkpoint from the current rolling state. @@ -960,7 +943,6 @@ def create_incremental_checkpoint(self) -> IncrementalCheckpoint: Returns: The created IncrementalCheckpoint object. """ - self._raise_if_checkpoints_disabled() if self._rolling_state is None or self._rolling_state.event_count == 0: raise ValueError("No rolling state to create checkpoint from") @@ -1017,7 +999,6 @@ def create_checkpoint(self) -> Checkpoint: Returns: The created Checkpoint object. """ - self._raise_if_checkpoints_disabled() self._load_rolling_state() try: last_event_timestamp = self.event_cache.get_latest_event_timestamp() @@ -1087,7 +1068,6 @@ def compact_checkpoints(self) -> Checkpoint: Returns: The compacted Checkpoint object. """ - self._raise_if_checkpoints_disabled() print("Compacting incremental checkpoints...") # Download existing full checkpoint (if exists) @@ -1188,9 +1168,9 @@ def try_create_checkpoint( Returns: The created checkpoint (incremental or compacted), or None. Always - None when this client has checkpoints turned off. + None when this client keeps no owner state. """ - if not self.use_checkpoints: + if not self.persist_owner_state: return None self._load_rolling_state() diff --git a/tests/unit/test_checkpoints.py b/tests/unit/test_checkpoints.py index cc12585dc8f..abe1ae05a6c 100644 --- a/tests/unit/test_checkpoints.py +++ b/tests/unit/test_checkpoints.py @@ -10,8 +10,6 @@ from pathlib import Path from uuid import uuid4 -import pytest - from syft.sync.checkpoints.checkpoint import ( Checkpoint, IncrementalCheckpoint, @@ -607,15 +605,17 @@ def test_local_do_changes_end_up_in_incremental_checkpoint(): ) -def test_no_checkpoints_written_when_disabled(): - """With use_checkpoints=False, nothing checkpoint-shaped reaches the drive. +def test_no_owner_state_written_when_disabled(): + """With persist_owner_state=False, no owner-only state reaches the drive. - The second half is the point of the flag: turning checkpoints off must not - cost correctness, so the DS still has to see every DO change. + 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, - use_checkpoints=False, + persist_owner_state=False, ) do_manager.datasite_owner_syncer.perm_context.open(".").grant_write_access( ds_manager.email @@ -632,6 +632,8 @@ def test_no_checkpoints_written_when_disabled(): 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 = ( @@ -641,7 +643,11 @@ def test_no_checkpoints_written_when_disabled(): ) assert not rolling_state_path.exists() - # Sync still works: the DO holds every file and the DS receives them back. + # 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 @@ -652,11 +658,16 @@ def test_no_checkpoints_written_when_disabled(): assert f"test{i}.txt" in ds_paths -def test_disabled_checkpoints_restore_via_all_events(): - """A cold start with checkpoints off rebuilds state from the events instead.""" +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, - use_checkpoints=False, + persist_owner_state=False, ) do_manager.datasite_owner_syncer.perm_context.open(".").grant_write_access( ds_manager.email @@ -668,24 +679,48 @@ def test_disabled_checkpoints_restore_via_all_events(): do_manager.sync() assert len(do_manager.datasite_owner_syncer.event_cache.file_hashes) == 2 - # Simulate a fresh login: no checkpoint exists to restore from. + # 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 len(do_manager.datasite_owner_syncer.event_cache.file_hashes) == 2 + assert downloads == 0 + assert len(syncer.event_cache.file_hashes) == 0 -def test_explicit_checkpoint_calls_raise_when_disabled(): - """An explicit request fails loudly rather than silently doing nothing.""" - _, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( +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, - use_checkpoints=False, + ) + do_manager.datasite_owner_syncer.perm_context.open(".").grant_write_access( + ds_manager.email ) - with pytest.raises(ValueError, match="use_checkpoints=False"): - do_manager.create_checkpoint() + 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) - # The automatic path stays silent instead - it is not a caller error. - assert do_manager.try_create_checkpoint(threshold=1) is None + 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 5499d4f728d..1ef1dbce756 100644 --- a/tests/unit/test_rolling_state.py +++ b/tests/unit/test_rolling_state.py @@ -247,14 +247,14 @@ def test_rolling_state_clear_resets_base_timestamp(): assert rs.last_event_timestamp is None -def test_no_rolling_state_when_checkpoints_disabled(): - """Rolling state is part of the checkpoint machinery, so it goes off too. +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( - use_checkpoints=False, + persist_owner_state=False, ) do_manager.datasite_owner_syncer.perm_context.open(".").grant_write_access( ds_manager.email @@ -268,5 +268,5 @@ def test_no_rolling_state_when_checkpoints_disabled(): syncer = do_manager.datasite_owner_syncer assert not (syncer.syftbox_folder / CACHE_DIR / ROLLING_STATE_FILENAME).exists() - # The in-memory state is initialised but never accumulates events. - assert syncer._rolling_state.event_count == 0 + # Nothing is tracked in memory either - there is nothing to track it for. + assert syncer._rolling_state is None