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

matrice_compute/instance_manager.py

@@ -3,10 +3,13 @@
 import json
 import logging
 import os
+import subprocess
 import threading
 import time
+from kafka import KafkaProducer
 from matrice_compute.actions_manager import ActionsManager
 from matrice_compute.actions_scaledown_manager import ActionsScaleDownManager
+from matrice_compute.compute_operations_handler import ComputeOperationsHandler
 from matrice_compute.instance_utils import (
     get_instance_info,
     get_decrypted_access_key_pair,
@@ -14,6 +17,7 @@ from matrice_compute.instance_utils import (
 from matrice_compute.resources_tracker import (
     MachineResourcesTracker,
     ActionsResourcesTracker,
+    KafkaResourceMonitor,
 )
 from matrice_compute.scaling import Scaling
 from matrice_compute.shutdown_manager import ShutdownManager
@@ -90,10 +94,55 @@ class InstanceManager:
         logging.info("InstanceManager initialized with machine resources tracker")
         self.actions_resources_tracker = ActionsResourcesTracker(self.scaling)
         logging.info("InstanceManager initialized with actions resources tracker")
+
+        # Initialize Kafka resource monitor using the same internal Kafka as scaling
+        try:
+            kafka_bootstrap = self.scaling.get_kafka_bootstrap_servers()
+            self.kafka_resource_monitor = KafkaResourceMonitor(
+                instance_id=os.environ.get("INSTANCE_ID"),
+                kafka_bootstrap=kafka_bootstrap,
+                interval_seconds=60
+            )
+            logging.info("InstanceManager initialized with Kafka resource monitor using internal Kafka: %s", kafka_bootstrap)
+        except (ValueError, Exception) as e:
+            logging.warning("Failed to initialize Kafka resource monitor: %s", e)
+            self.kafka_resource_monitor = None
+
+        # Initialize Compute Operations Handler for event-driven operations
+        # Uses EventListener from matrice_common for simplified Kafka consumption
+        try:
+            instance_id = os.environ.get("INSTANCE_ID")
+            self.compute_operations_handler = ComputeOperationsHandler(
+                actions_manager=self.actions_manager,
+                session=self.session,
+                scaling=self.scaling,
+                instance_id=instance_id
+            )
+            logging.info("InstanceManager initialized with Compute Operations Handler for instance ID: %s", instance_id)
+        except (ValueError, Exception) as e:
+            logging.warning("Failed to initialize Compute Operations Handler: %s", e)
+            self.compute_operations_handler = None
+
         self.poll_interval = 10
         # Note: encryption_key is set in _setup_env_credentials
+        
+        # Initialize container monitoring
+        self.container_monitor_thread = None
+        self.container_monitor_running = False
+        self.container_kafka_producer = None
+        
         logging.info("InstanceManager initialized.")
 
+        # report the resources at startup
+        try:
+            self.scaling.report_architecture_info()
+            logging.info("InstanceManager reported initial resources.")
+        except Exception as exc:
+            logging.error(
+                "Error reporting initial resources: %s",
+                str(exc),
+            )
+
     @log_errors(default_return=None, raise_exception=True, log_error=True)
     def _setup_env_credentials(
         self,
@@ -228,13 +277,13 @@ class InstanceManager:
             #         "Error in scale_down_manager auto_scaledown_actions: %s",
             #         str(exc),
             #     )
-            try:
-                self.machine_resources_tracker.update_available_resources()
-            except Exception as exc:
-                logging.error(
-                    "Error in machine_resources_tracker update_available_resources: %s",
-                    str(exc),
-                )
+            # try:
+            #     self.machine_resources_tracker.update_available_resources()
+            # except Exception as exc:
+            #     logging.error(
+            #         "Error in machine_resources_tracker update_available_resources: %s",
+            #         str(exc),
+            #     )
             try:
                 self.actions_resources_tracker.update_actions_resources()
             except Exception as exc:
@@ -245,6 +294,130 @@ class InstanceManager:
 
             time.sleep(self.poll_interval)
 
+    @log_errors(raise_exception=False, log_error=True)
+    def start_container_status_monitor(self):
+        """Start the background container status monitoring."""
+        if self.container_monitor_running:
+            logging.info("Container status monitor is already running")
+            return
+        
+        self.container_monitor_running = True
+        self.container_monitor_thread = threading.Thread(
+            target=self._container_status_monitor_worker,
+            daemon=True,
+            name="ContainerStatusMonitor"
+        )
+        self.container_monitor_thread.start()
+        logging.info("Started container status monitoring thread")
+
+    @log_errors(raise_exception=False, log_error=True)
+    def stop_container_status_monitor(self):
+        """Stop the background container status monitoring."""
+        if not self.container_monitor_running:
+            return
+        
+        logging.info("Stopping container status monitor...")
+        self.container_monitor_running = False
+        
+        if self.container_monitor_thread:
+            self.container_monitor_thread.join(timeout=10)
+        
+        if self.container_kafka_producer:
+            self.container_kafka_producer.close()
+            self.container_kafka_producer = None
+        
+        logging.info("Container status monitor stopped")
+
+    def _container_status_monitor_worker(self):
+        """Background worker function that monitors container status."""
+        # Initialize Kafka producer
+        try:
+            if self.scaling.enable_kafka:
+                bootstrap_servers = self.scaling.get_kafka_bootstrap_servers()
+                self.container_kafka_producer = KafkaProducer(
+                    bootstrap_servers=bootstrap_servers,
+                    value_serializer=lambda v: json.dumps(v).encode("utf-8"),
+                    max_block_ms=5000  # Timeout if Kafka is down
+                )
+                logging.info("Container status monitor: Kafka producer initialized")
+            else:
+                logging.warning("Container status monitor: Kafka is disabled, no monitoring will be performed")
+                return
+        except Exception as e:
+            logging.error("Container status monitor: Failed to initialize Kafka producer: %s", str(e))
+            return
+        
+        instance_id = os.environ.get("INSTANCE_ID")
+        topic_name = "compute_container_status"
+        
+        logging.info("Container status monitor started for instance: %s", instance_id)
+        
+        while self.container_monitor_running:
+            try:
+                # Get container status using docker ps -a
+                result = subprocess.run(
+                    ["docker", "ps", "-a", "--format", "json"],
+                    capture_output=True,
+                    text=True,
+                    timeout=30
+                )
+                
+                if result.returncode != 0:
+                    logging.error("Container status monitor: docker ps command failed: %s", result.stderr)
+                    time.sleep(30)  # Wait before retrying
+                    continue
+                
+                # Parse container information
+                containers = []
+                if result.stdout.strip():
+                    for line in result.stdout.strip().split('\n'):
+                        try:
+                            container_info = json.loads(line)
+                            containers.append({
+                                "container_id": container_info.get("ID", ""),
+                                "image": container_info.get("Image", ""),
+                                "command": container_info.get("Command", ""),
+                                "created": container_info.get("CreatedAt", ""),
+                                "status": container_info.get("Status", ""),
+                                "ports": container_info.get("Ports", ""),
+                                "names": container_info.get("Names", ""),
+                                "size": container_info.get("Size", ""),
+                                "state": container_info.get("State", ""),
+                                "labels": container_info.get("Labels", "")
+                            })
+                        except json.JSONDecodeError as e:
+                            logging.warning("Container status monitor: Failed to parse container info: %s", str(e))
+                            continue
+                
+                # Prepare message for Kafka
+                status_message = {
+                    "timestamp": time.time(),
+                    "instance_id": instance_id,
+                    "container_count": len(containers),
+                    "containers": containers
+                }
+                
+                # Send to Kafka
+                if self.container_kafka_producer:
+                    try:
+                        self.container_kafka_producer.send(topic_name, status_message)
+                        logging.debug("Container status monitor: Sent status for %d containers", len(containers))
+                    except Exception as e:
+                        logging.error("Container status monitor: Failed to send to Kafka: %s", str(e))
+                
+            except subprocess.TimeoutExpired:
+                logging.error("Container status monitor: docker ps command timed out")
+            except Exception as e:
+                logging.error("Container status monitor: Unexpected error: %s", str(e))
+            
+            # Wait 30 seconds before next check
+            for _ in range(30):
+                if not self.container_monitor_running:
+                    break
+                time.sleep(1)
+        
+        logging.info("Container status monitor worker stopped")
+
     @log_errors(default_return=(None, None), raise_exception=True)
     def start(self) -> tuple:
         """Start the instance manager threads.
@@ -252,6 +425,29 @@ class InstanceManager:
         Returns:
             tuple: (instance_manager_thread, actions_manager_thread)
         """
+        # Start Kafka resource monitor in background thread
+        if self.kafka_resource_monitor:
+            try:
+                self.kafka_resource_monitor.start()
+                logging.info("Started Kafka resource monitor")
+            except Exception as exc:
+                logging.error("Failed to start Kafka resource monitor: %s", str(exc))
+
+        # Start Compute Operations Handler in background thread
+        if self.compute_operations_handler:
+            try:
+                self.compute_operations_handler.start()
+                logging.info("Started Compute Operations Handler")
+            except Exception as exc:
+                logging.error("Failed to start Compute Operations Handler: %s", str(exc))
+
+        # Start Container Status Monitor in background thread
+        try:
+            self.start_container_status_monitor()
+            logging.info("Started Container Status Monitor")
+        except Exception as exc:
+            logging.error("Failed to start Container Status Monitor: %s", str(exc))
+
         # Create and start threads
         instance_manager_thread = threading.Thread(
             target=self.start_instance_manager,

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
@@ -402,8 +404,13 @@ class ActionsResourcesTracker:
                 new_args.extend(x.replace('"', "").replace("'", "") for x in arg.split(" "))
             return new_args
 
-        args_24 = [arg for arg in remove_quotation_marks(inspect_data["Args"]) if len(arg) == 24 and "pypi" not in arg]
-        action_record_id = args_24[-1] if args_24 else None
+        def is_valid_objectid(s: str) -> bool:
+            """Check if string is a valid MongoDB ObjectId (24 hex characters)"""
+            s = s.strip()
+            return len(s) == 24 and all(c in '0123456789abcdefABCDEF' for c in s)
+        
+        valid_objectids = [arg for arg in remove_quotation_marks(inspect_data["Args"]) if is_valid_objectid(arg)]
+        action_record_id = valid_objectids[-1] if valid_objectids else None
         if not action_record_id:
             logging.debug("No valid action_id found for the container. Container ID: %s, Args: %s", container.id, inspect_data["Args"])
         duration = calculate_time_difference(start_time, finish_time)
@@ -575,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()