Orm

.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

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",
-]
+]

requirements.txt

@@ -6,3 +6,4 @@ pyyaml
 flask
 ipdb
 ipython
+sqlalchemy

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()

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()

ventis/controller/cloud_provider_logic/EC2/_runtime.py

@@ -18,21 +18,20 @@
 import stat
 import subprocess
 import time
+from typing import Any
 
 import boto3
 
 from ventis.utils.redis_client import RedisClient
 
 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}. "

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"
 

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                                                  #
     # ------------------------------------------------------------------ #

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")
 

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),
+                },
+            )

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(

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(

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.
 #