diff --git a/.gitignore b/.gitignore index 8afded9..db416f0 100644 --- a/.gitignore +++ b/.gitignore @@ -40,11 +40,3 @@ docker_container/ !.env.example AWSCLIV2.pkg .python-version - -sandbox/ - -tests/run_ec2_e2e.sh -tests/test_ec2_e2e.py - -tests/run_ec2_e2e.sh -tests/test_ec2_e2e.py diff --git a/pyproject.toml b/pyproject.toml index 8d046bf..5e26029 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,7 @@ dependencies = [ "grpcio", "grpcio-tools", "redis", + "sqlalchemy", "pyyaml", "flask", ] @@ -50,4 +51,4 @@ allowed-unresolved-imports = [ "deploy", "*_stub", "*_agent_stub", -] \ No newline at end of file +] diff --git a/requirements.txt b/requirements.txt index 203234d..688673a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,4 @@ pyyaml flask ipdb ipython +sqlalchemy diff --git a/tests/test_runtime_ec2.py b/tests/test_runtime_ec2.py index 471c6dc..a8e3d64 100644 --- a/tests/test_runtime_ec2.py +++ b/tests/test_runtime_ec2.py @@ -110,7 +110,9 @@ def test_aws_clients_fails_when_required_fields_are_missing(self): ec2_runtime._aws_clients() def test_aws_clients_rejects_missing_ssh_private_key(self): - self.controller.config["ec2"]["ssh_private_key_path"] = "/tmp/missing-ventis-key" + self.controller.config["ec2"]["ssh_private_key_path"] = ( + "/tmp/missing-ventis-key" + ) with self.assertRaisesRegex(ValueError, "does not exist"): ec2_runtime._aws_clients() diff --git a/tests/test_runtime_sqlalchemy.py b/tests/test_runtime_sqlalchemy.py new file mode 100644 index 0000000..c0a80f1 --- /dev/null +++ b/tests/test_runtime_sqlalchemy.py @@ -0,0 +1,99 @@ +import os +import sys +import tempfile +import unittest + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from sqlalchemy import create_engine, text + +import ventis.controller.utils.sqlalchemy as sqlmod + + +class _FakeRedis: + def __init__(self, hashes): + self.hashes = hashes + + def scan_keys(self, pattern): + prefix = pattern.rstrip("*") + return [k for k in self.hashes if k.startswith(prefix)] + + def hgetall(self, name): + return dict(self.hashes.get(name, {})) + + def get(self, name): + return self.hashes.get(name) + + +_CREATE = """ +CREATE TABLE runtime_information ( + future_id TEXT PRIMARY KEY, + session_id TEXT, + workflow TEXT, + agent TEXT, + execution_time REAL, + cpu_resource REAL, + gpu_resource REAL, + created_at TEXT, + updated_at TEXT +) +""" + + +class RuntimeSqlalchemyTests(unittest.TestCase): + def setUp(self): + self.db = tempfile.NamedTemporaryFile(suffix=".db", delete=False) + self.db.close() + os.environ["VENTIS_DATABASE_URL"] = f"sqlite:///{self.db.name}" + sqlmod._engine = None + with create_engine(os.environ["VENTIS_DATABASE_URL"]).begin() as conn: + conn.execute(text(_CREATE)) + + def tearDown(self): + sqlmod._engine = None + os.unlink(self.db.name) + + def test_pull_and_upsert(self): + redis = _FakeRedis( + { + "future:abc": { + "id": "abc", + "request_id": "req1", + "agent": "AgentA", + "created_at": "1.0", + }, + "request:req1:workflow": "main", + "future:abc:consumers": {"x": "1"}, + } + ) + rows = sqlmod.pull_data(redis) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0]["future_id"], "abc") + + sqlmod.send_data(rows, {"AgentA": {"cpu": 2, "gpu": 1}}, redis) + with sqlmod._get_engine("").connect() as conn: + row = conn.execute( + text( + "SELECT execution_time, cpu_resource, gpu_resource, workflow " + "FROM runtime_information WHERE future_id='abc'" + ) + ).fetchone() + self.assertGreaterEqual(row[0], 0) + self.assertEqual(row[1], 2.0) + self.assertEqual(row[2], 1.0) + self.assertEqual(row[3], "main") + + rows[0]["finished_at"] = "9.0" + sqlmod.send_data(rows, {"AgentA": {"cpu": 2, "gpu": 1}}, redis) + with sqlmod._get_engine("").connect() as conn: + row = conn.execute( + text("SELECT * FROM runtime_information WHERE future_id='abc'") + ).fetchone() + self.assertEqual(row[4], 8.0) + self.assertEqual(row[8], "9.0") + for value in row: + self.assertNotIn(value, (None, "")) + + +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 fdca290..11723ae 100644 --- a/ventis/controller/cloud_provider_logic/EC2/_runtime.py +++ b/ventis/controller/cloud_provider_logic/EC2/_runtime.py @@ -18,6 +18,7 @@ import stat import subprocess import time +from typing import Any import boto3 @@ -25,14 +26,12 @@ CONTAINER_PORT = 50051 DEFAULT_SSH_KEY_PATH = os.path.expanduser("~/.ssh/ventis_ec2") -_controller = None +_controller: Any = None def _ssh_key_path(cfg): """Return the configured EC2 SSH identity after validating it locally.""" - key_path = os.path.expanduser( - cfg.get("ssh_private_key_path", DEFAULT_SSH_KEY_PATH) - ) + key_path = os.path.expanduser(cfg.get("ssh_private_key_path", DEFAULT_SSH_KEY_PATH)) if not os.path.isfile(key_path): raise ValueError( f"EC2 SSH private key does not exist: {key_path}. " diff --git a/ventis/controller/global_controller.py b/ventis/controller/global_controller.py index 81365a6..4615b52 100644 --- a/ventis/controller/global_controller.py +++ b/ventis/controller/global_controller.py @@ -16,6 +16,7 @@ from ventis.controller.instance_manager import InstanceManager from ventis.controller.utils.agent_specs import write_agent_specs from ventis.controller.utils.redis_utils import _wait_for_redis +from ventis.controller.utils.sqlalchemy import pull_data, send_data from ventis.utils.redis_client import RedisClient # Add generated grpc_stubs from the local project to the path @@ -392,12 +393,21 @@ def run(self): self.stop() def _poll_controllers(self): - """Check the health of each registered controller replica via its node's Redis.""" + """ + Check the health of each registered controller replica via its node's Redis. + Also retrieves the request calls made in each instance. + """ for instance in self.instance_manager.list_instances(): name = instance["agent_name"] host = instance["host"] port = instance["host_port"] node_redis = self._get_node_redis_for(host) + send_data( + pull_data(node_redis), + {c["name"]: c.get("resources", {}) for c in self.controllers}, + node_redis, + self.config.get("database", {}).get("url"), + ) agent_host = self._agent_host_key(host) status_key = f"controller:{agent_host}:{port}:status" diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index fbfb88d..1dd6895 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -415,7 +415,6 @@ def _execute_locally( # Propagate the request_id context into this worker thread if request_id: ventis_context.set_request_id(request_id) - if self.agent is None: logger.error("No agent loaded, cannot execute %s.%s", service, function) return @@ -462,6 +461,8 @@ 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()) + # ------------------------------------------------------------------ # # Request forwarding # # ------------------------------------------------------------------ # diff --git a/ventis/controller/local_controller_frontend.py b/ventis/controller/local_controller_frontend.py index 41c8515..c1d8c30 100644 --- a/ventis/controller/local_controller_frontend.py +++ b/ventis/controller/local_controller_frontend.py @@ -40,6 +40,9 @@ def __init__(self, my_endpoint="unknown"): def Execute(self, request, context): """Accept an Execute request and push it into the queue.""" logger.info(f"Received request: {request.resonse}") + data = json.loads(request.resonse) + future_id = data.get("future_id") + self.redis.hset(f"future:{future_id}", "agent", data.get("service")) self.request_queue.put(request.resonse) return local_controler_pb2.JsonResponse(resonse="Request queued successfully") diff --git a/ventis/controller/utils/sqlalchemy.py b/ventis/controller/utils/sqlalchemy.py new file mode 100644 index 0000000..37ff883 --- /dev/null +++ b/ventis/controller/utils/sqlalchemy.py @@ -0,0 +1,94 @@ +"""Pull future hashes from Redis and upsert runtime_information rows.""" + +import os +import time + +from sqlalchemy import create_engine, text +from ventis.utils.redis_client import RedisClient + +_engine = None + +_UPSERT = text( + """ + INSERT INTO runtime_information ( + future_id, session_id, workflow, agent, execution_time, + cpu_resource, gpu_resource, created_at, updated_at + ) VALUES ( + :future_id, :session_id, :workflow, :agent, :execution_time, + :cpu_resource, :gpu_resource, :created_at, :updated_at + ) + ON CONFLICT(future_id) DO UPDATE SET + session_id=excluded.session_id, + workflow=excluded.workflow, + agent=excluded.agent, + execution_time=excluded.execution_time, + cpu_resource=excluded.cpu_resource, + gpu_resource=excluded.gpu_resource, + created_at=excluded.created_at, + updated_at=excluded.updated_at + """ +) + + +def _get_engine(database_url): + global _engine + if _engine is None: + _engine = create_engine( + os.environ.get("VENTIS_DATABASE_URL", str(database_url)) + ) + return _engine + + +def pull_data(redis_client): + """Scan node Redis for future data""" + rows = [] + for key in redis_client.scan_keys("future:*"): + if key.count(":") != 1: + continue + data = redis_client.hgetall(key) + if data: + data["future_id"] = data.get("id") or key.split(":", 1)[1] + rows.append(data) + return rows + + +def send_data( + rows, + resources_by_agent=None, + redis_client: RedisClient | None = None, + database_url="", +): + """UPSERT rows and attach allocated cpu/gpu from resources_by_agent.""" + if not rows: + return + resources_by_agent = resources_by_agent or {} + with _get_engine(database_url).begin() as conn: + for raw in rows: + agent = raw.get("agent") + res = resources_by_agent.get(agent, {}) + fid = raw.get("future_id") + if not fid: + continue + session_id = raw.get("request_id") + workflow = ( + redis_client.get(f"request:{session_id}:workflow") + if redis_client is not None + else None + ) + start = float(raw.get("created_at") or 0) + end = float(raw.get("finished_at") or time.time()) + + conn.execute( + _UPSERT, + { + "future_id": fid, + "session_id": session_id, + "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), + }, + ) diff --git a/ventis/deploy.py b/ventis/deploy.py index 51bbed7..436fef0 100644 --- a/ventis/deploy.py +++ b/ventis/deploy.py @@ -111,6 +111,7 @@ def handle_workflow(): request_id = uuid.uuid4().hex status_key = f"request:{request_id}:status" redis_client.set(status_key, "pending") + redis_client.set(f"request:{request_id}:workflow", fn_name) # Dispatch the workflow in a background thread thread = threading.Thread( diff --git a/ventis/future.py b/ventis/future.py index 5e2d7dc..58f5649 100644 --- a/ventis/future.py +++ b/ventis/future.py @@ -109,6 +109,7 @@ def _submit_request(self): } ) request = local_controler_pb2.JsonResponse(resonse=request_payload) + self.redis.hset(f"future:{self.id}", "created_at", time.time()) try: self.response = stub.Execute(request) logger.debug( diff --git a/ventis/templates/config/global_controller.yaml b/ventis/templates/config/global_controller.yaml index 2001cb4..4056d96 100644 --- a/ventis/templates/config/global_controller.yaml +++ b/ventis/templates/config/global_controller.yaml @@ -38,6 +38,9 @@ redis: port: 6379 db: 0 +database: + url: sqlite:///ventis_runtime.db + # EC2 defaults for `provider: EC2` replicas. # Keep them here so `config/global_controller.yaml` stays the only source of truth. #