cloud_composer_autoscale.sh

#!/usr/bin/env bash
set -e

PROJECT=[provide your gcp project id]
COMPOSER_NAME=[provide your composer environment name]
COMPOSER_LOCATION=[provide the selected composer’s location e.g. us-central]
CLUSTER_ZONE=[provide the selected composer’s zone e.g. us-central1-a]

# Enable autoscaling in node level
GKE_CLUSTER=$(gcloud composer environments describe \
${COMPOSER_NAME} \
--location ${COMPOSER_LOCATION} \
--format="value(config.gkeCluster)" \
--project ${PROJECT} | \
grep -o '[^\/]*$')

gcloud container clusters update ${GKE_CLUSTER} --enable-autoscaling \
--min-nodes 1 \
--max-nodes 10 \
--zone ${CLUSTER_ZONE} \
--node-pool=default-pool \
--project ${PROJECT}

# Patch airflow-worker to make it a Guaranteed QoS Pod
gcloud container clusters get-credentials ${GKE_CLUSTER} --zone ${CLUSTER_ZONE} --project ${PROJECT}

AIRFLOW_WORKER_NS=$(kubectl get namespaces | grep composer | cut -d ' ' -f1)

kubectl patch deployment airflow-worker -n ${AIRFLOW_WORKER_NS} --patch "$(cat composer_airflow_worker_patch.yaml)"

# Enable autoscaling in worker level
sed "s/{AIRFLOW_WORKER_NS}/${AIRFLOW_WORKER_NS}/" composer_airflow_worker_hpa.yaml > composer_airflow_worker_hpa_temp.yaml

kubectl apply -f composer_airflow_worker_hpa_temp.yaml

# Remove bottleneck from airflow configuration
gcloud composer environments update $COMPOSER_NAME \
--update-airflow-configs=core-max_active_runs_per_dag=150 \
--update-airflow-configs=core-dag_concurrency=300 \
--update-airflow-configs=core-dagbag_import_timeout=120 \
--update-airflow-configs=core-parallelism=300 \
--location $COMPOSER_LOCATION \
--project $PROJECT || true

printf "\nAutoscale Enabled in composer environment: $COMPOSER_NAME project: $PROJECT\n"