class_docstr: A wrapper for Pipe to control the entire pipeline from a single process.
  Typical usecase would have rank 0 construct `PipeRPCWrapper` and run the
  training loop as normal, and all other ranks would call
  `torch.distributed.rpc.shutdown()`
  
  To run code on each worker, e.g. to run the optimizer, use `foreach_worker`
class_docstr_tok: ["A","wrapper","for","Pipe","to","control","the","entire","pipeline","from","a","single","process",".","Typical","usecase","would","have","rank","0","construct","`","PipeRPCWrapper","`","and","run","the","training","loop","as","normal",",","and","all","other","ranks","would","call","`","torch.distributed.rpc.shutdown","(",")","`","To","run","code","on","each","worker",",","e.g",".","to","run","the","optimizer",",","use","`","foreach_worker","`"]
code: def __init__(self, *args: Any, **kwargs: Any):
      super().__init__()
      self.group = (
          cast(ProcessGroup, kwargs.get("group"))
          or get_pipeline_parallel_group()
      )
      assert self.group.rank() == 0
      self.lock = Lock()
  
      if True:
          assert (
              self.group == get_pipeline_parallel_group()
          ), "Can't pickle groups, so group must be `get_pipeline_parallel_group()`"
          kwargs["group"] = None
      else:
          kwargs["group"] = self.group
  
      kwargs["input_device"] = torch.device(
          "cuda", torch.cuda.current_device()
      )
  
      self.model = AsyncPipe(*args, **kwargs)
      self.worker_map = kwargs["worker_map"]
      self._foreach_worker(
          self._register_remote_model, args=(args, kwargs)
      )
      self.model.cuda()
  
imports: from threading import Event, Lock, Thread
  from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
  import torch
  from torch import nn
  from torch.distributed import ProcessGroup, rpc
  from torch.distributed.distributed_c10d import _get_global_rank
  from fairscale.nn.model_parallel.initialize import get_pipeline_parallel_group
  from .async_pipe import AsyncPipe
  from .types import EVENT_LOOP_QUEUE, PipeMessage, TensorOrTensors
label_desc: Use image node_id 1 to create a new PipeRPCWrapper object from inherited base classes: nn with example: obj = PipeRPCWrapper()