diff --git a/python/ray/tests/test_metrics_agent.py b/python/ray/tests/test_metrics_agent.py index 86670b8a3..8e02c4ae3 100644 --- a/python/ray/tests/test_metrics_agent.py +++ b/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 diff --git a/python/ray/util/metrics.py b/python/ray/util/metrics.py index d287a503f..57a01cf7a 100644 --- a/python/ray/util/metrics.py +++ b/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",