diff --git a/.gitignore b/.gitignore index 0b2e769..de8776e 100644 --- a/.gitignore +++ b/.gitignore @@ -41,3 +41,6 @@ docker_container/ AWSCLIV2.pkg .python-version uv.lock + +Agent Artifacts +docs/ \ No newline at end of file diff --git a/tests/test_cli.py b/tests/test_cli.py index 6c81905..406b95d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -79,6 +79,7 @@ def test_deploy_runs_ec2_preflight_for_ec2_config( ensure_grpc.assert_called_once_with(os.getcwd()) preflight.assert_called_once_with(config, os.getcwd()) + controller.run.assert_called_once_with() @patch("ventis.cli._ensure_grpc_stubs_importable") @patch("ventis.cli._require_docker_for_ec2") diff --git a/tests/test_deploy.py b/tests/test_deploy.py index bea78fa..3f029cb 100644 --- a/tests/test_deploy.py +++ b/tests/test_deploy.py @@ -13,6 +13,7 @@ class _FakeRedis: def __init__(self): self.store = {} self.ttls = {} + self.hashes = {} def set(self, key, value): self.store[key] = value @@ -29,6 +30,12 @@ def expire(self, key, seconds): if key in self.store: self.ttls[key] = seconds + def hset_multiple(self, name, mapping): + self.hashes.setdefault(name, {}).update(mapping) + + def hgetall(self, name): + return dict(self.hashes.get(name, {})) + class _SyncThread: """Runs the target synchronously instead of on a real thread. None of @@ -338,5 +345,122 @@ def test_skips_the_fallback_when_the_database_is_not_configured(self): mock_get.assert_not_called() +class DeployLiveIdentityTests(unittest.TestCase): + """Bug E: VENTIS_PROJECT_ID/VENTIS_DATABASE_URL are Docker env vars frozen at container + launch. A GlobalController reload (SIGHUP) publishes the current project/database identity + to Redis (controller:identity); this container must read that fresh on every request + instead of trusting the env vars it booted with, or a project switch leaves it creating + session rows under the *old* project indefinitely.""" + + def setUp(self): + os.environ["VENTIS_DATABASE_URL"] = "postgresql://example/old-db" + os.environ["VENTIS_PROJECT_ID"] = "11111111-1111-1111-1111-111111111111" + + def tearDown(self): + os.environ.pop("VENTIS_DATABASE_URL", None) + os.environ.pop("VENTIS_PROJECT_ID", None) + + def test_a_value_already_in_redis_at_boot_overrides_the_env_var(self): + with patch.object(deploy_module, "upsert_session") as mock_upsert, \ + _deployed_app() as app: + app.fake_redis.hset_multiple( + deploy_module.IDENTITY_KEY, + { + "project_id": "22222222-2222-2222-2222-222222222222", + "database_url": "postgresql://example/new-db", + }, + ) + client = app.test_client() + client.post("/_noop_workflow", json={"x": 2}) + + first_call_args = mock_upsert.call_args_list[0].args + self.assertEqual(first_call_args[0], "postgresql://example/new-db") + self.assertEqual(first_call_args[1], "22222222-2222-2222-2222-222222222222") + + def test_a_switch_between_two_requests_is_picked_up_by_the_second_one(self): + """Directly reproduces the live incident: request A lands under the project that + was current when it was submitted; a reload happens (simulated here by the + controller writing a new value to the same Redis key a real SIGHUP would update); + request B, with no restart of this container in between, must land under the new + project -- not the one baked into this process's env vars at boot.""" + with patch.object(deploy_module, "upsert_session") as mock_upsert, \ + _deployed_app() as app: + client = app.test_client() + + client.post("/_noop_workflow", json={"x": 1}) + project_after_a = mock_upsert.call_args_list[0].args[1] + + app.fake_redis.hset_multiple( + deploy_module.IDENTITY_KEY, + { + "project_id": "22222222-2222-2222-2222-222222222222", + "database_url": "postgresql://example/new-db", + }, + ) + + client.post("/_noop_workflow", json={"x": 2}) + project_after_b = mock_upsert.call_args_list[-1].args[1] + + self.assertEqual(project_after_a, "11111111-1111-1111-1111-111111111111") + self.assertEqual(project_after_b, "22222222-2222-2222-2222-222222222222") + + def test_completion_and_failure_writes_also_use_the_live_value_not_the_boot_one(self): + """The running/completed/failed transitions are three separate call sites in + deploy.py -- a switch mid-request must not leave the later ones (written from the + background thread, after the switch) tagging with the value the request started + under.""" + with patch.object(deploy_module, "upsert_session") as mock_upsert, \ + _deployed_app(workflow_fn=_failing_workflow) as app: + app.fake_redis.hset_multiple( + deploy_module.IDENTITY_KEY, + { + "project_id": "22222222-2222-2222-2222-222222222222", + "database_url": "postgresql://example/new-db", + }, + ) + client = app.test_client() + client.post("/_failing_workflow", json={"x": 2}) + + projects_used = [call.args[1] for call in mock_upsert.call_args_list] + self.assertEqual( + projects_used, + [ + "22222222-2222-2222-2222-222222222222", + "22222222-2222-2222-2222-222222222222", + ], + ) + + def test_falls_back_to_the_boot_time_env_var_when_redis_has_no_identity_yet(self): + """A request racing the controller's own startup write must still get a value, + not silently skip the session upsert.""" + with patch.object(deploy_module, "upsert_session") as mock_upsert, \ + _deployed_app() as app: + client = app.test_client() + client.post("/_noop_workflow", json={"x": 2}) + + first_call_args = mock_upsert.call_args_list[0].args + self.assertEqual(first_call_args[0], "postgresql://example/old-db") + self.assertEqual(first_call_args[1], "11111111-1111-1111-1111-111111111111") + + def test_status_lookup_after_expiry_also_uses_the_live_value(self): + with patch.object(deploy_module, "get_session") as mock_get, \ + _deployed_app() as app: + app.fake_redis.hset_multiple( + deploy_module.IDENTITY_KEY, + { + "project_id": "22222222-2222-2222-2222-222222222222", + "database_url": "postgresql://example/new-db", + }, + ) + mock_get.return_value = None + app.test_client().get("/status/expired-id") + + mock_get.assert_called_once_with( + "postgresql://example/new-db", + "22222222-2222-2222-2222-222222222222", + "expired-id", + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_global_controller_cleanup.py b/tests/test_global_controller_cleanup.py new file mode 100644 index 0000000..5d2c264 --- /dev/null +++ b/tests/test_global_controller_cleanup.py @@ -0,0 +1,249 @@ +import json +import os +import sys +import unittest + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "grpc_stubs"))) + +from ventis.controller.global_controller import GlobalController +import local_controler_pb2 + + +class _FakeRedis: + def __init__(self, sets=None, hashes=None): + self.sets = sets or {} + self.hashes = hashes or {} + + def sadd(self, name, *values): + self.sets.setdefault(name, set()).update(values) + + def srem(self, name, *values): + self.sets.get(name, set()).difference_update(values) + + def smembers(self, name): + return set(self.sets.get(name, set())) + + def hgetall(self, name): + return dict(self.hashes.get(name, {})) + + +class _FakeStub: + def __init__(self): + self.calls = [] + + def Cleanup(self, request): + self.calls.append(json.loads(request.resonse)) + return local_controler_pb2.JsonResponse(resonse="Cleanup triggered") + + +class _SpyRedis(_FakeRedis): + """_FakeRedis that also records every srem call, so a test can assert + exactly which client a drain happened against (and that a client with + nothing completed is never touched at all).""" + + def __init__(self, sets=None): + super().__init__(sets) + self.srem_calls = [] + + def srem(self, name, *values): + self.srem_calls.append((name, values)) + super().srem(name, *values) + + +class _FailingStub: + def Cleanup(self, request): + raise RuntimeError("connection refused") + + +class _FakeInstanceManager: + def __init__(self, instances): + self._instances = instances + + def list_instances(self): + return self._instances + + +def _bare_controller(redis, instances, node_redis=None): + """Build a GlobalController without running its heavy __init__. + + `node_redis` optionally simulates the host -> RedisClient map that + _launch_redis_containers()/EC2 bootstrap populate. Passing None (the + default) leaves the attribute unset entirely, reproducing a controller + that never got node_redis set at all -- _trigger_cleanup must tolerate + this and fall back to `redis` alone. + """ + controller = GlobalController.__new__(GlobalController) + controller.redis = redis + controller.instance_manager = _FakeInstanceManager(instances) + controller._lc_stubs = {} + if node_redis is not None: + controller.node_redis = node_redis + return controller + + +class TriggerCleanupTests(unittest.TestCase): + def test_sends_one_batched_call_per_instance_not_per_request(self): + # A backlog of many completed requests must not multiply the number of + # gRPC calls per instance -- one call per instance carrying the whole + # batch, regardless of how large the backlog is. + completed = {f"req{i}" for i in range(25)} + expected = set(completed) # snapshot -- _FakeRedis aliases this set, and + # _trigger_cleanup drains "request:completed" via srem in place. + redis = _FakeRedis({"request:completed": completed}) + instances = [{"endpoint": f"host{i}:50051"} for i in range(3)] + controller = _bare_controller(redis, instances) + + stubs = {instance["endpoint"]: _FakeStub() for instance in instances} + controller._get_lc_stub = lambda endpoint: stubs[endpoint] + + controller._trigger_cleanup() + + for endpoint, stub in stubs.items(): + self.assertEqual( + len(stub.calls), 1, f"expected exactly one Cleanup call to {endpoint}" + ) + self.assertEqual(set(stub.calls[0]["request_ids"]), expected) + + # Drained after broadcasting, same as before. + self.assertEqual(redis.smembers("request:completed"), set()) + + def test_one_instance_failing_does_not_block_others_or_stop_draining(self): + completed = {"reqA", "reqB"} + expected = set(completed) # snapshot -- see note in the test above + redis = _FakeRedis({"request:completed": completed}) + instances = [{"endpoint": "good:50051"}, {"endpoint": "bad:50051"}] + controller = _bare_controller(redis, instances) + + good_stub = _FakeStub() + stubs = {"good:50051": good_stub, "bad:50051": _FailingStub()} + controller._get_lc_stub = lambda endpoint: stubs[endpoint] + + controller._trigger_cleanup() # must not raise + + self.assertEqual(len(good_stub.calls), 1) + self.assertEqual(set(good_stub.calls[0]["request_ids"]), expected) + self.assertEqual(redis.smembers("request:completed"), set()) + + def test_noop_when_nothing_completed(self): + redis = _FakeRedis() + instances = [{"endpoint": "host0:50051"}] + controller = _bare_controller(redis, instances) + + stub = _FakeStub() + controller._get_lc_stub = lambda endpoint: stub + + controller._trigger_cleanup() + self.assertEqual(stub.calls, []) + + def test_noop_when_no_instances_registered(self): + redis = _FakeRedis({"request:completed": {"req1"}}) + controller = _bare_controller(redis, []) + + # Should still drain the completed set even with nothing to broadcast to. + controller._trigger_cleanup() + self.assertEqual(redis.smembers("request:completed"), set()) + + +class MultiNodeTriggerCleanupTests(unittest.TestCase): + """Bug D: _trigger_cleanup must not be blind to non-localhost node Redis + instances -- each replica (local or EC2) records its own completions in + its own Redis, never centrally. See CLEANUP_FIX.md.""" + + def test_completed_request_only_on_non_localhost_node_gets_cleaned(self): + localhost_redis = _FakeRedis() + ec2_redis = _FakeRedis({"request:completed": {"reqE"}}) + node_redis = {"localhost": localhost_redis, "10.0.0.5": ec2_redis} + controller = _bare_controller( + localhost_redis, [{"endpoint": "wf:50051"}], node_redis=node_redis + ) + + stub = _FakeStub() + controller._get_lc_stub = lambda endpoint: stub + + controller._trigger_cleanup() + + self.assertEqual(len(stub.calls), 1) + self.assertEqual(set(stub.calls[0]["request_ids"]), {"reqE"}) + self.assertEqual(ec2_redis.smembers("request:completed"), set()) + + def test_requests_across_multiple_nodes_batched_into_one_call_per_instance(self): + localhost_redis = _FakeRedis({"request:completed": {"reqA", "reqB"}}) + ec2_redis_1 = _FakeRedis({"request:completed": {"reqC"}}) + ec2_redis_2 = _FakeRedis({"request:completed": {"reqD", "reqE"}}) + node_redis = { + "localhost": localhost_redis, + "ec2-1": ec2_redis_1, + "ec2-2": ec2_redis_2, + } + instances = [{"endpoint": f"host{i}:50051"} for i in range(3)] + controller = _bare_controller(localhost_redis, instances, node_redis=node_redis) + + stubs = {instance["endpoint"]: _FakeStub() for instance in instances} + controller._get_lc_stub = lambda endpoint: stubs[endpoint] + + controller._trigger_cleanup() + + expected = {"reqA", "reqB", "reqC", "reqD", "reqE"} + for endpoint, stub in stubs.items(): + self.assertEqual( + len(stub.calls), 1, f"expected exactly one Cleanup call to {endpoint}" + ) + self.assertEqual(set(stub.calls[0]["request_ids"]), expected) + + for redis in (localhost_redis, ec2_redis_1, ec2_redis_2): + self.assertEqual(redis.smembers("request:completed"), set()) + + def test_each_node_drained_from_its_own_client_not_cross_contaminated(self): + redis_with_data = _SpyRedis({"request:completed": {"reqX"}}) + redis_empty = _SpyRedis() + node_redis = {"a": redis_with_data, "b": redis_empty} + controller = _bare_controller( + redis_with_data, [{"endpoint": "wf:50051"}], node_redis=node_redis + ) + + stub = _FakeStub() + controller._get_lc_stub = lambda endpoint: stub + + controller._trigger_cleanup() + + self.assertEqual(redis_with_data.srem_calls, [("request:completed", ("reqX",))]) + self.assertEqual(redis_empty.srem_calls, []) + + def test_falls_back_to_self_redis_when_node_redis_attribute_missing(self): + # No node_redis kwarg at all -- the attribute genuinely doesn't exist, + # reproducing a controller built before _launch_redis_containers() ever + # ran. Behavior must match the pre-Bug-D single-redis path exactly. + completed = {"req1", "req2"} + redis = _FakeRedis({"request:completed": set(completed)}) + controller = _bare_controller(redis, [{"endpoint": "host0:50051"}]) + self.assertFalse(hasattr(controller, "node_redis")) + + stub = _FakeStub() + controller._get_lc_stub = lambda endpoint: stub + + controller._trigger_cleanup() + + self.assertEqual(set(stub.calls[0]["request_ids"]), completed) + self.assertEqual(redis.smembers("request:completed"), set()) + + def test_falls_back_to_self_redis_when_node_redis_is_empty_dict(self): + # node_redis present but empty -- the window right at the start of + # __init__, before _launch_redis_containers() populates it. + completed = {"req1"} + redis = _FakeRedis({"request:completed": set(completed)}) + controller = _bare_controller( + redis, [{"endpoint": "host0:50051"}], node_redis={} + ) + + stub = _FakeStub() + controller._get_lc_stub = lambda endpoint: stub + + controller._trigger_cleanup() + + self.assertEqual(set(stub.calls[0]["request_ids"]), completed) + self.assertEqual(redis.smembers("request:completed"), set()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_global_controller_identity.py b/tests/test_global_controller_identity.py new file mode 100644 index 0000000..e4337f7 --- /dev/null +++ b/tests/test_global_controller_identity.py @@ -0,0 +1,142 @@ +"""Bug E: a Workflow container's VENTIS_PROJECT_ID/VENTIS_DATABASE_URL env vars are frozen at +launch. _write_identity() publishes the controller's current project/database identity to +every node's Redis (mirroring the existing policy:rules/routing_table:* pattern) so deploy.py's +_current_identity() can read it live instead of trusting a boot-time env var. reload_config() +must call this too, or the fix is only half-applied -- see the sibling test in +test_global_controller_reload.py for the assign_project_id() half of the same bug class. +""" + +import os +import sys +import unittest +from types import SimpleNamespace + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from ventis.controller.global_controller import GlobalController + + +class _FakeRedis: + def __init__(self): + self.hashes = {} + + def hset_multiple(self, name, mapping): + self.hashes.setdefault(name, {}).update(mapping) + + def hgetall(self, name): + return dict(self.hashes.get(name, {})) + + +def _bare_controller(config, node_redis=None): + controller = GlobalController.__new__(GlobalController) + controller.config = config + controller.redis = _FakeRedis() + controller.node_redis = node_redis if node_redis is not None else {} + return controller + + +class WriteIdentityTests(unittest.TestCase): + def test_publishes_project_id_and_database_url_to_self_redis_when_no_node_redis(self): + controller = _bare_controller( + { + "project_id": "11111111-1111-1111-1111-111111111111", + "database": {"url": "postgresql://example/db"}, + } + ) + + controller._write_identity() + + self.assertEqual( + controller.redis.hgetall(GlobalController.IDENTITY_KEY), + { + "project_id": "11111111-1111-1111-1111-111111111111", + "database_url": "postgresql://example/db", + }, + ) + + def test_publishes_to_every_node_not_just_self_redis(self): + """The same lesson as Bug D: a Workflow replica on a remote EC2 host has its own + Redis, distinct from the controller's local one. Publishing to self.redis alone + would leave that replica reading nothing.""" + node_a = _FakeRedis() + node_b = _FakeRedis() + controller = _bare_controller( + { + "project_id": "11111111-1111-1111-1111-111111111111", + "database": {"url": "postgresql://example/db"}, + }, + node_redis={"localhost": node_a, "172.31.0.5": node_b}, + ) + + controller._write_identity() + + for node in (node_a, node_b): + self.assertEqual( + node.hgetall(GlobalController.IDENTITY_KEY)["project_id"], + "11111111-1111-1111-1111-111111111111", + ) + + def test_a_second_call_with_a_new_config_overwrites_the_published_value(self): + """Simulates what reload_config() does on a real SIGHUP: the same key must reflect + the *new* project after a switch, not just get set once at boot.""" + node = _FakeRedis() + controller = _bare_controller( + { + "project_id": "11111111-1111-1111-1111-111111111111", + "database": {"url": "postgresql://example/old-db"}, + }, + node_redis={"localhost": node}, + ) + controller._write_identity() + + controller.config = { + "project_id": "22222222-2222-2222-2222-222222222222", + "database": {"url": "postgresql://example/new-db"}, + } + controller._write_identity() + + self.assertEqual( + node.hgetall(GlobalController.IDENTITY_KEY), + { + "project_id": "22222222-2222-2222-2222-222222222222", + "database_url": "postgresql://example/new-db", + }, + ) + + def test_missing_project_id_or_database_publishes_safe_defaults(self): + controller = _bare_controller({}) + + controller._write_identity() + + self.assertEqual( + controller.redis.hgetall(GlobalController.IDENTITY_KEY), + {"project_id": "0", "database_url": ""}, + ) + + +class ReloadConfigWritesIdentityTests(unittest.TestCase): + """reload_config() itself must call _write_identity() -- not just have the method exist.""" + + def test_reload_config_republishes_identity(self): + node = _FakeRedis() + controller = _bare_controller( + {"project_id": "11111111-1111-1111-1111-111111111111", "agents": []}, + node_redis={"localhost": node}, + ) + controller.config_path = None + controller.instance_manager = SimpleNamespace(publish_routing_snapshot=lambda *_: None) + controller._load_config = lambda path: { + "project_id": "22222222-2222-2222-2222-222222222222", + "agents": [], + } + + controller.reload_config() + + self.assertEqual( + node.hgetall(GlobalController.IDENTITY_KEY)["project_id"], + "22222222-2222-2222-2222-222222222222", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_global_controller_redis_reuse.py b/tests/test_global_controller_redis_reuse.py new file mode 100644 index 0000000..1c1682d --- /dev/null +++ b/tests/test_global_controller_redis_reuse.py @@ -0,0 +1,102 @@ +"""Fix C: a restart must not unconditionally wipe and recreate each node's Redis container. + +_launch_redis_containers() used to `docker run` a fresh ventis-redis- container on every +__init__, unconditionally -- wiping every `agent_instance:*` record InstanceManager needs to +recognize already-running EC2 replicas as reusable. ensure_instances()'s dedup logic was already +correct; it was just fed an empty Redis on every restart, so it reprovisioned everything from +scratch and orphaned the previous replicas. The fix: check whether the existing container is +already healthy before recreating it. +""" + +import os +import sys +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from ventis.controller.global_controller import GlobalController + + +def _bare_controller(controllers): + controller = GlobalController.__new__(GlobalController) + controller.controllers = controllers + controller.redis_containers = {} + controller.node_redis = {} + controller.redis = None + return controller + + +class RedisContainerReuseTests(unittest.TestCase): + def _run(self, controller, inspect_stdout): + run_calls = [] + + def fake_run_cmd(cmd, host, user=None): + run_calls.append(cmd) + if cmd[:2] == ["docker", "inspect"]: + return SimpleNamespace(returncode=0, stdout=inspect_stdout, stderr="") + return SimpleNamespace(returncode=0, stdout="", stderr="") + + controller._run_cmd = fake_run_cmd + + with patch("ventis.controller.global_controller.RedisClient") as fake_redis_cls, patch( + "ventis.controller.global_controller._wait_for_redis" + ): + fake_redis_cls.return_value = MagicMock() + controller._launch_redis_containers() + + return [c for c in run_calls if c[:2] == ["docker", "run"]] + + def test_a_healthy_existing_container_is_reused_not_recreated(self): + controller = _bare_controller( + [{"name": "Workflow", "replicas": 1, "redis_port": 6379}] + ) + + docker_run_calls = self._run(controller, inspect_stdout="true\n") + + self.assertEqual( + docker_run_calls, [], "a healthy existing Redis container must not be recreated" + ) + self.assertIn("localhost", controller.redis_containers) + self.assertIn("localhost", controller.node_redis) + + def test_an_unhealthy_or_missing_container_still_gets_created(self): + controller = _bare_controller( + [{"name": "Workflow", "replicas": 1, "redis_port": 6379}] + ) + + docker_run_calls = self._run(controller, inspect_stdout="false\n") + + self.assertEqual( + len(docker_run_calls), 1, "a not-running container must still be (re)created" + ) + self.assertIn("localhost", controller.redis_containers) + + def test_reuse_check_probes_the_exact_expected_container_name(self): + controller = _bare_controller( + [{"name": "Workflow", "replicas": 1, "redis_port": 6379, "host": "10.0.0.5"}] + ) + + inspect_calls = [] + + def fake_run_cmd(cmd, host, user=None): + if cmd[:2] == ["docker", "inspect"]: + inspect_calls.append(cmd) + return SimpleNamespace(returncode=0, stdout="true\n", stderr="") + return SimpleNamespace(returncode=0, stdout="", stderr="") + + controller._run_cmd = fake_run_cmd + + with patch("ventis.controller.global_controller.RedisClient") as fake_redis_cls, patch( + "ventis.controller.global_controller._wait_for_redis" + ): + fake_redis_cls.return_value = MagicMock() + controller._launch_redis_containers() + + self.assertEqual(len(inspect_calls), 1) + self.assertIn("ventis-redis-10-0-0-5", inspect_calls[0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_global_controller_reload.py b/tests/test_global_controller_reload.py new file mode 100644 index 0000000..26d0662 --- /dev/null +++ b/tests/test_global_controller_reload.py @@ -0,0 +1,93 @@ +"""Bug A: reload_config() must re-sync the telemetry writer's cached project_id. + +__init__ calls assign_project_id() once, at boot. reload_config() -- the method meant to run +whenever the box's config changes -- rebuilt self.config/self.controllers/self.poll_interval and +republished the routing snapshot, but never called assign_project_id() again, so a box reclaimed +by a new project kept tagging every telemetry write with whichever project was live at process +start, forever. +""" + +import os +import sys +import tempfile +import unittest + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +import yaml + +import ventis.controller.utils.telemetry_logging as sqlmod +from ventis.controller.global_controller import GlobalController + + +class _FakeInstanceManager: + def publish_routing_snapshot(self, controllers): + pass + + +class _FakeRedis: + def hset_multiple(self, name, mapping): + pass + + +def _write_config(project_id): + f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) + yaml.safe_dump({"agents": [], "poll_interval": 5, "project_id": project_id}, f) + f.close() + return f.name + + +def _bare_controller(config_path): + """Build a GlobalController without running its heavy __init__, wired only with what + reload_config() actually touches -- now including node_redis/redis, since + reload_config() also calls _write_identity() (Bug E).""" + controller = GlobalController.__new__(GlobalController) + controller.config_path = config_path + controller.instance_manager = _FakeInstanceManager() + controller.node_redis = {} + controller.redis = _FakeRedis() + return controller + + +class ReloadConfigResyncTests(unittest.TestCase): + def setUp(self): + sqlmod._project_id = None + + def tearDown(self): + sqlmod._project_id = None + + def test_reload_config_resyncs_project_id_to_the_new_value(self): + config_path = _write_config("22222222-2222-2222-2222-222222222222") + try: + controller = _bare_controller(config_path) + sqlmod.assign_project_id("11111111-1111-1111-1111-111111111111") + + controller.reload_config() + + self.assertEqual(sqlmod._project_id, "22222222-2222-2222-2222-222222222222") + finally: + os.unlink(config_path) + + def test_reload_config_resyncs_on_every_switch_not_just_the_first(self): + """A box reclaimed A -> B -> A must resync every time, not just once.""" + config_a = _write_config("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa") + config_b = _write_config("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb") + try: + controller = _bare_controller(config_a) + + controller.reload_config() + self.assertEqual(sqlmod._project_id, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa") + + controller.config_path = config_b + controller.reload_config() + self.assertEqual(sqlmod._project_id, "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb") + + controller.config_path = config_a + controller.reload_config() + self.assertEqual(sqlmod._project_id, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa") + finally: + os.unlink(config_a) + os.unlink(config_b) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_instance_manager_runtime.py b/tests/test_instance_manager_runtime.py index 88e7577..df35efe 100644 --- a/tests/test_instance_manager_runtime.py +++ b/tests/test_instance_manager_runtime.py @@ -66,7 +66,9 @@ def _fake_controller(): node_redis={}, redis_containers={}, config={"poll_interval": 5}, - _run_cmd=MagicMock(return_value=SimpleNamespace(returncode=0)), + # stdout="" (not running) so the orphan-check `docker inspect` probe that now + # precedes `docker run` reads a real string instead of erroring on a missing attribute. + _run_cmd=MagicMock(return_value=SimpleNamespace(returncode=0, stdout="")), ) @@ -87,6 +89,34 @@ def _fake_runtime(**kwargs): class InstanceManagerRuntimeTests(unittest.TestCase): + def test_bootstrap_instance_removes_an_orphaned_running_container_before_recreating(self): + """We only get here when Redis has no agent_instance record for this replica -- but a + container with this exact name can still be running (e.g. Redis was wiped and rebuilt + empty while the container it used to track kept running). `docker run --name` would just + fail with "name already in use" in that case; treat it as an orphan and clear it first.""" + controller = _fake_controller() + calls = [] + + def fake_run_cmd(cmd, host, user=None): + calls.append(cmd) + if cmd[:2] == ["docker", "inspect"]: + return SimpleNamespace(returncode=0, stdout="true\n") + return SimpleNamespace(returncode=0, stdout="") + + controller._run_cmd = fake_run_cmd + manager = InstanceManager(controller, controller.redis) + + manager.ensure_instances([{"name": "Alpha", "provider": "local"}]) + + self.assertEqual( + [c[:3] for c in calls], + [ + ["docker", "inspect", "-f"], + ["docker", "rm", "-f"], + ["docker", "run", "-d"], + ], + ) + def test_local_instances_keep_default_host_and_increment_host_ports(self): controller = _fake_controller() manager = InstanceManager(controller, controller.redis) @@ -116,7 +146,8 @@ def test_local_instances_keep_default_host_and_increment_host_ports(self): self.assertEqual(beta["host"], "localhost") self.assertEqual(beta["host_port"], "8001") self.assertEqual( - controller._run_cmd.call_args_list[0].args, + # index [1], not [0]: the orphan-check `docker inspect` probe now runs first. + controller._run_cmd.call_args_list[1].args, ( [ "docker", diff --git a/tests/test_local_controller_cleanup.py b/tests/test_local_controller_cleanup.py new file mode 100644 index 0000000..0466bb7 --- /dev/null +++ b/tests/test_local_controller_cleanup.py @@ -0,0 +1,143 @@ +import json +import os +import sys +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "grpc_stubs"))) + +from ventis.controller.local_controller_frontend import LocalControllerServicer +import local_controler_pb2 + + +class _SyncThread: + """Stand-in for threading.Thread that runs its target immediately, inline. + + Cleanup() fires off a daemon Thread and returns without waiting for it, which + makes the real dispatch nondeterministic to assert on in a test. Swapping this + in for the module's Thread makes the dispatch happen synchronously instead. + """ + + def __init__(self, target=None, args=(), daemon=None): + self._target = target + self._args = args + + def start(self): + self._target(*self._args) + + +class CleanupDispatchTests(unittest.TestCase): + def test_batched_request_ids_dispatches_cleanup_for_each(self): + cleaned = [] + servicer = SimpleNamespace(_cleanup_request=lambda rid: cleaned.append(rid)) + request = local_controler_pb2.JsonResponse( + resonse=json.dumps({"request_ids": ["req1", "req2", "req3"]}) + ) + + with patch("ventis.controller.local_controller_frontend.Thread", _SyncThread): + LocalControllerServicer.Cleanup(servicer, request, context=None) + + self.assertEqual(cleaned, ["req1", "req2", "req3"]) + + def test_missing_ids_does_not_dispatch(self): + cleaned = [] + servicer = SimpleNamespace(_cleanup_request=lambda rid: cleaned.append(rid)) + request = local_controler_pb2.JsonResponse(resonse=json.dumps({})) + + with patch("ventis.controller.local_controller_frontend.Thread", _SyncThread): + LocalControllerServicer.Cleanup(servicer, request, context=None) + + self.assertEqual(cleaned, []) + + def test_old_single_request_id_payload_is_no_longer_supported(self): + # The only real caller (GlobalController._trigger_cleanup) always sends + # request_ids now, so the pre-batching {"request_id": ...} shape is + # intentionally treated the same as a missing/empty payload, not dispatched. + cleaned = [] + servicer = SimpleNamespace(_cleanup_request=lambda rid: cleaned.append(rid)) + request = local_controler_pb2.JsonResponse( + resonse=json.dumps({"request_id": "req-legacy"}) + ) + + with patch("ventis.controller.local_controller_frontend.Thread", _SyncThread): + LocalControllerServicer.Cleanup(servicer, request, context=None) + + self.assertEqual(cleaned, []) + + +class _FakeRedisStore: + """Enough of RedisClient's surface for _cleanup_request: strings, sets, setnx.""" + + def __init__(self, strings=None, sets=None): + self.strings = strings or {} + self.sets = sets or {} + + def setnx(self, key, value): + if key in self.strings: + return False + self.strings[key] = value + return True + + def smembers(self, name): + return set(self.sets.get(name, set())) + + def delete(self, *keys): + for key in keys: + self.strings.pop(key, None) + self.sets.pop(key, None) + + +def _bare_servicer(redis): + servicer = LocalControllerServicer.__new__(LocalControllerServicer) + servicer.redis = redis + servicer.my_endpoint = "test-node" + return servicer + + +class CleanupRequestTests(unittest.TestCase): + def test_cleanup_deletes_consolidated_future_hashes_and_bookkeeping(self): + redis = _FakeRedisStore( + sets={"request:req1:futures": {"fut1", "fut2"}}, + strings={ + "future:fut1": "x", + "future:fut2": "x", + "future:fut1:children": "x", + "future:fut1:consumers": "x", + }, + ) + servicer = _bare_servicer(redis) + + servicer._cleanup_request("req1") + + # Future-resolution bookkeeping: gone. + self.assertNotIn("request:req1:futures", redis.sets) + self.assertNotIn("future:fut1", redis.strings) + self.assertNotIn("future:fut2", redis.strings) + self.assertNotIn("future:fut1:children", redis.strings) + self.assertNotIn("future:fut1:consumers", redis.strings) + + def test_cleanup_still_deletes_affinity_bindings(self): + redis = _FakeRedisStore( + sets={"request:req1:futures": {"fut1"}}, + strings={"future:fut1": "x", "affinity:req1": "some-host"}, + ) + servicer = _bare_servicer(redis) + + servicer._cleanup_request("req1") + + self.assertNotIn("affinity:req1", redis.strings) + self.assertNotIn("future:fut1", redis.strings) + + def test_cleanup_releases_its_lock_even_with_no_futures(self): + redis = _FakeRedisStore(sets={"request:req1:futures": set()}) + servicer = _bare_servicer(redis) + + servicer._cleanup_request("req1") + + self.assertNotIn("request:req1:cleanup_lock", redis.strings) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_telemetry_logging.py b/tests/test_telemetry_logging.py index 5bf78d7..89e753f 100644 --- a/tests/test_telemetry_logging.py +++ b/tests/test_telemetry_logging.py @@ -28,6 +28,9 @@ def scan_keys(self, pattern): def hgetall(self, name): return dict(self.hashes.get(name, {})) + def hset(self, name, field, value): + self.hashes.setdefault(name, {})[field] = str(value) + def get(self, name): return self.hashes.get(name) diff --git a/ventis/cli.py b/ventis/cli.py index 2377c04..b43a6b3 100644 --- a/ventis/cli.py +++ b/ventis/cli.py @@ -419,6 +419,16 @@ def _signal_handler(sig, frame): signal.signal(signal.SIGTERM, _signal_handler) atexit.register(controller.cleanup) + # SIGHUP reloads config in place without tearing down any agent. + def _reload_handler(sig, frame): + logger.info("Received SIGHUP, reloading config...") + try: + controller.reload_config() + except Exception as e: + logger.error("Reload failed: %s", e) + + signal.signal(signal.SIGHUP, _reload_handler) + logger.info("Deploying from config: %s", config_path) controller.launch_docker_agents() controller._wait_for_healthy() diff --git a/ventis/controller/cloud_provider_logic/Local/_runtime.py b/ventis/controller/cloud_provider_logic/Local/_runtime.py index 7780ce4..963eef3 100644 --- a/ventis/controller/cloud_provider_logic/Local/_runtime.py +++ b/ventis/controller/cloud_provider_logic/Local/_runtime.py @@ -63,6 +63,17 @@ def bootstrap_instance(provisioned, spec, replica_index, agent_id): endpoint = routing_endpoint_for(provisioned) _require_controller().redis.set(f"controller:{endpoint}:agent_id", agent_id) + inspect = _require_controller()._run_cmd( + ["docker", "inspect", "-f", "{{.State.Running}}", runtime_id], host, user + ) + if inspect.returncode == 0 and inspect.stdout.strip() == "true": + logger.warning( + "No Redis record for %s but a container with that name is already running; " + "treating it as orphaned and recreating.", + runtime_id, + ) + _require_controller()._run_cmd(["docker", "rm", "-f", runtime_id], host, user) + cmd = [ "docker", "run", diff --git a/ventis/controller/global_controller.py b/ventis/controller/global_controller.py index 943c31d..b00cef6 100644 --- a/ventis/controller/global_controller.py +++ b/ventis/controller/global_controller.py @@ -11,6 +11,7 @@ import json import sys import os +from concurrent.futures import ThreadPoolExecutor import yaml from ventis.controller.instance_manager import InstanceManager @@ -57,6 +58,7 @@ class GlobalController(object): ROUTING_STATEFUL_KEY = "routing_table:stateful" SERVICES_SET_KEY = "routing_table:services" POLICY_RULES_KEY = "policy:rules" + IDENTITY_KEY = "controller:identity" # has controllers current project_id and database_url def __init__(self, config_path): self.config_path = config_path @@ -81,7 +83,7 @@ def __init__(self, config_path): self._lc_stubs = {} # endpoint -> gRPC stub self.instance_manager = InstanceManager(self) assign_project_id(self.config.get("project_id",0)) - + # Clean up any stale containers from previous runs self._cleanup_stale_containers() @@ -90,6 +92,7 @@ def __init__(self, config_path): write_agent_specs(self.config_path, self.redis) self._write_resource_specs() self._load_and_write_policies() + self._write_identity() self.instance_manager.publish_routing_snapshot(self.controllers) logger.info( "Global controller initialized with %d controller(s).", @@ -127,6 +130,13 @@ def _cleanup_stale_containers(self): for host, (user, container_names) in host_containers.items(): for container_name in container_names: try: + inspect = self._run_cmd( + ["docker", "inspect", "-f", "{{.State.Running}}", container_name], + host, + user, + ) + if inspect.returncode == 0 and inspect.stdout.strip() == "true": + continue # already running -- a live replica, not stale self._run_cmd(["docker", "rm", "-f", container_name], host, user) except Exception: pass # Container didn't exist, that's fine @@ -165,6 +175,8 @@ def reload_config(self): self.config = self._load_config(self.config_path) self.controllers = self.config.get("agents", []) self.poll_interval = self.config.get("poll_interval", 5) + assign_project_id(self.config.get("project_id", 0)) + self._write_identity() self.instance_manager.publish_routing_snapshot(self.controllers) def _write_resource_specs(self): @@ -216,6 +228,19 @@ def _load_and_write_policies(self): len(rules), ) + # Only relevant for demo purposes + def _write_identity(self): + """Publish the current project/database identity to every node's Redis.""" + payload = { + "project_id": str(self.config.get("project_id", 0)), + "database_url": self.config.get("database", {}).get("url") or "", + } + targets = list(self.node_redis.values()) or [self.redis] + for redis_client in targets: + redis_client.hset_multiple(self.IDENTITY_KEY, payload) + + logger.info("Identity (project %s) published to %d Redis instance(s).", payload["project_id"], len(targets)) + # Routing reads are direct Redis calls now that InstanceManager owns publication: # - self.redis.hgetall(self.ROUTING_ENDPOINTS_KEY) # - self.redis.hget(self.ROUTING_ENDPOINTS_KEY, service_name) @@ -228,14 +253,22 @@ def get_node_redis(self, host): # Redis container management # # ------------------------------------------------------------------ # - def _launch_redis_containers(self): - """ - Launch a Redis Docker container on each unique node. + def _redis_container_healthy(self, container_name, host, user, connect_host, redis_port): + """Check whether an existing Redis container is already up and answering.""" + inspect = self._run_cmd( + ["docker", "inspect", "-f", "{{.State.Running}}", container_name], host, user + ) + if inspect.returncode != 0 or inspect.stdout.strip() != "true": + return False + try: + probe = RedisClient(host=connect_host, port=redis_port) + _wait_for_redis(probe, host, redis_port, timeout=5, interval=1) + return True + except TimeoutError: + return False - Discovers unique hosts from the agent config and starts one - redis:alpine container per host. Creates a RedisClient instance - for each node so the global controller can query any node's Redis. - """ + def _launch_redis_containers(self): + """Launch a Redis container on each unique node, reusing one that's already healthy.""" # Collect unique nodes from all replica placements nodes = {} for ctrl in self.controllers: @@ -252,47 +285,51 @@ def _launch_redis_containers(self): redis_port = node_cfg["redis_port"] user = node_cfg["user"] container_name = f"ventis-redis-{host.replace('.', '-')}" + # For localhost, connect directly; for remote, connect via host IP + connect_host = "localhost" if host in ("localhost", "127.0.0.1") else host - cmd = [ - "docker", - "run", - "-d", - "--name", - container_name, - "-p", - f"{redis_port}:6379", - "redis:alpine", - ] + if self._redis_container_healthy(container_name, host, user, connect_host, redis_port): + logger.info("Reusing existing Redis container %s on %s", container_name, host) + self.redis_containers[host] = container_name + else: + cmd = [ + "docker", + "run", + "-d", + "--name", + container_name, + "-p", + f"{redis_port}:6379", + "redis:alpine", + ] - try: - result = self._run_cmd(cmd, host, user) - if result.returncode == 0: - self.redis_containers[host] = container_name - logger.info( - "Launched Redis container %s on %s:%d", - container_name, - host, - redis_port, - ) - else: + try: + result = self._run_cmd(cmd, host, user) + if result.returncode == 0: + self.redis_containers[host] = container_name + logger.info( + "Launched Redis container %s on %s:%d", + container_name, + host, + redis_port, + ) + else: + logger.critical( + "Failed to launch Redis on %s: %s", + host, + result.stderr.strip(), + ) + sys.exit(1) + except FileNotFoundError: logger.critical( - "Failed to launch Redis on %s: %s", - host, - result.stderr.strip(), + "Docker is not installed or not in PATH. Cannot launch Redis." ) sys.exit(1) - except FileNotFoundError: - logger.critical( - "Docker is not installed or not in PATH. Cannot launch Redis." - ) - sys.exit(1) - except Exception as e: - logger.critical("Failed to launch Redis on %s: %s", host, e) - sys.exit(1) + except Exception as e: + logger.critical("Failed to launch Redis on %s: %s", host, e) + sys.exit(1) # Create a RedisClient for this node - # For localhost, connect directly; for remote, connect via host IP - connect_host = "localhost" if host in ("localhost", "127.0.0.1") else host redis_client = RedisClient(host=connect_host, port=redis_port) _wait_for_redis(redis_client, host, redis_port) self.node_redis[host] = redis_client @@ -537,27 +574,50 @@ def _cleanup_loop(self): logger.warning("Cleanup loop encountered an error: %s", e) def _trigger_cleanup(self): - """Broadcast Cleanup gRPC to all local controllers for each completed request.""" - completed = self.redis.smembers("request:completed") - if not completed: + """Broadcast a batched Cleanup gRPC to all instances for every completed request, gathered from every node's Redis.""" + # Falls back to self.redis alone if node_redis is unset/empty. + node_redis_map = getattr(self, "node_redis", None) or {} + redis_clients = list(node_redis_map.values()) or [self.redis] + + completed_by_client = {} + all_completed = set() + for client in redis_clients: + completed = client.smembers("request:completed") + if completed: + completed_by_client[client] = completed + all_completed.update(completed) + + if not all_completed: return - for request_id in completed: - logger.info("Triggering cleanup for completed request %s", request_id) - for instance in self.instance_manager.list_instances(): - endpoint = instance["endpoint"] - try: - stub = self._get_lc_stub(endpoint) - payload = json.dumps({"request_id": request_id}) - stub.Cleanup(local_controler_pb2.JsonResponse(resonse=payload)) - logger.debug( - "Sent Cleanup for request %s to %s", request_id, endpoint - ) - except Exception as e: - logger.warning("Failed to trigger cleanup on %s: %s", endpoint, e) + payload = json.dumps({"request_ids": list(all_completed)}) - # Remove from completed set after broadcast - self.redis.srem("request:completed", request_id) + def _send(instance): + endpoint = instance["endpoint"] + try: + stub = self._get_lc_stub(endpoint) + stub.Cleanup(local_controler_pb2.JsonResponse(resonse=payload)) + logger.debug( + "Sent Cleanup batch of %d request(s) to %s", + len(all_completed), + endpoint, + ) + except Exception as e: + logger.warning("Failed to trigger cleanup on %s: %s", endpoint, e) + + instances = self.instance_manager.list_instances() + if instances: + with ThreadPoolExecutor(max_workers=len(instances)) as executor: + list(executor.map(_send, instances)) + + logger.info( + "Triggered cleanup for %d completed request(s) across %d node(s)", + len(all_completed), + len(completed_by_client), + ) + # Drain each node's own set from the same client it was read from. + for client, completed in completed_by_client.items(): + client.srem("request:completed", *completed) # ------------------------------------------------------------------ # # Runtime launching # @@ -683,6 +743,16 @@ def _signal_handler(sig, frame): signal.signal(signal.SIGINT, _signal_handler) signal.signal(signal.SIGTERM, _signal_handler) + + # Register config reload on SIGHUP and reload + def _reload_handler(sig, frame): + logger.info("Received SIGHUP, reloading config...") + try: + controller.reload_config() + except Exception as e: + logger.error("Reload failed: %s", e) + + signal.signal(signal.SIGHUP, _reload_handler) atexit.register(controller.cleanup) controller.launch_docker_agents() diff --git a/ventis/controller/local_controller_frontend.py b/ventis/controller/local_controller_frontend.py index 7274726..bf440c8 100644 --- a/ventis/controller/local_controller_frontend.py +++ b/ventis/controller/local_controller_frontend.py @@ -91,22 +91,26 @@ def WriteResult(self, request, context): return local_controler_pb2.JsonResponse(resonse="Result written") def Cleanup(self, request, context): - """Trigger async cleanup of all futures associated with a completed request.""" + """Trigger async cleanup for one or more completed requests.""" try: data = json.loads(request.resonse) - request_id = data.get("request_id") - if request_id: - Thread( - target=self._cleanup_request, args=(request_id,), daemon=True - ).start() + request_ids = data.get("request_ids") + + if request_ids: + # Process the cleanup batch asynchronously so the RPC returns immediately. + def _cleanup_batch(): + for request_id in request_ids: + self._cleanup_request(request_id) + + Thread(target=_cleanup_batch, daemon=True).start() else: - logger.warning("Cleanup: missing request_id in payload") + logger.warning("Cleanup: missing request_id(s) in payload") except Exception as e: logger.error("Cleanup: failed to parse payload: %s", e) return local_controler_pb2.JsonResponse(resonse="Cleanup triggered") def _cleanup_request(self, request_id): - """Delete all futures associated with a request from this node's Redis.""" + """Delete a request's consolidated future hashes and bookkeeping.""" # Atomically claim cleanup — prevents duplicate work when multiple LCs share a Redis lock_key = f"request:{request_id}:cleanup_lock" if not self.redis.setnx(lock_key, self.my_endpoint): diff --git a/ventis/controller/utils/session_logging.py b/ventis/controller/utils/session_logging.py index 6be176b..df8e4dc 100644 --- a/ventis/controller/utils/session_logging.py +++ b/ventis/controller/utils/session_logging.py @@ -39,17 +39,20 @@ """ ) -_engine = None +_engines = {} # resolved url -> Engine def _get_engine(database_url): - global _engine - if _engine is None: - url = os.environ.get("VENTIS_DATABASE_URL", str(database_url)) - if url.startswith("postgresql://"): - url = "postgresql+psycopg://" + url[len("postgresql://"):] - _engine = create_engine(url) - return _engine + """Return a cached Engine for `database_url`, building one on first use per resolved URL.""" + global _engines + url = os.environ.get("VENTIS_DATABASE_URL", str(database_url)) + if url.startswith("postgresql://"): + url = "postgresql+psycopg://" + url[len("postgresql://"):] + engine = _engines.get(url) + if engine is None: + engine = create_engine(url) + _engines[url] = engine + return engine def upsert_session( diff --git a/ventis/deploy.py b/ventis/deploy.py index b6ac721..2c2531b 100644 --- a/ventis/deploy.py +++ b/ventis/deploy.py @@ -47,6 +47,11 @@ def my_workflow(query: str): # How long a finished request's Redis keys stick around before Redis reclaims them. COMPLETED_TTL_SECONDS = 300 +# Written by GlobalController to every node's Redis. This value changes when +# the controller reloads, so session operations must look it up live instead +# of relying solely on the environment captured when this process started. +IDENTITY_KEY = "controller:identity" + # session.status is a Postgres enum whose value set ("running"/"failed"/"completed") # is owned by the database schema, not by us -- it cannot be renamed to match the # /status API's vocabulary ("running"/"error"/"done"). Translate between them here. @@ -57,6 +62,18 @@ def my_workflow(query: str): } +def _current_identity(redis_client, env_db_url, env_project_id): + """Looks at the controller:identity field again if either project_id/db_url got updated + + Meant for when a new DB wants to be used on the same global controller (staging -> prod), + Or if want to tag this workflow with a different project_id + """ + identity = redis_client.hgetall(IDENTITY_KEY) or {} + db_url = identity.get("database_url") or env_db_url + project_id = identity.get("project_id") or env_project_id + return db_url, project_id + + def deploy(workflow_fn, port=8080, host="0.0.0.0", redis_host=None, redis_port=None): """ Deploy a workflow function as a REST API endpoint. @@ -76,8 +93,10 @@ def deploy(workflow_fn, port=8080, host="0.0.0.0", redis_host=None, redis_port=N redis_port = redis_port or int(os.environ.get("VENTIS_REDIS_PORT", 6379)) redis_client = RedisClient(host=redis_host, port=redis_port) - db_url = os.environ.get("VENTIS_DATABASE_URL") - project_id = os.environ.get("VENTIS_PROJECT_ID") + # These are fallbacks only. _current_identity() reads the controller's + # current Redis value for every session transition and status fallback. + env_db_url = os.environ.get("VENTIS_DATABASE_URL") + env_project_id = os.environ.get("VENTIS_PROJECT_ID") fn_name = workflow_fn.__name__ app = Flask(f"ventis-{fn_name}") @@ -92,6 +111,9 @@ def _expire_request_keys(request_id): def _record_session(request_id, status, input_payload=None, output_payload=None): """Best-effort session upsert -- logs and swallows failures so a Postgres hiccup never takes down the request itself.""" + db_url, project_id = _current_identity( + redis_client, env_db_url, env_project_id + ) if not (db_url and project_id): return try: @@ -195,6 +217,9 @@ def _status_from_session(request_id): Used once a finished request's Redis keys have expired. Returns None when there is nothing to serve, so the caller can fall through to its 404. """ + db_url, project_id = _current_identity( + redis_client, env_db_url, env_project_id + ) if not (db_url and project_id): return None