feat: support engine metrics from PROMETHEUS_MULTIPROC_DIR in grpc mode

bindings/python/src/smg/serve.py

@@ -11,7 +11,9 @@
 import logging
 import os
 import random
+import shutil
 import signal
+import tempfile
 import socket
 import subprocess
 import sys
@@ -550,6 +552,7 @@ def __init__(self, backend: str, args: argparse.Namespace, backend_args: list[st
         self.launcher: WorkerLauncher = BACKEND_LAUNCHERS[backend]()
         self.workers: list[tuple[subprocess.Popen, int]] = []
         self._shutting_down = False
+        self._prometheus_dir: str | None = None
 
     # -- public API ---------------------------------------------------------
 
@@ -571,6 +574,15 @@ def run(self) -> None:
     def _launch_workers(self) -> None:
         ports = _find_available_ports(self.args.worker_base_port, self.args.data_parallel_size)
         host = self.args.worker_host
+
+        if getattr(self.args, "connection_mode", "grpc") == "grpc":
+            self._prometheus_dir = tempfile.mkdtemp(prefix="smg_prometheus_")
+            os.environ["PROMETHEUS_MULTIPROC_DIR"] = self._prometheus_dir
+            logger.info(
+                "Set PROMETHEUS_MULTIPROC_DIR=%s for gRPC metrics collection",
+                self._prometheus_dir,
+            )
+
         for dp_rank, port in enumerate(ports):
             env = self.launcher.gpu_env(self.args, dp_rank)
             proc = self.launcher.launch(self.args, self.backend_args, host, port, env)
@@ -641,6 +653,23 @@ def _cleanup_workers(self) -> None:
                 except (ProcessLookupError, OSError):
                     pass
 
+        self._cleanup_prometheus_dir()
+
+    def _cleanup_prometheus_dir(self) -> None:
+        """Remove the temporary prometheus multiprocess directory and its .db files."""
+        if self._prometheus_dir is None:
+            return
+        try:
+            shutil.rmtree(self._prometheus_dir)
+            logger.info("Cleaned up PROMETHEUS_MULTIPROC_DIR=%s", self._prometheus_dir)
+        except OSError as e:
+            logger.warning(
+                "Failed to clean up PROMETHEUS_MULTIPROC_DIR=%s: %s",
+                self._prometheus_dir,
+                e,
+            )
+        self._prometheus_dir = None
+
 
 # ---------------------------------------------------------------------------
 # Entry point

model_gateway/src/core/worker_manager.rs

@@ -17,7 +17,7 @@ use tokio::{
     sync::{watch, Mutex},
     task::JoinHandle,
 };
-use tracing::{debug, info};
+use tracing::{debug, info, warn};
 
 use crate::{
     core::{
@@ -85,6 +85,40 @@ impl IntoResponse for EngineMetricsResult {
     }
 }
 
+/// Collect gRPC worker metrics by aggregating `PROMETHEUS_MULTIPROC_DIR` via a python3 subprocess.
+async fn collect_prometheus_multiproc_metrics() -> Result<String, String> {
+    let dir = std::env::var("PROMETHEUS_MULTIPROC_DIR").map_err(|_| {
+        "PROMETHEUS_MULTIPROC_DIR not set; cannot collect metrics from gRPC workers".to_string()
+    })?;
+
+    let output = tokio::process::Command::new("python3")
+        .args([
+            "-c",
+            "import sys\n\
+             from prometheus_client import CollectorRegistry, generate_latest\n\
+             from prometheus_client.multiprocess import MultiProcessCollector\n\
+             registry = CollectorRegistry()\n\
+             MultiProcessCollector(registry)\n\
+             sys.stdout.buffer.write(generate_latest(registry))\n",
+        ])
+        .env("PROMETHEUS_MULTIPROC_DIR", &dir)
+        .output()
+        .await
+        .map_err(|e| format!("failed to run python3 prometheus collector: {e}"))?;
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        return Err(format!("python3 prometheus collector failed: {stderr}"));
+    }
+
+    let text = String::from_utf8(output.stdout)
+        .map_err(|e| format!("prometheus collector output is not valid UTF-8: {e}"))?;
+    if text.trim().is_empty() {
+        return Err("no metrics available from gRPC workers yet".to_string());
+    }
+    Ok(text)
+}
+
 pub struct WorkerManager;
 
 impl WorkerManager {
@@ -273,22 +307,52 @@ impl WorkerManager {
             return EngineMetricsResult::Err("No available workers".to_string());
         }
 
-        let responses = fan_out(&workers, client, "metrics", reqwest::Method::GET).await;
-
         let mut metric_packs = Vec::new();
-        for resp in responses {
-            if let Ok(r) = resp.result {
-                if r.status().is_success() {
-                    if let Ok(text) = r.text().await {
-                        metric_packs.push(MetricPack {
-                            labels: vec![("worker_addr".into(), resp.url)],
-                            metrics_text: text,
-                        });
+
+        let http_workers: Vec<_> = workers
+            .iter()
+            .filter(|w| matches!(w.connection_mode(), ConnectionMode::Http))
+            .cloned()
+            .collect();
+        let has_grpc = workers
+            .iter()
+            .any(|w| matches!(w.connection_mode(), ConnectionMode::Grpc));
+
+        if !http_workers.is_empty() {
+            let responses = fan_out(&http_workers, client, "metrics", reqwest::Method::GET).await;
+            for resp in responses {
+                if let Ok(r) = resp.result {
+                    if r.status().is_success() {
+                        if let Ok(text) = r.text().await {
+                            metric_packs.push(MetricPack {
+                                labels: vec![("worker_addr".into(), resp.url)],
+                                metrics_text: text,
+                            });
+                        }
                     }
                 }
             }
         }
 
+        if has_grpc {
+            match collect_prometheus_multiproc_metrics().await {
+                Ok(text) if !text.trim().is_empty() => {
+                    metric_packs.push(MetricPack {
+                        labels: vec![],
+                        metrics_text: text,
+                    });
+                }
+                Ok(_) => {
+                    // No metrics available yet from gRPC workers — skip silently
+                }
+                Err(e) => {
+                    warn!(
+                        "Failed to collect gRPC worker metrics from PROMETHEUS_MULTIPROC_DIR: {e}"
+                    );
+                }
+            }
+        }
+
         if metric_packs.is_empty() {
             return EngineMetricsResult::Err("All backend requests failed".to_string());
         }