Add enumerator to logging stats (#3073)

Summary:
Pull Request resolved: #3073

Pull Request resolved: #3018

Add enumerator to the stats logging. We will use this in the next diff to generate a hash for the planner input context.

Reviewed By: iamzainhuda

Differential Revision: D76278097

fbshipit-source-id: 9146284394c061b366d22742cc68f500fbeacf21

Author: aporialiao

torchrec/distributed/planner/stats.py

@@ -37,6 +37,7 @@
 )
 from torchrec.distributed.planner.types import (
     CriticalPathEstimate,
+    Enumerator,
     ParameterConstraints,
     Perf,
     ShardingOption,
@@ -157,6 +158,7 @@ def log(
         best_plan: List[ShardingOption],
         constraints: Optional[Dict[str, ParameterConstraints]] = None,
         sharders: Optional[List[ModuleSharder[nn.Module]]] = None,
+        enumerator: Optional[Enumerator] = None,
         debug: bool = True,
     ) -> None:
         """
@@ -1133,6 +1135,7 @@ def log(
         best_plan: List[ShardingOption],
         constraints: Optional[Dict[str, ParameterConstraints]] = None,
         sharders: Optional[List[ModuleSharder[nn.Module]]] = None,
+        enumerator: Optional[Enumerator] = None,
         debug: bool = True,
     ) -> None:
         pass

torchrec/distributed/planner/types.py

@@ -951,6 +951,7 @@ def log(
         best_plan: List[ShardingOption],
         constraints: Optional[Dict[str, ParameterConstraints]] = None,
         sharders: Optional[List[ModuleSharder[nn.Module]]] = None,
+        enumerator: Optional[Enumerator] = None,
         debug: bool = False,
     ) -> None:
         """