docstr: Simulates using the Checkpointer during distributed training with a sharded model.
docstr_tok: ["Simulates","using","the","Checkpointer","during","distributed","training","with","a","sharded","model","."]
code: def test_default_distributed_with_sharded_state(self):
      """
      Simulates using the Checkpointer during distributed training with a sharded model.
      """
      world_size = 2
      default_num_to_keep = 2
      num_epochs = 5
      target = [
          (e, 0)
          for e in range(num_epochs - default_num_to_keep, num_epochs)
      ]
  
      checkpointers = [
          Checkpointer(serialization_dir=self.TEST_DIR)
          for _ in range(world_size)
      ]
      for i, checkpointer in enumerate(checkpointers):
          checkpointer._rank = i
          checkpointer.state_is_sharded = True
  
      for epochs_completed in range(num_epochs):
          for batches_completed in [0, 5, 10]:
              for i, checkpointer in enumerate(checkpointers):
                  state = {
                      "epochs_completed": epochs_completed,
                      "batches_in_epoch_completed": batches_completed,
                      "rank": i,
                  }
                  checkpointer.maybe_save_checkpoint(
                      FakeTrainer(
                          model_state=state, training_state=state
                      ),
                      epochs_completed,
                      batches_completed,
                  )
  
      for i, checkpointer in enumerate(checkpointers):
          checkpoint = checkpointer.load_checkpoint()
          assert checkpoint is not None
          model_state, training_state = checkpoint
          assert model_state["rank"] == i
          assert training_state["rank"] == i
  
          models, training = self.retrieve_and_delete_saved(shard=i)
          assert models == training == target
  
file_path: allennlp/tests/training/checkpointer_test.py
imports: import os
  import time
  from typing import Optional
  from allennlp.common.testing import AllenNlpTestCase
  from allennlp.common.params import Params
  from allennlp.training import Checkpointer, Trainer
  from allennlp.training.trainer import TrainerCheckpoint
label_desc: Use image node_id 7 for calling the TestCheckpointer obj's underlying member method code with example usage: obj.test_default_distributed_with_sharded_state() without return types