Add checkpointing support to the SWIFT mechanism - #95
Add checkpointing support to the SWIFT mechanism#95copybara-service[bot] wants to merge 1 commit into
Conversation
4609771 to
585366d
Compare
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
585366d to
d893feb
Compare
MuhammadNiazAli
left a comment
There was a problem hiding this comment.
Good to see the docstring explicitly calls out the privacy implication "checkpointed state is resume-only... must be given the same protections as the input data." That's an important note since intermediate marginals/measurements are DP-relevant artifacts, not just performance cache; easy to overlook when adding checkpointing to a DP mechanism.
The full-resume path (cached_model + cached_measurements both present → skip straight to synthesis) mirrors the same pattern from PR #166's swift.py checkpoint version. The typing.cast(mbi.MarkovRandomField, cached_model) after the is not None check is reasonable — ckpt.load() presumably returns an untyped/generic object, so this narrows it for the type checker without a runtime check.
One question: assert self._select_rho is not None / assert self.measurement_rho is not None run before the checkpoint resume block. If those are only computed via the normal calibration path, does resuming from a checkpoint on a fresh run (e.g. same config, different process) still guarantee _select_rho/measurement_rho are set at this point, or could a full-resume hit those asserts before ever reaching calibration?
Also this PR (#95) and PR #166 both appear to add checkpointing to swift.py with different APIs (self.checkpointer here vs checkpoint_lib.Checkpointer(self.config.working_dir) there). Worth checking these two PRs aren't landing conflicting/duplicate implementations is one superseding the other?
Still only 3 of 5 files viewed here (haven't seen checkpoint.py or filesystem.py, which are probably where the actual Checkpointer class and its "private store" semantics are defined) holding off on approval until those are in.
Add checkpointing support to the SWIFT mechanism
Introduces
dpsynth.filesystem.FileSystem, a small pluggable wrapper overfilesystem operations (open/exists/makedirs/remove), and
dpsynth.checkpoint.Checkpointer, which persists and restores intermediatemechanism state as
.npzblobs viambi.save/mbi.load.The checkpointer separates two storage roles so that resume-only state is never
confused with inspectable output:
PrivateStoreholds resume-only state (exact marginals, noisymeasurements, 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.
PublicSinkholds DP-safe outputs intended to be inspected or consumeddownstream. It is write-only (egress) and is never read back by the library.
LocalDirStoreandLocalDirSinkimplement the two roles on top ofFileSystem(writing underprivate/andpublic/respectively);Checkpointer.local(working_dir)wires them up. The same two roles map cleanlyonto 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.
SWIFTMechanismgains an optionalcheckpointerfield. When it has a privatestore, 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.