feature: added threading to ventis deploy | bug fix: fixed schema misalignment for runtime_information table

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "grpcio-tools",
     "redis",
     "sqlalchemy",
+    "psycopg[binary]",
     "pyyaml",
     "flask",
 ]

requirements.txt

@@ -7,3 +7,4 @@ flask
 ipdb
 ipython
 sqlalchemy
+psycopg[binary]

tests/test_instance_manager_runtime.py

@@ -249,16 +249,23 @@ def test_manager_keeps_ec2_runtime_boundary_behavior(self):
             controller.redis_containers = {"10.0.0.30": "redis-box"}
             manager.remove_instance("EC2:Remote:0")
 
-        runtime.provision_instance.assert_called_once_with(
-            {
-                "name": "Remote",
-                "provider": "EC2",
-                "instance_type": "t3.small",
-                "redis_port": 6390,
-            },
-            0,
-            manager._next_host_port,
+        runtime.provision_instance.assert_called_once()
+        provision_args = runtime.provision_instance.call_args.args
+        self.assertEqual(
+            provision_args[:2],
+            (
+                {
+                    "name": "Remote",
+                    "provider": "EC2",
+                    "instance_type": "t3.small",
+                    "redis_port": 6390,
+                },
+                0,
+            ),
         )
+        # EC2 jobs never get a pre-reserved port; the callback stands in for
+        # one but should always resolve to None for this provider.
+        self.assertIsNone(provision_args[2]("10.0.0.30"))
         runtime.bootstrap_instance.assert_called_once_with(
             provisioned,
             {
@@ -324,17 +331,23 @@ def runtime_for(provider):
             )
 
         local_runtime.validate_config.assert_called_once_with()
-        local_runtime.provision_instance.assert_called_once_with(
-            {"name": "Local", "provider": "local"}, 0, manager._next_host_port
+        local_runtime.provision_instance.assert_called_once()
+        local_provision_args = local_runtime.provision_instance.call_args.args
+        self.assertEqual(
+            local_provision_args[:2], ({"name": "Local", "provider": "local"}, 0)
         )
+        # Local jobs get a pre-reserved port (8000 is the first free port).
+        self.assertEqual(local_provision_args[2]("localhost"), 8000)
         local_runtime.bootstrap_instance.assert_called_once_with(
             {}, {"name": "Local", "provider": "local"}, 0
         )
-        ec2_runtime.provision_instance.assert_called_once_with(
-            {"name": "Remote", "provider": "EC2", "instance_type": "t3.small"},
-            0,
-            manager._next_host_port,
+        ec2_runtime.provision_instance.assert_called_once()
+        ec2_provision_args = ec2_runtime.provision_instance.call_args.args
+        self.assertEqual(
+            ec2_provision_args[:2],
+            ({"name": "Remote", "provider": "EC2", "instance_type": "t3.small"}, 0),
         )
+        self.assertIsNone(ec2_provision_args[2]("10.0.0.30"))
         ec2_runtime.bootstrap_instance.assert_called_once_with(
             {}, {"name": "Remote", "provider": "EC2", "instance_type": "t3.small"}, 0
         )

tests/test_runtime_sqlalchemy.py

@@ -34,8 +34,8 @@ def get(self, name):
     execution_time REAL,
     cpu_resource REAL,
     gpu_resource REAL,
-    created_at TEXT,
-    updated_at TEXT
+    created_at REAL,
+    updated_at REAL
 )
 """
 
@@ -90,7 +90,7 @@ def test_pull_and_upsert(self):
                 text("SELECT * FROM runtime_information WHERE future_id='abc'")
             ).fetchone()
         self.assertEqual(row[4], 8.0)
-        self.assertEqual(row[8], "9.0")
+        self.assertEqual(row[8], 9.0)
         for value in row:
             self.assertNotIn(value, (None, ""))
 

ventis/controller/cloud_provider_logic/EC2/_runtime.py

@@ -12,6 +12,7 @@
 they can read config and reuse the controller's Docker/Redis logic.
 """
 
+import logging
 import os
 import shlex
 import socket
@@ -24,6 +25,8 @@
 
 from ventis.utils.redis_client import RedisClient
 
+logger = logging.getLogger(__name__)
+
 CONTAINER_PORT = 50051
 DEFAULT_SSH_KEY_PATH = os.path.expanduser("~/.ssh/ventis_ec2")
 _controller: Any = None
@@ -225,13 +228,20 @@ def _bootstrap_instance(host, spec, replica_index, cfg, redis_host, redis_port):
     if spec.get("type") == "workflow":
         port_args += ["-p", f"{spec.get('api_port', 8080)}:8080"]
 
+    logger.info("Transferring image %s to %s", image, host)
     result = subprocess.run(
+        "set -o pipefail; "
         f"docker save {shlex.quote(image)} | ssh -o StrictHostKeyChecking=no "
         f"-o IdentitiesOnly=yes -i {shlex.quote(key)} "
         f"{shlex.quote(f'{ssh_user}@{host}')} 'sudo docker load'",
         shell=True,
         capture_output=True,
         text=True,
+        executable="/bin/bash",
+    )
+    logger.info(
+        "docker save|load returncode=%s stdout=%s stderr=%s",
+        result.returncode, result.stdout, result.stderr,
     )
     if result.returncode != 0:
         raise RuntimeError(f"Failed to transfer image to {host}: {result.stderr}")

ventis/controller/instance_manager.py

@@ -7,6 +7,8 @@
 """
 
 import json
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 from ventis.controller.cloud_provider_logic.Local import _runtime as local_runtime
 
@@ -29,6 +31,8 @@ def redis(self):
     def ensure_instances(self, agent_specs):
         self._agent_specs = list(agent_specs)
         instances = []
+        existing = []
+        jobs = []
 
         for agent_spec in self._agent_specs:
             agent_name = agent_spec["name"]
@@ -45,22 +49,65 @@ def ensure_instances(self, agent_specs):
                 key = self._instance_key(provider, agent_name, replica_index)
                 instance = self.redis.hgetall(key)
 
-                if not instance:
-                    provisioned = runtime.provision_instance(
-                        agent_spec, replica_index, self._next_host_port
+                if instance and instance.get("runtime_id"):
+                    existing.append((agent_name, instance_id, instance))
+                    continue
+
+                reserved_port = None
+                if provider == "local":
+                    host = agent_spec.get("host", local_runtime.DEFAULT_HOST)
+                    reserved_port = self._next_host_port(
+                        host, key, agent_name, provider, replica_index
                     )
-                    instance = runtime.bootstrap_instance(
-                        provisioned, agent_spec, replica_index
+
+                jobs.append(
+                    {
+                        "agent_name": agent_name,
+                        "agent_spec": agent_spec,
+                        "runtime": runtime,
+                        "replica_index": replica_index,
+                        "instance_id": instance_id,
+                        "reserved_port": reserved_port,
+                    }
+                )
+
+        max_workers = int(os.environ.get("VENTIS_MAX_AGENT_INSTANCES", 8))
+        provisioned = []
+        if jobs:
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                future_to_job = {
+                    executor.submit(self._provision_one, job): job for job in jobs
+                }
+                for future in as_completed(future_to_job):
+                    job = future_to_job[future]
+                    instance = future.result()
+                    provisioned.append(
+                        (job["agent_name"], job["instance_id"], instance)
                     )
-                    self._write_instance(instance)
 
-                self._add_instance_to_agent(agent_name, instance_id)
-                self._track_runtime(agent_name, instance["runtime_id"])
-                instances.append(instance)
+        for agent_name, instance_id, instance in existing + provisioned:
+            self._add_instance_to_agent(agent_name, instance_id)
+            self._track_runtime(agent_name, instance["runtime_id"])
+            instances.append(instance)
 
         self.publish_routing_snapshot(self._agent_specs)
         return instances
 
+    def _provision_one(self, job):
+        runtime = job["runtime"]
+        agent_spec = job["agent_spec"]
+        replica_index = job["replica_index"]
+        reserved_port = job["reserved_port"]
+
+        next_host_port = lambda host: reserved_port
+
+        provisioned = runtime.provision_instance(
+            agent_spec, replica_index, next_host_port
+        )
+        instance = runtime.bootstrap_instance(provisioned, agent_spec, replica_index)
+        self._write_instance(instance)
+        return instance
+
     def _write_instance(self, instance):
         key = self._instance_key(
             instance["provider"], instance["agent_name"], int(instance["replica_index"])
@@ -128,7 +175,7 @@ def _track_runtime(self, agent_name, runtime_id):
         if runtime_id not in containers:
             containers.append(runtime_id)
 
-    def _next_host_port(self, host):
+    def _next_host_port(self, host, key, agent_name, provider, replica_index):
         used = {
             int(instance["host_port"])
             for instance in self.list_instances()
@@ -137,6 +184,17 @@ def _next_host_port(self, host):
         port = DEFAULT_HOST_PORT_START
         while port in used:
             port += 1
+
+        self.redis.hset_multiple(
+            key,
+            {
+                "agent_name": agent_name,
+                "provider": provider,
+                "replica_index": str(replica_index),
+                "host": host,
+                "host_port": str(port),
+            },
+        )
         return port
 
     @staticmethod

ventis/controller/utils/sqlalchemy.py

@@ -8,9 +8,11 @@
 
 _engine = None
 
+TABLE_NAME = "runtime_information"
+
 _UPSERT = text(
-    """
-    INSERT INTO runtime_information (
+    f"""
+    INSERT INTO {TABLE_NAME} (
         future_id, session_id, workflow, agent, execution_time,
         cpu_resource, gpu_resource, created_at, updated_at
     ) VALUES (
@@ -70,6 +72,8 @@ def send_data(
             if not fid:
                 continue
             session_id = raw.get("request_id")
+            if not session_id:
+                continue
             workflow = (
                 redis_client.get(f"request:{session_id}:workflow")
                 if redis_client is not None
@@ -90,7 +94,7 @@ def send_data(
                     "execution_time": end - start,
                     "cpu_resource": cpu_resource,
                     "gpu_resource": gpu_resource,
-                    "created_at": str(start),
-                    "updated_at": str(end),
+                    "created_at": start,
+                    "updated_at": end,
                 },
             )