Repository: huggingface/transformers
Pull request: feat(trainer): Just-in-time (JIT) asynchronous checkpointing using SIGTERM signals

# What does this PR do?

This PR adds Just-In-Time (JIT) asynchronous checkpointing to the Trainer API, enabling graceful handling of training interruptions in preemptible workloads. When enabled via `enable_jit_checkpoint=True`, the trainer
  automatically saves progress upon receiving a SIGTERM signal, preventing loss of training progress in shared cluster environments (e.g., Kueue).

## Motivation

Modern ML training increasingly relies on orchestration platforms like Kubeflow and workload managers like Kueue in shared Kubernetes clusters. These platforms require the ability to:

1. **Pause and Resume Jobs**: Dynamically reallocate resources by pausing long-running training jobs and resuming them later without losing progress
2. **Kueue Preemption**: Support preemptible workloads where lower-priority jobs can be terminated (via SIGTERM) to make room for higher-priority workloads
3. 
4. **Dynamic Scaling**: Enable training jobs to scale down/up based on cluster resource availability

Currently, if a training job is terminated before the next periodic checkpoint, all progress since the last checkpoint is lost. This wastes significant compute resources and training time.

## Implementation

  **New Components:**
    - `trainer_jit_checkpoint.py`: Core implementation with `CheckpointManager` and `JITCheckpointCallback`
    - Registers SIGTERM signal handler to intercept termination signals
    - Implements 3-second kill-wait period to distinguish SIGTERM from SIGKILL
    - Thread-based checkpoint execution to avoid blocking training loop

  **Modified Components:**
  - `trainer.py`: Automatically adds JITCheckpointCallback when `enable_jit_checkpoint=True`
  - `training_args.py`: New `enable_jit_checkpoint` boolean flag (configurable via `ENABLE_JIT_CHECKPOINT` environment variable)

**Environment Variable Support Rationale:**
 The new environment variables (`ENABLE_JIT_CHECKPOINT`, `SAVE_STRATEGY`, `SAVE_STEPS`, `SAVE_TOTAL_LIMIT`, `OUTPUT_DIR`, `RESUME_FROM_CHECKPOINT`) allow orchestration platforms like Kubeflow to automatically configure checkpointing behavior
  without requiring users to modify their training scripts. This enables:
  - Centralized checkpoint configuration through cluster operators
  - Dynamic checkpoint policies based on workload type
  - Seamless integration with existing training code

  **Key Features:**
  - Graceful training termination after checkpoint completion
  - Environment variable support for easier configuration in cluster environments
  - Compatible with existing checkpoint management (save_steps, save_strategy, etc.)

  ## Testing

  Comprehensive unit tests added in `tests/trainer/test_trainer_jit_checkpoint.py` covering:
  - CheckpointManager initialization and configuration
  - SIGTERM signal handler flow and duplicate signal handling
  - Callback integration with Trainer lifecycle hooks

Fixes #38961

## Who can review?

@SunMarc - This PR adds emergency checkpointing functionality to the Trainer class

## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [issue](https://github.com/huggingface/transformers/issues/38961)? Please add a link
      to it if that's the case.
- [x] Updated Trainer args doc.
- [x] Did you write any new necessary tests?


Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Changed files:
filename: src/transformers/trainer.py
patch: @@ -642,6 +642,16 @@ def __init__(
                 "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method."
             )
         default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to)
+
+        # Add JIT checkpoint callback if enabled
+        if self.args.enable_jit_checkpoint:
+            from .trainer_jit_checkpoint import JITCheckpointCallback
+
+            jit_callback = JITCheckpointCallback()
+            default_callbacks = default_callbacks + [jit_callback]
+            # Set trainer reference for JIT callback after initialization
+            jit_callback.set_trainer(self)
+
         callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks
         self.callback_handler = CallbackHandler(
             callbacks, self.model, self.processing_class, self.optimizer, self.lr_scheduler

filename: src/transformers/trainer_jit_checkpoint.py
patch: @@ -0,0 +1,126 @@
+import os
+import signal
+import threading
+from typing import Optional
+
+from .trainer_callback import TrainerCallback
+from .trainer_utils import PREFIX_CHECKPOINT_DIR
+from .utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class CheckpointManager:
+    def __init__(self, trainer, kill_wait: int = 3):
+        """
+        Initialize the CheckpointManager for Just-In-Time checkpoint handling.
+
+        Args:
+            trainer: The Trainer instance that will be used to save checkpoints when SIGTERM is received.
+            kill_wait (`int`, *optional*, defaults to 3): Grace period to distinguish between SIGTERM and SIGKILL.
+        """
+        self.trainer = trainer
+        self.is_checkpoint_requested = False
+        self._original_sigterm_handler = None
+        self.kill_wait = kill_wait
+
+    def setup_signal_handler(self):
+        self._original_sigterm_handler = signal.signal(signal.SIGTERM, self._sigterm_handler)
+        logger.info("JIT checkpoint signal handler registered for SIGTERM")
+
+    def _sigterm_handler(self, signum, frame):
+        if self.is_checkpoint_requested:
+            return
+
+        logger.info(f"SIGTERM received, will request JIT checkpoint after {self.kill_wait}s")
+        threading.Timer(self.kill_wait, self._enable_checkpoint).start()
+
+    def _enable_checkpoint(self):
+        logger.info("Kill wait period elapsed, requesting checkpoint")
+        self.is_checkpoint_requested = True
+
+    def execute_jit_checkpoint(self):
+        try:
+            # Set checkpoint flag to False to avoid multiple checkpoints getting triggered by other callbacks
+            self.is_checkpoint_requested = False
+
+            logger.info("Starting JIT checkpointing...")
+            current_step = self.trainer.state.global_step
+            logger.info(f"Saving JIT checkpoint at step {current_step}")
+
+            output_dir = self.trainer._get_output_dir(trial=None)
+            checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{current_step}"
+            checkpoint_path = os.path.join(output_dir, checkpoint_folder)
+
+            # Create checkpoint directory
+            os.makedirs(checkpoint_path, exist_ok=True)
+
+            # Create a sentinel file to indicate checkpointing is in progress
+            sentinel_file = os.path.join(output_dir, checkpoint_folder, "checkpoint-is-incomplete.txt")
+            with open(sentinel_file, "w") as f:
+                f.write(f"Checkpoint started at step {current_step} and in progress...")
+            logger.info(f"Created checkpoint progress sentinel marker file: {sentinel_file}")
+
+            # Invoke the trainer's checkpoint method directly
+            self.trainer._save_checkpoint(self.trainer.model, trial=None)
+
+            # Remove sentinel file upon successful checkpointing
+            if os.path.exists(sentinel_file):
+                os.remove(sentinel_file)
+                logger.info("Sentinel marker file removed")
+
+            logger.info("Immediate JIT checkpoint completed successfully")
+
+        except Exception as e:
+            logger.error(f"Failed to save JIT checkpoint: {e}")
+            raise
+
+
+class JITCheckpointCallback(TrainerCallback):
+    """
+    Callback for Just-In-Time checkpointing on SIGTERM signals.
+
+    When SIGTERM is received, the checkpoint manager sets `is_checkpoint_requested=True`.
+    The callbacks detect this flag and set `control.should_training_stop=True`, which signals
+    the Trainer's training loop to exit gracefully after saving the checkpoint.
+    """
+
+    def __init__(self):
+        self.trainer = None
+        self.jit_manager: Optional[CheckpointManager] = None
+
+    def set_trainer(self, trainer):
+        self.trainer = trainer
+        if trainer.args.enable_jit_checkpoint:
+            self.jit_manager = CheckpointManager(trainer=trainer)
+            self.jit_manager.setup_signal_handler()
+            logger.info("JIT checkpointing enabled")
+
+    def on_pre_optimizer_step(self, args, state, control, **kwargs):
+        if self.jit_manager and self.jit_manager.is_checkpoint_requested:
+            control.should_training_stop = True
+            self.jit_manager.execute_jit_checkpoint()
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        if self.jit_manager and self.jit_manager.is_checkpoint_requested:
+            control.should_training_stop = True
+            self.jit_manager.execute_jit_checkpoint()
+
+    def on_step_end(self, args, state, control, **kwargs):
+        if self.jit_manager and self.jit_manager.is_checkpoint_requested:
+            control.should_save = False
+            control.should_training_stop = True
+            self.jit_manager.execute_jit_checkpoint()
+
+    def on_epoch_end(self, args, state, control, **kwargs):
+        if self.jit_manager and self.jit_manager.is_checkpoint_requested:
+            control.should_save = False
+            control.should_training_stop = True
+            self.jit_manager.execute_jit_checkpoint()
+
+    def on_train_end(self, args, state, control, **kwargs):
+        #  Restore original SIGTERM handler
+        if self.jit_manager and self.jit_manager._original_sigterm_handler is not None:
+            signal.signal(signal.SIGTERM, self.jit_manager._original_sigterm_handler)
+            logger.info("Restored original SIGTERM handler after training completion")

filename: src/transformers/training_args.py
patch: @@ -340,6 +340,17 @@ class TrainingArguments:
             `save_total_limit=5` and `load_best_model_at_end`, the four last checkpoints will always be retained
             alongside the best model. When `save_total_limit=1` and `load_best_model_at_end`, it is possible that two
             checkpoints are saved: the last one and the best one (if they are different).
+        enable_jit_checkpoint (`bool`, *optional*, defaults to `False`):
+            Whether to enable Just-In-Time (JIT) checkpointing on SIGTERM signal. When enabled, training will
+            checkpoint upon receiving SIGTERM, allowing for graceful termination without losing
+            progress. This is particularly useful for shared clusters with preemptible workloads (e.g., Kueue).
+            **Important**: You must configure your orchestrator's graceful shutdown period to allow sufficient time
+            for checkpoint completion. For Kubernetes, set `terminationGracePeriodSeconds` in your job definition
+            (method varies by cloud-native trainer: Kubeflow, Ray, etc.). Note: the default is only 30 seconds,
+            which is typically insufficient. For Slurm, use `--signal=USR1@<seconds>` in your sbatch script to send
+            SIGTERM with adequate time before the job time limit. Calculate the required grace period as: longest
+            possible iteration time + checkpoint saving time. For example, if an iteration takes 2 minutes and
+            checkpoint saving takes 2 minutes, set at least 4 minutes (240 seconds) of grace time.
         save_safetensors (`bool`, *optional*, defaults to `True`):
             Use [safetensors](https://huggingface.co/docs/safetensors) saving and loading for state dicts instead of
             default `torch.load` and `torch.save`.
@@ -929,7 +940,23 @@ class TrainingArguments:
                 " for `save_total_limit=5` and `load_best_model_at_end=True`, the four last checkpoints will always be"
                 " retained alongside the best model. When `save_total_limit=1` and `load_best_model_at_end=True`,"
                 " it is possible that two checkpoints are saved: the last one and the best one (if they are different)."
-                " Default is unlimited checkpoints"
+                " Default is unlimited checkpoints."
+            )
+        },
+    )
+    enable_jit_checkpoint: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether to enable Just-In-Time (JIT) checkpointing on SIGTERM signal. "
+                "When enabled, training will checkpoint upon receiving SIGTERM, "
+                "allowing for graceful termination without losing progress. "
+                "This is particularly useful for shared clusters with preemptible workloads (Kueue). "
+                "IMPORTANT: You must configure your orchestrator's graceful shutdown period. "
+                "Kubernetes: set terminationGracePeriodSeconds (default 30s is insufficient!) in your job definition. "
+                "Slurm: use --signal=USR1@<seconds> in sbatch to send SIGTERM before time limit. "
+                "Calculate required grace period as: iteration time + checkpoint saving time. "
+                "Example: 2min iteration + 2min checkpoint = 240 seconds minimum."
             )
         },
     )

filename: tests/trainer/test_trainer_jit_checkpoint.py
patch: @@ -0,0 +1,406 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import signal
+import tempfile
+import unittest
+from unittest.mock import Mock, patch
+
+from transformers import TrainingArguments, is_torch_available
+from transformers.testing_utils import require_torch
+
+
+if is_torch_available():
+    from transformers.trainer_jit_checkpoint import CheckpointManager, JITCheckpointCallback
+
+    from .test_trainer import RegressionDataset, RegressionModelConfig, RegressionPreTrainedModel
+
+
+@require_torch
+class JITCheckpointTest(unittest.TestCase):
+    def setUp(self):
+        self.test_dir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        import shutil
+
+        shutil.rmtree(self.test_dir, ignore_errors=True)
+
+    def get_trainer(self, enable_jit=True):
+        """Helper method to create a trainer with JIT checkpointing enabled."""
+        from transformers import Trainer
+
+        model_config = RegressionModelConfig(a=1.5, b=2.5)
+        model = RegressionPreTrainedModel(model_config)
+
+        args = TrainingArguments(
+            output_dir=self.test_dir,
+            enable_jit_checkpoint=enable_jit,
+            per_device_train_batch_size=16,
+            learning_rate=0.1,
+            logging_steps=1,
+            num_train_epochs=1,
+            max_steps=10,
+            save_steps=10,
+        )
+
+        train_dataset = RegressionDataset(length=64)
+
+        return Trainer(model=model, args=args, train_dataset=train_dataset)
+
+    def test_checkpoint_manager_initialization(self):
+        """Test CheckpointManager initialization with different configurations."""
+        trainer = self.get_trainer()
+
+        # Test with default parameters
+        manager = CheckpointManager(trainer)
+        self.assertEqual(manager.trainer, trainer)
+        self.assertEqual(manager.kill_wait, 3)
+        self.assertFalse(manager.is_checkpoint_requested)
+
+        # Test with custom parameters
+        manager_custom = CheckpointManager(trainer, kill_wait=5)
+        self.assertEqual(manager_custom.kill_wait, 5)
+
+    def test_signal_handler_setup(self):
+        """Test signal handler setup and restoration."""
+        trainer = self.get_trainer()
+        manager = CheckpointManager(trainer)
+
+        # Store original handler
+        original_handler = signal.signal(signal.SIGTERM, signal.SIG_DFL)
+
+        try:
+            # Setup JIT signal handler
+            manager.setup_signal_handler()
+
+            # Verify handler is set
+            current_handler = signal.signal(signal.SIGTERM, signal.SIG_DFL)
+            self.assertNotEqual(current_handler, signal.SIG_DFL)
+
+            # Verify original handler is stored
+            self.assertIsNotNone(manager._original_sigterm_handler)
+
+        finally:
+            # Restore original handler
+            signal.signal(signal.SIGTERM, original_handler)
+
+    @patch("threading.Timer")
+    def test_sigterm_handler_flow(self, mock_timer):
+        """Test SIGTERM handler execution flow."""
+        trainer = self.get_trainer()
+        manager = CheckpointManager(trainer, kill_wait=2)
+
+        # Mock timer to prevent actual threading
+        mock_timer_instance = Mock()
+        mock_timer.return_value = mock_timer_instance
+
+        # Test first SIGTERM call
+        self.assertFalse(manager.is_checkpoint_requested)
+        manager._sigterm_handler(signal.SIGTERM, None)
+
+        # Verify checkpoint was NOT immediately requested (timer is used)
+        self.assertFalse(manager.is_checkpoint_requested)
+
+        # Verify timer was created with kill_wait period and correct callback
+        mock_timer.assert_called_once_with(2, manager._enable_checkpoint)
+        mock_timer_instance.start.assert_called_once()
+
+        # Manually trigger the timer callback to test flag setting
+        manager._enable_checkpoint()
+
+        # Verify checkpoint is now requested
+        self.assertTrue(manager.is_checkpoint_requested)
+
+        # Test second SIGTERM call (should be ignored)
+        mock_timer.reset_mock()
+        manager._sigterm_handler(signal.SIGTERM, None)
+
+        # Verify no additional timer was created
+        mock_timer.assert_not_called()
+
+    def test_toggle_checkpoint_flag(self):
+        """Test the toggle checkpoint flag method."""
+        trainer = self.get_trainer()
+        manager = CheckpointManager(trainer)
+
+        # Initially should not be requested
+        self.assertFalse(manager.is_checkpoint_requested)
+
+        # Toggle flag
+        manager._enable_checkpoint()
+
+        # Should now be requested
+        self.assertTrue(manager.is_checkpoint_requested)
+
+    def test_execute_jit_checkpoint(self):
+        """Test the checkpoint execution logic with sentinel file."""
+        from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+
+        trainer = self.get_trainer()
+        manager = CheckpointManager(trainer)
+
+        # Mock trainer's save checkpoint method
+        trainer._save_checkpoint = Mock()
+        trainer.state.global_step = 42
+
+        # Set checkpoint requested flag
+        manager.is_checkpoint_requested = True
+
+        # Execute checkpoint
+        manager.execute_jit_checkpoint()
+
+        # Verify checkpoint was called
+        trainer._save_checkpoint.assert_called_once_with(trainer.model, trial=None)
+
+        # Verify checkpoint flag was reset
+        self.assertFalse(manager.is_checkpoint_requested)
+
+        # Verify sentinel file was removed (should be in checkpoint-42 folder)
+        checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-42"
+        sentinel_file = os.path.join(self.test_dir, checkpoint_folder, "checkpoint-is-incomplete.txt")
+        self.assertFalse(os.path.exists(sentinel_file))
+
+    def test_execute_jit_checkpoint_sentinel_file_cleanup(self):
+        """Test that sentinel file is cleaned up after successful checkpoint."""
+        from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+
+        trainer = self.get_trainer()
+        manager = CheckpointManager(trainer)
+
+        # Mock trainer's save checkpoint method
+        trainer._save_checkpoint = Mock()
+        trainer.state.global_step = 42
+
+        checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-42"
+        sentinel_file = os.path.join(self.test_dir, checkpoint_folder, "checkpoint-is-incomplete.txt")
+
+        # Execute checkpoint
+        manager.execute_jit_checkpoint()
+
+        # Verify sentinel file doesn't exist after successful checkpoint
+        self.assertFalse(os.path.exists(sentinel_file))
+
+    def test_execute_jit_checkpoint_with_exception(self):
+        """Test checkpoint execution with exception handling."""
+        trainer = self.get_trainer()
+        manager = CheckpointManager(trainer)
+
+        # Mock trainer's save checkpoint method to raise exception
+        trainer._save_checkpoint = Mock(side_effect=Exception("Checkpoint failed"))
+        trainer.state.global_step = 42
+
+        # Test that exception is re-raised
+        with self.assertRaises(Exception) as context:
+            manager.execute_jit_checkpoint()
+
+        self.assertEqual(str(context.exception), "Checkpoint failed")
+
+        # Verify checkpoint flag was still reset to avoid multiple attempts
+        self.assertFalse(manager.is_checkpoint_requested)
+
+    def test_jit_checkpoint_callback_initialization(self):
+        """Test JITCheckpointCallback initialization."""
+        callback = JITCheckpointCallback()
+
+        self.assertIsNone(callback.trainer)
+        self.assertIsNone(callback.jit_manager)
+
+    def test_jit_checkpoint_callback_set_trainer_enabled(self):
+        """Test setting trainer with JIT checkpointing enabled."""
+        trainer = self.get_trainer(enable_jit=True)
+        callback = JITCheckpointCallback()
+
+        with patch.object(CheckpointManager, "setup_signal_handler") as mock_setup:
+            callback.set_trainer(trainer)
+
+            self.assertEqual(callback.trainer, trainer)
+            self.assertIsNotNone(callback.jit_manager)
+            self.assertIsInstance(callback.jit_manager, CheckpointManager)
+            mock_setup.assert_called_once()
+
+    def test_jit_checkpoint_callback_set_trainer_disabled(self):
+        """Test setting trainer with JIT checkpointing disabled."""
+        trainer = self.get_trainer(enable_jit=False)
+        callback = JITCheckpointCallback()
+
+        callback.set_trainer(trainer)
+
+        self.assertEqual(callback.trainer, trainer)
+        self.assertIsNone(callback.jit_manager)
+
+    def test_jit_checkpoint_callback_on_pre_optimizer_step(self):
+        """Test callback behavior during pre-optimizer step."""
+        trainer = self.get_trainer()
+        callback = JITCheckpointCallback()
+        callback.set_trainer(trainer)
+
+        # Mock control object
+        control = Mock()
+        control.should_training_stop = False
+
+        # Mock execute method
+        with patch.object(callback.jit_manager, "execute_jit_checkpoint") as mock_execute:
+            # Test when checkpoint not requested
+            callback.jit_manager.is_checkpoint_requested = False
+            callback.on_pre_optimizer_step(trainer.args, trainer.state, control)
+            self.assertFalse(control.should_training_stop)
+            mock_execute.assert_not_called()
+
+            # Test when checkpoint requested
+            callback.jit_manager.is_checkpoint_requested = True
+            callback.on_pre_optimizer_step(trainer.args, trainer.state, control)
+            self.assertTrue(control.should_training_stop)
+            mock_execute.assert_called_once()
+
+    def test_jit_checkpoint_callback_on_step_begin(self):
+        """Test callback behavior at step begin."""
+        trainer = self.get_trainer()
+        callback = JITCheckpointCallback()
+        callback.set_trainer(trainer)
+
+        # Mock control object
+        control = Mock()
+        control.should_training_stop = False
+
+        # Mock execute method
+        with patch.object(callback.jit_manager, "execute_jit_checkpoint") as mock_execute:
+            # Test when checkpoint not requested
+            callback.jit_manager.is_checkpoint_requested = False
+            callback.on_step_begin(trainer.args, trainer.state, control)
+            self.assertFalse(control.should_training_stop)
+            mock_execute.assert_not_called()
+
+            # Test when checkpoint requested
+            callback.jit_manager.is_checkpoint_requested = True
+            callback.on_step_begin(trainer.args, trainer.state, control)
+            self.assertTrue(control.should_training_stop)
+            mock_execute.assert_called_once()
+
+    def test_jit_checkpoint_callback_on_step_end(self):
+        """Test callback behavior at step end."""
+        trainer = self.get_trainer()
+        callback = JITCheckpointCallback()
+        callback.set_trainer(trainer)
+
+        # Mock control object
+        control = Mock()
+        control.should_training_stop = False
+        control.should_save = True
+
+        # Mock execute method
+        with patch.object(callback.jit_manager, "execute_jit_checkpoint") as mock_execute:
+            # Test when checkpoint not requested
+            callback.jit_manager.is_checkpoint_requested = False
+            callback.on_step_end(trainer.args, trainer.state, control)
+            self.assertFalse(control.should_training_stop)
+            mock_execute.assert_not_called()
+
+            # Reset control
+            control.should_save = True
+
+            # Test when checkpoint requested
+            callback.jit_manager.is_checkpoint_requested = True
+            callback.on_step_end(trainer.args, trainer.state, control)
+            self.assertTrue(control.should_training_stop)
+            self.assertFalse(control.should_save)
+            mock_execute.assert_called_once()
+
+    def test_jit_checkpoint_callback_on_epoch_end(self):
+        """Test callback behavior at epoch end."""
+        trainer = self.get_trainer()
+        callback = JITCheckpointCallback()
+        callback.set_trainer(trainer)
+
+        # Mock control object
+        control = Mock()
+        control.should_save = True
+        control.should_training_stop = False
+
+        # Mock execute method
+        with patch.object(callback.jit_manager, "execute_jit_checkpoint") as mock_execute:
+            # Test when checkpoint not requested
+            callback.jit_manager.is_checkpoint_requested = False
+            callback.on_epoch_end(trainer.args, trainer.state, control)
+            # should_save should remain unchanged when checkpoint not requested
+            self.assertTrue(control.should_save)
+            self.assertFalse(control.should_training_stop)
+            mock_execute.assert_not_called()
+
+            # Reset control
+            control.should_save = True
+            control.should_training_stop = False
+
+            # Test when checkpoint requested
+            callback.jit_manager.is_checkpoint_requested = True
+            callback.on_epoch_end(trainer.args, trainer.state, control)
+            self.assertFalse(control.should_save)
+            self.assertTrue(control.should_training_stop)
+            mock_execute.assert_called_once()
+
+    def test_jit_checkpoint_callback_on_train_end(self):
+        """Test signal handler restoration on training end."""
+        trainer = self.get_trainer()
+        callback = JITCheckpointCallback()
+
+        # Store original SIGTERM handler
+        original_handler = signal.signal(signal.SIGTERM, signal.SIG_DFL)
+
+        try:
+            callback.set_trainer(trainer)
+
+            # Verify signal handler was set up
+            self.assertIsNotNone(callback.jit_manager._original_sigterm_handler)
+
+            # Mock control object
+            control = Mock()
+
+            # Call on_train_end
+            callback.on_train_end(trainer.args, trainer.state, control)
+
+            # Verify signal handler was restored
+            current_handler = signal.signal(signal.SIGTERM, signal.SIG_DFL)
+            self.assertEqual(current_handler, callback.jit_manager._original_sigterm_handler)
+
+        finally:
+            # Restore original handler for cleanup
+            signal.signal(signal.SIGTERM, original_handler)
+
+    @patch("threading.Timer")
+    def test_kill_wait_period(self, mock_timer):
+        """Test the kill wait period functionality."""
+        trainer = self.get_trainer()
+        manager = CheckpointManager(trainer, kill_wait=5)
+
+        mock_timer_instance = Mock()
+        mock_timer.return_value = mock_timer_instance
+
+        manager._sigterm_handler(signal.SIGTERM, None)
+
+        # Verify Timer was created with the correct kill_wait period and callback
+        mock_timer.assert_called_once_with(5, manager._enable_checkpoint)
+        mock_timer_instance.start.assert_called_once()
+
+    def test_integration_with_trainer(self):
+        """Test integration of JIT checkpointing with Trainer."""
+        trainer = self.get_trainer(enable_jit=True)
+
+        # Check that JIT callback was added
+        jit_callbacks = [cb for cb in trainer.callback_handler.callbacks if isinstance(cb, JITCheckpointCallback)]
+        self.assertEqual(len(jit_callbacks), 1)
+
+        jit_callback = jit_callbacks[0]
+        self.assertIsNotNone(jit_callback.jit_manager)
+        self.assertEqual(jit_callback.trainer, trainer)