feat: add planning result logging (meta-pytorch#4022)

Summary:


Log planning SUCCEEDED/FAILED to `training_optimization_events` Scuba table from both `EmbeddingShardingPlanner` and `LinearProgrammingPlanner`.

Differential Revision: D97581563

Author: Ahmed Shuaibi

torchrec/distributed/logging_handlers.py

@@ -207,4 +207,14 @@ def detect_technique(items: List) -> OptimizationTechnique:  # type: ignore[type
     return OptimizationTechnique.NONE
 
 
+def log_planning_result(
+    planner_type: str,
+    technique: OptimizationTechnique = OptimizationTechnique.NONE,
+    error_message: Optional[str] = None,
+    **extra_metadata: str,
+) -> None:
+    """No-op OSS stub."""
+    pass
+
+
 _log_handlers: dict[str, logging.Handler] = defaultdict(logging.NullHandler)

torchrec/distributed/planner/planners.py

@@ -93,6 +93,9 @@ def decorator(func):
         return decorator
 
 
+from torchrec.distributed.logging_handlers import log_planning_result
+
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 
@@ -720,6 +723,13 @@ def plan(
                 )
 
             validate_rank_assignment(sharding_plan, self._topology)
+
+            log_planning_result(
+                planner_type=self.__class__.__name__,
+                num_proposals=str(self._num_proposals),
+                num_plans=str(self._num_plans),
+            )
+
             return sharding_plan
         else:
             global_storage_capacity = reduce(
@@ -779,6 +789,13 @@ def plan(
                     debug=self._debug,
                 )
 
+            log_planning_result(
+                planner_type=self.__class__.__name__,
+                error_message=str(last_planner_error),
+                num_proposals=str(self._num_proposals),
+                num_plans=str(self._num_plans),
+            )
+
             if not lowest_storage.fits_in(global_storage_constraints):
                 raise PlannerError(
                     error_type=PlannerErrorType.INSUFFICIENT_STORAGE,