Add checkpointing support to dpsynth using etils.epath and MBI pytree serialization:

dpsynth/__init__.py

@@ -16,6 +16,7 @@
 
 # pylint: disable=g-importing-member
 __version__ = '0.4.0'
+from absl import logging
 from dpsynth import api
 from dpsynth import constraints
 from dpsynth import discrete_mechanisms
@@ -35,6 +36,16 @@
 from dpsynth.domain import Schema
 from dpsynth.serialize import from_yaml
 from dpsynth.serialize import to_yaml
+import mbi
+
+# Route MBI callback logs (which use print-style formatting) to
+# absl.logging.info.
+if hasattr(mbi, 'callbacks') and hasattr(mbi.callbacks, 'set_log_fn'):
+  mbi.callbacks.set_log_fn(
+      lambda *args, sep=' ', **kwargs: logging.info(
+          sep.join(str(a) for a in args)
+      )
+  )
 
 ForeignKeyRelation = relational.ForeignKeyRelation
 MultiDataGenerationResult = relational.MultiDataGenerationResult

dpsynth/api.py

@@ -35,6 +35,7 @@
 
 import abc
 from collections.abc import Callable
+import dataclasses
 import functools
 from typing import Any
 
@@ -118,6 +119,24 @@ class MechanismConfig(abc.ABC):
 
   _registry: dict[str, type[MechanismConfig]] = {}
 
+  @property
+  def working_dir(self) -> str | None:
+    """Base directory path for checkpointing intermediate mechanism state."""
+    return None
+
+  def with_working_dir(self, working_dir: str | None) -> MechanismConfig:
+    """Returns a copy of the config with working_dir set if supported and unset."""
+    if self.working_dir is not None or working_dir is None:
+      return self
+    if dataclasses.is_dataclass(self):
+      try:
+        return dataclasses.replace(  # pyrefly: ignore[bad-specialization]
+            self, working_dir=working_dir
+        )
+      except (TypeError, ValueError):
+        return self
+    return self
+
   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/data_generation_v3.py

@@ -313,6 +313,7 @@ def __call__(
           m.compress(mappings, discrete.domain)  # pyrefly: ignore[bad-argument-type]
           for m in initial_measurements
       ]
+    logging.info('[DPSynth]: Compressed discrete domain:\n%s', discrete.domain)
 
     cfg = self.config.discrete_mechanism
     if hasattr(cfg, 'supporting_cliques'):
@@ -386,6 +387,8 @@ class TabularConfig(api.MechanismConfig):
       mbi.extensions.precompute_marginals) to compute marginals from Dataset.
     use_jax_for_generation: Whether to use JAX-accelerated generation (via
       mbi.extensions.synthetic_data) to generate synthetic data from the model.
+    working_dir: Base directory path for intermediate checkpoints (passed down
+      to the underlying discrete mechanism). If None, checkpointing is disabled.
   """
 
   domains: Mapping[str, domain.AttributeType] | None = None
@@ -396,6 +399,7 @@ class TabularConfig(api.MechanismConfig):
   compress_columns: bool = False
   use_jax_for_bincount: bool = False
   use_jax_for_generation: bool = False
+  working_dir: str | None = None
 
   def _compute_per_col_deltas(self, domains, delta):
     # Split delta across open-set columns, analogous to splitting zcdp_rho.
@@ -504,7 +508,11 @@ def configure(
         for col, init in inits.items()
     }
 
-    calibrated_discrete = self.discrete_mechanism.configure(
+    discrete_mechanism = self.discrete_mechanism.with_working_dir(
+        self.working_dir
+    )
+
+    calibrated_discrete = discrete_mechanism.configure(
         max_records_per_user=max_records_per_user,
         zcdp_rho=discrete_rho,
     )

dpsynth/discrete_mechanisms/aim.py

@@ -61,7 +61,7 @@ def _filter_candidates(
 def _worst_approximated(
     rng: np.random.Generator,
     candidates: Mapping[mbi.Clique, float],
-    answers: mbi.CliqueVector,
+    data: mbi.Dataset | mbi.CliqueVector,
     estimates: mbi.CliqueVector,
     eps: float,
     sigma: float,
@@ -72,7 +72,7 @@ def _worst_approximated(
   errors = {}
   for cl in candidates:
     wgt = candidates[cl]
-    diff = answers[cl].datavector() - estimates[cl].datavector()
+    diff = data.project(cl).datavector() - estimates[cl].datavector()
     bias = jnp.sqrt(2 / jnp.pi) * max_records_per_user * sigma * domain.size(cl)
     errors[cl] = wgt * (jnp.linalg.norm(diff, ord=1) - bias)
 
@@ -171,13 +171,11 @@ def __call__(
     rho_per_round = self.zcdp_rho / max_rounds
 
     #########################################################################
-    # Compile workload into candidate measurements, and precompute answers. #
+    # Compile workload into candidate measurements.                         #
     #########################################################################
     candidates = common.compiled_workload(
         data.domain, self.config.workload, self.config.max_marginal_size
     )
-    answers = mbi.CliqueVector.from_projectable(data, list(candidates))  # pyrefly: ignore[bad-argument-type]
-    logging.info('[AIM]: Calculated workload-query answers.')
 
     estimator = mbi.estimation.MirrorDescent(self.config.marginal_oracle)
     model = estimator.estimate(
@@ -215,7 +213,7 @@ def __call__(
         marginal_query = _worst_approximated(
             rng,
             small_candidates,
-            answers,
+            data,
             estimates,
             epsilon,
             sigma,

dpsynth/discrete_mechanisms/aim_gdp.py

@@ -64,24 +64,22 @@ def expected_size(cl):
 
 def _compute_dp_errors(
     rng: np.random.Generator,
-    answers: mbi.CliqueVector,
+    data: mbi.Dataset | mbi.CliqueVector,
     estimates: mbi.CliqueVector,
     gdp_budget: float,
-    subset: Iterable[mbi.Clique] | None = None,
+    subset: Iterable[mbi.Clique],
     max_records_per_user: int = 1,
 ) -> dict[mbi.Clique, float]:
   """Compute L1 error between the model answers and the true answers with DP."""
-  if subset is None:
-    subset = answers.cliques
-
+  clique_list = list(subset)
   # The L1 error of a marginal changes by at most ``max_records_per_user`` when
   # a single user (contributing up to that many records) is added or removed.
   per_candidate_sigma = max_records_per_user * accounting.gdp_gaussian_sigma(
-      gdp_budget / len(subset)  # pyrefly: ignore[bad-argument-type]
+      gdp_budget / len(clique_list)  # pyrefly: ignore[bad-argument-type]
   )
   result = {}
-  for cl in subset:
-    actual = answers[cl].datavector(flatten=True)
+  for cl in clique_list:
+    actual = data.project(cl).datavector(flatten=True)
     estimate = estimates[cl].datavector(flatten=True)
     error = jnp.linalg.norm(actual - estimate, ord=1)
     noise = rng.normal(loc=0, scale=per_candidate_sigma)
@@ -93,7 +91,7 @@ def _worst_approximated(
     rng: np.random.Generator,
     candidates: Mapping[mbi.Clique, float],
     errors: dict[mbi.Clique, float],  # will be updated in-place.
-    answers: mbi.CliqueVector,  # derived from sensitive data.
+    data: mbi.Dataset | mbi.CliqueVector,  # sensitive data.
     model: mbi.MarkovRandomField,
     select_budget: float,  # satisfies select_budget-GDP.
     measure_sigma: float,
@@ -118,10 +116,10 @@ def _worst_approximated(
   estimates = mbi.marginal_oracles.bulk_variable_elimination(
       model.potentials, subset, model.total  # pyrefly: ignore[bad-argument-type]
   )
-  # Only step that uses "answers", satisfies DP.
+  # Only step that uses "data", satisfies DP.
   current_errors = _compute_dp_errors(
       rng,
-      answers,
+      data,
       estimates,
       select_budget,
       subset,
@@ -240,13 +238,11 @@ def __call__(
     budget_per_round = budget_remaining / max_rounds
 
     #########################################################################
-    # Compile workload into candidate measurements, and precompute answers. #
+    # Compile workload into candidate measurements.                         #
     #########################################################################
     candidates = common.compiled_workload(
         data.domain, self.config.workload, self.config.max_marginal_size
     )
-    answers = mbi.CliqueVector.from_projectable(data, candidates)  # pyrefly: ignore[bad-argument-type]
-    logging.info('[AIM] Calculated workload-query answers.')
     domain = data.domain
 
     estimator = mbi.estimation.MirrorDescent(self.config.marginal_oracle)
@@ -297,7 +293,7 @@ def __call__(
             rng,
             candidates=small_candidates,
             errors=errors,
-            answers=answers,
+            data=data,
             model=model,
             select_budget=select_budget,
             measure_sigma=measure_sigma,

dpsynth/discrete_mechanisms/common.py

@@ -151,7 +151,9 @@ def precompute_marginals(
     *,
     use_jax: bool = False,
 ) -> mbi.CliqueVector:
-  """Computes marginals over cliques from a dataset, optionally using JAX."""
+  """Computes marginals over cliques from a Dataset, optionally using JAX."""
+  if not cliques:
+    return mbi.CliqueVector(data.domain, [], {})
   if use_jax:
     return mbi.extensions.precompute_marginals(
         data, cliques  # pyrefly: ignore[bad-argument-type]
@@ -231,8 +233,8 @@ def exponential_mechanism(
 
 def measure_marginals_with_noise(
     rng: np.random.Generator,
-    data: mbi.Projectable,
-    marginal_queries: list[tuple[str, ...]],
+    data: mbi.Dataset | mbi.CliqueVector,
+    marginal_queries: Sequence[mbi.Clique],
     gdp_sigma: float,
     weights: np.ndarray | None = None,
     max_records_per_user: int = 1,
@@ -418,7 +420,10 @@ def supporting_cliques(
     A list of cliques from the workload whose domain size is within the limit.
   """
   if workload is None:
-    cliques = list(itertools.combinations(domain.attributes, 3))
+    k = min(len(domain.attributes), 3)
+    cliques = (
+        list(itertools.combinations(domain.attributes, k)) if k > 0 else []
+    )
   elif isinstance(workload, Mapping):
     cliques = [tuple(cl) for cl in workload.keys()]
   else:
@@ -456,7 +461,10 @@ def compiled_workload(
   """
 
   if workload is None:
-    workload = list(itertools.combinations(domain.attributes, 3))
+    k = min(len(domain.attributes), 3)
+    workload = (
+        list(itertools.combinations(domain.attributes, k)) if k > 0 else []
+    )
 
   if not isinstance(workload, Mapping):
     workload = {tuple(cl): 1.0 for cl in workload}
@@ -477,7 +485,7 @@ def score(cl):
 
 
 def compute_independence_errors(
-    data: mbi.Projectable,
+    data: mbi.Dataset | mbi.CliqueVector,
     model: mbi.MarkovRandomField,
     cliques: Sequence[mbi.Clique],
 ) -> dict[mbi.Clique, float]: