docstr: IDTransformerGroup stores the IDTransformer for all sharded modules in a DMP module.
  
  Args:
      url: configuration for PS, e.g. redis://127.0.0.1:6379/?prefix=model.
      module: DMP module that need dynamic embedding.
      configs_dict: a dictionary that maps the module path of the sharded module to its embedding
          configs or embeddingbag configs. The plan of `module` should contain the module path
          in `configs_dict`.
      eviction_config: configuration for eviction policy. Default is `{"type": "mixed_lru_lfu"}`
      transform_config: configuration for the transformer. Default is `{"type": "naive"}`
      parallel: Whether the IDTransformerCollections will run parallel. When set to True,
          IDTransformerGroup will start a thread for each IDTransformerCollection.
  
  Example:
      class Model(nn.Module):
          def __init__(self, config1, config2):
              super().__init__()
              self.emb1 = EmbeddingCollection(tables=config1, device=torch.device("meta"))
              self.emb2 = EmbeddingCollection(tables=config2, device=torch.device("meta"))
              ...
  
          def forward(self, kjt1, kjt2):
              ...
  
      m = Model(config1, config2)
      m = DistributedModelParallel(m)
      transformers = IDTransformerGroup(
          "redis://127.0.0.1:6379/?prefix=model",
          m,
          { "emb1": config1, "emb2": config2 })
  
      for label, kjt1, kjt2 in dataset:
          kjts = transformers.transform({ "emb1": kjt1, "emb2": kjt2 })
          kjt1, kjt2 = kjts["emb1"], kjts["emb2"]
          output = m(kjt1, kjt2)
          ...
docstr_tok: ["IDTransformerGroup","stores","the","IDTransformer","for","all","sharded","modules","in","a","DMP","module",".","Args",":","url",":","configuration","for","PS",",","e.g",".","redis",":","\/\/127.0.0.1:6379\/","?","prefix=model",".","module",":","DMP","module","that","need","dynamic","embedding",".","configs_dict",":","a","dictionary","that","maps","the","module","path","of","the","sharded","module","to","its","embedding","configs","or","embeddingbag","configs",".","The","plan","of","`","module","`","should","contain","the","module","path","in","`","configs_dict","`",".","eviction_config",":","configuration","for","eviction","policy",".","Default","is","`","{","``","type","''",":","``","mixed_lru_lfu","''","}","`","transform_config",":","configuration","for","the","transformer",".","Default","is","`","{","``","type","''",":","``","naive","''","}","`","parallel",":","Whether","the","IDTransformerCollections","will","run","parallel",".","When","set","to","True",",","IDTransformerGroup","will","start","a","thread","for","each","IDTransformerCollection",".","Example",":","class","Model","(","nn.Module",")",":","def","__init__","(","self",",","config1",",","config2",")",":","super","(",")",".__init__","(",")","self.emb1","=","EmbeddingCollection","(","tables=config1",",","device=torch.device","(","``","meta","''",")",")","self.emb2","=","EmbeddingCollection","(","tables=config2",",","device=torch.device","(","``","meta","''",")",")","...","def","forward","(","self",",","kjt1",",","kjt2",")",":","...","m","=","Model","(","config1",",","config2",")","m","=","DistributedModelParallel","(","m",")","transformers","=","IDTransformerGroup","(","``","redis",":","\/\/127.0.0.1:6379\/","?","prefix=model","''",",","m",",","{","``","emb1","''",":","config1",",","``","emb2","''",":","config2","}",")","for","label",",","kjt1",",","kjt2","in","dataset",":","kjts","=","transformers.transform","(","{","``","emb1","''",":","kjt1",",","``","emb2","''",":","kjt2","}",")","kjt1",",","kjt2","=","kjts","[","``","emb1","''","]",",","kjts","[","``","emb2","''","]","output","=","m","(","kjt1",",","kjt2",")","..."]
code: def __init__(
      self,
      url,
      module: DistributedModelParallel,
      configs_dict: Dict[
          str, Union[List[EmbeddingBagConfig], List[EmbeddingConfig]]
      ],
      *,
      eviction_config=None,
      transform_config=None,
      ps_config=None,
      parallel=True,
  ):
      """
      IDTransformerGroup stores the IDTransformer for all sharded modules in a DMP module.
  
      Args:
          url: configuration for PS, e.g. redis://127.0.0.1:6379/?prefix=model.
          module: DMP module that need dynamic embedding.
          configs_dict: a dictionary that maps the module path of the sharded module to its embedding
              configs or embeddingbag configs. The plan of `module` should contain the module path
              in `configs_dict`.
          eviction_config: configuration for eviction policy. Default is `{"type": "mixed_lru_lfu"}`
          transform_config: configuration for the transformer. Default is `{"type": "naive"}`
          parallel: Whether the IDTransformerCollections will run parallel. When set to True,
              IDTransformerGroup will start a thread for each IDTransformerCollection.
  
      Example:
          class Model(nn.Module):
              def __init__(self, config1, config2):
                  super().__init__()
                  self.emb1 = EmbeddingCollection(tables=config1, device=torch.device("meta"))
                  self.emb2 = EmbeddingCollection(tables=config2, device=torch.device("meta"))
                  ...
  
              def forward(self, kjt1, kjt2):
                  ...
  
          m = Model(config1, config2)
          m = DistributedModelParallel(m)
          transformers = IDTransformerGroup(
              "redis://127.0.0.1:6379/?prefix=model",
              m,
              { "emb1": config1, "emb2": config2 })
  
          for label, kjt1, kjt2 in dataset:
              kjts = transformers.transform({ "emb1": kjt1, "emb2": kjt2 })
              kjt1, kjt2 = kjts["emb1"], kjts["emb2"]
              output = m(kjt1, kjt2)
              ...
      """
      self._parallel = parallel
  
      # get all sharded_modules from plan
      plan = module.plan
      sharded_modules = _get_sharded_modules_recursive(
          module.module, "", plan
      )
  
      self._id_transformer_collections: Dict[
          str, IDTransformerCollection
      ] = {}
      for path, configs in configs_dict.items():
          if path not in sharded_modules:
              raise ValueError(
                  f"`{path}` in configs dooes not match any sharded modules. "
                  f"Paths for current sharded modules are: {list(sharded_modules.keys())}."
              )
          sharded_module, params_plan = sharded_modules[path]
          ps_collection = PSCollection.fromModule(
              path, sharded_module, params_plan, url, ps_config
          )
          id_transformer_collection = IDTransformerCollection(
              configs, eviction_config, transform_config, ps_collection
          )
          self._id_transformer_collections[
              path
          ] = id_transformer_collection
  
      if self._parallel:
          self._threads = {}
          self._input_queues = {}
          self._output_queues = {}
          for (
              path,
              transformer,
          ) in self._id_transformer_collections.items():
              (
                  thread,
                  input_queue,
                  output_queue,
              ) = _create_transformer_thread(transformer)
              self._threads[path] = thread
              self._input_queues[path] = input_queue
              self._output_queues[path] = output_queue
  
file_path: torchrec/contrib/dynamic_embedding/src/torchrec_dynamic_embedding/id_transformer_group.py
imports: import queue
  import threading
  from typing import Dict, List, Union
  from torchrec import EmbeddingBagConfig, EmbeddingConfig, KeyedJaggedTensor
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from .id_transformer_collection import IDTransformerCollection
  from .ps import PSCollection
  from .utils import _get_sharded_modules_recursive
label_desc: Use image node_id 1 to create a new IDTransformerGroup object  with example: obj = IDTransformerGroup(url, module, configs_dict)