class_docstr: A distributed trainer for `AbstractComponent`s. Used mainly for supervised learning.
  
  Create a `DistributedModelTrainer` by passing a `AbstractNeuralModel` in the constructor.
  Invoke `train()` to initiate the training loop. The root `TNeuralModule` should return a scalar loss.
code: def _parallel_training_process(
      self,
      rank,
      world_size,
      training_data,
      validation_data,
      parallelize,
      patience,
      shuffle_training_data,
      validate_on_start: bool,
      worker_init: Optional[
          Callable[["DistributedModelTrainer", int, int], None]
      ] = None,
      start_epoch_idx: int = 0,
  ):
      assert (
          torch.cuda.is_available()
      ), "No CUDA available. Aborting training."
  
      os.environ["MASTER_ADDR"] = "localhost"
      os.environ["MASTER_PORT"] = "12355"
      dist.init_process_group(
          backend="nccl", rank=rank, world_size=world_size
      )
  
      if worker_init is not None:
          worker_init(self, rank, world_size)
  
      self.LOGGER.info(
          f"[{os.getpid()}]: world_size = {dist.get_world_size()}, "
          + f"rank = {dist.get_rank()}, backend={dist.get_backend()}"
      )
      device = dist.get_rank()
  
      self.neural_module.to(device)
      distributed_neural_module = DDP(
          self.neural_module, device_ids=[device]
      )
  
      training_data.set_rank(dist.get_rank(), dist.get_world_size())
      training_tensors = partial(
          self.model.tensorize_dataset,
          training_data,
          parallelize=parallelize,
          use_multiprocessing=True,
      )
      validation_data.set_rank(dist.get_rank(), dist.get_world_size())
      validation_tensors = partial(
          self.model.tensorize_dataset,
          validation_data,
          parallelize=parallelize,
          use_multiprocessing=True,
      )
  
      optimizer = self._create_optimizer(
          distributed_neural_module.parameters()
      )
  
      scheduler = (
          None
          if self._create_scheduler is None
          else self._create_scheduler(optimizer)
      )
  
      for hook in self._training_start_hooks:
          hook(self.model, distributed_neural_module.module, optimizer)
  
      if (
          self._target_metric_higher_is_better
          and self._target_metric is not None
      ):
          best_target_metric = -math.inf
      else:
          best_target_metric = math.inf
      if validate_on_start:
          target_metric, improved, _ = self._run_validation(
              distributed_neural_module,
              validation_tensors,
              start_epoch_idx,
              best_target_metric,
              device,
              parallelize,
          )
          assert improved
          self.LOGGER.info(
              f"Initial {self._target_metric or 'Loss'}: {target_metric}"
          )
          best_target_metric = target_metric
  
      num_epochs_not_improved: int = 0
  
      for epoch in range(start_epoch_idx, self._max_num_epochs):
          try:
              self._run_training(
                  distributed_neural_module,
                  training_tensors,
                  epoch,
                  device,
                  optimizer,
                  parallelize,
                  scheduler,
                  shuffle_training_data,
              )
          except Exception as e:
              self.LOGGER.exception("Error during training", exc_info=e)
              raise e
  
          # Save optimizer and epoch id for scheduler
          optimizer_state = optimizer.state_dict()
          if dist.get_rank() == 0:
              # Save only on rank 0
              torch.save(
                  {
                      "optimizer_state_dict": optimizer_state,
                      "epoch": epoch + 1,
                  },
                  self._checkpoint_location.with_suffix(
                      ".optimizerstate"
                  ),
              )
  
          (
              target_metric,
              target_metric_improved,
              validation_metrics,
          ) = self._run_validation(
              distributed_neural_module,
              validation_tensors,
              epoch,
              best_target_metric,
              device,
              parallelize,
          )
          if target_metric_improved:
              num_epochs_not_improved = 0
              if dist.get_rank() == 0:
                  self.LOGGER.info(
                      f"Best performance so far "
                      f"({self._target_metric or 'Loss'}: {target_metric:.3f} from {best_target_metric:.3f}). "
                      "Saving model checkpoint."
                  )
                  self._save_checkpoint()
  
              best_target_metric = target_metric
  
              for epoch_hook in self._improved_epoch_end_hooks:
                  epoch_hook(
                      self.model,
                      distributed_neural_module.module,
                      epoch,
                      validation_metrics,
                  )
          else:
              num_epochs_not_improved += 1
              if num_epochs_not_improved > patience:
                  self.LOGGER.warning(
                      f"The target metric has not improved for {num_epochs_not_improved} epochs . Stopping."
                  )
                  break
  
      dist.destroy_process_group()
  
imports: from typing_extensions import Final
  import json
  import logging
  import math
  import os
  import time
  import torch
  import torch.distributed
  import torch.multiprocessing
  from functools import partial
  from torch.nn.parallel import DistributedDataParallel
  from typing import Callable, Iterable, Optional, TypeVar
  from ptgnn.baseneuralmodel import ModelTrainer
  from ptgnn.baseneuralmodel.abstractneuralmodel import AbstractNeuralModel
  from ptgnn.baseneuralmodel.modulewithmetrics import ModuleWithMetrics
  from ptgnn.baseneuralmodel.utils.data import ShardedLazyDataIterable
label_desc: Use image node_id 5 for calling the DistributedModelTrainer obj's underlying member method code with example usage: obj._parallel_training_process(rank, world_size, training_data, validation_data, parallelize, patience, shuffle_training_data, validate_on_start, worker_init, start_epoch_idx) without return types