Ec2 connection/non static ports

.gitignore

@@ -41,4 +41,10 @@ docker_container/
 AWSCLIV2.pkg
 .python-version
 
-sandbox/
+sandbox/
+
+tests/run_ec2_e2e.sh
+tests/test_ec2_e2e.py
+
+tests/run_ec2_e2e.sh
+tests/test_ec2_e2e.py

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

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"},

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()}"

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

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,

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"):

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,
                 ],

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)),

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
 """