Fix: Migrate Deepmind Alphafold Dags to new environment.

GoogleCloudPlatform · Nov 12, 2024 · a93b330 · a93b330
1 parent 78138fe
commit a93b330
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 16 deletions.
diff --git a/datasets/deepmind/pipelines/alphafold/alphafold_dag.py b/datasets/deepmind/pipelines/alphafold/alphafold_dag.py
@@ -120,8 +120,9 @@
     create_and_run_sts_jobs_using_manifests = kubernetes_pod.KubernetesPodOperator(
         task_id="create_and_run_sts_jobs_using_manifests",
         name="create_and_run_sts_jobs_using_manifests",
-        namespace="composer",
-        service_account_name="datasets",
+        namespace="composer-user-workloads",
+        service_account_name="default",
+        config_file="/home/airflow/composer_kube_config",
         image_pull_policy="Always",
         image="{{ var.json.deepmind.alphafold.sts_jobs_generator }}",
         env_vars={
@@ -131,7 +132,6 @@
             "DESTINATION_BUCKET": "{{ var.json.deepmind.alphafold.destination_bucket }}",
             "GCP_PROJECT": "{{ var.value.gcp_project }}",
         },
-        resources={"request_memory": "128M", "request_cpu": "200m"},
     )
 
     # Copy manifests to public bucket

diff --git a/datasets/deepmind/pipelines/alphafold/pipeline.yaml b/datasets/deepmind/pipelines/alphafold/pipeline.yaml
@@ -138,8 +138,9 @@ dag:
       args:
         task_id: "create_and_run_sts_jobs_using_manifests"
         name: "create_and_run_sts_jobs_using_manifests"
-        namespace: "composer"
-        service_account_name: "datasets"
+        namespace: "composer-user-workloads"
+        service_account_name: "default"
+        config_file: "/home/airflow/composer_kube_config"
         image_pull_policy: "Always"
         image: "{{ var.json.deepmind.alphafold.sts_jobs_generator }}"
         env_vars:
@@ -148,9 +149,6 @@ dag:
           SOURCE_BUCKET: "{{ var.json.deepmind.alphafold.source_bucket }}"
           DESTINATION_BUCKET: "{{ var.json.deepmind.alphafold.destination_bucket }}"
           GCP_PROJECT: "{{ var.value.gcp_project }}"
-        resources:
-          request_memory: "128M"
-          request_cpu: "200m"
 
     - operator: "CloudDataTransferServiceGCSToGCSOperator"
       description: "Copy manifests to public bucket"

diff --git a/datasets/deepmind/pipelines/alphafold_v4/alphafold_v4_dag.py b/datasets/deepmind/pipelines/alphafold_v4/alphafold_v4_dag.py
@@ -116,8 +116,9 @@
     create_and_run_sts_jobs_using_manifests = kubernetes_pod.KubernetesPodOperator(
         task_id="create_and_run_sts_jobs_using_manifests",
         name="create_and_run_sts_jobs_using_manifests",
-        namespace="composer",
-        service_account_name="datasets",
+        namespace="composer-user-workloads",
+        service_account_name="default",
+        config_file="/home/airflow/composer_kube_config",
         image_pull_policy="Always",
         image="{{ var.json.deepmind.alphafold.sts_jobs_generator }}",
         env_vars={
@@ -127,7 +128,6 @@
             "DESTINATION_BUCKET": "{{ var.json.deepmind.alphafold.destination_bucket_v4 }}",
             "GCP_PROJECT": "{{ var.value.gcp_project }}",
         },
-        resources={"request_memory": "128M", "request_cpu": "200m"},
     )
 
     # Copy manifests to public bucket

diff --git a/datasets/deepmind/pipelines/alphafold_v4/pipeline.yaml b/datasets/deepmind/pipelines/alphafold_v4/pipeline.yaml
@@ -137,8 +137,9 @@ dag:
       args:
         task_id: "create_and_run_sts_jobs_using_manifests"
         name: "create_and_run_sts_jobs_using_manifests"
-        namespace: "composer"
-        service_account_name: "datasets"
+        namespace: "composer-user-workloads"
+        service_account_name: "default"
+        config_file: "/home/airflow/composer_kube_config"
         image_pull_policy: "Always"
         image: "{{ var.json.deepmind.alphafold.sts_jobs_generator }}"
         env_vars:
@@ -147,9 +148,6 @@ dag:
           SOURCE_BUCKET: "{{ var.json.deepmind.alphafold.source_bucket }}"
           DESTINATION_BUCKET: "{{ var.json.deepmind.alphafold.destination_bucket_v4 }}"
           GCP_PROJECT: "{{ var.value.gcp_project }}"
-        resources:
-          request_memory: "128M"
-          request_cpu: "200m"
 
     - operator: "CloudDataTransferServiceGCSToGCSOperator"
       description: "Copy manifests to public bucket"