Add Python enum configs and KJT builder for enrichment (#5464)

Summary:

X-link: meta-pytorch/torchrec#3850

X-link: facebookresearch/FBGEMM#2439

CONTEXT: The enrichment configuration in EnrichmentPolicy uses raw strings for enrichment_type and response_format, which is error-prone and lacks type safety. Additionally, there is no utility to extract unhashed IDs from KJT features for enrichment queries.

WHAT: Strengthen enrichment configuration with Python enums and add a KJT builder utility.
- Added EnrichmentType enum (IGR_LASER_EMBEDDING, IGR_LASER_SID) and EnrichmentResponseFormat enum (JSON, THRIFT_FLOAT, THRIFT_INT64) in split_table_batched_embeddings_ops_common.py
- Updated EnrichmentPolicy to use enum types instead of strings
- Added enrichment_policy field to KVZCHTBEConfig for config propagation
- Convert enum values to int when passing to C++ TorchScript layer in training.py
- Added build_embedding_cache_write_kjt() in kvzch_utils.py to extract hashed/unhashed feature pairs from KJT and encode unhashed IDs as float32 weights for enrichment queries
- Wired enrichment_policy through batched_embedding_kernel.py to KVZCHParams

Differential Revision: D95883280

Author: EddyLXJ

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_common.py

@@ -233,17 +233,28 @@ def validate(self) -> None:
             )
 
 
+class EnrichmentType(enum.IntEnum):
+    IGR_LASER_EMBEDDING = 0
+    IGR_LASER_SID = 1
+
+
+class EnrichmentResponseFormat(enum.IntEnum):
+    JSON = 0
+    THRIFT_FLOAT = 1
+    THRIFT_INT64 = 2
+
+
 class EnrichmentPolicy(NamedTuple):
-    # Model and method identifier, e.g. "igr_laser_embedding", "igr_laser_sid"
-    enrichment_type: str = ""
+    # Model and method identifier
+    enrichment_type: EnrichmentType = EnrichmentType.IGR_LASER_EMBEDDING
     # External provider name (e.g. Laser provider name)
     provider_name: str = ""
     # Client identifier for the external service
     client_id: str = ""
     # Dimension of data returned by the source
     enrichment_dim: int = 0
-    # Deserialization format: "json", "thrift_float", "thrift_int64"
-    response_format: str = "json"
+    # Deserialization format
+    response_format: EnrichmentResponseFormat = EnrichmentResponseFormat.JSON
 
 
 class KVZCHParams(NamedTuple):
@@ -294,6 +305,8 @@ class KVZCHTBEConfig(NamedTuple):
     optimizer_type_for_st: Optional[str] = None
     # [DO NOT USE] This is for st publish only, do not set it in your config
     optimizer_state_dtypes_for_st: Optional[FrozenSet[Tuple[str, int]]] = None
+    # Enrichment policy for embedding cache enrichment from external sources (e.g. Laser)
+    enrichment_policy: Optional[EnrichmentPolicy] = None
 
 
 class BackendType(enum.IntEnum):

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -773,15 +773,15 @@ def __init__(
             if (
                 self.kv_zch_params
                 and self.kv_zch_params.enrichment_policy is not None
-                and self.kv_zch_params.enrichment_policy.enrichment_type
+                and self.kv_zch_params.enrichment_policy.enrichment_type is not None
             ):
                 ep = self.kv_zch_params.enrichment_policy
                 enrichment_config = torch.classes.fbgemm.EnrichmentConfig(
-                    ep.enrichment_type,
+                    ep.enrichment_type.value,
                     ep.provider_name,
                     ep.client_id,
                     ep.enrichment_dim,
-                    ep.response_format,
+                    ep.response_format.value,
                 )
             self._ssd_db = torch.classes.fbgemm.DramKVEmbeddingCacheWrapper(
                 self.cache_row_dim,