Skip to content

Add checkpointing support to the SWIFT mechanism - #95

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
cl/952333775
Open

Add checkpointing support to the SWIFT mechanism#95
copybara-service[bot] wants to merge 1 commit into
mainfrom
cl/952333775

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Jul 22, 2026

Copy link
Copy Markdown

Add checkpointing support to the SWIFT mechanism

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.

@copybara-service
copybara-service Bot force-pushed the cl/952333775 branch 2 times, most recently from 4609771 to 585366d Compare July 29, 2026 01:04
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

@MuhammadNiazAli MuhammadNiazAli left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant