[RLlib] Add Prometheus metrics to the selected RLlib components

rllib/algorithms/appo/utils.py

@@ -3,16 +3,20 @@
 Luo et al. 2020
 https://arxiv.org/pdf/1912.00167
 """
-from collections import deque
 import threading
 import time
+from collections import deque
 
 import numpy as np
 
 from ray.rllib.models.catalog import ModelCatalog
 from ray.rllib.models.modelv2 import ModelV2
 from ray.rllib.utils.annotations import OldAPIStack
-
+from ray.rllib.utils.metrics.ray_metrics import (
+    DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+    TimerAndPrometheusLogger,
+)
+from ray.util.metrics import Counter, Histogram
 
 POLICY_SCOPE = "func"
 TARGET_POLICY_SCOPE = "target_func"
@@ -43,43 +47,77 @@ def __init__(self, num_batches: int, iterations_per_batch: int):
 
         self._rng = np.random.default_rng()
 
+        # Ray metrics
+        self._metrics_circular_buffer_add_time = Histogram(
+            name="rllib_utils_circular_buffer_add_time",
+            description="Time spent in CircularBuffer.add()",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_circular_buffer_add_time.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
+        self._metrics_circular_buffer_add_ts_dropped = Counter(
+            name="rllib_utils_circular_buffer_add_ts_dropped_counter",
+            description="Total number of env steps dropped by the CircularBuffer.",
+            tag_keys=("rllib",),
+        )
+        self._metrics_circular_buffer_add_ts_dropped.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
+        self._metrics_circular_buffer_sample_time = Histogram(
+            name="rllib_utils_circular_buffer_sample_time",
+            description="Time spent in CircularBuffer.sample()",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_circular_buffer_sample_time.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
     def add(self, batch):
         # Add buffer and k=0 information to the deque.
-        with self._lock:
-            dropped_entry = self._buffer[0]
-            for _ in range(self.iterations_per_batch):
-                self._buffer.append(batch)
-                self._indices.add(self._offset)
-                self._indices.discard(self._offset - self._NxK)
-                self._offset += 1
-            self._num_added += 1
-
-        # A valid entry (w/ a batch whose k has not been reach K yet) was dropped.
-        dropped_ts = 0
-        if dropped_entry is not None:
-            dropped_ts = dropped_entry.env_steps()
+        with TimerAndPrometheusLogger(self._metrics_circular_buffer_add_time):
+            with self._lock:
+                dropped_entry = self._buffer[0]
+                for _ in range(self.iterations_per_batch):
+                    self._buffer.append(batch)
+                    self._indices.add(self._offset)
+                    self._indices.discard(self._offset - self._NxK)
+                    self._offset += 1
+                self._num_added += 1
+
+            # A valid entry (w/ a batch whose k has not been reach K yet) was dropped.
+            dropped_ts = 0
+            if dropped_entry is not None:
+                dropped_ts = dropped_entry.env_steps()
+                if dropped_ts > 0:
+                    self._metrics_circular_buffer_add_ts_dropped.inc(value=dropped_ts)
 
         return dropped_ts
 
     def sample(self):
         # Only initially, the buffer may be empty -> Just wait for some time.
-        while len(self) == 0:
-            time.sleep(0.0001)
-
-        # Sample a random buffer index.
-        with self._lock:
-            idx = self._rng.choice(list(self._indices))
-            actual_buffer_idx = idx - self._offset + self._NxK
-            batch = self._buffer[actual_buffer_idx]
-            assert batch is not None, (
-                idx,
-                actual_buffer_idx,
-                self._offset,
-                self._indices,
-                [b is None for b in self._buffer],
-            )
-            self._buffer[actual_buffer_idx] = None
-            self._indices.discard(idx)
+        with TimerAndPrometheusLogger(self._metrics_circular_buffer_sample_time):
+            while len(self) == 0:
+                time.sleep(0.0001)
+
+            # Sample a random buffer index.
+            with self._lock:
+                idx = self._rng.choice(list(self._indices))
+                actual_buffer_idx = idx - self._offset + self._NxK
+                batch = self._buffer[actual_buffer_idx]
+                assert batch is not None, (
+                    idx,
+                    actual_buffer_idx,
+                    self._offset,
+                    self._indices,
+                    [b is None for b in self._buffer],
+                )
+                self._buffer[actual_buffer_idx] = None
+                self._indices.discard(idx)
 
         # Return the sampled batch.
         return batch