matrice-compute 0.1.26__py3-none-any.whl → 0.1.27__py3-none-any.whl

matrice_compute/resources_tracker.py

@@ -5,6 +5,8 @@ This module contains classes for tracking machine and action resources.
 import os
 import subprocess
 import logging
+import threading
+import json
 from datetime import datetime, timezone
 import psutil
 import docker
@@ -580,3 +582,261 @@ class MachineResourcesTracker:
                 "Error in updating available resources: %s",
                 err,
             )
+
+
+class KafkaResourceMonitor:
+    """
+    Monitors system resources and publishes them to Kafka in a separate thread.
+    This class provides thread-safe start/stop operations for resource monitoring.
+    """
+
+    def __init__(
+        self,
+        instance_id: Optional[str] = None,
+        kafka_bootstrap: Optional[str] = None,
+        interval_seconds: int = 60,
+    ):
+        """
+        Initialize KafkaResourceMonitor.
+
+        Args:
+            instance_id: Instance identifier for Kafka topic. Defaults to INSTANCE_ID env var.
+            kafka_bootstrap: Kafka bootstrap servers. Required - should be obtained from Scaling.get_kafka_bootstrap_servers().
+            interval_seconds: Interval between resource checks in seconds. Defaults to 60.
+        """
+        self.instance_id = instance_id or os.getenv("INSTANCE_ID")
+        if not self.instance_id:
+            raise ValueError("instance_id must be provided or INSTANCE_ID env var must be set")
+
+        if not kafka_bootstrap:
+            raise ValueError("kafka_bootstrap must be provided - use Scaling.get_kafka_bootstrap_servers() to get internal Kafka config")
+
+        self.kafka_bootstrap = kafka_bootstrap
+        self.interval_seconds = interval_seconds
+        self.topic_name = "compute_instance_resource_utilization"
+
+        self._stop_event = threading.Event()
+        self._monitor_thread: Optional[threading.Thread] = None
+        self._producer = None
+        self._is_running = False
+
+    @staticmethod
+    def get_all_gpu_memory() -> Dict[int, tuple]:
+        """
+        Get GPU memory usage and total for all GPUs.
+
+        Returns:
+            Dict[int, tuple]: Dictionary mapping GPU ID to (used_gb, total_gb).
+                             Returns empty dict if nvidia-smi is not available.
+        """
+        gpu_usage = {}
+
+        try:
+            cmd = [
+                "nvidia-smi",
+                "--query-gpu=index,memory.used,memory.total",
+                "--format=csv,noheader,nounits"
+            ]
+            result = subprocess.check_output(cmd, stderr=subprocess.DEVNULL, timeout=5)
+            lines = result.decode().strip().split("\n")
+
+            for line in lines:
+                gpu_id_str, mem_used_mb_str, mem_total_mb_str = line.split(",")
+                gpu_id = int(gpu_id_str.strip())
+                mem_used_gb = int(mem_used_mb_str.strip()) / 1024  # MB → GB
+                mem_total_gb = int(mem_total_mb_str.strip()) / 1024  # MB → GB
+                gpu_usage[gpu_id] = (round(mem_used_gb, 2), round(mem_total_gb, 2))
+
+        except Exception as e:
+            logging.debug("Failed to get GPU memory info: %s", e)
+            return {}
+
+        return gpu_usage
+
+    @staticmethod
+    def get_all_storage_info() -> Dict[str, float]:
+        """
+        Get free storage space for all mounted drives.
+
+        Returns:
+            Dict[str, float]: Dictionary mapping mount point to free storage space in GB.
+        """
+        storage_info = {}
+
+        try:
+            # Get all disk partitions
+            partitions = psutil.disk_partitions()
+            
+            for partition in partitions:
+                try:
+                    # Get usage statistics for this partition
+                    usage = psutil.disk_usage(partition.mountpoint)
+                    
+                    # Convert bytes to GB
+                    free_gb = usage.free / (1024 ** 3)
+                    
+                    storage_info[partition.mountpoint] = round(free_gb, 2)
+                    
+                except PermissionError:
+                    # Skip drives that we can't access (common on Windows)
+                    logging.debug("Permission denied accessing %s", partition.mountpoint)
+                    continue
+                except Exception as e:
+                    logging.debug("Error getting storage info for %s: %s", partition.mountpoint, e)
+                    continue
+                    
+        except Exception as e:
+            logging.debug("Failed to get storage info: %s", e)
+            return {}
+
+        return storage_info
+
+    def get_stats(self) -> Tuple[float, int, float, float, Dict[int, tuple], Dict[str, float]]:
+        """
+        Collect current system resource statistics.
+
+        Returns:
+            Tuple[float, int, float, float, Dict[int, tuple], Dict[str, float]]: 
+            CPU usage %, CPU cores, RAM total GB, RAM used GB, GPU memory dict (used, total), Free storage dict
+        """
+        cpu_usage = psutil.cpu_percent(interval=1)
+        cpu_cores = psutil.cpu_count(logical=True)  # Total logical CPU cores
+
+        mem = psutil.virtual_memory()
+        ram_total = mem.total / (1024 ** 3)
+        ram_used = mem.used / (1024 ** 3)
+
+        gpu_usage = self.get_all_gpu_memory()
+        storage_info = self.get_all_storage_info()
+
+        return cpu_usage, cpu_cores, ram_total, ram_used, gpu_usage, storage_info
+
+    def _monitor_worker(self):
+        """
+        Worker function that runs in a separate thread to monitor and publish resources.
+        """
+        try:
+            from kafka import KafkaProducer
+
+            self._producer = KafkaProducer(
+                bootstrap_servers=self.kafka_bootstrap,
+                value_serializer=lambda v: json.dumps(v).encode("utf-8"),
+                retries=5,
+            )
+            logging.info("Kafka resource monitor started. Publishing to topic: %s", self.topic_name)
+
+        except ImportError:
+            logging.error("kafka-python not installed. Install with: pip install kafka-python")
+            return
+        except Exception as e:
+            logging.error("Failed to initialize Kafka producer: %s", e)
+            return
+
+        while not self._stop_event.is_set():
+            try:
+                cpu, cpu_cores, total, used, gpus, storage = self.get_stats()
+
+                # Format GPU info for output: {0: {"used_gb": x, "total_gb": y}, ...}
+                gpu_memory_gb = {k: {"used_gb": v[0], "total_gb": v[1]} for k, v in gpus.items()}
+                payload = {
+                    "instance_id": self.instance_id,
+                    "cpu_usage_percent": round(cpu, 2),
+                    "cpu_cores": cpu_cores,
+                    "ram_total_gb": round(total, 2),
+                    "ram_used_gb": round(used, 2),
+                    "gpu_memory_gb": gpu_memory_gb,  # dict: {0: {used_gb, total_gb}, ...}
+                    "free_storage_gb": storage,  # dict: {"/": 50.5, "C:": 123.4}
+                    "timestamp": datetime.now(timezone.utc).isoformat()
+                }
+
+                self._producer.send(self.topic_name, payload)
+                self._producer.flush()
+
+                logging.debug("Published resource stats: %s", payload)
+
+            except Exception as e:
+                logging.error("Error in resource monitor loop: %s", e)
+
+            # Wait for interval or until stop event is set
+            if self._stop_event.wait(self.interval_seconds):
+                break
+
+        # Cleanup
+        if self._producer:
+            try:
+                self._producer.close()
+            except Exception as e:
+                logging.debug("Error closing Kafka producer: %s", e)
+
+        logging.info("Kafka resource monitor stopped.")
+
+    @log_errors(raise_exception=False, log_error=True)
+    def start(self):
+        """
+        Start the resource monitoring thread.
+
+        Returns:
+            bool: True if started successfully, False otherwise.
+        """
+        if self._is_running:
+            logging.warning("Kafka resource monitor is already running.")
+            return False
+
+        self._stop_event.clear()
+        self._monitor_thread = threading.Thread(
+            target=self._monitor_worker,
+            daemon=True,
+            name="KafkaResourceMonitor"
+        )
+        self._monitor_thread.start()
+        self._is_running = True
+
+        logging.info("Started Kafka resource monitor thread.")
+        return True
+
+    @log_errors(raise_exception=False, log_error=True)
+    def stop(self, timeout: int = 10):
+        """
+        Stop the resource monitoring thread gracefully.
+
+        Args:
+            timeout: Maximum time to wait for thread to stop in seconds.
+
+        Returns:
+            bool: True if stopped successfully, False otherwise.
+        """
+        if not self._is_running:
+            logging.warning("Kafka resource monitor is not running.")
+            return False
+
+        logging.info("Stopping Kafka resource monitor...")
+        self._stop_event.set()
+
+        if self._monitor_thread and self._monitor_thread.is_alive():
+            self._monitor_thread.join(timeout=timeout)
+
+            if self._monitor_thread.is_alive():
+                logging.error("Kafka resource monitor thread did not stop within timeout.")
+                return False
+
+        self._is_running = False
+        logging.info("Kafka resource monitor stopped successfully.")
+        return True
+
+    def is_running(self) -> bool:
+        """
+        Check if the resource monitor is currently running.
+
+        Returns:
+            bool: True if running, False otherwise.
+        """
+        return self._is_running
+
+    def __enter__(self):
+        """Context manager entry."""
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.stop()