class_docstr: A callback which periodically saves the application state during training using `TorchSnapshot <https://pytorch.org/torchsnapshot/>`_.
  
  This callback supplements the application state provided by :class:`torchtnt.unit.AppStateMixin`
  with the train progress, train dataloader (if applicable), and random number generator state.
  
  If used with :func:`torchtnt.framework.fit`, this class will also save the evaluation progress state.
  
  Checkpoints will be saved under ``dirpath/epoch_{epoch}_step_{step}`` where step is the *total* number of training and evaluation steps completed across all epochs.
  
  Args:
      dirpath: Parent directory to save snapshots to.
      save_every_n_train_steps: Frequency of steps with which to save snapshots during the train epoch. If None, no intra-epoch snapshots are generated.
      save_every_n_epochs: Frequency of epochs with which to save snapshots during training. If None, no end-of-epoch snapshots are generated.
      keep_last_n_checkpoints: Number of most recent checkpoints to keep. If None, all checkpoints are kept. If an excess of existing checkpoints are present, the oldest ones will be deleted to clean the difference.
      process_group: The process group on which the ranks will communicate on. default: ``None`` (the entire world)
      async_checkpoint: Whether to perform asynchronous snapshotting. Default: ``True``.
      replicated: A glob-pattern of replicated key names that indicate which application state entries have the same state across all processes.
          For more information, see https://pytorch.org/torchsnapshot/main/api_reference.html#torchsnapshot.Snapshot.take.
  
          .. warning:: The replication property is safer to not set, and should only be used if really needed.
                       Things like metrics, grad_scalers, etc should not be marked as replicated as they
                       may contain different values across processes. If unsure, leave this field unset.
  
      storage_options: Additional keyword options for the storage plugin to use, to be passed to `torchsnapshot.Snapshot <https://pytorch.org/torchsnapshot/stable/api_reference.html#torchsnapshot.Snapshot>`_.
          See each storage plugin's documentation for customizations.
      knob_options: Additional keyword options for the snapshot knobs
  
  Note:
      If torch.distributed is available and default process group is initialized, the constructor will call a collective operation for rank 0 to broadcast the dirpath to all other ranks
  
  Note:
      If checkpointing FSDP model, you can set state_dict type calling `set_state_dict_type <https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.set_state_dict_type>`_ prior to starting training.
class_docstr_tok: ["A","callback","which","periodically","saves","the","application","state","during","training","using","`","TorchSnapshot","<","https",":","\/\/pytorch.org\/torchsnapshot\/",">","`","_",".","This","callback","supplements","the","application","state","provided","by",":","class",":","`","torchtnt.unit.AppStateMixin","`","with","the","train","progress",",","train","dataloader","(","if","applicable",")",",","and","random","number","generator","state",".","If","used","with",":","func",":","`","torchtnt.framework.fit","`",",","this","class","will","also","save","the","evaluation","progress","state",".","Checkpoints","will","be","saved","under","``","dirpath\/epoch_","{","epoch","}","_step_","{","step","}","``","where","step","is","the","*","total","*","number","of","training","and","evaluation","steps","completed","across","all","epochs",".","Args",":","dirpath",":","Parent","directory","to","save","snapshots","to",".","save_every_n_train_steps",":","Frequency","of","steps","with","which","to","save","snapshots","during","the","train","epoch",".","If","None",",","no","intra-epoch","snapshots","are","generated",".","save_every_n_epochs",":","Frequency","of","epochs","with","which","to","save","snapshots","during","training",".","If","None",",","no","end-of-epoch","snapshots","are","generated",".","keep_last_n_checkpoints",":","Number","of","most","recent","checkpoints","to","keep",".","If","None",",","all","checkpoints","are","kept",".","If","an","excess","of","existing","checkpoints","are","present",",","the","oldest","ones","will","be","deleted","to","clean","the","difference",".","process_group",":","The","process","group","on","which","the","ranks","will","communicate","on",".","default",":","``","None","``","(","the","entire","world",")","async_checkpoint",":","Whether","to","perform","asynchronous","snapshotting",".","Default",":","``","True","``",".","replicated",":","A","glob-pattern","of","replicated","key","names","that","indicate","which","application","state","entries","have","the","same","state","across","all","processes",".","For","more","information",",","see","https",":","\/\/pytorch.org\/torchsnapshot\/main\/api_reference.html","#","torchsnapshot.Snapshot.take",".","..","warning",":",":","The","replication","property","is","safer","to","not","set",",","and","should","only","be","used","if","really","needed",".","Things","like","metrics",",","grad_scalers",",","etc","should","not","be","marked","as","replicated","as","they","may","contain","different","values","across","processes",".","If","unsure",",","leave","this","field","unset",".","storage_options",":","Additional","keyword","options","for","the","storage","plugin","to","use",",","to","be","passed","to","`","torchsnapshot.Snapshot","<","https",":","\/\/pytorch.org\/torchsnapshot\/stable\/api_reference.html","#","torchsnapshot.Snapshot",">","`","_",".","See","each","storage","plugin","'s","documentation","for","customizations",".","knob_options",":","Additional","keyword","options","for","the","snapshot","knobs","Note",":","If","torch.distributed","is","available","and","default","process","group","is","initialized",",","the","constructor","will","call","a","collective","operation","for","rank","0","to","broadcast","the","dirpath","to","all","other","ranks","Note",":","If","checkpointing","FSDP","model",",","you","can","set","state_dict","type","calling","`","set_state_dict_type","<","https",":","\/\/pytorch.org\/docs\/stable\/fsdp.html","#","torch.distributed.fsdp.FullyShardedDataParallel.set_state_dict_type",">","`","_","prior","to","starting","training","."]
docstr: Validate there's no key collision for the app state.
docstr_tok: ["Validate","there","'s","no","key","collision","for","the","app","state","."]
code: def on_train_start(self, state: State, unit: TTrainUnit) -> None:
      """Validate there's no key collision for the app state."""
      app_state = _prepare_app_state(unit)
      _check_app_state_collision(app_state)
  
      super().on_train_start(state, unit)
  
imports: from __future__ import annotations
  import logging
  from contextlib import contextmanager, ExitStack
  from typing import Any, Dict, Generator, Iterable, List, Optional, Set, Union
  import torch.distributed
  from torchtnt.framework.callbacks._checkpoint_utils import _prepare_app_state, _prepare_app_state_for_checkpoint, _prepare_app_state_for_restore, _TRAIN_DL_STATE_KEY
  from torchtnt.framework.callbacks.base_checkpointer import BaseCheckpointer
  from torchtnt.framework.callbacks.checkpointer_types import KnobOptions, RestoreOptions
  from torchtnt.framework.state import State
  from torchtnt.framework.unit import AppStateMixin, TEvalUnit, TPredictUnit, TTrainData, TTrainUnit
  from torchtnt.framework.utils import get_timing_context
  from torchtnt.utils.optimizer import init_optim_state
  from torchtnt.utils.rank_zero_log import rank_zero_info, rank_zero_warn
  from torchtnt.utils.stateful import Stateful
label_desc: Use image node_id 2 for calling the TorchSnapshotSaver obj's underlying member method code with example usage: obj.on_train_start(state, unit) without return types