Add RawIdTrackerWrapper within TBE to access tracked ids and raw ids (meta-pytorch#3506)

Summary:

This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP.

internal
This is needed to support MPZCH modules for Raw embedding streaming.
Mode details :
https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld

Reviewed By: chouxi

Differential Revision: D84925177

Author: aliafzal

torchrec/distributed/batched_embedding_kernel.py

@@ -17,6 +17,7 @@
 from math import sqrt
 from typing import (
     Any,
+    Callable,
     cast,
     Dict,
     Generic,
@@ -70,6 +71,7 @@
     GroupedEmbeddingConfig,
     ShardedEmbeddingTable,
 )
+from torchrec.distributed.model_tracker.types import IndexedLookup
 from torchrec.distributed.shards_wrapper import LocalShardsWrapper
 from torchrec.distributed.types import (
     Shard,
@@ -80,6 +82,7 @@
     TensorProperties,
 )
 from torchrec.distributed.utils import append_prefix, none_throws
+
 from torchrec.modules.embedding_configs import (
     CountBasedEvictionPolicy,
     CountTimestampMixedEvictionPolicy,
@@ -97,13 +100,30 @@
 )
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 RES_ENABLED_TABLES_STR = "res_enabled_tables"
 RES_STORE_SHARDS_STR = "res_store_shards"
 ENABLE_RAW_EMBEDDING_STREAMING_STR = "enable_raw_embedding_streaming"
 
 
+class RawIdTrackerWrapper:
+    def __init__(
+        self,
+        get_indexed_lookups: Callable[
+            [List[str], Optional[str]],
+            Dict[str, List[torch.Tensor]],
+        ],
+        delete: Callable[
+            [int],
+            None,
+        ],
+    ) -> None:
+        self.get_indexed_lookups = get_indexed_lookups
+        self.delete = delete
+
+
 def _populate_res_params(config: GroupedEmbeddingConfig) -> Tuple[bool, RESParams]:
     # populate res_params, which is used for raw embedding streaming
     # here only populates the params available in fused_params and TBE configs
@@ -2526,6 +2546,7 @@ def __init__(
         self._lengths_per_emb: List[int] = []
         self.table_name_to_count: Dict[str, int] = {}
         self._param_per_table: Dict[str, TableBatchedEmbeddingSlice] = {}
+        self._raw_id_tracker_wrapper: Optional[RawIdTrackerWrapper] = None
 
         for idx, table_config in enumerate(self._config.embedding_tables):
             self._local_rows.append(table_config.local_rows)
@@ -2668,6 +2689,22 @@ def named_parameters_by_table(
         for name, param in self._param_per_table.items():
             yield name, param
 
+    def init_raw_id_tracker(
+        self,
+        get_indexed_lookups: Callable[
+            [List[str], Optional[str]],
+            Dict[str, List[torch.Tensor]],
+        ],
+        delete: Callable[
+            [int],
+            None,
+        ],
+    ) -> None:
+        if isinstance(self._emb_module, SplitTableBatchedEmbeddingBagsCodegen):
+            self._raw_id_tracker_wrapper = RawIdTrackerWrapper(
+                get_indexed_lookups, delete
+            )
+
 
 class KeyValueEmbeddingBag(BaseBatchedEmbeddingBag[torch.Tensor], FusedOptimizerModule):
     def __init__(