From fe41fae16942c101b2dfd6e97b72d184913de88c Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Tue, 14 Jul 2026 14:22:17 -0700 Subject: [PATCH 1/5] Made the workflows API ports configurable from the .yaml file --- ventis/cli.py | 1 + ventis/controller/cloud_provider_logic/EC2/README.md | 5 +++-- ventis/controller/cloud_provider_logic/EC2/_runtime.py | 2 +- ventis/controller/cloud_provider_logic/Local/_runtime.py | 3 +-- ventis/controller/utils/agent_specs.py | 1 + ventis/stub_generator.py | 6 +++--- ventis/templates/config/global_controller.yaml | 1 + 7 files changed, 11 insertions(+), 8 deletions(-) diff --git a/ventis/cli.py b/ventis/cli.py index a5345f5..65a5ad5 100644 --- a/ventis/cli.py +++ b/ventis/cli.py @@ -245,6 +245,7 @@ def cmd_build(args): stub_paths, output_dir=docker_context, grpc_stubs_dir=grpc_stubs_dir, + api_port=agent_cfg.get("api_port", 8080), ) image_name = f"ventis-{agent_name.lower()}" diff --git a/ventis/controller/cloud_provider_logic/EC2/README.md b/ventis/controller/cloud_provider_logic/EC2/README.md index 1977ba4..60d04c4 100644 --- a/ventis/controller/cloud_provider_logic/EC2/README.md +++ b/ventis/controller/cloud_provider_logic/EC2/README.md @@ -3,8 +3,9 @@ What you need to run agents on EC2. - AMI with Docker installed. - Used Ubuntu base image - Python/pip too, but only needed for the global controller --Security group allowing port 50051, 6379, 8080, and 22 communication within the security group --Passwordless ssh set up. +- Security group allowing port 50051, 6379, and 22 communication within the security group +- Also need the port your workflow API's are open to allowed +- Passwordless ssh set up. Steps: diff --git a/ventis/controller/cloud_provider_logic/EC2/_runtime.py b/ventis/controller/cloud_provider_logic/EC2/_runtime.py index c5554b7..1cfaf4f 100644 --- a/ventis/controller/cloud_provider_logic/EC2/_runtime.py +++ b/ventis/controller/cloud_provider_logic/EC2/_runtime.py @@ -242,7 +242,7 @@ def _bootstrap_instance(host, spec, replica_index, cfg, redis_host, redis_port): key = os.path.expanduser("~/.ssh/ventis_ec2") port_args = ["-p", f"{CONTAINER_PORT}:{CONTAINER_PORT}"] if spec.get("type") == "workflow": - port_args += ["-p", "8080:8080"] + port_args += ["-p", f"{spec.get('api_port', 8080)}:8080"] result = subprocess.run( f"docker save {image} | ssh -o StrictHostKeyChecking=no -i {key} " diff --git a/ventis/controller/cloud_provider_logic/Local/_runtime.py b/ventis/controller/cloud_provider_logic/Local/_runtime.py index 3391de8..4cc4c8e 100644 --- a/ventis/controller/cloud_provider_logic/Local/_runtime.py +++ b/ventis/controller/cloud_provider_logic/Local/_runtime.py @@ -12,7 +12,6 @@ DEFAULT_HOST = "localhost" CONTAINER_PORT = 50051 -WORKFLOW_API_PORT = 8080 PROVIDER = "local" _controller = None @@ -81,7 +80,7 @@ def bootstrap_instance(provisioned, spec, replica_index): f"VENTIS_REDIS_PORT={spec.get('redis_port', 6379)}", ] if ctrl_type == "workflow": - cmd.extend(["-p", f"{WORKFLOW_API_PORT}:8080"]) + cmd.extend(["-p", f"{spec.get('api_port', 8080)}:8080"]) if resources.get("cpu"): cmd.extend(["--cpus", str(resources["cpu"])]) if resources.get("memory"): diff --git a/ventis/controller/utils/agent_specs.py b/ventis/controller/utils/agent_specs.py index 2b580f7..764e73a 100644 --- a/ventis/controller/utils/agent_specs.py +++ b/ventis/controller/utils/agent_specs.py @@ -14,6 +14,7 @@ def write_agent_specs(config_path, redis_client): f"agent:{name}:", { "class": spec_class, + "api_port": json.dumps(agent.get("api_port", 8080)), "resources": json.dumps(agent.get("resources", {})), "replicas": json.dumps(agent.get("replicas", 1)), "stateful": json.dumps(agent.get("stateful", False)), diff --git a/ventis/stub_generator.py b/ventis/stub_generator.py index 09d8112..83f37a1 100644 --- a/ventis/stub_generator.py +++ b/ventis/stub_generator.py @@ -366,7 +366,7 @@ def generate_docker( def generate_workflow_docker( - workflow_file, stub_files, output_dir=None, grpc_stubs_dir=None + workflow_file, stub_files, output_dir=None, grpc_stubs_dir=None, api_port=8080 ): """ Generate a Docker build context for a workflow. @@ -459,7 +459,7 @@ def start_lc(): f.write(launcher) # ---- Dockerfile ------------------------------------------------------ - dockerfile = """FROM python:3.11-slim + dockerfile = f"""FROM python:3.11-slim WORKDIR /app @@ -469,7 +469,7 @@ def start_lc(): COPY . . EXPOSE 50051 -EXPOSE 8080 +EXPOSE {api_port} CMD python workflow_launcher.py """ diff --git a/ventis/templates/config/global_controller.yaml b/ventis/templates/config/global_controller.yaml index d38dbd0..44a9b73 100644 --- a/ventis/templates/config/global_controller.yaml +++ b/ventis/templates/config/global_controller.yaml @@ -26,6 +26,7 @@ agents: replicas: 1 type: workflow redis_port: 6379 + api_port: 8080 # Only needed for workflows, defaults to 8080 if not filled workflow_file: workflows/example_workflow.py provider: EC2 instance_type: t2.nano From 0c8274383d5b29d10ec5aeec375e84ba61834e3d Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Tue, 14 Jul 2026 16:08:07 -0700 Subject: [PATCH 2/5] cleaned up some md/test descriptions --- README.md | 2 +- tests/test_runtime_ec2.py | 19 +------- .../cloud_provider_logic/EC2/README.md | 2 +- .../cloud_provider_logic/EC2/_runtime.py | 43 ------------------- 4 files changed, 4 insertions(+), 62 deletions(-) diff --git a/README.md b/README.md index 330e339..a1be390 100644 --- a/README.md +++ b/README.md @@ -75,7 +75,7 @@ Edit `config/global_controller.yaml` in your project directory to list the agent ventis build ``` #### Step 2.1 (Only if performing distributed deployment): -If you are deploying agents and tools to multiple hosts, you need to make sure that the hosts are reachable from the machine where you are running the deploy command. To enable this please ensure that you have passwordless ssh access to the hosts. A guide to enable passwordless ssh access can be found [here](https://www.redhat.com/en/blog/passwordless-ssh). +If you are deploying agents and tools to multiple hosts, make sure the hosts are reachable from the machine where you are running the deploy command and that SSH key-based access is already configured. A guide to set that up can be found [here](https://www.redhat.com/en/blog/passwordless-ssh). #### Step 3: Deploy the project diff --git a/tests/test_runtime_ec2.py b/tests/test_runtime_ec2.py index a68b9f2..de214c8 100644 --- a/tests/test_runtime_ec2.py +++ b/tests/test_runtime_ec2.py @@ -1,7 +1,5 @@ -import base64 import os import sys -import tempfile import unittest from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -60,17 +58,8 @@ def terminate_instances(self, InstanceIds): 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.client_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}, @@ -99,7 +88,6 @@ def setUp(self): def tearDown(self): self.client_patch.stop() ec2_runtime._controller = self.original_controller - ec2_runtime.DEFAULT_SSH_KEY_PATH = self.original_key_path def _make_client(self, service_name, region_name=None): self.client_calls.append( @@ -114,7 +102,7 @@ def test_aws_clients_fails_when_required_fields_are_missing(self): with self.assertRaisesRegex(ValueError, "Missing EC2 config"): ec2_runtime._aws_clients() - def test_provision_uses_userdata_and_ec2_client(self): + def test_provision_uses_ec2_client(self): spec = { "name": "Tagged", "provider": "EC2", @@ -127,10 +115,7 @@ def test_provision_uses_userdata_and_ec2_client(self): 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.assertNotIn("UserData", request) self.assertEqual( request["TagSpecifications"][0]["Tags"][0], {"Key": "Name", "Value": "ventis-Tagged-2"}, diff --git a/ventis/controller/cloud_provider_logic/EC2/README.md b/ventis/controller/cloud_provider_logic/EC2/README.md index 60d04c4..ead8e22 100644 --- a/ventis/controller/cloud_provider_logic/EC2/README.md +++ b/ventis/controller/cloud_provider_logic/EC2/README.md @@ -5,7 +5,7 @@ What you need to run agents on EC2. - Python/pip too, but only needed for the global controller - Security group allowing port 50051, 6379, and 22 communication within the security group - Also need the port your workflow API's are open to allowed -- Passwordless ssh set up. +- SSH access set up from the deploy machine. Steps: diff --git a/ventis/controller/cloud_provider_logic/EC2/_runtime.py b/ventis/controller/cloud_provider_logic/EC2/_runtime.py index 1cfaf4f..6950513 100644 --- a/ventis/controller/cloud_provider_logic/EC2/_runtime.py +++ b/ventis/controller/cloud_provider_logic/EC2/_runtime.py @@ -12,7 +12,6 @@ they can read config and reuse the controller's Docker/Redis logic. """ -import base64 import os import socket import subprocess @@ -43,57 +42,15 @@ def _aws_clients(): return cfg, boto3.client("ec2", region_name=cfg["region"]) -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": [ From 45597dca274bbe96144e93b700c99c4d4f574799 Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Tue, 14 Jul 2026 16:03:23 -0700 Subject: [PATCH 3/5] agent updates --- tests/run_ec2_e2e.sh | 23 +++++ tests/test_ec2_e2e.py | 226 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 249 insertions(+) create mode 100755 tests/run_ec2_e2e.sh create mode 100644 tests/test_ec2_e2e.py diff --git a/tests/run_ec2_e2e.sh b/tests/run_ec2_e2e.sh new file mode 100755 index 0000000..3bf2949 --- /dev/null +++ b/tests/run_ec2_e2e.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Set these values here (or override them in the invoking environment). +# These settings are intentionally explicit because this test provisions AWS resources. +: "${VENTIS_E2E_ENABLED:=1}" # Safety switch: prevents accidental real AWS provisioning. +: "${VENTIS_CONTROLLER_HOST:=44.200.194.190}" +: "${VENTIS_CONTROLLER_USER:=ubuntu}" +: "${VENTIS_CONTROLLER_PRIVATE_KEY:=~/Downloads/saakec2.pem}" +: "${VENTIS_WORKER_PRIVATE_KEY:=~/.ssh/id_ed25519}" +: "${VENTIS_AWS_REGION:=us-east-1}" +: "${VENTIS_AMI_ID:=ami-0fca83057091173ee}" +: "${VENTIS_SUBNET_ID:=subnet-084ff6499737ff809}" +: "${VENTIS_SECURITY_GROUP_IDS:=sg-0bc353412827553fd}" +: "${VENTIS_WORKER_SSH_USER:=ubuntu}" +: "${VENTIS_INSTANCE_TYPE:=t2.nano}" +export VENTIS_E2E_ENABLED VENTIS_CONTROLLER_HOST VENTIS_CONTROLLER_USER \ + VENTIS_CONTROLLER_PRIVATE_KEY VENTIS_WORKER_PRIVATE_KEY VENTIS_AWS_REGION VENTIS_AMI_ID \ + VENTIS_SUBNET_ID VENTIS_SECURITY_GROUP_IDS VENTIS_WORKER_SSH_USER \ + VENTIS_INSTANCE_TYPE + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +exec python3 "$ROOT/tests/test_ec2_e2e.py" diff --git a/tests/test_ec2_e2e.py b/tests/test_ec2_e2e.py new file mode 100644 index 0000000..8cf4bf3 --- /dev/null +++ b/tests/test_ec2_e2e.py @@ -0,0 +1,226 @@ +"""Opt-in live EC2 test for the Ventis template. + +This module intentionally has no pytest dependency beyond what the repository +already uses. It is normally run directly through ``run_ec2_e2e.sh``. +""" + +import json +import os +import shlex +import subprocess +import sys +import time +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +TIMEOUT = int(os.environ.get("VENTIS_E2E_TIMEOUT", "900")) +def settings(): + names = ( + "host", "user", "key", "worker_key", "region", "ami_id", "subnet_id", + "security_group_ids", "ssh_user", + ) + env_names = { + "host": "VENTIS_CONTROLLER_HOST", + "user": "VENTIS_CONTROLLER_USER", + "key": "VENTIS_CONTROLLER_PRIVATE_KEY", + "worker_key": "VENTIS_WORKER_PRIVATE_KEY", + "region": "VENTIS_AWS_REGION", + "ami_id": "VENTIS_AMI_ID", + "subnet_id": "VENTIS_SUBNET_ID", + "security_group_ids": "VENTIS_SECURITY_GROUP_IDS", + "ssh_user": "VENTIS_WORKER_SSH_USER", + } + result = {key: os.environ.get(env_names[key]) for key in names} + missing = [key for key, value in result.items() if not value] + if missing: + raise RuntimeError("Missing launcher settings: " + ", ".join(missing)) + result["security_group_ids"] = [x.strip() for x in result["security_group_ids"].split(",") if x.strip()] + if not result["security_group_ids"]: + raise RuntimeError("VENTIS_SECURITY_GROUP_IDS must not be empty") + result["instance_type"] = os.environ.get("VENTIS_INSTANCE_TYPE", "t2.nano") + result["key"] = os.path.expanduser(result["key"]) + if not Path(result["key"]).is_file(): + raise RuntimeError(f"Controller private key does not exist: {result['key']}") + return result + + +class Remote: + def __init__(self, cfg): + self.base = ["ssh", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=no", + "-i", cfg["key"], f"{cfg['user']}@{cfg['host']}"] + + def run(self, command, check=True, timeout=120): + try: + return subprocess.run( + self.base + [command], + text=True, + capture_output=True, + check=check, + timeout=timeout, + ) + except subprocess.CalledProcessError as exc: + details = (exc.stderr or exc.stdout or "").strip() + raise RuntimeError( + f"Remote command failed with exit status {exc.returncode}: {command}\n" + f"{details}" + ) from exc + + def output(self, command): + return self.run(command).stdout.strip() + + +def wait_until(predicate, description, timeout=TIMEOUT, interval=5): + deadline = time.time() + timeout + last = None + while time.time() < deadline: + try: + value = predicate() + if value: + return value + last = value + except Exception as exc: # transient AWS/SSH availability is expected + last = exc + time.sleep(interval) + raise AssertionError(f"Timed out waiting for {description}: {last}") + + +def records(remote, project): + code = ("import json,redis; r=redis.Redis(host='127.0.0.1',port=6379,decode_responses=True); " + "print(json.dumps([r.hgetall(k) for k in r.scan_iter('agent_instance:*')]))") + command = ( + f"cd {shlex.quote(project)} && source ../.venv/bin/activate && " + f"python -c {shlex.quote(code)}" + ) + return json.loads(remote.output(command)) + + +def main(): + if os.environ.get("VENTIS_E2E_ENABLED") != "1": + print("EC2 E2E skipped (set VENTIS_E2E_ENABLED=1 to enable).") + return 0 + cfg = settings() + remote = Remote(cfg) + temp = f"/tmp/ventis-ec2-e2e-{os.getpid()}-{int(time.time())}" + deploy_pid = None + captured = set() + try: + print("[1/8] Logging into Global Controller and checking prerequisites...", flush=True) + remote.run("command -v python3 && command -v docker && docker info >/dev/null && python3 -c 'import boto3; print(boto3.client(\"sts\").get_caller_identity()[\"Arn\"])'") + print(" Global Controller login and prerequisites succeeded.", flush=True) + + print("[2/8] Staging current source on Global Controller...", flush=True) + remote.run(f"rm -rf {shlex.quote(temp)} && mkdir -p {shlex.quote(temp)}") + archive = subprocess.Popen(["tar", "-czf", "-", "--exclude=.git", "."], cwd=ROOT, stdout=subprocess.PIPE) + ssh = subprocess.Popen(remote.base + [f"tar -xzf - -C {shlex.quote(temp)}"], stdin=archive.stdout, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=False) + archive.stdout.close(); _, err = ssh.communicate(timeout=120); archive.wait() + if ssh.returncode or archive.returncode: + raise RuntimeError(f"Could not stage source: {err.decode(errors='replace')}") + print(" Source staging succeeded.", flush=True) + + print("[3/8] Creating virtual environment and installing Ventis...", flush=True) + project = f"{temp}/project" + remote.run( + f"cd {shlex.quote(temp)} && python3 -m venv .venv && " + ". .venv/bin/activate && pip install -q -e . && " + ".venv/bin/python -m ventis.cli new-project project" + ) + print(" Virtual environment and editable install succeeded.", flush=True) + config = {"agents": [ + {"name": "ExampleAgent", "replicas": 1, "redis_port": 6379, "instance_type": cfg["instance_type"], "entrypoint": "agents/example_agent.py", "provider": "EC2"}, + {"name": "Workflow", "replicas": 1, "type": "workflow", "redis_port": 6379, "workflow_file": "workflows/example_workflow.py", "provider": "EC2", "instance_type": cfg["instance_type"]}], + "poll_interval": 5, "redis": {"host": "localhost", "port": 6379, "db": 0}, + "ec2": {"region": cfg["region"], "ami_id": cfg["ami_id"], "subnet_id": cfg["subnet_id"], "security_group_ids": cfg["security_group_ids"], "ssh_user": cfg["ssh_user"]}} + encoded = json.dumps(config) + config_code = ( + "import json,yaml; yaml.safe_dump(json.loads(" + repr(encoded) + "), " + "open(" + repr(project + "/config/global_controller.yaml") + ", 'w'))" + ) + remote.run( + f"cd {shlex.quote(project)} && ../.venv/bin/python -c " + + shlex.quote(config_code) + ) + print("[4/8] Running ventis build...", flush=True) + remote.run(f"cd {shlex.quote(project)} && ../.venv/bin/python -m ventis.cli build", timeout=TIMEOUT) + remote.run(f"cd {shlex.quote(project)} && docker image inspect ventis-exampleagent ventis-workflow >/dev/null") + print(" ventis build succeeded and both images exist.", flush=True) + + print("[5/8] Starting ventis deploy on Global Controller...", flush=True) + log = f"{temp}/deploy.log" + worker_key = cfg["worker_key"] + remote_worker_key = ( + worker_key.replace("~/", "$HOME/", 1) + if worker_key.startswith("~/") + else shlex.quote(worker_key) + ) + proc = remote.run( + f"cd {shlex.quote(project)} && export VENTIS_EC2_SSH_KEY=" + f"{remote_worker_key} && " + "nohup ../.venv/bin/python -m ventis.cli deploy " + f"{shlex.quote(log)} 2>&1 & echo $!", + timeout=30, + ) + deploy_pid = int(proc.stdout.strip().splitlines()[-1]) + print(f" ventis deploy started (PID {deploy_pid}).", flush=True) + def ready(): + found = records(remote, project) + selected = [x for x in found if x.get("agent_name") in {"ExampleAgent", "Workflow"}] + if len(selected) == 2: + captured.update(x["ec2_instance_id"] for x in selected); return selected + return False + selected = wait_until(ready, "both EC2 service records") + print(" ExampleAgent and Workflow EC2 instances are ready.", flush=True) + agent = next(x for x in selected if x["agent_name"] == "ExampleAgent") + workflow = next(x for x in selected if x["agent_name"] == "Workflow") + remote.run( + f"ssh -o StrictHostKeyChecking=no -i {remote_worker_key} " + f"{shlex.quote(cfg['ssh_user'])}@{shlex.quote(agent['host'])} " + "'sudo docker ps --format {{.Names}}' | grep -q ventis-ec2-exampleagent" + ) + print("[6/8] Verified ExampleAgent container on an EC2 worker.", flush=True) + workflow_url = f"http://{workflow['host']}:8080" + response = remote.output(f"curl -fsS -X POST {shlex.quote(workflow_url + '/main')} -H 'Content-Type: application/json' -d '{{\"name\":\"World\"}}'") + request_id = json.loads(response)["request_id"] + def done(): + data = json.loads(remote.output(f"curl -fsS {shlex.quote(workflow_url + '/status/' + request_id)}")) + return data if data.get("status") in {"done", "error"} else False + result = wait_until(done, "workflow completion", timeout=180) + assert result.get("status") == "done" and result.get("result") == {"greeting": "Hello, World! I'm the ExampleAgent."}, result + print("[7/8] Workflow completed with the expected greeting.", flush=True) + finally: + if deploy_pid: + print("[8/8] Stopping deploy and cleaning up Ventis containers...", flush=True) + remote.run(f"kill -TERM {deploy_pid}", check=False) + wait_until(lambda: remote.run(f"kill -0 {deploy_pid}", check=False).returncode != 0, "deploy exit", timeout=120) + wait_until(lambda: "ventis-redis" not in remote.output("docker ps --format '{{.Names}}'") + and "ventis-ec2" not in remote.output("docker ps --format '{{.Names}}'"), + "Ventis container cleanup", timeout=120) + print(" Ventis containers and Redis cleaned up successfully.", flush=True) + if captured: + ids = repr(sorted(captured)) + cleanup_code = ( + "import boto3; c=boto3.client('ec2', region_name=" + + repr(cfg["region"]) + + "); ids=" + ids + + "; c.terminate_instances(InstanceIds=ids)" + ) + remote.run("python3 -c " + shlex.quote(cleanup_code), check=False) + def terminated(): + check_code = ( + "import boto3,json; c=boto3.client('ec2', region_name=" + + repr(cfg["region"]) + + "); r=c.describe_instances(InstanceIds=" + ids + "); " + + "print(json.dumps(not any(i.get('State', {}).get('Name') not in " + + "{'terminated', 'shutting-down'} for x in r['Reservations'] " + + "for i in x['Instances'])))" + ) + return json.loads(remote.output("python3 -c " + shlex.quote(check_code))) + wait_until(terminated, "captured workers termination", timeout=300) + print(" Captured EC2 workers terminated successfully.", flush=True) + remote.run(f"rm -rf {shlex.quote(temp)}", check=False) + print(" Remote temporary test directory removed.", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 652262485bce5085c9c1312370605afd9a104770 Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Wed, 15 Jul 2026 15:26:20 -0700 Subject: [PATCH 4/5] removed generating keys, assuming keys are baked into AMI for EC2 --- .gitignore | 6 +- tests/test_ec2_e2e.py | 120 +++++++++++++++--- tests/test_runtime_ec2.py | 19 +++ .../cloud_provider_logic/EC2/README.md | 50 ++++++-- .../cloud_provider_logic/EC2/_runtime.py | 34 ++++- ventis/controller/global_controller.py | 9 +- .../templates/config/global_controller.yaml | 8 +- 7 files changed, 208 insertions(+), 38 deletions(-) diff --git a/.gitignore b/.gitignore index 86b3c9b..507b959 100644 --- a/.gitignore +++ b/.gitignore @@ -41,4 +41,8 @@ docker_container/ AWSCLIV2.pkg .python-version -sandbox/ \ No newline at end of file +sandbox/ + +tests/run_ec2_e2e.sh +tests/test_ec2_e2e.py + diff --git a/tests/test_ec2_e2e.py b/tests/test_ec2_e2e.py index 8cf4bf3..6afb0a2 100644 --- a/tests/test_ec2_e2e.py +++ b/tests/test_ec2_e2e.py @@ -14,10 +14,19 @@ ROOT = Path(__file__).resolve().parents[1] TIMEOUT = int(os.environ.get("VENTIS_E2E_TIMEOUT", "900")) + + def settings(): names = ( - "host", "user", "key", "worker_key", "region", "ami_id", "subnet_id", - "security_group_ids", "ssh_user", + "host", + "user", + "key", + "worker_key", + "region", + "ami_id", + "subnet_id", + "security_group_ids", + "ssh_user", ) env_names = { "host": "VENTIS_CONTROLLER_HOST", @@ -37,7 +46,6 @@ def settings(): result["security_group_ids"] = [x.strip() for x in result["security_group_ids"].split(",") if x.strip()] if not result["security_group_ids"]: raise RuntimeError("VENTIS_SECURITY_GROUP_IDS must not be empty") - result["instance_type"] = os.environ.get("VENTIS_INSTANCE_TYPE", "t2.nano") result["key"] = os.path.expanduser(result["key"]) if not Path(result["key"]).is_file(): raise RuntimeError(f"Controller private key does not exist: {result['key']}") @@ -69,7 +77,43 @@ def output(self, command): return self.run(command).stdout.strip() -def wait_until(predicate, description, timeout=TIMEOUT, interval=5): +def install_worker_key(remote, cfg): + """Install the worker key where the EC2 runtime expects it.""" + worker_key = Path(os.path.expanduser(cfg["worker_key"])) + if not worker_key.is_file(): + raise RuntimeError(f"Worker private key does not exist: {worker_key}") + + remote.run("mkdir -p ~/.ssh && chmod 700 ~/.ssh") + target = f"{cfg['user']}@{cfg['host']}:~/.ssh/ventis_ec2" + result = subprocess.run( + [ + "scp", + "-o", + "BatchMode=yes", + "-o", + "StrictHostKeyChecking=no", + "-i", + cfg["key"], + str(worker_key), + target, + ], + text=True, + capture_output=True, + check=False, + ) + if result.returncode != 0: + details = (result.stderr or result.stdout or "").strip() + raise RuntimeError(f"Could not install worker key: {details}") + remote.run("chmod 600 ~/.ssh/ventis_ec2") + + +def wait_until( + predicate, + description, + timeout=TIMEOUT, + interval=5, + fatal_exceptions=(), +): deadline = time.time() + timeout last = None while time.time() < deadline: @@ -79,6 +123,8 @@ def wait_until(predicate, description, timeout=TIMEOUT, interval=5): return value last = value except Exception as exc: # transient AWS/SSH availability is expected + if fatal_exceptions and isinstance(exc, fatal_exceptions): + raise last = exc time.sleep(interval) raise AssertionError(f"Timed out waiting for {description}: {last}") @@ -94,6 +140,15 @@ def records(remote, project): return json.loads(remote.output(command)) +def _tail_log(remote, log_path, lines=200): + result = remote.run(f"test -f {shlex.quote(log_path)} && tail -n {lines} {shlex.quote(log_path)} || true", check=False) + return (result.stdout or result.stderr or "").strip() + + +def _deploy_exited(remote, deploy_pid): + return remote.run(f"kill -0 {deploy_pid}", check=False).returncode != 0 + + def main(): if os.environ.get("VENTIS_E2E_ENABLED") != "1": print("EC2 E2E skipped (set VENTIS_E2E_ENABLED=1 to enable).") @@ -110,7 +165,14 @@ def main(): print("[2/8] Staging current source on Global Controller...", flush=True) remote.run(f"rm -rf {shlex.quote(temp)} && mkdir -p {shlex.quote(temp)}") - archive = subprocess.Popen(["tar", "-czf", "-", "--exclude=.git", "."], cwd=ROOT, stdout=subprocess.PIPE) + # A virtualenv is machine-specific. In particular, macOS/uv virtualenvs + # contain symlinks to the local interpreter, which are invalid on the + # Linux controller and can make ``.venv/bin/python3`` disappear. + archive = subprocess.Popen( + ["tar", "-czf", "-", "--exclude=.git", "--exclude=.venv", "."], + cwd=ROOT, + stdout=subprocess.PIPE, + ) ssh = subprocess.Popen(remote.base + [f"tar -xzf - -C {shlex.quote(temp)}"], stdin=archive.stdout, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=False) archive.stdout.close(); _, err = ssh.communicate(timeout=120); archive.wait() @@ -121,20 +183,24 @@ def main(): print("[3/8] Creating virtual environment and installing Ventis...", flush=True) project = f"{temp}/project" remote.run( - f"cd {shlex.quote(temp)} && python3 -m venv .venv && " + f"cd {shlex.quote(temp)} && rm -rf .venv && python3 -m venv .venv && " ". .venv/bin/activate && pip install -q -e . && " - ".venv/bin/python -m ventis.cli new-project project" + "mkdir -p project && cp -R ventis/templates/. project/" ) print(" Virtual environment and editable install succeeded.", flush=True) - config = {"agents": [ - {"name": "ExampleAgent", "replicas": 1, "redis_port": 6379, "instance_type": cfg["instance_type"], "entrypoint": "agents/example_agent.py", "provider": "EC2"}, - {"name": "Workflow", "replicas": 1, "type": "workflow", "redis_port": 6379, "workflow_file": "workflows/example_workflow.py", "provider": "EC2", "instance_type": cfg["instance_type"]}], - "poll_interval": 5, "redis": {"host": "localhost", "port": 6379, "db": 0}, - "ec2": {"region": cfg["region"], "ami_id": cfg["ami_id"], "subnet_id": cfg["subnet_id"], "security_group_ids": cfg["security_group_ids"], "ssh_user": cfg["ssh_user"]}} - encoded = json.dumps(config) config_code = ( - "import json,yaml; yaml.safe_dump(json.loads(" + repr(encoded) + "), " - "open(" + repr(project + "/config/global_controller.yaml") + ", 'w'))" + "from pathlib import Path\n" + "import yaml\n" + f"path = Path({project + '/config/global_controller.yaml'!r})\n" + "config = yaml.safe_load(path.read_text())\n" + "config['ec2'].update({\n" + f" 'region': {cfg['region']!r},\n" + f" 'ami_id': {cfg['ami_id']!r},\n" + f" 'subnet_id': {cfg['subnet_id']!r},\n" + f" 'security_group_ids': {cfg['security_group_ids']!r},\n" + f" 'ssh_user': {cfg['ssh_user']!r},\n" + "})\n" + "path.write_text(yaml.safe_dump(config, sort_keys=False))\n" ) remote.run( f"cd {shlex.quote(project)} && ../.venv/bin/python -c " @@ -146,6 +212,8 @@ def main(): print(" ventis build succeeded and both images exist.", flush=True) print("[5/8] Starting ventis deploy on Global Controller...", flush=True) + install_worker_key(remote, cfg) + print(" Worker private key installed as ~/.ssh/ventis_ec2.", flush=True) log = f"{temp}/deploy.log" worker_key = cfg["worker_key"] remote_worker_key = ( @@ -157,18 +225,29 @@ def main(): f"cd {shlex.quote(project)} && export VENTIS_EC2_SSH_KEY=" f"{remote_worker_key} && " "nohup ../.venv/bin/python -m ventis.cli deploy " - f"{shlex.quote(log)} 2>&1 & echo $!", - timeout=30, + f"{shlex.quote(log)} 2>&1 & deploy_pid=$!; " + "printf '%s\\n' \"$deploy_pid\"", + timeout=300, ) deploy_pid = int(proc.stdout.strip().splitlines()[-1]) print(f" ventis deploy started (PID {deploy_pid}).", flush=True) + def ready(): + if _deploy_exited(remote, deploy_pid): + raise RuntimeError( + "ventis deploy exited before the EC2 instances became ready:\n" + + _tail_log(remote, log) + ) found = records(remote, project) selected = [x for x in found if x.get("agent_name") in {"ExampleAgent", "Workflow"}] if len(selected) == 2: captured.update(x["ec2_instance_id"] for x in selected); return selected return False - selected = wait_until(ready, "both EC2 service records") + selected = wait_until( + ready, + "both EC2 service records", + fatal_exceptions=(RuntimeError,), + ) print(" ExampleAgent and Workflow EC2 instances are ready.", flush=True) agent = next(x for x in selected if x["agent_name"] == "ExampleAgent") workflow = next(x for x in selected if x["agent_name"] == "Workflow") @@ -217,6 +296,11 @@ def terminated(): return json.loads(remote.output("python3 -c " + shlex.quote(check_code))) wait_until(terminated, "captured workers termination", timeout=300) print(" Captured EC2 workers terminated successfully.", flush=True) + elif deploy_pid: + log_tail = _tail_log(remote, log) + if log_tail: + print(" Deploy log tail:", flush=True) + print(log_tail, flush=True) remote.run(f"rm -rf {shlex.quote(temp)}", check=False) print(" Remote temporary test directory removed.", flush=True) return 0 diff --git a/tests/test_runtime_ec2.py b/tests/test_runtime_ec2.py index de214c8..471c6dc 100644 --- a/tests/test_runtime_ec2.py +++ b/tests/test_runtime_ec2.py @@ -1,5 +1,6 @@ import os import sys +import tempfile import unittest from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -58,6 +59,10 @@ def terminate_instances(self, InstanceIds): class EC2RuntimeTests(unittest.TestCase): def setUp(self): self.original_controller = ec2_runtime._controller + key_file = tempfile.NamedTemporaryFile(delete=False) + key_file.close() + self.key_path = key_file.name + os.chmod(self.key_path, 0o600) self.fake_client = _FakeEC2Client() self.client_calls = [] self.controller = SimpleNamespace( @@ -69,6 +74,7 @@ def setUp(self): "security_group_ids": ["sg-123456"], "region": "us-east-1", "ssh_user": "ubuntu", + "ssh_private_key_path": self.key_path, "public_ip_timeout": 1, }, }, @@ -87,6 +93,7 @@ def setUp(self): def tearDown(self): self.client_patch.stop() + os.unlink(self.key_path) ec2_runtime._controller = self.original_controller def _make_client(self, service_name, region_name=None): @@ -102,6 +109,18 @@ def test_aws_clients_fails_when_required_fields_are_missing(self): with self.assertRaisesRegex(ValueError, "Missing EC2 config"): ec2_runtime._aws_clients() + def test_aws_clients_rejects_missing_ssh_private_key(self): + self.controller.config["ec2"]["ssh_private_key_path"] = "/tmp/missing-ventis-key" + + with self.assertRaisesRegex(ValueError, "does not exist"): + ec2_runtime._aws_clients() + + def test_aws_clients_rejects_insecure_ssh_private_key_permissions(self): + os.chmod(self.key_path, 0o644) + + with self.assertRaisesRegex(ValueError, "insecure permissions 0644"): + ec2_runtime._aws_clients() + def test_provision_uses_ec2_client(self): spec = { "name": "Tagged", diff --git a/ventis/controller/cloud_provider_logic/EC2/README.md b/ventis/controller/cloud_provider_logic/EC2/README.md index ead8e22..fdec0ba 100644 --- a/ventis/controller/cloud_provider_logic/EC2/README.md +++ b/ventis/controller/cloud_provider_logic/EC2/README.md @@ -1,20 +1,48 @@ What you need to run agents on EC2. -- AMI with Docker installed. - - Used Ubuntu base image - - Python/pip too, but only needed for the global controller -- Security group allowing port 50051, 6379, and 22 communication within the security group -- Also need the port your workflow API's are open to allowed -- SSH access set up from the deploy machine. +For local controller +- AMI with the following installed: + - Docker + - Public key in authorizedkeys in ~/.ssh + +For global controller +- Instance with all of the following installed: + - Both things in local controller + - Python 3.10+ + - Ventis folder + - pip requirements installed in env + - pip install -e . --break-system-packages + - Private key labeled as ventis_ec2 inside ~/.ssh + - Private key complementing local public key + + +For convenience, you can use the global controller as the +AMI for local to save some time + +AWS Specific Permissions +- IAM Role to attach to global controller instance allowing + - ec2:RunInstance,TerminateInstance,DescribeInstances,CreateTags +- SSH permissions for entering global controller +- Security group allowing port 50051, 6379, and 22 TCP communication within security group +- Also need the port your workflow API's are open to be allowed + Steps: -1. SSH into EC2 machine and clone this repo in +1. SSH into EC2 Global Controller Instance + +2. Create your app + +3. Change the agents/configs/workflow folders to suit your needs + +4. Run ventis build + ventis deploy + +For cleanup, use ventis clean to clean stubs/containers + +If encountering permission errors with the keys, run this to give key more permissions if blocked +chmod 700 ~/.ssh +chmod 600 ~/.ssh/ventis_ec2 -2. Install and activate venv -3. Create the app, or navigate to ventis/templates -4. Change the global_controller.yaml configs to suit your setup -5. Run ventis build + ventis deploy diff --git a/ventis/controller/cloud_provider_logic/EC2/_runtime.py b/ventis/controller/cloud_provider_logic/EC2/_runtime.py index 6950513..fdca290 100644 --- a/ventis/controller/cloud_provider_logic/EC2/_runtime.py +++ b/ventis/controller/cloud_provider_logic/EC2/_runtime.py @@ -13,7 +13,9 @@ """ import os +import shlex import socket +import stat import subprocess import time @@ -26,6 +28,28 @@ _controller = None +def _ssh_key_path(cfg): + """Return the configured EC2 SSH identity after validating it locally.""" + key_path = os.path.expanduser( + cfg.get("ssh_private_key_path", DEFAULT_SSH_KEY_PATH) + ) + if not os.path.isfile(key_path): + raise ValueError( + f"EC2 SSH private key does not exist: {key_path}. " + "Set ec2.ssh_private_key_path to the controller key." + ) + if not os.access(key_path, os.R_OK): + raise ValueError(f"EC2 SSH private key is not readable: {key_path}") + + mode = stat.S_IMODE(os.stat(key_path).st_mode) + if mode & 0o077: + raise ValueError( + f"EC2 SSH private key has insecure permissions {mode:04o}: {key_path}. " + "Use chmod 600 (or 400)." + ) + return key_path + + def _aws_clients(): """Return validated EC2 config and EC2 client.""" cfg = _controller.config.get("ec2", {}) @@ -39,6 +63,7 @@ def _aws_clients(): missing = [field for field in required if not cfg.get(field)] if missing: raise ValueError(f"Missing EC2 config: {', '.join(sorted(missing))}") + _ssh_key_path(cfg) return cfg, boto3.client("ec2", region_name=cfg["region"]) @@ -161,7 +186,7 @@ def _bootstrap_instance(host, spec, replica_index, cfg, redis_host, redis_port): """Run the agent container over SSH.""" ssh_user = cfg["ssh_user"] - for _ in range(30): + for _ in range(20): result = _controller._run_cmd(["true"], host, user=ssh_user) if result.returncode == 0: break @@ -196,14 +221,15 @@ def _bootstrap_instance(host, spec, replica_index, cfg, redis_host, redis_port): agent_name = spec["name"] image = f"ventis-{agent_name.lower()}" container_name = f"ventis-ec2-{agent_name.lower()}-{replica_index}" - key = os.path.expanduser("~/.ssh/ventis_ec2") + key = _ssh_key_path(cfg) port_args = ["-p", f"{CONTAINER_PORT}:{CONTAINER_PORT}"] if spec.get("type") == "workflow": port_args += ["-p", f"{spec.get('api_port', 8080)}:8080"] result = subprocess.run( - f"docker save {image} | ssh -o StrictHostKeyChecking=no -i {key} " - f"{ssh_user}@{host} 'sudo docker load'", + f"docker save {shlex.quote(image)} | ssh -o StrictHostKeyChecking=no " + f"-o IdentitiesOnly=yes -i {shlex.quote(key)} " + f"{shlex.quote(f'{ssh_user}@{host}')} 'sudo docker load'", shell=True, capture_output=True, text=True, diff --git a/ventis/controller/global_controller.py b/ventis/controller/global_controller.py index 4a266dc..81365a6 100644 --- a/ventis/controller/global_controller.py +++ b/ventis/controller/global_controller.py @@ -509,6 +509,11 @@ def _run_cmd(self, cmd, host, user=None): if is_local: return subprocess.run(cmd, capture_output=True, text=True) else: + ssh_key_path = os.path.expanduser( + self.config.get("ec2", {}).get( + "ssh_private_key_path", "~/.ssh/ventis_ec2" + ) + ) ssh_target = f"{user}@{host}" if user else host remote_cmd = " ".join(cmd) if cmd and cmd[0] == "docker": @@ -519,9 +524,11 @@ def _run_cmd(self, cmd, host, user=None): "-o", "StrictHostKeyChecking=no", "-o", + "IdentitiesOnly=yes", + "-o", "ConnectTimeout=10", "-i", - os.path.expanduser("~/.ssh/ventis_ec2"), + ssh_key_path, ssh_target, remote_cmd, ], diff --git a/ventis/templates/config/global_controller.yaml b/ventis/templates/config/global_controller.yaml index 44a9b73..2001cb4 100644 --- a/ventis/templates/config/global_controller.yaml +++ b/ventis/templates/config/global_controller.yaml @@ -5,7 +5,7 @@ agents: - name: ExampleAgent replicas: 1 redis_port: 6379 - instance_type: t2.nano + instance_type: t3.micro resources: cpu: 1 memory: 512 @@ -20,7 +20,7 @@ agents: memory: 2048 entrypoint: agents/vllm_agent.py provider: EC2 - instance_type: t2.nano + instance_type: t3.micro - name: Workflow replicas: 1 @@ -29,7 +29,7 @@ agents: api_port: 8080 # Only needed for workflows, defaults to 8080 if not filled workflow_file: workflows/example_workflow.py provider: EC2 - instance_type: t2.nano + instance_type: t3.micro poll_interval: 5 @@ -50,3 +50,5 @@ ec2: security_group_ids: - sg-0123456789abcdef0 ssh_user: ubuntu + # Private key on the global controller used to SSH to EC2 workers. + ssh_private_key_path: ~/.ssh/ventis_ec2 From efe72ea06d7346b6543e2c5572f572ddc6a3c871 Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Wed, 15 Jul 2026 15:32:21 -0700 Subject: [PATCH 5/5] removed local testing files, nothing critical included --- .gitignore | 2 + tests/run_ec2_e2e.sh | 23 ---- tests/test_ec2_e2e.py | 310 ------------------------------------------ 3 files changed, 2 insertions(+), 333 deletions(-) delete mode 100755 tests/run_ec2_e2e.sh delete mode 100644 tests/test_ec2_e2e.py diff --git a/.gitignore b/.gitignore index 507b959..8afded9 100644 --- a/.gitignore +++ b/.gitignore @@ -46,3 +46,5 @@ sandbox/ tests/run_ec2_e2e.sh tests/test_ec2_e2e.py +tests/run_ec2_e2e.sh +tests/test_ec2_e2e.py diff --git a/tests/run_ec2_e2e.sh b/tests/run_ec2_e2e.sh deleted file mode 100755 index 3bf2949..0000000 --- a/tests/run_ec2_e2e.sh +++ /dev/null @@ -1,23 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# Set these values here (or override them in the invoking environment). -# These settings are intentionally explicit because this test provisions AWS resources. -: "${VENTIS_E2E_ENABLED:=1}" # Safety switch: prevents accidental real AWS provisioning. -: "${VENTIS_CONTROLLER_HOST:=44.200.194.190}" -: "${VENTIS_CONTROLLER_USER:=ubuntu}" -: "${VENTIS_CONTROLLER_PRIVATE_KEY:=~/Downloads/saakec2.pem}" -: "${VENTIS_WORKER_PRIVATE_KEY:=~/.ssh/id_ed25519}" -: "${VENTIS_AWS_REGION:=us-east-1}" -: "${VENTIS_AMI_ID:=ami-0fca83057091173ee}" -: "${VENTIS_SUBNET_ID:=subnet-084ff6499737ff809}" -: "${VENTIS_SECURITY_GROUP_IDS:=sg-0bc353412827553fd}" -: "${VENTIS_WORKER_SSH_USER:=ubuntu}" -: "${VENTIS_INSTANCE_TYPE:=t2.nano}" -export VENTIS_E2E_ENABLED VENTIS_CONTROLLER_HOST VENTIS_CONTROLLER_USER \ - VENTIS_CONTROLLER_PRIVATE_KEY VENTIS_WORKER_PRIVATE_KEY VENTIS_AWS_REGION VENTIS_AMI_ID \ - VENTIS_SUBNET_ID VENTIS_SECURITY_GROUP_IDS VENTIS_WORKER_SSH_USER \ - VENTIS_INSTANCE_TYPE - -ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" -exec python3 "$ROOT/tests/test_ec2_e2e.py" diff --git a/tests/test_ec2_e2e.py b/tests/test_ec2_e2e.py deleted file mode 100644 index 6afb0a2..0000000 --- a/tests/test_ec2_e2e.py +++ /dev/null @@ -1,310 +0,0 @@ -"""Opt-in live EC2 test for the Ventis template. - -This module intentionally has no pytest dependency beyond what the repository -already uses. It is normally run directly through ``run_ec2_e2e.sh``. -""" - -import json -import os -import shlex -import subprocess -import sys -import time -from pathlib import Path - -ROOT = Path(__file__).resolve().parents[1] -TIMEOUT = int(os.environ.get("VENTIS_E2E_TIMEOUT", "900")) - - -def settings(): - names = ( - "host", - "user", - "key", - "worker_key", - "region", - "ami_id", - "subnet_id", - "security_group_ids", - "ssh_user", - ) - env_names = { - "host": "VENTIS_CONTROLLER_HOST", - "user": "VENTIS_CONTROLLER_USER", - "key": "VENTIS_CONTROLLER_PRIVATE_KEY", - "worker_key": "VENTIS_WORKER_PRIVATE_KEY", - "region": "VENTIS_AWS_REGION", - "ami_id": "VENTIS_AMI_ID", - "subnet_id": "VENTIS_SUBNET_ID", - "security_group_ids": "VENTIS_SECURITY_GROUP_IDS", - "ssh_user": "VENTIS_WORKER_SSH_USER", - } - result = {key: os.environ.get(env_names[key]) for key in names} - missing = [key for key, value in result.items() if not value] - if missing: - raise RuntimeError("Missing launcher settings: " + ", ".join(missing)) - result["security_group_ids"] = [x.strip() for x in result["security_group_ids"].split(",") if x.strip()] - if not result["security_group_ids"]: - raise RuntimeError("VENTIS_SECURITY_GROUP_IDS must not be empty") - result["key"] = os.path.expanduser(result["key"]) - if not Path(result["key"]).is_file(): - raise RuntimeError(f"Controller private key does not exist: {result['key']}") - return result - - -class Remote: - def __init__(self, cfg): - self.base = ["ssh", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=no", - "-i", cfg["key"], f"{cfg['user']}@{cfg['host']}"] - - def run(self, command, check=True, timeout=120): - try: - return subprocess.run( - self.base + [command], - text=True, - capture_output=True, - check=check, - timeout=timeout, - ) - except subprocess.CalledProcessError as exc: - details = (exc.stderr or exc.stdout or "").strip() - raise RuntimeError( - f"Remote command failed with exit status {exc.returncode}: {command}\n" - f"{details}" - ) from exc - - def output(self, command): - return self.run(command).stdout.strip() - - -def install_worker_key(remote, cfg): - """Install the worker key where the EC2 runtime expects it.""" - worker_key = Path(os.path.expanduser(cfg["worker_key"])) - if not worker_key.is_file(): - raise RuntimeError(f"Worker private key does not exist: {worker_key}") - - remote.run("mkdir -p ~/.ssh && chmod 700 ~/.ssh") - target = f"{cfg['user']}@{cfg['host']}:~/.ssh/ventis_ec2" - result = subprocess.run( - [ - "scp", - "-o", - "BatchMode=yes", - "-o", - "StrictHostKeyChecking=no", - "-i", - cfg["key"], - str(worker_key), - target, - ], - text=True, - capture_output=True, - check=False, - ) - if result.returncode != 0: - details = (result.stderr or result.stdout or "").strip() - raise RuntimeError(f"Could not install worker key: {details}") - remote.run("chmod 600 ~/.ssh/ventis_ec2") - - -def wait_until( - predicate, - description, - timeout=TIMEOUT, - interval=5, - fatal_exceptions=(), -): - deadline = time.time() + timeout - last = None - while time.time() < deadline: - try: - value = predicate() - if value: - return value - last = value - except Exception as exc: # transient AWS/SSH availability is expected - if fatal_exceptions and isinstance(exc, fatal_exceptions): - raise - last = exc - time.sleep(interval) - raise AssertionError(f"Timed out waiting for {description}: {last}") - - -def records(remote, project): - code = ("import json,redis; r=redis.Redis(host='127.0.0.1',port=6379,decode_responses=True); " - "print(json.dumps([r.hgetall(k) for k in r.scan_iter('agent_instance:*')]))") - command = ( - f"cd {shlex.quote(project)} && source ../.venv/bin/activate && " - f"python -c {shlex.quote(code)}" - ) - return json.loads(remote.output(command)) - - -def _tail_log(remote, log_path, lines=200): - result = remote.run(f"test -f {shlex.quote(log_path)} && tail -n {lines} {shlex.quote(log_path)} || true", check=False) - return (result.stdout or result.stderr or "").strip() - - -def _deploy_exited(remote, deploy_pid): - return remote.run(f"kill -0 {deploy_pid}", check=False).returncode != 0 - - -def main(): - if os.environ.get("VENTIS_E2E_ENABLED") != "1": - print("EC2 E2E skipped (set VENTIS_E2E_ENABLED=1 to enable).") - return 0 - cfg = settings() - remote = Remote(cfg) - temp = f"/tmp/ventis-ec2-e2e-{os.getpid()}-{int(time.time())}" - deploy_pid = None - captured = set() - try: - print("[1/8] Logging into Global Controller and checking prerequisites...", flush=True) - remote.run("command -v python3 && command -v docker && docker info >/dev/null && python3 -c 'import boto3; print(boto3.client(\"sts\").get_caller_identity()[\"Arn\"])'") - print(" Global Controller login and prerequisites succeeded.", flush=True) - - print("[2/8] Staging current source on Global Controller...", flush=True) - remote.run(f"rm -rf {shlex.quote(temp)} && mkdir -p {shlex.quote(temp)}") - # A virtualenv is machine-specific. In particular, macOS/uv virtualenvs - # contain symlinks to the local interpreter, which are invalid on the - # Linux controller and can make ``.venv/bin/python3`` disappear. - archive = subprocess.Popen( - ["tar", "-czf", "-", "--exclude=.git", "--exclude=.venv", "."], - cwd=ROOT, - stdout=subprocess.PIPE, - ) - ssh = subprocess.Popen(remote.base + [f"tar -xzf - -C {shlex.quote(temp)}"], stdin=archive.stdout, - stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=False) - archive.stdout.close(); _, err = ssh.communicate(timeout=120); archive.wait() - if ssh.returncode or archive.returncode: - raise RuntimeError(f"Could not stage source: {err.decode(errors='replace')}") - print(" Source staging succeeded.", flush=True) - - print("[3/8] Creating virtual environment and installing Ventis...", flush=True) - project = f"{temp}/project" - remote.run( - f"cd {shlex.quote(temp)} && rm -rf .venv && python3 -m venv .venv && " - ". .venv/bin/activate && pip install -q -e . && " - "mkdir -p project && cp -R ventis/templates/. project/" - ) - print(" Virtual environment and editable install succeeded.", flush=True) - config_code = ( - "from pathlib import Path\n" - "import yaml\n" - f"path = Path({project + '/config/global_controller.yaml'!r})\n" - "config = yaml.safe_load(path.read_text())\n" - "config['ec2'].update({\n" - f" 'region': {cfg['region']!r},\n" - f" 'ami_id': {cfg['ami_id']!r},\n" - f" 'subnet_id': {cfg['subnet_id']!r},\n" - f" 'security_group_ids': {cfg['security_group_ids']!r},\n" - f" 'ssh_user': {cfg['ssh_user']!r},\n" - "})\n" - "path.write_text(yaml.safe_dump(config, sort_keys=False))\n" - ) - remote.run( - f"cd {shlex.quote(project)} && ../.venv/bin/python -c " - + shlex.quote(config_code) - ) - print("[4/8] Running ventis build...", flush=True) - remote.run(f"cd {shlex.quote(project)} && ../.venv/bin/python -m ventis.cli build", timeout=TIMEOUT) - remote.run(f"cd {shlex.quote(project)} && docker image inspect ventis-exampleagent ventis-workflow >/dev/null") - print(" ventis build succeeded and both images exist.", flush=True) - - print("[5/8] Starting ventis deploy on Global Controller...", flush=True) - install_worker_key(remote, cfg) - print(" Worker private key installed as ~/.ssh/ventis_ec2.", flush=True) - log = f"{temp}/deploy.log" - worker_key = cfg["worker_key"] - remote_worker_key = ( - worker_key.replace("~/", "$HOME/", 1) - if worker_key.startswith("~/") - else shlex.quote(worker_key) - ) - proc = remote.run( - f"cd {shlex.quote(project)} && export VENTIS_EC2_SSH_KEY=" - f"{remote_worker_key} && " - "nohup ../.venv/bin/python -m ventis.cli deploy " - f"{shlex.quote(log)} 2>&1 & deploy_pid=$!; " - "printf '%s\\n' \"$deploy_pid\"", - timeout=300, - ) - deploy_pid = int(proc.stdout.strip().splitlines()[-1]) - print(f" ventis deploy started (PID {deploy_pid}).", flush=True) - - def ready(): - if _deploy_exited(remote, deploy_pid): - raise RuntimeError( - "ventis deploy exited before the EC2 instances became ready:\n" - + _tail_log(remote, log) - ) - found = records(remote, project) - selected = [x for x in found if x.get("agent_name") in {"ExampleAgent", "Workflow"}] - if len(selected) == 2: - captured.update(x["ec2_instance_id"] for x in selected); return selected - return False - selected = wait_until( - ready, - "both EC2 service records", - fatal_exceptions=(RuntimeError,), - ) - print(" ExampleAgent and Workflow EC2 instances are ready.", flush=True) - agent = next(x for x in selected if x["agent_name"] == "ExampleAgent") - workflow = next(x for x in selected if x["agent_name"] == "Workflow") - remote.run( - f"ssh -o StrictHostKeyChecking=no -i {remote_worker_key} " - f"{shlex.quote(cfg['ssh_user'])}@{shlex.quote(agent['host'])} " - "'sudo docker ps --format {{.Names}}' | grep -q ventis-ec2-exampleagent" - ) - print("[6/8] Verified ExampleAgent container on an EC2 worker.", flush=True) - workflow_url = f"http://{workflow['host']}:8080" - response = remote.output(f"curl -fsS -X POST {shlex.quote(workflow_url + '/main')} -H 'Content-Type: application/json' -d '{{\"name\":\"World\"}}'") - request_id = json.loads(response)["request_id"] - def done(): - data = json.loads(remote.output(f"curl -fsS {shlex.quote(workflow_url + '/status/' + request_id)}")) - return data if data.get("status") in {"done", "error"} else False - result = wait_until(done, "workflow completion", timeout=180) - assert result.get("status") == "done" and result.get("result") == {"greeting": "Hello, World! I'm the ExampleAgent."}, result - print("[7/8] Workflow completed with the expected greeting.", flush=True) - finally: - if deploy_pid: - print("[8/8] Stopping deploy and cleaning up Ventis containers...", flush=True) - remote.run(f"kill -TERM {deploy_pid}", check=False) - wait_until(lambda: remote.run(f"kill -0 {deploy_pid}", check=False).returncode != 0, "deploy exit", timeout=120) - wait_until(lambda: "ventis-redis" not in remote.output("docker ps --format '{{.Names}}'") - and "ventis-ec2" not in remote.output("docker ps --format '{{.Names}}'"), - "Ventis container cleanup", timeout=120) - print(" Ventis containers and Redis cleaned up successfully.", flush=True) - if captured: - ids = repr(sorted(captured)) - cleanup_code = ( - "import boto3; c=boto3.client('ec2', region_name=" - + repr(cfg["region"]) - + "); ids=" + ids - + "; c.terminate_instances(InstanceIds=ids)" - ) - remote.run("python3 -c " + shlex.quote(cleanup_code), check=False) - def terminated(): - check_code = ( - "import boto3,json; c=boto3.client('ec2', region_name=" - + repr(cfg["region"]) - + "); r=c.describe_instances(InstanceIds=" + ids + "); " - + "print(json.dumps(not any(i.get('State', {}).get('Name') not in " - + "{'terminated', 'shutting-down'} for x in r['Reservations'] " - + "for i in x['Instances'])))" - ) - return json.loads(remote.output("python3 -c " + shlex.quote(check_code))) - wait_until(terminated, "captured workers termination", timeout=300) - print(" Captured EC2 workers terminated successfully.", flush=True) - elif deploy_pid: - log_tail = _tail_log(remote, log) - if log_tail: - print(" Deploy log tail:", flush=True) - print(log_tail, flush=True) - remote.run(f"rm -rf {shlex.quote(temp)}", check=False) - print(" Remote temporary test directory removed.", flush=True) - return 0 - - -if __name__ == "__main__": - sys.exit(main())