swift_megatron, syncer: resume / rejoin correctness and graceful syncer stop (3/3)

docs/PROTOCOL.md

@@ -593,6 +593,18 @@ coordinator always rewrites `--checkpoint-path` at the final quiescent cut
 before terminal delivery. Thus a total-step count not divisible by the
 periodic interval still leaves an atomic final checkpoint.
 
+SIGTERM or SIGINT stops the syncer gracefully: it launches no new round,
+gives the rounds still gathering `--shutdown-timeout-s` (default 600) to
+close (a round with no delta whose islands are all disconnected is abandoned
+at once), then abandons the rest (they have not touched the state, so a resumed
+run launches those steps again), commits every computed round, writes
+`--checkpoint-path` synchronously (after any background checkpoint) and exits
+0 with `stop: checkpoint written step=<t>`. The checkpoint's global step is
+then the last `outer step` on the event tape, and `--resume` continues the
+tape at `t + 1`. A signal before training starts exits at once (nothing was
+merged), and a second signal exits immediately without a checkpoint. The
+terminal cut is not marked final.
+
 `--checkpoint-at-steps T1,T2,...` (launcher: `--rl-sync-checkpoint-at-steps`,
 asynchronous full-parameter rounds) also writes the cut at each listed global
 step `t` (one fragment merge, the `step=` of the `outer step` log, the unit of

syncer/src/main.rs

@@ -19,6 +19,7 @@ mod protocol;
 mod sample_relay;
 mod server;
 mod sha256;
+mod shutdown;
 mod state;
 mod subscriber;
 
@@ -158,6 +159,13 @@ struct Args {
     /// during the write) host memory peaks at about twice params + momentum.
     #[arg(long)]
     background_checkpoint: bool,
+    /// Write --checkpoint-path once before the first round (after INIT or
+    /// --resume), so the rounds before the first --checkpoint-every one are
+    /// covered too. Skipped when --resume found that checkpoint (it is the
+    /// state just loaded). On a background thread with
+    /// --background-checkpoint.
+    #[arg(long)]
+    checkpoint_on_start: bool,
     /// Policy sweep: write per-round checkpoints off the scheduler; each
     /// round is broadcast only after its checkpoint is durable.
     #[arg(long, default_value_t = false)]
@@ -167,6 +175,13 @@ struct Args {
     /// sessions are refused.
     #[arg(long)]
     bcast_delta_full_every: Option<u64>,
+    /// On SIGTERM / SIGINT: stop launching rounds, give the rounds still
+    /// gathering this many seconds to close (then drop them; they have not
+    /// touched the state), commit every computed round, write
+    /// --checkpoint-path synchronously and exit 0. A second signal exits at
+    /// once without a checkpoint.
+    #[arg(long, default_value_t = 600)]
+    shutdown_timeout_s: u64,
     /// Resume from --checkpoint-path if it exists.
     #[arg(long, default_value_t = false)]
     resume: bool,
@@ -505,6 +520,7 @@ fn main() -> anyhow::Result<()> {
         checkpoint_every: args.checkpoint_every,
         checkpoint_at_steps: args.checkpoint_at_steps,
         background_checkpoint: args.background_checkpoint,
+        checkpoint_on_start: args.checkpoint_on_start,
         async_checkpoint: args.async_checkpoint,
         bcast_delta_full_every: args.bcast_delta_full_every,
         resume: args.resume,
@@ -518,6 +534,9 @@ fn main() -> anyhow::Result<()> {
         island_hmac_key,
         expected_island_contract,
         syncer_epoch: args.syncer_epoch,
+        shutdown: Some(std::sync::Arc::new(shutdown::Shutdown::new(std::time::Duration::from_secs(
+            args.shutdown_timeout_s,
+        )))),
         membership,
         recovery,
         inline_resend: args.inline_resend,