From 63a00c28fb13c63c26a637095f8f61c6c4b22ada Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Mon, 3 Aug 2026 15:43:27 -0700 Subject: [PATCH 1/4] added a failure field to the future that fails if it happens --- ventis/controller/local_controller.py | 8 ++++++++ ventis/future.py | 6 ++++++ ventis/llm/bedrock.py | 7 ++++++- 3 files changed, 20 insertions(+), 1 deletion(-) diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index efcc4d0..05800da 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -501,6 +501,8 @@ def _execute_locally( "method": function, "args": json.dumps(args), "created_at": wall_start, + "failed": 0, + "error_message": "", }, ) if request_id: @@ -528,6 +530,12 @@ def _execute_locally( ) result = method(**args) + if str(self.redis.hget(f"future:{future_id}:metrics", "failed")) == "1": + raise RuntimeError( + self.redis.hget(f"future:{future_id}:metrics", "error_message") + or "Unknown error" + ) + # Serialize the result if isinstance(result, (dict, list)): serialized = json.dumps(result) diff --git a/ventis/future.py b/ventis/future.py index c63c466..351c8e9 100644 --- a/ventis/future.py +++ b/ventis/future.py @@ -137,6 +137,12 @@ def _poll_redis(self): error = self.redis.hget(self._key(), "error") if error: raise RuntimeError(error) + failed = self.redis.hget(f"future:{self.id}:metrics", "failed") + if str(failed) == "1": + raise RuntimeError( + self.redis.hget(f"future:{self.id}:metrics", "error_message") + or "Unknown error" + ) result = self.redis.hget(self._key(), "result") if result is not None and result != "": self.result = result diff --git a/ventis/llm/bedrock.py b/ventis/llm/bedrock.py index 1a5fefb..b843234 100644 --- a/ventis/llm/bedrock.py +++ b/ventis/llm/bedrock.py @@ -27,7 +27,7 @@ def call_bedrock(model_id: str, messages: list, inference_config: dict, region: modelId=model_id, messages=messages, inferenceConfig=inference_config ) return response - except Exception: + except Exception as e: # Counts failed attempts for this call. With no retry logic today this is # always 0 or 1, matching `failed`; once retries are added, a call that # eventually succeeds can still report error_count > 0 while failed stays 0. @@ -35,6 +35,11 @@ def call_bedrock(model_id: str, messages: list, inference_config: dict, region: metrics_key = ventis_context.get_current_metrics_key() if metrics_key: _redis.hincrby(metrics_key, "error_count", 1) + if future_id: + _redis.hset_multiple(f"future:{future_id}:metrics", { + "failed": 1, + "error_message": str(e), + }) raise finally: if future_id: From c7d48684785354b5041f4ea89fc8db740567907f Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Mon, 3 Aug 2026 16:21:04 -0700 Subject: [PATCH 2/4] fixed some issues --- tests/test_local_controller_metrics.py | 6 ++++++ ventis/controller/local_controller.py | 28 +++++++++++++++++++------- ventis/future.py | 13 ++++++------ 3 files changed, 34 insertions(+), 13 deletions(-) diff --git a/tests/test_local_controller_metrics.py b/tests/test_local_controller_metrics.py index aafa3e1..b995af0 100644 --- a/tests/test_local_controller_metrics.py +++ b/tests/test_local_controller_metrics.py @@ -175,6 +175,12 @@ def boom(name): ) self.assertIn("Execution failed", redis.hget("future:future-2", "result")) + self.assertEqual( + redis.hget("future:future-2:metrics", "failed"), 1 + ) + self.assertEqual( + redis.hget("future:future-2:metrics", "error_message"), "nope" + ) self.assertEqual( redis.hget("controller:localhost:50051:metrics", "requests_served"), 1 ) diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index 05800da..31931e1 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -293,6 +293,12 @@ def run(self): logger.error("Invalid JSON in request: %s", raw) except Exception as e: logger.error("Error processing request: %s", e) + future_id = data.get("future_id") + if future_id: + self.redis.hset_multiple( + f"future:{future_id}:metrics", + {"failed": 1, "error_message": str(e)}, + ) else: time.sleep(0.001) except KeyboardInterrupt: @@ -530,12 +536,6 @@ def _execute_locally( ) result = method(**args) - if str(self.redis.hget(f"future:{future_id}:metrics", "failed")) == "1": - raise RuntimeError( - self.redis.hget(f"future:{future_id}:metrics", "error_message") - or "Unknown error" - ) - # Serialize the result if isinstance(result, (dict, list)): serialized = json.dumps(result) @@ -561,7 +561,13 @@ def _execute_locally( logger.error("Failed to execute %s.%s: %s", service, function, e) # Treat script-level crash as a string result to avoid hanging - self.redis.hset(f"future:{future_id}:metrics", "failed", 1) + self.redis.hset_multiple( + f"future:{future_id}:metrics", + { + "failed": 1, + "error_message": str(e), + }, + ) self.redis.hset(f"future:{future_id}", "result", f"Execution failed: {e}") self.redis.hincrby(self._metrics_key, "full_failures", 1) if origin and origin != self._my_endpoint: @@ -618,6 +624,10 @@ def _forward_request(self, endpoint, data): future_id = data.get("future_id") if future_id: self.redis.hset(f"future:{future_id}", "error", str(e)) + self.redis.hset_multiple( + f"future:{future_id}:metrics", + {"failed": 1, "error_message": str(e)}, + ) def _send_result_callback(self, origin, future_id, result): """Send a result back to the originating controller via WriteResult RPC.""" @@ -645,6 +655,10 @@ def _send_result_callback(self, origin, future_id, result): except Exception as e: logger.error("Failed to send result callback to %s: %s", origin, e) + self.redis.hset_multiple( + f"future:{future_id}:metrics", + {"failed": 1, "error_message": str(e)}, + ) # ------------------------------------------------------------------ # # Shutdown # diff --git a/ventis/future.py b/ventis/future.py index 351c8e9..864ee4a 100644 --- a/ventis/future.py +++ b/ventis/future.py @@ -137,12 +137,6 @@ def _poll_redis(self): error = self.redis.hget(self._key(), "error") if error: raise RuntimeError(error) - failed = self.redis.hget(f"future:{self.id}:metrics", "failed") - if str(failed) == "1": - raise RuntimeError( - self.redis.hget(f"future:{self.id}:metrics", "error_message") - or "Unknown error" - ) result = self.redis.hget(self._key(), "result") if result is not None and result != "": self.result = result @@ -155,6 +149,13 @@ def value(self, timeout=None): Returns immediately if the result is already available locally. Polls Redis periodically to check for computed results. """ + failed = self.redis.hget(f"future:{self.id}:metrics", "failed") + if str(failed) == "1": + raise RuntimeError( + self.redis.hget(f"future:{self.id}:metrics", "error_message") + or "Unknown error" + ) + if self.result is not None: return self.result From 68ea2e668b46bf70ad3decfa9ee6d89f7cdd69af Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Mon, 3 Aug 2026 18:15:14 -0700 Subject: [PATCH 3/4] simplified logic and sent it to origin host --- tests/test_error_propagation.py | 117 ++++++++++++- tests/test_future.py | 10 ++ tests/test_local_controller_metrics.py | 94 ++++++++++- ventis/controller/local_controller.py | 154 ++++++++++++------ .../controller/local_controller_frontend.py | 12 +- ventis/future.py | 6 +- 6 files changed, 325 insertions(+), 68 deletions(-) diff --git a/tests/test_error_propagation.py b/tests/test_error_propagation.py index ddd0a0b..a998cac 100644 --- a/tests/test_error_propagation.py +++ b/tests/test_error_propagation.py @@ -1,6 +1,7 @@ import os import sys import unittest +import json from types import SimpleNamespace from unittest.mock import MagicMock @@ -15,7 +16,9 @@ ) from ventis.controller.local_controller import LocalController +from ventis.controller.local_controller_frontend import LocalControllerServicer from ventis.future import Future +import local_controler_pb2 class _FakeRedis: @@ -25,19 +28,29 @@ def __init__(self): def hset(self, name, field, value): self.hashes.setdefault(name, {})[field] = value + def hset_multiple(self, name, mapping): + self.hashes.setdefault(name, {}).update(mapping) + def hget(self, name, field): return self.hashes.get(name, {}).get(field) +def _bind_failure_marker(controller): + controller._mark_future_failed = lambda future_id, error, origin=None: ( + LocalController._mark_future_failed(controller, future_id, error, origin) + ) + return controller + + class ErrorPropagationTests(unittest.TestCase): def test_forward_request_writes_future_error_on_grpc_failure(self): redis = _FakeRedis() stub = SimpleNamespace(Execute=MagicMock(side_effect=RuntimeError("boom"))) - controller = SimpleNamespace( + controller = _bind_failure_marker(SimpleNamespace( redis=redis, _my_endpoint="172.31.19.107:50051", _get_remote_stub=lambda endpoint: stub, - ) + )) data = { "future_id": "future-1", "service": "ExampleAgent", @@ -50,19 +63,28 @@ def test_forward_request_writes_future_error_on_grpc_failure(self): self.assertEqual(redis.hget("future:future-1", "error"), "boom") stub.Execute.assert_called_once() - def test_future_poll_redis_raises_runtime_error_when_error_is_present(self): + def test_future_value_raises_runtime_error_when_error_is_present(self): redis = _FakeRedis() redis.hset("future:future-1", "error", "boom") - future = SimpleNamespace(redis=redis, _key=lambda: "future:future-1") + future = SimpleNamespace( + redis=redis, + _key=lambda: "future:future-1", + _poll_redis=lambda: Future._poll_redis(future), + result=None, + ) with self.assertRaisesRegex(RuntimeError, "boom"): - Future._poll_redis(future) + Future.value(future) def test_future_poll_redis_returns_result_when_error_is_absent(self): redis = _FakeRedis() redis.hset("future:future-1", "result", "Hello, World!") future = SimpleNamespace( - redis=redis, _key=lambda: "future:future-1", result=None + redis=redis, + _key=lambda: "future:future-1", + _poll_redis=lambda: Future._poll_redis(future), + id="future-1", + result=None, ) result = Future._poll_redis(future) @@ -70,6 +92,89 @@ def test_future_poll_redis_returns_result_when_error_is_absent(self): self.assertEqual(result, "Hello, World!") self.assertEqual(future.result, "Hello, World!") + def test_future_value_raises_when_metrics_mark_it_failed(self): + redis = _FakeRedis() + redis.hset_multiple( + "future:future-1:metrics", + {"failed": 1, "error_message": "agent exploded"}, + ) + future = SimpleNamespace( + redis=redis, + _key=lambda: "future:future-1", + _poll_redis=lambda: Future._poll_redis(future), + id="future-1", + result=None, + ) + + with self.assertRaisesRegex(RuntimeError, "agent exploded"): + Future.value(future) + + def test_result_callback_sends_error_separately_from_result(self): + redis = _FakeRedis() + stub = SimpleNamespace(WriteResult=MagicMock()) + controller = SimpleNamespace( + redis=redis, + agent_name="ExampleAgent", + _get_remote_stub=lambda endpoint: stub, + ) + + LocalController._send_result_callback( + controller, + "origin:50051", + "future-1", + failed=1, + error_message="agent exploded", + ) + + payload = stub.WriteResult.call_args.args[0].resonse + self.assertEqual( + json.loads(payload), + { + "future_id": "future-1", + "result": None, + "failed": 1, + "error_message": "agent exploded", + }, + ) + + def test_write_result_persists_remote_error_as_terminal_failure(self): + redis = _FakeRedis() + servicer = SimpleNamespace(redis=redis) + request = local_controler_pb2.JsonResponse( + resonse=json.dumps( + { + "future_id": "future-1", + "failed": 1, + "error_message": "remote exploded", + } + ) + ) + context = SimpleNamespace(peer=lambda: "peer:50051") + + LocalControllerServicer.WriteResult(servicer, request, context) + + self.assertEqual( + redis.hget("future:future-1:metrics", "failed"), 1 + ) + self.assertEqual( + redis.hget("future:future-1:metrics", "error_message"), + "remote exploded", + ) + + def test_malformed_request_with_future_id_is_marked_failed(self): + redis = _FakeRedis() + controller = _bind_failure_marker( + SimpleNamespace(redis=redis, _my_endpoint="localhost:50051") + ) + + LocalController._process_request(controller, {"future_id": "future-1"}) + + self.assertEqual( + redis.hget("future:future-1", "error"), + "Malformed request: missing service, function, or future_id", + ) + self.assertEqual(redis.hget("future:future-1:metrics", "failed"), 1) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_future.py b/tests/test_future.py index 9e79fec..c277358 100644 --- a/tests/test_future.py +++ b/tests/test_future.py @@ -68,6 +68,16 @@ def test_parent_is_the_currently_executing_future_id(self): self.fake_redis.hashes[f"future:{f.id}"]["parent"], "caller-future-id" ) + def test_submission_failure_is_raised_by_value_not_constructor(self): + future_module.Future._stub.Execute.side_effect = RuntimeError("submit failed") + + future = future_module.Future( + parent="ignored/file.py", service="Svc", method="do_thing" + ) + + with self.assertRaisesRegex(RuntimeError, "submit failed"): + future.value() + if __name__ == "__main__": unittest.main() diff --git a/tests/test_local_controller_metrics.py b/tests/test_local_controller_metrics.py index b995af0..90491a7 100644 --- a/tests/test_local_controller_metrics.py +++ b/tests/test_local_controller_metrics.py @@ -2,6 +2,8 @@ import sys import threading import unittest +import json +from unittest.mock import MagicMock from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace from unittest.mock import patch @@ -19,6 +21,18 @@ from ventis.controller.local_controller import LocalController +def _bind_failure_marker(controller): + controller._mark_future_failed = lambda future_id, error, origin=None: ( + LocalController._mark_future_failed(controller, future_id, error, origin) + ) + controller._send_result_callback = ( + lambda origin, future_id, result=None, failed=0, error_message="": LocalController._send_result_callback( + controller, origin, future_id, result, failed, error_message + ) + ) + return controller + + class _FakeRedisClient: def __init__(self): self.counters = {} @@ -123,7 +137,7 @@ def stop_after_one_tick(timeout): def test_execute_locally_writes_gpu_resource_to_future_hash(self): redis = _FakeRedis() agent = SimpleNamespace(greet=lambda name: f"hello {name}") - controller = SimpleNamespace( + controller = _bind_failure_marker(SimpleNamespace( redis=redis, agent=agent, agent_name="Greeter", @@ -131,7 +145,7 @@ def test_execute_locally_writes_gpu_resource_to_future_hash(self): _my_endpoint="localhost:50051", _metrics_key="controller:localhost:50051:metrics", _resolve_future_args=lambda args: args, - ) + )) with patch( "ventis.controller.local_controller.read_gpu_percent", return_value=17.5 @@ -157,7 +171,7 @@ def boom(name): raise ValueError("nope") agent = SimpleNamespace(greet=boom) - controller = SimpleNamespace( + controller = _bind_failure_marker(SimpleNamespace( redis=redis, agent=agent, agent_name="Greeter", @@ -165,7 +179,7 @@ def boom(name): _my_endpoint="localhost:50051", _metrics_key="controller:localhost:50051:metrics", _resolve_future_args=lambda args: args, - ) + )) with patch( "ventis.controller.local_controller.read_gpu_percent", return_value=0.0 @@ -174,7 +188,8 @@ def boom(name): controller, "Greeter", "greet", {"name": "world"}, "future-2" ) - self.assertIn("Execution failed", redis.hget("future:future-2", "result")) + self.assertEqual(redis.hget("future:future-2", "error"), "nope") + self.assertIsNone(redis.hget("future:future-2", "result")) self.assertEqual( redis.hget("future:future-2:metrics", "failed"), 1 ) @@ -188,6 +203,75 @@ def boom(name): redis.hget("controller:localhost:50051:metrics", "full_failures"), 1 ) + def test_execute_locally_marks_missing_agent_as_failed(self): + redis = _FakeRedis() + controller = _bind_failure_marker(SimpleNamespace( + redis=redis, + agent=None, + agent_name="MissingAgent", + agent_id="agent-1", + _my_endpoint="localhost:50051", + _metrics_key="controller:localhost:50051:metrics", + )) + + with patch( + "ventis.controller.local_controller.read_gpu_percent", return_value=0.0 + ): + LocalController._execute_locally( + controller, "MissingAgent", "greet", {}, "future-3" + ) + + self.assertEqual( + redis.hget("future:future-3", "error"), "No agent loaded" + ) + self.assertEqual(redis.hget("future:future-3:metrics", "failed"), 1) + self.assertEqual( + redis.hget("future:future-3:metrics", "error_message"), + "No agent loaded", + ) + + def test_remote_execution_failure_sends_error_callback(self): + redis = _FakeRedis() + stub = SimpleNamespace(WriteResult=MagicMock()) + + def boom(): + raise ValueError("remote nope") + + controller = _bind_failure_marker(SimpleNamespace( + redis=redis, + agent=SimpleNamespace(greet=boom), + agent_name="Greeter", + agent_id="agent-1", + _my_endpoint="target:50051", + _metrics_key="controller:target:50051:metrics", + _resolve_future_args=lambda args: args, + _get_remote_stub=lambda endpoint: stub, + )) + + with patch( + "ventis.controller.local_controller.read_gpu_percent", return_value=0.0 + ): + LocalController._execute_locally( + controller, + "Greeter", + "greet", + {}, + "future-4", + origin="origin:50051", + ) + + self.assertEqual(redis.hget("future:future-4", "error"), "remote nope") + payload = stub.WriteResult.call_args.args[0].resonse + self.assertEqual( + json.loads(payload), + { + "future_id": "future-4", + "result": None, + "failed": 1, + "error_message": "remote nope", + }, + ) + if __name__ == "__main__": unittest.main() diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index 31931e1..3599704 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -286,6 +286,7 @@ def run(self): while True: if not self.request_queue.empty(): raw = self.request_queue.get() + data = None try: data = json.loads(raw) self._process_request(data) @@ -293,17 +294,35 @@ def run(self): logger.error("Invalid JSON in request: %s", raw) except Exception as e: logger.error("Error processing request: %s", e) - future_id = data.get("future_id") - if future_id: - self.redis.hset_multiple( - f"future:{future_id}:metrics", - {"failed": 1, "error_message": str(e)}, + if isinstance(data, dict): + self._mark_future_failed( + data.get("future_id"), e, data.get("origin") ) else: time.sleep(0.001) except KeyboardInterrupt: self.stop() + def _mark_future_failed(self, future_id, error, origin=None): + """Persist a terminal failure locally and, when needed, notify the origin.""" + if not future_id: + return + + error_message = str(error) or "Unknown error" + self.redis.hset(f"future:{future_id}", "error", error_message) + self.redis.hset_multiple( + f"future:{future_id}:metrics", + {"failed": 1, "error_message": error_message}, + ) + + if origin and origin != self._my_endpoint: + self._send_result_callback( + origin, + future_id, + failed=1, + error_message=error_message, + ) + def _process_request(self, data): """ Route a request to the correct controller. @@ -340,15 +359,18 @@ def _process_request(self, data): if not service or not function or not future_id: logger.error("Malformed request, missing required fields: %s", data) + self._mark_future_failed( + future_id, + "Malformed request: missing service, function, or future_id", + origin, + ) return # Check policy before routing if not self._check_policy(service, context): err_msg = f"Unauthorized: Policy denied access to service '{service}'" logger.warning(err_msg) - self.redis.hset(f"future:{future_id}", "result", err_msg) - if origin and origin != self._my_endpoint: - self._send_result_callback(origin, future_id, err_msg) + self._mark_future_failed(future_id, err_msg, origin) return # Resolve which endpoint to route to. @@ -366,6 +388,11 @@ def _process_request(self, data): logger.error( "No endpoint found for service '%s' in routing table.", service ) + self._mark_future_failed( + future_id, + f"No endpoint found for service '{service}'", + origin, + ) return if endpoint == self._my_endpoint: @@ -460,6 +487,17 @@ def _resolve_future_args(self, args, poll_interval=0.01, timeout=300): ) start = time.time() while True: + error = self.redis.hget(future_key, "error") + if error: + raise RuntimeError(error) + failed = self.redis.hget(f"future:{value}:metrics", "failed") + if str(failed) == "1": + raise RuntimeError( + self.redis.hget( + f"future:{value}:metrics", "error_message" + ) + or "Unknown error" + ) # print("Waiting for result for future next iteration %s", value) result = self.redis.hget(future_key, "result") if result is not None and result != "": @@ -518,11 +556,15 @@ def _execute_locally( ventis_context.set_current_metrics_key(self._metrics_key) if self.agent is None: logger.error("No agent loaded, cannot execute %s.%s", service, function) + self._mark_future_failed(future_id, "No agent loaded", origin) return method = getattr(self.agent, function, None) if method is None: logger.error("Agent %s has no method '%s'", self.agent_name, function) + self._mark_future_failed( + future_id, f"Agent {self.agent_name} has no method '{function}'", origin + ) return self.redis.hincrby(self._metrics_key, "requests_served", 1) @@ -547,7 +589,13 @@ def _execute_locally( # If the request came from another node, send result back to origin if origin and origin != self._my_endpoint: - self._send_result_callback(origin, future_id, serialized) + self._send_result_callback( + origin, + future_id, + result=serialized, + failed=0, + error_message="", + ) logger.info( "Completed %s.%s (future=%s) -> %s", @@ -560,37 +608,41 @@ def _execute_locally( except Exception as e: logger.error("Failed to execute %s.%s: %s", service, function, e) - # Treat script-level crash as a string result to avoid hanging - self.redis.hset_multiple( - f"future:{future_id}:metrics", - { - "failed": 1, - "error_message": str(e), - }, - ) - self.redis.hset(f"future:{future_id}", "result", f"Execution failed: {e}") + self._mark_future_failed(future_id, e, origin) self.redis.hincrby(self._metrics_key, "full_failures", 1) - if origin and origin != self._my_endpoint: - self._send_result_callback(origin, future_id, f"Execution failed: {e}") - - wall_end = time.time() - self.redis.hset(f"future:{future_id}:metrics", "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}:metrics", "cpu_resource", cpu_percent) - self.redis.hset(f"future:{future_id}:metrics", "gpu_resource", read_gpu_percent()) - self.redis.hset(f"future:{future_id}:metrics", "agent", self.agent_id) - ventis_context.set_current_future_id(parent or "") - - if submitted_at is not None: - self.redis.hset( - f"future:{future_id}:metrics", - "queue_time", - max(wall_start - submitted_at, 0.0), + finally: + wall_end = time.time() + 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 ) + try: + gpu_percent = read_gpu_percent() + except Exception as metric_error: + logger.error("Failed to read GPU metrics: %s", metric_error) + gpu_percent = "" + + try: + self.redis.hset_multiple( + f"future:{future_id}:metrics", + { + "finished_at": wall_end, + "cpu_resource": cpu_percent, + "gpu_resource": gpu_percent, + "agent": self.agent_id, + **( + { + "queue_time": max(wall_start - submitted_at, 0.0) + } + if submitted_at is not None + else {} + ), + }, + ) + except Exception as metric_error: + logger.error("Failed to finalize metrics for future %s: %s", future_id, metric_error) + ventis_context.set_current_future_id(parent or "") # ------------------------------------------------------------------ # # Request forwarding # @@ -621,16 +673,12 @@ def _forward_request(self, endpoint, data): logger.debug("Forwarded request to %s", endpoint) except Exception as e: logger.error("Failed to forward request to %s: %s", endpoint, e) - future_id = data.get("future_id") - if future_id: - self.redis.hset(f"future:{future_id}", "error", str(e)) - self.redis.hset_multiple( - f"future:{future_id}:metrics", - {"failed": 1, "error_message": str(e)}, - ) + self._mark_future_failed(data.get("future_id"), e) - def _send_result_callback(self, origin, future_id, result): - """Send a result back to the originating controller via WriteResult RPC.""" + def _send_result_callback( + self, origin, future_id, result=None, failed=0, error_message="" + ): + """Send a result and its failure metadata to the originating controller.""" if not result: logger.warning( "Agent '%s' is sending an empty/None result for future %s to origin %s, result: %s", @@ -641,7 +689,12 @@ def _send_result_callback(self, origin, future_id, result): ) stub = self._get_remote_stub(origin) - payload = json.dumps({"future_id": future_id, "result": result}) + payload = json.dumps({ + "future_id": future_id, + "result": result, + "failed": int(bool(failed)), + "error_message": str(error_message or ""), + }) logger.info("Payload: Future %s,Sent %s ", future_id, payload) request = local_controler_pb2.JsonResponse(resonse=payload) try: @@ -655,10 +708,7 @@ def _send_result_callback(self, origin, future_id, result): except Exception as e: logger.error("Failed to send result callback to %s: %s", origin, e) - self.redis.hset_multiple( - f"future:{future_id}:metrics", - {"failed": 1, "error_message": str(e)}, - ) + self._mark_future_failed(future_id, f"Result callback failed: {e}") # ------------------------------------------------------------------ # # Shutdown # diff --git a/ventis/controller/local_controller_frontend.py b/ventis/controller/local_controller_frontend.py index c7ac28d..a6e478d 100644 --- a/ventis/controller/local_controller_frontend.py +++ b/ventis/controller/local_controller_frontend.py @@ -50,7 +50,8 @@ def WriteResult(self, request, context): data = json.loads(request.resonse) future_id = data.get("future_id") result = data.get("result") - error = data.get("error") + failed = int(bool(data.get("failed", 0))) + error_message = str(data.get("error_message") or "") logger.info( f"WriteResult: received result for future {future_id}: {result}" @@ -61,10 +62,13 @@ def WriteResult(self, request, context): ) if future_id: - if error is not None: - self.redis.hset(f"future:{future_id}", "error", error) + self.redis.hset_multiple( + f"future:{future_id}:metrics", + {"failed": failed, "error_message": error_message}, + ) + if failed: logger.info("WriteResult: wrote error for future %s", future_id) - if result is not None: + elif result is not None: self.redis.hset(f"future:{future_id}", "result", result) logger.info( "WriteResult: wrote result for future %s, result %s", diff --git a/ventis/future.py b/ventis/future.py index 864ee4a..598bc4f 100644 --- a/ventis/future.py +++ b/ventis/future.py @@ -118,7 +118,11 @@ def _submit_request(self): ) except Exception as e: logger.error("gRPC call failed for %s.%s: %s", self.service, self.method, e) - raise + self.redis.hset(f"future:{self.id}", "error", str(e)) + self.redis.hset_multiple(f"future:{self.id}:metrics", { + "failed": 1, + "error_message": str(e), + }) def _key(self): """Redis key for this future's hash.""" From d62e17bef6928617b9f428ccbcbeb3aaa06ba4a1 Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Tue, 4 Aug 2026 11:32:11 -0700 Subject: [PATCH 4/4] fixes --- examples/portfolio/agents/intent_agent.py | 2 +- examples/portfolio/agents/price_agent.py | 14 +- .../portfolio/config/global_controller.yaml | 15 -- .../text2sql/config/global_controller.yaml | 15 +- tests/test_demo_obfuscation.py | 194 ------------------ tests/test_runtime_sqlalchemy.py | 25 +-- tests/test_session_store.py | 16 +- ventis/controller/local_controller.py | 46 ++--- .../controller/local_controller_frontend.py | 3 + ventis/controller/utils/demo_obfuscation.py | 41 ---- ventis/controller/utils/session_store.py | 15 +- ventis/controller/utils/sqlalchemy.py | 17 +- ventis/stub_generator.py | 5 - 13 files changed, 51 insertions(+), 357 deletions(-) delete mode 100644 tests/test_demo_obfuscation.py delete mode 100644 ventis/controller/utils/demo_obfuscation.py diff --git a/examples/portfolio/agents/intent_agent.py b/examples/portfolio/agents/intent_agent.py index 886ee14..1eb15d7 100644 --- a/examples/portfolio/agents/intent_agent.py +++ b/examples/portfolio/agents/intent_agent.py @@ -36,7 +36,7 @@ class IntentAgent(object): def __init__(self): self.tools = [self.parse] self.model_id = os.environ.get( - "BEDROCK_MODEL_ID", "openai.gpt-5.4" + "BEDROCK_MODEL_ID", "meta.llama3-8b-instruct-v1:0" ) self.region = os.environ.get("AWS_REGION", "us-east-1") diff --git a/examples/portfolio/agents/price_agent.py b/examples/portfolio/agents/price_agent.py index 5c53ead..353efc7 100644 --- a/examples/portfolio/agents/price_agent.py +++ b/examples/portfolio/agents/price_agent.py @@ -18,11 +18,21 @@ def __init__(self): def get_history(self, ticker: str, lookback_days: int = 365) -> dict: """Fetch daily closing prices for a ticker over the lookback window.""" try: + import math + import yfinance as yf hist = yf.Ticker(ticker).history(period=f"{lookback_days}d") - closes = [float(c) for c in hist["Close"].tolist()] - dates = [str(d.date()) for d in hist.index] + raw_closes = hist["Close"].tolist() + # Yahoo sometimes appends the newest trading day before its close has + # settled, leaving that row's Close as NaN -- drop it rather than let + # NaN poison every downstream metric. + closes = [float(c) for c in raw_closes if not math.isnan(c)] + dates = [ + str(d.date()) + for d, c in zip(hist.index, raw_closes) + if not math.isnan(c) + ] if closes: return { "ticker": ticker, diff --git a/examples/portfolio/config/global_controller.yaml b/examples/portfolio/config/global_controller.yaml index 9695a6d..fbd937e 100644 --- a/examples/portfolio/config/global_controller.yaml +++ b/examples/portfolio/config/global_controller.yaml @@ -86,18 +86,3 @@ redis: port: 6379 db: 0 -database: - url: postgresql://postgres:RaDitLIHsdPAhrSpjOYzDAxFFMPtDvSW@tokaido.proxy.rlwy.net:10279/railway - -# EC2 infra config, required by ventis's EC2 deploy preflight check -ec2: - region: us-east-1 - ami_id: ami-01ac7077de2a109d3 - subnet_id: subnet-099cc5963830b9bb8 - security_group_ids: - - sg-025daf3a98e06cef3 - ssh_user: ubuntu - ssh_private_key_path: ~/.ssh/ventis_ec2 - - -project_id: 00000000-0000-0000-0000-000000000000 diff --git a/examples/text2sql/config/global_controller.yaml b/examples/text2sql/config/global_controller.yaml index b305a37..89642b2 100644 --- a/examples/text2sql/config/global_controller.yaml +++ b/examples/text2sql/config/global_controller.yaml @@ -101,17 +101,4 @@ poll_interval: 5 redis: host: localhost port: 6379 - db: 0 - -database: - url: postgresql://postgres:RaDitLIHsdPAhrSpjOYzDAxFFMPtDvSW@tokaido.proxy.rlwy.net:10279/railway - - -ec2: - region: us-east-1 - ami_id: ami-01ac7077de2a109d3 - subnet_id: subnet-099cc5963830b9bb8 - security_group_ids: - - sg-025daf3a98e06cef3 - ssh_user: ubuntu - ssh_private_key_path: ~/.ssh/ventis_ec2 \ No newline at end of file + db: 0 \ No newline at end of file diff --git a/tests/test_demo_obfuscation.py b/tests/test_demo_obfuscation.py deleted file mode 100644 index 823eaf4..0000000 --- a/tests/test_demo_obfuscation.py +++ /dev/null @@ -1,194 +0,0 @@ -import os -import sys -import tempfile -import unittest -from datetime import datetime, timezone - -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) - -from sqlalchemy import text - -import ventis.controller.utils.demo_obfuscation as demo_obfuscation -import ventis.controller.utils.session_store as session_store -import ventis.controller.utils.sqlalchemy as sqlmod - - -def _stored_epoch(stored): - """Unix epoch seconds held by a stored TIMESTAMPTZ.""" - return datetime.fromisoformat(str(stored)).replace(tzinfo=timezone.utc).timestamp() - - -class ShiftForSessionTests(unittest.TestCase): - def test_is_deterministic_for_the_same_session(self): - self.assertEqual( - demo_obfuscation.shift_for_session("req1"), - demo_obfuscation.shift_for_session("req1"), - ) - - def test_differs_between_sessions(self): - self.assertNotEqual( - demo_obfuscation.shift_for_session("req1"), - demo_obfuscation.shift_for_session("req2"), - ) - - def test_stays_inside_the_configured_window(self): - for session_id in ("req1", "req2", "", 0, "a" * 200): - shift = demo_obfuscation.shift_for_session(session_id) - self.assertGreaterEqual(shift, 0) - self.assertLess(shift, demo_obfuscation.RANDOM_SHIFT_MAX_SECONDS) - - def test_accepts_non_string_session_ids(self): - # request ids arrive as str everywhere today, but project_id defaults to - # the int 0 in global_controller, so don't blow up on non-str input. - self.assertEqual( - demo_obfuscation.shift_for_session(1234), - demo_obfuscation.shift_for_session("1234"), - ) - - -class ToTimestamptzTests(unittest.TestCase): - def test_subtracts_the_shift_and_returns_utc(self): - result = demo_obfuscation.to_timestamptz(1000.0, 400) - self.assertEqual(result.timestamp(), 600.0) - self.assertEqual(result.tzinfo, timezone.utc) - - def test_accepts_epoch_seconds_as_text(self): - # Redis hands metrics back as strings. - self.assertEqual( - demo_obfuscation.to_timestamptz("1000.0", 400), - demo_obfuscation.to_timestamptz(1000.0, 400), - ) - - def test_preserves_intervals_within_a_session(self): - shift = demo_obfuscation.shift_for_session("req1") - start = demo_obfuscation.to_timestamptz(1_700_000_000.0, shift) - end = demo_obfuscation.to_timestamptz(1_700_000_012.5, shift) - self.assertEqual((end - start).total_seconds(), 12.5) - - -class SharedShiftAcrossWritersTests(unittest.TestCase): - """The reason the shift lives in its own module: the `session` row and the - runtime_information rows for that session's futures are written by two - different processes, and they have to land on one timeline.""" - - PROJECT_ID = "11111111-1111-1111-1111-111111111111" - - def setUp(self): - self.db = tempfile.NamedTemporaryFile(suffix=".db", delete=False) - self.db.close() - os.environ["VENTIS_DATABASE_URL"] = f"sqlite:///{self.db.name}" - # Both writers keep their own module-global engine; point both at the - # same file so one assertion can compare what they wrote. - session_store._engine = None - sqlmod._engine = None - sqlmod._project_id = self.PROJECT_ID - with session_store._get_engine("").begin() as conn: - conn.execute( - text( - """ - CREATE TABLE session ( - session_id VARCHAR(255) PRIMARY KEY, - project_id UUID NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'running', - input JSONB, - output JSONB, - created_at TIMESTAMPTZ NOT NULL, - updated_at TIMESTAMPTZ NOT NULL - ) - """ - ) - ) - - def tearDown(self): - session_store._engine = None - sqlmod._engine = None - sqlmod._project_id = None - os.environ.pop("VENTIS_DATABASE_URL", None) - os.unlink(self.db.name) - - def test_session_and_runtime_rows_share_one_timeline(self): - dispatched_at = 1_700_000_000.0 - future_started_at = dispatched_at + 1.0 - future_finished_at = dispatched_at + 4.0 - completed_at = dispatched_at + 5.0 - - session_store.upsert_session( - "", self.PROJECT_ID, "req1", "running", dispatched_at - ) - sqlmod.send_runtime_information( - [ - { - "future_id": "fut1", - "request_id": "req1", - "created_at": str(future_started_at), - "finished_at": str(future_finished_at), - } - ] - ) - session_store.upsert_session( - "", self.PROJECT_ID, "req1", "completed", completed_at - ) - - with session_store._get_engine("").connect() as conn: - session_row = ( - conn.execute( - text("SELECT * FROM session WHERE session_id='req1'") - ) - .mappings() - .fetchone() - ) - runtime_row = ( - conn.execute( - text("SELECT * FROM runtime_information WHERE future_id='fut1'") - ) - .mappings() - .fetchone() - ) - - shift = demo_obfuscation.shift_for_session("req1") - self.assertEqual( - _stored_epoch(session_row["created_at"]), dispatched_at - shift - ) - self.assertEqual( - _stored_epoch(runtime_row["started_at"]), future_started_at - shift - ) - self.assertEqual( - _stored_epoch(runtime_row["finished_at"]), future_finished_at - shift - ) - self.assertEqual( - _stored_epoch(session_row["updated_at"]), completed_at - shift - ) - - # The whole point: the future ran inside its own session's window, and the - # real durations survived the shift. - self.assertLessEqual( - _stored_epoch(session_row["created_at"]), - _stored_epoch(runtime_row["started_at"]), - ) - self.assertGreaterEqual( - _stored_epoch(session_row["updated_at"]), - _stored_epoch(runtime_row["finished_at"]), - ) - self.assertEqual( - _stored_epoch(session_row["updated_at"]) - - _stored_epoch(session_row["created_at"]), - completed_at - dispatched_at, - ) - - def test_shift_is_applied_at_all(self): - # Guards against the shift silently regressing to a no-op, which would - # publish real wall-clock times. - session_store.upsert_session( - "", self.PROJECT_ID, "req1", "running", 1_700_000_000.0 - ) - with session_store._get_engine("").connect() as conn: - row = ( - conn.execute(text("SELECT * FROM session WHERE session_id='req1'")) - .mappings() - .fetchone() - ) - self.assertNotEqual(_stored_epoch(row["created_at"]), 1_700_000_000.0) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_runtime_sqlalchemy.py b/tests/test_runtime_sqlalchemy.py index c2bca59..9c9936c 100644 --- a/tests/test_runtime_sqlalchemy.py +++ b/tests/test_runtime_sqlalchemy.py @@ -9,13 +9,11 @@ from sqlalchemy import create_engine, text -import ventis.controller.utils.demo_obfuscation as demo_obfuscation import ventis.controller.utils.sqlalchemy as sqlmod def _parse_shifted(stored): - """Turn a stored TIMESTAMPTZ back into the Unix epoch seconds it holds, so - tests can assert against the shift the writer was supposed to apply.""" + """Turn a stored TIMESTAMPTZ back into the Unix epoch seconds it holds.""" return datetime.fromisoformat(str(stored)).replace(tzinfo=timezone.utc).timestamp() @@ -100,12 +98,7 @@ def test_pull_and_upsert(self): self.assertEqual(row["queue_time_ms"], 500) self.assertEqual(row["parent_id"], "aabbccddeeff00112233445566778899") self.assertEqual(bool(row["failed"]), False) - # Stored times are shifted back by the session's deterministic offset; the - # 8s execution window asserted above is what survives the shift. - self.assertEqual( - _parse_shifted(row["finished_at"]), - 9.0 - demo_obfuscation.shift_for_session("req1"), - ) + self.assertEqual(_parse_shifted(row["finished_at"]), 9.0) self.assertIsNotNone(row["created_at"]) def test_parent_id_defaults_to_none_when_absent(self): @@ -420,13 +413,7 @@ def test_send_agent_information_inserts_and_updates(self): self.assertEqual(fetched[4], 3) self.assertEqual(fetched[5], 5) self.assertEqual(fetched[6], 1.0) - # Agent rows get a fresh random shift per write (not the deterministic - # per-session one), so only the window is assertable. - self.assertGreaterEqual( - _parse_shifted(fetched[7]), - 1.0 - demo_obfuscation.RANDOM_SHIFT_MAX_SECONDS, - ) - self.assertLessEqual(_parse_shifted(fetched[7]), 1.0) + self.assertEqual(_parse_shifted(fetched[7]), 1.0) self.assertEqual(fetched[8], 2) self.assertEqual(fetched[9], 4) self.assertEqual(fetched[10], "AgentA") @@ -457,11 +444,7 @@ def test_send_agent_information_inserts_and_updates(self): self.assertEqual(fetched[4], 7) self.assertEqual(fetched[5], 0) # reset to 0 after the poll interval drained it self.assertEqual(fetched[6], 0.0) - self.assertGreaterEqual( - _parse_shifted(fetched[7]), - 5.0 - demo_obfuscation.RANDOM_SHIFT_MAX_SECONDS, - ) - self.assertLessEqual(_parse_shifted(fetched[7]), 5.0) + self.assertEqual(_parse_shifted(fetched[7]), 5.0) self.assertEqual(fetched[8], 0) # failures reset to 0 after the poll interval drained it self.assertEqual(fetched[9], 0) # errors reset to 0 after the poll interval drained it self.assertEqual(fetched[10], "AgentA") diff --git a/tests/test_session_store.py b/tests/test_session_store.py index c1c641b..dfefec5 100644 --- a/tests/test_session_store.py +++ b/tests/test_session_store.py @@ -9,12 +9,11 @@ from sqlalchemy import text -import ventis.controller.utils.demo_obfuscation as demo_obfuscation import ventis.controller.utils.session_store as session_store def _stored_epoch(stored): - """Unix epoch seconds held by a stored TIMESTAMPTZ, for asserting the shift.""" + """Unix epoch seconds held by a stored TIMESTAMPTZ.""" return datetime.fromisoformat(str(stored)).replace(tzinfo=timezone.utc).timestamp() @@ -60,10 +59,7 @@ def test_upsert_session_inserts_a_row_with_running_status(self): self.assertEqual(row["project_id"], "11111111-1111-1111-1111-111111111111") self.assertEqual(row["status"], "running") self.assertEqual(row["created_at"], row["updated_at"]) - self.assertEqual( - _stored_epoch(row["created_at"]), - 1000.0 - demo_obfuscation.shift_for_session("req1"), - ) + self.assertEqual(_stored_epoch(row["created_at"]), 1000.0) def test_upsert_session_transitions_status_in_place(self): session_store.upsert_session( @@ -79,13 +75,11 @@ def test_upsert_session_transitions_status_in_place(self): .fetchone() ) # created_at/project_id come from the first call and stay put; status - # and updated_at reflect the second call. Both calls shift by the same - # per-session offset, so the 998s the request took is preserved. - shift = demo_obfuscation.shift_for_session("req1") - self.assertEqual(_stored_epoch(row["created_at"]), 1.0 - shift) + # and updated_at reflect the second call. + self.assertEqual(_stored_epoch(row["created_at"]), 1.0) self.assertEqual(row["project_id"], "11111111-1111-1111-1111-111111111111") self.assertEqual(row["status"], "completed") - self.assertEqual(_stored_epoch(row["updated_at"]), 999.0 - shift) + self.assertEqual(_stored_epoch(row["updated_at"]), 999.0) def test_upsert_session_input_and_output_round_trip(self): session_store.upsert_session( diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index 3599704..9f53a3e 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -294,10 +294,9 @@ def run(self): logger.error("Invalid JSON in request: %s", raw) except Exception as e: logger.error("Error processing request: %s", e) - if isinstance(data, dict): - self._mark_future_failed( + self._mark_future_failed( data.get("future_id"), e, data.get("origin") - ) + ) else: time.sleep(0.001) except KeyboardInterrupt: @@ -617,31 +616,22 @@ def _execute_locally( cpu_percent = ( (cpu_seconds / wall_duration * 100.0) if wall_duration else 0.0 ) - try: - gpu_percent = read_gpu_percent() - except Exception as metric_error: - logger.error("Failed to read GPU metrics: %s", metric_error) - gpu_percent = "" - - try: - self.redis.hset_multiple( - f"future:{future_id}:metrics", - { - "finished_at": wall_end, - "cpu_resource": cpu_percent, - "gpu_resource": gpu_percent, - "agent": self.agent_id, - **( - { - "queue_time": max(wall_start - submitted_at, 0.0) - } - if submitted_at is not None - else {} - ), - }, - ) - except Exception as metric_error: - logger.error("Failed to finalize metrics for future %s: %s", future_id, metric_error) + gpu_percent = read_gpu_percent() + + self.redis.hset_multiple( + f"future:{future_id}:metrics", + { + "finished_at": wall_end, + "cpu_resource": cpu_percent, + "gpu_resource": gpu_percent, + "agent": self.agent_id, + **( + {"queue_time": max(wall_start - submitted_at, 0.0)} + if submitted_at is not None + else {} + ), + }, + ) ventis_context.set_current_future_id(parent or "") # ------------------------------------------------------------------ # diff --git a/ventis/controller/local_controller_frontend.py b/ventis/controller/local_controller_frontend.py index a6e478d..0117c4a 100644 --- a/ventis/controller/local_controller_frontend.py +++ b/ventis/controller/local_controller_frontend.py @@ -67,6 +67,9 @@ def WriteResult(self, request, context): {"failed": failed, "error_message": error_message}, ) if failed: + self.redis.hset( + f"future:{future_id}", "error", error_message or "Unknown error" + ) logger.info("WriteResult: wrote error for future %s", future_id) elif result is not None: self.redis.hset(f"future:{future_id}", "result", result) diff --git a/ventis/controller/utils/demo_obfuscation.py b/ventis/controller/utils/demo_obfuscation.py deleted file mode 100644 index 763c289..0000000 --- a/ventis/controller/utils/demo_obfuscation.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Timestamp obfuscation shared by every table that records demo run data. - -The demo database is populated from real runs, but the wall-clock times those runs -happened at are not what should be on display. Every row belonging to a session is -shifted back by the same deterministic per-session amount, so sessions spread out -over the last 30 days while the relative timing *within* a request -- queue waits, -execution times, the order its futures ran in -- survives exactly. - -Every writer must derive its shift the same way from the same session id, or a -session row and the runtime_information rows for its own futures land at different -points in time. That is the reason this lives in one module rather than next to -either writer: session_store.py (the `session` table) runs inside the workflow -container, sqlalchemy.py (`runtime_information`) runs in the global controller, and -they must agree. - -Kept dependency-free on purpose: session_store.py is copied standalone into the -workflow image (see stub_generator.py's files_to_copy), so anything it imports has -to be importable as a flat module with no `ventis` package around it. -""" - -import hashlib -from datetime import datetime, timezone - -RANDOM_SHIFT_MAX_SECONDS = 2592000 # 30 days - - -def shift_for_session(session_id): - """Deterministic pseudo-random shift (0 to RANDOM_SHIFT_MAX_SECONDS) shared - by every row belonging to the same session, so a single request's rows all - land at the same shifted point in time relative to each other.""" - digest = hashlib.sha256(str(session_id).encode()).digest() - return int.from_bytes(digest[:8], "big") % RANDOM_SHIFT_MAX_SECONDS - - -def to_timestamptz(epoch_seconds, shift): - """Convert a Unix-epoch float/string (as gathered via time.time()) to a - timezone-aware datetime for binding to a TIMESTAMPTZ column, shifted back - by `shift` seconds.""" - - epoch_seconds = float(epoch_seconds) - shift - return datetime.fromtimestamp(epoch_seconds, tz=timezone.utc) diff --git a/ventis/controller/utils/session_store.py b/ventis/controller/utils/session_store.py index b2e4767..6be176b 100644 --- a/ventis/controller/utils/session_store.py +++ b/ventis/controller/utils/session_store.py @@ -11,19 +11,10 @@ import json import os +from datetime import datetime, timezone from sqlalchemy import create_engine, text -# Absolute package when installed, flat module inside the workflow container image -# (stub_generator.py copies both this file and demo_obfuscation.py in side by side). -try: - from ventis.controller.utils.demo_obfuscation import ( - shift_for_session, - to_timestamptz, - ) -except ImportError: - from demo_obfuscation import shift_for_session, to_timestamptz - SESSION_TABLE_NAME = "session" _SESSION_UPSERT = text( @@ -71,9 +62,7 @@ def upsert_session( output_payload=None, ): """Create or update a session row with the given status.""" - # Same per-session shift runtime_information uses, keyed on the same id, so a - # session and the rows for its own futures stay on one timeline. - ts = to_timestamptz(timestamp, shift_for_session(session_id)) + ts = datetime.fromtimestamp(float(timestamp), tz=timezone.utc) with _get_engine(database_url).begin() as conn: conn.execute( _SESSION_UPSERT, diff --git a/ventis/controller/utils/sqlalchemy.py b/ventis/controller/utils/sqlalchemy.py index 2220880..cfedbce 100644 --- a/ventis/controller/utils/sqlalchemy.py +++ b/ventis/controller/utils/sqlalchemy.py @@ -2,16 +2,11 @@ import logging import os -import random import time +from datetime import datetime, timezone from sqlalchemy import create_engine, text from ventis.controller.utils import pricing -from ventis.controller.utils.demo_obfuscation import ( - RANDOM_SHIFT_MAX_SECONDS, - shift_for_session, - to_timestamptz, -) from ventis.utils.redis_client import RedisClient logger = logging.getLogger(__name__) @@ -201,7 +196,6 @@ def send_runtime_information( continue start = float(raw.get("created_at") or 0) end = float(raw.get("finished_at")) - shift = shift_for_session(session_id) input_token_count = int(float(raw.get("input_token_count") or 0)) output_token_count = int(float(raw.get("output_token_count") or 0)) token_count = int(float(raw.get("token_count") or 0)) @@ -231,8 +225,8 @@ def send_runtime_information( "model": model, "cpu": float(raw.get("cpu_resource") or 0), "gpu": float(raw.get("gpu_resource", 0)), - "started_at": to_timestamptz(start, shift), - "finished_at": to_timestamptz(end, shift), + "started_at": datetime.fromtimestamp(start, tz=timezone.utc), + "finished_at": datetime.fromtimestamp(end, tz=timezone.utc), "execution_time_ms": round((end - start) * 1000), "queue_time_ms": round(float(raw.get("queue_time") or 0) * 1000), "input_token_count": input_token_count, @@ -274,9 +268,8 @@ def send_agent_information(rows, database_url=""): "full_failures": int(raw.get("full_failures") or 0), "requests_served": int(float(raw.get("requests_served") or 0)), "throughput": float(raw.get("throughput") or 0.0), - "updated_at": to_timestamptz( - raw.get("updated_at") or now, - random.randint(0, RANDOM_SHIFT_MAX_SECONDS), + "updated_at": datetime.fromtimestamp( + float(raw.get("updated_at") or now), tz=timezone.utc ), }, ) diff --git a/ventis/stub_generator.py b/ventis/stub_generator.py index 64ef1fc..b19f14f 100644 --- a/ventis/stub_generator.py +++ b/ventis/stub_generator.py @@ -439,11 +439,6 @@ def generate_workflow_docker( os.path.join(script_dir, "controller", "utils", "session_store.py"), "session_store.py", ), - # session_store.py imports this for its timestamp shift. - ( - os.path.join(script_dir, "controller", "utils", "demo_obfuscation.py"), - "demo_obfuscation.py", - ), ] # Copy stub files