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.
docstr: The training-validation loop for `AbstractNeuralModel`s.
  
  :param training_data: An iterable that yields the training data for a given shard. Note
      that the data iterator should provide the `set_rank(rank, world_size)` that filters
      the data appropriately.
  :param validation_data: the validation set, as in `training_data`.
  :param validate_on_start: Whether to run a validation loop on start
  :param patience: The number of iterations before early stopping kicks in.
  :param initialize_metadata: If true, initialize the metadata from the training_data. Otherwise,
      assume that the model that is being trained has its metadata already initialized.
  :param parallelize: Bool indicating whether to run in parallel
  :param shuffle_training_data: shuffle the incoming data from `training_data`.
  :param start_epoch_idx: the idx of the first epoch in this training loop (used for resuming).
docstr_tok: ["The","training-validation","loop","for","`","AbstractNeuralModel","`","s",".",":","param","training_data",":","An","iterable","that","yields","the","training","data","for","a","given","shard",".","Note","that","the","data","iterator","should","provide","the","`","set_rank","(","rank",",","world_size",")","`","that","filters","the","data","appropriately",".",":","param","validation_data",":","the","validation","set",",","as","in","`","training_data","`",".",":","param","validate_on_start",":","Whether","to","run","a","validation","loop","on","start",":","param","patience",":","The","number","of","iterations","before","early","stopping","kicks","in",".",":","param","initialize_metadata",":","If","true",",","initialize","the","metadata","from","the","training_data",".","Otherwise",",","assume","that","the","model","that","is","being","trained","has","its","metadata","already","initialized",".",":","param","parallelize",":","Bool","indicating","whether","to","run","in","parallel",":","param","shuffle_training_data",":","shuffle","the","incoming","data","from","`","training_data","`",".",":","param","start_epoch_idx",":","the","idx","of","the","first","epoch","in","this","training","loop","(","used","for","resuming",")","."]
code: def distributed_train(
      self,
      world_size: int,
      training_data: ShardedLazyDataIterable[TRawDatapoint],
      validation_data: ShardedLazyDataIterable[TRawDatapoint],
      *,
      validate_on_start: bool = True,
      patience: int = 5,
      initialize_metadata: bool = True,
      parallelize: bool = True,
      shuffle_training_data: bool = True,
      worker_init: Optional[
          Callable[["DistributedModelTrainer", int, int], None]
      ] = None,
      start_epoch_idx: int = 0,
  ) -> None:
      """
      The training-validation loop for `AbstractNeuralModel`s.
  
      :param training_data: An iterable that yields the training data for a given shard. Note
          that the data iterator should provide the `set_rank(rank, world_size)` that filters
          the data appropriately.
      :param validation_data: the validation set, as in `training_data`.
      :param validate_on_start: Whether to run a validation loop on start
      :param patience: The number of iterations before early stopping kicks in.
      :param initialize_metadata: If true, initialize the metadata from the training_data. Otherwise,
          assume that the model that is being trained has its metadata already initialized.
      :param parallelize: Bool indicating whether to run in parallel
      :param shuffle_training_data: shuffle the incoming data from `training_data`.
      :param start_epoch_idx: the idx of the first epoch in this training loop (used for resuming).
      """
      assert torch.distributed.is_available()
  
      if initialize_metadata:
          training_data.set_rank(
              0, 1
          )  # No sharding currently possible during metadata loading.
          self.load_metadata_and_create_network(
              training_data, parallelize, False
          )
  
      self.LOGGER.info(
          "Model has %s trainable parameters.",
          sum(
              param.numel()
              for param in self.neural_module.parameters(recurse=True)
              if param.requires_grad
          ),
      )
  
      ### Distributed code starts here
      # Spawn processes here
      mp.spawn(
          self._parallel_training_process,
          args=(
              world_size,
              training_data,
              validation_data,
              parallelize,
              patience,
              shuffle_training_data,
              validate_on_start,
              worker_init,
              start_epoch_idx,
          ),
          nprocs=world_size,
          join=True,
      )
      self._restore_checkpoint()
  
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 4 for calling the DistributedModelTrainer obj's underlying member method code with example usage: obj.distributed_train(world_size, training_data, validation_data) without return types