diff --git a/dpsynth/api.py b/dpsynth/api.py index 146ed28..0b1c559 100644 --- a/dpsynth/api.py +++ b/dpsynth/api.py @@ -39,6 +39,7 @@ from typing import Any import dp_accounting +from etils import epath class CalibratedMechanism(abc.ABC): @@ -118,6 +119,11 @@ class MechanismConfig(abc.ABC): _registry: dict[str, type[MechanismConfig]] = {} + @property + def working_dir(self) -> epath.PathLike | None: + """Base directory path for checkpointing intermediate mechanism state.""" + return None + def __init_subclass__(cls, **kwargs: Any): super().__init_subclass__(**kwargs) MechanismConfig._registry[cls.__name__] = cls diff --git a/dpsynth/checkpoint.py b/dpsynth/checkpoint.py new file mode 100644 index 0000000..e82a0b5 --- /dev/null +++ b/dpsynth/checkpoint.py @@ -0,0 +1,93 @@ +# 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. + +"""Checkpointing utilities for long-running mechanism synthesis. + +Provides :class:`Checkpointer`, which serializes and deserializes intermediate +mechanism state (e.g. exact marginals, noisy measurements, graphical models) +using :mod:`mbi` pytree serialization on top of :mod:`etils.epath`. +""" + +from __future__ import annotations + +import dataclasses +import io +from typing import Any + +from etils import epath +import mbi + + +@dataclasses.dataclass(frozen=True) +class Checkpointer: + """Saves and restores intermediate mechanism state as .npz checkpoints. + + When ``working_dir`` is None (the default), all save/load operations are + no-ops, allowing callers to disable checkpointing without branching. + When ``working_dir`` is provided, intermediate mechanism state is persisted + directly under that directory as .npz files using ``mbi.save`` and + ``mbi.load``. + + Attributes: + working_dir: Base directory path for checkpoint files (supports local, + Cloud, and remote paths via epath.Path). If None, checkpointing is + disabled. + """ + + working_dir: epath.PathLike | None = None + + @property + def path(self) -> epath.Path | None: + """The resolved working directory path, or None if disabled.""" + return ( + epath.Path(self.working_dir) if self.working_dir is not None else None + ) + + def save(self, name: str, obj: Any) -> None: + """Saves an object to the working directory (no-op if disabled). + + Args: + name: Filename to write the object to (e.g. 'model.npz'). + obj: A JAX pytree to serialize (e.g. a CliqueVector, model, or list of + measurements). + """ + if self.path is None: + return + self.path.mkdir(parents=True, exist_ok=True) + buf = io.BytesIO() + mbi.save(obj, buf) + (self.path / name).write_bytes(buf.getvalue()) + + def load(self, name: str) -> Any | None: + """Loads an object from the working directory, or None if absent/disabled. + + Args: + name: Filename of the checkpointed object. + + Returns: + The deserialized object, or None if checkpointing is disabled or the + file does not exist. + """ + if self.path is None: + return None + target = self.path / name + if not target.exists(): + return None + return mbi.load(io.BytesIO(target.read_bytes())) + + def exists(self, name: str) -> bool: + """Returns True if the named checkpoint file exists.""" + if self.path is None: + return False + return (self.path / name).exists() diff --git a/dpsynth/discrete_mechanisms/swift.py b/dpsynth/discrete_mechanisms/swift.py index 7bcdf6a..639fcff 100644 --- a/dpsynth/discrete_mechanisms/swift.py +++ b/dpsynth/discrete_mechanisms/swift.py @@ -35,10 +35,12 @@ from absl import logging import dp_accounting from dpsynth import api +from dpsynth import checkpoint as checkpoint_lib from dpsynth.discrete_mechanisms import accounting from dpsynth.discrete_mechanisms import clique_tree from dpsynth.discrete_mechanisms import common from dpsynth.discrete_mechanisms import swift_utils +from etils import epath import mbi import networkx as nx import numpy as np @@ -59,6 +61,8 @@ class SWIFTConfig(api.MechanismConfig): pgm_iters: Number of mirror descent iterations for PGM estimation. select_budget_frac: Fraction of the total budget used for selecting which marginals to measure. + working_dir: Base directory path for intermediate checkpoints (e.g. exact + marginals, noisy measurements, model). If None, checkpointing is disabled. """ workload: Mapping[mbi.Clique, float] | Iterable[mbi.Clique] | None = None @@ -67,6 +71,7 @@ class SWIFTConfig(api.MechanismConfig): pgm_iters: int = 10_000 marginal_oracle: mbi.MarginalOracle | None = None select_budget_frac: float = 0.1 + working_dir: epath.PathLike | None = None def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: """Returns the workload cliques filtered by max_marginal_size.""" @@ -108,6 +113,29 @@ def __call__( ) -> common.DiscreteMechanismResult: common.validate_initial_measurements(initial_measurements) phase_times = {} + checkpointer = checkpoint_lib.Checkpointer(self.config.working_dir) + + # 1. Full resume: if model and measurements already exist, skip + # straight to synthesis. + if checkpointer.exists('model.npz') and checkpointer.exists( + 'measurements.npz' + ): + logging.info('[SWIFT] Resuming from checkpointed model and measurements.') + final_model = checkpointer.load('model.npz') + measurements = checkpointer.load('measurements.npz') + assert final_model is not None and measurements is not None + total_src = initial_measurements if initial_measurements else measurements + rows = mbi.estimation.minimum_variance_unbiased_total(total_src) # pyrefly: ignore[bad-argument-type] + rows = int(round(max(rows, 1))) + syn = mbi.extensions.synthetic_data(final_model, rows) # pyrefly: ignore[bad-argument-type] + diagnostics = common.clique_stats(final_model) + diagnostics.phase_times = phase_times + return common.DiscreteMechanismResult( + synthetic_data=syn, + measurements=measurements, + model=final_model, + diagnostics=diagnostics, + ) select_gdp_budget = self.gdp_budget * self.config.select_budget_frac measure_gdp_budget = self.gdp_budget - select_gdp_budget @@ -123,8 +151,15 @@ def __call__( ) logging.info('[SWIFT] %d candidates.', len(candidates)) - with common.timed(phase_times, 'from_projectable'): - answers = mbi.CliqueVector.from_projectable(data, candidates) # pyrefly: ignore[bad-argument-type] + # 2. Exact marginals resume: load if cached, otherwise compute and save. + if checkpointer.exists('marginals.npz'): + logging.info('[SWIFT] Resuming exact marginals from checkpoint.') + answers = checkpointer.load('marginals.npz') + assert answers is not None + else: + with common.timed(phase_times, 'from_projectable'): + answers = mbi.CliqueVector.from_projectable(data, candidates) # pyrefly: ignore[bad-argument-type] + checkpointer.save('marginals.npz', answers) with common.timed(phase_times, 'initial_mirror_descent'): estimator = mbi.estimation.MirrorDescent(self.config.marginal_oracle) @@ -193,6 +228,7 @@ def __call__( max_records_per_user=self.max_records_per_user, ) measurements = list(initial_measurements) + new_measurements + checkpointer.save('measurements.npz', measurements) logging.info('[SWIFT] Finished measurements.') ######################################################## @@ -210,6 +246,7 @@ def __call__( callback_fn=mbi.callbacks.default(measurements, data.domain), constraints=constraints, ) + checkpointer.save('model.npz', final_model) logging.info('[SWIFT] Estimated final model.') t0 = time.time() @@ -218,11 +255,13 @@ def __call__( syn = mbi.extensions.synthetic_data(final_model, rows) # pyrefly: ignore[bad-argument-type] logging.info('[SWIFT] Generated %d synthetic records.', rows) + diagnostics = common.clique_stats(final_model) + diagnostics.phase_times = phase_times return common.DiscreteMechanismResult( synthetic_data=syn, measurements=measurements, model=final_model, - diagnostics=common.clique_stats(final_model), + diagnostics=diagnostics, ) diff --git a/tests/checkpoint_test.py b/tests/checkpoint_test.py new file mode 100644 index 0000000..990cd5b --- /dev/null +++ b/tests/checkpoint_test.py @@ -0,0 +1,117 @@ +# 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. + +"""Unit tests for dpsynth.checkpoint.""" + +import pathlib +from absl.testing import absltest +from dpsynth import checkpoint as checkpoint_lib +from etils import epath +import jax.numpy as jnp +import mbi +import numpy as np + + +class CheckpointerTest(absltest.TestCase): + + def test_noop_when_disabled(self): + ckpt = checkpoint_lib.Checkpointer(working_dir=None) + self.assertIsNone(ckpt.path) + self.assertFalse(ckpt.exists('model.npz')) + self.assertIsNone(ckpt.load('model.npz')) + + # Saving should be a no-op and not raise. + domain = mbi.Domain.fromdict({'a': 2, 'b': 3}) + cliques = [('a',), ('b',)] + potentials = mbi.CliqueVector.zeros(domain, cliques) + ckpt.save('model.npz', potentials) + self.assertFalse(ckpt.exists('model.npz')) + self.assertIsNone(ckpt.load('model.npz')) + + def test_save_and_load_roundtrip(self): + working_dir = self.create_tempdir().full_path + ckpt = checkpoint_lib.Checkpointer(working_dir=working_dir) + + domain = mbi.Domain.fromdict({'a': 2, 'b': 3}) + cliques = [('a',), ('a', 'b')] + potentials = mbi.CliqueVector.zeros(domain, cliques) + potentials[('a',)] = jnp.array([1.0, 2.0]) + marginals = mbi.CliqueVector.zeros(domain, cliques) + mrf = mbi.MarkovRandomField( + potentials=potentials, marginals=marginals, total=10.0 + ) + + self.assertFalse(ckpt.exists('model.npz')) + ckpt.save('model.npz', mrf) + self.assertTrue(ckpt.exists('model.npz')) + + loaded = ckpt.load('model.npz') + self.assertIsInstance(loaded, mbi.MarkovRandomField) + np.testing.assert_allclose(loaded.potentials[('a',)], potentials[('a',)]) + self.assertEqual(loaded.total, 10.0) + + def test_save_and_load_linear_measurements(self): + working_dir = self.create_tempdir().full_path + ckpt = checkpoint_lib.Checkpointer(working_dir=working_dir) + + measurements = [ + mbi.LinearMeasurement(np.array([5.0, 10.0]), ('a',), stddev=1.0), + mbi.LinearMeasurement( + np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]), ('a', 'b'), stddev=0.5 + ), + ] + ckpt.save('measurements.npz', measurements) + self.assertTrue(ckpt.exists('measurements.npz')) + + loaded = ckpt.load('measurements.npz') + self.assertLen(loaded, 2) + self.assertEqual(loaded[0].clique, ('a',)) + self.assertEqual(loaded[0].stddev, 1.0) + np.testing.assert_allclose( + loaded[0].noisy_measurement, measurements[0].noisy_measurement + ) + self.assertEqual(loaded[1].clique, ('a', 'b')) + self.assertEqual(loaded[1].stddev, 0.5) + np.testing.assert_allclose( + loaded[1].noisy_measurement, measurements[1].noisy_measurement + ) + + def test_load_nonexistent_returns_none(self): + working_dir = self.create_tempdir().full_path + ckpt = checkpoint_lib.Checkpointer(working_dir=working_dir) + self.assertIsNone(ckpt.load('nonexistent.npz')) + + def test_accepts_different_path_types(self): + temp_dir = self.create_tempdir().full_path + + # str + ckpt_str = checkpoint_lib.Checkpointer(working_dir=temp_dir) + ckpt_str.save('test_str.npz', {'v': jnp.array([1, 2, 3])}) + self.assertTrue(ckpt_str.exists('test_str.npz')) + + # pathlib.Path + ckpt_pathlib = checkpoint_lib.Checkpointer( + working_dir=pathlib.Path(temp_dir) + ) + self.assertTrue(ckpt_pathlib.exists('test_str.npz')) + loaded = ckpt_pathlib.load('test_str.npz') + np.testing.assert_array_equal(loaded['v'], [1, 2, 3]) + + # epath.Path + ckpt_epath = checkpoint_lib.Checkpointer(working_dir=epath.Path(temp_dir)) + self.assertTrue(ckpt_epath.exists('test_str.npz')) + + +if __name__ == '__main__': + absltest.main() diff --git a/tests/discrete_mechanisms/swift_test.py b/tests/discrete_mechanisms/swift_test.py index b11afb4..27cf8e2 100644 --- a/tests/discrete_mechanisms/swift_test.py +++ b/tests/discrete_mechanisms/swift_test.py @@ -19,6 +19,7 @@ from dpsynth.discrete_mechanisms import common from dpsynth.discrete_mechanisms import swift from dpsynth.discrete_mechanisms import swift_utils +from etils import epath import mbi import networkx as nx import numpy as np @@ -146,6 +147,75 @@ def test_fits_one_way_marginals(self): actual = result.model.project([col]).datavector() np.testing.assert_allclose(actual, expected, atol=1) + def test_checkpointing_saves_and_resumes(self): + temp_dir = self.create_tempdir().full_path + data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [2, 3, 4]), N=500) + initial = [ + mbi.LinearMeasurement( + data.project((c,)).datavector(), (c,), stddev=0.01 + ) + for c in data.domain + ] + + # 1. Cold run: should save marginals, measurements, and model + config = swift.SWIFTConfig(pgm_iters=100, working_dir=temp_dir).configure( + zcdp_rho=1000 + ) + result1 = config( + np.random.default_rng(0), data, initial_measurements=initial + ) + + ckpt_marginals = epath.Path(temp_dir) / 'marginals.npz' + ckpt_measurements = epath.Path(temp_dir) / 'measurements.npz' + ckpt_model = epath.Path(temp_dir) / 'model.npz' + + self.assertTrue(ckpt_marginals.exists()) + self.assertTrue(ckpt_measurements.exists()) + self.assertTrue(ckpt_model.exists()) + + # 2. Resume run: should load model and measurements and skip to synthesis + result2 = config( + np.random.default_rng(1), data, initial_measurements=initial + ) + self.assertEqual(result2.synthetic_data.records, 500) + for cl in result1.model.potentials.cliques: + np.testing.assert_allclose( + result2.model.potentials[cl].values, + result1.model.potentials[cl].values, + ) + + def test_checkpointing_resumes_from_marginals(self): + temp_dir = self.create_tempdir().full_path + data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [2, 3, 4]), N=500) + initial = [ + mbi.LinearMeasurement( + data.project((c,)).datavector(), (c,), stddev=0.01 + ) + for c in data.domain + ] + + # First run up to marginals + config1 = swift.SWIFTConfig(pgm_iters=100, working_dir=temp_dir).configure( + zcdp_rho=1000 + ) + config1(np.random.default_rng(0), data, initial_measurements=initial) + + # Delete measurements and model, keep marginals + (epath.Path(temp_dir) / 'measurements.npz').unlink() + (epath.Path(temp_dir) / 'model.npz').unlink() + + # Second run: should reuse marginals and re-estimate + config2 = swift.SWIFTConfig(pgm_iters=100, working_dir=temp_dir).configure( + zcdp_rho=1000 + ) + result = config2( + np.random.default_rng(0), data, initial_measurements=initial + ) + + self.assertTrue((epath.Path(temp_dir) / 'measurements.npz').exists()) + self.assertTrue((epath.Path(temp_dir) / 'model.npz').exists()) + self.assertEqual(result.synthetic_data.records, 500) + if __name__ == '__main__': absltest.main()