[fix] Decouple metrics exporter from controller internals

- Replace direct controller access with snapshot-based collection:
  controller pushes plain-dict snapshots via update_controller_snapshot(),
  metrics thread reads dicts instead of live objects/tensors.
- Make measure() fault-tolerant: wrap all prometheus calls in try/except
  so metrics failures never block the controller request loop.
- Add dedicated snapshot daemon thread in controller.
- Update tests to use snapshot-based API.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: 2 people

tests/test_metrics.py

@@ -16,7 +16,6 @@
 """Unit tests for the Prometheus metrics exporter (transfer_queue.metrics)."""
 
 import time
-from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
@@ -32,54 +31,35 @@
 
 
 # ---------------------------------------------------------------------------
-# Helpers — lightweight fakes that satisfy TQMetricsExporter without Ray
+# Helpers — build snapshot dicts that TQMetricsExporter.update_controller_snapshot expects
 # ---------------------------------------------------------------------------
 
 
-def _make_fake_partition(
-    partition_id: str,
+def _make_partition_snapshot(
     total_samples: int = 10,
-    total_fields: int = 2,
-    allocated_samples: int = 16,
     produced_ratio: float = 0.5,
     consumption: dict | None = None,
-):
-    """Return an object that mimics ``DataPartitionStatus.get_statistics()``."""
+) -> dict:
+    """Return a partition snapshot dict."""
     consumption_stats = {}
     if consumption:
         for task, progress in consumption.items():
-            consumption_stats[task] = {
-                "consumed_samples": int(total_samples * progress),
-                "consumption_progress": progress,
-            }
+            consumption_stats[task] = {"consumption_progress": progress}
 
-    stats = {
-        "partition_id": partition_id,
-        "created_at": time.time(),
+    return {
         "total_samples_num": total_samples,
-        "total_fields_num": total_fields,
-        "allocated_samples_num": allocated_samples,
-        "allocated_fields_num": total_fields,
-        "registered_tasks": list((consumption or {}).keys()),
-        "produced_samples": int(total_samples * produced_ratio),
         "production_progress": produced_ratio,
-        "field_statistics": {},
         "consumption_statistics": consumption_stats,
     }
 
-    partition = MagicMock()
-    partition.get_statistics.return_value = stats
-    return partition
-
 
-def _make_fake_controller(partitions=None, allocated=10, reusable=2):
-    """Return a lightweight stand-in for ``TransferQueueController``."""
-    ctrl = SimpleNamespace()
-    ctrl.partitions = partitions or {}
-    ctrl.index_manager = SimpleNamespace()
-    ctrl.index_manager.allocated_indexes = set(range(allocated))
-    ctrl.index_manager.reusable_indexes = list(range(reusable))
-    return ctrl
+def _make_snapshot(partitions=None, allocated=10, reusable=2) -> dict:
+    """Return a controller metrics snapshot dict."""
+    return {
+        "partitions": partitions or {},
+        "global_index_allocated": allocated,
+        "global_index_reusable": reusable,
+    }
 
 
 # ---------------------------------------------------------------------------
@@ -90,8 +70,7 @@ def _make_fake_controller(partitions=None, allocated=10, reusable=2):
 class TestMetricDefinitions:
     def test_all_metrics_are_registered(self):
         """Verify that all expected metric families exist in the exporter's registry."""
-        ctrl = _make_fake_controller()
-        exporter = TQMetricsExporter(ctrl)
+        exporter = TQMetricsExporter()
 
         expected_prefixes = [
             "tq_controller_uptime_seconds",
@@ -123,9 +102,9 @@ def test_all_metrics_are_registered(self):
 
 class TestControllerMetricsCollection:
     def test_collect_empty_controller(self):
-        """Collect metrics from a controller with no partitions — should not raise."""
-        ctrl = _make_fake_controller(partitions={}, allocated=0, reusable=0)
-        exporter = TQMetricsExporter(ctrl)
+        """Collect metrics from an empty snapshot — should not raise."""
+        exporter = TQMetricsExporter()
+        exporter.update_controller_snapshot(_make_snapshot(partitions={}, allocated=0, reusable=0))
         exporter.collect_controller_metrics()
 
         assert exporter.partitions_total._value.get() == 0
@@ -134,15 +113,12 @@ def test_collect_empty_controller(self):
 
     def test_collect_with_partitions(self):
         """Partition-level metrics are populated correctly."""
-        p1 = _make_fake_partition(
-            "train_0", total_samples=20, produced_ratio=0.8, consumption={"gen": 0.5}
-        )
-        p2 = _make_fake_partition(
-            "train_1", total_samples=10, produced_ratio=1.0, consumption={"gen": 1.0, "train": 0.3}
-        )
-        ctrl = _make_fake_controller(partitions={"train_0": p1, "train_1": p2}, allocated=30, reusable=5)
+        p1 = _make_partition_snapshot(total_samples=20, produced_ratio=0.8, consumption={"gen": 0.5})
+        p2 = _make_partition_snapshot(total_samples=10, produced_ratio=1.0, consumption={"gen": 1.0, "train": 0.3})
+        snapshot = _make_snapshot(partitions={"train_0": p1, "train_1": p2}, allocated=30, reusable=5)
 
-        exporter = TQMetricsExporter(ctrl)
+        exporter = TQMetricsExporter()
+        exporter.update_controller_snapshot(snapshot)
         exporter.collect_controller_metrics()
 
         assert exporter.partitions_total._value.get() == 2
@@ -163,8 +139,8 @@ def test_collect_with_partitions(self):
 
     def test_uptime_increases(self):
         """Controller uptime should be positive after collection."""
-        ctrl = _make_fake_controller()
-        exporter = TQMetricsExporter(ctrl)
+        exporter = TQMetricsExporter()
+        exporter.update_controller_snapshot(_make_snapshot())
         time.sleep(0.05)
         exporter.collect_controller_metrics()
         assert exporter.controller_uptime._value.get() > 0
@@ -177,8 +153,7 @@ def test_uptime_increases(self):
 
 class TestMeasureContextManager:
     def test_measure_records_count_and_duration(self):
-        ctrl = _make_fake_controller()
-        exporter = TQMetricsExporter(ctrl)
+        exporter = TQMetricsExporter()
 
         with exporter.measure("GET_META"):
             time.sleep(0.01)
@@ -192,8 +167,7 @@ def test_measure_records_count_and_duration(self):
         assert hist._sum.get() > 0
 
     def test_measure_records_errors(self):
-        ctrl = _make_fake_controller()
-        exporter = TQMetricsExporter(ctrl)
+        exporter = TQMetricsExporter()
 
         with pytest.raises(ValueError):
             with exporter.measure("BAD_OP"):
@@ -204,8 +178,7 @@ def test_measure_records_errors(self):
         assert exporter.request_total.labels(op_type="BAD_OP")._value.get() == 1.0
 
     def test_multiple_ops_tracked_independently(self):
-        ctrl = _make_fake_controller()
-        exporter = TQMetricsExporter(ctrl)
+        exporter = TQMetricsExporter()
 
         for _ in range(3):
             with exporter.measure("GET_META"):
@@ -226,15 +199,13 @@ def test_multiple_ops_tracked_independently(self):
 class TestStorageMetricsCollection:
     def test_collect_with_no_storage_units(self):
         """No storage units registered — collect should be a no-op."""
-        ctrl = _make_fake_controller()
-        exporter = TQMetricsExporter(ctrl)
+        exporter = TQMetricsExporter()
         # Should not raise
         exporter.collect_storage_metrics()
 
     def test_storage_metrics_populated_on_success(self):
         """Verify storage gauges are set when _query_storage_unit returns data."""
-        ctrl = _make_fake_controller()
-        exporter = TQMetricsExporter(ctrl)
+        exporter = TQMetricsExporter()
 
         fake_su_info = MagicMock()
         fake_su_info.id = "SU_001"
@@ -259,8 +230,7 @@ def test_storage_metrics_populated_on_success(self):
 
     def test_storage_metrics_handles_query_failure(self):
         """If a storage unit query fails, other units should still be collected."""
-        ctrl = _make_fake_controller()
-        exporter = TQMetricsExporter(ctrl)
+        exporter = TQMetricsExporter()
 
         su1 = MagicMock()
         su1.id = "SU_001"

transfer_queue/controller.py

@@ -2146,6 +2146,43 @@ def register_sampler(
 
     # ==================== Metrics API ====================
 
+    def _build_metrics_snapshot(self) -> dict:
+        """Build a plain-dict snapshot of controller state for the metrics exporter.
+
+        The snapshot contains only primitive / dict values — no references to
+        live controller objects — so the metrics thread can read it safely.
+        """
+        partitions_data: dict = {}
+        for pid, partition in list(self.partitions.items()):
+            try:
+                stats = partition.get_statistics()
+                partitions_data[pid] = {
+                    "total_samples_num": stats["total_samples_num"],
+                    "production_progress": stats.get("production_progress", 0),
+                    "consumption_statistics": {
+                        task_name: {"consumption_progress": cstats.get("consumption_progress", 0)}
+                        for task_name, cstats in stats.get("consumption_statistics", {}).items()
+                    },
+                }
+            except Exception:
+                pass
+
+        return {
+            "partitions": partitions_data,
+            "global_index_allocated": len(self.index_manager.allocated_indexes),
+            "global_index_reusable": len(self.index_manager.reusable_indexes),
+        }
+
+    def _push_metrics_snapshot(self) -> None:
+        """Push a fresh metrics snapshot to the exporter (called from controller threads)."""
+        if self._metrics is None:
+            return
+        try:
+            snapshot = self._build_metrics_snapshot()
+            self._metrics.update_controller_snapshot(snapshot)
+        except Exception as e:
+            logger.debug(f"[{self.controller_id}]: Failed to push metrics snapshot: {e}")
+
     def start_metrics(self) -> str:
         """Initialize and start the Prometheus metrics exporter.
 
@@ -2160,11 +2197,27 @@ def start_metrics(self) -> str:
             return self._metrics_endpoint
         from transfer_queue.metrics import TQMetricsExporter
 
-        self._metrics = TQMetricsExporter(self)
+        self._metrics = TQMetricsExporter()
         self._metrics_endpoint = self._metrics.start(node_ip=self._node_ip)
+        # Launch a daemon thread that periodically pushes controller state
+        # snapshots to the exporter, keeping them process-isolated.
+        self._metrics_snapshot_thread = Thread(
+            target=self._metrics_snapshot_loop,
+            name="TQMetricsSnapshotThread",
+            daemon=True,
+        )
+        self._metrics_snapshot_thread.start()
         logger.info(f"[{self.controller_id}]: Prometheus metrics exporter started on {self._metrics_endpoint}")
         return self._metrics_endpoint
 
+    def _metrics_snapshot_loop(self) -> None:
+        """Periodically push a metrics snapshot to the exporter."""
+        from transfer_queue.metrics import TQ_METRICS_COLLECT_INTERVAL
+
+        while True:
+            self._push_metrics_snapshot()
+            time.sleep(TQ_METRICS_COLLECT_INTERVAL)
+
     def register_storage_units_for_metrics(self, storage_unit_infos: dict) -> None:
         """Register storage unit ZMQ endpoints so the metrics collector can query them.