Add on-demand full-state checkpointing for OpenShift AI / KubeFlow preemption (#686)

* Add on-demand full-state checkpointing for OpenShift AI / KubeFlow preemption

Implements signal-driven checkpoint-and-exit for distributed training jobs
running in OpenShift AI as KubeFlow training jobs or multi-node bare metal.

When `on_demand_checkpointing=True` is set in TrainingArgs:

- Parent process (run_training) installs handlers for SIGTERM, SIGINT,
  SIGUSR1, SIGUSR2, SIGXCPU, and SIGHUP — covering all signals
  Kubernetes/OpenShift sends before the hard SIGKILL.
- On signal receipt, a trigger file is atomically written to /dev/shm
  (tmpfs, shared within the pod, zero disk I/O).
- Worker processes check for the trigger file after each optimizer step
  via an all_reduce(MAX) collective, ensuring global consensus across
  all ranks on all nodes.
- When any rank detects the trigger, all ranks collectively save a
  full-state distributed checkpoint (model + optimizer + LR scheduler)
  then exit gracefully.
- Parent waits up to 300s for workers to complete the checkpoint before
  proceeding with normal shutdown.

https://claude.ai/code/session_01HSxsk7SnMULJxy7uafe7t3

* Address review feedback for on-demand checkpointing

- Fix mypy error: properly type _original_handlers dict with _SignalHandler
  type alias instead of bare object
- Fix ruff/isort: remove duplicate comment, fix import ordering
- Namespace trigger file with rdzv_id as job_id so concurrent jobs sharing
  /dev/shm don't interfere with each other
- Recompute subprocess failure status after forced termination to avoid
  stale exit code
- Gate consensus log message to rank 0 to reduce log noise on large jobs

* Check for on-demand checkpoint after each minibatch backward

Move the checkpoint request check from after the full optimizer step to
after each minibatch's backward pass inside BatchLossManager.process_batch.
This ensures the system responds within one fwd+bwd cycle (~1-2s) even
when gradient accumulation spans many minibatches, giving more time to
save before Kubernetes sends SIGKILL after the grace period.

The check is passed as an optional interrupt_check callback to keep
checkpoint-specific logic out of BatchLossManager. When triggered, the
batch loop breaks early and the training loop saves the checkpoint
immediately, skipping the optimizer step to preserve the pre-step model
state for exact resumption.

* Add diagnostic note when on-demand checkpoint fails to save in time

When the training subprocess fails after an on-demand checkpoint signal
was received, the error message now includes guidance to increase
terminationGracePeriodSeconds or reduce fwd/bwd pass time so the
checkpoint check fires before SIGKILL arrives.

* Fix help text: checkpoint check happens after each minibatch backward

Update --on_demand_checkpointing help text and TrainingArgs description
to accurately state that workers check for the trigger file after each
minibatch backward pass, not after each training step.

* Add checkpoint checks at 5 synchronization points per training step

Expand on-demand checkpointing to check for a trigger at five points:
1. Before each minibatch forward pass
2. Before each minibatch backward pass
3. After each minibatch backward pass (existing)
4. Before the optimizer step
5. After the optimizer step

This minimizes the latency between a termination signal arriving and
the checkpoint being saved, which is critical when the SIGKILL grace
period is short (e.g. 30s on OpenShift/Kubernetes).

Also cleans up the save-and-exit logic in train() by extracting a
_save_and_exit() helper to eliminate three nearly identical blocks,
and fixes _compute_average_loss to handle the case where the
minibatch loop is interrupted before any forward pass completes.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* test: add unit tests for on-demand checkpointing

24 tests covering:
- Trigger file helpers (path generation, write, exists, remove)
- ParentSignalHandler (install, handle, idempotency, uninstall, real signal)
- check_checkpoint_requested (trigger detection, cleanup, all_reduce consensus)
- BatchLossManager interrupt handling (all 3 check points, early exit, float loss)

* fix: correct help text for on_demand_checkpointing sync point count

* Add stale trigger cleanup and exact mid-epoch resume

Two fixes to the on-demand checkpointing feature:

1. Stale trigger file cleanup: ParentSignalHandler.install() now checks
   for and removes any existing trigger file before installing signal
   handlers. If the file exists before handlers are installed, it's from
   a previous run that was killed before workers could clean it up.
   Prevents a new training job from immediately checkpointing and exiting.

2. Exact mid-epoch resume: save_on_demand_checkpoint() now persists
   global_step in the checkpoint metadata alongside current_epoch and
   samples_seen. On resume, load_latest_full_state() detects the
   global_step field and sets last_step accordingly, so the training
   loop fast-forwards to the exact step within the epoch. Without this,
   mid-epoch checkpoints would skip to the next epoch on resume, losing
   remaining steps.

Tested with Qwen2-1.5B-Instruct on 2 GPUs: interrupted at step 19/25,
checkpoint saved with global_step=19, resumed and completed steps 20-25.

* Fix ruff formatting in utils.py

* Add documentation for on-demand checkpointing feature

* Document manual trigger file creation for on-demand checkpointing

* Improve manual trigger docs: clarify job ID requirement

* Simplify trigger file: remove job_id namespacing

Drop the job_id suffix from the trigger file path. The file is now
always /dev/shm/instructlab_checkpoint_requested with no suffix.

The namespacing was defensive against concurrent jobs sharing /dev/shm,
but in practice Kubernetes pods each get their own /dev/shm. This makes
manual triggering trivial: touch /dev/shm/instructlab_checkpoint_requested

* Fix pylint: remove unnecessary lambda wrapper

* fix: update unit tests for simplified trigger file API (no job_id)

* fix: add compat shim for Nemotron's is_flash_attn_greater_or_equal_2_10

The function was renamed to is_flash_attn_greater_or_equal in
transformers 5.x, but Nemotron's HF Hub remote code still imports
the old name. Inject a compatibility wrapper before model loading.

* fix: make unit tests work without CUDA

- Patch torch.tensor in check_checkpoint_requested tests to avoid
  CUDA device creation (CI runs without GPUs)
- Override BatchLossManager.torch_device to CPU in fixture
  so _prepare_model_inputs doesn't trigger CUDA init

---------

Co-authored-by: Claude <noreply@anthropic.com>

Author: RobotSail

docs/on_demand_checkpointing.md

@@ -0,0 +1,167 @@
+# On-Demand Checkpointing
+
+On-demand checkpointing enables graceful checkpoint-and-exit when termination
+signals are received during training. It is designed for environments like
+OpenShift AI and KubeFlow where training jobs can be preempted at any time.
+
+## How It Works
+
+When enabled, the system installs signal handlers in the parent (launcher)
+process that catch termination signals before the hard SIGKILL. When a signal
+arrives:
+
+1. The parent writes a trigger file to `/dev/shm` (a fast, node-local tmpfs).
+2. Worker processes check for the trigger file at multiple synchronization
+   points during each training step.
+3. Workers coordinate via `all_reduce` so that if any rank on any node
+   detects the trigger, all ranks agree to save.
+4. A full-state checkpoint (model + optimizer + LR scheduler) is saved.
+5. Workers exit cleanly.
+
+On resume, the training loop detects the mid-epoch checkpoint, restores the
+full training state, and fast-forwards to the exact step where training was
+interrupted.
+
+## Signals Handled
+
+The following signals are intercepted (SIGKILL cannot be caught):
+
+| Signal | Source |
+|--------|--------|
+| SIGTERM | Kubernetes graceful shutdown (default) |
+| SIGINT | Ctrl-C / some job controllers |
+| SIGUSR1 | Custom preemption controllers |
+| SIGUSR2 | Custom preemption controllers |
+| SIGXCPU | CPU time limit exceeded (resource quotas) |
+| SIGHUP | Terminal disconnect / some eviction paths |
+
+## Usage
+
+### Python API
+
+```python
+from instructlab.training.config import TorchrunArgs, TrainingArgs
+from instructlab.training import run_training
+
+torch_args = TorchrunArgs(
+    nproc_per_node=8,
+    nnodes=1,
+    node_rank=0,
+    rdzv_id=12345,
+    rdzv_endpoint="127.0.0.1:29500",
+)
+
+train_args = TrainingArgs(
+    model_path="Qwen/Qwen2-1.5B-Instruct",
+    data_path="./data.jsonl",
+    data_output_dir="./processed",
+    ckpt_output_dir="./checkpoints",
+    num_epochs=3,
+    on_demand_checkpointing=True,  # Enable the feature
+    # ... other training args
+)
+
+run_training(torch_args, train_args)
+```
+
+### CLI
+
+```bash
+torchrun --nproc-per-node=8 \
+    instructlab/training/main_ds.py \
+    --model_name_or_path Qwen/Qwen2-1.5B-Instruct \
+    --data_path ./data.jsonl \
+    --output_dir ./checkpoints \
+    --on_demand_checkpointing \
+    ...
+```
+
+## Resume Behavior
+
+When a checkpoint saved by on-demand checkpointing is found in the output
+directory, the training loop automatically:
+
+1. Loads the full optimizer and LR scheduler state from the checkpoint.
+2. Reads `global_step` from the checkpoint metadata to determine where
+   training was interrupted.
+3. Resumes at the **same epoch** and fast-forwards to the exact step,
+   skipping already-completed batches.
+
+This differs from epoch-boundary checkpoints, which resume at the start of
+the next epoch.
+
+### Checkpoint Metadata
+
+On-demand checkpoints store additional metadata compared to epoch-boundary
+checkpoints:
+
+```json
+{
+    "current_epoch": 0,
+    "samples_seen": 144,
+    "global_step": 19
+}
+```
+
+The `global_step` field is what distinguishes an on-demand checkpoint from an
+epoch-boundary one. When present, the resume logic keeps `current_epoch`
+unchanged and sets `last_step = global_step` to enable fast-forwarding.
+
+## Multi-Node Training
+
+The trigger file mechanism works correctly across multiple nodes:
+
+- The trigger file lives on `/dev/shm`, which is node-local. Each node's
+  parent process writes its own trigger file when it receives a signal.
+- Workers use `all_reduce(MAX)` to synchronize: if any rank on any node
+  detects a trigger, all ranks on all nodes agree to save.
+- The checkpoint itself is saved to the shared filesystem (the configured
+  `ckpt_output_dir`), accessible by all nodes on resume.
+
+## Stale Trigger Files
+
+If a previous training run was killed before workers could clean up the
+trigger file, the new run's `ParentSignalHandler` detects and removes it
+during initialization. This prevents a new job from immediately
+checkpointing and exiting due to a leftover trigger from a prior run.
+
+## Manually Triggering a Checkpoint
+
+You can trigger a checkpoint-and-exit without sending a signal by writing
+the trigger file directly. This is useful for debugging, testing, or
+integration with custom orchestration that doesn't use Unix signals.
+
+The trigger file is always at a fixed path. To trigger a checkpoint
+(e.g. via `kubectl exec` into the training pod):
+
+```bash
+touch /dev/shm/instructlab_checkpoint_requested
+```
+
+Workers check for the trigger file at each synchronization point in the
+training loop (multiple times per step). Once any rank on any node detects
+it, all ranks coordinate via `all_reduce` to save a checkpoint and exit.
+
+You only need to write the file on **one node** — the `all_reduce` ensures
+all nodes participate even if they don't see the file locally.
+
+From Python:
+
+```python
+from instructlab.training.on_demand_checkpoint import write_trigger_file
+
+write_trigger_file()
+```
+
+## Kubernetes / OpenShift Configuration
+
+To give workers enough time to save a checkpoint before the hard SIGKILL,
+increase `terminationGracePeriodSeconds` in your pod spec:
+
+```yaml
+spec:
+  terminationGracePeriodSeconds: 300  # 5 minutes
+```
+
+The default of 30 seconds may not be enough for large models. The checkpoint
+save time depends on model size, number of GPUs, and filesystem speed.

src/instructlab/training/batch_loss_manager.py

@@ -7,7 +7,8 @@
 """
 
 # Standard
-from dataclasses import dataclass
+from collections.abc import Callable
+from dataclasses import dataclass, field
 import logging
 
 # Third Party
@@ -33,6 +34,7 @@ class BatchMetrics:
     accumulated_aux_loss: torch.Tensor | None
     grad_accum_steps: int
     num_minibatches: int
+    interrupted: bool = field(default=False)
 
 
 class BatchLossManager:
@@ -62,12 +64,22 @@ def __init__(self, model, accelerator, world_size: int, local_rank: int):
         self.local_rank: int = local_rank
         self.torch_device = torch.device("cuda", local_rank)
 
-    def process_batch(self, batch: list[CollatedItem]) -> tuple[BatchMetrics, float]:
+    def process_batch(
+        self,
+        batch: list[CollatedItem],
+        interrupt_check: Callable[[], bool] | None = None,
+    ) -> tuple[BatchMetrics, float]:
         """
         Process a batch of minibatches, computing losses and accumulating gradients.
 
         Args:
             batch: List of minibatches to process
+            interrupt_check: Optional callback invoked at three points per
+                minibatch: before forward, before backward, and after
+                backward. If it returns ``True`` at any point, processing
+                stops early and ``BatchMetrics.interrupted`` is set. Used by
+                on-demand checkpointing to react as quickly as possible
+                instead of waiting for the full optimizer step.
 
         Returns:
             tuple: (BatchMetrics, average_loss_across_ranks)
@@ -82,9 +94,15 @@ def process_batch(self, batch: list[CollatedItem]) -> tuple[BatchMetrics, float]
         accumulated_loss = 0.0
         accumulated_aux_loss = 0.0
         grad_accum_steps = 0
+        interrupted = False
 
         # process each minibatch
         for mb in batch:
+            # Check for on-demand checkpoint before forward
+            if interrupt_check is not None and interrupt_check():
+                interrupted = True
+                break
+
             # extract minibatch-specific info
             micro_batch_size = mb["num_samples"]
             total_length = mb["total_length"]
@@ -96,10 +114,16 @@ def process_batch(self, batch: list[CollatedItem]) -> tuple[BatchMetrics, float]
             # prepare model inputs
             model_inputs = self._prepare_model_inputs(mb)
 
-            # compute loss and backward pass
+            # compute loss (forward pass)
             scaled_loss, raw_losses = self.model.compute_loss(
                 model_inputs, self.world_size, batch_num_loss_counted_tokens
             )
+
+            # Check for on-demand checkpoint before backward
+            if interrupt_check is not None and interrupt_check():
+                interrupted = True
+                break
+
             self.accelerator.backward(scaled_loss)
 
             # accumulate losses
@@ -108,6 +132,11 @@ def process_batch(self, batch: list[CollatedItem]) -> tuple[BatchMetrics, float]
             if raw_losses.aux_loss is not None:
                 accumulated_aux_loss += raw_losses.aux_loss
 
+            # Check for on-demand checkpoint after backward
+            if interrupt_check is not None and interrupt_check():
+                interrupted = True
+                break
+
         # reduce metrics across ranks
         batch_total_samples, batch_total_length = self._reduce_metrics(
             batch_total_samples, batch_total_length
@@ -127,6 +156,7 @@ def process_batch(self, batch: list[CollatedItem]) -> tuple[BatchMetrics, float]
             accumulated_aux_loss=accumulated_aux_loss,
             grad_accum_steps=grad_accum_steps,
             num_minibatches=num_minibatches,
+            interrupted=interrupted,
         )
 
         return metrics, avg_loss_across_ranks
@@ -165,8 +195,8 @@ def _reduce_metrics(
 
     def _compute_average_loss(
         self,
-        accumulated_loss: torch.Tensor,
-        accumulated_aux_loss: torch.Tensor | None,
+        accumulated_loss: torch.Tensor | float,
+        accumulated_aux_loss: torch.Tensor | float | None,
         batch_num_loss_counted_tokens: int,
     ) -> float:
         """Compute average loss across all ranks for metrics logging."""
@@ -177,11 +207,16 @@ def _compute_average_loss(
         if accumulated_aux_loss is not None:
             total_batch_loss += accumulated_aux_loss
 
+        # Extract scalar value — accumulated_loss may be a plain float if the
+        # minibatch loop was interrupted before any forward pass completed.
+        if isinstance(total_batch_loss, torch.Tensor):
+            loss_value = total_batch_loss.detach().item()
+        else:
+            loss_value = float(total_batch_loss)
+
         # reduce across ranks
         avg_loss_across_ranks = self.accelerator.reduce(
-            torch.tensor(
-                total_batch_loss.detach().item(), device=self.accelerator.device
-            ),
+            torch.tensor(loss_value, device=self.accelerator.device),
             reduction="mean",
         ).item()
 

src/instructlab/training/config.py

@@ -361,6 +361,19 @@ class TrainingArgs(BaseModel):
         description="How often to evaluate validation loss (in training steps). Required when validation_split > 0.",
     )
 
+    on_demand_checkpointing: bool = Field(
+        default=False,
+        description=(
+            "Enable on-demand full-state checkpointing triggered by Unix signals. "
+            "When enabled, the parent process intercepts termination signals "
+            "(SIGTERM, SIGINT, SIGUSR1, SIGUSR2, SIGXCPU, SIGHUP) and writes a "
+            "trigger file to /dev/shm. Worker processes check for this trigger "
+            "after each minibatch backward pass and collectively save a distributed "
+            "checkpoint before exiting gracefully. Designed for OpenShift AI / "
+            "KubeFlow training jobs where preemption signals must be handled."
+        ),
+    )
+
     @model_validator(mode="after")
     def validate_validation_config(self):
         if not 0.0 <= self.validation_split < 1.0: