docstr: Override local_size after pg construction because unit test device count is
  larger than local_size setup. This can be problematic for twrw because we have
  ShardedTensor placement check.
  
  TODO (T108556130) Mock out functions in comm.py instead of overriding env vars
docstr_tok: ["Override","local_size","after","pg","construction","because","unit","test","device","count","is","larger","than","local_size","setup",".","This","can","be","problematic","for","twrw","because","we","have","ShardedTensor","placement","check",".","TODO","(","T108556130",")","Mock","out","functions","in","comm.py","instead","of","overriding","env","vars"]
code: def __enter__(self):
      """
      Override local_size after pg construction because unit test device count is
      larger than local_size setup. This can be problematic for twrw because we have
      ShardedTensor placement check.
  
      TODO (T108556130) Mock out functions in comm.py instead of overriding env vars
      """
  
      os.environ["LOCAL_WORLD_SIZE"] = str(
          self.local_size or self.world_size
      )
      if self.local_size is not None:
          os.environ["LOCAL_RANK"] = str(self.rank % self.local_size)
  
      self.pg = init_distributed_single_host(
          rank=self.rank,
          world_size=self.world_size,
          backend=self.backend,
          local_size=self.local_size,
      )
      return self
  
file_path: torchrec/torchrec/distributed/test_utils/multi_process.py
imports: import multiprocessing
  import os
  import unittest
  from typing import Any, Callable, Dict, List, Optional
  import torch
  import torch.distributed
  from torchrec.distributed.comm import _CROSS_PG, _INTRA_PG
  from torchrec.test_utils import get_free_port, init_distributed_single_host, seed_and_log
label_desc: Use image node_id 2 for calling the MultiProcessContext obj's underlying member method code with example usage: obj.__enter__() and returns: self