From 26141ffbf3c4d235a851aea7f500acf7c57a3177 Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Tue, 21 Jul 2026 13:30:45 -0700 Subject: [PATCH] feature: added threading to ventis deploy | bug fix: fixed schema misalignment for runtime_information table --- pyproject.toml | 1 + requirements.txt | 1 + tests/test_instance_manager_runtime.py | 43 ++++++---- tests/test_runtime_sqlalchemy.py | 60 +++++++++++++- .../cloud_provider_logic/EC2/_runtime.py | 10 +++ ventis/controller/instance_manager.py | 78 ++++++++++++++++--- ventis/controller/local_controller.py | 12 ++- ventis/controller/utils/sqlalchemy.py | 20 +++-- 8 files changed, 189 insertions(+), 36 deletions(-) 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 c0a80f1..2651092 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,10 +90,64 @@ 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, "")) + def test_send_data_skips_rows_without_a_session_id(self): + redis = _FakeRedis( + { + "future:no_session": { + "id": "no_session", + "request_id": "", + "agent": "AgentA", + "created_at": "1.0", + }, + } + ) + rows = sqlmod.pull_data(redis) + self.assertEqual(len(rows), 1) + + sqlmod.send_data(rows, {"AgentA": {"cpu": 2, "gpu": 1}}, redis) + with sqlmod._get_engine("").connect() as conn: + count = conn.execute( + text( + "SELECT COUNT(*) FROM runtime_information " + "WHERE future_id='no_session'" + ) + ).scalar() + self.assertEqual(count, 0) + + def test_observed_cpu_overrides_config_and_gpu_uses_config(self): + redis = _FakeRedis( + { + "future:xyz": { + "id": "xyz", + "request_id": "req2", + "agent": "AgentB", + "created_at": "10.0", + "finished_at": "12.0", + "execution_time": "1.25", + "cpu_resource": "37.5", + "gpu_resource": "256", + }, + "request:req2:workflow": "wf2", + } + ) + + rows = sqlmod.pull_data(redis) + sqlmod.send_data(rows, {"AgentB": {"cpu": 8, "gpu": 0}}, redis) + with sqlmod._get_engine("").connect() as conn: + row = conn.execute( + text( + "SELECT execution_time, cpu_resource, gpu_resource " + "FROM runtime_information WHERE future_id='xyz'" + ) + ).fetchone() + self.assertEqual(row[0], 2.0) + self.assertEqual(row[1], 37.5) + self.assertEqual(row[2], 0.0) + if __name__ == "__main__": unittest.main() 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/local_controller.py b/ventis/controller/local_controller.py index 1dd6895..e67bb5c 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -412,6 +412,9 @@ def _execute_locally( self, service, function, args, future_id, origin=None, request_id=None ): """Execute a request on the local agent and write the result to Redis.""" + wall_start = time.time() + thread_cpu_start = time.thread_time() + # Propagate the request_id context into this worker thread if request_id: ventis_context.set_request_id(request_id) @@ -461,7 +464,14 @@ def _execute_locally( if origin and origin != self._my_endpoint: self._send_result_callback(origin, future_id, f"Execution failed: {e}") - self.redis.hset(f"future:{future_id}", "finished_at", time.time()) + wall_end = time.time() + self.redis.hset(f"future:{future_id}", "finished_at", wall_end) + + wall_duration = max(wall_end - wall_start, 0.0) + cpu_seconds = max(time.thread_time() - thread_cpu_start, 0.0) + cpu_percent = (cpu_seconds / wall_duration * 100.0) if wall_duration else 0.0 + + self.redis.hset(f"future:{future_id}", "cpu_resource", cpu_percent) # ------------------------------------------------------------------ # # Request forwarding # diff --git a/ventis/controller/utils/sqlalchemy.py b/ventis/controller/utils/sqlalchemy.py index 37ff883..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 ( @@ -58,7 +60,7 @@ def send_data( redis_client: RedisClient | None = None, database_url="", ): - """UPSERT rows and attach allocated cpu/gpu from resources_by_agent.""" + """UPSERT rows with observed CPU and configured GPU resource values.""" if not rows: return resources_by_agent = resources_by_agent or {} @@ -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 @@ -77,6 +81,8 @@ def send_data( ) start = float(raw.get("created_at") or 0) end = float(raw.get("finished_at") or time.time()) + cpu_resource = float(raw.get("cpu_resource") or res.get("cpu", 0)) + gpu_resource = float(res.get("gpu", 0)) conn.execute( _UPSERT, @@ -86,9 +92,9 @@ def send_data( "workflow": workflow, "agent": agent, "execution_time": end - start, - "cpu_resource": float(res.get("cpu", 0)), - "gpu_resource": float(res.get("gpu", 0)), - "created_at": str(start), - "updated_at": str(end), + "cpu_resource": cpu_resource, + "gpu_resource": gpu_resource, + "created_at": start, + "updated_at": end, }, )