matrice-compute 0.1.1__tar.gz → 0.1.12__tar.gz

{matrice_compute-0.1.1 → matrice_compute-0.1.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: matrice_compute
-Version: 0.1.1
+Version: 0.1.12
 Summary: Common server utilities for Matrice.ai services
 Author-email: "Matrice.ai" <dipendra@matrice.ai>
 License-Expression: MIT

{matrice_compute-0.1.1 → matrice_compute-0.1.12}/matrice_compute.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: matrice_compute
-Version: 0.1.1
+Version: 0.1.12
 Summary: Common server utilities for Matrice.ai services
 Author-email: "Matrice.ai" <dipendra@matrice.ai>
 License-Expression: MIT

{matrice_compute-0.1.1 → matrice_compute-0.1.12}/src/matrice_compute/__init__.py

@@ -3,7 +3,7 @@
 
 from matrice_common.utils import dependencies_check
 
-dependencies_check(["docker", "psutil", "cryptography", "notebook", "aiohttp"])
+dependencies_check(["docker", "psutil", "cryptography", "notebook", "aiohttp", "kafka-python"])
 from matrice_compute.instance_manager import InstanceManager  # noqa: E402
 
 __all__ = ["InstanceManager"]

{matrice_compute-0.1.1 → matrice_compute-0.1.12}/src/matrice_compute/action_instance.py

@@ -348,9 +348,14 @@ class ActionInstance:
             "-v /var/run/docker.sock:/var/run/docker.sock" if mount_docker_sock else "",
         ]
         pypi_index = f"https://{'test.' if env != 'prod' else ''}pypi.org/simple/"
+        
         pkgs = ["matrice_common", "matrice"]
         pkgs.extend(extra_pkgs)
-        pip_install_matrice = f"pip install --upgrade --force-reinstall --index-url {pypi_index} {' '.join(pkgs)}"
+        if env == 'dev':
+            pkgs = [pkg + ">=1.0.0" for pkg in pkgs]
+            pip_install_matrice = f"pip install --pre --upgrade --force-reinstall --index-url {pypi_index} {' '.join(pkgs)}"
+        else:
+            pip_install_matrice = f"pip install --upgrade --force-reinstall --index-url {pypi_index} {' '.join(pkgs)}"
         pip_install_requirements = (
             "if [ -f requirements.txt ]; then pip install -r requirements.txt; fi "
         )
@@ -1490,6 +1495,12 @@ def kafka_setup_execute(self: ActionInstance):
     # Build the docker command directly to match user's pattern
     pypi_index = f"https://{'test.' if env != 'prod' else ''}pypi.org/simple/"
 
+    if env == 'dev':
+        pypi_index = f"https://test.pypi.org/simple/ --pre"
+        pkgs = f"matrice_common>=1.0.0 matrice>=1.0.0"
+    else:
+        pkgs = f"matrice_common matrice"
+
     cmd = (
         f"docker run -p {host_port}:{container_port} "
         f"{env_args} "
@@ -1499,7 +1510,7 @@ def kafka_setup_execute(self: ActionInstance):
         f"source venv/bin/activate && "
         f"/opt/kafka/bin/startup.sh & "
         f"if [ -f requirements.txt ]; then venv/bin/python3 -m pip install -r requirements.txt; fi && "
-        f"venv/bin/python3 -m pip install --upgrade --force-reinstall --index-url {pypi_index} matrice_common matrice && "
+        f"venv/bin/python3 -m pip install --upgrade --force-reinstall --index-url {pypi_index} {pkgs} && "
         f"sleep 20 && "
         f'venv/bin/python3 main.py {self.action_record_id} {host_port}"'
     )

{matrice_compute-0.1.1 → matrice_compute-0.1.12}/src/matrice_compute/actions_manager.py

@@ -196,6 +196,7 @@ class ActionsManager:
     def start_actions_manager(self) -> None:
         """Start the actions manager main loop."""
         while True:
+            waiting_time = self.poll_interval  # Default wait time
             try:
                 mem_usage = get_mem_usage()
                 logging.info("Memory usage: %d", mem_usage)

{matrice_compute-0.1.1 → matrice_compute-0.1.12}/src/matrice_compute/instance_manager.py

@@ -91,7 +91,7 @@ class InstanceManager:
         self.actions_resources_tracker = ActionsResourcesTracker(self.scaling)
         logging.info("InstanceManager initialized with actions resources tracker")
         self.poll_interval = 10
-        self.encryption_key = None
+        # Note: encryption_key is set in _setup_env_credentials
         logging.info("InstanceManager initialized.")
 
     @log_errors(default_return=None, raise_exception=True, log_error=True)
@@ -220,13 +220,13 @@ class InstanceManager:
                     "Error in shutdown_manager handle_shutdown: %s",
                     str(exc),
                 )
-            try:
-                self.scale_down_manager.auto_scaledown_actions()
-            except Exception as exc:
-                logging.error(
-                    "Error in scale_down_manager auto_scaledown_actions: %s",
-                    str(exc),
-                )
+            # try:
+            #     self.scale_down_manager.auto_scaledown_actions()
+            # except Exception as exc:
+            #     logging.error(
+            #         "Error in scale_down_manager auto_scaledown_actions: %s",
+            #         str(exc),
+            #     )
             try:
                 self.machine_resources_tracker.update_available_resources()
             except Exception as exc:

{matrice_compute-0.1.1 → matrice_compute-0.1.12}/src/matrice_compute/scaling.py

@@ -99,7 +99,7 @@ class Scaling:
             "Getting downscaled ids for instance %s",
             self.instance_id,
         )
-        path = f"/v1/scaling/down_scaled_ids/{self.instance_id}"
+        path = f"/v1/compute/down_scaled_ids/{self.instance_id}"
         resp = self.rpc.get(path=path)
         return self.handle_response(
             resp,
@@ -295,7 +295,7 @@ class Scaling:
             payload=payload,
             request_topic=self.kafka_config["api_request_topic"],
             response_topic=self.kafka_config["api_response_topic"],
-            timeout=600
+            timeout=60
         )
         if kafka_response_received:
             return data, error, message
@@ -347,7 +347,7 @@ class Scaling:
             payload=payload,
             request_topic=self.kafka_config["api_request_topic"],
             response_topic=self.kafka_config["api_response_topic"],
-            timeout=600
+            timeout=60
         )
         if kafka_response_received:
             return data, error, message
@@ -380,7 +380,7 @@ class Scaling:
             payload=payload,
             request_topic=self.kafka_config["api_request_topic"],
             response_topic=self.kafka_config["api_response_topic"],
-            timeout=600
+            timeout=60
         )
 
         if kafka_response_received:
@@ -427,7 +427,7 @@ class Scaling:
             payload=payload,
             request_topic=self.kafka_config["scaling_request_topic"],
             response_topic=self.kafka_config["scaling_response_topic"],
-            timeout=600
+            timeout=60
         )
 
         if kafka_response_received:
@@ -460,7 +460,7 @@ class Scaling:
                 payload=payload,
                 request_topic=self.kafka_config["api_request_topic"],
                 response_topic=self.kafka_config["api_response_topic"],
-                timeout=600
+                timeout=60
             )
 
         if kafka_response_received:
@@ -782,7 +782,7 @@ class Scaling:
         Returns:
             Tuple of (data, error, message) from API response
         """
-        path = f"/v1/scaling/get_internal_api_key/{action_id}/{self.instance_id}"
+        path = f"/v1/actions/get_internal_api_key/{action_id}/{self.instance_id}"
         resp = self.rpc.get(path=path)
         return self.handle_response(
             resp,
@@ -807,7 +807,7 @@ class Scaling:
             logging.error("%s: %s", message, error)
         return data, error, message
 
-    def _send_kafka_request(self, api, payload, request_topic, response_topic, timeout=600):
+    def _send_kafka_request(self, api, payload, request_topic, response_topic, timeout=60):
         """
         Helper to send a request to Kafka and wait for a response.
         Returns (data, error, message, kafka_response_received) where kafka_response_received is True if a response was received (even if error), False if transport error/timeout.
@@ -844,20 +844,27 @@ class Scaling:
             return None, f"Kafka producer error: {e}", "Kafka send failed", False
         try:
             start = time.time()
-            for message in consumer:
-                msg = message.value
-                if msg.get("correlationId") == correlation_id:
-                    consumer.close()
-                    # Always treat a received response as final, even if error
-                    return self.handle_kafka_response(
-                        msg,
-                        f"Fetched via Kafka for {api}",
-                        f"Kafka error response for {api}"
-                    ) + (True,)
-                if time.time() - start > timeout:
-                    break
+            while time.time() - start < timeout:
+                # Poll for messages with a short timeout to avoid blocking forever
+                message_batch = consumer.poll(timeout_ms=1000)
+                if message_batch:
+                    for topic_partition, messages in message_batch.items():
+                        for message in messages:
+                            print("trying to fetch message")
+                            msg = message.value
+                            if msg.get("correlationId") == correlation_id:
+                                consumer.close()
+                                # Always treat a received response as final, even if error
+                                return self.handle_kafka_response(
+                                    msg,
+                                    f"Fetched via Kafka for {api}",
+                                    f"Kafka error response for {api}"
+                                ) + (True,)
+                else:
+                    print(f"No messages received, waiting... ({time.time() - start:.1f}s/{timeout}s)")
+            
             consumer.close()
-            logging.warning("Kafka response timeout for %s", api)
+            logging.warning("Kafka response timeout for %s after %d seconds", api, timeout)
             return None, "Kafka response timeout", "Kafka response timeout", False
         except Exception as e:
             logging.error("Kafka consumer error: %s", e)