diff --git a/.gitignore b/.gitignore index 86b3c9b..8afded9 100644 --- a/.gitignore +++ b/.gitignore @@ -41,4 +41,10 @@ docker_container/ AWSCLIV2.pkg .python-version -sandbox/ \ No newline at end of file +sandbox/ + +tests/run_ec2_e2e.sh +tests/test_ec2_e2e.py + +tests/run_ec2_e2e.sh +tests/test_ec2_e2e.py 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..471c6dc 100644 --- a/tests/test_runtime_ec2.py +++ b/tests/test_runtime_ec2.py @@ -1,4 +1,3 @@ -import base64 import os import sys import tempfile @@ -60,17 +59,12 @@ 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 + 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.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}, @@ -80,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, }, }, @@ -98,8 +93,8 @@ def setUp(self): def tearDown(self): self.client_patch.stop() + os.unlink(self.key_path) 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 +109,19 @@ 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_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", "provider": "EC2", @@ -127,10 +134,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/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..fdec0ba 100644 --- a/ventis/controller/cloud_provider_logic/EC2/README.md +++ b/ventis/controller/cloud_provider_logic/EC2/README.md @@ -1,19 +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, 8080, and 22 communication within the security group --Passwordless ssh set up. +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 c5554b7..fdca290 100644 --- a/ventis/controller/cloud_provider_logic/EC2/_runtime.py +++ b/ventis/controller/cloud_provider_logic/EC2/_runtime.py @@ -12,9 +12,10 @@ they can read config and reuse the controller's Docker/Redis logic. """ -import base64 import os +import shlex import socket +import stat import subprocess import time @@ -27,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", {}) @@ -40,60 +63,19 @@ 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"]) -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": [ @@ -204,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 @@ -239,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", "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} " - 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/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/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/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..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,15 +20,16 @@ agents: memory: 2048 entrypoint: agents/vllm_agent.py provider: EC2 - instance_type: t2.nano + instance_type: t3.micro - name: Workflow 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 + instance_type: t3.micro poll_interval: 5 @@ -49,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