diff --git a/tests/test_runtime_sqlalchemy.py b/tests/test_runtime_sqlalchemy.py index c0a80f1..fc725ea 100644 --- a/tests/test_runtime_sqlalchemy.py +++ b/tests/test_runtime_sqlalchemy.py @@ -94,6 +94,36 @@ def test_pull_and_upsert(self): for value in row: self.assertNotIn(value, (None, "")) + 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/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..922a178 100644 --- a/ventis/controller/utils/sqlalchemy.py +++ b/ventis/controller/utils/sqlalchemy.py @@ -58,7 +58,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 {} @@ -77,6 +77,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,8 +88,8 @@ 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)), + "cpu_resource": cpu_resource, + "gpu_resource": gpu_resource, "created_at": str(start), "updated_at": str(end), },