Cleaner MetricModule interface (#3554)

Summary:

Utilities for CPUOffloadedRecMetricModule and RecMetricModule.

Also raise exceptions in the main thread if any of the background threads. Added unit tests.

Simplify the core metric types:
- MetricsResult = Dict[str, MetricValue]: sync metrics computation
- MetricsFuture = concurrent.futures.Future[MetricsResult]: for async computation
- MetricsOutput = Union[MetricsResult, MetricsFuture]: Either a MetricsResult, or a MetricsFuture

Introduce a metrics_output_util to handle the logic between futures and dicts. Users can schedule callbacks via
`on_metrics_ready()`

Introduce `device` argument to RecMetricModule constructor. It is a noop for the standard metric module, but CPUOffloadedRecMetricModule requires it to determine whether to perform GPU to CPU transfers.

Differential Revision: D87110900

Author: jeffkbkim

torchrec/metrics/cpu_comms_metric_module.py

@@ -9,6 +9,8 @@
 import logging
 from typing import Any, cast, Dict
 
+import torch
+
 from torch import nn
 
 from torch.profiler import record_function
@@ -48,7 +50,8 @@ def __init__(
         """
         All arguments are the same as RecMetricModule
         """
-
+        # Ensure device is set to CPU
+        kwargs["device"] = torch.device("cpu")
         super().__init__(*args, **kwargs)
 
         rec_metrics_clone = self._clone_rec_metrics()
@@ -106,9 +109,6 @@ def _load_metric_states(
         Uses aggregated states.
         """
 
-        # All update() calls were done prior. Clear previous computed state.
-        # Otherwise, we get warnings that compute() was called before
-        # update() which is not the case.
         computation = cast(RecMetricComputation, computation)
         set_update_called(computation)
         computation._computed = None
@@ -157,8 +157,9 @@ def _clone_rec_metrics(self) -> RecMetricList:
 
 def set_update_called(computation: RecMetricComputation) -> None:
     """
-    Set _update_called to True for RecMetricComputation.
-    This is a workaround for torchmetrics 1.0.3+.
+    All update() calls were done prior. Clear previous computed state.
+    Otherwise, we get warnings that compute() was called before
+    update() which is not the case.
     """
     try:
         computation._update_called = True

torchrec/metrics/cpu_offloaded_metric_module.py

@@ -22,7 +22,7 @@
     MetricUpdateJob,
     SynchronizationMarker,
 )
-from torchrec.metrics.metric_module import MetricValue, RecMetricModule
+from torchrec.metrics.metric_module import MetricsFuture, MetricsResult, RecMetricModule
 from torchrec.metrics.metric_state_snapshot import MetricStateSnapshot
 from torchrec.metrics.model_utils import parse_task_model_outputs
 from torchrec.metrics.rec_metric import RecMetricException
@@ -74,7 +74,9 @@ def __init__(
             - compute_queue_size: Maximum size of the update queue. Default is 100.
         """
         super().__init__(*args, **kwargs)
-        self._shutdown_event = threading.Event()
+        self._shutdown_event: threading.Event = threading.Event()
+        self._captured_exception_event: threading.Event = threading.Event()
+        self._captured_exception: Optional[Exception] = None
 
         self.update_queue: queue.Queue[
             Union[MetricUpdateJob, SynchronizationMarker]
@@ -132,8 +134,16 @@ def _update_rec_metrics(
         if self._shutdown_event.is_set():
             raise RecMetricException("metric processor thread is shut down.")
 
+        if self._captured_exception_event.is_set():
+            assert self._captured_exception is not None
+            raise self._captured_exception
+
         try:
-            cpu_model_out, transfer_completed_event = self._transfer_to_cpu(model_out)
+            cpu_model_out, transfer_completed_event = (
+                self._transfer_to_cpu(model_out)
+                if self.device == torch.device("cuda")
+                else (model_out, None)
+            )
             self.update_queue.put_nowait(
                 MetricUpdateJob(
                     model_out=cpu_model_out,
@@ -191,31 +201,25 @@ def _process_metric_update_job(self, metric_update_job: MetricUpdateJob) -> None
         """
 
         with record_function("## CPUOffloadedRecMetricModule:update ##"):
-            try:
+            if metric_update_job.transfer_completed_event is not None:
                 metric_update_job.transfer_completed_event.synchronize()
-                labels, predictions, weights, required_inputs = (
-                    parse_task_model_outputs(
-                        self.rec_tasks,
-                        metric_update_job.model_out,
-                        self.get_required_inputs(),
-                    )
-                )
-                if required_inputs:
-                    metric_update_job.kwargs["required_inputs"] = required_inputs
-
-                self.rec_metrics.update(
-                    predictions=predictions,
-                    labels=labels,
-                    weights=weights,
-                    **metric_update_job.kwargs,
-                )
-
-                if self.throughput_metric:
-                    self.throughput_metric.update()
+            labels, predictions, weights, required_inputs = parse_task_model_outputs(
+                self.rec_tasks,
+                metric_update_job.model_out,
+                self.get_required_inputs(),
+            )
+            if required_inputs:
+                metric_update_job.kwargs["required_inputs"] = required_inputs
+
+            self.rec_metrics.update(
+                predictions=predictions,
+                labels=labels,
+                weights=weights,
+                **metric_update_job.kwargs,
+            )
 
-            except Exception as e:
-                logger.exception("Error processing metric update: %s", e)
-                raise e
+            if self.throughput_metric:
+                self.throughput_metric.update()
 
     @override
     def shutdown(self) -> None:
@@ -248,30 +252,34 @@ def shutdown(self) -> None:
         logger.info("CPUOffloadedRecMetricModule has been successfully shutdown.")
 
     @override
-    def compute(self) -> Dict[str, MetricValue]:
+    def compute(self) -> MetricsResult:
         raise RecMetricException(
-            "compute() is not supported in CPUOffloadedRecMetricModule. Use async_compute() instead."
+            "CPUOffloadedRecMetricModule does not support compute(). Use async_compute() instead."
         )
 
     @override
-    def async_compute(
-        self, future: concurrent.futures.Future[Dict[str, MetricValue]]
-    ) -> None:
+    def async_compute(self) -> MetricsFuture:
         """
         Entry point for asynchronous metric compute. It enqueues a synchronization marker
         to the update queue.
 
-        Args:
+        Returns:
             future: Pre-created future where the computed metrics will be set.
         """
+        metrics_future = concurrent.futures.Future()
         if self._shutdown_event.is_set():
-            future.set_exception(
+            metrics_future.set_exception(
                 RecMetricException("metric processor thread is shut down.")
             )
-            return
+            return metrics_future
+
+        if self._captured_exception_event.is_set():
+            assert self._captured_exception is not None
+            raise self._captured_exception
 
-        self.update_queue.put_nowait(SynchronizationMarker(future))
+        self.update_queue.put_nowait(SynchronizationMarker(metrics_future))
         self.update_queue_size_logger.add(self.update_queue.qsize())
+        return metrics_future
 
     def _process_synchronization_marker(
         self, synchronization_marker: SynchronizationMarker
@@ -304,7 +312,7 @@ def _process_synchronization_marker(
 
     def _process_metric_compute_job(
         self, metric_compute_job: MetricComputeJob
-    ) -> Dict[str, MetricValue]:
+    ) -> MetricsResult:
         """
         Process a metric compute job:
         1. Comms module performs all gather
@@ -355,6 +363,8 @@ def _update_loop(self) -> None:
                 self._do_work(self.update_queue)
             except Exception as e:
                 logger.exception(f"Exception in update loop: {e}")
+                self._captured_exception_event.set()
+                self._captured_exception = e
                 raise e
 
         remaining = self._flush_remaining_work(self.update_queue)
@@ -372,6 +382,8 @@ def _compute_loop(self) -> None:
                 self._do_work(self.compute_queue)
             except Exception as e:
                 logger.exception(f"Exception in compute loop: {e}")
+                self._captured_exception_event.set()
+                self._captured_exception = e
                 raise e
 
         remaining = self._flush_remaining_work(self.compute_queue)

torchrec/metrics/metric_job_types.py

@@ -8,7 +8,7 @@
 # pyre-strict
 
 import concurrent
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 import torch
 from torchrec.metrics.metric_module import MetricValue
@@ -26,7 +26,7 @@ class MetricUpdateJob:
     def __init__(
         self,
         model_out: Dict[str, torch.Tensor],
-        transfer_completed_event: torch.cuda.Event,
+        transfer_completed_event: Optional[torch.cuda.Event],
         kwargs: Dict[str, Any],
     ) -> None:
         """
@@ -37,7 +37,9 @@ def __init__(
         """
 
         self.model_out: Dict[str, torch.Tensor] = model_out
-        self.transfer_completed_event: torch.cuda.Event = transfer_completed_event
+        self.transfer_completed_event: Optional[torch.cuda.Event] = (
+            transfer_completed_event
+        )
         self.kwargs: Dict[str, Any] = kwargs
 
 

torchrec/metrics/metric_module.py

@@ -117,6 +117,9 @@
 
 
 MetricValue = Union[torch.Tensor, float]
+MetricsResult = Dict[str, MetricValue]
+MetricsFuture = concurrent.futures.Future[MetricsResult]
+MetricsOutput = Union[MetricsResult, MetricsFuture]
 
 
 class StateMetric(abc.ABC):
@@ -125,7 +128,7 @@ class StateMetric(abc.ABC):
     """
 
     @abc.abstractmethod
-    def get_metrics(self) -> Dict[str, MetricValue]:
+    def get_metrics(self) -> MetricsResult:
         pass
 
 
@@ -189,6 +192,7 @@ def __init__(
         self,
         batch_size: int,
         world_size: int,
+        device: torch.device,
         rec_tasks: Optional[List[RecTaskInfo]] = None,
         rec_metrics: Optional[RecMetricList] = None,
         throughput_metric: Optional[ThroughputMetric] = None,
@@ -205,6 +209,7 @@ def __init__(
         self.trained_batches: int = 0
         self.batch_size = batch_size
         self.world_size = world_size
+        self.device = device
         self.oom_count = 0
         self.compute_count = 0
 
@@ -315,12 +320,12 @@ def _adjust_compute_interval(self) -> None:
     def should_compute(self) -> bool:
         return self.trained_batches % self.compute_interval_steps == 0
 
-    def compute(self) -> Dict[str, MetricValue]:
+    def compute(self) -> MetricsResult:
         r"""compute() is called when the global metrics are required, usually
         right before logging the metrics results to the data sink.
         """
         self.compute_count += 1
-        ret: Dict[str, MetricValue] = {}
+        ret: MetricsResult = {}
         with record_function("## RecMetricModule:compute ##"):
             if self.rec_metrics:
                 self._adjust_compute_interval()
@@ -337,11 +342,11 @@ def compute(self) -> Dict[str, MetricValue]:
                     )
         return ret
 
-    def local_compute(self) -> Dict[str, MetricValue]:
+    def local_compute(self) -> MetricsResult:
         r"""local_compute() is called when per-trainer metrics are required. It's
         can be used for debugging. Currently only rec_metrics is supported.
         """
-        ret: Dict[str, MetricValue] = {}
+        ret: MetricsResult = {}
         if self.rec_metrics:
             ret.update(self.rec_metrics.local_compute())
         return ret
@@ -492,9 +497,7 @@ def load_pre_compute_states(
     def shutdown(self) -> None:
         logger.info("Initiating graceful shutdown...")
 
-    def async_compute(
-        self, future: concurrent.futures.Future[Dict[str, MetricValue]]
-    ) -> None:
+    def async_compute(self) -> MetricsFuture:
         raise RecMetricException("async_compute is not supported in RecMetricModule")
 
 
@@ -610,6 +613,7 @@ def generate_metric_module(
     metrics = metric_class(
         batch_size=batch_size,
         world_size=world_size,
+        device=device,
         rec_tasks=metrics_config.rec_tasks,
         rec_metrics=rec_metrics,
         throughput_metric=throughput_metric,