diff --git a/pyproject.toml b/pyproject.toml index 5e26029..2bde8ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ dependencies = [ "grpcio-tools", "redis", "sqlalchemy", + "psycopg[binary]", "pyyaml", "flask", ] diff --git a/requirements.txt b/requirements.txt index 688673a..488f274 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,4 @@ flask ipdb ipython sqlalchemy +psycopg[binary] diff --git a/tests/test_instance_manager_runtime.py b/tests/test_instance_manager_runtime.py index b066acd..ba09958 100644 --- a/tests/test_instance_manager_runtime.py +++ b/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 ) diff --git a/tests/test_runtime_sqlalchemy.py b/tests/test_runtime_sqlalchemy.py index fc725ea..6d6ed1d 100644 --- a/tests/test_runtime_sqlalchemy.py +++ b/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, "")) diff --git a/ventis/controller/cloud_provider_logic/EC2/_runtime.py b/ventis/controller/cloud_provider_logic/EC2/_runtime.py index 11723ae..60cd224 100644 --- a/ventis/controller/cloud_provider_logic/EC2/_runtime.py +++ b/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}") diff --git a/ventis/controller/instance_manager.py b/ventis/controller/instance_manager.py index e28d7fc..cb4db16 100644 --- a/ventis/controller/instance_manager.py +++ b/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 diff --git a/ventis/controller/utils/sqlalchemy.py b/ventis/controller/utils/sqlalchemy.py index 922a178..313171f 100644 --- a/ventis/controller/utils/sqlalchemy.py +++ b/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, }, )