From 2738056c979f6bc317cb2a7de87fb5930ca97ad7 Mon Sep 17 00:00:00 2001
From: Abdelhak Marouane <am0089@uah.com>
Date: Thu, 1 Feb 2024 09:15:59 -0600
Subject: [PATCH] Make workers CPU and memory configurable

---
 infrastructure/ecs_services/airflow_scheduler.tf | 12 +++++++-----
 infrastructure/ecs_services/ecr.tf               |  2 --
 infrastructure/ecs_services/varaibles.tf         | 13 +++++++++++++
 infrastructure/main.tf                           |  3 +++
 infrastructure/varaibles.tf                      | 14 ++++++++++++++
 5 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/infrastructure/ecs_services/airflow_scheduler.tf b/infrastructure/ecs_services/airflow_scheduler.tf
index 88ea7a3..594ac68 100644
--- a/infrastructure/ecs_services/airflow_scheduler.tf
+++ b/infrastructure/ecs_services/airflow_scheduler.tf
@@ -35,10 +35,11 @@ resource "aws_ssm_parameter" "airflow_ecs_cloudwatch_agent_config" {
   )
 }
 
+
 resource "aws_ecs_task_definition" "airflow_scheduler" {
   family             = "${var.prefix}-scheduler"
-  cpu                = 1024
-  memory             = 2048
+  cpu                = var.scheduler_cpu
+  memory             = var.scheduler_memory
   execution_role_arn = aws_iam_role.ecs_task_execution_role.arn
   task_role_arn      = aws_iam_role.airflow_task.arn
   network_mode       = "awsvpc"
@@ -51,8 +52,8 @@ resource "aws_ecs_task_definition" "airflow_scheduler" {
     {
       name   = "scheduler"
       image  = join(":", [aws_ecr_repository.airflow.repository_url, "latest"])
-      cpu    = 1024
-      memory = 2048
+      cpu    = var.scheduler_cpu
+      memory = var.scheduler_memory
       healthcheck = {
         command = [
           "CMD-SHELL",
@@ -114,6 +115,7 @@ resource "aws_security_group" "airflow_scheduler_service" {
   }
 }
 
+
 resource "aws_ecs_service" "airflow_scheduler" {
   name = "${var.prefix}-scheduler"
   depends_on = [ null_resource.build_ecr_image ,  aws_ecr_repository.airflow ]
@@ -126,7 +128,7 @@ resource "aws_ecs_service" "airflow_scheduler" {
   }
   deployment_maximum_percent         = 200
   deployment_minimum_healthy_percent = 100
-  desired_count                      = 1
+  desired_count                      = var.number_of_schedulers
   lifecycle {
     ignore_changes = [desired_count]
   }
diff --git a/infrastructure/ecs_services/ecr.tf b/infrastructure/ecs_services/ecr.tf
index 927f163..2b662ca 100644
--- a/infrastructure/ecs_services/ecr.tf
+++ b/infrastructure/ecs_services/ecr.tf
@@ -73,8 +73,6 @@ resource "null_resource" "build_worker_ecr_image" {
   triggers = {
     build_path_worker = sha1(join("", [for f in fileset(local.worker_build_path, "**") : filesha1("${local.worker_build_path}/${f}")]))
     dag_folder_path       = sha1(join("", [for f in fileset(local.dag_folder_path, "**") : filesha1("${local.dag_folder_path}/${f}")]))
-
-
   }
 
   provisioner "local-exec" {
diff --git a/infrastructure/ecs_services/varaibles.tf b/infrastructure/ecs_services/varaibles.tf
index e4ae190..0730f19 100644
--- a/infrastructure/ecs_services/varaibles.tf
+++ b/infrastructure/ecs_services/varaibles.tf
@@ -64,3 +64,16 @@ variable "custom_worker_policy_statement" {
 
 }
 
+variable "number_of_schedulers" {
+   type = number
+}
+
+variable "scheduler_cpu" {
+
+  type = number
+}
+variable "scheduler_memory" {
+  type = number
+
+}
+
diff --git a/infrastructure/main.tf b/infrastructure/main.tf
index efaafed..afc5556 100644
--- a/infrastructure/main.tf
+++ b/infrastructure/main.tf
@@ -74,6 +74,9 @@ module "ecs_services" {
   worker_memory                    = var.worker_memory #4096 * 2
 
   custom_worker_policy_statement = var.custom_worker_policy_statement
+  number_of_schedulers           = var.number_of_schedulers
+  scheduler_memory               = var.scheduler_memory
+  scheduler_cpu                  = var.scheduler_cpu
 }
 
 resource "null_resource" "airflow_create_airflow_user" {
diff --git a/infrastructure/varaibles.tf b/infrastructure/varaibles.tf
index f418fee..cc3245d 100644
--- a/infrastructure/varaibles.tf
+++ b/infrastructure/varaibles.tf
@@ -130,4 +130,18 @@ variable "custom_worker_policy_statement" {
   }))
   default = []
 
+}
+
+variable "number_of_schedulers" {
+  type = number
+  default = 1
+}
+
+variable "scheduler_cpu" {
+  type = number
+  default = 1024
+}
+variable "scheduler_memory" {
+  type = number
+  default = 2048
 }
\ No newline at end of file