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"

once following line executes, it puts a $ as prefix as namespace value -
sed "s/{AIRFLOW_WORKER_NS}/${AIRFLOW_WORKER_NS}/" composer_airflow_worker_hpa.yaml > composer_airflow_worker_hpa_temp.yaml

like this - $composer-1-16-6-airflow-1-10-15-82d6yh15
and on kubectl -apply it throws following error -
Error from server (NotFound): error when creating "composer_airflow_worker_hpa_temp.yaml": namespaces "$composer-1-16-6-airflow-1-10-15-82d6yh15" not found

can you please help me fixing this?

I was able to fix it by just removing the ${} from the ${AIRFLOW_WORKER_NS} word in composer_airflow_worker_hpa.yaml.
Thank you for this awesome script anyway. :)

Awesome. Glad you are able to solve it

Hey @joshuahendinata , would you mind clarifying something?

When I manually edit the worker/scheduler .yml file, it actually goes back to the initial configs once I update our composer or airflow. Applying the changes this way, through kubectl, will the changes be permanent or the same thing will happen?

@andre-abi I have not tried checking the config after composer/airflow update but it's a good assumption that the changes will be replaced by default configuration from composer every time you update composer/airflow