Adding util func to get inference threshold for feature score eviction. (meta-pytorch#3321)

Summary:

Since the eviction threshold for  feature score eviction is dynamic changed, we have to recalculate the threshold in publish to evict based on max inference id pre rank to avoid inference oom.

Reviewed By: emlin

Differential Revision: D80880936

Author: EddyLXJ

torchrec/modules/embedding_configs.py

@@ -222,6 +222,9 @@ class FeatureScoreBasedEvictionPolicy(VirtualTableEvictionPolicy):
     eviction_ttl_mins: int = (
         0  # if not 0, means we will use timestamp based policy but not feature score policy
     )
+    max_inference_id_num_per_rank: int = (
+        0  # max number of inference ids per rank, default is max_training_id_num_per_rank
+    )
     inference_eviction_feature_score_threshold: Optional[float] = (
         None  # 0 means no eviction
     )
@@ -232,6 +235,8 @@ def __post_init__(self) -> None:
             self.inference_eviction_feature_score_threshold = 0
         if self.inference_eviction_ttl_mins is None:
             self.inference_eviction_ttl_mins = self.eviction_ttl_mins
+        if self.max_inference_id_num_per_rank == 0:
+            self.max_inference_id_num_per_rank = self.max_training_id_num_per_rank
 
 
 @dataclass