feat(tools): run precompute engine on dedicated named tokio runtime and add per-thread CPU attribution to process monitor

Author: milindsrivastava1997

asap-query-engine/src/main.rs

@@ -196,7 +196,7 @@ async fn main() -> Result<()> {
     // check_config() already enforces the ingest source is compatible (http_remote_write or csv).
     let mut pe_engine_handle: Option<PrecomputeEngineHandle> = None;
 
-    let precompute_handle = if config.streaming_engine == StreamingEngine::Precompute {
+    let _precompute_runtime = if config.streaming_engine == StreamingEngine::Precompute {
         let precompute_config = PrecomputeEngineConfig {
             num_workers: config.precompute_engine.num_workers,
             allowed_lateness_ms: config.precompute_engine.allowed_lateness_ms,
@@ -260,11 +260,18 @@ async fn main() -> Result<()> {
             spawn_memory_diagnostics(diag_store, Some(worker_diagnostics)).await;
         });
 
-        Some(tokio::spawn(async move {
+        let rt = tokio::runtime::Builder::new_multi_thread()
+            .thread_name("pc-worker")
+            .worker_threads(config.precompute_engine.num_workers)
+            .enable_all()
+            .build()
+            .expect("failed to build precompute runtime");
+        rt.spawn(async move {
             if let Err(e) = pe.run().await {
                 error!("Precompute engine error: {}", e);
             }
-        }))
+        });
+        Some(rt)
     } else {
         let diag_store = store.clone();
         tokio::spawn(async move {
@@ -457,12 +464,6 @@ async fn main() -> Result<()> {
         let _ = handle.await;
     }
 
-    if let Some(handle) = precompute_handle {
-        info!("Shutting down precompute engine...");
-        handle.abort();
-        let _ = handle.await;
-    }
-
     info!("Shutdown complete");
     Ok(())
 }

asap-tools/experiments/classes/process_monitor.py

@@ -1,9 +1,33 @@
 import multiprocessing
+import time
 import psutil
 import traceback
 from typing import List, Any
 from classes.ProcessMonitorHook import ProcessMonitorHook, ProcessMetricSnapshot
 
+_PRECOMPUTE_THREAD_PREFIX = "pc-worker"
+
+
+def _read_thread_cpu(pid: int) -> dict:
+    """
+    Returns {tid: (thread_name, cpu_seconds)} for all threads of pid.
+    cpu_seconds = user_time + system_time from psutil; name from /proc/[pid]/task/[tid]/comm.
+    Silently skips threads that disappear mid-read.
+    """
+    result = {}
+    try:
+        threads = psutil.Process(pid).threads()
+    except (psutil.NoSuchProcess, psutil.AccessDenied):
+        return result
+    for t in threads:
+        try:
+            with open(f"/proc/{pid}/task/{t.id}/comm") as f:
+                name = f.read().strip()
+            result[t.id] = (name, t.user_time + t.system_time)
+        except (FileNotFoundError, psutil.NoSuchProcess):
+            pass
+    return result
+
 
 class MyMonitor(multiprocessing.Process):
     def __init__(
@@ -34,11 +58,20 @@ def __init__(
             self.pid_monitor_map[pid] = {m: [] for m in self.monitors}
             self.pid_monitor_map[pid]["keyword"] = keyword
 
+        self._prev_thread_jiffies: dict = {}
+        self._prev_poll_monotonic: float = 0.0
+
+        for pid in self.pids_to_monitor:
+            self.pid_monitor_map[pid]["precompute_cpu_percent"] = []
+            self.pid_monitor_map[pid]["query_cpu_percent"] = []
+
     def add_child_pid_to_map(self, pid, child_pid):
         self.pid_monitor_map[child_pid] = {m: [] for m in self.monitors}
         self.pid_monitor_map[child_pid]["keyword"] = self.pid_monitor_map[pid][
             "keyword"
         ]
+        self.pid_monitor_map[child_pid]["precompute_cpu_percent"] = []
+        self.pid_monitor_map[child_pid]["query_cpu_percent"] = []
 
     def init_hooks(self):
         """
@@ -70,6 +103,45 @@ def close_hooks(self):
                 hook.close()
         return
 
+    def _compute_thread_group_cpu(self, pid: int, elapsed: float):
+        """
+        Reads current per-thread CPU seconds for pid, diffs against previous snapshot,
+        and appends precompute_cpu_percent / query_cpu_percent to pid_monitor_map.
+
+        CPU% is on the same scale as psutil's cpu_percent: can exceed 100% on
+        multi-core systems (e.g. 2 fully loaded cores → ~200%).
+        """
+        current = _read_thread_cpu(pid)
+        prev = self._prev_thread_jiffies.get(pid, {})
+
+        if not prev:
+            self._prev_thread_jiffies[pid] = current
+            self.pid_monitor_map[pid]["precompute_cpu_percent"].append(0.0)
+            self.pid_monitor_map[pid]["query_cpu_percent"].append(0.0)
+            return
+
+        precompute_seconds = 0.0
+        query_seconds = 0.0
+
+        for tid, (name, cpu_secs) in current.items():
+            prev_secs = prev.get(tid, (None, 0.0))[1]
+            delta = max(0.0, cpu_secs - prev_secs)
+            if name.startswith(_PRECOMPUTE_THREAD_PREFIX):
+                precompute_seconds += delta
+            else:
+                query_seconds += delta
+
+        if elapsed > 0:
+            precompute_pct = (precompute_seconds / elapsed) * 100.0
+            query_pct = (query_seconds / elapsed) * 100.0
+        else:
+            precompute_pct = 0.0
+            query_pct = 0.0
+
+        self.pid_monitor_map[pid]["precompute_cpu_percent"].append(precompute_pct)
+        self.pid_monitor_map[pid]["query_cpu_percent"].append(query_pct)
+        self._prev_thread_jiffies[pid] = current
+
     def update_pid_monitor_map(self, p) -> List[ProcessMetricSnapshot]:
         # if p.pid not in self.pid_monitor_map:
         #     self.pid_monitor_map[p.pid] = {m: [] for m in self.monitors}
@@ -97,18 +169,27 @@ def run(self):
         #       of the list of hooks
         self.init_hooks()
         self.pipe.send("ready")
-        stop = False
+
+        self._prev_poll_monotonic = time.monotonic()
+        for pid in self.pids_to_monitor:
+            self._prev_thread_jiffies[pid] = _read_thread_cpu(pid)
 
         try:
             while True:
-                iteration_info = []  # list of process snapshots from this iteration
+                now = time.monotonic()
+                elapsed = now - self._prev_poll_monotonic
+                self._prev_poll_monotonic = now
+
+                iteration_info = []
                 for pid, p in self.psutil_handles.items():
                     iteration_info += self.update_pid_monitor_map(p)
+                    self._compute_thread_group_cpu(pid, elapsed)
                     if self.include_children:
                         for child in p.children(recursive=True):
                             if child.pid not in self.pid_monitor_map:
                                 self.add_child_pid_to_map(pid, child.pid)
                             iteration_info += self.update_pid_monitor_map(child)
+                            self._compute_thread_group_cpu(child.pid, elapsed)
 
                 self.update_hooks(iteration_info)
                 stop = self.pipe.poll(self.interval)