"""Enhanced metrics collection for task tracking.
This module provides comprehensive performance metrics collection during node execution,
including CPU usage, memory consumption, I/O operations, and custom metrics.
Design Purpose:
- Enable real-time performance monitoring during node execution
- Integrate seamlessly with TaskManager and visualization components
- Support both synchronous and asynchronous execution contexts
Upstream Dependencies:
- Runtime engines (local.py, parallel.py, docker.py) use this to collect metrics
- TaskManager uses this to store performance data
Downstream Consumers:
- Visualization components use collected metrics for performance graphs
- Export utilities include metrics in workflow reports
"""
import asyncio
import threading
import time
from collections.abc import Callable
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any
try:
import psutil
PSUTIL_AVAILABLE = True
except ImportError:
psutil = None # type: ignore[assignment]
PSUTIL_AVAILABLE = False
[docs]
class MetricsCollector:
"""Collects performance metrics during task execution.
This class provides context managers for collecting detailed performance
metrics during node execution, with support for both process-level and
system-level monitoring.
Usage:
>>> collector = MetricsCollector()
>>> with collector.collect() as metrics:
... # Execute node code here
... pass
>>> performance_data = metrics.result()
"""
[docs]
def __init__(
self,
sampling_interval: float = 0.1,
enable_resource_monitoring: bool | None = None,
):
"""Initialize metrics collector.
Args:
sampling_interval: How often to sample metrics (seconds)
enable_resource_monitoring: Whether to enable psutil-based resource monitoring.
``None`` (default) auto-detects: enabled iff psutil is installed.
``True`` explicitly opts in (warns if psutil is missing).
``False`` disables; only duration is tracked. P0D-001: Dramatically
reduces per-node overhead by avoiding thread creation/join.
"""
self.sampling_interval = sampling_interval
# Auto-detect when caller didn't specify; explicit opt-in still warns.
explicit_opt_in = enable_resource_monitoring is True
if enable_resource_monitoring is None:
enable_resource_monitoring = PSUTIL_AVAILABLE
# P0D-001: Only enable psutil monitoring when both available AND requested
self._monitoring_enabled = PSUTIL_AVAILABLE and enable_resource_monitoring
if not PSUTIL_AVAILABLE and explicit_opt_in:
import warnings
warnings.warn(
"psutil not available — resource monitoring requested but disabled. "
"Install with: pip install 'kailash[monitoring]'",
stacklevel=2,
)
[docs]
@contextmanager
def collect(self, node_id: str | None = None):
"""Context manager for collecting metrics during execution.
Args:
node_id: Optional node identifier for tracking
Yields:
MetricsContext: Context object with result() method
"""
context = MetricsContext(
node_id=node_id,
sampling_interval=self.sampling_interval,
monitoring_enabled=self._monitoring_enabled,
)
try:
context.start()
yield context
finally:
context.stop()
[docs]
async def collect_async(self, coro, node_id: str | None = None):
"""Collect metrics for async execution.
Args:
coro: Coroutine to execute
node_id: Optional node identifier
Returns:
Tuple of (result, metrics)
"""
context = MetricsContext(
node_id=node_id,
sampling_interval=self.sampling_interval,
monitoring_enabled=self._monitoring_enabled,
)
try:
context.start()
result = await coro
return result, context.result()
finally:
context.stop()
class MetricsContext:
"""Context for collecting metrics during a specific execution."""
def __init__(
self, node_id: str | None, sampling_interval: float, monitoring_enabled: bool
):
self.node_id = node_id
self.sampling_interval = sampling_interval
self.monitoring_enabled = monitoring_enabled
self.start_time: float | None = None
self.end_time: float | None = None
self.process: Any | None = None
self.initial_io: Any | None = None
self.initial_memory: float | None = None
self.peak_memory: float = 0.0
self.cpu_samples: list = []
self.monitoring_thread: threading.Thread | None = None
self._stop_monitoring = threading.Event()
def start(self):
"""Start metrics collection."""
self.start_time = time.time()
if self.monitoring_enabled and psutil is not None:
try:
self.process = psutil.Process()
assert self.process is not None
mem_info = self.process.memory_info()
self.initial_memory = mem_info.rss / 1024 / 1024 # MB
self.peak_memory = (
self.initial_memory if self.initial_memory is not None else 0.0
)
# Get initial I/O counters if available
if hasattr(self.process, "io_counters"):
try:
self.initial_io = self.process.io_counters()
except (psutil.AccessDenied, AttributeError):
self.initial_io = None
# Start monitoring thread
self._stop_monitoring.clear()
self.monitoring_thread = threading.Thread(
target=self._monitor_resources
)
self.monitoring_thread.daemon = True
self.monitoring_thread.start()
except Exception:
# Fallback if process monitoring fails
self.monitoring_enabled = False
def stop(self):
"""Stop metrics collection."""
self.end_time = time.time()
if self.monitoring_enabled and self.monitoring_thread:
self._stop_monitoring.set()
self.monitoring_thread.join(timeout=1.0)
def _monitor_resources(self):
"""Monitor resources in background thread."""
if self.process is None:
return
while not self._stop_monitoring.is_set():
try:
# Sample CPU usage
cpu = self.process.cpu_percent(interval=None)
if cpu > 0: # Filter out initial 0 readings
self.cpu_samples.append(cpu)
# Track peak memory
mem_info = self.process.memory_info()
memory = mem_info.rss / 1024 / 1024 # MB
self.peak_memory = max(self.peak_memory, memory)
except Exception:
break
self._stop_monitoring.wait(self.sampling_interval)
def result(self) -> PerformanceMetrics:
"""Get collected metrics."""
metrics = PerformanceMetrics()
# Calculate duration
if self.start_time and self.end_time:
metrics.duration = self.end_time - self.start_time
if self.monitoring_enabled and self.process and psutil is not None:
try:
# CPU usage (average of samples)
if self.cpu_samples:
metrics.cpu_percent = sum(self.cpu_samples) / len(self.cpu_samples)
# Memory metrics
metrics.memory_mb = self.peak_memory
if self.initial_memory:
current_memory = self.process.memory_info().rss / 1024 / 1024
metrics.memory_delta_mb = current_memory - self.initial_memory
# I/O metrics
if self.initial_io and hasattr(self.process, "io_counters"):
try:
current_io = self.process.io_counters()
metrics.io_read_bytes = (
current_io.read_bytes - self.initial_io.read_bytes
)
metrics.io_write_bytes = (
current_io.write_bytes - self.initial_io.write_bytes
)
metrics.io_read_count = (
current_io.read_count - self.initial_io.read_count
)
metrics.io_write_count = (
current_io.write_count - self.initial_io.write_count
)
except (psutil.AccessDenied, AttributeError):
pass
# Thread and context switch info
try:
metrics.thread_count = self.process.num_threads()
if hasattr(self.process, "num_ctx_switches"):
ctx = self.process.num_ctx_switches()
metrics.context_switches = ctx.voluntary + ctx.involuntary
except (psutil.AccessDenied, AttributeError):
pass
except (psutil.NoSuchProcess, psutil.AccessDenied):
pass
return metrics
def add_custom_metric(self, name: str, value: Any):
"""Add a custom metric."""
if not hasattr(self, "_custom_metrics"):
self._custom_metrics = {}
self._custom_metrics[name] = value
def get_custom_metrics(self) -> dict[str, Any]:
"""Get custom metrics."""
return getattr(self, "_custom_metrics", {})
# Global collector instance for convenience
default_collector = MetricsCollector()
def collect_metrics(func: Callable | None = None, *, node_id: str | None = None):
"""Decorator for collecting metrics on function execution.
Can be used as @collect_metrics or @collect_metrics(node_id="my_node")
Args:
func: Function to wrap
node_id: Optional node identifier
Returns:
Wrapped function that returns (result, metrics) tuple
"""
def decorator(f):
if asyncio.iscoroutinefunction(f):
async def async_wrapper(*args, **kwargs):
result, metrics = await default_collector.collect_async(
f(*args, **kwargs), node_id=node_id
)
return result, metrics
return async_wrapper
else:
def sync_wrapper(*args, **kwargs):
with default_collector.collect(node_id=node_id) as context:
result = f(*args, **kwargs)
return result, context.result()
return sync_wrapper
if func is None:
# Called with arguments: @collect_metrics(node_id="...")
return decorator
else:
# Called without arguments: @collect_metrics
return decorator(func)