[metrics] Check that all tag_keys are set when recording (#13420)

python/ray/tests/test_metrics.py

@@ -141,6 +141,6 @@ def test_multi_node_metrics_export_port_discovery(ray_start_cluster):
 
 
 if __name__ == "__main__":
-    import sys
     import pytest
+    import sys
     sys.exit(pytest.main(["-v", __file__]))

python/ray/tests/test_metrics_agent.py

@@ -199,8 +199,8 @@ def test_basic_custom_metrics(metric_mock):
     # -- Count --
     count = Count("count", tag_keys=("a", ))
     count._metric = metric_mock
-    count.record(1)
-    metric_mock.record.assert_called_with(1, tags={})
+    count.record(1, {"a": "1"})
+    metric_mock.record.assert_called_with(1, tags={"a": "1"})
 
     # -- Gauge --
     gauge = Gauge("gauge", description="gauge")
@@ -212,11 +212,6 @@ def test_basic_custom_metrics(metric_mock):
     histogram = Histogram(
         "hist", description="hist", boundaries=[1.0, 3.0], tag_keys=("a", "b"))
     histogram._metric = metric_mock
-    histogram.record(4)
-    metric_mock.record.assert_called_with(4, tags={})
-    tags = {"a": "3"}
-    histogram.record(10, tags=tags)
-    metric_mock.record.assert_called_with(10, tags=tags)
     tags = {"a": "10", "b": "b"}
     histogram.record(8, tags=tags)
     metric_mock.record.assert_called_with(8, tags=tags)
@@ -243,10 +238,6 @@ def test_custom_metrics_default_tags(metric_mock):
         })
     histogram._metric = metric_mock
 
-    # Check default tags.
-    histogram.record(4)
-    metric_mock.record.assert_called_with(4, tags={"b": "b"})
-
     # Check specifying non-default tags.
     histogram.record(10, tags={"a": "a"})
     metric_mock.record.assert_called_with(10, tags={"a": "a", "b": "b"})
@@ -301,17 +292,40 @@ def test_metrics_override_shouldnt_warn(ray_start_regular, log_pubsub):
             assert "Attempt to register measure" not in line
 
 
-def test_custom_metrics_tag_validation(ray_start_regular_shared):
+def test_custom_metrics_validation(ray_start_regular_shared):
+    # Missing tag(s) from tag_keys.
+    metric = Count("name", tag_keys=("a", "b"))
+    metric.set_default_tags({"a": "1"})
+
+    metric.record(1.0, {"b": "2"})
+    metric.record(1.0, {"a": "1", "b": "2"})
+
+    with pytest.raises(ValueError):
+        metric.record(1.0)
+
+    with pytest.raises(ValueError):
+        metric.record(1.0, {"a": "2"})
+
+    # Extra tag not in tag_keys.
+    metric = Count("name", tag_keys=("a", ))
+    with pytest.raises(ValueError):
+        metric.record(1.0, {"a": "1", "b": "2"})
+
+    # tag_keys must be tuple.
     with pytest.raises(TypeError):
         Count("name", tag_keys="a")
+    # tag_keys must be strs.
     with pytest.raises(TypeError):
         Count("name", tag_keys=(1, ))
 
     metric = Count("name", tag_keys=("a", ))
+    # Set default tag that isn't in tag_keys.
     with pytest.raises(ValueError):
         metric.set_default_tags({"a": "1", "c": "2"})
+    # Default tag value must be str.
     with pytest.raises(TypeError):
         metric.set_default_tags({"a": 1})
+    # Tag value must be str.
     with pytest.raises(TypeError):
         metric.record(1.0, {"a": 1})
 

python/ray/util/metrics.py

@@ -75,6 +75,8 @@ class Metric:
     def record(self, value: float, tags: dict = None) -> None:
         """Record the metric point of the metric.
 
+        Tags passed in will take precedence over the metric's default tags.
+
         Args:
             value(float): The value to be recorded as a metric point.
         """
@@ -85,9 +87,22 @@ class Metric:
                     raise TypeError(
                         f"Tag values must be str, got {type(val)}.")
 
-        default_tag_copy = self._default_tags.copy()
-        default_tag_copy.update(tags or {})
-        self._metric.record(value, tags=default_tag_copy)
+        final_tags = {}
+        tags_copy = tags.copy() if tags else {}
+        for tag_key in self._tag_keys:
+            # Prefer passed tags over default tags.
+            if tags is not None and tag_key in tags:
+                final_tags[tag_key] = tags_copy.pop(tag_key)
+            elif tag_key in self._default_tags:
+                final_tags[tag_key] = self._default_tags[tag_key]
+            else:
+                raise ValueError(f"Missing value for tag key {tag_key}.")
+
+        if len(tags_copy) > 0:
+            raise ValueError(
+                f"Unrecognized tag keys: {list(tags_copy.keys())}.")
+
+        self._metric.record(value, tags=final_tags)
 
     @property
     def info(self) -> Dict[str, Any]:
@@ -156,7 +171,7 @@ class Histogram(Metric):
         if boundaries is None or len(boundaries) == 0:
             raise ValueError(
                 "boundaries argument should be provided when using the "
-                "Histogram class. EX) Histgoram(boundaries=[1.0, 2.0])")
+                "Histogram class. e.g., Histogram(boundaries=[1.0, 2.0])")
         self.boundaries = boundaries
         self._metric = CythonHistogram(self._name, self._description,
                                        self._unit, self.boundaries,