code: def _run_e2e_test(self, compile, model_type, async_op=False):
      model, optim = self._create_model(compile, ModelType.NONE)
      _train(model, optim, train_steps=2)
  
      dist_model, dist_optim = self._create_model(compile, model_type)
      _train(dist_model, dist_optim, train_steps=2)
  
      original_stateful_obj = (
          TestStatefulObj()
      )  # tests arbitrary saving/loading
      sd = {
          "model": dist_model,
          "optimizer": dist_optim,
          "s": original_stateful_obj,
      }
  
      checkpointer = _FileSystemCheckpointer(self.temp_dir)
      if async_op:
          f = checkpointer.async_save(state_dict=sd)
          t = time.monotonic()
          while not f.done():
              time.sleep(1)
              print(f"still waiting... {time.monotonic() - t}")
  
          f.result()
      else:
          checkpointer.save(state_dict=sd)
  
      loaded_stateful_obj = TestStatefulObj()
      dist_model, dist_optim = self._create_model(compile, model_type)
  
      loaded_stateful_obj = TestStatefulObj()
      dist_model, dist_optim = self._create_model(compile, model_type)
  
      checkpointer.load(
          state_dict={
              "model": dist_model,
              "optimizer": dist_optim,
              "s": loaded_stateful_obj,
          }
      )
  
      self.assertEqual(original_stateful_obj, loaded_stateful_obj)
  
      # train one more step on both models
      loss = _train(model, optim, train_steps=1)
      dist_loss = _train(dist_model, dist_optim, train_steps=1)
      self.assertEqual(loss, dist_loss)
  
      dist_msd, dist_osd = get_state_dict(
          dist_model, optimizers=dist_optim
      )
      model_sd, optim_sd = get_state_dict(model, optimizers=optim)
  
      self._verify_msd(model_sd, dist_msd)
      self._verify_osd_by_load(
          model, optim, self._optim(model), dist_osd
      )
  
imports: import time
  from enum import auto, Enum
  import torch
  import torch.distributed
  import torch.distributed.checkpoint
  import torch.nn
  import torch.nn.functional
  from torch.distributed._tensor.device_mesh import init_device_mesh
  from torch.distributed.checkpoint.filesystem import _FileSystemCheckpointer
  from torch.distributed.checkpoint.state_dict import _patch_model_state_dict, _patch_optimizer_state_dict, get_state_dict
  from torch.distributed.distributed_c10d import ReduceOp
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp.api import ShardingStrategy
  from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module, RowwiseParallel
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, skip_if_lt_x_gpu, with_comms
  from torch.testing._internal.distributed.checkpoint_utils import with_temp_dir
  from torch.testing._internal.distributed.common_state_dict import VerifyStateDictMixin
label_desc: Use image node_id 6 for calling the TestE2ESaveAndLoad obj's underlying member method code with example usage: obj._run_e2e_test(compile, model_type, async_op) without return types