docstr: Each epoch/pass, for each partition pair, loads in embeddings and edgelist
  from disk, runs HOGWILD training on them, and writes partitions back to disk.
docstr_tok: ["Each","epoch\/pass",",","for","each","partition","pair",",","loads","in","embeddings","and","edgelist","from","disk",",","runs","HOGWILD","training","on","them",",","and","writes","partitions","back","to","disk","."]
code: def __init__(  # noqa
      self,
      config: ConfigSchema,
      model: Optional[MultiRelationEmbedder] = None,
      trainer: Optional[AbstractBatchProcessor] = None,
      evaluator: Optional[AbstractBatchProcessor] = None,
      rank: Rank = SINGLE_TRAINER,
      subprocess_init: Optional[Callable[[], None]] = None,
      stats_handler: StatsHandler = NOOP_STATS_HANDLER,
  ):
      """Each epoch/pass, for each partition pair, loads in embeddings and edgelist
      from disk, runs HOGWILD training on them, and writes partitions back to disk.
      """
      tag_logs_with_process_name(f"Trainer-{rank}")
      self.config = config
      if config.verbose > 0:
          import pprint
  
          pprint.PrettyPrinter().pprint(config.to_dict())
  
      logger.info("Loading entity counts...")
      entity_storage = ENTITY_STORAGES.make_instance(config.entity_path)
      entity_counts: Dict[str, List[int]] = {}
      for entity, econf in config.entities.items():
          entity_counts[entity] = []
          for part in range(econf.num_partitions):
              entity_counts[entity].append(
                  entity_storage.load_count(entity, part)
              )
  
      # Figure out how many lhs and rhs partitions we need
      holder = self.holder = EmbeddingHolder(config)
  
      logger.debug(
          f"nparts {holder.nparts_lhs} {holder.nparts_rhs} "
          f"types {holder.lhs_partitioned_types} {holder.rhs_partitioned_types}"
      )
  
      # We know ahead of time that we wil need 1-2 storages for each embedding type,
      # as well as the max size of this storage (num_entities x D).
      # We allocate these storages n advance in `embedding_storage_freelist`.
      # When we need storage for an entity type, we pop it from this free list,
      # and then add it back when we 'delete' the embedding table.
      embedding_storage_freelist: Dict[
          EntityName, Set[torch.FloatStorage]
      ] = defaultdict(set)
      for entity_type, counts in entity_counts.items():
          max_count = max(counts)
          if holder.nparts_lhs == 1 and holder.nparts_rhs == 1:
              num_sides = 1
          else:
              num_sides = (
                  (
                      1
                      if entity_type in holder.lhs_partitioned_types
                      else 0
                  )
                  + (
                      1
                      if entity_type in holder.rhs_partitioned_types
                      else 0
                  )
                  + (
                      1
                      if entity_type
                      in (
                          holder.lhs_unpartitioned_types
                          | holder.rhs_unpartitioned_types
                      )
                      else 0
                  )
              )
          for _ in range(num_sides):
              embedding_storage_freelist[entity_type].add(
                  allocate_shared_tensor(
                      (max_count, config.entity_dimension(entity_type)),
                      dtype=torch.float,
                  ).storage()
              )
  
      # create the handlers, threads, etc. for distributed training
      if config.num_machines > 1 or config.num_partition_servers > 0:
          if not 0 <= rank < config.num_machines:
              raise RuntimeError("Invalid rank for trainer")
          if not td.is_available():
              raise RuntimeError(
                  "The installed PyTorch version doesn't provide "
                  "distributed training capabilities."
              )
          ranks = ProcessRanks.from_num_invocations(
              config.num_machines, config.num_partition_servers
          )
  
          if config.num_partition_servers > config.num_machines:
              num_ps_groups = config.num_partition_servers * (
                  config.num_groups_per_sharded_partition_server + 1
              )
          else:
              num_ps_groups = config.num_groups_for_partition_server
          groups: List[List[int]] = [ranks.trainers]  # barrier group
          groups += [
              ranks.trainers + ranks.partition_servers
          ] * num_ps_groups  # ps groups
          group_idxs_for_partition_servers = range(1, len(groups))
  
          if rank == SINGLE_TRAINER:
              logger.info("Setup lock server...")
              start_server(
                  LockServer(
                      num_clients=len(ranks.trainers),
                      nparts_lhs=holder.nparts_lhs,
                      nparts_rhs=holder.nparts_rhs,
                      entities_lhs=holder.lhs_partitioned_types,
                      entities_rhs=holder.rhs_partitioned_types,
                      entity_counts=entity_counts,
                      init_tree=config.distributed_tree_init_order,
                      stats_handler=stats_handler,
                  ),
                  process_name="LockServer",
                  init_method=config.distributed_init_method,
                  world_size=ranks.world_size,
                  server_rank=ranks.lock_server,
                  groups=groups,
                  subprocess_init=subprocess_init,
              )
  
          self.bucket_scheduler = DistributedBucketScheduler(
              server_rank=ranks.lock_server,
              client_rank=ranks.trainers[rank],
          )
  
          logger.info("Setup param server...")
          start_server(
              ParameterServer(num_clients=len(ranks.trainers)),
              process_name=f"ParamS-{rank}",
              init_method=config.distributed_init_method,
              world_size=ranks.world_size,
              server_rank=ranks.parameter_servers[rank],
              groups=groups,
              subprocess_init=subprocess_init,
          )
  
          parameter_sharer = ParameterSharer(
              process_name=f"ParamC-{rank}",
              client_rank=ranks.parameter_clients[rank],
              all_server_ranks=ranks.parameter_servers,
              init_method=config.distributed_init_method,
              world_size=ranks.world_size,
              groups=groups,
              subprocess_init=subprocess_init,
          )
  
          if config.num_partition_servers == -1:
              start_server(
                  ParameterServer(
                      num_clients=len(ranks.trainers),
                      group_idxs=group_idxs_for_partition_servers,
                      log_stats=True,
                  ),
                  process_name=f"PartS-{rank}",
                  init_method=config.distributed_init_method,
                  world_size=ranks.world_size,
                  server_rank=ranks.partition_servers[rank],
                  groups=groups,
                  subprocess_init=subprocess_init,
              )
  
          groups = init_process_group(
              rank=ranks.trainers[rank],
              world_size=ranks.world_size,
              init_method=config.distributed_init_method,
              groups=groups,
          )
          trainer_group, *groups_for_partition_servers = groups
          self.barrier_group = trainer_group
  
          if len(ranks.partition_servers) > 0:
              if config.num_partition_servers > config.num_machines:
                  partition_client = ShardedPartitionClient(
                      ranks.partition_servers,
                      config.num_groups_per_sharded_partition_server,
                      config.partition_shard_size,
                      config.entities,
                      groups=groups_for_partition_servers,
                      log_stats=True,
                  )
              else:
                  partition_client = PartitionClient(
                      ranks.partition_servers,
                      groups=groups_for_partition_servers,
                      log_stats=True,
                  )
          else:
              partition_client = None
      else:
          self.barrier_group = None
          self.bucket_scheduler = SingleMachineBucketScheduler(
              holder.nparts_lhs,
              holder.nparts_rhs,
              config.bucket_order,
              stats_handler,
          )
          parameter_sharer = None
          partition_client = None
          hide_distributed_logging()
  
      # fork early for HOGWILD threads
      logger.info("Creating workers...")
      self.num_workers = get_num_workers(config.workers)
      self.pool = create_pool(
          self.num_workers,
          subprocess_name=f"TWorker-{rank}",
          subprocess_init=subprocess_init,
      )
  
      checkpoint_manager = CheckpointManager(
          config.checkpoint_path,
          rank=rank,
          num_machines=config.num_machines,
          partition_client=partition_client,
          subprocess_name=f"BackgRW-{rank}",
          subprocess_init=subprocess_init,
      )
      self.checkpoint_manager = checkpoint_manager
      checkpoint_manager.register_metadata_provider(
          ConfigMetadataProvider(config)
      )
      if rank == 0:
          checkpoint_manager.write_config(config)
  
      num_edge_chunks = get_num_edge_chunks(config)
  
      self.iteration_manager = IterationManager(
          config.num_epochs,
          config.edge_paths,
          num_edge_chunks,
          iteration_idx=checkpoint_manager.checkpoint_version,
      )
      checkpoint_manager.register_metadata_provider(
          self.iteration_manager
      )
  
      logger.info("Initializing global model...")
      if model is None:
          model = make_model(config)
      model.share_memory()
      loss_fn = LOSS_FUNCTIONS.get_class(config.loss_fn)(
          margin=config.margin
      )
      relation_weights = [
          relation.weight for relation in config.relations
      ]
      if trainer is None:
          trainer = Trainer(
              model_optimizer=make_optimizer(
                  config, model.parameters(), False
              ),
              loss_fn=loss_fn,
              relation_weights=relation_weights,
          )
      if evaluator is None:
          eval_overrides = {}
          if config.eval_num_batch_negs is not None:
              eval_overrides[
                  "num_batch_negs"
              ] = config.eval_num_batch_negs
          if config.eval_num_uniform_negs is not None:
              eval_overrides[
                  "num_uniform_negs"
              ] = config.eval_num_uniform_negs
  
          evaluator = RankingEvaluator(
              loss_fn=loss_fn,
              relation_weights=relation_weights,
              overrides=eval_overrides,
          )
  
      if config.init_path is not None:
          self.loadpath_manager = CheckpointManager(config.init_path)
      else:
          self.loadpath_manager = None
  
      # load model from checkpoint or loadpath, if available
      state_dict, optim_state = checkpoint_manager.maybe_read_model()
      if state_dict is None and self.loadpath_manager is not None:
          (
              state_dict,
              optim_state,
          ) = self.loadpath_manager.maybe_read_model()
      if state_dict is not None:
          model.load_state_dict(state_dict, strict=False)
      if optim_state is not None:
          trainer.model_optimizer.load_state_dict(optim_state)
  
      logger.debug("Loading unpartitioned entities...")
      for entity in (
          holder.lhs_unpartitioned_types
          | holder.rhs_unpartitioned_types
      ):
          count = entity_counts[entity][0]
          s = embedding_storage_freelist[entity].pop()
          dimension = config.entity_dimension(entity)
          embs = torch.FloatTensor(s).view(-1, dimension)[:count]
          embs, optimizer = self._load_embeddings(
              entity, UNPARTITIONED, out=embs
          )
          holder.unpartitioned_embeddings[entity] = embs
          trainer.unpartitioned_optimizers[entity] = optimizer
  
      # start communicating shared parameters with the parameter server
      if parameter_sharer is not None:
          shared_parameters: Set[int] = set()
          for name, param in model.named_parameters():
              if id(param) in shared_parameters:
                  continue
              shared_parameters.add(id(param))
              key = f"model.{name}"
              logger.info(
                  f"Adding {key} ({param.numel()} params) to parameter server"
              )
              parameter_sharer.set_param(key, param.data)
          for entity, embs in holder.unpartitioned_embeddings.items():
              key = f"entity.{entity}"
              logger.info(
                  f"Adding {key} ({embs.numel()} params) to parameter server"
              )
              parameter_sharer.set_param(key, embs.data)
  
      # store everything in self
      self.model = model
      self.trainer = trainer
      self.evaluator = evaluator
      self.rank = rank
      self.entity_counts = entity_counts
      self.embedding_storage_freelist = embedding_storage_freelist
      self.stats_handler = stats_handler
  
      self.strict = False
  
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 1 to create a new TrainingCoordinator object  with example: obj = TrainingCoordinator(config, model, trainer, evaluator, rank, subprocess_init, stats_handler)