Internal observability agent (Datadog) + Cleanup Brainstore vars (#69)

* Default to true

* Remove index validate vars. Add extra writer vars

* Add support for internal datadog agent

* Pass extra env vars for the writer

* Add back in missing var

* Uncomment so it is less confusing. Remove vars

* Point at the alias rather than the version.

Pointing at the version makes terraform apply slow. Pointing at the alias is faster since the provisioned config happens in the background.

* Comment

* Validate against the local module, not main

* Add quarantine lambda delete script

Author: mdeeks

examples/braintrust-data-plane/main.tf

@@ -4,103 +4,113 @@ module "braintrust-data-plane" {
   source = "github.com/braintrustdata/terraform-braintrust-data-plane"
   # Append '?ref=<version_tag>' to lock to a specific version of the module.
 
+  ### Examples below are shown with the module defaults. You do not have to uncomment them
+  ### unless you want to change the default value.
+  ### The default values are for production-sized deployments.
+
   # This is primarily used for tagging and naming resources in your AWS account.
   # Do not change this after deployment. RDS and S3 resources can not be renamed.
   deployment_name = "braintrust"
 
   # Add your organization name from the Braintrust UI here
   braintrust_org_name = ""
 
-  ### Service Configuration
-  # The maximum number of concurrent executions to reserve and constrain Braintrust lambdas to.
-  # If you run Braintrust in a shared account you should set these to a reasonable limit to avoid
-  # impacting other non-Braintrust Lambdas. AWS has a global shared limit of 1000 concurrent executions per account.
-  # By default these are unlimited which is ideal for dedicated AWS account.
-  # Recommended 100 to 1000 for production in a shared account.
-  # api_handler_reserved_concurrent_executions = 100
-  # ai_proxy_reserved_concurrent_executions    = 100
-
-  # The number API Handler instances to provision and keep alive. This reduces cold start times and improves latency, with some increase in cost.
-  # api_handler_provisioned_concurrency   = 0
-
   ### Postgres configuration
-  # postgres_instance_type                = "db.r8g.2xlarge"
+  # Changing this will incur a short downtime.
+  postgres_instance_type = "db.r8g.2xlarge"
 
-  # Storage size (in GB) for the RDS instance.
-  # postgres_storage_size                 = 1000
+  # Initial storage size (in GB) for the RDS instance.
+  postgres_storage_size = 1000
   # Maximum storage size (in GB) to allow the RDS instance to auto-scale to.
-  # postgres_max_storage_size             = 4000
+  postgres_max_storage_size = 10000
 
   # Storage type for the RDS instance. Recommended io2 for large production deployments.
-  # postgres_storage_type                 = "gp3"
+  postgres_storage_type = "gp3"
 
   # Storage IOPS for the RDS instance. Only applicable if storage_type is io1, io2, or gp3.
-  # Recommended 15000 for production. Default for gp3 is 3000.
-  # postgres_storage_iops                 = 10000
+  # Recommended 15000 for production.
+  postgres_storage_iops = 15000
 
   # Throughput for the RDS instance. Only applicable if storage_type is gp3.
   # Recommended 500 for production if you are using gp3. Leave blank for io1 or io2
-  # postgres_storage_throughput           = 500
+  postgres_storage_throughput = 500
 
   # PostgreSQL engine version for the RDS instance.
-  # postgres_version                      = "15"
+  postgres_version = "15"
 
   # Automatic upgrades of PostgreSQL minor engine version.
   # If true, AWS will automatically upgrade the minor version of the PostgreSQL engine for you.
   # Note: Don't include the minor version in your postgres_version if you want to use this.
   # If false, you will need to manually upgrade the minor version of the PostgreSQL engine.
-  # postgres_auto_minor_version_upgrade   = true
+  postgres_auto_minor_version_upgrade = true
 
   # Multi-AZ RDS instance. Enabling increases cost but provides higher availability.
-  # Recommended for critical production environments.
-  # postgres_multi_az                     = true
+  # Recommended for critical production environments. Doubles the cost of the RDS instance.
+  # postgres_multi_az                     = false
 
   ### Brainstore configuration
   # The license key for the Brainstore instance. You can get this from the Braintrust UI in Settings > API URL.
   brainstore_license_key = var.brainstore_license_key
 
   # The number of Brainstore reader instances to provision
   # Recommended Graviton instance type with 16GB of memory
-  # brainstore_instance_count = 2
-  #  brainstore_instance_type  = "c8gd.4xlarge"
+  brainstore_instance_count = 2
+  brainstore_instance_type  = "c8gd.4xlarge"
 
   # The number of dedicated Brainstore writer nodes to create
   # Recommended Graviton instance type with 32GB of memory
-  # brainstore_writer_instance_count = 1
-  # brainstore_writer_instance_type  = "c8gd.8xlarge"
+  brainstore_writer_instance_count = 1
+  brainstore_writer_instance_type  = "c8gd.8xlarge"
+
 
   ### Redis configuration
+
   # Default is acceptable for typical production deployments.
-  # redis_instance_type                   = "cache.t4g.medium"
+  redis_instance_type = "cache.t4g.medium"
 
   # Redis engine version
-  # redis_version                         = "7.0"
+  redis_version = "7.0"
+
 
   ### Network configuration
-  # CIDR block for the VPC. You might need to adjust this so it does not conflict with any
-  # other VPC CIDR blocks you intend to peer with Braintrust
+  # WARNING: You should choose these values carefully after discussing with your networking team.
+  # Changing them after the fact is not possible and will require a complete rebuild of your Braintrust deployment.
+
+  # CIDR block for the VPC. The core Braintrust services will be deployed in this VPC.
+  # You might need to adjust this so it does not conflict with any other VPC CIDR blocks you intend to peer with Braintrust.
   # vpc_cidr                             = "10.175.0.0/21"
 
   # CIDR block for the Quarantined VPC. This is used to run user defined functions in an isolated environment.
+  # You might need to adjust this so it does not conflict with any other VPC CIDR blocks you intend to peer with Braintrust
   # quarantine_vpc_cidr                   = "10.175.8.0/21"
 
-  ### Advanced configuration
-  # List of origins to whitelist for CORS
-  # whitelisted_origins                   = []
-
-  # Custom domain name for the CloudFront distribution
-  # custom_domain                       = null
-
-  # ARN of the ACM certificate for the custom domain
-  # custom_certificate_arn              = null
 
-  # The maximum number of requests per user allowed in the time frame specified by outbound_rate_limit_window_minutes. Setting to 0 will disable rate limits
-  # outbound_rate_limit_max_requests      = 0
+  ### Advanced configuration
 
-  # The time frame in minutes over which rate per-user rate limits are accumulated
-  # outbound_rate_limit_window_minutes    = 1
+  # The maximum number of concurrent executions to reserve and constrain Braintrust lambdas to.
+  # If you run Braintrust in a dedicated account you can leave these at "-1" (unlimited).
+  # If you run Braintrust in a shared account you should set these to a reasonable limit to avoid
+  # impacting other non-Braintrust Lambdas. Recommended 100 to 1000 for production in a shared account.
+  # api_handler_reserved_concurrent_executions = -1
+  # ai_proxy_reserved_concurrent_executions    = -1
+
+  # Uncomment these to set extra environment variables for the services.
+  # Only use this when instructed to by the Braintrust team.
+  # brainstore_extra_env_vars = {}
+  #
+  # brainstore_extra_env_vars_writer = {}
+  #
+  # service_extra_env_vars = {
+  #   APIHandler               = {}
+  #   AIProxy                  = {}
+  #   CatchupETL               = {}
+  #   MigrateDatabaseFunction  = {}
+  #   QuarantineWarmupFunction = {}
+  # }
+
+
+  ### Braintrust Remote Support
 
-  ### Braintrust Support
   # Enable sharing of Cloudwatch logs with Braintrust staff
   # enable_braintrust_support_logs_access = true
 

main.tf

@@ -164,26 +164,29 @@ module "brainstore" {
   source = "./modules/brainstore"
   count  = var.enable_brainstore ? 1 : 0
 
-  deployment_name                          = var.deployment_name
-  instance_count                           = var.brainstore_instance_count
-  instance_type                            = var.brainstore_instance_type
-  instance_key_pair_name                   = var.brainstore_instance_key_pair_name
-  port                                     = var.brainstore_port
-  license_key                              = var.brainstore_license_key
-  version_override                         = var.brainstore_version_override
-  s3_bucket_retention_days                 = var.brainstore_s3_bucket_retention_days
-  extra_env_vars                           = var.brainstore_extra_env_vars
-  writer_instance_count                    = var.brainstore_writer_instance_count
-  writer_instance_type                     = var.brainstore_writer_instance_type
-  brainstore_disable_optimization_worker   = var.brainstore_disable_optimization_worker
-  brainstore_vacuum_all_objects            = var.brainstore_vacuum_all_objects
-  brainstore_enable_index_validation       = var.brainstore_enable_index_validation
-  brainstore_index_validation_only_deletes = var.brainstore_index_validation_only_deletes
-  database_host                            = module.database.postgres_database_address
-  database_port                            = module.database.postgres_database_port
-  database_secret_arn                      = module.database.postgres_database_secret_arn
-  redis_host                               = module.redis.redis_endpoint
-  redis_port                               = module.redis.redis_port
+  deployment_name                        = var.deployment_name
+  instance_count                         = var.brainstore_instance_count
+  instance_type                          = var.brainstore_instance_type
+  instance_key_pair_name                 = var.brainstore_instance_key_pair_name
+  port                                   = var.brainstore_port
+  license_key                            = var.brainstore_license_key
+  version_override                       = var.brainstore_version_override
+  s3_bucket_retention_days               = var.brainstore_s3_bucket_retention_days
+  extra_env_vars                         = var.brainstore_extra_env_vars
+  extra_env_vars_writer                  = var.brainstore_extra_env_vars_writer
+  writer_instance_count                  = var.brainstore_writer_instance_count
+  writer_instance_type                   = var.brainstore_writer_instance_type
+  brainstore_disable_optimization_worker = var.brainstore_disable_optimization_worker
+  brainstore_vacuum_all_objects          = var.brainstore_vacuum_all_objects
+  database_host                          = module.database.postgres_database_address
+  database_port                          = module.database.postgres_database_port
+  database_secret_arn                    = module.database.postgres_database_secret_arn
+  redis_host                             = module.redis.redis_endpoint
+  redis_port                             = module.redis.redis_port
+
+  internal_observability_api_key  = var.internal_observability_api_key
+  internal_observability_env_name = var.internal_observability_env_name
+  internal_observability_region   = var.internal_observability_region
 
   vpc_id            = module.main_vpc.vpc_id
   security_group_id = module.main_vpc.default_security_group_id

mise.toml

@@ -8,8 +8,27 @@ terraform-docs = "latest"
 [tasks]
 lint = ["terraform fmt -recursive", "tflint --recursive"]
 setup = ["pre-commit install", "tflint --init"]
-validate = [
-    "terraform init && terraform validate",
-    "cd examples/braintrust-data-plane && terraform init && terraform validate",
-]
 precommit = ["pre-commit run --all-files"]
+
+[tasks.validate]
+description = "Validate the Terraform module and example code"
+run = """
+#!/usr/bin/env bash
+set -e
+echo "Validating module"
+terraform init
+terraform validate
+
+echo "Validating example code"
+cd examples/braintrust-data-plane
+# Override the module source to point to the local module.
+# '*_override.tf' is a lesser known native terraform feature
+trap 'rm -f main_override.tf' EXIT
+cat <<EOF> main_override.tf
+module "braintrust-data-plane" {
+  source = "../../"
+}
+EOF
+terraform init
+terraform validate
+"""

modules/brainstore/main-writer.tf

@@ -36,24 +36,25 @@ resource "aws_launch_template" "brainstore_writer" {
   }
 
   user_data = base64encode(templatefile("${path.module}/templates/user_data.sh.tpl", {
-    aws_region                               = data.aws_region.current.name
-    deployment_name                          = var.deployment_name
-    database_secret_arn                      = var.database_secret_arn
-    database_host                            = var.database_host
-    database_port                            = var.database_port
-    redis_host                               = var.redis_host
-    redis_port                               = var.redis_port
-    brainstore_port                          = var.port
-    brainstore_s3_bucket                     = aws_s3_bucket.brainstore.id
-    brainstore_license_key                   = var.license_key
-    brainstore_version_override              = var.version_override == null ? "" : var.version_override
-    brainstore_release_version               = local.brainstore_release_version
-    brainstore_disable_optimization_worker   = var.brainstore_disable_optimization_worker
-    brainstore_vacuum_all_objects            = var.brainstore_vacuum_all_objects
-    brainstore_enable_index_validation       = var.brainstore_enable_index_validation
-    brainstore_index_validation_only_deletes = var.brainstore_index_validation_only_deletes
-    is_dedicated_writer_node                 = "true"
-    extra_env_vars                           = var.extra_env_vars
+    aws_region                             = data.aws_region.current.name
+    deployment_name                        = var.deployment_name
+    database_secret_arn                    = var.database_secret_arn
+    database_host                          = var.database_host
+    database_port                          = var.database_port
+    redis_host                             = var.redis_host
+    redis_port                             = var.redis_port
+    brainstore_port                        = var.port
+    brainstore_s3_bucket                   = aws_s3_bucket.brainstore.id
+    brainstore_license_key                 = var.license_key
+    brainstore_version_override            = var.version_override == null ? "" : var.version_override
+    brainstore_release_version             = local.brainstore_release_version
+    brainstore_disable_optimization_worker = var.brainstore_disable_optimization_worker
+    brainstore_vacuum_all_objects          = var.brainstore_vacuum_all_objects
+    is_dedicated_writer_node               = "true"
+    extra_env_vars                         = var.extra_env_vars_writer
+    internal_observability_api_key         = var.internal_observability_api_key
+    internal_observability_env_name        = var.internal_observability_env_name
+    internal_observability_region          = var.internal_observability_region
   }))
 
   tags = merge({

modules/brainstore/main.tf

@@ -56,12 +56,13 @@ resource "aws_launch_template" "brainstore" {
     brainstore_version_override = var.version_override == null ? "" : var.version_override
     brainstore_release_version  = local.brainstore_release_version
     # Important note: if there are no dedicated writer nodes, this node serves as a read/writer node
-    brainstore_disable_optimization_worker   = local.has_writer_nodes ? true : var.brainstore_disable_optimization_worker
-    brainstore_vacuum_all_objects            = local.has_writer_nodes ? false : var.brainstore_vacuum_all_objects
-    brainstore_enable_index_validation       = var.brainstore_enable_index_validation
-    brainstore_index_validation_only_deletes = var.brainstore_index_validation_only_deletes
-    is_dedicated_writer_node                 = "false"
-    extra_env_vars                           = var.extra_env_vars
+    brainstore_disable_optimization_worker = local.has_writer_nodes ? true : var.brainstore_disable_optimization_worker
+    brainstore_vacuum_all_objects          = local.has_writer_nodes ? false : var.brainstore_vacuum_all_objects
+    is_dedicated_writer_node               = "false"
+    extra_env_vars                         = var.extra_env_vars
+    internal_observability_api_key         = var.internal_observability_api_key
+    internal_observability_env_name        = var.internal_observability_env_name
+    internal_observability_region          = var.internal_observability_region
   }))
 
   tags = merge({

modules/brainstore/templates/user_data.sh.tpl

@@ -104,8 +104,6 @@ BRAINSTORE_CACHE_DIR=/mnt/tmp/brainstore
 BRAINSTORE_LICENSE_KEY=${brainstore_license_key}
 BRAINSTORE_DISABLE_OPTIMIZATION_WORKER=${brainstore_disable_optimization_worker}
 BRAINSTORE_VACUUM_OBJECT_ALL=${brainstore_vacuum_all_objects}
-BRAINSTORE_INDEX_WRITER_VALIDATE=${brainstore_enable_index_validation}
-BRAINSTORE_INDEX_WRITER_VALIDATE_ONLY_DELETES=${brainstore_index_validation_only_deletes}
 NO_COLOR=1
 AWS_DEFAULT_REGION=${aws_region}
 AWS_REGION=${aws_region}
@@ -119,6 +117,42 @@ if [ "${is_dedicated_writer_node}" = "true" ]; then
   echo '0 * * * * root /usr/bin/docker restart brainstore > /var/log/brainstore-restart.log 2>&1' > /etc/cron.d/restart-brainstore
 fi
 
+if [ -n "${internal_observability_api_key}" ]; then
+  if [ -n "${internal_observability_env_name}" ]; then
+    export DD_ENV="${internal_observability_env_name}"
+  fi
+  # Install Datadog Agent
+  export DD_API_KEY="${internal_observability_api_key}"
+  export DD_SITE="${internal_observability_region}.datadoghq.com"
+  export DD_APM_INSTRUMENTATION_ENABLED=host
+  export DD_APM_INSTRUMENTATION_LIBRARIES=java:1,python:3,js:5,php:1,dotnet:3
+  bash -c "$(curl -L https://install.datadoghq.com/scripts/install_script_agent7.sh)"
+  usermod -a -G docker dd-agent
+
+  cat <<EOF > /etc/datadog-agent/environment
+DD_OTLP_CONFIG_RECEIVER_PROTOCOLS_HTTP_ENDPOINT=0.0.0.0:4318
+DD_COLLECT_EC2_TAGS=true
+DD_COLLECT_EC2_TAGS_USE_IMDS=true
+EOF
+  # Configure Datadog Agent to collect Docker logs
+  cat <<EOF >> /etc/datadog-agent/datadog.yaml
+logs_enabled: true
+listeners:
+    - name: docker
+config_providers:
+    - name: docker
+      polling: true
+logs_config:
+    container_collect_all: true
+EOF
+  # Configure Brainstore to send traces to Datadog
+  cat <<EOF >> /etc/brainstore.env
+BRAINSTORE_OTLP_HTTP_ENDPOINT=http://localhost:4318
+EOF
+  # Restart Datadog Agent to pick up new configuration
+  systemctl restart datadog-agent
+fi
+
 BRAINSTORE_RELEASE_VERSION=${brainstore_release_version}
 BRAINSTORE_VERSION_OVERRIDE=${brainstore_version_override}
 BRAINSTORE_VERSION=$${BRAINSTORE_VERSION_OVERRIDE:-$${BRAINSTORE_RELEASE_VERSION}}