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","`"]
docstr: Call `callback` on each worker with the `ctx` and model local to that
  worker. e.g.
  def register_optimizer(ctx, model):
      args, kwargs = ctx
      model.optimizer = torch.optim.SGD(model.parameters(), *args, **kwargs)
  
  pipe_model = PipeRPCWrapper( ... )
  
  pipe_model.foreach_worker(
      register_optimizer,
      ([], {"lr" : 0.01, "momentum" : 0.9})
  )
docstr_tok: ["Call","`","callback","`","on","each","worker","with","the","`","ctx","`","and","model","local","to","that","worker",".","e.g",".","def","register_optimizer","(","ctx",",","model",")",":","args",",","kwargs","=","ctx","model.optimizer","=","torch.optim.SGD","(","model.parameters","(",")",",","*","args",",","*","*","kwargs",")","pipe_model","=","PipeRPCWrapper","(","...",")","pipe_model.foreach_worker","(","register_optimizer",",","(","[","]",",","{","``","lr","''",":","0.01",",","``","momentum","''",":","0.9","}",")",")"]
code: def foreach_worker(
      self,
      callback: Callable[[Any, AsyncPipe], None],
      ctx: Any = None,
      *,
      include_self: bool = False,
  ) -> None:
      """Call `callback` on each worker with the `ctx` and model local to that
      worker. e.g.
      def register_optimizer(ctx, model):
          args, kwargs = ctx
          model.optimizer = torch.optim.SGD(model.parameters(), *args, **kwargs)
  
      pipe_model = PipeRPCWrapper( ... )
  
      pipe_model.foreach_worker(
          register_optimizer,
          ([], {"lr" : 0.01, "momentum" : 0.9})
      )
      """
  
      self._foreach_worker(callback_with_model, args=(callback, ctx))
  
      if include_self:
          with self.model.lock:
              callback(ctx, self.model)
  
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 4 for calling the PipeRPCWrapper obj's underlying member method code with example usage: obj.foreach_worker(callback, ctx) without return types