[copybara-service]

Keep the finalize thread from rewinding the preemption sync counter.

CheckpointManager._finalize_checkpoint runs on the async finalize thread, for
a step older than the training step. It called reached_preemption(step) only
to decide whether to record /jax/checkpoint/write/preempt/duration_saved_secs.
That call reaches JAX ReachedSyncPoint, which stores the step as the
process's preemption sync counter. When a preemption notice arrives, JAX picks
the highest counter plus one as the sync step. With the counter moved back, a
process can already be past the sync step and never save at it. In multi-process
JAX the other processes then wait for it at the save barrier. On Pathways the
old step can latch the preemption, so the training thread never sees it.

save() already knows whether a step is a preemption save. It now records the
step, and _finalize_checkpoint reads that record instead of calling
reached_preemption(). _finalize clears the record on every path, including
a cancelled save. The metric fires in the same cases, and each save makes one
fewer coordination call.

A new multi-process test fixes the order of events with key-value handshakes.
Without the fix the primary misses the sync step and the other processes save
one step early. A second test checks that the metric is recorded once, only on
the primary.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.