[style] Apply ruff lint/format fixes and resolve mypy type errors

Auto-fixed by pre-commit: remove unused imports, Optional -> X | None,
zip() strict=False, line length formatting. Also add proper type
annotation for _metrics in SimpleStorageUnit to fix mypy errors.

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

Author: 2 people

tests/test_metrics.py

@@ -132,20 +132,21 @@ def test_collect_with_partitions(self):
 
         # Check partition-level gauges
         assert exporter.partition_samples.labels(partition_id="train_0")._value.get() == 20
-        assert exporter.partition_production_progress.labels(
-            partition_id="train_0", task_name="gen"
-        )._value.get() == 0.8
-        assert exporter.partition_consumption_progress.labels(
-            partition_id="train_0", task_name="gen"
-        )._value.get() == 0.5
+        assert (
+            exporter.partition_production_progress.labels(partition_id="train_0", task_name="gen")._value.get() == 0.8
+        )
+        assert (
+            exporter.partition_consumption_progress.labels(partition_id="train_0", task_name="gen")._value.get() == 0.5
+        )
 
         assert exporter.partition_samples.labels(partition_id="train_1")._value.get() == 10
-        assert exporter.partition_production_progress.labels(
-            partition_id="train_1", task_name="gen"
-        )._value.get() == 1.0
-        assert exporter.partition_consumption_progress.labels(
-            partition_id="train_1", task_name="train"
-        )._value.get() == 0.3
+        assert (
+            exporter.partition_production_progress.labels(partition_id="train_1", task_name="gen")._value.get() == 1.0
+        )
+        assert (
+            exporter.partition_consumption_progress.labels(partition_id="train_1", task_name="train")._value.get()
+            == 0.3
+        )
 
     def test_uptime_increases(self):
         """Controller uptime should be positive after collection."""

transfer_queue/interface.py

@@ -18,7 +18,7 @@
 import subprocess
 import time
 from importlib import resources
-from typing import Any, Callable, Optional
+from typing import Any, Callable
 from urllib.parse import urlparse
 
 import ray
@@ -384,7 +384,7 @@ def close():
 
 
 # ==================== Metrics API ====================
-def get_metrics_endpoint() -> Optional[str]:
+def get_metrics_endpoint() -> str | None:
     """Return the Prometheus metrics endpoint address (``host:port``), or *None* if metrics are disabled.
 
     Works from any process — the endpoint is stored in the Controller's config

transfer_queue/metrics.py

@@ -17,7 +17,7 @@
 import time
 from contextlib import contextmanager
 from threading import Thread
-from typing import Any, Optional
+from typing import Any
 from uuid import uuid4
 
 import psutil
@@ -334,25 +334,23 @@ def collect_storage_metrics(self) -> None:
                 active = metrics.get("active_keys", 0)
                 self.storage_capacity.labels(storage_unit_id=label).set(capacity)
                 self.storage_active_keys.labels(storage_unit_id=label).set(active)
-                self.storage_utilization.labels(storage_unit_id=label).set(
-                    active / capacity if capacity > 0 else 0.0
-                )
+                self.storage_utilization.labels(storage_unit_id=label).set(active / capacity if capacity > 0 else 0.0)
                 self.storage_memory_rss.labels(storage_unit_id=label).set(metrics.get("process_rss_bytes", 0))
 
                 # Per-operation request stats
                 for op_type, op_data in metrics.get("op_stats", {}).items():
-                    self.storage_request_ops.labels(
-                        storage_unit_id=label, op_type=op_type
-                    ).set(op_data.get("request_count", 0))
-                    self.storage_request_latency_avg.labels(
-                        storage_unit_id=label, op_type=op_type
-                    ).set(op_data.get("latency_avg", 0))
-                    self.storage_request_latency_p50.labels(
-                        storage_unit_id=label, op_type=op_type
-                    ).set(op_data.get("latency_p50", 0))
-                    self.storage_request_latency_p99.labels(
-                        storage_unit_id=label, op_type=op_type
-                    ).set(op_data.get("latency_p99", 0))
+                    self.storage_request_ops.labels(storage_unit_id=label, op_type=op_type).set(
+                        op_data.get("request_count", 0)
+                    )
+                    self.storage_request_latency_avg.labels(storage_unit_id=label, op_type=op_type).set(
+                        op_data.get("latency_avg", 0)
+                    )
+                    self.storage_request_latency_p50.labels(storage_unit_id=label, op_type=op_type).set(
+                        op_data.get("latency_p50", 0)
+                    )
+                    self.storage_request_latency_p99.labels(storage_unit_id=label, op_type=op_type).set(
+                        op_data.get("latency_p99", 0)
+                    )
             except Exception as e:
                 logger.warning(f"Failed to collect metrics from storage unit {su_id}: {e}")
 
@@ -375,7 +373,7 @@ def _get_or_create_socket(self, su_id: str, su_info: ZMQServerInfo) -> zmq.Socke
         self._zmq_sockets[su_id] = sock
         return sock
 
-    def _query_storage_unit(self, su_info: ZMQServerInfo, su_id: str) -> Optional[dict[str, Any]]:
+    def _query_storage_unit(self, su_info: ZMQServerInfo, su_id: str) -> dict[str, Any] | None:
         """Send a synchronous GET_METRICS request to a single storage unit."""
         try:
             sock = self._get_or_create_socket(su_id, su_info)

transfer_queue/storage/simple_storage.py

@@ -17,7 +17,7 @@
 import time
 import weakref
 from threading import Event, Thread
-from typing import Any
+from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
 import psutil
@@ -38,6 +38,9 @@
     get_node_ip_address,
 )
 
+if TYPE_CHECKING:
+    from transfer_queue.metrics import TQMetricsExporter
+
 logger = get_logger(__name__)
 
 TQ_STORAGE_POLLER_TIMEOUT = int(os.environ.get("TQ_STORAGE_POLLER_TIMEOUT", 5))  # in seconds
@@ -172,7 +175,7 @@ def __init__(self, storage_unit_size: int):
         self.proxy_thread: Thread | None = None
         self.worker_thread: Thread | None = None
 
-        self._metrics = None
+        self._metrics: TQMetricsExporter | None = None
 
         self._init_zmq_socket()
         self._start_process_put_get()
@@ -559,9 +562,11 @@ def _quantile_from_cumulative(hist, cumulative_counts: list[float], q: float) ->
         target = q * total
         prev_bound = 0.0
         prev_cumulative = 0.0
-        for bound, cum_count in zip(hist._upper_bounds, cumulative_counts):
+        for bound, cum_count in zip(hist._upper_bounds, cumulative_counts, strict=False):
             if cum_count >= target:
-                fraction = (target - prev_cumulative) / (cum_count - prev_cumulative) if cum_count > prev_cumulative else 0
+                fraction = (
+                    (target - prev_cumulative) / (cum_count - prev_cumulative) if cum_count > prev_cumulative else 0
+                )
                 return prev_bound + (bound - prev_bound) * fraction
             prev_bound = bound
             prev_cumulative = cum_count