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

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/data_generation_v3.py

@@ -30,6 +30,7 @@
 from dpsynth.local_mode import initialization
 from dpsynth.local_mode import primitives
 from dpsynth.local_mode import vectorized_transformations as vtx
+from etils import epath
 import mbi
 import numpy as np
 import pandas as pd
@@ -368,6 +369,8 @@ class TabularConfig(api.MechanismConfig):
     cross_attribute_constraints: Constraints to enforce on generated data.
     compress_columns: Whether to compress rare categories (< 3*sigma) for
       CategoricalAttribute columns not present in constraints.
+    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
@@ -376,6 +379,7 @@ class TabularConfig(api.MechanismConfig):
   init_budget_fraction: float = 0.1
   cross_attribute_constraints: Sequence[constraints.Constraint] = ()
   compress_columns: bool = False
+  working_dir: epath.PathLike | None = None
 
   def _compute_per_col_deltas(self, domains, delta):
     # Split delta across open-set columns, analogous to splitting zcdp_rho.
@@ -484,7 +488,21 @@ def configure(
         for col, init in inits.items()
     }
 
-    calibrated_discrete = self.discrete_mechanism.configure(
+    discrete_mechanism = self.discrete_mechanism
+    if (
+        self.working_dir is not None
+        and dataclasses.is_dataclass(discrete_mechanism)
+        and any(
+            f.name == 'working_dir'
+            for f in dataclasses.fields(discrete_mechanism)
+        )
+        and discrete_mechanism.working_dir is None
+    ):
+      discrete_mechanism = dataclasses.replace(  # pyrefly: ignore[bad-specialization]
+          discrete_mechanism, 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/common.py

@@ -384,7 +384,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 = list(workload.keys())
   else:
@@ -422,7 +425,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 = {cl: 1.0 for cl in workload}

dpsynth/discrete_mechanisms/discrete.py

@@ -26,11 +26,14 @@
 from collections.abc import Sequence
 import dataclasses
 
+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 common
 from dpsynth.discrete_mechanisms import mst
+from etils import epath
 import mbi
 import numpy as np
 
@@ -45,20 +48,36 @@ class DiscreteConfig(api.MechanismConfig):
     one_way_budget_fraction: Fraction of zCDP budget for one-way marginals.
     constraints: Default MBI constraints to enforce. Can be overridden at call
       time via the ``constraints`` kwarg on ``DiscreteMechanism.__call__``.
+    working_dir: Base directory path for intermediate checkpoints. If None,
+      checkpointing is disabled.
   """
 
   mechanism: api.MechanismConfig = mst.MSTConfig()
   compress_columns: bool | Sequence[str] = False
   one_way_budget_fraction: float = 0.1
   constraints: Sequence[mbi.Constraint] = ()
+  working_dir: epath.PathLike | None = None
 
   def configure(self, _=None, *, zcdp_rho, delta=0, max_records_per_user=1):
     """Configures the synthesizer with a zCDP budget."""
     api.validate_max_records_per_user(max_records_per_user)
 
+    inner_mechanism = self.mechanism
+    if (
+        self.working_dir is not None
+        and dataclasses.is_dataclass(inner_mechanism)
+        and any(
+            f.name == 'working_dir' for f in dataclasses.fields(inner_mechanism)
+        )
+        and inner_mechanism.working_dir is None
+    ):
+      inner_mechanism = dataclasses.replace(  # pyrefly: ignore[bad-specialization]
+          inner_mechanism, working_dir=self.working_dir
+      )
+
     one_way_rho = zcdp_rho * self.one_way_budget_fraction
-    remaining_rho = zcdp_rho * (1 - self.one_way_budget_fraction)
-    inner = self.mechanism.configure(
+    remaining_rho = zcdp_rho - one_way_rho
+    inner = inner_mechanism.configure(
         zcdp_rho=remaining_rho,
         delta=delta,
         max_records_per_user=max_records_per_user,
@@ -127,20 +146,30 @@ def __call__(
     if constraints is None:
       constraints = self.config.constraints
 
+    checkpointer = checkpoint_lib.Checkpointer(self.config.working_dir)
+
     if initial_measurements is not None:
       measurements = list(initial_measurements)
     elif self.one_way_gdp_budget > 0:
-      one_way_cliques = [(a,) for a in data.domain]
-      if hasattr(data, 'cliques'):
-        supported = common.downward_closure(data.cliques)
-        one_way_cliques = [cl for cl in one_way_cliques if cl in supported]
-      measurements = common.measure_marginals_with_noise(
-          rng=rng,
-          data=data,  # pyrefly: ignore[bad-argument-type]
-          marginal_queries=one_way_cliques,  # pyrefly: ignore[bad-argument-type]
-          gdp_sigma=accounting.gdp_gaussian_sigma(self.one_way_gdp_budget),
-          max_records_per_user=self.max_records_per_user,
-      )
+      if checkpointer.exists('one_way_measurements.npz'):
+        logging.info(
+            '[DiscreteMechanism] Resuming one-way measurements from checkpoint.'
+        )
+        measurements = checkpointer.load('one_way_measurements.npz')
+        assert measurements is not None
+      else:
+        one_way_cliques = [(a,) for a in data.domain]
+        if hasattr(data, 'cliques'):
+          supported = common.downward_closure(data.cliques)
+          one_way_cliques = [cl for cl in one_way_cliques if cl in supported]
+        measurements = common.measure_marginals_with_noise(
+            rng=rng,
+            data=data,  # pyrefly: ignore[bad-argument-type]
+            marginal_queries=one_way_cliques,  # pyrefly: ignore[bad-argument-type]
+            gdp_sigma=accounting.gdp_gaussian_sigma(self.one_way_gdp_budget),
+            max_records_per_user=self.max_records_per_user,
+        )
+        checkpointer.save('one_way_measurements.npz', measurements)
     else:
       measurements = []