matrice-compute 0.1.44__py3-none-any.whl → 0.1.46__py3-none-any.whl

matrice_compute/action_instance.py

@@ -21,6 +21,7 @@ from matrice_compute.scaling import (
     Scaling,
 )
 from matrice_common.utils import log_errors
+from typing import cast
 
 
 class ActionInstance:
@@ -369,6 +370,8 @@ class ActionInstance:
             "MATRICE_SECRET_ACCESS_KEY": self.matrice_secret_access_key,
             "MATRICE_ACCESS_KEY_ID": self.matrice_access_key_id,
         }
+        if os.environ.get("MATRICE_BASE_URL"):
+            env_vars["MATRICE_BASE_URL"] = os.environ["MATRICE_BASE_URL"]
         if self.get_hugging_face_token(model_key):
             env_vars["HUGGING_FACE_ACCESS_TOKEN"] = self.get_hugging_face_token(
                 model_key
@@ -563,6 +566,7 @@ class ActionInstance:
                     action_id,
                     model_codebase_url,
                     model_codebase_requirements_url,
+                    scaling=self.scaling,
                 )
 
         # Setup Docker credentials
@@ -872,7 +876,8 @@ class ActionInstance:
             "bg-job-scheduler",
             "DKR_CMD",
             "OK",
-            f"Start docker container with command: {cmd.replace(self.matrice_access_key_id, 'MATRICE_ACCESS_KEY_ID').replace(self.matrice_secret_access_key, 'MATRICE_SECRET_ACCESS_KEY')}",
+            f"Start docker container with command: "
+            f"{cmd.replace(cast(str, self.matrice_access_key_id), 'MATRICE_ACCESS_KEY_ID').replace(cast(str, self.matrice_secret_access_key), 'MATRICE_SECRET_ACCESS_KEY')}",
         )
 
     @log_errors(raise_exception=False, log_error=False)
@@ -1140,6 +1145,8 @@ def database_setup_execute(self: ActionInstance):
         f"-v {dbPath}:{dbPath} "
         f"--name {self.action_record_id}_{self.action_type} "
         f"-v /var/run/docker.sock:/var/run/docker.sock "
+        f"-v /etc/matrice/service-config.yaml:/etc/matrice/service-config.yaml "
+        f'-e SERVICE_CONFIG_PATH="/etc/matrice/service-config.yaml" '
         f"-e ACTION_RECORD_ID={self.action_record_id} "
         f"-e MATRICE_ACCESS_KEY_ID={self.matrice_access_key_id} "
         f"-e MATRICE_SECRET_ACCESS_KEY={self.matrice_secret_access_key} "
@@ -1274,9 +1281,11 @@ def inference_ws_server_execute(self: ActionInstance):
     worker_cmd = (
         f"docker run -d --pull=always --net=host "
         f"--name {self.action_record_id}_{self.action_type} "
+        f"-v /etc/matrice/service-config.yaml:/etc/matrice/service-config.yaml "
         f'-e ENV="{os.environ.get("ENV", "prod")}" '
         f'-e MATRICE_SECRET_ACCESS_KEY="{self.matrice_secret_access_key}" '
         f'-e MATRICE_ACCESS_KEY_ID="{self.matrice_access_key_id}" '
+        f'-e SERVICE_CONFIG_PATH="/etc/matrice/service-config.yaml" '
         f' --restart=unless-stopped '
         f"{image} "
         f"./app "
@@ -1503,12 +1512,14 @@ def redis_setup_execute(self: ActionInstance):
     # bg-redis management container with --net=host (Port: 8082)
     cmd = (
         f"docker run --net=host "
-        f"-e REDIS_URL={shlex.quote(env_vars['REDIS_URL'])} "
-        f"-e REDIS_PASSWORD={shlex.quote(env_vars['REDIS_PASSWORD'])} "
-        f"-e MATRICE_ACCESS_KEY_ID={shlex.quote(self.matrice_access_key_id)} "
-        f"-e MATRICE_SECRET_ACCESS_KEY={shlex.quote(self.matrice_secret_access_key)} "
+        f"-e REDIS_URL={shlex.quote(cast(str, env_vars['REDIS_URL']))} "
+        f"-e REDIS_PASSWORD={shlex.quote(cast(str, env_vars['REDIS_PASSWORD']))} "
+        f"-e MATRICE_ACCESS_KEY_ID={shlex.quote(cast(str, self.matrice_access_key_id))} "
+        f"-e MATRICE_SECRET_ACCESS_KEY={shlex.quote(cast(str, self.matrice_secret_access_key))} "
+        f'-e SERVICE_CONFIG_PATH="/etc/matrice/service-config.yaml" '
         f"-e ENV={shlex.quote(os.environ.get('ENV', 'prod'))} "
         f"-v /var/run/docker.sock:/var/run/docker.sock "
+        f"-v /etc/matrice/service-config.yaml:/etc/matrice/service-config.yaml "
         f"--shm-size=30G --pull=always "
         f"{self.docker_container} "
         f"{self.action_record_id} "
@@ -1869,10 +1880,12 @@ def inference_tracker_setup_execute(self: ActionInstance):
         f"docker run -d --pull=always --net=host "
          f"--name {self.action_record_id}_{self.action_type} "
         f"-v matrice_myvol:/matrice_data "
+        f"-v /etc/matrice/service-config.yaml:/etc/matrice/service-config.yaml "
         f'-e ENV="{os.environ.get("ENV", "prod")}" '
         f'-e MATRICE_SECRET_ACCESS_KEY="{self.matrice_secret_access_key}" '
         f'-e MATRICE_ACCESS_KEY_ID="{self.matrice_access_key_id}" '
         f'-e ACTION_ID="{self.action_record_id}" '
+        f'-e SERVICE_CONFIG_PATH="/etc/matrice/service-config.yaml" '
         f' --restart=unless-stopped '
         f"{image}"
     )
@@ -1912,11 +1925,13 @@ def video_storage_setup_execute(self: ActionInstance):
     worker_cmd = (
         f"docker run -d --pull=always --net=host "
         f"--name {self.action_record_id}_{self.action_type} "
-        f"-v {media_storage_path}:/storage "  
+        f"-v {media_storage_path}:/storage " 
+        f"-v /etc/matrice/service-config.yaml:/etc/matrice/service-config.yaml "
         f'-e ENV="{os.environ.get("ENV", "prod")}" '
         f'-e MATRICE_SECRET_ACCESS_KEY="{self.matrice_secret_access_key}" '
         f'-e MATRICE_ACCESS_KEY_ID="{self.matrice_access_key_id}" '
         f'-e ACTION_ID="{self.action_record_id}" '
+        f'-e SERVICE_CONFIG_PATH="/etc/matrice/service-config.yaml" '
         f' --restart=unless-stopped '
         f"{image}"
     )

matrice_compute/actions_manager.py

@@ -3,6 +3,7 @@
 import logging
 import os
 import time
+from typing import Any
 from matrice_compute.action_instance import (
     ActionInstance,
 )
@@ -41,7 +42,7 @@ class ActionsManager:
         Returns:
             list: List of fetched actions
         """
-        actions = []
+        actions: list[Any] = []
         logging.info("Polling backend for new jobs")
         result = self.scaling.assign_jobs(has_gpu())
         if result is None:

matrice_compute/actions_scaledown_manager.py

@@ -37,6 +37,11 @@ class ActionsScaleDownManager:
         if down_scaled_jobs:
             for container in containers:
                 container_id = container.id
+                if container_id is None:
+                    logging.warning(
+                        "Skipping container with missing id while inspecting."
+                    )
+                    continue
                 inspect_data = self.docker_client.api.inspect_container(container_id)
                 action_record_id = next(
                     (arg for arg in inspect_data["Args"] if len(arg) == 24),

matrice_compute/instance_manager.py

@@ -6,6 +6,7 @@ import os
 import subprocess
 import threading
 import time
+from typing import Any, Optional
 from kafka import KafkaProducer
 from matrice_compute.actions_manager import ActionsManager
 from matrice_compute.actions_scaledown_manager import ActionsScaleDownManager
@@ -31,6 +32,22 @@ class InstanceManager:
 
     Now includes auto streaming capabilities for specified deployment IDs.
     """
+    # Instance attributes for type checking
+    scaling: Scaling
+    current_actions: dict[Any, Any]
+    actions_manager: ActionsManager
+    scale_down_manager: ActionsScaleDownManager
+    shutdown_manager: ShutdownManager
+    machine_resources_tracker: MachineResourcesTracker
+    actions_resources_tracker: ActionsResourcesTracker
+    kafka_resource_monitor: Optional[KafkaResourceMonitor]
+    container_resource_monitor: Optional[ContainerResourceMonitor]
+    compute_operations_handler: Optional[ComputeOperationsHandler]
+    poll_interval: int
+    container_monitor_thread: Optional[threading.Thread]
+    container_monitor_running: bool
+    container_kafka_producer: Optional[KafkaProducer]
+    encryption_key: str
 
     def __init__(
         self,
@@ -87,7 +104,7 @@ class InstanceManager:
             logging.info("InstanceManager updated Jupyter token")
         else:
             logging.warning("No Jupyter token found in environment variables")
-        self.current_actions = {}
+        self.current_actions: dict[Any, Any] = {}
         self.actions_manager = ActionsManager(self.scaling)
         logging.info("InstanceManager initialized with actions manager")
         self.scale_down_manager = ActionsScaleDownManager(self.scaling)
@@ -100,6 +117,7 @@ class InstanceManager:
         logging.info("InstanceManager initialized with actions resources tracker")
 
         # Initialize Kafka resource monitor using the same internal Kafka as scaling
+        self.kafka_resource_monitor = None
         try:
             kafka_bootstrap = self.scaling.get_kafka_bootstrap_servers()
             self.kafka_resource_monitor = KafkaResourceMonitor(
@@ -113,6 +131,7 @@ class InstanceManager:
             self.kafka_resource_monitor = None
 
         # Initialize Container resource monitor using the same internal Kafka as scaling
+        self.container_resource_monitor = None
         try:
             kafka_bootstrap = self.scaling.get_kafka_bootstrap_servers()
             self.container_resource_monitor = ContainerResourceMonitor(
@@ -127,13 +146,14 @@ class InstanceManager:
 
         # Initialize Compute Operations Handler for event-driven operations
         # Uses EventListener from matrice_common for simplified Kafka consumption
+        self.compute_operations_handler = None
         try:
-            instance_id = os.environ.get("INSTANCE_ID")
+            instance_id_env = os.environ.get("INSTANCE_ID") or ""
             self.compute_operations_handler = ComputeOperationsHandler(
                 actions_manager=self.actions_manager,
                 session=self.session,
                 scaling=self.scaling,
-                instance_id=instance_id
+                instance_id=instance_id_env
             )
             logging.info("InstanceManager initialized with Compute Operations Handler for instance ID: %s", instance_id)
         except (ValueError, Exception) as e:
@@ -225,10 +245,10 @@ class InstanceManager:
             raise Exception(
                 "SERVICE_PROVIDER and INSTANCE_ID must be set as environment variables or passed as arguments"
             )
-        self.encryption_key = manual_instance_info["MATRICE_ENCRYPTION_KEY"]
+        self.encryption_key = str(manual_instance_info["MATRICE_ENCRYPTION_KEY"] or "")
 
-        access_key = manual_instance_info["MATRICE_ACCESS_KEY_ID"]
-        secret_key = manual_instance_info["MATRICE_SECRET_ACCESS_KEY"]
+        access_key = str(manual_instance_info["MATRICE_ACCESS_KEY_ID"] or "")
+        secret_key = str(manual_instance_info["MATRICE_SECRET_ACCESS_KEY"] or "")
 
         if (  # Keys are not encrypted
             self.encryption_key

matrice_compute/instance_utils.py

@@ -15,9 +15,9 @@ from cryptography.hazmat.primitives.ciphers import (
 )
 from cryptography.hazmat.backends import default_backend
 from matrice_common.utils import log_errors
+from typing import Optional, Tuple
 
-
-def get_instance_info(service_provider: str = None, instance_id: str = None) -> tuple:
+def get_instance_info(service_provider: Optional[str] = None, instance_id: Optional[str] = None) -> tuple:
     """
     Get instance provider and ID information.
 
@@ -450,7 +450,7 @@ def get_disk_space_usage() -> list:
 
 
 @log_errors(default_return=None, raise_exception=False)
-def get_max_file_system() -> str:
+def get_max_file_system() -> Optional[str]:
     """
     Get filesystem with maximum available space.
 
@@ -1060,7 +1060,7 @@ def get_decrypted_access_key_pair(
     enc_access_key: str,
     enc_secret_key: str,
     encryption_key: str = "",
-) -> tuple:
+) -> Tuple[Optional[str], Optional[str]]:
     """
     Get decrypted access key pair.
 
@@ -1072,7 +1072,7 @@ def get_decrypted_access_key_pair(
     Returns:
         tuple: (access_key, secret_key) strings
     """
-    encryption_key = encryption_key or os.environ.get("MATRICE_ENCRYPTION_KEY")
+    encryption_key = encryption_key or os.environ.get("MATRICE_ENCRYPTION_KEY", "")
     if not encryption_key:
         logging.warning("Encryption key is not set, Will assume that the keys are not encrypted")
         return enc_access_key, enc_secret_key
@@ -1107,7 +1107,7 @@ def get_encrypted_access_key_pair(
     access_key: str,
     secret_key: str,
     encryption_key: str = "",
-) -> tuple:
+) -> Tuple[Optional[str], Optional[str]]:
     """
     Get encrypted access key pair.
 
@@ -1119,7 +1119,7 @@ def get_encrypted_access_key_pair(
     Returns:
         tuple: (encrypted_access_key, encrypted_secret_key) strings
     """
-    encryption_key = encryption_key or os.environ.get("MATRICE_ENCRYPTION_KEY")
+    encryption_key = encryption_key or os.environ.get("MATRICE_ENCRYPTION_KEY", "")
     if not encryption_key:
         logging.warning("Encryption key is not set, returning unencrypted keys")
         return access_key, secret_key
@@ -1155,7 +1155,7 @@ def get_encrypted_access_key_pair(
     
     return encoded_access_key, encoded_secret_key
 
-def _get_private_ip() -> str:
+def _get_private_ip() -> Optional[str]:
     """
     Get the actual private/LAN IP address using UDP socket trick.
     This works reliably even in Docker, NAT, VPN, etc.