feat: AWS Terraform Blueprint Improvements

- Added support for restoring from snapshot - Added snapshotting volume before destroy - Bumped default Chroma version to 0.4.12 - Added support for persistent volume mounts (in fstab) - Improved tagging - Made source ranges for the Security Group configurable - Made external Chroma port configurable - Switched to use a GH startup script. Refs: chroma-core#1172
amikos-tech · Sep 22, 2023 · ffddd63 · ffddd63
1 parent 35991bf
commit ffddd63
Show file tree

Hide file tree

Showing 3 changed files with 107 additions and 50 deletions.
diff --git a/examples/deployments/aws-terraform/README.md b/examples/deployments/aws-terraform/README.md
@@ -47,15 +47,28 @@ ssh-keygen -t RSA -b 4096 -C "Chroma AWS Key" -N "" -f ./chroma-aws && chmod 400
 Set up your Terraform variables and deploy your instance:
 
 ```bash
-export TF_VAR_AWS_ACCESS_KEY=<AWS_ACCESS_KEY> #take note of this as it must be present in all of the subsequent steps
-export TF_VAR_AWS_SECRET_ACCESS_KEY=<AWS_SECRET_ACCESS_KEY> #take note of this as it must be present in all of the subsequent steps
-export TF_ssh_public_key="./chroma-aws.pub" #path to the public key you generated above (or can be different if you want to use your own key)
-export TF_ssh_private_key="./chroma-aws" #path to the private key you generated above (or can be different if you want to use your own key) - used for formatting the Chroma data volume
-export TF_VAR_chroma_release=0.4.8 #set the chroma release to deploy
-export TF_VAR_region="us-west-1" # AWS region to deploy the chroma instance to
-export TF_VAR_public_access="true" #enable public access to the chroma instance on port 8000
-export TF_VAR_enable_auth="true" #enable basic auth for the chroma instance
-export TF_VAR_auth_type="token" #The auth type to use for the chroma instance (token or basic)
+#AWS access key
+export TF_VAR_AWS_ACCESS_KEY=<AWS_ACCESS_KEY>
+#AWS secret access key
+export TF_VAR_AWS_SECRET_ACCESS_KEY=<AWS_SECRET_ACCESS_KEY>
+#path to the public key you generated above (or can be different if you want to use your own key)
+export TF_ssh_public_key="./chroma-aws.pub"
+#path to the private key you generated above (or can be different if you want to use your own key) - used for formatting the Chroma data volume
+export TF_ssh_private_key="./chroma-aws"
+#set the chroma release to deploy
+export TF_VAR_chroma_release=0.4.12
+# AWS region to deploy the chroma instance to
+export TF_VAR_region="us-west-1"
+#enable public access to the chroma instance on port 8000
+export TF_VAR_public_access="true"
+#enable basic auth for the chroma instance
+export TF_VAR_enable_auth="true"
+#The auth type to use for the chroma instance (token or basic)
+export TF_VAR_auth_type="token"
+#optional - if you want to restore from a snapshot
+export TF_VAR_chroma_data_restore_from_snapshot_id=""
+#optional - if you want to snapshot the data volume before destroying the instance
+export TF_VAR_chroma_data_volume_snapshot_before_destroy="true"
 terraform apply -auto-approve
 ```
 > Note: Basic Auth is supported by Chroma v0.4.7+
@@ -77,60 +90,59 @@ curl -v http://$instance_public_ip:8000/api/v1/heartbeat
 
 #### 4.1 Checking Auth
 
-##### Basic
-When basic auth is enabled you can check the get the credentials from Terraform state by running:
+##### Token
+When token auth is enabled you can check the get the credentials from Terraform state by running:
 
 ```bash
-terraform output chroma_auth_basic
+terraform output chroma_auth_token
 ```
 
 You should see something of the form:
 
 ```bash
-chroma:VuA8I}QyNrm0@QLq
+PVcQ4qUUnmahXwUgAf3UuYZoMlos6MnF
 ```
 
 You can then export these credentials:
 
 ```bash
-export CHROMA_AUTH=$(terraform output chroma_auth_basic | sed 's/"//g')
+export CHROMA_AUTH=$(terraform output chroma_auth_token | sed 's/"//g')
 ```
 
 Using the credentials:
 
 ```bash
-curl -v http://$instance_public_ip:8000/api/v1/collections -u "${CHROMA_AUTH}"
+curl -v http://$instance_public_ip:8000/api/v1/collections -H "Authorization: Bearer ${CHROMA_AUTH}"
 ```
 
-> Note: Without `-u` you should be getting 401 Unauthorized response
-
-
-##### Token
-When token auth is enabled you can check the get the credentials from Terraform state by running:
+##### Basic
+When basic auth is enabled you can check the get the credentials from Terraform state by running:
 
 ```bash
-terraform output chroma_auth_token
+terraform output chroma_auth_basic
 ```
 
 You should see something of the form:
 
 ```bash
-PVcQ4qUUnmahXwUgAf3UuYZoMlos6MnF
+chroma:VuA8I}QyNrm0@QLq
 ```
 
 You can then export these credentials:
 
 ```bash
-export CHROMA_AUTH=$(terraform output chroma_auth_token | sed 's/"//g')
+export CHROMA_AUTH=$(terraform output chroma_auth_basic | sed 's/"//g')
 ```
 
 Using the credentials:
 
 ```bash
-curl -v http://$instance_public_ip:8000/api/v1/collections -H "Authorization: Bearer ${CHROMA_AUTH}"
+curl -v http://$instance_public_ip:8000/api/v1/collections -u "${CHROMA_AUTH}"
 ```
 
-#### 4.2 SSH to your instance
+> Note: Without `-u` you should be getting 401 Unauthorized response
+
+#### 4.2 Connect (ssh) to your instance
 
 
 To SSH to your instance:

diff --git a/examples/deployments/aws-terraform/chroma.tf b/examples/deployments/aws-terraform/chroma.tf
@@ -26,16 +26,16 @@ resource "aws_security_group" "chroma_sg" {
     from_port   = 22
     to_port     = 22
     protocol    = "tcp"
-    cidr_blocks = ["0.0.0.0/0"]
+    cidr_blocks = var.mgmt_source_ranges
   }
 
   dynamic "ingress" {
     for_each = var.public_access ? [1] : []
     content {
-      from_port   = 8000
+      from_port   = var.chroma_port
       to_port     = 8000
       protocol    = "tcp"
-      cidr_blocks = ["0.0.0.0/0"]
+      cidr_blocks = var.source_ranges
     }
   }
 
@@ -47,9 +47,7 @@ resource "aws_security_group" "chroma_sg" {
     ipv6_cidr_blocks = ["::/0"]
   }
 
-  tags = {
-    Name = "chroma"
-  }
+  tags = local.tags
 }
 
 resource "aws_key_pair" "chroma-keypair" {
@@ -83,17 +81,9 @@ resource "aws_instance" "chroma_instance" {
   key_name        = "chroma-keypair"
   security_groups = [aws_security_group.chroma_sg.name]
 
-  user_data = templatefile("${path.module}/startup.sh", {
-    chroma_release         = var.chroma_release,
-    enable_auth            = var.enable_auth,
-    auth_type              = var.auth_type,
-    basic_auth_credentials = "${local.basic_auth_credentials.username}:${local.basic_auth_credentials.password}",
-    token_auth_credentials = random_password.chroma_token.result,
-  })
+  user_data = data.template_file.user_data.rendered
 
-  tags = {
-    Name = "chroma"
-  }
+  tags = local.tags
 
   ebs_block_device {
     device_name = "/dev/sda1"
@@ -105,29 +95,33 @@ resource "aws_instance" "chroma_instance" {
 resource "aws_ebs_volume" "chroma-volume" {
   availability_zone = aws_instance.chroma_instance.availability_zone
   size              = var.chroma_data_volume_size
+  final_snapshot = var.chroma_data_volume_snapshot_before_destroy
+  snapshot_id = var.chroma_data_restore_from_snapshot_id
 
-  tags = {
-    Name = "chroma"
-  }
+  tags = local.tags
 
-  lifecycle {
-    prevent_destroy = var.prevent_chroma_data_volume_delete # size in GBs
-  }
 }
 
 locals {
   cleaned_volume_id = replace(aws_ebs_volume.chroma-volume.id, "-", "")
 }
 
+locals {
+  restore_from_snapshot = length(var.chroma_data_restore_from_snapshot_id) == 0 ? false : true
+}
+
 resource "aws_volume_attachment" "chroma_volume_attachment" {
   device_name = "/dev/sdh"
   volume_id   = aws_ebs_volume.chroma-volume.id
   instance_id = aws_instance.chroma_instance.id
   provisioner "remote-exec" {
     inline = [
-      "export VOLUME_ID=${local.cleaned_volume_id} && sudo mkfs -t ext4 /dev/$(lsblk -o +SERIAL | grep $VOLUME_ID | awk '{print $1}')",
+      "if [ -z \"${local.restore_from_snapshot}\"  ]; then export VOLUME_ID=${local.cleaned_volume_id} && sudo mkfs -t ext4 /dev/$(lsblk -o +SERIAL | grep $VOLUME_ID | awk '{print $1}'); fi",
       "sudo mkdir /chroma-data",
-      "export VOLUME_ID=${local.cleaned_volume_id} && sudo mount /dev/$(lsblk -o +SERIAL | grep $VOLUME_ID | awk '{print $1}') /chroma-data"
+      "export VOLUME_ID=${local.cleaned_volume_id} && sudo mount /dev/$(lsblk -o +SERIAL | grep $VOLUME_ID | awk '{print $1}') /chroma-data",
+      "export VOLUME_ID=${local.cleaned_volume_id} && cat <<EOF | sudo tee /etc/fstab >> /dev/null",
+      "/dev/$(lsblk -o +SERIAL | grep $VOLUME_ID | awk '{print $1}') /chroma-data ext4 defaults,nofail,discard 0 0",
+      "EOF",
     ]
 
     connection {

diff --git a/examples/deployments/aws-terraform/variables.tf b/examples/deployments/aws-terraform/variables.tf
@@ -1,7 +1,24 @@
 variable "chroma_release" {
   description = "The chroma release to deploy"
   type        = string
-  default     = "0.4.8"
+  default     = "0.4.12"
+}
+
+#TODO this should be updated to point to https://raw.githubusercontent.com/chroma-core/chroma/main/examples/deployments/common/startup.sh in the repo
+data "http" "startup_script_remote" {
+  url = "https://raw.githubusercontent.com/chroma-core/chroma/main/examples/deployments/aws-terraform/startup.sh"
+}
+
+data "template_file" "user_data" {
+  template = data.http.startup_script_remote.response_body
+
+  vars = {
+    chroma_release         = var.chroma_release
+    enable_auth            = var.enable_auth
+    auth_type              = var.auth_type
+    basic_auth_credentials = "${local.basic_auth_credentials.username}:${local.basic_auth_credentials.password}"
+    token_auth_credentials = random_password.chroma_token.result
+  }
 }
 
 variable "region" {
@@ -62,6 +79,10 @@ locals {
   token_auth_credentials = {
     token = random_password.chroma_token.result
   }
+  tags = [
+    "chroma",
+    "release-${replace(var.chroma_release, ".", "")}",
+  ]
 }
 
 variable "ssh_public_key" {
@@ -92,3 +113,33 @@ variable "prevent_chroma_data_volume_delete" {
     type        = bool
     default     = false
 }
+
+variable "chroma_data_volume_snapshot_before_destroy" {
+    description = "Take a snapshot of the chroma data volume before destroying it"
+    type        = bool
+    default     = false
+}
+
+variable "chroma_data_restore_from_snapshot_id" {
+    description = "Restore the chroma data volume from a snapshot"
+    type        = string
+    default     = null
+}
+
+variable "chroma_port" {
+  default     = "8000"
+  description = "The port that chroma listens on"
+  type        = string
+}
+
+variable "source_ranges" {
+  default     = ["0.0.0.0/0", "::/0"]
+  type        = list(string)
+  description = "List of CIDR ranges to allow through the firewall"
+}
+
+variable "mgmt_source_ranges" {
+  default     = ["0.0.0.0/0", "::/0"]
+  type        = list(string)
+  description = "List of CIDR ranges to allow for management of the Chroma instance. This is used for SSH incoming traffic filtering"
+}