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_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 aafa3e1..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,14 @@ 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 + ) + self.assertEqual( + redis.hget("future:future-2:metrics", "error_message"), "nope" + ) self.assertEqual( redis.hget("controller:localhost:50051:metrics", "requests_served"), 1 ) @@ -182,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/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 efcc4d0..9f53a3e 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,11 +294,34 @@ def run(self): logger.error("Invalid JSON in request: %s", raw) except Exception as e: logger.error("Error processing request: %s", e) + 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. @@ -334,15 +358,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. @@ -360,6 +387,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: @@ -454,6 +486,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 != "": @@ -501,6 +544,8 @@ def _execute_locally( "method": function, "args": json.dumps(args), "created_at": wall_start, + "failed": 0, + "error_message": "", }, ) if request_id: @@ -510,11 +555,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) @@ -539,7 +588,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", @@ -552,31 +607,32 @@ 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(f"future:{future_id}:metrics", "failed", 1) - 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 "") + 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 + ) + gpu_percent = read_gpu_percent() - if submitted_at is not None: - self.redis.hset( + self.redis.hset_multiple( f"future:{future_id}:metrics", - "queue_time", - max(wall_start - submitted_at, 0.0), + { + "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 "") # ------------------------------------------------------------------ # # Request forwarding # @@ -607,12 +663,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._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", @@ -623,7 +679,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: @@ -637,6 +698,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._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..0117c4a 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,16 @@ 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: + self.redis.hset( + f"future:{future_id}", "error", error_message or "Unknown error" + ) 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/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/future.py b/ventis/future.py index c63c466..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.""" @@ -149,6 +153,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 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: 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