From a93b33023f4518127c2b0fee4fec486ee8b34a02 Mon Sep 17 00:00:00 2001
From: nlarge-google <nlarge@xwf.google.com>
Date: Tue, 12 Nov 2024 20:24:39 +0000
Subject: [PATCH] Fix: Migrate Deepmind Alphafold Dags to new environment.

---
 datasets/deepmind/pipelines/alphafold/alphafold_dag.py    | 6 +++---
 datasets/deepmind/pipelines/alphafold/pipeline.yaml       | 8 +++-----
 .../deepmind/pipelines/alphafold_v4/alphafold_v4_dag.py   | 6 +++---
 datasets/deepmind/pipelines/alphafold_v4/pipeline.yaml    | 8 +++-----
 4 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/datasets/deepmind/pipelines/alphafold/alphafold_dag.py b/datasets/deepmind/pipelines/alphafold/alphafold_dag.py
index 2de49272c..366e948b2 100644
--- 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
index f4c8a3662..6c17c7e57 100644
--- 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
index 6204bb3bd..ff2f21fdc 100644
--- 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
index e718f4ee3..c417298e3 100644
--- 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"