Add checkpointing support for long-running mechanism synthesis in dpsynth:

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

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()

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,
     )
 
 

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()