repo_description: 🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training. 
repo_topics: audio deep-learning deepseek gemma glm hacktoberfest llm machine-learning model-hub natural-language-processing nlp pretrained-models python pytorch pytorch-transformers qwen speech-recognition transformer vlm
code: import os
  import signal
  import threading
  
  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: CheckpointManager | None = 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")
  