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
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.

0 participants