diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..a5c47a7 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,21 @@ +.git +.omx +.pytest_cache +.venv +.env +.env.* +!.env.example +__pycache__ +*.pyc +*.pyo +*.pyd +*.swp +*.swo +*~ +._* +Thumbs.db +.DS_Store +AWSCLIV2.pkg +docker_container +grpc_stubs +stubs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..54b9e9a --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,42 @@ +name: CI + +on: + pull_request: + branches: [ main ] + workflow_dispatch: + + +permissions: + contents: read + + +jobs: + test: + runs-on: ubuntu-latest + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up uv + uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + + - name: "Run Ruff: Lint" + run: uvx ruff check . + + - name: "Run Ruff: Format Check" + run: uvx ruff format --check . + + - name: "Install dependencies for Ty" + run: | + uv venv + uv pip install -e . + + - name: "Run Ty: Type Check" + run: uvx ty check + +# If you want to run this locally, install act and run "act pull_request" +# For fixing Ruff lint errors: uvx ruff check --fix . +# For fixing Ruff formatting errors: uvx ruff format . \ No newline at end of file diff --git a/.gitignore b/.gitignore index ee8f58f..86b3c9b 100644 --- a/.gitignore +++ b/.gitignore @@ -24,6 +24,7 @@ env/ # OS files .DS_Store Thumbs.db +._* # Generated stubs stubs/ @@ -32,3 +33,12 @@ docker_container/ # Logs *.log + +# Local env / machine artifacts +.env +.env.* +!.env.example +AWSCLIV2.pkg +.python-version + +sandbox/ \ No newline at end of file diff --git a/README.md b/README.md index f643ee5..330e339 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ git clone https://github.com/your-repo/ventis.git cd ventis pip install -e . ``` -Note: Installation of ventis only needs to be done on the machine where you are running the deploy command. It does not need to be installed on the remote hosts where the agents are deployed. Ventis will automatically push code and environment to the remote hosts. +Note: Installation of ventis only needs to be done on the machine where you are running the deploy command. It does not need to be installed on the remote hosts where the agents are deployed. Ventis runs the built container images on the target hosts; for remote EC2 deployments, make sure the image is already available on the host. ### 2. Prerequisites @@ -68,7 +68,7 @@ cp -r ../examples/* ./ ## Deployment Guide #### Step 1: Configure the Global Controller -Edit `examples/config/global_controller.yaml` to list the agents you want to deploy, their hosts, ports, and resource limits. +Edit `config/global_controller.yaml` in your project directory to list the agents you want to deploy, their `provider`, `replicas`, and resource limits. #### Step 2: Build the project ```bash @@ -89,10 +89,10 @@ Upon running the deploy command, ventis automatically generates a REST API endpo Users can send requests to this endpoint to trigger the workflow. For this example, workflow to send a request - ```bash -curl -X POST http://localhost:8080/finance_workflow/run \ +curl -X POST http://localhost:8080/main \ -H "Content-Type: application/json" \ -d '{ - "query": "What is the current stock price of Apple?" + "ticker": "AAPL" }' ``` The request is asynchronous. To get the result, you use the following URL- @@ -105,7 +105,7 @@ curl http://localhost:8080/status/ Remove all generated stub and gRPC files: ```bash -make clean +ventis clean ``` ### Harnessing the power of Ventis diff --git a/examples/agents/finance_agent.py b/examples/agents/finance_agent.py index e760ae2..46fb923 100644 --- a/examples/agents/finance_agent.py +++ b/examples/agents/finance_agent.py @@ -1,4 +1,6 @@ from vllm_agent_stub import VllmAgentStub + + # Example of a simple finance agent class FinanceAgent(object): def __init__(self): @@ -19,15 +21,16 @@ def get_company_name(self, ticker: str) -> str: def run(self, query: str) -> str: # company = self.get_company_name(query) # price = self.get_stock_price(company) - + prompt = f"The company is {query} and the stock price is . Please write a short, professional response." - + # Call the VLLM agent remotely and wait for the result # .value() blocks until the future completes via Redis response = self.vllm.generate(prompt).value() # print(response.value()) return response + if __name__ == "__main__": agent = FinanceAgent() - print(agent.run("What is the stock price of Apple?")) \ No newline at end of file + print(agent.run("What is the stock price of Apple?")) diff --git a/examples/agents/market_agent.py b/examples/agents/market_agent.py index 775937b..f1576f3 100644 --- a/examples/agents/market_agent.py +++ b/examples/agents/market_agent.py @@ -1,5 +1,6 @@ # Market Research Agent + class MarketResearchAgent(object): def __init__(self): self.tools = [self.get_market_trend, self.get_sector_analysis] @@ -19,4 +20,4 @@ def get_competitor_list(self, company: str) -> list: def run(self, query: str) -> str: trend = self.get_market_trend(query) analysis = self.get_sector_analysis(query) - return f"{analysis} Trend: {trend['trend']}" \ No newline at end of file + return f"{analysis} Trend: {trend['trend']}" diff --git a/examples/agents/vllm_agent.py b/examples/agents/vllm_agent.py index 3b2ee6b..8b91a09 100644 --- a/examples/agents/vllm_agent.py +++ b/examples/agents/vllm_agent.py @@ -1,5 +1,6 @@ # Simple VLLM Agent Example + class VllmAgent(object): def __init__(self): self.tools = [self.generate] @@ -10,11 +11,12 @@ def __init__(self): def generate(self, prompt: str) -> str: """Generates a response using an LLM model based on the given prompt.""" print(f"VllmAgent: Received prompt: '{prompt}'") - + # Simulated LLM generation synthetic_response = f"This is an LLM generated response to: '{prompt}'" return synthetic_response + if __name__ == "__main__": agent = VllmAgent() print(agent.generate("What is the stock price?")) diff --git a/examples/config/global_controller.ec2_smoke.yaml b/examples/config/global_controller.ec2_smoke.yaml new file mode 100644 index 0000000..acf632f --- /dev/null +++ b/examples/config/global_controller.ec2_smoke.yaml @@ -0,0 +1,31 @@ +# EC2 variant of the main example config: +# - exactly 2 agent instances +# - generic agent image + bind-mounted project code +# - t2.nano for lowest-cost MVP validation + +agents: + - name: MarketResearchAgent + replicas: 2 + redis_port: 6379 + instance_type: t2.nano + resources: + cpu: 1 + memory: 256 + entrypoint: agents/market_agent.py + provider: EC2 + + +poll_interval: 5 + +redis: + host: localhost + port: 6379 + db: 0 + +ec2: + region: us-east-1 + ami_id: ami-0123456789abcdef0 + subnet_id: subnet-0123456789abcdef0 + security_group_ids: + - sg-0123456789abcdef0 + ssh_user: ubuntu diff --git a/examples/config/global_controller.yaml b/examples/config/global_controller.yaml index 8e8539d..6e07e24 100644 --- a/examples/config/global_controller.yaml +++ b/examples/config/global_controller.yaml @@ -3,70 +3,63 @@ agents: - name: FinanceAgent - # Replicas can be placed on different hosts with explicit host/port. - # Alternatively, use `replicas: 3` as shorthand to launch 3 instances - # on the same host with sequential ports starting from `port`. - replicas: - - host: localhost - port: 8051 - - host: localhost - port: 8052 - redis_port: 6379 # Redis port on this node - # user: sagarwal # SSH user for remote hosts (omit for localhost) - # Stateful agents get session affinity: all calls within a single - # request_id are routed to the same instance. + replicas: 2 + redis_port: 6379 # Redis port on this node stateful: true - # Resource limits per instance resources: - cpu: 1 # Number of CPU cores - memory: 512 # Memory in MB - # Path to the agent entrypoint script + cpu: 1 + memory: 512 entrypoint: agents/finance_agent.py + provider: local - name: MarketResearchAgent - host: localhost - port: 8053 - redis_port: 6379 - # user: sagarwal replicas: 1 + redis_port: 6379 resources: cpu: 1 memory: 512 entrypoint: agents/market_agent.py + provider: local - name: VllmAgent - host: localhost - port: 8054 - redis_port: 6379 replicas: 1 + redis_port: 6379 resources: cpu: 2 memory: 2048 entrypoint: agents/vllm_agent.py + provider: local - name: Workflow - host: localhost - port: 8050 # LC gRPC port (exposed to host) + replicas: 1 type: workflow - api_port: 8080 # Flask REST API port (exposed to host) redis_port: 6379 - replicas: 1 workflow_file: workflows/example_workflow.py + provider: local -# Polling interval in seconds poll_interval: 5 -# Redis connection redis: host: localhost port: 6379 db: 0 -# Docker image registry (optional). -# If set, `ventis deploy` will push images to this registry locally and -# pull them on remote nodes before starting containers. -# If omitted, images are shipped to remote nodes via `docker save | ssh ... docker load`. +# EC2 defaults for `provider: EC2` replicas. +# Security group must allow inbound TCP 22, 50051, and 6379 from the +# global controller host's source IP. +# +# ec2: +# region: us-east-1 +# ami_id: ami-0123456789abcdef0 +# subnet_id: subnet-0123456789abcdef0 +# security_group_ids: +# - sg-0123456789abcdef0 +# ssh_user: ubuntu + + +# Docker image registry (legacy; no longer used). +# Remote nodes are expected to already have the image before `ventis deploy`. # # registry: -# url: myregistry.example.com:5000 # Registry host:port -# user: sagarwal # (optional) SSH user used when pulling on remote nodes +# url: myregistry.example.com:5000 +# user: sagarwal diff --git a/pyproject.toml b/pyproject.toml index 16e5a56..8d046bf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,6 +4,7 @@ version = "0.1.0" description = "Distributed agent orchestration framework" requires-python = ">=3.10" dependencies = [ + "boto3", "grpcio", "grpcio-tools", "redis", @@ -26,3 +27,27 @@ ventis = [ "templates/**/*", "controller/proto/*.proto", ] + + + +[tool.ty.environment] +python = ".venv" + +[tool.ty.src] +include = ["ventis"] +exclude = [ + "ventis/templates/**", + "ventis/stub_generator.py", +] + +[tool.ty.analysis] +allowed-unresolved-imports = [ + "local_controler_pb2", + "local_controler_pb2_grpc", + "local_controller_frontend", + "redis_client", + "ventis_context", + "deploy", + "*_stub", + "*_agent_stub", +] \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index de56b8e..203234d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +boto3 grpcio grpcio-tools redis diff --git a/session-manager-plugin.pkg b/session-manager-plugin.pkg new file mode 100644 index 0000000..adc3cea Binary files /dev/null and b/session-manager-plugin.pkg differ diff --git a/tests/README.md b/tests/README.md index c609a0e..09394de 100644 --- a/tests/README.md +++ b/tests/README.md @@ -4,10 +4,11 @@ This directory contains an automated end-to-end testing suite for Ventis. It is ## 1. Automated Test Runner (`run_tests.sh`) This script automates the entire testing lifecycle by interacting with the `ventis` CLI: +0. Runs a small pytest suite from this `tests/` directory. 1. Scaffolds a new temporary project using `ventis new-project`. 2. Compiles the project using `ventis build`. 3. Launches the project using `ventis deploy` in the background. -4. Waits for the GlobalController and all agent sidecars to become healthy. +4. Waits for the deployed workflow endpoint to become reachable, then gives the agents a few extra seconds to register. 5. Runs the Python integration and performance scripts. 6. **Cleanup:** Automatically terminates the deployment and cleans up the temporary directory upon success or failure. @@ -19,7 +20,7 @@ To run the complete suite: ## 2. Functional Integration Validation (`test_integration.py`) Verifies that Ventis correctly passes data and dependencies between chained agents. - Dispatches a single query to the deployed `/main` endpoint. -- Polls the `/status` endpoint until completion. +- Polls the `/status/` endpoint until completion. - Validates the output payload structure and ensures that data successfully flowed through `FinanceAgent`, `MarketResearchAgent`, and `VllmAgent`. To run manually against an already-deployed Ventis instance: diff --git a/tests/run_tests.sh b/tests/run_tests.sh index e9e1dc7..e9f8386 100755 --- a/tests/run_tests.sh +++ b/tests/run_tests.sh @@ -5,6 +5,11 @@ echo "===========================================" echo " Ventis Integration & Performance Tests" echo "===========================================" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" &> /dev/null && pwd)" + +echo ">> 0. Running small pytest suite..." +python3 -m pytest "$SCRIPT_DIR" + TEST_DIR="/tmp/ventis_test_env_$$" PROJECT_NAME="ventis_test" @@ -25,6 +30,8 @@ cd "$TEST_DIR" echo ">> 1. Generating new project..." ventis new-project $PROJECT_NAME cd $PROJECT_NAME +grep -v 'gpu:' config/global_controller.yaml > config/global_controller.yaml.tmp +mv config/global_controller.yaml.tmp config/global_controller.yaml echo ">> 2. Building agents (ventis build)..." ventis build @@ -44,16 +51,14 @@ sleep 5 echo ">> Deployment healthy! Running test suite." ORIG_CWD=$(pwd) -# Assuming the script was called from inside the ventis repo root -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" &> /dev/null && pwd)" echo "-------------------------------------------" echo ">> Running Integration Tests..." -python "$SCRIPT_DIR/test_integration.py" || exit 1 +python3 "$SCRIPT_DIR/test_integration.py" || exit 1 echo "-------------------------------------------" echo ">> Running Performance/Load Tests..." -python "$SCRIPT_DIR/test_performance.py" --concurrent 5 --total 20 || exit 1 +python3 "$SCRIPT_DIR/test_performance.py" --concurrent 5 --total 20 || exit 1 echo "===========================================" echo " All Tests Passed Successfully!" diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..8bec94b --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,177 @@ +import os +import sys +import tempfile +import types +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import yaml + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from ventis import cli + + +class CliDeployTests(unittest.TestCase): + def _fake_controller_module(self, controller): + module = types.ModuleType("ventis.controller.global_controller") + module.GlobalController = lambda _config_path: controller + return module + + @patch("atexit.register") + @patch("signal.signal") + @patch("ventis.cli._ensure_grpc_stubs_importable") + @patch("ventis.cli._preflight_ec2_deploy") + def test_deploy_skips_ec2_preflight_for_local_config( + self, + preflight, + ensure_grpc, + _signal_patch, + _atexit_patch, + ): + controller = MagicMock() + controller_module = self._fake_controller_module(controller) + args = SimpleNamespace(config="config/global_controller.yaml") + config = {"agents": [{"name": "LocalAgent", "provider": "local"}]} + + with ( + patch("ventis.cli.os.path.isfile", return_value=True), + patch("ventis.cli._load_config", return_value=config), + patch.dict( + sys.modules, {"ventis.controller.global_controller": controller_module} + ), + ): + cli.cmd_deploy(args) + + preflight.assert_not_called() + ensure_grpc.assert_called_once_with(os.getcwd()) + controller.launch_docker_agents.assert_called_once_with() + controller._wait_for_healthy.assert_called_once_with() + controller.run.assert_called_once_with() + + @patch("atexit.register") + @patch("signal.signal") + @patch("ventis.cli._ensure_grpc_stubs_importable") + @patch("ventis.cli._preflight_ec2_deploy") + def test_deploy_runs_ec2_preflight_for_ec2_config( + self, + preflight, + ensure_grpc, + _signal_patch, + _atexit_patch, + ): + controller = MagicMock() + controller_module = self._fake_controller_module(controller) + args = SimpleNamespace(config="config/global_controller.yaml") + config = {"agents": [{"name": "Ec2Agent", "provider": "EC2"}]} + + with ( + patch("ventis.cli.os.path.isfile", return_value=True), + patch("ventis.cli._load_config", return_value=config), + patch.dict( + sys.modules, {"ventis.controller.global_controller": controller_module} + ), + ): + cli.cmd_deploy(args) + + ensure_grpc.assert_called_once_with(os.getcwd()) + preflight.assert_called_once_with(config, os.getcwd()) + + @patch("ventis.cli._ensure_grpc_stubs_importable") + @patch("ventis.cli._require_docker_for_ec2") + def test_preflight_does_not_require_ssh_fields(self, require_docker, ensure_grpc): + config = { + "ec2": { + "ami_id": "ami-123", + "subnet_id": "subnet-123", + "security_group_ids": ["sg-123"], + "region": "us-east-1", + } + } + + cli._preflight_ec2_deploy(config, os.getcwd()) + + require_docker.assert_called_once_with("deploy") + ensure_grpc.assert_called_once_with(os.getcwd()) + + +class CliBuildTests(unittest.TestCase): + def test_build_does_not_build_global_controller_image(self): + with tempfile.TemporaryDirectory() as tmpdir: + project_dir = Path(tmpdir) + (project_dir / "config").mkdir() + (project_dir / "agents").mkdir() + (project_dir / "workflows").mkdir() + (project_dir / "docker").mkdir() + (project_dir / "docker" / "global-controller.Dockerfile").write_text( + "FROM scratch\n" + ) + (project_dir / "agents" / "example_agent.py").write_text("print('ok')\n") + (project_dir / "workflows" / "example_workflow.py").write_text( + "print('ok')\n" + ) + agent_yaml = project_dir / "agents" / "example_agent.yaml" + agent_yaml.write_text("agent:\n name: ExampleAgent\n") + config_path = project_dir / "config" / "global_controller.yaml" + config_path.write_text( + yaml.safe_dump( + { + "agents": [ + { + "name": "ExampleAgent", + "entrypoint": "agents/example_agent.py", + "provider": "local", + }, + { + "name": "Workflow", + "type": "workflow", + "workflow_file": "workflows/example_workflow.py", + "provider": "local", + }, + ] + } + ) + ) + + args = SimpleNamespace(config=str(config_path)) + docker_calls = [] + + def fake_run(cmd, check): + docker_calls.append(cmd) + return SimpleNamespace(returncode=0) + + with ( + patch( + "ventis.cli._get_package_dir", + return_value=str(project_dir / "package"), + ), + patch( + "ventis.cli.glob.glob", + side_effect=[[str(agent_yaml)], ["proto/a.proto"]], + ), + patch("ventis.stub_generator.generate_stub"), + patch("ventis.stub_generator.generate_docker"), + patch("ventis.stub_generator.generate_workflow_docker"), + patch("ventis.cli.subprocess.run", side_effect=fake_run), + patch.dict(os.environ, {}, clear=False), + ): + cwd = os.getcwd() + os.chdir(project_dir) + try: + cli.cmd_build(args) + finally: + os.chdir(cwd) + + flattened = [" ".join(call) for call in docker_calls] + self.assertFalse( + any("global-controller.Dockerfile" in call for call in flattened) + ) + self.assertEqual( + sum(call[:2] == ["docker", "build"] for call in docker_calls), 2 + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_instance_manager_runtime.py b/tests/test_instance_manager_runtime.py new file mode 100644 index 0000000..8862f90 --- /dev/null +++ b/tests/test_instance_manager_runtime.py @@ -0,0 +1,352 @@ +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.cloud_provider_logic.Local import _runtime as local_runtime +from ventis.controller.instance_manager import InstanceManager + + +class _FakeRedis: + def __init__(self): + self.strings = {} + self.hashes = {} + self.sets = {} + + def set(self, key, value): + self.strings[key] = value + + def get(self, key): + return self.strings.get(key) + + def delete(self, *keys): + for key in keys: + self.strings.pop(key, None) + self.hashes.pop(key, None) + self.sets.pop(key, None) + + 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 hgetall(self, name): + return dict(self.hashes.get(name, {})) + + def hdel(self, name, field): + self.hashes.setdefault(name, {}).pop(field, None) + + def sadd(self, name, *values): + self.sets.setdefault(name, set()).update(values) + + def srem(self, name, *values): + self.sets.setdefault(name, set()).difference_update(values) + + def smembers(self, name): + return set(self.sets.get(name, set())) + + def scan_keys(self, pattern): + prefix = pattern.rstrip("*") + keys = set(self.strings) | set(self.hashes) | set(self.sets) + return [key for key in sorted(keys) if key.startswith(prefix)] + + +def _fake_controller(): + redis = _FakeRedis() + return SimpleNamespace( + redis=redis, + containers={}, + node_redis={}, + redis_containers={}, + _run_cmd=MagicMock(return_value=SimpleNamespace(returncode=0)), + ) + + +def _fake_runtime(**kwargs): + runtime = SimpleNamespace( + validate_config=MagicMock(), + provision_instance=MagicMock(return_value={}), + bootstrap_instance=MagicMock(return_value={}), + terminate_instance=MagicMock(), + routing_endpoint_for=MagicMock( + side_effect=lambda instance: instance["endpoint"] + ), + _controller=None, + ) + for key, value in kwargs.items(): + setattr(runtime, key, value) + return runtime + + +class InstanceManagerRuntimeTests(unittest.TestCase): + def test_local_instances_keep_default_host_and_increment_host_ports(self): + controller = _fake_controller() + manager = InstanceManager(controller, controller.redis) + + alpha = manager.ensure_instances([{"name": "Alpha", "provider": "local"}])[0] + beta = manager.ensure_instances([{"name": "Beta", "provider": "local"}])[0] + + self.assertEqual( + alpha, + { + "agent_name": "Alpha", + "provider": "local", + "replica_index": "0", + "host": "localhost", + "host_port": "8000", + "container_port": "50051", + "endpoint": "localhost:8000", + "redis_host": "host.docker.internal", + "redis_port": "6379", + "runtime_id": "ventis-local-alpha-0", + }, + ) + self.assertEqual(beta["host"], "localhost") + self.assertEqual(beta["host_port"], "8001") + self.assertEqual( + controller._run_cmd.call_args_list[0].args, + ( + [ + "docker", + "run", + "-d", + "-it", + "--add-host=host.docker.internal:host-gateway", + "--name", + "ventis-local-alpha-0", + "-p", + "8000:50051", + "-e", + "VENTIS_AGENT_PORT=8000", + "-e", + "VENTIS_AGENT_HOST=host.docker.internal", + "-e", + "VENTIS_REDIS_HOST=host.docker.internal", + "-e", + "VENTIS_REDIS_PORT=6379", + "ventis-alpha", + ], + "localhost", + None, + ), + ) + + def test_local_workflow_and_resource_flags_stay_the_same(self): + controller = _fake_controller() + manager = InstanceManager(controller, controller.redis) + + manager.ensure_instances( + [ + { + "name": "Workflow", + "provider": "local", + "type": "workflow", + "resources": {"cpu": 2, "memory": 1024, "gpu": 1}, + } + ] + ) + + self.assertEqual( + controller._run_cmd.call_args.args, + ( + [ + "docker", + "run", + "-d", + "-it", + "--add-host=host.docker.internal:host-gateway", + "--name", + "ventis-local-workflow-0", + "-p", + "8000:50051", + "-e", + "VENTIS_AGENT_PORT=8000", + "-e", + "VENTIS_AGENT_HOST=host.docker.internal", + "-e", + "VENTIS_REDIS_HOST=host.docker.internal", + "-e", + "VENTIS_REDIS_PORT=6379", + "-p", + "8080:8080", + "--cpus", + "2", + "--memory", + "1024m", + "--gpus", + "1", + "ventis-workflow", + ], + "localhost", + None, + ), + ) + + def test_local_remove_instance_still_removes_the_same_container(self): + controller = _fake_controller() + manager = InstanceManager(controller, controller.redis) + manager.ensure_instances([{"name": "Alpha", "provider": "local"}]) + + controller._run_cmd.reset_mock() + manager.remove_instance("local:Alpha:0") + + self.assertEqual( + controller._run_cmd.call_args.args, + (["docker", "rm", "-f", "ventis-local-alpha-0"], "localhost", None), + ) + self.assertEqual(controller.redis.hgetall("agent_instance:local:Alpha:0"), {}) + self.assertEqual(controller.containers["Alpha"], []) + + def test_manager_keeps_ec2_runtime_boundary_behavior(self): + controller = _fake_controller() + manager = InstanceManager(controller, controller.redis) + + provisioned = { + "host": "10.0.0.30", + "runtime_id": "ventis-ec2-remote-0--i-test1", + "ec2_instance_id": "i-test1", + "redis_port": 6390, + } + instance = { + "agent_name": "Remote", + "provider": "EC2", + "replica_index": "0", + "host": "10.0.0.30", + "host_port": "50051", + "container_port": "50051", + "endpoint": "10.0.0.30:50051", + "redis_host": "10.0.0.30", + "redis_port": "6390", + "runtime_id": "ventis-ec2-remote-0--i-test1", + "ec2_instance_id": "i-test1", + } + + runtime = _fake_runtime( + provision_instance=MagicMock(return_value=provisioned), + bootstrap_instance=MagicMock(return_value=instance), + ) + + with patch.object(manager, "_provider_runtime", return_value=runtime): + created = manager.ensure_instances( + [ + { + "name": "Remote", + "provider": "EC2", + "instance_type": "t3.small", + "redis_port": 6390, + } + ] + )[0] + controller.node_redis = {"10.0.0.30": _FakeRedis()} + controller.redis_containers = {"10.0.0.30": "redis-box"} + manager.remove_instance("EC2:Remote:0") + + runtime.validate_config.assert_called_once_with() + runtime.provision_instance.assert_called_once_with( + { + "name": "Remote", + "provider": "EC2", + "instance_type": "t3.small", + "redis_port": 6390, + }, + 0, + manager._next_host_port, + ) + runtime.bootstrap_instance.assert_called_once_with( + provisioned, + { + "name": "Remote", + "provider": "EC2", + "instance_type": "t3.small", + "redis_port": 6390, + }, + 0, + ) + runtime.terminate_instance.assert_called_once_with(instance) + self.assertEqual(created, instance) + + def test_manager_uses_same_runtime_contract_for_local_and_ec2(self): + controller = _fake_controller() + manager = InstanceManager(controller, controller.redis) + + local_instance = { + "agent_name": "Local", + "provider": "local", + "replica_index": "0", + "host": "localhost", + "host_port": "8000", + "container_port": "50051", + "endpoint": "localhost:8000", + "redis_host": "host.docker.internal", + "redis_port": "6379", + "runtime_id": "ventis-local-local-0", + } + ec2_instance = { + "agent_name": "Remote", + "provider": "EC2", + "replica_index": "0", + "host": "10.0.0.30", + "host_port": "50051", + "container_port": "50051", + "endpoint": "10.0.0.30:50051", + "redis_host": "10.0.0.30", + "redis_port": "6379", + "runtime_id": "ventis-ec2-remote-0--i-test1", + "ec2_instance_id": "i-test1", + } + + local_runtime = _fake_runtime( + bootstrap_instance=MagicMock(return_value=local_instance), + routing_endpoint_for=MagicMock(return_value="host.docker.internal:8000"), + ) + ec2_runtime = _fake_runtime( + bootstrap_instance=MagicMock(return_value=ec2_instance), + routing_endpoint_for=MagicMock(return_value="10.0.0.30:50051"), + ) + + def runtime_for(provider): + return ec2_runtime if provider == "EC2" else local_runtime + + with patch.object(manager, "_provider_runtime", side_effect=runtime_for): + manager.ensure_instances( + [ + {"name": "Local", "provider": "local"}, + {"name": "Remote", "provider": "EC2", "instance_type": "t3.small"}, + ] + ) + + local_runtime.validate_config.assert_called_once_with() + local_runtime.provision_instance.assert_called_once_with( + {"name": "Local", "provider": "local"}, 0, manager._next_host_port + ) + local_runtime.bootstrap_instance.assert_called_once_with( + {}, {"name": "Local", "provider": "local"}, 0 + ) + ec2_runtime.validate_config.assert_called_once_with() + ec2_runtime.provision_instance.assert_called_once_with( + {"name": "Remote", "provider": "EC2", "instance_type": "t3.small"}, + 0, + manager._next_host_port, + ) + ec2_runtime.bootstrap_instance.assert_called_once_with( + {}, {"name": "Remote", "provider": "EC2", "instance_type": "t3.small"}, 0 + ) + + def test_local_provider_runtime_does_not_require_ec2_import(self): + controller = _fake_controller() + manager = InstanceManager(controller, controller.redis) + + runtime = manager._provider_runtime("local") + + self.assertIs(runtime, local_runtime) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_integration.py b/tests/test_integration.py index fac98fd..a2e3747 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -1,51 +1,127 @@ import requests -import time import sys +import time +import unittest +from types import SimpleNamespace +from unittest.mock import patch + def run_integration_test(): base_url = "http://localhost:8080" print(f"Submitting query to {base_url}/main...") - + response = requests.post(f"{base_url}/main", json={"ticker": "MSFT"}) - + if response.status_code != 202: print(f"Error submitting request: HTTP {response.status_code}") print(response.text) sys.exit(1) - + data = response.json() req_id = data.get("request_id") print(f"Got Request ID: {req_id}") - + max_wait = 30 elapsed = 0 - + while elapsed < max_wait: status_res = requests.get(f"{base_url}/status/{req_id}").json() status = status_res.get("status") - + if status == "done": result = status_res.get("result", {}) print(f"\nWorkflow Completed! Result: {result}") - + # Validation assertions - assert "MSFT" in result.get("company_name", ""), "Missing expected company name." - assert "This is an LLM generated response to" in result.get("competitors", ""), "VllmAgent response formatting missing." - assert result.get("stock_price") == 100.0, "FinanceAgent did not return 100.0" - + assert "MSFT" in result.get("company_name", ""), ( + "Missing expected company name." + ) + assert "This is an LLM generated response to" in result.get( + "competitors", "" + ), "VllmAgent response formatting missing." + assert result.get("stock_price") == 100.0, ( + "FinanceAgent did not return 100.0" + ) + print("\nIntegration test passed. All validations successful.") sys.exit(0) - + if status == "error": print(f"Workflow hit an error: {status_res.get('error')}") sys.exit(1) - + print(f"Status: {status} ... waiting") time.sleep(1) elapsed += 1 - + print(f"Timed out after {max_wait}s waiting for workflow completion.") sys.exit(1) + if __name__ == "__main__": run_integration_test() + + +class IntegrationScriptTests(unittest.TestCase): + def test_run_integration_test_exits_zero_on_expected_result(self): + submit_response = SimpleNamespace( + status_code=202, + json=lambda: {"request_id": "req-1"}, + ) + status_response = SimpleNamespace( + json=lambda: { + "status": "done", + "result": { + "company_name": "MSFT Corp", + "competitors": "This is an LLM generated response to competitors", + "stock_price": 100.0, + }, + } + ) + + with patch(f"{__name__}.requests.post", return_value=submit_response): + with patch(f"{__name__}.requests.get", return_value=status_response): + with self.assertRaises(SystemExit) as raised: + run_integration_test() + + self.assertEqual(raised.exception.code, 0) + + def test_run_integration_test_exits_one_on_submit_failure(self): + submit_response = SimpleNamespace(status_code=500, text="boom") + + with patch(f"{__name__}.requests.post", return_value=submit_response): + with self.assertRaises(SystemExit) as raised: + run_integration_test() + + self.assertEqual(raised.exception.code, 1) + + def test_run_integration_test_exits_one_on_workflow_error(self): + submit_response = SimpleNamespace( + status_code=202, + json=lambda: {"request_id": "req-1"}, + ) + status_response = SimpleNamespace( + json=lambda: {"status": "error", "error": "boom"} + ) + + with patch(f"{__name__}.requests.post", return_value=submit_response): + with patch(f"{__name__}.requests.get", return_value=status_response): + with self.assertRaises(SystemExit) as raised: + run_integration_test() + + self.assertEqual(raised.exception.code, 1) + + def test_run_integration_test_times_out(self): + submit_response = SimpleNamespace( + status_code=202, + json=lambda: {"request_id": "req-1"}, + ) + status_response = SimpleNamespace(json=lambda: {"status": "running"}) + + with patch(f"{__name__}.requests.post", return_value=submit_response): + with patch(f"{__name__}.requests.get", return_value=status_response): + with patch(f"{__name__}.time.sleep", return_value=None): + with self.assertRaises(SystemExit) as raised: + run_integration_test() + + self.assertEqual(raised.exception.code, 1) diff --git a/tests/test_performance.py b/tests/test_performance.py index d1f876a..9d998d6 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -4,6 +4,10 @@ import sys import requests import statistics +import unittest +from types import SimpleNamespace +from unittest.mock import patch + def dispatch_request(session, base_url, payload): start_time = time.time() @@ -11,10 +15,15 @@ def dispatch_request(session, base_url, payload): response = session.post(f"{base_url}/main", json=payload, timeout=5) response.raise_for_status() req_id = response.json().get("request_id") - return {"status": "dispatched", "req_id": req_id, "latency": time.time() - start_time} + return { + "status": "dispatched", + "req_id": req_id, + "latency": time.time() - start_time, + } except Exception as e: return {"status": "error", "error": str(e), "latency": time.time() - start_time} + def poll_request(session, base_url, req_id): start_time = time.time() while True: @@ -23,56 +32,79 @@ def poll_request(session, base_url, req_id): if res.get("status") in ["done", "error"]: return {"status": res["status"], "latency": time.time() - start_time} except Exception as e: - return {"status": "error", "error": str(e), "latency": time.time() - start_time} + return { + "status": "error", + "error": str(e), + "latency": time.time() - start_time, + } time.sleep(0.5) + def run_performance_test(concurrent_users, total_requests): base_url = "http://localhost:8080" - print(f"Starting Performance Test: {total_requests} total requests across {concurrent_users} workers.") - + print( + f"Starting Performance Test: {total_requests} total requests across {concurrent_users} workers." + ) + session = requests.Session() - + dispatch_results = [] - + # Phase 1: Dispatch barrage print("\n[Phase 1] Queuing requests...") global_start = time.time() - with concurrent.futures.ThreadPoolExecutor(max_workers=concurrent_users) as executor: - futures = [executor.submit(dispatch_request, session, base_url, {"ticker": f"TICK{i}"}) for i in range(total_requests)] + with concurrent.futures.ThreadPoolExecutor( + max_workers=concurrent_users + ) as executor: + futures = [ + executor.submit(dispatch_request, session, base_url, {"ticker": f"TICK{i}"}) + for i in range(total_requests) + ] for future in concurrent.futures.as_completed(futures): dispatch_results.append(future.result()) - + dispatch_success = [r for r in dispatch_results if r["status"] == "dispatched"] print(f"Successfully queued {len(dispatch_success)} / {total_requests} requests.") - + # Phase 2: Poll till completion print("\n[Phase 2] Waiting for resolution...") poll_results = [] - with concurrent.futures.ThreadPoolExecutor(max_workers=concurrent_users) as executor: - futures = [executor.submit(poll_request, session, base_url, r["req_id"]) for r in dispatch_success] + with concurrent.futures.ThreadPoolExecutor( + max_workers=concurrent_users + ) as executor: + futures = [ + executor.submit(poll_request, session, base_url, r["req_id"]) + for r in dispatch_success + ] for future in concurrent.futures.as_completed(futures): poll_results.append(future.result()) - + global_time = time.time() - global_start - + # Analysis completed = [r for r in poll_results if r["status"] == "done"] failed = [r for r in poll_results if r["status"] == "error"] - + if completed: latencies = [r["latency"] for r in completed] avg_lat = statistics.mean(latencies) - p95_lat = statistics.quantiles(latencies, n=20)[18] if len(latencies) >= 20 else max(latencies) - + p95_lat = ( + statistics.quantiles(latencies, n=20)[18] + if len(latencies) >= 20 + else max(latencies) + ) + print("\n--- Analytics Report ---") print(f"Total Time Elapsed: {global_time:.2f} s") print(f"Completed Successfully: {len(completed)}") - print(f"Failed / Errors: {len(failed) + (total_requests - len(dispatch_success))}") + print( + f"Failed / Errors: {len(failed) + (total_requests - len(dispatch_success))}" + ) print(f"Throughput (RPS): {len(completed) / global_time:.2f} req/s") print(f"Avg End-to-End Latency: {avg_lat:.2f} s") print(f"P95 End-to-End Latency: {p95_lat:.2f} s") print("------------------------\n") - + if len(failed) > 0: sys.exit(1) sys.exit(0) @@ -80,10 +112,85 @@ def run_performance_test(concurrent_users, total_requests): print("\nAll requests failed. Performance cannot be calculated.") sys.exit(1) + if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--concurrent", type=int, default=10, help="Number of parallel workers") - parser.add_argument("--total", type=int, default=50, help="Total number of requests to generate") + parser.add_argument( + "--concurrent", type=int, default=10, help="Number of parallel workers" + ) + parser.add_argument( + "--total", type=int, default=50, help="Total number of requests to generate" + ) args = parser.parse_args() - + run_performance_test(args.concurrent, args.total) + + +class PerformanceScriptTests(unittest.TestCase): + def test_dispatch_request_returns_request_id_on_success(self): + response = SimpleNamespace( + raise_for_status=lambda: None, + json=lambda: {"request_id": "req-1"}, + ) + session = SimpleNamespace(post=lambda *args, **kwargs: response) + + result = dispatch_request(session, "http://test", {"ticker": "MSFT"}) + + self.assertEqual(result["status"], "dispatched") + self.assertEqual(result["req_id"], "req-1") + + def test_dispatch_request_returns_error_on_exception(self): + session = SimpleNamespace( + post=lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("boom")) + ) + + result = dispatch_request(session, "http://test", {"ticker": "MSFT"}) + + self.assertEqual(result["status"], "error") + self.assertIn("boom", result["error"]) + + def test_poll_request_returns_done_status(self): + response = SimpleNamespace(json=lambda: {"status": "done"}) + session = SimpleNamespace(get=lambda *args, **kwargs: response) + + result = poll_request(session, "http://test", "req-1") + + self.assertEqual(result["status"], "done") + + def test_poll_request_returns_error_on_exception(self): + session = SimpleNamespace( + get=lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("boom")) + ) + + result = poll_request(session, "http://test", "req-1") + + self.assertEqual(result["status"], "error") + self.assertIn("boom", result["error"]) + + def test_run_performance_test_exits_zero_when_all_requests_complete(self): + class Session: + def post(self, *args, **kwargs): + return SimpleNamespace( + raise_for_status=lambda: None, + json=lambda: {"request_id": "req-1"}, + ) + + def get(self, *args, **kwargs): + return SimpleNamespace(json=lambda: {"status": "done"}) + + with patch(f"{__name__}.requests.Session", return_value=Session()): + with self.assertRaises(SystemExit) as raised: + run_performance_test(concurrent_users=1, total_requests=1) + + self.assertEqual(raised.exception.code, 0) + + def test_run_performance_test_exits_one_when_dispatch_fails(self): + class Session: + def post(self, *args, **kwargs): + raise RuntimeError("boom") + + with patch(f"{__name__}.requests.Session", return_value=Session()): + with self.assertRaises(SystemExit) as raised: + run_performance_test(concurrent_users=1, total_requests=1) + + self.assertEqual(raised.exception.code, 1) diff --git a/tests/test_redis_utils.py b/tests/test_redis_utils.py new file mode 100644 index 0000000..789b04d --- /dev/null +++ b/tests/test_redis_utils.py @@ -0,0 +1,30 @@ +import unittest +import os +import sys +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from ventis.controller.utils.redis_utils import _wait_for_redis + + +class WaitForRedisTests(unittest.TestCase): + @patch("ventis.controller.utils.redis_utils.time.sleep") + @patch("ventis.controller.utils.redis_utils.time.time", return_value=0) + def test_timeout_message_stays_the_same(self, mock_time, mock_sleep): + redis_client = MagicMock() + + with self.assertRaisesRegex( + TimeoutError, + r"Timed out connecting to Redis at 10\.0\.0\.1:6379\. " + r"For EC2 runtimes, ensure the instance security group allows inbound " + r"TCP 6379 from the global controller host\.", + ): + _wait_for_redis(redis_client, "10.0.0.1", 6379, timeout=0, interval=1) + + redis_client.set.assert_not_called() + mock_sleep.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_runtime_ec2.py b/tests/test_runtime_ec2.py new file mode 100644 index 0000000..e08e390 --- /dev/null +++ b/tests/test_runtime_ec2.py @@ -0,0 +1,287 @@ +import base64 +import os +import sys +import tempfile +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.cloud_provider_logic.EC2 import _runtime as ec2_runtime + + +class _FakeWaiter: + def __init__(self): + self.calls = [] + + def wait(self, InstanceIds): + self.calls.append(list(InstanceIds)) + + +class _FakeEC2Client: + def __init__(self, public_ip="54.10.20.30", private_ip="10.0.0.30"): + self.public_ip = public_ip + self.private_ip = private_ip + self.instances = {} + self.run_requests = [] + self.terminate_requests = [] + self.waiter = _FakeWaiter() + + def run_instances(self, **kwargs): + self.run_requests.append(kwargs) + instance_id = f"i-test{len(self.run_requests)}" + self.instances[instance_id] = { + "InstanceId": instance_id, + "State": {"Name": "running"}, + "PrivateIpAddress": self.private_ip, + "PublicIpAddress": self.public_ip, + } + return {"Instances": [{"InstanceId": instance_id}]} + + def get_waiter(self, name): + assert name == "instance_running" + return self.waiter + + def describe_instances(self, InstanceIds): + return { + "Reservations": [ + {"Instances": [self.instances[instance_id]]} + for instance_id in InstanceIds + if instance_id in self.instances + ] + } + + def terminate_instances(self, InstanceIds): + self.terminate_requests.append(list(InstanceIds)) + return {} + + +class _FakeSession: + def __init__(self, ec2_client, region_name="us-east-1", credentials=True): + self._ec2_client = ec2_client + self.region_name = region_name + self._credentials = object() if credentials else None + self.client_calls = [] + + def get_credentials(self): + return self._credentials + + def client(self, service_name, region_name=None): + self.client_calls.append( + { + "service_name": service_name, + "region_name": region_name, + } + ) + if service_name == "ec2": + return self._ec2_client + raise AssertionError(service_name) + + +class EC2RuntimeTests(unittest.TestCase): + def setUp(self): + self.original_controller = ec2_runtime._controller + self.original_key_path = ec2_runtime.DEFAULT_SSH_KEY_PATH + self.fake_client = _FakeEC2Client() + self.fake_session = _FakeSession(self.fake_client) + self.session_calls = [] + self.key_dir = tempfile.mkdtemp() + self.private_key = os.path.join(self.key_dir, "ventis_ec2") + self.public_key = self.private_key + ".pub" + with open(self.private_key, "w") as f: + f.write("PRIVATE") + with open(self.public_key, "w") as f: + f.write("ssh-ed25519 AAAA test@ventis\n") + ec2_runtime.DEFAULT_SSH_KEY_PATH = self.private_key + self.controller = SimpleNamespace( + config={ + "redis": {"host": "redis.internal", "port": 6379}, + "ec2": { + "ami_id": "ami-123456", + "subnet_id": "subnet-123456", + "security_group_ids": ["sg-123456"], + "region": "us-east-1", + "ssh_user": "ubuntu", + "profile": "ventis-profile", + "aws_access_key_id": "AKIA_TEST", + "aws_secret_access_key": "secret", + "aws_session_token": "token", + "public_ip_timeout": 1, + }, + }, + registry_url=None, + _run_cmd=MagicMock( + return_value=SimpleNamespace(returncode=0, stderr="", stdout="") + ), + ) + ec2_runtime._controller = self.controller + self.session_patch = patch.object( + ec2_runtime.boto3, + "Session", + side_effect=self._make_session, + ) + self.session_patch.start() + + def tearDown(self): + self.session_patch.stop() + ec2_runtime._controller = self.original_controller + ec2_runtime.DEFAULT_SSH_KEY_PATH = self.original_key_path + + def _make_session(self, **kwargs): + self.session_calls.append(kwargs) + return self.fake_session + + def test_validate_config_fails_when_required_fields_are_missing(self): + self.controller.config["ec2"].pop("ami_id") + + with self.assertRaisesRegex(ValueError, "Missing EC2 config"): + ec2_runtime.validate_config() + + def test_provision_uses_userdata_and_configured_session(self): + spec = { + "name": "Tagged", + "provider": "EC2", + "instance_type": "t3.small", + "redis_port": 6390, + } + + provisioned = ec2_runtime.provision_instance(spec, 2) + + request = self.fake_client.run_requests[0] + self.assertEqual(request["ImageId"], "ami-123456") + self.assertNotIn("KeyName", request) + self.assertIn("UserData", request) + userdata = base64.b64decode(request["UserData"]).decode() + self.assertIn("ssh-ed25519 AAAA test@ventis", userdata) + self.assertIn("/home/ubuntu/.ssh", userdata) + self.assertEqual( + request["TagSpecifications"][0]["Tags"][0], + {"Key": "Name", "Value": "ventis-Tagged-2"}, + ) + self.assertEqual(self.fake_client.waiter.calls, [["i-test1"]]) + self.assertEqual(provisioned["host"], "10.0.0.30") + self.assertEqual( + self.fake_session.client_calls, + [ + { + "service_name": "ec2", + "region_name": "us-east-1", + }, + ], + ) + self.assertEqual( + self.session_calls, + [ + { + "region_name": "us-east-1", + "profile_name": "ventis-profile", + "aws_access_key_id": "AKIA_TEST", + "aws_secret_access_key": "secret", + "aws_session_token": "token", + } + ], + ) + + def test_provision_and_bootstrap_instance_return_runtime_record(self): + spec = { + "name": "Tagged", + "provider": "EC2", + "instance_type": "t3.small", + "redis_port": 6390, + } + + with ( + patch.object(ec2_runtime, "_bootstrap_instance"), + patch.object(ec2_runtime, "_check_controller_health", return_value=True), + ): + provisioned = ec2_runtime.provision_instance(spec, 2) + instance = ec2_runtime.bootstrap_instance(provisioned, spec, 2) + + self.assertEqual(instance["host"], "10.0.0.30") + self.assertEqual(instance["endpoint"], "10.0.0.30:50051") + self.assertEqual(instance["redis_host"], "10.0.0.30") + self.assertEqual(instance["redis_port"], "6390") + self.assertIn("--i-test1", instance["runtime_id"]) + + def test_bootstrap_instance_terminates_instance_when_bootstrap_fails(self): + spec = {"name": "Broken", "provider": "EC2", "instance_type": "t3.small"} + provisioned = ec2_runtime.provision_instance(spec, 0) + + with ( + patch.object( + ec2_runtime, "_bootstrap_instance", side_effect=RuntimeError("boom") + ), + self.assertRaisesRegex(RuntimeError, "boom"), + ): + ec2_runtime.bootstrap_instance(provisioned, spec, 0) + + self.assertEqual(self.fake_client.terminate_requests, [["i-test1"]]) + self.assertEqual(self.session_calls[-1], self.session_calls[0]) + + def test_bootstrap_uses_ssh_user(self): + spec = {"name": "Tagged", "provider": "EC2", "redis_port": 6390} + with patch.object(ec2_runtime.time, "sleep"): + ec2_runtime._bootstrap_instance( + "10.0.0.30", + spec, + 2, + self.controller.config["ec2"], + redis_host="10.0.0.30", + redis_port=6390, + ) + + self.assertTrue(self.controller._run_cmd.called) + for call in self.controller._run_cmd.call_args_list: + self.assertEqual(call.kwargs["user"], "ubuntu") + # SSH wait + single docker run + self.assertEqual(self.controller._run_cmd.call_count, 2) + self.assertEqual( + self.controller._run_cmd.call_args_list[-1].args[0][0], "docker" + ) + + def test_bootstrap_instance_terminates_instance_when_health_check_fails(self): + spec = {"name": "Broken", "provider": "EC2", "instance_type": "t3.small"} + provisioned = ec2_runtime.provision_instance(spec, 0) + + with ( + patch.object(ec2_runtime, "_bootstrap_instance"), + patch.object( + ec2_runtime, + "_check_controller_health", + side_effect=TimeoutError("boom"), + ), + self.assertRaises(TimeoutError), + ): + ec2_runtime.bootstrap_instance(provisioned, spec, 0) + + self.assertEqual(self.fake_client.terminate_requests, [["i-test1"]]) + self.assertEqual(self.session_calls[-1], self.session_calls[0]) + + def test_terminate_instance_still_cleans_host_side_maps(self): + spec = {"name": "Tagged", "provider": "EC2", "instance_type": "t3.small"} + provisioned = ec2_runtime.provision_instance(spec, 0) + self.controller.redis_containers = {"10.0.0.30": "redis-box"} + self.controller.node_redis = {"10.0.0.30": object()} + + ec2_runtime.terminate_instance( + { + "runtime_id": provisioned["runtime_id"], + "host": "10.0.0.30", + } + ) + + self.assertEqual(self.fake_client.terminate_requests, [["i-test1"]]) + self.assertNotIn("10.0.0.30", self.controller.redis_containers) + self.assertNotIn("10.0.0.30", self.controller.node_redis) + + def test_health_check_error_message_mentions_ec2_runtime_endpoint(self): + with patch.object(ec2_runtime.time, "time", side_effect=[0, 999]): + with self.assertRaisesRegex( + TimeoutError, "EC2 runtime endpoint never became reachable" + ): + ec2_runtime._check_controller_health("10.0.0.30:50051", timeout=1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_stateful_affinity.py b/tests/test_stateful_affinity.py deleted file mode 100644 index c781696..0000000 --- a/tests/test_stateful_affinity.py +++ /dev/null @@ -1,286 +0,0 @@ -""" -Tests for stateful agent session affinity. - -These tests verify that: -1. The routing table stores multiple endpoints per agent (JSON list). -2. Stateful agents get affinity bindings (same instance per request_id). -3. Stateless agents get random routing (no affinity keys written). -4. Affinity keys are cleaned up when a request completes. - -Requirements: - - A running Redis instance on localhost:6379. - - The ventis package installed (pip install -e .). - -Run: - python -m pytest tests/test_stateful_affinity.py -v -""" - -import json -import os -import sys -import uuid - -import pytest - -# Ensure project paths are importable -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "utils")) -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ventis", "controller")) - -from redis_client import RedisClient - - -# ------------------------------------------------------------------ # -# Fixtures # -# ------------------------------------------------------------------ # - -REDIS_HOST = os.environ.get("VENTIS_TEST_REDIS_HOST", "localhost") -REDIS_PORT = int(os.environ.get("VENTIS_TEST_REDIS_PORT", 6379)) - -ROUTING_ENDPOINTS_KEY = "routing_table:endpoints" -ROUTING_STATEFUL_KEY = "routing_table:stateful" - - -@pytest.fixture -def redis(): - """Provide a RedisClient connected to a test Redis, flushed before each test.""" - client = RedisClient(host=REDIS_HOST, port=REDIS_PORT) - # Flush only keys in the test namespace to avoid clobbering production data - for key in client.scan_keys("routing_table:*"): - client.delete(key) - for key in client.scan_keys("affinity:*"): - client.delete(key) - for key in client.scan_keys("request:*"): - client.delete(key) - for key in client.scan_keys("future:*"): - client.delete(key) - yield client - - -def _seed_routing_table(redis, agents): - """Write a routing table to Redis. - - Args: - redis: RedisClient instance. - agents: list of dicts with keys: name, endpoints, stateful (optional). - """ - endpoints_table = {} - stateful_table = {} - for agent in agents: - endpoints_table[agent["name"]] = json.dumps(agent["endpoints"]) - if agent.get("stateful"): - stateful_table[agent["name"]] = "true" - redis.hset_multiple(ROUTING_ENDPOINTS_KEY, endpoints_table) - if stateful_table: - redis.hset_multiple(ROUTING_STATEFUL_KEY, stateful_table) - - -# ------------------------------------------------------------------ # -# Tests — Routing table format # -# ------------------------------------------------------------------ # - -class TestRoutingTable: - """Verify the routing table stores all replica endpoints.""" - - def test_single_replica_stored_as_list(self, redis): - """An agent with 1 replica should have a 1-element JSON list.""" - _seed_routing_table(redis, [ - {"name": "AgentA", "endpoints": ["host:5001"]}, - ]) - raw = redis.hget(ROUTING_ENDPOINTS_KEY, "AgentA") - endpoints = json.loads(raw) - assert endpoints == ["host:5001"] - - def test_multiple_replicas_stored_as_list(self, redis): - """An agent with 3 replicas should have a 3-element JSON list.""" - _seed_routing_table(redis, [ - {"name": "AgentA", "endpoints": ["host:5001", "host:5002", "host:5003"]}, - ]) - raw = redis.hget(ROUTING_ENDPOINTS_KEY, "AgentA") - endpoints = json.loads(raw) - assert len(endpoints) == 3 - assert "host:5001" in endpoints - assert "host:5002" in endpoints - assert "host:5003" in endpoints - - def test_stateful_flag_written(self, redis): - """Stateful agents should have their flag set in routing_table:stateful.""" - _seed_routing_table(redis, [ - {"name": "StatefulAgent", "endpoints": ["host:5001"], "stateful": True}, - {"name": "StatelessAgent", "endpoints": ["host:6001"]}, - ]) - assert redis.hget(ROUTING_STATEFUL_KEY, "StatefulAgent") == "true" - assert redis.hget(ROUTING_STATEFUL_KEY, "StatelessAgent") is None - - -# ------------------------------------------------------------------ # -# Tests — Affinity resolution # -# ------------------------------------------------------------------ # - -class TestAffinityResolution: - """Verify the _resolve_endpoint logic for stateful vs stateless agents.""" - - def _resolve_endpoint(self, redis, service, request_id): - """Pure-function equivalent of LocalController._resolve_endpoint. - - Extracted here so we can test without spinning up the full LC. - """ - import random as _random - - endpoints_json = redis.hget(ROUTING_ENDPOINTS_KEY, service) - if not endpoints_json: - return None - endpoints = json.loads(endpoints_json) - if not endpoints: - return None - - is_stateful = redis.hget(ROUTING_STATEFUL_KEY, service) == "true" - - if is_stateful and request_id: - affinity_key = f"affinity:{request_id}" - existing = redis.hget(affinity_key, service) - if existing: - return existing - chosen = _random.choice(endpoints) - redis.hset(affinity_key, service, chosen) - return chosen - else: - return _random.choice(endpoints) - - def test_stateful_affinity_is_sticky(self, redis): - """Multiple calls with the same request_id should return the same endpoint.""" - _seed_routing_table(redis, [ - {"name": "FA", "endpoints": ["h:5001", "h:5002", "h:5003"], "stateful": True}, - ]) - rid = uuid.uuid4().hex - - first = self._resolve_endpoint(redis, "FA", rid) - assert first in ["h:5001", "h:5002", "h:5003"] - - # Subsequent calls must return the same endpoint - for _ in range(10): - assert self._resolve_endpoint(redis, "FA", rid) == first - - def test_stateful_different_requests_can_differ(self, redis): - """Different request_ids can (but don't have to) map to different instances.""" - _seed_routing_table(redis, [ - {"name": "FA", "endpoints": ["h:5001", "h:5002", "h:5003"], "stateful": True}, - ]) - results = set() - for _ in range(50): - rid = uuid.uuid4().hex - results.add(self._resolve_endpoint(redis, "FA", rid)) - - # With 50 random picks across 3 endpoints, we should see at least 2 - # (statistically near-certain) - assert len(results) >= 2, ( - f"Expected at least 2 distinct endpoints across 50 requests, got {results}" - ) - - def test_stateful_binding_persisted_in_redis(self, redis): - """The affinity binding should exist as a Redis key.""" - _seed_routing_table(redis, [ - {"name": "FA", "endpoints": ["h:5001", "h:5002"], "stateful": True}, - ]) - rid = uuid.uuid4().hex - chosen = self._resolve_endpoint(redis, "FA", rid) - - affinity_key = f"affinity:{rid}" - assert redis.hget(affinity_key, "FA") == chosen - - def test_stateless_no_affinity_key(self, redis): - """Stateless agents should NOT create affinity keys in Redis.""" - _seed_routing_table(redis, [ - {"name": "SL", "endpoints": ["h:6001", "h:6002", "h:6003"]}, - ]) - rid = uuid.uuid4().hex - self._resolve_endpoint(redis, "SL", rid) - - affinity_key = f"affinity:{rid}" - assert redis.hget(affinity_key, "SL") is None - - def test_stateless_returns_valid_endpoint(self, redis): - """Stateless agents should always return one of the available endpoints.""" - _seed_routing_table(redis, [ - {"name": "SL", "endpoints": ["h:6001", "h:6002"]}, - ]) - for _ in range(20): - result = self._resolve_endpoint(redis, "SL", uuid.uuid4().hex) - assert result in ["h:6001", "h:6002"] - - def test_unknown_service_returns_none(self, redis): - """Looking up a service not in the routing table should return None.""" - _seed_routing_table(redis, [ - {"name": "AgentA", "endpoints": ["h:5001"]}, - ]) - assert self._resolve_endpoint(redis, "NonExistent", uuid.uuid4().hex) is None - - -# ------------------------------------------------------------------ # -# Tests — Cleanup # -# ------------------------------------------------------------------ # - -class TestAffinityCleanup: - """Verify that affinity keys are cleaned up when a request completes.""" - - def test_affinity_keys_cleaned_up(self, redis): - """After cleanup, affinity: keys should be gone.""" - rid = uuid.uuid4().hex - # Simulate affinity bindings for a request - redis.hset(f"affinity:{rid}", "AgentA", "h:5001") - redis.hset(f"affinity:{rid}", "AgentB", "h:6001") - - # Verify they exist - assert redis.hget(f"affinity:{rid}", "AgentA") == "h:5001" - assert redis.hget(f"affinity:{rid}", "AgentB") == "h:6001" - - # Simulate cleanup (what _cleanup_request does) - redis.delete(f"affinity:{rid}") - - # Verify they're gone - assert redis.hget(f"affinity:{rid}", "AgentA") is None - assert redis.hget(f"affinity:{rid}", "AgentB") is None - - def test_cleanup_does_not_affect_other_requests(self, redis): - """Cleaning up one request's affinity should not touch another's.""" - rid_a = uuid.uuid4().hex - rid_b = uuid.uuid4().hex - - redis.hset(f"affinity:{rid_a}", "Agent", "h:5001") - redis.hset(f"affinity:{rid_b}", "Agent", "h:5002") - - # Clean up only rid_a - redis.delete(f"affinity:{rid_a}") - - # rid_a gone, rid_b untouched - assert redis.hget(f"affinity:{rid_a}", "Agent") is None - assert redis.hget(f"affinity:{rid_b}", "Agent") == "h:5002" - - -# ------------------------------------------------------------------ # -# Tests — scan_keys utility # -# ------------------------------------------------------------------ # - -class TestScanKeys: - """Verify the scan_keys helper on RedisClient.""" - - def test_scan_keys_finds_matching(self, redis): - """scan_keys should return all keys matching a glob pattern.""" - prefix = uuid.uuid4().hex[:8] - redis.set(f"test:{prefix}:a", "1") - redis.set(f"test:{prefix}:b", "2") - redis.set(f"test:{prefix}:c", "3") - redis.set(f"test:other:d", "4") - - keys = redis.scan_keys(f"test:{prefix}:*") - assert len(keys) == 3 - assert f"test:{prefix}:a" in keys - assert f"test:{prefix}:b" in keys - assert f"test:{prefix}:c" in keys - - # Cleanup - redis.delete(f"test:{prefix}:a", f"test:{prefix}:b", f"test:{prefix}:c", f"test:other:d") - - def test_scan_keys_returns_empty_for_no_match(self, redis): - """scan_keys should return an empty list when nothing matches.""" - keys = redis.scan_keys("nonexistent_pattern_xyz:*") - assert keys == [] diff --git a/ventis/cli.py b/ventis/cli.py index ebbdd88..a5345f5 100644 --- a/ventis/cli.py +++ b/ventis/cli.py @@ -17,12 +17,21 @@ logging.basicConfig(level=logging.INFO) logger = logging.getLogger("ventis") +DEFAULT_DOCKER_PLATFORM = "linux/amd64" +DEFAULT_CONFIG_PATH = "config/global_controller.yaml" +EC2_REQUIRED_CONFIG_KEYS = ( + "ami_id", + "subnet_id", + "security_group_ids", + "region", +) # ------------------------------------------------------------------ # # Helpers # # ------------------------------------------------------------------ # + def _get_templates_dir(): """Return the absolute path to the bundled templates directory.""" return os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates") @@ -36,14 +45,79 @@ def _get_package_dir(): def _load_config(config_path): """Load a YAML config file.""" import yaml + with open(config_path, "r") as f: return yaml.safe_load(f) +def _docker_platform(): + """Return the target Docker platform for portable runtime images.""" + return os.environ.get("VENTIS_DOCKER_PLATFORM", DEFAULT_DOCKER_PLATFORM) + + +def _docker_build_cmd(*args): + """Build a Docker build command with an explicit target platform.""" + return ["docker", "build", "--platform", _docker_platform(), *args] + + +def _docker_available(): + if not shutil.which("docker"): + return False + + try: + result = subprocess.run( + ["docker", "info"], + capture_output=True, + text=True, + check=False, + ) + except OSError: + return False + + return result.returncode == 0 + + +def _require_docker_for_ec2(command_name): + if _docker_available(): + return + raise RuntimeError( + f"EC2-backed `ventis {command_name}` requires local Docker, but Docker is unavailable " + "or unreachable." + ) + + +def _ensure_grpc_stubs_importable(project_dir): + grpc_stubs_dir = os.path.join(project_dir, "grpc_stubs") + if grpc_stubs_dir not in sys.path: + sys.path.insert(0, grpc_stubs_dir) + + try: + __import__("local_controler_pb2") + __import__("local_controler_pb2_grpc") + except ImportError as exc: + raise RuntimeError( + "Deploy failed: generated grpc_stubs are missing or not importable. " + "Run `ventis build` on this host first." + ) from exc + + +def _preflight_ec2_deploy(config, project_dir): + ec2_cfg = config.get("ec2", {}) + missing = [key for key in EC2_REQUIRED_CONFIG_KEYS if not ec2_cfg.get(key)] + if missing: + raise RuntimeError( + f"EC2 deploy preflight failed: missing ec2 config keys: {', '.join(sorted(missing))}" + ) + + _require_docker_for_ec2("deploy") + _ensure_grpc_stubs_importable(project_dir) + + # ------------------------------------------------------------------ # # ventis new-project # # ------------------------------------------------------------------ # + def cmd_new_project(args): """Scaffold a new Ventis project.""" project_name = args.name @@ -76,6 +150,7 @@ def cmd_new_project(args): # ventis build # # ------------------------------------------------------------------ # + def cmd_build(args): """ Generate stubs, compile gRPC protos, generate Docker contexts, @@ -100,11 +175,11 @@ def cmd_build(args): stubs_dir = os.path.join(project_dir, "stubs") os.makedirs(stubs_dir, exist_ok=True) - # Add repo_root to sys.path so 'ventis.stub_generator' can be imported - repo_root = os.path.dirname(package_dir) - sys.path.insert(0, repo_root) - - from ventis.stub_generator import generate_stub, generate_docker, generate_workflow_docker + from ventis.stub_generator import ( + generate_stub, + generate_docker, + generate_workflow_docker, + ) yaml_files = glob.glob(os.path.join(agents_dir, "*.yaml")) if not yaml_files: @@ -129,16 +204,21 @@ def cmd_build(args): for proto_file in proto_files: logger.info("Compiling gRPC proto: %s", proto_file) - subprocess.run([ - sys.executable, "-m", "grpc_tools.protoc", - f"-I{proto_dir}", - f"--python_out={grpc_stubs_dir}", - f"--grpc_python_out={grpc_stubs_dir}", - proto_file, - ], check=True) + subprocess.run( + [ + sys.executable, + "-m", + "grpc_tools.protoc", + f"-I{proto_dir}", + f"--python_out={grpc_stubs_dir}", + f"--grpc_python_out={grpc_stubs_dir}", + proto_file, + ], + check=True, + ) # -------------------------------------------------------------- # - # Step 3: Generate Docker contexts and build images # + # Step 4: Generate Docker contexts and build images # # -------------------------------------------------------------- # for agent_cfg in agents: agent_name = agent_cfg["name"] @@ -148,7 +228,9 @@ def cmd_build(args): # Workflow container workflow_file = agent_cfg.get("workflow_file") if not workflow_file: - logger.warning("Skipping workflow '%s': no workflow_file specified", agent_name) + logger.warning( + "Skipping workflow '%s': no workflow_file specified", agent_name + ) continue workflow_path = os.path.join(project_dir, workflow_file) @@ -167,13 +249,17 @@ def cmd_build(args): image_name = f"ventis-{agent_name.lower()}" logger.info("Building Docker image: %s", image_name) - subprocess.run(["docker", "build", "-t", image_name, docker_context], check=True) + subprocess.run( + _docker_build_cmd("-t", image_name, docker_context), check=True + ) else: # Agent container entrypoint = agent_cfg.get("entrypoint") if not entrypoint: - logger.warning("Skipping agent '%s': no entrypoint specified", agent_name) + logger.warning( + "Skipping agent '%s': no entrypoint specified", agent_name + ) continue agent_file = os.path.join(project_dir, entrypoint) @@ -185,6 +271,7 @@ def cmd_build(args): matching_yaml = None for yaml_path in yaml_files: import yaml + with open(yaml_path) as f: ydata = yaml.safe_load(f) if ydata.get("agent", {}).get("name") == agent_name: @@ -192,7 +279,10 @@ def cmd_build(args): break if not matching_yaml: - logger.warning("No YAML definition found for agent '%s', skipping Docker", agent_name) + logger.warning( + "No YAML definition found for agent '%s', skipping Docker", + agent_name, + ) continue docker_context = os.path.join(project_dir, "docker_container", agent_name) @@ -207,7 +297,9 @@ def cmd_build(args): image_name = f"ventis-{agent_name.lower()}" logger.info("Building Docker image: %s", image_name) - subprocess.run(["docker", "build", "-t", image_name, docker_context], check=True) + subprocess.run( + _docker_build_cmd("-t", image_name, docker_context), check=True + ) logger.info("Build complete.") @@ -216,6 +308,7 @@ def cmd_build(args): # ventis deploy # # ------------------------------------------------------------------ # + def cmd_deploy(args): """ Launch the Global Controller, which starts Redis containers, @@ -229,12 +322,16 @@ def cmd_deploy(args): logger.error("Config file not found: %s", config_path) sys.exit(1) - # Ensure imports resolve - package_dir = _get_package_dir() - repo_root = os.path.dirname(package_dir) - sys.path.insert(0, repo_root) - # Add the project's grpc_stubs to path so global controller can find them - sys.path.insert(0, os.path.join(os.getcwd(), "grpc_stubs")) + config = _load_config(config_path) + project_dir = os.getcwd() + + _ensure_grpc_stubs_importable(project_dir) + + if any( + agent.get("provider", "local").upper() == "EC2" + for agent in config.get("agents", []) + ): + _preflight_ec2_deploy(config, project_dir) from ventis.controller.global_controller import GlobalController @@ -260,6 +357,7 @@ def _signal_handler(sig, frame): # ventis clean # # ------------------------------------------------------------------ # + def cmd_clean(args): """ Remove generated stubs, gRPC files, and Docker build contexts. @@ -277,6 +375,7 @@ def cmd_clean(args): logger.info("Cleaning %s...", path) if os.path.isdir(path): import shutil + shutil.rmtree(path) else: os.remove(path) @@ -288,6 +387,7 @@ def cmd_clean(args): # Main entry point # # ------------------------------------------------------------------ # + def main(): parser = argparse.ArgumentParser( prog="ventis", @@ -309,9 +409,10 @@ def main(): help="Generate stubs, compile protos, and build Docker images", ) build.add_argument( - "-c", "--config", - default="config/global_controller.yaml", - help="Path to global controller config (default: config/global_controller.yaml)", + "-c", + "--config", + default=DEFAULT_CONFIG_PATH, + help=f"Path to global controller config (default: {DEFAULT_CONFIG_PATH})", ) build.set_defaults(func=cmd_build) @@ -321,9 +422,10 @@ def main(): help="Launch agents via the Global Controller", ) deploy.add_argument( - "-c", "--config", - default="config/global_controller.yaml", - help="Path to global controller config (default: config/global_controller.yaml)", + "-c", + "--config", + default=DEFAULT_CONFIG_PATH, + help=f"Path to global controller config (default: {DEFAULT_CONFIG_PATH})", ) deploy.set_defaults(func=cmd_deploy) diff --git a/ventis/controller/cloud_provider_logic/EC2/_runtime.py b/ventis/controller/cloud_provider_logic/EC2/_runtime.py new file mode 100644 index 0000000..f6d306b --- /dev/null +++ b/ventis/controller/cloud_provider_logic/EC2/_runtime.py @@ -0,0 +1,322 @@ +""" +EC2 runtime helpers for Ventis. + +This module is the EC2-specific backend for `provider: EC2` agents. +It does four things: +1. validate the EC2 config +2. create an EC2 instance +3. start the agent container on that instance +4. clean up the EC2 instance if startup fails + +The global controller sets `_controller` before calling these helpers so +they can read config and reuse the controller's Docker/Redis logic. +""" + +import base64 +import os +import socket +import subprocess +import time + +import boto3 + +CONTAINER_PORT = 50051 +DEFAULT_SSH_KEY_PATH = os.path.expanduser("~/.ssh/ventis_ec2") +_controller = None + + +def _require_controller(): + """Return the global controller, or raise if it was never set.""" + if _controller is None: + raise RuntimeError("EC2 runtime controller is not configured.") + return _controller + + +def _aws_clients(): + """Return validated EC2 config and EC2 client.""" + cfg = _require_controller().config.get("ec2", {}) + required = [ + "ami_id", + "subnet_id", + "security_group_ids", + "region", + "ssh_user", + ] + missing = [field for field in required if not cfg.get(field)] + if missing: + raise ValueError(f"Missing EC2 config: {', '.join(sorted(missing))}") + if not isinstance(cfg["security_group_ids"], list) or not all( + cfg["security_group_ids"] + ): + raise ValueError("EC2 security_group_ids must be a non-empty list.") + + session_kwargs = {"region_name": cfg["region"]} + for key in ( + "profile", + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + ): + value = cfg.get(key) + if value: + session_kwargs["profile_name" if key == "profile" else key] = value + + session = boto3.Session(**session_kwargs) + if not session.region_name: + raise ValueError("EC2 region must be configured.") + if session.get_credentials() is None: + raise ValueError("AWS credentials are not available for the EC2 runtime.") + return cfg, session.client("ec2", region_name=session.region_name) + + +def validate_config(): + """Check that the EC2 config has the fields needed to launch instances.""" + cfg, _ = _aws_clients() + return cfg + + +def _ensure_ssh_keypair(): + """Create ~/.ssh/ventis_ec2 if missing and return the public key.""" + private = DEFAULT_SSH_KEY_PATH + public = private + ".pub" + if not os.path.exists(private): + key_dir = os.path.dirname(private) + if key_dir: + os.makedirs(key_dir, exist_ok=True) + subprocess.run( + [ + "ssh-keygen", + "-t", + "ed25519", + "-f", + private, + "-N", + "", + "-C", + "ventis-ec2", + ], + check=True, + capture_output=True, + ) + with open(public) as f: + return f.read().strip() + + +def _userdata(ssh_user, pubkey): + """Build base64 UserData that installs the public key on the worker.""" + script = ( + "#!/bin/bash\n" + f"mkdir -p /home/{ssh_user}/.ssh\n" + f"echo '{pubkey}' >> /home/{ssh_user}/.ssh/authorized_keys\n" + f"chown -R {ssh_user}:{ssh_user} /home/{ssh_user}/.ssh\n" + f"chmod 700 /home/{ssh_user}/.ssh\n" + f"chmod 600 /home/{ssh_user}/.ssh/authorized_keys\n" + ) + return base64.b64encode(script.encode()).decode() + + +def provision_instance(spec, replica_index, next_host_port=None): + """Launch one EC2 instance for an agent replica and wait for its IPs.""" + cfg, client = _aws_clients() + pubkey = _ensure_ssh_keypair() + agent_name = spec["name"] + request = { + "ImageId": cfg["ami_id"], + "InstanceType": spec["instance_type"], + "SubnetId": cfg["subnet_id"], + "SecurityGroupIds": cfg["security_group_ids"], + "UserData": _userdata(cfg["ssh_user"], pubkey), + "MinCount": 1, + "MaxCount": 1, + "TagSpecifications": [ + { + "ResourceType": "instance", + "Tags": [ + {"Key": "Name", "Value": f"ventis-{agent_name}-{replica_index}"}, + {"Key": "CreatedBy", "Value": "EC2 Fast Launch"}, + ], + }, + { + "ResourceType": "volume", + "Tags": [ + { + "Key": "CreatedBy", + "Value": "EC2 Fast Launch", + } + ], + }, + ], + } + + response = client.run_instances(**request) + instance_id = response["Instances"][0]["InstanceId"] + runtime_id = f"ventis-ec2-{agent_name.lower()}-{replica_index}--{instance_id}" + client.get_waiter("instance_running").wait(InstanceIds=[instance_id]) + + deadline = time.time() + cfg.get("public_ip_timeout", 120) + instance = None + while time.time() < deadline: + response = client.describe_instances(InstanceIds=[instance_id]) + for reservation in response.get("Reservations", []): + for candidate in reservation.get("Instances", []): + if candidate.get("InstanceId") == instance_id: + instance = candidate + break + if instance: + break + if instance and ( + instance.get("PrivateIpAddress") or instance.get("PublicIpAddress") + ): + break + time.sleep(2) + + host = ( + instance.get("PrivateIpAddress") or instance.get("PublicIpAddress") + if instance + else None + ) + if not host: + raise RuntimeError( + f"EC2 instance {instance_id} does not have a reachable IP address." + ) + + redis_cfg = _require_controller().config.get("redis", {}) + record = { + "host": host, + "runtime_id": runtime_id, + "ec2_instance_id": instance_id, + "redis_host": host, + "redis_port": spec.get("redis_port", redis_cfg.get("port", 6379)), + } + return record + + +def bootstrap_instance(provisioned, spec, replica_index): + """Start the agent container on the new EC2 host and return its record.""" + cfg = validate_config() + host = provisioned["host"] + runtime_id = provisioned["runtime_id"] + redis_cfg = _require_controller().config.get("redis", {}) + redis_host = provisioned.get("redis_host", host) + redis_port = provisioned.get( + "redis_port", spec.get("redis_port", redis_cfg.get("port", 6379)) + ) + + try: + _bootstrap_instance( + host, + spec, + replica_index, + cfg, + redis_host=redis_host, + redis_port=redis_port, + ) + _check_controller_health( + f"{host}:{CONTAINER_PORT}", + timeout=cfg.get("controller_health_timeout", 180), + ) + return { + "agent_name": spec["name"], + "provider": "EC2", + "replica_index": str(replica_index), + "host": host, + "host_port": str(CONTAINER_PORT), + "container_port": str(CONTAINER_PORT), + "endpoint": f"{host}:{CONTAINER_PORT}", + "redis_host": redis_host, + "redis_port": str(redis_port), + "runtime_id": runtime_id, + "ec2_instance_id": provisioned.get("ec2_instance_id"), + } + except Exception: + terminate_instance(provisioned) + raise + + +def _bootstrap_instance( + host, spec, replica_index, cfg, redis_host=None, redis_port=None +): + """Run the agent container over SSH.""" + controller = _require_controller() + ssh_user = cfg["ssh_user"] + + for _ in range(30): + result = controller._run_cmd(["true"], host, user=ssh_user) + if result.returncode == 0: + break + time.sleep(2) + else: + raise TimeoutError(f"SSH never became ready on {host}") + + agent_name = spec["name"] + image = f"ventis-{agent_name.lower()}" + redis_cfg = controller.config.get("redis", {}) + redis_host = redis_host or redis_cfg.get("host", "localhost") + redis_port = redis_port or spec.get("redis_port", redis_cfg.get("port", 6379)) + container_name = f"ventis-ec2-{agent_name.lower()}-{replica_index}" + + cmd = [ + "docker", + "run", + "-d", + "-it", + "--restart", + "unless-stopped", + "--add-host=host.docker.internal:host-gateway", + "--name", + container_name, + "-p", + f"{CONTAINER_PORT}:{CONTAINER_PORT}", + "-e", + f"VENTIS_REDIS_HOST={redis_host}", + "-e", + f"VENTIS_REDIS_PORT={redis_port}", + "-e", + f"VENTIS_AGENT_HOST={host}", + "-e", + f"VENTIS_AGENT_PORT={CONTAINER_PORT}", + image, + ] + result = controller._run_cmd(cmd, host, user=ssh_user) + if result.returncode != 0: + raise RuntimeError( + f"SSH bootstrap failed on {host}: {(result.stderr or result.stdout or '').strip()}" + ) + + +def _check_controller_health(endpoint, timeout=None): + """Wait until the launched container accepts TCP connections.""" + host, port = endpoint.split(":") + deadline = time.time() + ( + timeout + or _require_controller() + .config.get("ec2", {}) + .get("controller_health_timeout", 180) + ) + while time.time() < deadline: + try: + with socket.create_connection((host, int(port)), timeout=2): + return True + except OSError: + time.sleep(2) + raise TimeoutError(f"EC2 runtime endpoint never became reachable at {endpoint}.") + + +def terminate_instance(instance): + """Delete the EC2 instance that belongs to a runtime id.""" + runtime_id = instance.get("runtime_id") if isinstance(instance, dict) else instance + if not runtime_id or "--" not in runtime_id: + raise ValueError(f"Invalid EC2 runtime id: {runtime_id}") + + host = instance.get("host") if isinstance(instance, dict) else None + if host: + getattr(_controller, "redis_containers", {}).pop(host, None) + getattr(_controller, "node_redis", {}).pop(host, None) + + _, client = _aws_clients() + client.terminate_instances(InstanceIds=[runtime_id.rsplit("--", 1)[1]]) + + +def routing_endpoint_for(instance): + """Return the gRPC endpoint string used for routing to this instance.""" + return instance["endpoint"] diff --git a/ventis/controller/cloud_provider_logic/Local/_runtime.py b/ventis/controller/cloud_provider_logic/Local/_runtime.py new file mode 100644 index 0000000..3391de8 --- /dev/null +++ b/ventis/controller/cloud_provider_logic/Local/_runtime.py @@ -0,0 +1,132 @@ +""" +Local runtime helpers for Ventis. + +This module is the local-provider backend for `provider: local` agents. +It keeps the existing Docker launch/teardown behavior while letting +InstanceManager stay focused on orchestration and persistence. +""" + +import logging + +logger = logging.getLogger(__name__) + +DEFAULT_HOST = "localhost" +CONTAINER_PORT = 50051 +WORKFLOW_API_PORT = 8080 +PROVIDER = "local" +_controller = None + + +def _require_controller(): + if _controller is None: + raise RuntimeError("Local runtime controller is not configured.") + return _controller + + +def _is_local_host(host): + return host in {"localhost", "127.0.0.1"} + + +def _container_routing_host(host): + return "host.docker.internal" if _is_local_host(host) else host + + +def validate_config(): + return None + + +def provision_instance(spec, replica_index, next_host_port): + host = spec.get("host", DEFAULT_HOST) + host_port = int(spec.get("host_port", spec.get("port", next_host_port(host)))) + agent_name = spec["name"] + + return { + "provider": PROVIDER, + "host": host, + "host_port": host_port, + "redis_host": _container_routing_host(host), + "runtime_id": f"ventis-{PROVIDER}-{agent_name.lower()}-{replica_index}", + "user": spec.get("user"), + } + + +def bootstrap_instance(provisioned, spec, replica_index): + agent_name = spec["name"] + resources = spec.get("resources", {}) + ctrl_type = spec.get("type", "agent") + image = f"ventis-{agent_name.lower()}" + host = provisioned["host"] + host_port = provisioned["host_port"] + user = provisioned.get("user") + redis_host = provisioned["redis_host"] + runtime_id = provisioned["runtime_id"] + + cmd = [ + "docker", + "run", + "-d", + "-it", + "--add-host=host.docker.internal:host-gateway", + "--name", + runtime_id, + "-p", + f"{host_port}:{CONTAINER_PORT}", + "-e", + f"VENTIS_AGENT_PORT={host_port}", + "-e", + f"VENTIS_AGENT_HOST={redis_host}", + "-e", + f"VENTIS_REDIS_HOST={redis_host}", + "-e", + f"VENTIS_REDIS_PORT={spec.get('redis_port', 6379)}", + ] + if ctrl_type == "workflow": + cmd.extend(["-p", f"{WORKFLOW_API_PORT}:8080"]) + if resources.get("cpu"): + cmd.extend(["--cpus", str(resources["cpu"])]) + if resources.get("memory"): + cmd.extend(["--memory", f"{resources['memory']}m"]) + if resources.get("gpu"): + cmd.extend(["--gpus", str(resources["gpu"])]) + cmd.append(image) + + result = _require_controller()._run_cmd(cmd, host, user) + if result.returncode != 0: + raise RuntimeError(f"Failed to launch {runtime_id}") + + instance = { + "agent_name": agent_name, + "provider": PROVIDER, + "replica_index": str(replica_index), + "host": host, + "host_port": str(host_port), + "container_port": str(CONTAINER_PORT), + "endpoint": f"{host}:{host_port}", + "redis_host": redis_host, + "redis_port": str(spec.get("redis_port", 6379)), + "runtime_id": runtime_id, + } + if user: + instance["user"] = user + logger.info("Runtime ready: %s -> %s", runtime_id, instance["endpoint"]) + return instance + + +def terminate_instance(instance): + runtime_id = instance.get("runtime_id") + if not runtime_id: + return + + result = _require_controller()._run_cmd( + ["docker", "rm", "-f", runtime_id], + instance.get("host", DEFAULT_HOST), + instance.get("user"), + ) + if result.returncode != 0: + logger.warning("Failed to remove runtime %s", runtime_id) + + +def routing_endpoint_for(instance): + host = instance.get("host") + port = instance["host_port"] + return f"{_container_routing_host(host)}:{port}" diff --git a/ventis/controller/global_controller.py b/ventis/controller/global_controller.py index 58655c1..4a266dc 100644 --- a/ventis/controller/global_controller.py +++ b/ventis/controller/global_controller.py @@ -13,7 +13,9 @@ import os import yaml - +from ventis.controller.instance_manager import InstanceManager +from ventis.controller.utils.agent_specs import write_agent_specs +from ventis.controller.utils.redis_utils import _wait_for_redis from ventis.utils.redis_client import RedisClient # Add generated grpc_stubs from the local project to the path @@ -22,13 +24,18 @@ import local_controler_pb2_grpc import grpc -print(f"DEBUG: Loading gRPC stubs from: {local_controler_pb2_grpc.__file__}") -print(f"DEBUG: LocalControllerStub attributes: {[a for a in dir(local_controler_pb2_grpc.LocalControllerStub) if not a.startswith('_')]}") - logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) +def _is_local_host(host): + return host in {"localhost", "127.0.0.1"} + + +def _container_routing_host(host): + return "host.docker.internal" if _is_local_host(host) else host + + class GlobalController(object): """ Daemon that manages a routing table across multiple local controller instances. @@ -60,28 +67,26 @@ def __init__(self, config_path): self.cleanup_interval = self.config.get("cleanup_interval", 10) self.controllers = self.config.get("agents", []) self.running = False - self.processes = {} # name -> [Popen, ...] self.containers = {} # name -> [container_name, ...] self.redis_containers = {} # host -> container_name self.node_redis = {} # host -> RedisClient self._last_status = {} # (host, port) -> last known status - self._lc_stubs = {} # endpoint -> gRPC stub - self._shipped_images = set() # (image, host) already shipped this session - - # Registry config: if set, use push/pull; otherwise fall back to SSH pipe - registry_cfg = self.config.get("registry", {}) - self.registry_url = registry_cfg.get("url") # e.g. "myregistry.example.com:5000" - self.registry_user = registry_cfg.get("user") + self._lc_stubs = {} # endpoint -> gRPC stub + self.instance_manager = InstanceManager(self) # Clean up any stale containers from previous runs self._cleanup_stale_containers() # Launch Redis on each unique node, then write routing table and policies self._launch_redis_containers() - self._build_routing_table() + write_agent_specs(self.config_path, self.redis) self._write_resource_specs() self._load_and_write_policies() - logger.info("Global controller initialized with %d controller(s).", len(self.controllers)) + self.instance_manager.publish_routing_snapshot(self.controllers) + logger.info( + "Global controller initialized with %d controller(s).", + len(self.controllers), + ) # Start background cleanup thread self._cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True) @@ -107,19 +112,14 @@ def _cleanup_stale_containers(self): for i, (host, port) in enumerate(placements): if host not in host_containers: host_containers[host] = (user, set()) - - # Redis container for this host host_containers[host][1].add(f"ventis-redis-{host.replace('.', '-')}") - # Agent container host_containers[host][1].add(f"ventis-{name.lower()}-{i}") # Try to remove each one on its respective host for host, (user, container_names) in host_containers.items(): for container_name in container_names: try: - self._run_cmd( - ["docker", "rm", "-f", container_name], host, user - ) + self._run_cmd(["docker", "rm", "-f", container_name], host, user) except Exception: pass # Container didn't exist, that's fine @@ -135,116 +135,53 @@ def _load_config(config_path): with open(config_path, "r") as f: return yaml.safe_load(f) - def reload_config(self): - """Reload the config file and rebuild the routing table.""" - logger.info("Reloading config from %s", self.config_path) - self.config = self._load_config(self.config_path) - self.controllers = self.config.get("agents", []) - self.poll_interval = self.config.get("poll_interval", 5) - self._build_routing_table() - - # ------------------------------------------------------------------ # - # Routing table # - # ------------------------------------------------------------------ # - @staticmethod def _get_replica_placements(ctrl): - """Normalize the ``replicas`` field into a list of (host, port) tuples. - - Supports two formats in the YAML config: - - 1. **Integer shorthand** — ``replicas: 3`` - Uses the agent's ``host`` and sequential ports starting at ``port``. - - 2. **Explicit list** — each entry specifies its own ``host`` and ``port``: - :: - - replicas: - - host: node1 - port: 8051 - - host: node2 - port: 8052 - """ + """Normalize replicas into a list of (host, port) placements.""" replicas = ctrl.get("replicas", 1) default_host = ctrl.get("host", "localhost") base_port = ctrl.get("port", 50051) if isinstance(replicas, int): return [(default_host, base_port + i) for i in range(replicas)] - elif isinstance(replicas, list): + if isinstance(replicas, list): return [ (r.get("host", default_host), r.get("port", base_port)) for r in replicas ] - else: - return [(default_host, base_port)] - - def _build_routing_table(self): - """Write the routing table to Redis on every node. + return [(default_host, base_port)] - For each agent, stores a JSON list of all replica endpoints under - ``routing_table:endpoints``. Agents marked ``stateful: true`` are - recorded in ``routing_table:stateful`` so local controllers can - enforce session affinity. - """ - endpoints_table = {} # name → JSON list of endpoints - stateful_table = {} # name → "true" (only for stateful agents) - - for ctrl in self.controllers: - name = ctrl["name"] - stateful = ctrl.get("stateful", False) - placements = self._get_replica_placements(ctrl) - - endpoints = [] - for host, port in placements: - # Use host.docker.internal for localhost so Docker containers - # can reach each other through the host's port mappings. - rt_host = "host.docker.internal" if host in ("localhost", "127.0.0.1") else host - endpoints.append(f"{rt_host}:{port}") - - endpoints_table[name] = json.dumps(endpoints) - if stateful: - stateful_table[name] = "true" - - # Write to every node's Redis so each local controller can look up - # the full routing table from its own Redis instance. - targets = list(self.node_redis.values()) if self.node_redis else [self.redis] - for redis_client in targets: - if endpoints_table: - redis_client.hset_multiple(self.ROUTING_ENDPOINTS_KEY, endpoints_table) - if stateful_table: - redis_client.hset_multiple(self.ROUTING_STATEFUL_KEY, stateful_table) - - existing = redis_client.smembers(self.SERVICES_SET_KEY) - for stale in existing - set(endpoints_table.keys()): - redis_client.srem(self.SERVICES_SET_KEY, stale) - for name in endpoints_table.keys(): - redis_client.sadd(self.SERVICES_SET_KEY, name) - - logger.info("Routing table written to %d Redis instance(s): %s", - len(targets), endpoints_table) - self._on_routing_table_updated(endpoints_table) + def reload_config(self): + """Reload the config file and rebuild the routing table.""" + logger.info("Reloading config from %s", self.config_path) + self.config = self._load_config(self.config_path) + self.controllers = self.config.get("agents", []) + self.poll_interval = self.config.get("poll_interval", 5) + self.instance_manager.publish_routing_snapshot(self.controllers) def _write_resource_specs(self): """Write the per-agent resource specs to Redis.""" for ctrl in self.controllers: name = ctrl["name"] resources = ctrl.get("resources", {}) - placements = self._get_replica_placements(ctrl) - - self.redis.hset_multiple(f"agent:{name}:resources", { - "cpu": str(resources.get("cpu", 1)), - "memory": str(resources.get("memory", 512)), - "replicas": str(len(placements)), - }) + self.redis.hset_multiple( + f"agent:{name}:resources", + { + "cpu": str(resources.get("cpu", 1)), + "memory": str(resources.get("memory", 512)), + "replicas": str(int(ctrl.get("replicas", 1))), + }, + ) - def _load_and_write_policies(self): - """Load policy rules from config/policy.yaml and write to all Redis instances.""" + def _load_policy_rules(self): + """Load policy rules from config/policy.yaml.""" config_dir = os.path.dirname(os.path.abspath(self.config_path)) policy_path = os.path.join(config_dir, "policy.yaml") if not os.path.isfile(policy_path): - logger.info("No policy file found at %s, skipping policy setup.", policy_path) + logger.info( + "No policy file found at %s, skipping policy setup.", policy_path + ) return with open(policy_path, "r") as f: @@ -255,23 +192,25 @@ def _load_and_write_policies(self): # Sort rules by specificity: most match keys first # This way the local controller can iterate and use the first matching rule. rules.sort(key=lambda r: len(r.get("match", {})), reverse=True) + return rules + def _load_and_write_policies(self): + """Load policy rules and publish them to every host Redis.""" + rules = self._load_policy_rules() + targets = list(self.node_redis.values()) or [self.redis] rules_json = json.dumps(rules) - - # Write to every node's Redis - targets = list(self.node_redis.values()) if self.node_redis else [self.redis] for redis_client in targets: - redis_client.set(self.POLICY_RULES_KEY, rules_json) + redis_client.set("policy:rules", rules_json) - logger.info("Policy rules written to %d Redis instance(s): %d rule(s)", len(targets), len(rules)) - - def get_routing_table(self): - """Read the current routing table from Redis.""" - return self.redis.hgetall(self.ROUTING_TABLE_KEY) + logger.info( + "Policy rules written to %d Redis instance(s): %d rule(s)", + len(targets), + len(rules), + ) - def get_endpoint(self, service_name): - """Look up the endpoint for a given service.""" - return self.redis.hget(self.ROUTING_TABLE_KEY, service_name) + # 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) def get_node_redis(self, host): """Get the RedisClient for a specific node.""" @@ -307,9 +246,13 @@ def _launch_redis_containers(self): container_name = f"ventis-redis-{host.replace('.', '-')}" cmd = [ - "docker", "run", "-d", - "--name", container_name, - "-p", f"{redis_port}:6379", + "docker", + "run", + "-d", + "--name", + container_name, + "-p", + f"{redis_port}:6379", "redis:alpine", ] @@ -319,16 +262,21 @@ def _launch_redis_containers(self): self.redis_containers[host] = container_name logger.info( "Launched Redis container %s on %s:%d", - container_name, host, redis_port, + container_name, + host, + redis_port, ) else: logger.critical( "Failed to launch Redis on %s: %s", - host, result.stderr.strip(), + host, + result.stderr.strip(), ) sys.exit(1) except FileNotFoundError: - logger.critical("Docker is not installed or not in PATH. Cannot launch Redis.") + 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) @@ -337,9 +285,9 @@ def _launch_redis_containers(self): # 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 - self.node_redis[host] = RedisClient( - host=connect_host, port=redis_port, - ) + redis_client = RedisClient(host=connect_host, port=redis_port) + _wait_for_redis(redis_client, host, redis_port) + self.node_redis[host] = redis_client # Update the primary redis client to the local node's Redis if "localhost" in self.node_redis: @@ -349,16 +297,16 @@ def _launch_redis_containers(self): def _stop_redis_containers(self): """Stop and remove all launched Redis containers.""" + nodes = {} + for ctrl in self.controllers: + if ctrl.get("provider", "local").upper() == "EC2": + continue + user = ctrl.get("user") + redis_port = ctrl.get("redis_port", 6379) + for host, _port in self._get_replica_placements(ctrl): + nodes.setdefault(host, {"user": user, "redis_port": redis_port}) for host, container_name in self.redis_containers.items(): - # Find the SSH user for this host by checking all replica placements - user = None - for ctrl in self.controllers: - for rhost, _rport in self._get_replica_placements(ctrl): - if rhost == host: - user = ctrl.get("user") - break - if user is not None: - break + user = nodes.get(host, {}).get("user") try: self._run_cmd(["docker", "stop", container_name], host, user) self._run_cmd(["docker", "rm", container_name], host, user) @@ -379,7 +327,7 @@ def _get_node_redis_for(self, host): def _agent_host_key(self, host): """Return the host string as seen by Docker containers (for status key matching).""" - return "host.docker.internal" if host in ("localhost", "127.0.0.1") else host + return _container_routing_host(host) def _wait_for_healthy(self, timeout=30, interval=2): """ @@ -390,15 +338,16 @@ def _wait_for_healthy(self, timeout=30, interval=2): interval: Seconds between checks. """ deadline = time.time() + timeout - # Build a list of every individual replica to check - pending = [] - for c in self.controllers: - name = c["name"] - for host, port in self._get_replica_placements(c): - pending.append((name, host, port)) + pending = [ + (instance["agent_name"], instance["host"], instance["host_port"]) + for instance in self.instance_manager.list_instances() + ] - logger.info("Waiting for %d replica(s) to become healthy (timeout=%ds)...", - len(pending), timeout) + logger.info( + "Waiting for %d replica(s) to become healthy (timeout=%ds)...", + len(pending), + timeout, + ) while pending and time.time() < deadline: still_pending = [] @@ -419,7 +368,10 @@ def _wait_for_healthy(self, timeout=30, interval=2): for name, host, port in pending: logger.warning( "Controller %s (%s:%s) not ready after %ds.", - name, host, port, timeout, + name, + host, + port, + timeout, ) # ------------------------------------------------------------------ # @@ -441,33 +393,40 @@ def run(self): def _poll_controllers(self): """Check the health of each registered controller replica via its node's Redis.""" - for ctrl in self.controllers: - name = ctrl["name"] - for host, port in self._get_replica_placements(ctrl): - node_redis = self._get_node_redis_for(host) - agent_host = self._agent_host_key(host) - status_key = f"controller:{agent_host}:{port}:status" - - status = node_redis.get(status_key) or "unknown" - prev = self._last_status.get((host, port)) - - if status != prev: - if status == "healthy": - logger.info("Controller %s (%s:%s) is now healthy.", name, host, port) - self._on_controller_healthy(name, host, port) - else: - logger.warning( - "Controller %s (%s:%s) status changed: %s -> %s", - name, host, port, prev or "(none)", status, - ) - self._on_controller_unhealthy(name, host, port) - self._last_status[(host, port)] = status + for instance in self.instance_manager.list_instances(): + name = instance["agent_name"] + host = instance["host"] + port = instance["host_port"] + node_redis = self._get_node_redis_for(host) + agent_host = self._agent_host_key(host) + status_key = f"controller:{agent_host}:{port}:status" + + status = node_redis.get(status_key) or "unknown" + prev = self._last_status.get((host, port)) + + if status != prev: + if status == "healthy": + logger.info( + "Controller %s (%s:%s) is now healthy.", name, host, port + ) + self._on_controller_healthy(name, host, port) else: - # No change — healthy stays quiet, unhealthy stays quiet too - if status == "healthy": - self._on_controller_healthy(name, host, port) - else: - self._on_controller_unhealthy(name, host, port) + logger.warning( + "Controller %s (%s:%s) status changed: %s -> %s", + name, + host, + port, + prev or "(none)", + status, + ) + self._on_controller_unhealthy(name, host, port) + self._last_status[(host, port)] = status + else: + # No change — healthy stays quiet, unhealthy stays quiet too + if status == "healthy": + self._on_controller_healthy(name, host, port) + else: + self._on_controller_unhealthy(name, host, port) # ------------------------------------------------------------------ # # Extensibility hooks — override in subclasses # @@ -493,7 +452,9 @@ def _get_lc_stub(self, endpoint): """Get or create a cached gRPC stub for a local controller endpoint.""" if endpoint not in self._lc_stubs: channel = grpc.insecure_channel(endpoint) - self._lc_stubs[endpoint] = local_controler_pb2_grpc.LocalControllerStub(channel) + self._lc_stubs[endpoint] = local_controler_pb2_grpc.LocalControllerStub( + channel + ) return self._lc_stubs[endpoint] def _cleanup_loop(self): @@ -513,115 +474,23 @@ def _trigger_cleanup(self): for request_id in completed: logger.info("Triggering cleanup for completed request %s", request_id) - for ctrl in self.controllers: - for host, port in self._get_replica_placements(ctrl): - endpoint = f"{host}:{port}" - 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) + 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) # Remove from completed set after broadcast self.redis.srem("request:completed", request_id) # ------------------------------------------------------------------ # - # Agent launching # - # ------------------------------------------------------------------ # - def launch_agents(self): - """ - Launch all agents defined in the config. - - For each controller entry, spawn `replicas` number of subprocesses - using the configured entrypoint script. Each replica gets assigned - a port starting from the controller's base port. - """ - project_root = os.path.abspath( - os.path.join(os.path.dirname(__file__), "..", "..") - ) - - for ctrl in self.controllers: - name = ctrl["name"] - placements = self._get_replica_placements(ctrl) - entrypoint = ctrl.get("entrypoint") - - if not entrypoint: - logger.warning("No entrypoint for %s, skipping launch.", name) - continue - - entrypoint_path = os.path.join(project_root, entrypoint) - if not os.path.isfile(entrypoint_path): - logger.error("Entrypoint not found: %s", entrypoint_path) - continue - - self.processes[name] = [] - for host, port in placements: - proc = self._launch_single_agent(name, entrypoint_path, port, ctrl, host) - if proc: - self.processes[name].append(proc) - - total = sum(len(procs) for procs in self.processes.values()) - logger.info("Launched %d agent process(es) across %d service(s).", - total, len(self.processes)) - - def _launch_single_agent(self, name, entrypoint_path, port, ctrl, host=None): - """ - Launch a single agent subprocess. - - Args: - name: Service/agent name. - entrypoint_path: Absolute path to the agent script. - port: Port number for this instance. - ctrl: Full controller config dict. - - Returns: - The Popen object, or None on failure. - """ - resources = ctrl.get("resources", {}) - env = os.environ.copy() - env["VENTIS_AGENT_NAME"] = name - env["VENTIS_AGENT_PORT"] = str(port) - env["VENTIS_AGENT_CPU"] = str(resources.get("cpu", 1)) - env["VENTIS_AGENT_MEMORY"] = str(resources.get("memory", 512)) - - try: - proc = subprocess.Popen( - [sys.executable, entrypoint_path, "--port", str(port)], - env=env, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - logger.info("Launched %s (pid=%d) on port %d", name, proc.pid, port) - - # Record status in Redis - host = host or ctrl.get("host", "localhost") - self.redis.set(f"controller:{host}:{port}:status", "healthy") - self.redis.set(f"controller:{host}:{port}:pid", str(proc.pid)) - - return proc - except Exception as e: - logger.error("Failed to launch %s on port %d: %s", name, port, e) - return None - - def _stop_agents(self): - """Terminate all launched agent subprocesses.""" - for name, procs in self.processes.items(): - for proc in procs: - if proc.poll() is None: # still running - logger.info("Terminating %s (pid=%d)", name, proc.pid) - proc.terminate() - try: - proc.wait(timeout=5) - except subprocess.TimeoutExpired: - logger.warning("Killing %s (pid=%d)", name, proc.pid) - proc.kill() - self.processes.clear() - logger.info("All agent processes stopped.") - - # ------------------------------------------------------------------ # - # Docker launching # + # Runtime launching # # ------------------------------------------------------------------ # def _run_cmd(self, cmd, host, user=None): @@ -636,223 +505,58 @@ def _run_cmd(self, cmd, host, user=None): Returns: subprocess.CompletedProcess """ - is_local = host in ("localhost", "127.0.0.1") + is_local = _is_local_host(host) if is_local: return subprocess.run(cmd, capture_output=True, text=True) else: ssh_target = f"{user}@{host}" if user else host remote_cmd = " ".join(cmd) + if cmd and cmd[0] == "docker": + remote_cmd = f"sudo {remote_cmd}" return subprocess.run( - ["ssh", ssh_target, remote_cmd], - capture_output=True, text=True, + [ + "ssh", + "-o", + "StrictHostKeyChecking=no", + "-o", + "ConnectTimeout=10", + "-i", + os.path.expanduser("~/.ssh/ventis_ec2"), + ssh_target, + remote_cmd, + ], + capture_output=True, + text=True, ) - def _ensure_image_on_host(self, image, host, user): - """ - Ensure `image` is available on `host` before running a container. - - Strategy: - - If a registry is configured: tag + push to registry, then pull on - the remote host via SSH. - - Fallback (no registry): stream the image over SSH using - ``docker save | ssh ... docker load``. - - Images are only shipped once per (image, host) pair per session. - Does nothing for localhost. - """ - if host in ("localhost", "127.0.0.1"): - return # already on the local Docker engine - - if (image, host) in self._shipped_images: - logger.debug("Image %s already shipped to %s this session, skipping.", image, host) - return - - if self.registry_url: - self._ship_image_registry(image, host, user) - else: - self._ship_image_ssh(image, host, user) - - self._shipped_images.add((image, host)) - - def _ship_image_registry(self, image, host, user): - """ - Push image to the configured registry on the local host, then - instruct the remote host to pull it. - """ - remote_tag = f"{self.registry_url}/{image}:latest" - logger.info("Tagging %s -> %s for registry push...", image, remote_tag) - subprocess.run(["docker", "tag", image, remote_tag], check=True) - - logger.info("Pushing %s to registry %s...", image, self.registry_url) - subprocess.run(["docker", "push", remote_tag], check=True) - - logger.info("Pulling %s on %s from registry...", image, host) - result = self._run_cmd(["docker", "pull", remote_tag], host, user) - if result.returncode != 0: - raise RuntimeError( - f"Failed to pull {remote_tag} on {host}: {result.stderr.strip()}" + def launch_docker_agents(self): + """Launch all configured runtimes through InstanceManager.""" + try: + instances = self.instance_manager.ensure_instances(self.controllers) + except FileNotFoundError: + logger.critical( + "Docker is not installed or not in PATH. Cannot launch agents." ) + self._stop_redis_containers() + sys.exit(1) + except Exception as e: + logger.critical("Failed to launch configured runtimes: %s", e) + self._stop_docker_agents() + self._stop_redis_containers() + sys.exit(1) - # Retag to the expected short name on the remote so container launch works - retag_result = self._run_cmd( - ["docker", "tag", remote_tag, f"{image}:latest"], - host, user, - ) - if retag_result.returncode != 0: - logger.warning("Failed to retag %s on %s: %s", remote_tag, host, retag_result.stderr.strip()) - - logger.info("Image %s is ready on %s (via registry).", image, host) - - def _ship_image_ssh(self, image, host, user): - """ - Stream image to remote host using `docker save | ssh docker load`. - Used as a fallback when no registry is configured. - """ - ssh_target = f"{user}@{host}" if user else host logger.info( - "Shipping image %s to %s via SSH pipe (no registry configured)...", - image, host, - ) - save_proc = subprocess.Popen( - ["docker", "save", image], - stdout=subprocess.PIPE, - ) - load_proc = subprocess.Popen( - ["ssh", ssh_target, "docker load"], - stdin=save_proc.stdout, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, + "Launched %d Docker container(s) across %d service(s).", + len(instances), + len({instance["agent_name"] for instance in instances}), ) - save_proc.stdout.close() - _, stderr = load_proc.communicate() - save_proc.wait() - - if load_proc.returncode != 0: - raise RuntimeError( - f"Failed to ship image {image} to {host} via SSH: {stderr.decode().strip()}" - ) - logger.info("Image %s shipped to %s successfully.", image, host) - - def launch_docker_agents(self): - """ - Launch all agents as Docker containers. - - For each agent in the config, runs `docker run` either locally or - via SSH on the specified host. Spawns `replicas` containers per agent, - each on an incrementing port from the base port. - - Assumes Docker images are pre-built (via `make docker`). - Image name convention: ventis- - """ - for ctrl in self.controllers: - name = ctrl["name"] - default_host = ctrl.get("host", "localhost") - user = ctrl.get("user") - resources = ctrl.get("resources", {}) - ctrl_type = ctrl.get("type", "agent") - placements = self._get_replica_placements(ctrl) - - image = f"ventis-{name.lower()}" - self.containers[name] = [] - - for i, (host, port) in enumerate(placements): - container_name = f"ventis-{name.lower()}-{i}" - - # Containers can't reach host's Redis via "localhost"; - # use host.docker.internal to route to the Docker host. - redis_host_for_container = "host.docker.internal" if host in ("localhost", "127.0.0.1") else host - - cmd = [ - "docker", "run", "-d", "-it", - "--add-host=host.docker.internal:host-gateway", - "--name", container_name, - "-p", f"{port}:50051", - "-e", f"VENTIS_AGENT_PORT={port}", - "-e", f"VENTIS_AGENT_HOST={'host.docker.internal' if host in ('localhost', '127.0.0.1') else host}", - "-e", f"VENTIS_REDIS_HOST={redis_host_for_container}", - "-e", f"VENTIS_REDIS_PORT={ctrl.get('redis_port', 6379)}", - ] - - # Workflow containers also expose the REST API port - if ctrl_type == "workflow": - api_port = ctrl.get("api_port", 8080) - cmd.extend(["-p", f"{api_port}:8080"]) - - # Apply resource limits - cpu = resources.get("cpu") - memory = resources.get("memory") - gpu = resources.get("gpu") - if cpu: - cmd.extend(["--cpus", str(cpu)]) - if memory: - cmd.extend(["--memory", f"{memory}m"]) - if gpu: - # Provide the specific count or identifier requested (e.g., '1', '2', 'all') - cmd.extend(["--gpus", str(gpu)]) - - cmd.append(image) - - try: - # Ensure the image exists on the target host before launching - replica_user = user # TODO: per-replica user support - self._ensure_image_on_host(image, host, replica_user) - result = self._run_cmd(cmd, host, replica_user) - if result.returncode == 0: - container_id = result.stdout.strip()[:12] - self.containers[name].append(container_name) - logger.info( - "Launched container %s (%s) on %s:%d", - container_name, container_id, host, port, - ) - else: - logger.critical( - "Failed to launch %s on %s:%d: %s", - container_name, host, port, result.stderr.strip(), - ) - # Remove the failed container left in "Created" state - self._run_cmd(["docker", "rm", "-f", container_name], host, user) - self._stop_docker_agents() - self._stop_redis_containers() - sys.exit(1) - except FileNotFoundError: - logger.critical("Docker is not installed or not in PATH. Cannot launch agents.") - self._stop_redis_containers() - sys.exit(1) - except Exception as e: - logger.critical( - "Failed to launch %s on %s:%d: %s", - container_name, host, port, e, - ) - self._run_cmd(["docker", "rm", "-f", container_name], host, user) - self._stop_docker_agents() - self._stop_redis_containers() - sys.exit(1) - - total = sum(len(c) for c in self.containers.values()) - logger.info("Launched %d Docker container(s) across %d service(s).", - total, len(self.containers)) def _stop_docker_agents(self): - """Stop and remove all launched Docker containers.""" - for ctrl in self.controllers: - name = ctrl["name"] - user = ctrl.get("user") - placements = self._get_replica_placements(ctrl) - - containers_for_agent = self.containers.get(name, []) - for i, container_name in enumerate(containers_for_agent): - # Match each container to its placement host - if i < len(placements): - host = placements[i][0] - else: - host = ctrl.get("host", "localhost") - - try: - self._run_cmd(["docker", "stop", container_name], host, user) - self._run_cmd(["docker", "rm", container_name], host, user) - logger.info("Stopped and removed %s on %s", container_name, host) - except Exception as e: - logger.warning("Failed to stop %s on %s: %s", container_name, host, e) + """Stop and remove all launched runtimes.""" + for instance in self.instance_manager.list_instances(): + self.instance_manager.remove_instance( + self.instance_manager._instance_id_from_record(instance) + ) self.containers.clear() logger.info("All Docker containers stopped.") @@ -885,7 +589,8 @@ def stop(self): parser = argparse.ArgumentParser(description="Ventis Global Controller daemon.") parser.add_argument( - "-c", "--config", + "-c", + "--config", default=default_config, help="Path to the YAML config file (default: config/global_controller.yaml)", ) diff --git a/ventis/controller/instance_manager.py b/ventis/controller/instance_manager.py new file mode 100644 index 0000000..7c8b141 --- /dev/null +++ b/ventis/controller/instance_manager.py @@ -0,0 +1,201 @@ +""" +Coordinate agent runtime instances for the controller. + +This file decides whether each agent replica should run locally or on EC2, +starts missing instances, records their runtime metadata in Redis, and +publishes the routing data other parts of Ventis use to reach those agents. +""" + +import json + +from ventis.controller.cloud_provider_logic.Local import _runtime as local_runtime + +DEFAULT_HOST_PORT_START = 8000 + + +class InstanceManager: + ROUTING_ENDPOINTS_KEY = "routing_table:endpoints" + ROUTING_STATEFUL_KEY = "routing_table:stateful" + SERVICES_SET_KEY = "routing_table:services" + + def __init__(self, controller, redis_client=None): + self.controller = controller + self._redis = redis_client + + @property + def redis(self): + return self._redis or self.controller.redis + + def ensure_instances(self, agent_specs): + self._agent_specs = list(agent_specs) + instances = [] + + for agent_spec in self._agent_specs: + agent_name = agent_spec["name"] + provider = agent_spec.get("provider", "local") + runtime = self._provider_runtime(provider) + self.controller.containers.setdefault(agent_name, []) + + runtime.validate_config() + + for replica_index in range(int(agent_spec.get("replicas", 1))): + instance_id = self._instance_id(provider, agent_name, replica_index) + key = self._instance_key(provider, agent_name, replica_index) + instance = self.redis.hgetall(key) + + if not instance: + provisioned = runtime.provision_instance( + agent_spec, replica_index, self._next_host_port + ) + instance = runtime.bootstrap_instance( + provisioned, agent_spec, replica_index + ) + self._write_instance(instance) + + self._add_instance_to_agent(agent_name, instance_id) + self._track_runtime(agent_name, instance["runtime_id"]) + instances.append(instance) + + self.publish_routing_snapshot(self._agent_specs) + return instances + + def _write_instance(self, instance): + key = self._instance_key( + instance["provider"], instance["agent_name"], int(instance["replica_index"]) + ) + mapping = { + "agent_name": instance["agent_name"], + "provider": instance["provider"], + "replica_index": str(instance["replica_index"]), + "host": instance["host"], + "host_port": str(instance["host_port"]), + "container_port": str(instance["container_port"]), + "endpoint": instance["endpoint"], + "redis_host": instance["redis_host"], + "redis_port": str(instance["redis_port"]), + "runtime_id": instance["runtime_id"], + } + if instance.get("ec2_instance_id"): + mapping["ec2_instance_id"] = instance["ec2_instance_id"] + if instance.get("user"): + mapping["user"] = instance["user"] + self.redis.hset_multiple(key, mapping) + + def _add_instance_to_agent(self, agent_name, instance_id): + self.redis.sadd(f"agent:{agent_name}:instances", instance_id) + + def remove_instance(self, instance_id): + key = f"agent_instance:{instance_id}" + instance = self.redis.hgetall(key) + if not instance: + return + + self._destroy_runtime(instance) + self.redis.delete(key) + self.redis.srem(f"agent:{instance['agent_name']}:instances", instance_id) + self.controller.containers[instance["agent_name"]] = [ + runtime_id + for runtime_id in self.controller.containers.get(instance["agent_name"], []) + if runtime_id != instance["runtime_id"] + ] + self.publish_routing_snapshot( + getattr(self, "_agent_specs", getattr(self.controller, "controllers", [])) + ) + + def list_instances(self, agent_name=None): + if agent_name: + instance_ids = sorted(self.redis.smembers(f"agent:{agent_name}:instances")) + return [ + instance + for instance_id in instance_ids + if (instance := self.redis.hgetall(f"agent_instance:{instance_id}")) + ] + + return [ + instance + for key in sorted(self.redis.scan_keys("agent_instance:*")) + if (instance := self.redis.hgetall(key)) + ] + + def _destroy_runtime(self, instance): + runtime = self._provider_runtime(instance.get("provider", "local")) + runtime.terminate_instance(instance) + + def _track_runtime(self, agent_name, runtime_id): + containers = self.controller.containers.setdefault(agent_name, []) + if runtime_id not in containers: + containers.append(runtime_id) + + def _next_host_port(self, host): + used = { + int(instance["host_port"]) + for instance in self.list_instances() + if instance.get("host") == host and instance.get("host_port") + } + port = DEFAULT_HOST_PORT_START + while port in used: + port += 1 + return port + + @staticmethod + def _instance_id(provider, agent_name, replica_index): + return f"{provider}:{agent_name}:{replica_index}" + + @classmethod + def _instance_key(cls, provider, agent_name, replica_index): + return f"agent_instance:{cls._instance_id(provider, agent_name, replica_index)}" + + def _instance_id_from_record(self, instance): + return self._instance_id( + instance["provider"], instance["agent_name"], int(instance["replica_index"]) + ) + + def _provider_runtime(self, provider): + if provider.upper() == "EC2": + from ventis.controller.cloud_provider_logic.EC2 import _runtime as runtime + else: + runtime = local_runtime + runtime._controller = self.controller + return runtime + + def publish_routing_snapshot(self, agent_specs): + services = {agent_spec["name"] for agent_spec in agent_specs} + stateful = { + agent_spec["name"] + for agent_spec in agent_specs + if agent_spec.get("stateful", False) + } + targets = list(getattr(self.controller, "node_redis", {}).values()) or [ + self.redis + ] + + for redis_client in targets: + hdel = getattr(redis_client, "hdel", None) or redis_client.client.hdel + existing_services = redis_client.smembers(self.SERVICES_SET_KEY) + for stale in existing_services - services: + redis_client.srem(self.SERVICES_SET_KEY, stale) + hdel(self.ROUTING_STATEFUL_KEY, stale) + hdel(self.ROUTING_ENDPOINTS_KEY, stale) + for service in services: + redis_client.sadd(self.SERVICES_SET_KEY, service) + if service in stateful: + redis_client.hset(self.ROUTING_STATEFUL_KEY, service, "true") + else: + hdel(self.ROUTING_STATEFUL_KEY, service) + endpoints = [ + self._routing_endpoint_for(item) + for item in sorted( + self.list_instances(service), + key=lambda item: int(item["replica_index"]), + ) + ] + if endpoints: + redis_client.hset( + self.ROUTING_ENDPOINTS_KEY, service, json.dumps(endpoints) + ) + else: + hdel(self.ROUTING_ENDPOINTS_KEY, service) + + def _routing_endpoint_for(self, instance): + runtime = self._provider_runtime(instance.get("provider", "local")) + return runtime.routing_endpoint_for(instance) diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index 98e1c2b..ee8090d 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -48,11 +48,11 @@ def __init__(self, port=50051): self.agent_host = os.environ.get("VENTIS_AGENT_HOST", "localhost") self.agent_name = os.environ.get("VENTIS_AGENT_NAME") self.agent_file = os.environ.get("VENTIS_AGENT_FILE") - + # Public port is how the routing table and other nodes know us; # internally the gRPC server binds to `port` (50051 inside Docker). self.public_port = os.environ.get("VENTIS_AGENT_PORT", str(port)) - + self._my_endpoint = f"{self.agent_host}:{self.public_port}" self.server, self.servicer = start_server(port, my_endpoint=self._my_endpoint) @@ -67,7 +67,7 @@ def __init__(self, port=50051): # Cache for gRPC stubs to remote controllers self._remote_channels = {} # endpoint -> grpc.Channel - self._remote_stubs = {} # endpoint -> LocalControllerStub + self._remote_stubs = {} # endpoint -> LocalControllerStub # Policy rules cache (loaded lazily from Redis) self._policy_rules = None @@ -78,23 +78,29 @@ def __init__(self, port=50051): max_instances = int(os.environ.get("VENTIS_MAX_AGENT_INSTANCES", 8)) self._executor = ThreadPoolExecutor(max_workers=max_instances) - logger.info("Local controller initialized at %s (max_agent_instances=%d), reported healthy to Redis.", self._my_endpoint, max_instances) - + logger.info( + "Local controller initialized at %s (max_agent_instances=%d), reported healthy to Redis.", + self._my_endpoint, + max_instances, + ) + # Load the agent class dynamically self.agent = self._load_agent() def _load_agent(self): """Dynamically load and instantiate the agent class.""" if not self.agent_name or not self.agent_file: - logger.warning("VENTIS_AGENT_NAME or VENTIS_AGENT_FILE not set. Running without an agent.") + logger.warning( + "VENTIS_AGENT_NAME or VENTIS_AGENT_FILE not set. Running without an agent." + ) return None agent_module_name = self.agent_file.replace(".py", "") - + # We assume the agent file is in the same directory as the local controller (e.g. copied by Docker) # or in the current working directory. agent_path = os.path.abspath(str(self.agent_file)) - + if not os.path.exists(agent_path): logger.error(f"Agent file not found at {agent_path}") return None @@ -102,19 +108,27 @@ def _load_agent(self): try: spec = importlib.util.spec_from_file_location(agent_module_name, agent_path) if spec is None or getattr(spec, "loader", None) is None: - logger.error(f"Cannot find spec or loader for module {agent_module_name} at {agent_path}") + logger.error( + f"Cannot find spec or loader for module {agent_module_name} at {agent_path}" + ) return None - + loader = spec.loader + assert loader is not None + module = importlib.util.module_from_spec(spec) sys.modules[agent_module_name] = module - spec.loader.exec_module(module) - + loader.exec_module(module) + agent_class = getattr(module, self.agent_name) agent_instance = agent_class() - logger.info(f"Successfully loaded and instantiated agent: {self.agent_name}") + logger.info( + f"Successfully loaded and instantiated agent: {self.agent_name}" + ) return agent_instance except Exception as e: - logger.error(f"Failed to load agent {self.agent_name} from {agent_path}: {e}") + logger.error( + f"Failed to load agent {self.agent_name} from {agent_path}: {e}" + ) return None # ------------------------------------------------------------------ # @@ -156,7 +170,11 @@ def _check_policy(self, service, context): return service in access # No rule matched at all - logger.warning("No policy rule matched for context=%s, denying access to %s", context, service) + logger.warning( + "No policy rule matched for context=%s, denying access to %s", + context, + service, + ) return False # ------------------------------------------------------------------ # @@ -192,12 +210,16 @@ def _resolve_endpoint(self, service, request_id): affinity_key = f"affinity:{request_id}" existing = self.redis.hget(affinity_key, service) if existing: - logger.debug("Affinity hit: %s -> %s (request %s)", service, existing, request_id) + logger.debug( + "Affinity hit: %s -> %s (request %s)", service, existing, request_id + ) return existing # No existing binding — pick randomly and persist to Hash chosen = random.choice(endpoints) self.redis.hset(affinity_key, service, chosen) - logger.info("Affinity set: %s -> %s (request %s)", service, chosen, request_id) + logger.info( + "Affinity set: %s -> %s (request %s)", service, chosen, request_id + ) return chosen else: # Stateless: pick randomly @@ -275,26 +297,50 @@ def _process_request(self, data): # Resolve which endpoint to route to endpoint = self._resolve_endpoint(service, request_id) if not endpoint: - logger.error("No endpoint found for service '%s' in routing table.", service) + logger.error( + "No endpoint found for service '%s' in routing table.", service + ) return if endpoint == self._my_endpoint: - self._executor.submit(self._execute_locally, service, function, args, future_id, origin, request_id) + self._executor.submit( + self._execute_locally, + service, + function, + args, + future_id, + origin, + request_id, + ) else: # Register the target as a consumer for any Future args # so results get pushed to its Redis via WriteResult. for key, value in args.items(): - if isinstance(value, str) and len(value) == 32 and all(c in "0123456789abcdefABCDEF" for c in value): + if ( + isinstance(value, str) + and len(value) == 32 + and all(c in "0123456789abcdefABCDEF" for c in value) + ): future_key = f"future:{value}" if self.redis.hget(future_key, "id") is not None: self.redis.sadd(f"{future_key}:consumers", endpoint) - logger.info("Registered %s as consumer of future %s (arg '%s')", endpoint, value, key) + logger.info( + "Registered %s as consumer of future %s (arg '%s')", + endpoint, + value, + key, + ) # If the result is already available, push it immediately. # This handles the race where _notify_consumers already ran. existing_result = self.redis.hget(future_key, "result") if existing_result is not None and existing_result != "": - logger.info("Future %s already resolved, pushing value %s to %s", value, existing_result, endpoint) + logger.info( + "Future %s already resolved, pushing value %s to %s", + value, + existing_result, + endpoint, + ) self._send_result_callback(endpoint, value, existing_result) # Build comprehensive outward baggage so the receiver gets all context and routing descisions @@ -314,7 +360,13 @@ def _process_request(self, data): elif request_id: data["baggage"]["affinities"][service] = endpoint - logger.info("Forwarding %s.%s (future=%s) to %s", service, function, future_id, endpoint) + logger.info( + "Forwarding %s.%s (future=%s) to %s", + service, + function, + future_id, + endpoint, + ) self._forward_request(endpoint, data) def _resolve_future_args(self, args, poll_interval=0.01, timeout=300): @@ -326,9 +378,17 @@ def _resolve_future_args(self, args, poll_interval=0.01, timeout=300): resolved = {} for key, value in args.items(): # Check if this arg value is a UUID hex string identifying a future - if isinstance(value, str) and len(value) == 32 and all(c in "0123456789abcdefABCDEF" for c in value): + if ( + isinstance(value, str) + and len(value) == 32 + and all(c in "0123456789abcdefABCDEF" for c in value) + ): future_key = f"future:{value}" - logger.info("Arg '%s' looks like a Future UUID (%s), waiting for result...", key, value) + logger.info( + "Arg '%s' looks like a Future UUID (%s), waiting for result...", + key, + value, + ) start = time.time() while True: # print("Waiting for result for future next iteration %s", value) @@ -348,7 +408,9 @@ def _resolve_future_args(self, args, poll_interval=0.01, timeout=300): resolved[key] = value return resolved - def _execute_locally(self, service, function, args, future_id, origin=None, request_id=None): + def _execute_locally( + self, service, function, args, future_id, origin=None, request_id=None + ): """Execute a request on the local agent and write the result to Redis.""" # Propagate the request_id context into this worker thread if request_id: @@ -367,7 +429,9 @@ def _execute_locally(self, service, function, args, future_id, origin=None, requ # Resolve any Future IDs in the args before executing args = self._resolve_future_args(args) - logger.info("Executing %s.%s (future=%s) locally", service, function, future_id) + logger.info( + "Executing %s.%s (future=%s) locally", service, function, future_id + ) result = method(**args) # Serialize the result @@ -383,10 +447,16 @@ def _execute_locally(self, service, function, args, future_id, origin=None, requ if origin and origin != self._my_endpoint: self._send_result_callback(origin, future_id, serialized) - logger.info("Completed %s.%s (future=%s) -> %s", service, function, future_id, serialized) + logger.info( + "Completed %s.%s (future=%s) -> %s", + service, + function, + future_id, + serialized, + ) 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}", "result", f"Execution failed: {e}") if origin and origin != self._my_endpoint: @@ -421,7 +491,13 @@ def _forward_request(self, endpoint, data): def _send_result_callback(self, origin, future_id, result): """Send a result back to the originating controller via WriteResult RPC.""" if not result: - logger.warning("Agent '%s' is sending an empty/None result for future %s to origin %s, result: %s", self.agent_name, future_id, origin, result) + logger.warning( + "Agent '%s' is sending an empty/None result for future %s to origin %s, result: %s", + self.agent_name, + future_id, + origin, + result, + ) stub = self._get_remote_stub(origin) payload = json.dumps({"future_id": future_id, "result": result}) @@ -429,7 +505,12 @@ def _send_result_callback(self, origin, future_id, result): request = local_controler_pb2.JsonResponse(resonse=payload) try: stub.WriteResult(request) - logger.info("Sent result callback to %s for future %s, result %s", origin, future_id, result) + logger.info( + "Sent result callback to %s for future %s, result %s", + origin, + future_id, + result, + ) except Exception as e: logger.error("Failed to send result callback to %s: %s", origin, e) @@ -448,6 +529,7 @@ def stop(self): if __name__ == "__main__": import argparse + parser = argparse.ArgumentParser() parser.add_argument("--port", type=int, default=50051) args = parser.parse_args() diff --git a/ventis/controller/local_controller_frontend.py b/ventis/controller/local_controller_frontend.py index e00fb6f..41c8515 100644 --- a/ventis/controller/local_controller_frontend.py +++ b/ventis/controller/local_controller_frontend.py @@ -3,19 +3,12 @@ import grpc from concurrent import futures -import subprocess import os -import signal -from collections import defaultdict -from threading import Lock, Thread +from threading import Thread import json -import time -import redis -import traceback import queue import logging import sys -import os # Add generated grpc_stubs to the path (Docker context copies them directly to /app) sys.path.insert(0, ".") @@ -57,18 +50,26 @@ def WriteResult(self, request, context): future_id = data.get("future_id") result = data.get("result") error = data.get("error") - - logger.info(f"WriteResult: received result for future {future_id}: {result}") + + logger.info( + f"WriteResult: received result for future {future_id}: {result}" + ) if not result: - logger.warning(f"WriteResult received empty/None result for future {future_id} from {context.peer()}") - + logger.warning( + f"WriteResult received empty/None result for future {future_id} from {context.peer()}" + ) + if future_id: if error is not None: self.redis.hset(f"future:{future_id}", "error", error) logger.info("WriteResult: wrote error for future %s", future_id) if result is not None: self.redis.hset(f"future:{future_id}", "result", result) - logger.info("WriteResult: wrote result for future %s, result %s", future_id, result) + logger.info( + "WriteResult: wrote result for future %s, result %s", + future_id, + result, + ) else: logger.error("WriteResult: missing future_id in %s", data) except Exception as e: @@ -81,7 +82,9 @@ def Cleanup(self, request, context): 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() + Thread( + target=self._cleanup_request, args=(request_id,), daemon=True + ).start() else: logger.warning("Cleanup: missing request_id in payload") except Exception as e: @@ -93,7 +96,10 @@ def _cleanup_request(self, request_id): # 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): - logger.info("Cleanup for request %s already claimed by another LC, skipping.", request_id) + logger.info( + "Cleanup for request %s already claimed by another LC, skipping.", + request_id, + ) return try: @@ -105,13 +111,17 @@ def _cleanup_request(self, request_id): keys_to_delete = [futures_key] for fid in future_ids: - keys_to_delete.extend([ - f"future:{fid}", - f"future:{fid}:children", - f"future:{fid}:consumers", - ]) + keys_to_delete.extend( + [ + f"future:{fid}", + f"future:{fid}:children", + f"future:{fid}:consumers", + ] + ) self.redis.delete(*keys_to_delete) - logger.info("Cleaned up %d future(s) for request %s", len(future_ids), request_id) + logger.info( + "Cleaned up %d future(s) for request %s", len(future_ids), request_id + ) # Clean up affinity bindings for this request self.redis.delete(f"affinity:{request_id}") @@ -126,9 +136,7 @@ def start_server(port=50051, my_endpoint="unknown"): servicer = LocalControllerServicer(my_endpoint=my_endpoint) server = grpc.server(futures.ThreadPoolExecutor(max_workers=1)) - local_controler_pb2_grpc.add_LocalControllerServicer_to_server( - servicer, server - ) + local_controler_pb2_grpc.add_LocalControllerServicer_to_server(servicer, server) server.add_insecure_port(f"[::]:{port}") server.start() logger.info(f"Local controller frontend started on port {port}") @@ -143,4 +151,3 @@ def start_server(port=50051, my_endpoint="unknown"): except KeyboardInterrupt: logger.info("Shutting down server...") server.stop(0) - diff --git a/ventis/controller/utils/__init__.py b/ventis/controller/utils/__init__.py new file mode 100644 index 0000000..3db8269 --- /dev/null +++ b/ventis/controller/utils/__init__.py @@ -0,0 +1 @@ +# Ventis Controller Utility helpers diff --git a/ventis/controller/utils/agent_specs.py b/ventis/controller/utils/agent_specs.py new file mode 100644 index 0000000..2b580f7 --- /dev/null +++ b/ventis/controller/utils/agent_specs.py @@ -0,0 +1,23 @@ +import json + +import yaml + + +def write_agent_specs(config_path, redis_client): + with open(config_path, "r") as f: + config = yaml.safe_load(f) or {} + + for agent in config.get("agents", []): + name = agent["name"] + spec_class = "WorkflowSpec" if agent.get("type") == "workflow" else "AgentSpec" + redis_client.hset_multiple( + f"agent:{name}:", + { + "class": spec_class, + "resources": json.dumps(agent.get("resources", {})), + "replicas": json.dumps(agent.get("replicas", 1)), + "stateful": json.dumps(agent.get("stateful", False)), + "redis_port": str(agent.get("redis_port", 6379)), + "provider": agent.get("provider", "local"), + }, + ) diff --git a/ventis/controller/utils/redis_utils.py b/ventis/controller/utils/redis_utils.py new file mode 100644 index 0000000..4bdc265 --- /dev/null +++ b/ventis/controller/utils/redis_utils.py @@ -0,0 +1,20 @@ +import time + + +def _wait_for_redis(redis_client, host, port, timeout=30, interval=1): + """Wait until Redis accepts commands, surfacing network issues clearly.""" + deadline = time.time() + timeout + last_error = None + while time.time() < deadline: + try: + redis_client.set("__ventis_redis_healthcheck__", "ok") + return + except Exception as exc: + last_error = exc + time.sleep(interval) + + raise TimeoutError( + f"Timed out connecting to Redis at {host}:{port}. " + "For EC2 runtimes, ensure the instance security group allows inbound " + f"TCP {port} from the global controller host." + ) from last_error diff --git a/ventis/deploy.py b/ventis/deploy.py index ed48538..51bbed7 100644 --- a/ventis/deploy.py +++ b/ventis/deploy.py @@ -23,7 +23,6 @@ def my_workflow(query: str): import json import logging import os -import sys import threading import traceback import uuid @@ -121,8 +120,12 @@ def handle_workflow(): ) thread.start() - logger.info("Queued request %s for workflow '%s' with args: %s", - request_id, fn_name, kwargs) + logger.info( + "Queued request %s for workflow '%s' with args: %s", + request_id, + fn_name, + kwargs, + ) return jsonify({"request_id": request_id}), 202 @@ -151,7 +154,9 @@ def get_status(request_id): return jsonify(response), 200 - logger.info("Deploying workflow '%s' at http://%s:%d/%s", fn_name, host, port, fn_name) + logger.info( + "Deploying workflow '%s' at http://%s:%d/%s", fn_name, host, port, fn_name + ) logger.info("Status endpoint: GET http://%s:%d/status/", host, port) app.run(host=host, port=port) diff --git a/ventis/future.py b/ventis/future.py index 3165d81..0dbc856 100644 --- a/ventis/future.py +++ b/ventis/future.py @@ -6,6 +6,7 @@ import logging import grpc + try: import ventis.ventis_context as ventis_context except ImportError: @@ -25,6 +26,7 @@ logger = logging.getLogger(__name__) + # defines the future object which will be returned by each function call class Future(object): redis = RedisClient( @@ -55,10 +57,10 @@ def __init__(self, parent, service, method, args=None): method: method to be called in service args: arguments to be passed to the method """ - - # initial value of future object + + # initial value of future object self.id = uuid.uuid4().hex - + # Grab the request_id from the thread-local context (set by deploy) self.request_id = ventis_context.get_request_id() @@ -72,17 +74,20 @@ def __init__(self, parent, service, method, args=None): self.args = args or {} self.children = [] self.consumers = [] - # For simplicity I am making a decision here, the future value only be sent back to the parent + # For simplicity I am making a decision here, the future value only be sent back to the parent # Store scalar fields in a Redis Hash: future: - self.redis.hset_multiple(self._key(), { - "id": self.id, - "request_id": self.request_id or "", - "result": "", - "parent": self.parent, - "service": self.service, - "method": self.method, - "args": json.dumps(self.args), - }) + self.redis.hset_multiple( + self._key(), + { + "id": self.id, + "request_id": self.request_id or "", + "result": "", + "parent": self.parent, + "service": self.service, + "method": self.method, + "args": json.dumps(self.args), + }, + ) # Register this future under its request so local controllers can clean it up if self.request_id: self.redis.sadd(f"request:{self.request_id}:futures", self.id) @@ -94,17 +99,21 @@ def __init__(self, parent, service, method, args=None): def _submit_request(self): """Send the gRPC request to the local controller.""" stub = self._get_stub() - request_payload = json.dumps({ - "service": self.service, - "function": self.method, - "args": self.args, - "future_id": self.id, - "request_id": self.request_id, - }) + request_payload = json.dumps( + { + "service": self.service, + "function": self.method, + "args": self.args, + "future_id": self.id, + "request_id": self.request_id, + } + ) request = local_controler_pb2.JsonResponse(resonse=request_payload) try: self.response = stub.Execute(request) - logger.debug("Submitted %s.%s (future=%s)", self.service, self.method, self.id) + logger.debug( + "Submitted %s.%s (future=%s)", self.service, self.method, self.id + ) except Exception as e: logger.error("gRPC call failed for %s.%s: %s", self.service, self.method, e) raise @@ -126,7 +135,7 @@ def _poll_redis(self): result = self.redis.hget(self._key(), "result") if result is not None and result != "": self.result = result - + return self.result def value(self, timeout=None): @@ -146,7 +155,9 @@ def value(self, timeout=None): while self._poll_redis() is None and time.time() - start_time < timeout: time.sleep(0.01) if self.result is None: - raise TimeoutError(f"Future value not available within {timeout} seconds") + raise TimeoutError( + f"Future value not available within {timeout} seconds" + ) self.calculated = True # Push result to all consumers @@ -168,7 +179,6 @@ def _get_consumers(self): """Return the list of consumers from Redis.""" return self.redis.smembers(self._consumers_key()) - def _notify_consumers(self): """Push this future's result to all registered consumer endpoints via gRPC WriteResult.""" consumers = self._get_consumers() @@ -177,15 +187,26 @@ def _notify_consumers(self): for endpoint in consumers: try: if not self.result: - logger.warning("Future %s is notifying consumer %s with an empty/None result", self.id, endpoint) + logger.warning( + "Future %s is notifying consumer %s with an empty/None result", + self.id, + endpoint, + ) channel = grpc.insecure_channel(endpoint) stub = local_controler_pb2_grpc.LocalControllerStub(channel) payload = json.dumps({"future_id": self.id, "result": self.result}) request = local_controler_pb2.JsonResponse(resonse=payload) stub.WriteResult(request) - logger.info("Notified consumer %s with result for future %s", endpoint, self.id) + logger.info( + "Notified consumer %s with result for future %s", endpoint, self.id + ) except Exception as e: - logger.error("Failed to notify consumer %s for future %s: %s", endpoint, self.id, e) + logger.error( + "Failed to notify consumer %s for future %s: %s", + endpoint, + self.id, + e, + ) def _add_consumer(self, consumer): """Add a consumer.""" diff --git a/ventis/stub_generator.py b/ventis/stub_generator.py index f1cbdbe..09d8112 100644 --- a/ventis/stub_generator.py +++ b/ventis/stub_generator.py @@ -262,7 +262,9 @@ def _format_source(source): return "\n".join(formatted) + "\n" -def generate_docker(yaml_path, agent_file, output_dir=None, grpc_stubs_dir=None, stub_files=None): +def generate_docker( + yaml_path, agent_file, output_dir=None, grpc_stubs_dir=None, stub_files=None +): """ Generate a minimal Docker build context for an agent. @@ -302,8 +304,14 @@ def generate_docker(yaml_path, agent_file, output_dir=None, grpc_stubs_dir=None, (os.path.abspath(agent_file), os.path.basename(agent_file)), (os.path.join(script_dir, "future.py"), "future.py"), (os.path.join(script_dir, "ventis_context.py"), "ventis_context.py"), - (os.path.join(script_dir, "controller", "local_controller.py"), "local_controller.py"), - (os.path.join(script_dir, "controller", "local_controller_frontend.py"), "local_controller_frontend.py"), + ( + os.path.join(script_dir, "controller", "local_controller.py"), + "local_controller.py", + ), + ( + os.path.join(script_dir, "controller", "local_controller_frontend.py"), + "local_controller_frontend.py", + ), (os.path.join(script_dir, "utils", "redis_client.py"), "redis_client.py"), ] @@ -318,9 +326,7 @@ def generate_docker(yaml_path, agent_file, output_dir=None, grpc_stubs_dir=None, if os.path.isdir(grpc_stubs_dir): for fname in os.listdir(grpc_stubs_dir): if fname.endswith(".py"): - files_to_copy.append( - (os.path.join(grpc_stubs_dir, fname), fname) - ) + files_to_copy.append((os.path.join(grpc_stubs_dir, fname), fname)) for src, dst in files_to_copy: if os.path.isfile(src): @@ -329,7 +335,10 @@ def generate_docker(yaml_path, agent_file, output_dir=None, grpc_stubs_dir=None, print(f" Warning: source file not found, skipping: {src}") # Copy the YAML definition too - shutil.copy2(os.path.abspath(yaml_path), os.path.join(output_dir, os.path.basename(yaml_path))) + shutil.copy2( + os.path.abspath(yaml_path), + os.path.join(output_dir, os.path.basename(yaml_path)), + ) # ---- Dockerfile ------------------------------------------------------ agent_basename = os.path.basename(agent_file) @@ -356,7 +365,9 @@ def generate_docker(yaml_path, agent_file, output_dir=None, grpc_stubs_dir=None, return output_dir -def generate_workflow_docker(workflow_file, stub_files, output_dir=None, grpc_stubs_dir=None): +def generate_workflow_docker( + workflow_file, stub_files, output_dir=None, grpc_stubs_dir=None +): """ Generate a Docker build context for a workflow. @@ -394,24 +405,26 @@ def generate_workflow_docker(workflow_file, stub_files, output_dir=None, grpc_st (os.path.join(script_dir, "future.py"), "future.py"), (os.path.join(script_dir, "ventis_context.py"), "ventis_context.py"), (os.path.join(script_dir, "deploy.py"), "deploy.py"), - (os.path.join(script_dir, "controller", "local_controller.py"), "local_controller.py"), - (os.path.join(script_dir, "controller", "local_controller_frontend.py"), "local_controller_frontend.py"), + ( + os.path.join(script_dir, "controller", "local_controller.py"), + "local_controller.py", + ), + ( + os.path.join(script_dir, "controller", "local_controller_frontend.py"), + "local_controller_frontend.py", + ), (os.path.join(script_dir, "utils", "redis_client.py"), "redis_client.py"), ] # Copy stub files for stub_file in stub_files: - files_to_copy.append( - (os.path.abspath(stub_file), os.path.basename(stub_file)) - ) + files_to_copy.append((os.path.abspath(stub_file), os.path.basename(stub_file))) # Copy gRPC generated stubs if they exist if os.path.isdir(grpc_stubs_dir): for fname in os.listdir(grpc_stubs_dir): if fname.endswith(".py"): - files_to_copy.append( - (os.path.join(grpc_stubs_dir, fname), fname) - ) + files_to_copy.append((os.path.join(grpc_stubs_dir, fname), fname)) for src, dst in files_to_copy: if os.path.isfile(src): @@ -446,7 +459,7 @@ def start_lc(): f.write(launcher) # ---- Dockerfile ------------------------------------------------------ - dockerfile = f"""FROM python:3.11-slim + dockerfile = """FROM python:3.11-slim WORKDIR /app @@ -482,7 +495,8 @@ def start_lc(): help="Path to the YAML agent definition file (default: examples/finance_agent.yaml)", ) parser.add_argument( - "-o", "--output", + "-o", + "--output", default=None, help="Output path for the generated stub file (default: stubs/_stub.py)", ) @@ -536,4 +550,3 @@ def start_lc(): if not args.workflow_file: parser.error("--workflow-file is required when using --workflow") generate_workflow_docker(args.workflow_file, args.stub_files) - diff --git a/ventis/templates/agents/vllm_agent.py b/ventis/templates/agents/vllm_agent.py index 42223a3..924f027 100644 --- a/ventis/templates/agents/vllm_agent.py +++ b/ventis/templates/agents/vllm_agent.py @@ -1,5 +1,6 @@ # Simple VLLM Agent Example + class VllmAgent(object): def __init__(self): self.tools = [self.generate] @@ -7,11 +8,12 @@ def __init__(self): def generate(self, prompt: str) -> str: """Generates a response using an LLM model based on the given prompt.""" print(f"VllmAgent: Received prompt: '{prompt}'") - + # Simulated LLM generation synthetic_response = f"This is an LLM generated response to: '{prompt}'" return synthetic_response + if __name__ == "__main__": agent = VllmAgent() print(agent.generate("What is the stock price?")) diff --git a/ventis/templates/config/global_controller.yaml b/ventis/templates/config/global_controller.yaml index 74167bf..473e1cc 100644 --- a/ventis/templates/config/global_controller.yaml +++ b/ventis/templates/config/global_controller.yaml @@ -3,40 +3,47 @@ agents: - name: ExampleAgent - host: localhost - port: 8051 - redis_port: 6379 replicas: 1 + redis_port: 6379 + instance_type: t2.nano resources: cpu: 1 memory: 512 entrypoint: agents/example_agent.py + provider: EC2 - name: VllmAgent - host: localhost - port: 8052 - redis_port: 6379 replicas: 1 + redis_port: 6379 resources: cpu: 2 memory: 2048 - gpu: 1 entrypoint: agents/vllm_agent.py + provider: EC2 - name: Workflow - host: localhost - port: 8050 + replicas: 1 type: workflow - api_port: 8080 redis_port: 6379 - replicas: 1 workflow_file: workflows/example_workflow.py + provider: EC2 -# Polling interval in seconds poll_interval: 5 -# Redis connection redis: host: localhost port: 6379 db: 0 + +# EC2 defaults for `provider: EC2` replicas. +# Keep them here so `config/global_controller.yaml` stays the only source of truth. +# +ec2: + # Security group must allow inbound TCP 22, 50051, and 6379 from the + # global controller host's source IP. + region: us-east-1 + ami_id: ami-0123456789abcdef0 + subnet_id: subnet-0123456789abcdef0 + security_group_ids: + - sg-0123456789abcdef0 + ssh_user: ubuntu diff --git a/ventis/utils/redis_client.py b/ventis/utils/redis_client.py index 4f97c61..f05c500 100644 --- a/ventis/utils/redis_client.py +++ b/ventis/utils/redis_client.py @@ -17,7 +17,7 @@ def get(self, key): """Get a value by key from Redis. Returns None if key does not exist.""" value = self.client.get(key) if value is not None: - return value.decode("utf-8") + return self._decode(value) return None def delete(self, *keys): @@ -42,13 +42,13 @@ def hget(self, name, field): """Get a single field from a hash. Returns None if field does not exist.""" value = self.client.hget(name, field) if value is not None: - return value.decode("utf-8") + return self._decode(value) return None def hgetall(self, name): """Get all fields and values from a hash.""" data = self.client.hgetall(name) - return {k.decode("utf-8"): v.decode("utf-8") for k, v in data.items()} + return {self._decode(k): self._decode(v) for k, v in data.items()} # --- Set operations --- @@ -62,7 +62,7 @@ def srem(self, name, *values): def smembers(self, name): """Get all members of a set.""" - return {v.decode("utf-8") for v in self.client.smembers(name)} + return {self._decode(v) for v in self.client.smembers(name)} # --- Scan operations --- @@ -72,7 +72,14 @@ def scan_keys(self, pattern): cursor = 0 while True: cursor, batch = self.client.scan(cursor, match=pattern, count=100) - keys.extend(k.decode("utf-8") for k in batch) + keys.extend(self._decode(k) for k in batch) if cursor == 0: break return keys + + # --- Helper --- + + def _decode(self, value: bytes | str) -> str: + if isinstance(value, bytes): + return value.decode("utf-8") + return value diff --git a/ventis/ventis_context.py b/ventis/ventis_context.py index c6fa76a..2a81b8b 100644 --- a/ventis/ventis_context.py +++ b/ventis/ventis_context.py @@ -3,10 +3,12 @@ # Thread-local storage for request context _local = threading.local() + def set_request_id(request_id: str): """Set the current request ID for this thread.""" _local.request_id = request_id + def get_request_id() -> str: - """Get the current request ID for this thread, or None if not set.""" - return getattr(_local, "request_id", None) + """Get the current request ID for this thread, or an empty string if not set.""" + return getattr(_local, "request_id", "")