matrice-compute 0.1.1__py3-none-any.whl

matrice_compute/shutdown_manager.py

@@ -0,0 +1,314 @@
+"""Module providing shutdown_manager functionality."""
+
+import logging
+import time
+import os
+import sys
+import platform
+import subprocess
+import signal
+from typing import Optional, Tuple
+from matrice_common.utils import log_errors
+from matrice_compute.scaling import Scaling
+
+
+class ShutdownManager:
+    """Class for managing compute instance shutdown."""
+
+    def __init__(self, scaling: Scaling) -> None:
+        """
+        Initialize ShutdownManager.
+
+        Args:
+            scaling (Scaling): Scaling instance to manage shutdown.
+        """
+        self.scaling = scaling
+        self.launch_time: float = time.time()
+        self.last_no_queued_time: Optional[float] = None
+        self.shutdown_threshold: int = 500  # Idle time threshold in seconds
+        self.launch_duration: int = 1  # Launch duration in minutes
+        self.instance_source: str = "auto"
+        self.encryption_key: Optional[str] = None
+        self.reserved_instance: Optional[bool] = None
+        self.shutdown_attempts: int = 0
+        self.max_shutdown_attempts: int = 3
+        self.force_shutdown_attempts: int = 0
+        self.max_force_shutdown_attempts: int = 2
+        self.launch_duration_seconds: int = self.launch_duration * 60  # Convert minutes to seconds
+        self._load_shutdown_configuration()
+
+    @log_errors(raise_exception=False, log_error=True)
+    def _load_shutdown_configuration(self) -> None:
+        """
+        Load shutdown configuration from AWS secrets and initialize parameters.
+
+        This method retrieves shutdown configuration details from the scaling service
+        and updates the instance attributes accordingly.
+        """
+        response, error, message = self.scaling.get_shutdown_details()
+        if error is None:
+            self.shutdown_threshold = response.get("shutdownThreshold", 500)
+            self.launch_duration = response.get("launchDuration", 1)
+            self.instance_source = response.get("instanceSource", "auto")
+            self.encryption_key = response.get("encryptionKey")
+        self.launch_duration_seconds = self.launch_duration * 60  # Convert minutes to seconds
+        self.reserved_instance = self.instance_source == "reserved"
+        logging.info(
+            "Loaded shutdown configuration: threshold=%s, duration=%s, source=%s, reserved=%s",
+            self.shutdown_threshold,
+            self.launch_duration,
+            self.instance_source,
+            self.reserved_instance,
+        )
+
+    def _check_root(self) -> bool:
+        """
+        Check if the current process has root privileges.
+
+        Returns:
+            bool: True if the process has root privileges, False otherwise.
+        """
+        if hasattr(os, "geteuid") and os.geteuid() != 0:
+            logging.error("Shutdown requires root privileges.")
+            return False
+        return True
+
+    def _execute_shutdown_command(self) -> bool:
+        """
+        Execute system shutdown command with multiple fallbacks.
+
+        This method attempts to shut down the system using various commands
+        depending on the operating system.
+
+        Returns:
+            bool: True if any shutdown command succeeded, False otherwise.
+        """
+        self._check_root()
+        system = platform.system().lower()
+        
+        # Define shutdown commands in order of preference (most graceful first)
+        shutdown_commands = []
+        
+        if system == "linux":
+            shutdown_commands = [
+                ["shutdown", "now"],  # Standard Linux shutdown
+                ["systemctl", "poweroff"],  # Systemd poweroff
+                ["systemctl", "poweroff", "--force"],  # Force systemd poweroff
+                ["halt", "-f"],  # Force halt
+                ["poweroff", "-f"],  # Force poweroff
+                ["init", "0"],  # Init level 0 (shutdown)
+                ["telinit", "0"],  # Alternative init command
+            ]
+        elif system == "windows":
+            shutdown_commands = [
+                ["shutdown", "/s", "/t", "0"],  # Windows shutdown
+                ["shutdown", "/s", "/f", "/t", "0"],  # Windows force shutdown
+                ["shutdown", "/p"],  # Windows immediate poweroff
+            ]
+        elif system == "darwin":  # macOS
+            shutdown_commands = [
+                ["shutdown", "-h", "now"],  # macOS shutdown
+                ["halt"],  # macOS halt
+                ["sudo", "shutdown", "-h", "now"],  # Sudo shutdown
+            ]
+        else:
+            # Generic Unix-like fallbacks
+            shutdown_commands = [
+                ["shutdown", "-h", "now"],
+                ["halt"],
+                ["poweroff"],
+                ["init", "0"],
+            ]
+        
+        # Try each command in sequence
+        for cmd in shutdown_commands:
+            try:
+                logging.info("Attempting shutdown with command: %s", " ".join(cmd))
+                result = subprocess.run(
+                    cmd, 
+                    capture_output=True, 
+                    text=True, 
+                    timeout=30,
+                    check=False
+                )
+                
+                if result.returncode == 0:
+                    logging.info("Shutdown command succeeded: %s", " ".join(cmd))
+                    return True
+                else:
+                    logging.warning(
+                        "Shutdown command failed with return code %d: %s. STDERR: %s", 
+                        result.returncode, 
+                        " ".join(cmd),
+                        result.stderr
+                    )
+            except subprocess.TimeoutExpired:
+                logging.warning("Shutdown command timed out: %s", " ".join(cmd))
+            except FileNotFoundError:
+                logging.warning("Shutdown command not found: %s", " ".join(cmd))
+            except Exception as e:
+                logging.warning("Shutdown command failed: %s. Error: %s", " ".join(cmd), str(e))
+        
+        # If all standard commands failed, try more aggressive methods
+        return self._try_aggressive_shutdown()
+
+    def _try_aggressive_shutdown(self) -> bool:
+        """
+        Try more aggressive shutdown methods when standard commands fail.
+
+        Returns:
+            bool: True if any aggressive shutdown method succeeded, False otherwise.
+        """
+        logging.warning("Standard shutdown commands failed, trying aggressive methods")
+        try:
+            system = platform.system().lower()
+            if system == "linux":
+                try:
+                    subprocess.run("echo 1 > /proc/sys/kernel/sysrq", shell=True, check=False, timeout=5)
+                except Exception:
+                    pass
+
+                aggressive_commands = [
+                    "sync",
+                    "echo 4 > /proc/acpi/sleep",
+                    "echo s > /proc/sysrq-trigger",
+                    "echo u > /proc/sysrq-trigger",
+                    "echo o > /proc/sysrq-trigger",
+                    "echo b > /proc/sysrq-trigger",
+                ]
+                for cmd in aggressive_commands:
+                    try:
+                        logging.info("Trying aggressive shutdown: %s", cmd)
+                        result = subprocess.run(cmd, shell=True, check=False, timeout=10)
+                        if result.returncode == 0:
+                            logging.info("Aggressive shutdown command succeeded")
+                            time.sleep(2)
+                            return True
+                    except Exception as e:
+                        logging.debug("Aggressive command failed: %s", str(e))
+        except Exception as e:
+            logging.error("Error in aggressive shutdown methods: %s", str(e))
+        return False
+    
+
+    @log_errors(raise_exception=True, log_error=True)
+    def do_cleanup_and_shutdown(self) -> bool:
+        """Clean up resources and shut down the instance.
+        
+        This method attempts a coordinated shutdown with multiple fallback strategies:
+        1. API call to notify the scaling service
+        2. Graceful OS shutdown command
+        3. Aggressive shutdown methods if needed
+        4. Emergency forced shutdown as last resort
+        
+        Returns:
+            bool: True if shutdown was initiated successfully, False otherwise
+        """
+        max_retries = self.max_shutdown_attempts
+        
+        for attempt in range(1, max_retries + 1):
+            try:
+                logging.info("Shutdown attempt %d of %d", attempt, max_retries)
+                
+                # Step 1: Notify scaling service of shutdown
+                logging.info("Notifying scaling service of instance shutdown")
+                try:
+                    response = self.scaling.stop_instance()
+                    
+                    # Handle case where stop_instance returns None or unexpected format
+                    if response is None:
+                        result, error, message = None, "API returned None", "No response from stop_instance API"
+                    elif isinstance(response, tuple) and len(response) == 3:
+                        result, error, message = response
+                    else:
+                        result, error, message = None, "Invalid response format", f"Unexpected response format: {response}"
+                        
+                except Exception as api_error:
+                    result, error, message = None, str(api_error), "Exception during API call"
+                
+                if error:
+                    logging.error("Failed to notify scaling service (attempt %d): %s", attempt, error)
+                    if attempt < max_retries:
+                        logging.info("Retrying in 5 seconds...")
+                        time.sleep(5)
+                        continue
+                    else:
+                        logging.warning("Proceeding with shutdown despite API notification failure")
+                else:
+                    logging.info("Scaling service notified successfully: %s", message)
+                
+                # Step 2: Attempt graceful system shutdown
+                logging.info("Initiating graceful system shutdown")
+                shutdown_success = self._execute_shutdown_command()
+                
+                if shutdown_success:
+                    logging.info("Graceful shutdown command executed successfully")
+                    # Give the system time to process the shutdown
+                    time.sleep(10)
+                    # If we reach here, graceful shutdown may have failed
+                    logging.warning("System did not shut down gracefully, trying aggressive methods")
+                
+                # Step 3: Try aggressive shutdown methods
+                logging.warning("Attempting aggressive shutdown methods")
+                aggressive_success = self._try_aggressive_shutdown()
+                
+                if aggressive_success:
+                    logging.info("Aggressive shutdown initiated")
+                    time.sleep(5)
+                
+                return True
+                
+            except Exception as e:
+                logging.error("Critical error during shutdown attempt %d: %s", attempt, str(e)) 
+        return False
+
+    @log_errors(raise_exception=False, log_error=True)
+    def handle_shutdown(self, tasks_running: bool) -> None:
+        """Check idle time and trigger shutdown if threshold is exceeded.
+
+        Args:
+            tasks_running: Boolean indicating if there are running tasks
+        """
+        # CRITICAL: Check if this is a reserved instance that should not be shut down
+        # if self.reserved_instance:
+        #     logging.debug("Reserved instance detected, skipping shutdown check")
+        #     return
+
+        # Update idle time tracking
+        if tasks_running:
+            self.last_no_queued_time = None
+            logging.info("Tasks are running, resetting idle timer")
+        elif self.last_no_queued_time is None:
+            self.last_no_queued_time = time.time()
+            logging.info("No tasks running, starting idle timer")
+
+        if self.last_no_queued_time is not None:
+            idle_time = time.time() - self.last_no_queued_time
+            launch_time_passed = (time.time() - self.launch_time) > self.launch_duration_seconds
+
+            # Log current status
+            logging.info(
+                "Time since last action: %s seconds. Time left to shutdown: %s seconds.",
+                idle_time,
+                max(0, self.shutdown_threshold - idle_time),
+            )
+
+            # Check if we should shut down
+            if idle_time <= self.shutdown_threshold:
+                return
+
+            if not launch_time_passed:
+                logging.info(
+                "Instance not shutting down yet. Launch duration: %s seconds, elapsed: %s seconds",
+                    self.launch_duration_seconds,
+                    time.time() - self.launch_time,
+                )
+                return
+
+            logging.info(
+                "Idle time %s seconds exceeded threshold %s seconds. Shutting down.",
+                idle_time,
+                self.shutdown_threshold
+            )
+
+            self.do_cleanup_and_shutdown()

matrice_compute/task_utils.py

@@ -0,0 +1,77 @@
+"""Module providing task_utils functionality."""
+
+import os
+import shutil
+import urllib.request
+import zipfile
+from typing import Optional
+from matrice_common.utils import log_errors
+
+
+@log_errors(raise_exception=True, log_error=True)
+def setup_workspace_and_run_task(
+    work_fs: str,
+    action_id: str,
+    model_codebase_url: str,
+    model_codebase_requirements_url: Optional[str] = None,
+) -> None:
+    """Set up workspace and run task with provided parameters.
+
+    Args:
+        work_fs (str): Working filesystem path.
+        action_id (str): Unique identifier for the action.
+        model_codebase_url (str): URL to download model codebase from.
+        model_codebase_requirements_url (Optional[str]): URL to download requirements from. Defaults to None.
+
+    Returns:
+        None
+    """
+    workspace_dir = f"{work_fs}/{action_id}"
+    codebase_zip_path = f"{workspace_dir}/file.zip"
+    requirements_txt_path = f"{workspace_dir}/requirements.txt"
+    if os.path.exists(workspace_dir):
+        return
+    os.makedirs(workspace_dir, exist_ok=True)
+    
+    # Download codebase ZIP file
+    urllib.request.urlretrieve(model_codebase_url, codebase_zip_path)
+    
+    # Extract ZIP file with overwrite
+    with zipfile.ZipFile(codebase_zip_path, 'r') as zip_ref:
+        zip_ref.extractall(workspace_dir)
+    
+    # Move files from subdirectories to workspace root (equivalent to rsync -av)
+    for root, dirs, files in os.walk(workspace_dir):
+        # Skip the workspace_dir itself to avoid moving files to themselves
+        if root == workspace_dir:
+            continue
+        
+        for file in files:
+            src_file = os.path.join(root, file)
+            dst_file = os.path.join(workspace_dir, file)
+            
+            # If destination file exists, overwrite it (equivalent to rsync -av behavior)
+            if os.path.exists(dst_file):
+                os.remove(dst_file)
+            
+            shutil.move(src_file, dst_file)
+        
+        # Remove empty directories after moving files
+        for dir_name in dirs:
+            dir_path = os.path.join(root, dir_name)
+            if os.path.exists(dir_path) and not os.listdir(dir_path):
+                os.rmdir(dir_path)
+    
+    # Clean up any remaining empty subdirectories
+    for root, dirs, files in os.walk(workspace_dir, topdown=False):
+        if root == workspace_dir:
+            continue
+        if not files and not dirs:
+            try:
+                os.rmdir(root)
+            except OSError:
+                pass  # Directory might not be empty or might not exist
+    
+    # Download requirements.txt if URL is provided
+    if model_codebase_requirements_url:
+        urllib.request.urlretrieve(model_codebase_requirements_url, requirements_txt_path)

matrice_compute-0.1.1.dist-info/METADATA

@@ -0,0 +1,28 @@
+Metadata-Version: 2.4
+Name: matrice_compute
+Version: 0.1.1
+Summary: Common server utilities for Matrice.ai services
+Author-email: "Matrice.ai" <dipendra@matrice.ai>
+License-Expression: MIT
+Keywords: matrice,common,utilities,pyarmor,obfuscated
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: OS Independent
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: MacOS
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Typing :: Typed
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE.txt
+Dynamic: license-file
+Dynamic: requires-python
+
+## matrice\_instance\_manager
+

matrice_compute-0.1.1.dist-info/RECORD

@@ -0,0 +1,17 @@
+matrice_compute/__init__.py,sha256=bV-Vtb0-Sa3K1uXaAc8qP7NK6pf354OU1K19aQSt3VI,288
+matrice_compute/action_instance.py,sha256=CGRvWzkYeMHP-_GGqrWAgQ_4suqrFovlWH-NkY30p9g,57600
+matrice_compute/actions_manager.py,sha256=Nbosuf6kFisl2vbsgd5rZiYjdX4za8X588aut0OoKpI,7713
+matrice_compute/actions_scaledown_manager.py,sha256=pJ0nduNwHWZ10GnqJNx0Ok7cVWabQ_M8E2Vb9pH3A_k,2002
+matrice_compute/instance_manager.py,sha256=mGA2JsvRPWGQESKcZoc0k_xCB40fmsQDv2ojCT3kf1g,10060
+matrice_compute/instance_utils.py,sha256=tIFVUi8HJPy4GY-jtfVx2zIgmXNta7s3jCIRzBga1hI,21977
+matrice_compute/prechecks.py,sha256=W9YmNF3RcLhOf4U8WBlExvFqDw1aGWSNTlJtA73lbDQ,17196
+matrice_compute/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+matrice_compute/resources_tracker.py,sha256=My26LPglDHcQcTkxxiXwpfdqkpEAt3clrqJ-k1fAl1M,17878
+matrice_compute/scaling.py,sha256=UqxvSZtypELWmD4EBKS1X-vX3tJnSpgOO6_b4XOe2Uk,31175
+matrice_compute/shutdown_manager.py,sha256=0MYV_AqygqR9NEntYf7atUC-PbWXyNkm1f-8c2aizgA,13234
+matrice_compute/task_utils.py,sha256=ML9uTrYQiWgEMJitYxoGlVOa9KUXNKV_WqnousOTK6k,2762
+matrice_compute-0.1.1.dist-info/licenses/LICENSE.txt,sha256=_uQUZpgO0mRYL5-fPoEvLSbNnLPv6OmbeEDCHXhK6Qc,1066
+matrice_compute-0.1.1.dist-info/METADATA,sha256=zdb-JyIdO6bT2yko4QB9LgECmENwsu4TEkHlXpowSJI,1037
+matrice_compute-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+matrice_compute-0.1.1.dist-info/top_level.txt,sha256=63Plr3L1GzBUWZO5JZaFkiv8IcB10xUPU-9w3i6ptvE,16
+matrice_compute-0.1.1.dist-info/RECORD,,

matrice_compute-0.1.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

matrice_compute-0.1.1.dist-info/licenses/LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Matrice.ai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

matrice_compute-0.1.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+matrice_compute