code: def _maybe_write_checkpoint(
      self,
      epoch_idx: int,
      edge_path_idx: int,
      edge_chunk_idx: int,
      current_index: int,
  ) -> None:
      config = self.config
  
      # Preserving a checkpoint requires two steps:
      # - create a snapshot (w/ symlinks) after it's first written;
      # - don't delete it once the following one is written.
      # These two happen in two successive iterations of the main loop: the
      # one just before and the one just after the epoch boundary.
      preserve_old_checkpoint = should_preserve_old_checkpoint(
          self.iteration_manager,
          config.checkpoint_preservation_interval,
      )
      preserve_new_checkpoint = should_preserve_old_checkpoint(
          self.iteration_manager + 1,
          config.checkpoint_preservation_interval,
      )
  
      # Write metadata: for multiple machines, write from rank-0
      logger.info(
          f"Finished epoch {epoch_idx + 1} / {self.iteration_manager.num_epochs}, "
          f"edge path {edge_path_idx + 1} / {self.iteration_manager.num_edge_paths}, "
          f"edge chunk {edge_chunk_idx + 1} / "
          f"{self.iteration_manager.num_edge_chunks}"
      )
      if self.rank == 0:
          for (
              entity,
              embs,
          ) in self.holder.unpartitioned_embeddings.items():
              logger.info(f"Writing {entity} embeddings")
              optimizer = self.trainer.unpartitioned_optimizers[entity]
              self.checkpoint_manager.write(
                  entity,
                  UNPARTITIONED,
                  embs.detach(),
                  optimizer.state_dict(),
                  unpartitioned=True,
              )
  
          logger.info("Writing the metadata")
          state_dict: ModuleStateDict = self.model.state_dict()
          self.checkpoint_manager.write_model(
              state_dict, self.trainer.model_optimizer.state_dict()
          )
  
          logger.info("Writing the training stats")
          all_stats_dicts: List[Dict[str, Any]] = []
          bucket_eval_stats_list = []
          chunk_stats_dict = {
              "epoch_idx": epoch_idx,
              "edge_path_idx": edge_path_idx,
              "edge_chunk_idx": edge_chunk_idx,
          }
          for stats in self.bucket_scheduler.get_stats_for_pass():
              stats_dict = {
                  "lhs_partition": stats.lhs_partition,
                  "rhs_partition": stats.rhs_partition,
                  "index": stats.index,
                  "stats": stats.train.to_dict(),
              }
              if stats.eval_before is not None:
                  stats_dict[
                      "eval_stats_before"
                  ] = stats.eval_before.to_dict()
                  bucket_eval_stats_list.append(stats.eval_before)
  
              if stats.eval_after is not None:
                  stats_dict[
                      "eval_stats_after"
                  ] = stats.eval_after.to_dict()
  
              stats_dict.update(chunk_stats_dict)
              all_stats_dicts.append(stats_dict)
  
          if len(bucket_eval_stats_list) != 0:
              eval_stats_chunk_avg = Stats.average_list(
                  bucket_eval_stats_list
              )
              self.stats_handler.on_stats(
                  index=current_index,
                  eval_stats_chunk_avg=eval_stats_chunk_avg,
              )
              chunk_stats_dict["index"] = current_index
              chunk_stats_dict[
                  "eval_stats_chunk_avg"
              ] = eval_stats_chunk_avg.to_dict()
              all_stats_dicts.append(chunk_stats_dict)
  
          self.checkpoint_manager.append_stats(all_stats_dicts)
  
      logger.info("Writing the checkpoint")
      self.checkpoint_manager.write_new_version(
          config, self.entity_counts, self.embedding_storage_freelist
      )
  
      dist_logger.info(
          "Waiting for other workers to write their parts of the checkpoint"
      )
      self._barrier()
      dist_logger.info("All parts of the checkpoint have been written")
  
      logger.info("Switching to the new checkpoint version")
      self.checkpoint_manager.switch_to_new_version()
  
      dist_logger.info(
          "Waiting for other workers to switch to the new checkpoint version"
      )
      self._barrier()
      dist_logger.info(
          "All workers have switched to the new checkpoint version"
      )
  
      # After all the machines have finished committing
      # checkpoints, we either remove the old checkpoints
      # or we preserve it
      if preserve_new_checkpoint:
          # Add 1 so the index is a multiple of the interval, it looks nicer.
          self.checkpoint_manager.preserve_current_version(
              config, epoch_idx + 1
          )
      if not preserve_old_checkpoint:
          self.checkpoint_manager.remove_old_version(config)
  
imports: import logging
  import math
  import time
  from collections import defaultdict
  from functools import partial
  from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple
  import torch
  import torch.distributed
  from torch.optim import Optimizer
  from torchbiggraph.async_adagrad import AsyncAdagrad
  from torchbiggraph.batching import AbstractBatchProcessor, call, process_in_batches
  from torchbiggraph.bucket_scheduling import BucketStats, DistributedBucketScheduler, LockServer, SingleMachineBucketScheduler
  from torchbiggraph.checkpoint_manager import CheckpointManager, ConfigMetadataProvider, MetadataProvider, PartitionClient, ShardedPartitionClient
  from torchbiggraph.config import ConfigSchema
  from torchbiggraph.distributed import init_process_group, ProcessRanks, start_server
  from torchbiggraph.edgelist import EdgeList
  from torchbiggraph.eval import RankingEvaluator
  from torchbiggraph.graph_storages import EDGE_STORAGES, ENTITY_STORAGES
  from torchbiggraph.losses import AbstractLossFunction, LOSS_FUNCTIONS
  from torchbiggraph.model import make_model, MultiRelationEmbedder
  from torchbiggraph.parameter_sharing import ParameterServer, ParameterSharer
  from torchbiggraph.row_adagrad import RowAdagrad
  from torchbiggraph.stats import Stats, StatsHandler
  from torchbiggraph.types import Bucket, EntityName, FloatTensorType, ModuleStateDict, Partition, Rank, SINGLE_TRAINER, UNPARTITIONED
  from torchbiggraph.util import allocate_shared_tensor, BucketLogger, create_pool, DummyOptimizer, EmbeddingHolder, fast_approx_rand, get_async_result, get_num_workers, hide_distributed_logging, round_up_to_nearest_multiple, split_almost_equally, tag_logs_with_process_name
label_desc: Use image node_id 9 for calling the TrainingCoordinator obj's underlying member method code with example usage: obj._maybe_write_checkpoint(epoch_idx, edge_path_idx, edge_chunk_idx, current_index) without return types