From d893febd2e0b2cb1a5e5d45b5ba3dc81c9bad7aa Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Wed, 22 Jul 2026 14:29:59 -0700 Subject: [PATCH] Add checkpointing support to the SWIFT mechanism Introduces `dpsynth.filesystem.FileSystem`, a small pluggable wrapper over filesystem operations (open/exists/makedirs/remove), and `dpsynth.checkpoint.Checkpointer`, which persists and restores intermediate mechanism state as `.npz` blobs via `mbi.save` / `mbi.load`. The checkpointer separates two storage roles so that resume-only state is never confused with inspectable output: - A `PrivateStore` holds resume-only state (exact marginals, noisy measurements, the estimated model). It is read back to resume a preempted run and is never intended for human inspection; it may contain sensitive intermediates and must be given the same protections as the input data. - A `PublicSink` holds DP-safe outputs intended to be inspected or consumed downstream. It is write-only (egress) and is never read back by the library. `LocalDirStore` and `LocalDirSink` implement the two roles on top of `FileSystem` (writing under `private/` and `public/` respectively); `Checkpointer.local(working_dir)` wires them up. The same two roles map cleanly onto a Trusted Execution Environment's release APIs (a private store onto the recovery-info channel, a public sink onto the unencrypted-release channel), so the abstraction doubles as a TEE hook. `SWIFTMechanism` gains an optional `checkpointer` field. When it has a private store, SWIFT persists resume-only state and resumes from the latest completed phase after a preemption: a fully completed run (model + noisy measurements) resumes straight to synthesis, and exact marginals are reused to skip recomputation. Sensitive intermediates are deleted once they are no longer needed. When the checkpointer has no stores (the default) every operation is a no-op, so existing behavior is unchanged. PiperOrigin-RevId: 952333775 --- dpsynth/checkpoint.py | 232 +++++++++++++++++++++++++++ dpsynth/discrete_mechanisms/swift.py | 49 +++++- dpsynth/filesystem.py | 74 +++++++++ tests/checkpoint_test.py | 143 +++++++++++++++++ tests/filesystem_test.py | 65 ++++++++ 5 files changed, 561 insertions(+), 2 deletions(-) create mode 100644 dpsynth/checkpoint.py create mode 100644 dpsynth/filesystem.py create mode 100644 tests/checkpoint_test.py create mode 100644 tests/filesystem_test.py diff --git a/dpsynth/checkpoint.py b/dpsynth/checkpoint.py new file mode 100644 index 00000000..bccdd1d7 --- /dev/null +++ b/dpsynth/checkpoint.py @@ -0,0 +1,232 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Checkpoint manager for saving and resuming intermediate mechanism state. + +Storage is separated into two roles so that resume-only state is never confused +with inspectable output: + +- A :class:`PrivateStore` holds resume-only state (e.g. exact marginals, noisy + measurements, the estimated model). It is read back to resume a preempted run + and is never intended for human inspection. It may contain sensitive + intermediates, so it must live somewhere with appropriate protections. +- A :class:`PublicSink` holds DP-safe outputs that are meant to be inspected or + consumed downstream. It is write-only (egress); the library never reads it + back. + +The default local implementations (:class:`LocalDirStore`, +:class:`LocalDirSink`) +are backed by the :mod:`dpsynth.filesystem` abstraction. The same two roles map +cleanly onto a Trusted Execution Environment's release APIs: a private store +onto the recovery-info channel, and a public sink onto the unencrypted-release +channel. +""" + +from __future__ import annotations + +from collections.abc import Sequence +import dataclasses +import io +import os +from typing import Any, Protocol, cast + +from dpsynth import filesystem +import mbi + + +class PrivateStore(Protocol): + """Resume-only, opaque checkpoint storage within the trust boundary. + + Holds internal state needed to resume a preempted run. Contents are never + meant for human inspection and may include sensitive intermediates (e.g. + exact marginals). In a TEE this maps to the recovery-info channel. + """ + + def put(self, name: str, data: bytes) -> None: + ... + + def get(self, name: str) -> bytes | None: + ... + + def delete(self, names: Sequence[str]) -> None: + ... + + +class PublicSink(Protocol): + """DP-safe, inspectable output storage (write-only egress). + + Holds outputs intended for downstream consumption or human inspection. In a + TEE this maps to the unencrypted-release channel, which cannot be read back. + """ + + def export(self, name: str, data: bytes) -> None: + ... + + +@dataclasses.dataclass(frozen=True) +class LocalDirStore: + """A :class:`PrivateStore` backed by a filesystem directory. + + Attributes: + root: Directory under which checkpoint files are written. + fs: Filesystem abstraction for all I/O. Defaults to the local filesystem. + """ + + root: str + fs: filesystem.FileSystem = dataclasses.field( + default_factory=filesystem.FileSystem + ) + + def put(self, name: str, data: bytes) -> None: + self.fs.makedirs(self.root) + with self.fs.open(os.path.join(self.root, name), 'wb') as f: + f.write(data) + + def get(self, name: str) -> bytes | None: + path = os.path.join(self.root, name) + if not self.fs.exists(path): + return None + with self.fs.open(path, 'rb') as f: + return cast(bytes, f.read()) + + def delete(self, names: Sequence[str]) -> None: + for name in names: + path = os.path.join(self.root, name) + if self.fs.exists(path): + self.fs.remove(path) + + +@dataclasses.dataclass(frozen=True) +class LocalDirSink: + """A :class:`PublicSink` backed by a filesystem directory. + + Attributes: + root: Directory under which exported files are written. + fs: Filesystem abstraction for all I/O. Defaults to the local filesystem. + """ + + root: str + fs: filesystem.FileSystem = dataclasses.field( + default_factory=filesystem.FileSystem + ) + + def export(self, name: str, data: bytes) -> None: + self.fs.makedirs(self.root) + with self.fs.open(os.path.join(self.root, name), 'wb') as f: + f.write(data) + + +@dataclasses.dataclass +class Checkpointer: + """Saves resume state privately and exports DP-safe outputs publicly. + + A ``Checkpointer`` separates two storage roles: + + - ``private``: a :class:`PrivateStore` for resume-only state (e.g. exact + marginals, noisy measurements, the estimated model). This data is used + solely to resume a preempted run and is never meant for human inspection. + - ``public``: a :class:`PublicSink` for DP-safe outputs that are intended to + be inspected or consumed downstream. + + Either role may be ``None`` to disable it. When ``private`` is ``None`` the + resume methods (:meth:`save`, :meth:`load`, :meth:`cleanup`) are no-ops and + :meth:`load` returns ``None``, so mechanisms behave exactly as if + checkpointing were disabled. When ``public`` is ``None`` :meth:`export` is a + no-op. + + Objects are serialized as ``.npz`` blobs via ``mbi.save`` / ``mbi.load``, + which round-trip arbitrary JAX pytrees (e.g. ``CliqueVector``, + ``MarkovRandomField``, and lists of ``LinearMeasurement``). + + Attributes: + private: Store for resume-only checkpoint state, or None to disable. + public: Sink for DP-safe inspectable outputs, or None to disable. + """ + + private: PrivateStore | None = None + public: PublicSink | None = None + + @classmethod + def local( + cls, + working_dir: str, + fs: filesystem.FileSystem | None = None, + ) -> Checkpointer: + """Returns a Checkpointer backed by ``private/`` and ``public/`` subdirs. + + Args: + working_dir: Root directory under which ``private/`` and ``public/`` + subdirectories are created lazily on first write. + fs: Filesystem abstraction for all I/O. Defaults to the local filesystem. + """ + fs = fs if fs is not None else filesystem.FileSystem() + return cls( + private=LocalDirStore(os.path.join(working_dir, 'private'), fs), + public=LocalDirSink(os.path.join(working_dir, 'public'), fs), + ) + + def load(self, name: str) -> Any | None: + """Loads resume state, returning None if absent or private is disabled. + + Args: + name: Filename of the checkpointed object. + + Returns: + The deserialized object, or None if there is no private store or the + object has not been checkpointed. + """ + if self.private is None: + return None + data = self.private.get(name) + if data is None: + return None + return mbi.load(io.BytesIO(data)) + + def save(self, name: str, obj: Any) -> None: + """Saves resume state to the private store (no-op if disabled). + + Args: + name: Filename to write the object to. + obj: A JAX pytree to serialize (e.g. a CliqueVector, model, or list of + measurements). + """ + if self.private is None: + return + buf = io.BytesIO() + mbi.save(obj, buf) + self.private.put(name, buf.getvalue()) + + def cleanup(self, names: Sequence[str]) -> None: + """Deletes named resume files from the private store, ignoring missing ones. + + Args: + names: Filenames to delete. Files that do not exist are skipped. + """ + if self.private is None: + return + self.private.delete(names) + + def export(self, name: str, obj: Any) -> None: + """Exports a DP-safe object to the public sink (no-op if disabled). + + Args: + name: Filename to export the object to. + obj: A JAX pytree to serialize. Callers are responsible for ensuring that + only DP-safe data is exported publicly. + """ + if self.public is None: + return + buf = io.BytesIO() + mbi.save(obj, buf) + self.public.export(name, buf.getvalue()) diff --git a/dpsynth/discrete_mechanisms/swift.py b/dpsynth/discrete_mechanisms/swift.py index 4b60a85a..05b53b3f 100644 --- a/dpsynth/discrete_mechanisms/swift.py +++ b/dpsynth/discrete_mechanisms/swift.py @@ -35,6 +35,7 @@ from absl import logging import dp_accounting +from dpsynth import checkpoint as checkpoint_lib from dpsynth.discrete_mechanisms import accounting from dpsynth.discrete_mechanisms import base from dpsynth.discrete_mechanisms import clique_tree @@ -60,6 +61,13 @@ class SWIFTMechanism(base.DiscreteMechanism): marginals to measure. one_way_budget_frac: Alias for one_way_budget_fraction. Kept for backward compatibility. + checkpointer: Checkpoint manager for saving and resuming intermediate state. + When it has a private store, SWIFT persists resume-only state (exact + marginals, noisy measurements, the estimated model) and resumes from the + latest completed phase on restart. All checkpointed state is resume-only + and written to the private store, which may contain sensitive + intermediates and must be given the same protections as the input data. By + default the checkpointer is disabled and SWIFT runs in a single pass. """ workload: Mapping[mbi.Clique, float] | Iterable[mbi.Clique] | None = None @@ -68,6 +76,9 @@ class SWIFTMechanism(base.DiscreteMechanism): pgm_iters: int = 25_000 select_budget_frac: float = 0.1 one_way_budget_fraction: float = 0.1 + checkpointer: checkpoint_lib.Checkpointer = dataclasses.field( + default_factory=checkpoint_lib.Checkpointer + ) # Internal state set by configure. _select_rho: float | None = dataclasses.field(default=None, repr=False) @@ -99,6 +110,24 @@ def _run(self, rng, data, measurements, constraints, phase_times): assert self._select_rho is not None assert self.measurement_rho is not None + ckpt = self.checkpointer + + # If a completed run was checkpointed (model + noisy measurements), resume + # straight to synthesis. Both are resume-only state in the private store. + # Domain decompression and result assembly are handled by the base + # __call__. + cached_model = ckpt.load('model.npz') + cached_measurements = ckpt.load('measurements.npz') + if cached_model is not None and cached_measurements is not None: + logging.info('[SWIFT] Resuming from model + measurements checkpoint.') + final_model = typing.cast(mbi.MarkovRandomField, cached_model) + measurements = cached_measurements + rows = int(mbi.estimation.minimum_variance_unbiased_total(measurements)) + syn = mbi.extensions.synthetic_data(final_model, rows) + logging.info('[SWIFT] Generated %d synthetic records.', rows) + ckpt.cleanup(['marginals.npz']) + return final_model, syn, measurements + # Budgets in GDP units, derived from the zCDP allocation set by configure. gdp_budget = accounting.zcdp_to_gdp(self._select_rho + self.measurement_rho) @@ -113,8 +142,15 @@ def _run(self, rng, data, measurements, constraints, phase_times): ) logging.info('[SWIFT] %d candidates.', len(candidates)) - with common.timed(phase_times, 'from_projectable'): - answers = mbi.CliqueVector.from_projectable(data, candidates) + # Exact marginals are sensitive; checkpoint them in the private store so a + # preempted run can skip the expensive recomputation on restart. + answers = ckpt.load('marginals.npz') + if answers is None: + with common.timed(phase_times, 'from_projectable'): + answers = mbi.CliqueVector.from_projectable(data, candidates) + ckpt.save('marginals.npz', answers) + else: + logging.info('[SWIFT] Loaded %d cached marginals.', len(answers.cliques)) domain = data.domain with common.timed(phase_times, 'initial_mirror_descent'): @@ -174,6 +210,11 @@ def _run(self, rng, data, measurements, constraints, phase_times): measurements.extend(new_measurements) logging.info('[SWIFT] Finished measurements.') + # Checkpoint the noisy measurements for resume. The exact marginals are no + # longer needed once the measurements are saved. + ckpt.save('measurements.npz', measurements) + ckpt.cleanup(['marginals.npz']) + ######################################################## # Estimate the model using all measurements # ######################################################## @@ -197,6 +238,10 @@ def _run(self, rng, data, measurements, constraints, phase_times): assert isinstance(final_model, mbi.MarkovRandomField) logging.info('[SWIFT] Estimated final model.') + # Checkpoint the estimated model for resume so synthesis can restart + # without re-estimating. + ckpt.save('model.npz', final_model) + if synth_future is not None: t0 = time.time() try: diff --git a/dpsynth/filesystem.py b/dpsynth/filesystem.py new file mode 100644 index 00000000..4a972369 --- /dev/null +++ b/dpsynth/filesystem.py @@ -0,0 +1,74 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Filesystem abstraction for checkpoint I/O. + +Provides a simple ``FileSystem`` dataclass that wraps a handful of filesystem +operations (open, exists, makedirs, remove) so that library code can checkpoint +intermediate state without hardcoding any specific storage backend. + +Local filesystem behavior is used by default. To read and write on a remote +or networked filesystem, construct a ``FileSystem`` with callables from the +appropriate storage client:: + + fs = FileSystem( + open=client.open, + exists=client.exists, + makedirs=client.makedirs, + remove=client.remove, + ) +""" + +from __future__ import annotations + +import builtins +from collections.abc import Callable +import dataclasses +import os +from typing import Any, IO + + +@dataclasses.dataclass(frozen=True) +class FileSystem: + """Pluggable filesystem for checkpoint I/O. + + Wraps the operations that the library needs for checkpointing: + + - ``open(path, mode) -> file``: Open a file for reading or writing. + - ``exists(path) -> bool``: Check whether a path exists. + - ``makedirs(path) -> None``: Create a directory (and parents). + - ``remove(path) -> None``: Delete a file. + + The defaults use Python's built-in ``open`` and the ``os`` module, so + local-filesystem checkpointing works out of the box with no extra arguments. + To use a remote or networked filesystem, construct a ``FileSystem`` with + the appropriate callables. + + The library does **not** handle encryption or access control. When saving + sensitive data (e.g., exact marginals), it is the caller's responsibility to + ensure that the directory has appropriate protections. + + Attributes: + open: Callable that opens a file given (path, mode). + exists: Callable that checks whether a path exists. + makedirs: Callable that creates a directory and its parents. + remove: Callable that deletes a file. + """ + + open: Callable[..., IO[Any]] = builtins.open + exists: Callable[[str], bool] = os.path.exists + makedirs: Callable[..., None] = dataclasses.field( + default_factory=lambda: lambda path: os.makedirs(path, exist_ok=True) + ) + remove: Callable[[str], None] = os.remove diff --git a/tests/checkpoint_test.py b/tests/checkpoint_test.py new file mode 100644 index 00000000..6921e882 --- /dev/null +++ b/tests/checkpoint_test.py @@ -0,0 +1,143 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import tempfile + +from absl.testing import absltest +from dpsynth import checkpoint as checkpoint_lib +from dpsynth import filesystem +import jax.numpy as jnp +import numpy as np + + +class CheckpointerTest(absltest.TestCase): + + def test_disabled_by_default(self): + """A Checkpointer with no stores is a no-op for every method.""" + ckpt = checkpoint_lib.Checkpointer() + ckpt.save('x.npz', {'a': jnp.arange(3)}) # No-op. + self.assertIsNone(ckpt.load('x.npz')) + ckpt.export('x.npz', {'a': jnp.arange(3)}) # No-op. + ckpt.cleanup(['x.npz']) # Should not raise. + + def test_save_load_roundtrip(self): + obj = {'a': jnp.arange(3), 'b': jnp.ones((2, 2))} + with tempfile.TemporaryDirectory() as tmpdir: + ckpt = checkpoint_lib.Checkpointer.local(tmpdir) + self.assertIsNone(ckpt.load('x.npz')) + ckpt.save('x.npz', obj) + loaded = ckpt.load('x.npz') + np.testing.assert_array_equal(loaded['a'], obj['a']) + np.testing.assert_array_equal(loaded['b'], obj['b']) + + def test_local_uses_private_and_public_subdirs(self): + """Resume state lands under private/ and exports land under public/.""" + with tempfile.TemporaryDirectory() as tmpdir: + ckpt = checkpoint_lib.Checkpointer.local(tmpdir) + ckpt.save('model.npz', {'a': jnp.arange(2)}) + ckpt.export('synthetic.npz', {'b': jnp.arange(2)}) + self.assertTrue( + os.path.isfile(os.path.join(tmpdir, 'private', 'model.npz')) + ) + self.assertTrue( + os.path.isfile(os.path.join(tmpdir, 'public', 'synthetic.npz')) + ) + + def test_public_exports_are_not_visible_to_load(self): + """load reads only resume state; it never reads back public exports.""" + with tempfile.TemporaryDirectory() as tmpdir: + ckpt = checkpoint_lib.Checkpointer.local(tmpdir) + ckpt.export('synthetic.npz', {'b': jnp.arange(2)}) + self.assertIsNone(ckpt.load('synthetic.npz')) + + def test_export_disabled_when_public_is_none(self): + with tempfile.TemporaryDirectory() as tmpdir: + ckpt = checkpoint_lib.Checkpointer( + private=checkpoint_lib.LocalDirStore(os.path.join(tmpdir, 'private')) + ) + ckpt.export('x.npz', {'a': jnp.arange(1)}) # No-op, no public sink. + self.assertFalse(os.path.isdir(os.path.join(tmpdir, 'public'))) + + def test_cleanup_removes_files_and_tolerates_missing(self): + with tempfile.TemporaryDirectory() as tmpdir: + ckpt = checkpoint_lib.Checkpointer.local(tmpdir) + ckpt.save('marginals.npz', {'a': jnp.arange(2)}) + self.assertIsNotNone(ckpt.load('marginals.npz')) + # Present and missing files can be requested together. + ckpt.cleanup(['marginals.npz', 'missing.npz']) + self.assertIsNone(ckpt.load('marginals.npz')) + + def test_routes_io_through_filesystem(self): + """The local stores perform all I/O through the provided FileSystem.""" + events = [] + real = filesystem.FileSystem() + + def recording_open(path, mode): + events.append(('open', mode)) + return real.open(path, mode) + + with tempfile.TemporaryDirectory() as tmpdir: + fs = filesystem.FileSystem( + open=recording_open, + exists=real.exists, + makedirs=real.makedirs, + remove=real.remove, + ) + ckpt = checkpoint_lib.Checkpointer.local(tmpdir, fs=fs) + ckpt.save('x.npz', {'a': jnp.arange(1)}) + self.assertIsNotNone(ckpt.load('x.npz')) + self.assertIn(('open', 'wb'), events) + self.assertIn(('open', 'rb'), events) + + def test_accepts_custom_stores(self): + """Custom in-memory stores can back the two roles (e.g. a TEE adapter).""" + + class MemStore: + """A minimal PrivateStore backed by a dict.""" + + def __init__(self): + self.blobs = {} + + def put(self, name, data): + self.blobs[name] = data + + def get(self, name): + return self.blobs.get(name) + + def delete(self, names): + for name in names: + self.blobs.pop(name, None) + + class MemSink: + """A minimal PublicSink backed by a dict.""" + + def __init__(self): + self.released = {} + + def export(self, name, data): + self.released[name] = data + + store, sink = MemStore(), MemSink() + ckpt = checkpoint_lib.Checkpointer(private=store, public=sink) + ckpt.save('model.npz', {'a': jnp.arange(3)}) + ckpt.export('synthetic.npz', {'b': jnp.arange(3)}) + self.assertIn('model.npz', store.blobs) + self.assertIn('synthetic.npz', sink.released) + loaded = ckpt.load('model.npz') + np.testing.assert_array_equal(loaded['a'], jnp.arange(3)) + + +if __name__ == '__main__': + absltest.main() diff --git a/tests/filesystem_test.py b/tests/filesystem_test.py new file mode 100644 index 00000000..d613080d --- /dev/null +++ b/tests/filesystem_test.py @@ -0,0 +1,65 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import tempfile + +from absl.testing import absltest +from dpsynth import filesystem + + +class FileSystemTest(absltest.TestCase): + + def test_default_local_roundtrip(self): + """Default FileSystem reads, writes, and removes on local disk.""" + fs = filesystem.FileSystem() + with tempfile.TemporaryDirectory() as tmpdir: + subdir = os.path.join(tmpdir, 'a', 'b') + fs.makedirs(subdir) + self.assertTrue(fs.exists(subdir)) + + path = os.path.join(subdir, 'test.bin') + self.assertFalse(fs.exists(path)) + + with fs.open(path, 'wb') as f: + f.write(b'hello') + self.assertTrue(fs.exists(path)) + + with fs.open(path, 'rb') as f: + self.assertEqual(f.read(), b'hello') + + fs.remove(path) + self.assertFalse(fs.exists(path)) + + def test_custom_callables(self): + """FileSystem dispatches to the provided callables.""" + calls = [] + fs = filesystem.FileSystem( + open=lambda path, mode: calls.append(('open', path, mode)), + exists=lambda path: bool(calls.append(('exists', path))), + makedirs=lambda path: calls.append(('makedirs', path)), + remove=lambda path: calls.append(('remove', path)), + ) + fs.makedirs('/fake/dir') + fs.exists('/fake/path') + fs.open('/fake/file', 'rb') + fs.remove('/fake/file') + + self.assertEqual( + [c[0] for c in calls], ['makedirs', 'exists', 'open', 'remove'] + ) + + +if __name__ == '__main__': + absltest.main()