Repository: AgnostiqHQ/covalent-ec2-plugin (language: python)

Issue:
Update EC2 plugin to provision EC2 key pair
## Overview

To fully deploy the EC2 plugin this will not require infra to be provisioned as EC2 plugin is responsible for creating it's required resources, however it will need to also provision an EC2 key pair that matches the supplied SSH key.

## Acceptance Criteria

- [ ] Upsert the SSH key provided as a EC2 key pair as part of the executor
- [ ] Update functional test pipeline to support changes 


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index fa6dae5..a25cdcd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+## Changed
+
+- Now instead of throwing an error when an ssh key doesn't exist, EC2Executor will create the key pair, link it to AWS, and then use it to connect to the instance
+
+- Updated covalent version to use a variable instead of a hard-coded value
+
 ## [0.13.1] - 2023-02-09
 
 ### Fixed
diff --git a/covalent_ec2_plugin/__init__.py b/covalent_ec2_plugin/__init__.py
index e69de29..03ac871 100644
--- a/covalent_ec2_plugin/__init__.py
+++ b/covalent_ec2_plugin/__init__.py
@@ -0,0 +1,21 @@
+# Copyright 2023 Agnostiq Inc.
+#
+# This file is part of Covalent.
+#
+# Licensed under the GNU Affero General Public License 3.0 (the "License").
+# A copy of the License may be obtained with this software package or at
+#
+#      https://www.gnu.org/licenses/agpl-3.0.en.html
+#
+# Use of this file is prohibited except in compliance with the License. Any
+# modifications or derivative works of this file must retain this copyright
+# notice, and modified files must contain a notice indicating that they have
+# been altered from the originals.
+#
+# Covalent is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE. See the License for more details.
+#
+# Relief from the License may be granted by purchasing a commercial license.
+
+from .ec2 import EC2Executor
diff --git a/covalent_ec2_plugin/ec2.py b/covalent_ec2_plugin/ec2.py
index e413e35..bd2519c 100644
--- a/covalent_ec2_plugin/ec2.py
+++ b/covalent_ec2_plugin/ec2.py
@@ -54,6 +54,12 @@
     }
 )
 
+EC2_KEYPAIR_NAME = "covalent-ec2-executor-keypair"
+EC2_SSH_DIR = "~/.ssh/covalent"
+
+# TODO: Remove this once AWSExecutor has a `covalent_version` attribute
+TEMP_COVALENT_VERSION = "0.221.1rc0"
+
 
 class EC2Executor(SSHExecutor, AWSExecutor):
     """
@@ -119,6 +125,10 @@ def __init__(
         self.vpc = vpc or get_config("executors.ec2.vpc")
         self.subnet = subnet or get_config("executors.ec2.subnet")
 
+        # TODO: Remove this once AWSExecutor has a `covalent_version` attribute
+        # Setting covalent version to be used in the EC2 instance
+        self.covalent_version = TEMP_COVALENT_VERSION
+
     async def _run_async_subprocess(self, cmd: List[str], cwd=None, log_output: bool = False):
 
         proc = await asyncio.create_subprocess_shell(
@@ -151,9 +161,7 @@ async def _run_async_subprocess(self, cmd: List[str], cwd=None, log_output: bool
         return proc, stdout, stderr
 
     def _get_tf_statefile_path(self, task_metadata: Dict) -> str:
-        state_file = os.path.join(
-            self.cache_dir, f"{task_metadata['dispatch_id']}-{task_metadata['node_id']}.tfstate"
-        )
+        state_file = f"{self.cache_dir}/ec2-{task_metadata['dispatch_id']}-{task_metadata['node_id']}.tfstate"
         return state_file
 
     async def _get_tf_output(self, var: str, state_file: str) -> str:
@@ -168,21 +176,45 @@ async def setup(self, task_metadata: Dict) -> None:
 
         """
 
-        state_file = self._get_tf_statefile_path(task_metadata)
-
         # Init Terraform (doing this in a blocking manner to avoid race conditions during init, better way would to be use asyncio
         # locks or to ensure that terraform init is run just once)
         subprocess.run(["terraform init"], cwd=self._TF_DIR, shell=True, check=True)
 
+        state_file = self._get_tf_statefile_path(task_metadata)
+
         boto_session = boto3.Session(**self.boto_session_options())
         profile = boto_session.profile_name
         region = boto_session.region_name
 
-        # moved validaiton of ssh_key_file here so SSH executor calls _validate_credentials from AWSExecutor
-        if not Path(self.ssh_key_file).expanduser().resolve().exists():
-            raise FileNotFoundError(
-                f"The SSH key file (associated with EC2 key pair) '{self.ssh_key_file}' does not exist. Please set ssh_key_file executor argument."
-            )
+        ec2 = boto_session.client("ec2")
+        self.key_name = EC2_KEYPAIR_NAME
+
+        # Create dir if it doesn't exist
+        ec2_ssh_dir = Path(EC2_SSH_DIR).expanduser().resolve()
+        ec2_ssh_dir.mkdir(parents=True, exist_ok=True)
+
+        self.ssh_key_file = str(ec2_ssh_dir / f"{self.key_name}.pem")
+
+        # Try to import the key pair/ssh key file that might've been created earlier
+        # If those don't exist, create the key pair and save the key material to the ssh_key_file
+        if not Path(self.ssh_key_file).exists():
+            try:
+                key_pair = ec2.create_key_pair(KeyName=self.key_name)
+            except Exception as e:
+                if e.response["Error"]["Code"] != "InvalidKeyPair.Duplicate":
+                    raise
+
+                app_log.warning(
+                    f"Key pair {self.key_name} already exists, deleting and creating a new one"
+                )
+                ec2.delete_key_pair(KeyName=self.key_name)
+                key_pair = ec2.create_key_pair(KeyName=self.key_name)
+
+            with open(self.ssh_key_file, "w") as f:
+                f.write(str(key_pair["KeyMaterial"]))
+
+            # Set permissions on the key file to 400
+            os.chmod(self.ssh_key_file, 0o400)
 
         # Apply Terraform Plan
         base_cmd = [
@@ -195,14 +227,12 @@ async def setup(self, task_metadata: Dict) -> None:
         self.infra_vars = [
             f"-var=aws_region={region}",
             f"-var=aws_profile={profile}",
-            "-var=name=covalent-task-{dispatch_id}-{node_id}".format(
-                dispatch_id=task_metadata["dispatch_id"],
-                node_id=task_metadata["node_id"],
-            ),
+            f"-var=name=covalent-ec2-task-{task_metadata['dispatch_id']}-{task_metadata['node_id']}",
             f"-var=instance_type={self.instance_type}",
             f"-var=disk_size={self.volume_size}",
             f"-var=key_file={self.ssh_key_file}",
             f"-var=key_name={self.key_name}",
+            f"-var=covalent_version={self.covalent_version}",
         ]
 
         if self.credentials_file:
@@ -246,3 +276,7 @@ async def teardown(self, task_metadata: Dict) -> None:
         app_log.debug(f"Running teardown Terraform command: {cmd}")
 
         await self._run_async_subprocess(cmd, cwd=self._TF_DIR, log_output=True)
+
+        # Delete the state file
+        os.remove(state_file)
+        os.remove(f"{state_file}.backup")
diff --git a/covalent_ec2_plugin/infra/main.tf b/covalent_ec2_plugin/infra/main.tf
index 2bd3e8e..7137030 100644
--- a/covalent_ec2_plugin/infra/main.tf
+++ b/covalent_ec2_plugin/infra/main.tf
@@ -38,7 +38,8 @@ data "aws_ami" "ubuntu" {
   owners = ["099720109477"]
 }
 
-resource "aws_instance" "covalent_svc_instance" {
+
+resource "aws_instance" "covalent_ec2_instance" {
 
   ami           = data.aws_ami.ubuntu.id
   instance_type = var.instance_type
@@ -79,7 +80,9 @@ resource "null_resource" "deps_install" {
       "echo \"conda activate covalent\" >> $HOME/.bashrc",
       "conda activate covalent",
       "echo 'Installing Covalent...'",
-      "pip install covalent==0.202.0",
+
+      # TODO: Update to a variable version
+      "pip install covalent==${var.covalent_version}",
       "chmod +x /tmp/script.sh",
       "sudo bash /tmp/script.sh",
       "echo ok"
@@ -90,6 +93,6 @@ resource "null_resource" "deps_install" {
     type        = "ssh"
     user        = local.username
     private_key = file(var.key_file) # Path to a valid key file
-    host        = aws_instance.covalent_svc_instance.public_ip
+    host        = aws_instance.covalent_ec2_instance.public_ip
   }
 }
diff --git a/covalent_ec2_plugin/infra/networking.tf b/covalent_ec2_plugin/infra/networking.tf
index 4c5adda..ffe87cf 100644
--- a/covalent_ec2_plugin/infra/networking.tf
+++ b/covalent_ec2_plugin/infra/networking.tf
@@ -33,7 +33,7 @@ module "vpc" {
   ]
   private_subnets = []
 
-  enable_nat_gateway   = true
+  enable_nat_gateway   = false
   single_nat_gateway   = false
   enable_dns_hostnames = true
 }
diff --git a/covalent_ec2_plugin/infra/outputs.tf b/covalent_ec2_plugin/infra/outputs.tf
index 0005852..7aeca04 100644
--- a/covalent_ec2_plugin/infra/outputs.tf
+++ b/covalent_ec2_plugin/infra/outputs.tf
@@ -1,5 +1,5 @@
 output "hostname" {
-  value       = aws_instance.covalent_svc_instance.public_dns
+  value       = aws_instance.covalent_ec2_instance.public_dns
   description = "Hostname assigned to the EC2 instance"
 }
 
@@ -7,10 +7,12 @@ output "username" {
   value       = local.username
   description = "Username used to connect to the EC2 instance"
 }
+
 output "python3_path" {
   value = var.python3_path
   description = "Python 3 path in the EC2 instance"
 }
+
 output "remote_cache" {
   value = var.remote_cache
   description = "Location of Covalent cache data and state files in the EC2 instance"
diff --git a/covalent_ec2_plugin/infra/variables.tf b/covalent_ec2_plugin/infra/variables.tf
index 4281b1f..aa81380 100644
--- a/covalent_ec2_plugin/infra/variables.tf
+++ b/covalent_ec2_plugin/infra/variables.tf
@@ -19,7 +19,7 @@
 # Relief from the License may be granted by purchasing a commercial license.
 
 variable "name" {
-  default     = "covalent-svc"
+  default     = "covalent-executor-ec2"
   description = "Name used to prefix AWS resources"
 }
 
@@ -82,3 +82,8 @@ variable "key_file" {
   default     = ""
   description = "Private key used for SSH provisioner"
 }
+
+variable "covalent_version" {
+  default     = "0.202.0"
+  description = "Covalent version to install on the EC2 instance"
+}


Test patch:
diff --git a/tests/ec2_test.py b/tests/ec2_test.py
index 2131ed3..fd014df 100644
--- a/tests/ec2_test.py
+++ b/tests/ec2_test.py
@@ -53,90 +53,109 @@ def test_init_ec2_executor(executor):
 
 
 @pytest.mark.asyncio
-async def test_teardown(executor: ec2.EC2Executor, mocker: mock, tmp_path: Path):
+async def test_setup(executor: ec2.EC2Executor, mocker: mock, tmp_path: Path):
+    """Test validation of key file and setup."""
+
+    MOCK_TF_VAR_OUTPUT = "mocked_tf_output"
+
     mock_task_metadata = {"dispatch_id": "123", "node_id": 1}
 
     state_file = tmp_path / "state.tfstate"
-    not_existing_state_file = tmp_path / "dne.tfstate"
     state_file.touch()
 
-    executor.infra_vars = ["-var='mock_var=123'"]
+    boto3_mock = mocker.patch("covalent_ec2_plugin.ec2.boto3")
+    subprocess_mock: mock.MagicMock = mocker.patch("covalent_ec2_plugin.ec2.subprocess")
 
     run_async_process_mock = mock.AsyncMock()
+    get_tf_output_mock = mock.AsyncMock()
+    get_tf_output_mock.return_value = MOCK_TF_VAR_OUTPUT
+
     mocker.patch(
         "covalent_ec2_plugin.ec2.EC2Executor._run_async_subprocess",
         side_effect=run_async_process_mock,
     )
-
-    # test failure if tfstate does not exist
     mocker.patch(
-        "covalent_ec2_plugin.ec2.EC2Executor._get_tf_statefile_path",
-        return_value=str(not_existing_state_file),
+        "covalent_ec2_plugin.ec2.EC2Executor._get_tf_output",
+        side_effect=get_tf_output_mock,
     )
-    with pytest.raises(FileNotFoundError):
-        await executor.teardown(mock_task_metadata)
 
     mocker.patch(
         "covalent_ec2_plugin.ec2.EC2Executor._get_tf_statefile_path", return_value=str(state_file)
     )
-    await executor.teardown(mock_task_metadata)
 
-    run_async_process_mock.assert_called_once()
+    ec2_client_mock = boto3_mock.Session.return_value.client.return_value
 
+    mock_key_name = ec2.EC2_KEYPAIR_NAME
 
-@pytest.mark.parametrize(
-    "does_ssh_key_exist",
-    [(False), (True)],
-)
-@pytest.mark.asyncio
-async def test_setup(executor: ec2.EC2Executor, mocker: mock, does_ssh_key_exist, tmp_path: Path):
-    """Test validation of key file and setup."""
+    ec2_ssh_dir = Path(ec2.EC2_SSH_DIR).expanduser().resolve()
 
-    MOCK_TF_VAR_OUTPUT = "mocked_tf_output"
+    mocker.patch("covalent_ec2_plugin.ec2.Path.mkdir", return_value=ec2_ssh_dir)
+    mock_ssh_key_file = str(ec2_ssh_dir / f"{mock_key_name}.pem")
+
+    mocked_key_pair = {"KeyMaterial": "mocked_key_material"}
+    ec2_client_mock.create_key_pair.return_value = mocked_key_pair
+
+    os_chmod_mock = mocker.patch("covalent_ec2_plugin.ec2.os.chmod")
+
+    mocker.patch("covalent_ec2_plugin.ec2.Path.exists", return_value=False)
+
+    with mock.patch("builtins.open", mock.mock_open()) as mocked_open:
+
+        await executor.setup(mock_task_metadata)
+
+        ec2_client_mock.create_key_pair.assert_called_with(KeyName=mock_key_name)
 
+        file_handle = mocked_open()
+        file_handle.write.assert_called_with(mocked_key_pair["KeyMaterial"])
+
+    assert executor.username == MOCK_TF_VAR_OUTPUT
+    assert executor.hostname == MOCK_TF_VAR_OUTPUT
+    assert executor.remote_cache == MOCK_TF_VAR_OUTPUT
+    subprocess_mock.run.assert_called_once()
+    run_async_process_mock.assert_called_once()
+
+    os_chmod_mock.assert_called_with(mock_ssh_key_file, 0o400)
+
+
+@pytest.mark.asyncio
+async def test_teardown(executor: ec2.EC2Executor, mocker: mock, tmp_path: Path):
     mock_task_metadata = {"dispatch_id": "123", "node_id": 1}
-    ssh_key_file = tmp_path / "ssh_key"
-    executor.ssh_key_file = str(ssh_key_file)
 
     state_file = tmp_path / "state.tfstate"
+    not_existing_state_file = str(tmp_path / "dne.tfstate")
+
     state_file.touch()
+    state_file = str(state_file)
 
-    boto3_mock = mocker.patch("covalent_ec2_plugin.ec2.boto3")
-    subprocess_mock: mock.MagicMock = mocker.patch("covalent_ec2_plugin.ec2.subprocess")
+    executor.infra_vars = ["-var='mock_var=123'"]
 
     run_async_process_mock = mock.AsyncMock()
-    get_tf_output_mock = mock.AsyncMock()
-    get_tf_output_mock.return_value = MOCK_TF_VAR_OUTPUT
-
     mocker.patch(
         "covalent_ec2_plugin.ec2.EC2Executor._run_async_subprocess",
         side_effect=run_async_process_mock,
     )
+
+    # test failure if tfstate does not exist
     mocker.patch(
-        "covalent_ec2_plugin.ec2.EC2Executor._get_tf_output",
-        side_effect=get_tf_output_mock,
+        "covalent_ec2_plugin.ec2.EC2Executor._get_tf_statefile_path",
+        return_value=not_existing_state_file,
     )
+    with pytest.raises(FileNotFoundError):
+        await executor.teardown(mock_task_metadata)
 
     mocker.patch(
-        "covalent_ec2_plugin.ec2.EC2Executor._get_tf_statefile_path", return_value=str(state_file)
+        "covalent_ec2_plugin.ec2.EC2Executor._get_tf_statefile_path",
+        return_value=state_file,
     )
 
-    if does_ssh_key_exist:
-        ssh_key_file.touch()
-
-    if does_ssh_key_exist:
-        try:
-            await executor.setup(mock_task_metadata)
-            assert executor.username == MOCK_TF_VAR_OUTPUT
-            assert executor.hostname == MOCK_TF_VAR_OUTPUT
-            assert executor.remote_cache == MOCK_TF_VAR_OUTPUT
-            subprocess_mock.run.assert_called_once()
-            run_async_process_mock.assert_called_once()
-        except FileNotFoundError:
-            pytest.fail("Setup should not throw an error if the ssh key file exists.")
-    else:
-        with pytest.raises(FileNotFoundError):
-            await executor.setup(mock_task_metadata)
+    mock_os_remove = mocker.patch("covalent_ec2_plugin.ec2.os.remove")
+
+    await executor.teardown(mock_task_metadata)
+
+    run_async_process_mock.assert_called_once()
+
+    mock_os_remove.assert_any_call(state_file)
+    mock_os_remove.assert_any_call(f"{state_file}.backup")
 
 
 def test_upload_task():