parameterize scripts for autopilot

cloudbuild-ap.yaml

@@ -55,15 +55,15 @@ steps:
 
         # Get kube config
         gcloud container clusters get-credentials \
-        github-ci-autopilot \
+        $_CLUSTER_NAME \
         --location us-central1 \
         --project $PROJECT_ID 
 
         cd /workspace/applications/ray/
         terraform apply \
         -var-file=workloads.tfvars \
         -var=project_id=$PROJECT_ID \
-        -var=cluster_name=github-ci-autopilot \
+        -var=cluster_name=$_CLUSTER_NAME \
         -var=cluster_location=$_REGION \
         -var=ray_namespace=ml-$SHORT_SHA-ap \
         -var=gcp_service_account=ray-sa-$SHORT_SHA-ap \
@@ -95,7 +95,7 @@ steps:
         terraform destroy \
         -var-file=workloads.tfvars \
         -var=project_id=$PROJECT_ID \
-        -var=cluster_name=github-ci-autopilot \
+        -var=cluster_name=$_CLUSTER_NAME \
         -var=cluster_location=$_REGION \
         -var=ray_namespace=ml-$SHORT_SHA-ap \
         -var=gcp_service_account=ray-sa-$SHORT_SHA-ap \
@@ -114,13 +114,13 @@ steps:
         set -e
 
         cd /workspace/modules/jupyter/tests
-        python3 change_jupyter_config_autopilot.py
+        python3 change_jupyter_config.py autopilot
 
         cd /workspace/applications/jupyter
         terraform apply \
         -var-file=workloads-without-iap.example.tfvars \
         -var=project_id=$PROJECT_ID \
-        -var=cluster_name=github-ci-autopilot \
+        -var=cluster_name=$_CLUSTER_NAME \
         -var=namespace=ml-$SHORT_SHA-ap \
         -var=workload_identity_service_account=jupyter-sa-$SHORT_SHA-ap \
         -var=gcs_bucket=gke-aieco-jupyter-$SHORT_SHA-ap \
@@ -133,7 +133,7 @@ steps:
         echo "HOST URL is " $(cat /workspace/jupyterhub_host_url.txt)
 
         cd /workspace/modules/jupyter/tests
-        python3 test_hub_autopilot.py $(cat /workspace/jupyterhub_host_url.txt)
+        python3 test_hub.py $(cat /workspace/jupyterhub_host_url.txt) autopilot
         echo "pass" > /workspace/jupyterhub_test_result.txt
     allowFailure: true
 
@@ -149,7 +149,7 @@ steps:
         terraform destroy \
         -var-file=workloads-without-iap.example.tfvars \
         -var=project_id=$PROJECT_ID \
-        -var=cluster_name=github-ci-autopilot \
+        -var=cluster_name=$_CLUSTER_NAME \
         -var=namespace=ml-$SHORT_SHA-ap \
         -var=workload_identity_service_account=jupyter-sa-$SHORT_SHA-ap \
         -var=gcs_bucket=gke-aieco-jupyter-$SHORT_SHA-ap \
@@ -168,20 +168,20 @@ steps:
 
         # Get kube config
         gcloud container clusters get-credentials \
-        github-ci-autopilot \
+        $_CLUSTER_NAME \
         --location us-central1 \
         --project $PROJECT_ID 
 
         cd /workspace/modules/jupyter/tests
-        python3 change_jupyter_config_autopilot.py
+        python3 change_jupyter_config.py autopilot
 
         cd /workspace/applications/rag/
         terraform apply \
         -var-file=workloads.tfvars \
         -var=jupyter_add_auth=false \
         -var=frontend_add_auth=false \
         -var=project_id=$PROJECT_ID \
-        -var=cluster_name=github-ci-autopilot \
+        -var=cluster_name=$_CLUSTER_NAME \
         -var=kubernetes_namespace=rag-$SHORT_SHA-ap \
         -var=gcs_bucket=gke-aieco-rag-$SHORT_SHA-ap \
         -var=ray_service_account=ray-sa-$SHORT_SHA-ap \
@@ -209,7 +209,7 @@ steps:
 
         # Validate Jupyterhub: Test Hub
         cd /workspace/modules/jupyter/tests
-        python3 test_hub_autopilot.py $(cat /workspace/rag_jupyterhub_host_url.txt)
+        python3 test_hub.py $(cat /workspace/rag_jupyterhub_host_url.txt) autopilot
         echo "pass" > /workspace/rag_jupyterhub_test_result.txt
 
         # Validate RAG: Test rag frontend
@@ -235,7 +235,7 @@ steps:
         terraform destroy \
         -var-file=workloads.tfvars \
         -var=project_id=$PROJECT_ID \
-        -var=cluster_name=github-ci-autopilot \
+        -var=cluster_name=$_CLUSTER_NAME \
         -var=kubernetes_namespace=rag-$SHORT_SHA-ap \
         -var=gcs_bucket=gke-aieco-rag-$SHORT_SHA-ap \
         -var=ray_service_account=ray-sa-$SHORT_SHA-ap \
@@ -298,6 +298,7 @@ steps:
 substitutions:
   _REGION: us-central1
   _USER_NAME: github
+  _CLUSTER_NAME: github-ci-autopilot
 options:
   substitutionOption: 'ALLOW_LOOSE'
 

cloudbuild.yaml

@@ -252,27 +252,6 @@ steps:
       - |
         set -e
 
-        cd /workspace/applications/jupyter/
-        terraform destroy \
-        -var-file=workloads-without-iap.example.tfvars \
-        -var=project_id=$PROJECT_ID \
-        -var=cluster_name=ml-$SHORT_SHA-$_PR_NUMBER-cluster \
-        -var=namespace=ml-$SHORT_SHA \
-        -var=workload_identity_service_account=jupyter-sa-$SHORT_SHA \
-        -var=gcs_bucket=gke-aieco-jupyter-$SHORT_SHA \
-        -auto-approve -no-color
-
-        cd /workspace/applications/ray/
-        terraform destroy \
-        -var-file=workloads.tfvars \
-        -var=project_id=$PROJECT_ID \
-        -var=cluster_name=ml-$SHORT_SHA-$_PR_NUMBER-cluster \
-        -var=cluster_location=$_REGION \
-        -var=ray_namespace=ml-$SHORT_SHA \
-        -var=gcp_service_account=ray-sa-$SHORT_SHA \
-        -var=gcs_bucket=gke-aieco-ray-$SHORT_SHA \
-        -auto-approve -no-color
-
         cd /workspace/applications/rag/
         terraform destroy \
         -var-file=workloads.tfvars \

modules/jupyter/tests/change_jupyter_config.py

@@ -1,9 +1,17 @@
 import yaml
+import sys
 
-with open("../jupyter_config/config-selfauth.yaml", "r") as yaml_file:
+config_file = "../jupyter_config/config-selfauth.yaml"
+if len(sys.argv) == 2:
+    autopilot = (sys.argv[1] == "autopilot")
+    if autopilot:
+        config_file = "../jupyter_config/config-selfauth-autopilot.yaml"
+
+
+with open(config_file, "r") as yaml_file:
     data = yaml.safe_load(yaml_file)
 
 data["hub"]["config"]["DummyAuthenticator"]["password"] = "dummy"
 
-with open("../jupyter_config/config-selfauth.yaml", 'w') as yaml_file:
-    yaml.dump(data, yaml_file)
+with open(config_file, 'w') as yaml_file:
+    yaml.dump(data, yaml_file)

modules/jupyter/tests/test_hub.py

@@ -4,6 +4,12 @@
 
 from packaging.version import Version as V
 
+config_file = "../jupyter_config/config-selfauth.yaml"
+if len(sys.argv) == 3:
+    autopilot = (sys.argv[2] == "autopilot")
+    if autopilot:
+        config_file = "../jupyter_config/config-selfauth-autopilot.yaml"
+
 
 def test_hub_up(hub_url):
     r = requests.get(hub_url)
@@ -32,7 +38,7 @@ def test_hub_login(hub_url):
     from /jupyter_config/config.yaml. After successfully login, user will be 
     redirected to /hub/spawn.
     """
-    with open("../jupyter_config/config-selfauth.yaml", "r") as yaml_file:
+    with open(config_file, "r") as yaml_file:
         data = yaml.safe_load(yaml_file)
 
     username = data["hub"]["config"]["Authenticator"]["admin_users"][0]