[Metrics] Fix serialization for custom metrics (#13571)

python/ray/tests/test_metrics_agent.py

@@ -37,6 +37,8 @@ def _setup_cluster_for_test(ray_start_cluster):
     def f():
         counter = Count("test_counter", description="desc")
         counter.record(1)
+        counter = ray.get(ray.put(counter))  # Test serialization.
+        counter.record(1)
         ray.get(worker_should_exit.wait.remote())
 
     @ray.remote
@@ -44,6 +46,7 @@ def _setup_cluster_for_test(ray_start_cluster):
         async def ping(self):
             histogram = Histogram(
                 "test_histogram", description="desc", boundaries=[0.1, 1.6])
+            histogram = ray.get(ray.put(histogram))  # Test serialization.
             histogram.record(1.5)
             ray.get(worker_should_exit.wait.remote())
 
@@ -100,7 +103,7 @@ def test_metrics_export_end_to_end(_setup_cluster_for_test):
         test_counter_sample = [
             m for m in metric_samples if "test_counter" in m.name
         ][0]
-        assert test_counter_sample.value == 1.0
+        assert test_counter_sample.value == 2.0
 
         test_driver_counter_sample = [
             m for m in metric_samples if "test_driver_counter" in m.name

python/ray/util/metrics.py

@@ -147,6 +147,11 @@ class Count(Metric):
         self._metric = CythonCount(self._name, self._description, self._unit,
                                    self._tag_keys)
 
+    def __reduce__(self):
+        deserializer = Count
+        serialized_data = (self._name, self._description, self._tag_keys)
+        return deserializer, serialized_data
+
 
 class Histogram(Metric):
     """Histogram distribution of metric points.
@@ -177,6 +182,12 @@ class Histogram(Metric):
                                        self._unit, self.boundaries,
                                        self._tag_keys)
 
+    def __reduce__(self):
+        deserializer = Histogram
+        serialized_data = (self._name, self._description, self.boundaries,
+                           self._tag_keys)
+        return deserializer, serialized_data
+
     @property
     def info(self):
         """Return information about histogram metric."""
@@ -204,6 +215,11 @@ class Gauge(Metric):
         self._metric = CythonGauge(self._name, self._description, self._unit,
                                    self._tag_keys)
 
+    def __reduce__(self):
+        deserializer = Gauge
+        serialized_data = (self._name, self._description, self._tag_keys)
+        return deserializer, serialized_data
+
 
 __all__ = [
     "Count",