From 2d8160a509d72c44f6d71dfbb227649f3072c072 Mon Sep 17 00:00:00 2001
From: Aftab Alam <afalam@google.com>
Date: Thu, 14 Nov 2024 06:56:22 +0000
Subject: [PATCH] Fix: Migration changes for human_variant_annotation DAG

---
 .../pipelines/clinvar/clinvar_dag.py          | 24 ++++++++++-----
 .../pipelines/clinvar/pipeline.yaml           | 30 ++++++++++++-------
 .../pipelines/db_snp/db_snp_dag.py            |  2 +-
 3 files changed, 38 insertions(+), 18 deletions(-)

diff --git a/datasets/human_variant_annotation/pipelines/clinvar/clinvar_dag.py b/datasets/human_variant_annotation/pipelines/clinvar/clinvar_dag.py
index 0ad21eda1..31d474932 100644
--- a/datasets/human_variant_annotation/pipelines/clinvar/clinvar_dag.py
+++ b/datasets/human_variant_annotation/pipelines/clinvar/clinvar_dag.py
@@ -1,4 +1,4 @@
-# Copyright 2021 Google LLC
+# Copyright 2022 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -38,8 +38,9 @@
         task_id="clinvar_vcf_grch37",
         startup_timeout_seconds=600,
         name="name_basics",
-        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.human_variant_annotation.container_registry.run_csv_transform_kub }}",
         env_vars={
@@ -50,7 +51,11 @@
             "TARGET_GCS_FOLDER": "data/human_variant_annotation/clinVar-vcf_GRCh37/",
             "PIPELINE": "clinvar",
         },
-        resources={"limit_memory": "1G", "limit_cpu": "1"},
+        container_resources={
+            "memory": {"request": "32Gi"},
+            "cpu": {"request": "2"},
+            "ephemeral-storage": {"request": "10Gi"},
+        },
     )
 
     # Task to run a GoogleCloudStorageToGoogleCloudStorageOperator
@@ -69,8 +74,9 @@
         task_id="clinvar_vcf_grch38",
         startup_timeout_seconds=600,
         name="name_basics",
-        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.human_variant_annotation.container_registry.run_csv_transform_kub }}",
         env_vars={
@@ -81,7 +87,11 @@
             "TARGET_GCS_FOLDER": "data/human_variant_annotation/clinVar-vcf_GRCh38/",
             "PIPELINE": "db_snp",
         },
-        resources={"limit_memory": "1G", "limit_cpu": "1"},
+        container_resources={
+            "memory": {"request": "32Gi"},
+            "cpu": {"request": "2"},
+            "ephemeral-storage": {"request": "10Gi"},
+        },
     )
 
     # Task to run a GoogleCloudStorageToGoogleCloudStorageOperator
diff --git a/datasets/human_variant_annotation/pipelines/clinvar/pipeline.yaml b/datasets/human_variant_annotation/pipelines/clinvar/pipeline.yaml
index 0a316cc89..df7306d0d 100644
--- a/datasets/human_variant_annotation/pipelines/clinvar/pipeline.yaml
+++ b/datasets/human_variant_annotation/pipelines/clinvar/pipeline.yaml
@@ -35,8 +35,9 @@ dag:
         task_id: "clinvar_vcf_grch37"
         startup_timeout_seconds: 600
         name: "name_basics"
-        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.human_variant_annotation.container_registry.run_csv_transform_kub }}"
         env_vars:
@@ -46,9 +47,13 @@ dag:
           GCS_BUCKET: "{{ var.value.composer_bucket }}"
           TARGET_GCS_FOLDER: "data/human_variant_annotation/clinVar-vcf_GRCh37/"
           PIPELINE: "clinvar"
-        resources:
-          limit_memory: "1G"
-          limit_cpu: "1"
+        container_resources:
+          memory:
+            request: "32Gi"
+          cpu:
+            request: "2"
+          ephemeral-storage:
+            request: "10Gi"
 
     - operator: "GoogleCloudStorageToGoogleCloudStorageOperator"
       description: "Task to run a GoogleCloudStorageToGoogleCloudStorageOperator"
@@ -67,8 +72,9 @@ dag:
         task_id: "clinvar_vcf_grch38"
         startup_timeout_seconds: 600
         name: "name_basics"
-        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.human_variant_annotation.container_registry.run_csv_transform_kub }}"
         env_vars:
@@ -78,9 +84,13 @@ dag:
           GCS_BUCKET: "{{ var.value.composer_bucket }}"
           TARGET_GCS_FOLDER: "data/human_variant_annotation/clinVar-vcf_GRCh38/"
           PIPELINE: "db_snp"
-        resources:
-          limit_memory: "1G"
-          limit_cpu: "1"
+        container_resources:
+          memory:
+            request: "32Gi"
+          cpu:
+            request: "2"
+          ephemeral-storage:
+            request: "10Gi"
 
     - operator: "GoogleCloudStorageToGoogleCloudStorageOperator"
       description: "Task to run a GoogleCloudStorageToGoogleCloudStorageOperator"
diff --git a/datasets/human_variant_annotation/pipelines/db_snp/db_snp_dag.py b/datasets/human_variant_annotation/pipelines/db_snp/db_snp_dag.py
index 00f5ce899..a26cf9dae 100644
--- a/datasets/human_variant_annotation/pipelines/db_snp/db_snp_dag.py
+++ b/datasets/human_variant_annotation/pipelines/db_snp/db_snp_dag.py
@@ -1,4 +1,4 @@
-# Copyright 2021 Google LLC
+# Copyright 2022 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.