Update to the latest test owners (GoogleCloudPlatform#1054)

Author: severus-ho

dags/common/test_owner.py

@@ -44,6 +44,7 @@ class Team(enum.Enum):
 
 # PYTORCH
 MANFEI_B = "manfeiBai"
+BHAVYA_B = "bhavya01"
 
 # MaxText
 TONY_C = "tonyjohnchen"
@@ -92,7 +93,18 @@ class Team(enum.Enum):
 
 # MAP_REPRODUCIBILITY
 GUNJAN_J = "gunjanj007"
+BRYAN_W = "bwuu"
 
 # Bite
 Maggie_Z = "jiya-zhang"
 Andrew_S = "asall"
+
+# Dashboard
+SEVERUS_H = "severus-ho"
+
+# maxtext_pathways
+JULIE_K = "JulieKuo"
+
+# TPU Observability
+YUNA_T = "yuna-tzeng"
+QUINN_M = "QuinnMMcGarry"

dags/dashboard/configs/export_config.py

@@ -11,6 +11,8 @@
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.exceptions import AirflowException
 
+from dags.common import test_owner
+
 # Prefix for exported files in GCS bucket
 GCS_PREFIX = "airflow_exports"
 
@@ -143,6 +145,7 @@ class AirflowTable:
 def get_export_operator(source_table: AirflowTable):
   return PythonOperator(
       task_id=f"export_{source_table.table_name}",
+      owner=test_owner.SEVERUS_H,
       python_callable=export_table_schema_and_data,
       op_kwargs={"table": source_table},
   )
@@ -155,6 +158,7 @@ def get_gcs_to_bq_operator(
   table_name = source_table.table_name
   return GCSToBigQueryOperator(
       task_id=f"load_{table_name}_to_bq",
+      owner=test_owner.SEVERUS_H,
       bucket=source_bucket,
       source_objects=[f"{GCS_PREFIX}/{table_name}_part_*.json"],
       destination_project_dataset_table=destination_table,

dags/dashboard/xlml_to_buganizer.py

@@ -17,6 +17,8 @@
 from google.api_core.exceptions import NotFound
 from google.cloud import bigquery, container_v1
 
+from dags.common import test_owner
+
 # --- Airflow DAG Schedule ---
 SCHEDULED_TIME = "0 */4 * * *"  # Per 4 hours
 
@@ -309,7 +311,7 @@ def fetch_clusters_status(credentials, clusters: List[Any]) -> List[Any]:
 # ================================================================
 # Step 6: Airflow Tasks
 # ================================================================
-@task
+@task(owner=test_owner.SEVERUS_H)
 def pull_clusters_status() -> List[Any]:
   credentials, _ = google.auth.default(
       scopes=["https://www.googleapis.com/auth/cloud-platform"]
@@ -318,7 +320,7 @@ def pull_clusters_status() -> List[Any]:
   return fetch_clusters_status(credentials, target_clusters)
 
 
-@task
+@task(owner=test_owner.SEVERUS_H)
 def insert_gsheet_rows_task(gspread_rows_to_insert):
   if GSPREAD_INSERT_ENABLED and gspread_rows_to_insert:
     context = get_current_context()

dags/mantaray/build_mantaray_docker.py

@@ -15,6 +15,8 @@
 """DAG to build nightly mantaray docker image: gcr.io/tpu-prod-env-multipod/mantaray_maxtext_tpu:nightly"""
 
 from airflow import models
+
+from dags.common import test_owner
 from xlml.utils import mantaray
 import datetime
 from dags import composer_env
@@ -28,5 +30,7 @@
       start_date=datetime.datetime(2024, 9, 4),
       catchup=False,
   ) as dag:
-    run_workload = mantaray.build_docker_image()
+    run_workload = mantaray.build_docker_image.override(
+        owner=test_owner.BHAVYA_B
+    )()
     run_workload

dags/mantaray/run_mantaray_jobs.py

@@ -374,7 +374,9 @@ def vllm_request_rate_inf_nightly_on_v6e_4_persistent_TPUVM():
           start_date=datetime.datetime(2024, 4, 22),
           catchup=False,
       ) as dag:
-        run_workload = mantaray.run_workload(
+        run_workload = mantaray.run_workload.override(
+            owner=test_owner.BHAVYA_B
+        )(
             workload_file_name=job["file_name"],
         )
     run_workload

dags/map_reproducibility/a3mega/gpt3_175b_nemo.py

@@ -19,6 +19,7 @@
 from airflow import models
 
 from dags import composer_env
+from dags.common import test_owner
 from dags.map_reproducibility.utils.common_utils import get_scheduled_time, run_nemo_workload, run_workload_with_quarantine
 
 
@@ -61,6 +62,7 @@
   run_workload_with_quarantine(
       test_name=DAG_ID,
       workload_function=run_nemo_workload,
+      owner=test_owner.BRYAN_W,
       hypercomputer=HYPERCOMPUTER,
       model_id=MODEL_ID,
       framework=FRAMEWORK,

dags/map_reproducibility/a3mega/llama_3_1_70b_nemo.py

@@ -18,6 +18,7 @@
 
 from airflow import models
 from dags import composer_env
+from dags.common import test_owner
 from dags.map_reproducibility.utils.common_utils import get_scheduled_time, run_workload_with_quarantine
 from dags.map_reproducibility.utils.common_utils import run_nemo_workload
 
@@ -61,6 +62,7 @@
   run_workload_with_quarantine(
       test_name=DAG_ID,
       workload_function=run_nemo_workload,
+      owner=test_owner.BRYAN_W,
       hypercomputer=HYPERCOMPUTER,
       model_id=MODEL_ID,
       framework=FRAMEWORK,

dags/map_reproducibility/a3mega/llama_3_1_70b_nemo_gcsfuse.py

@@ -17,6 +17,8 @@
 import datetime
 
 from airflow import models
+
+from dags.common import test_owner
 from dags.map_reproducibility.utils.common_utils import run_nemo_workload, run_workload_with_quarantine
 
 MODEL_ID = "llama3-1-70b"
@@ -59,6 +61,7 @@
   run_workload_with_quarantine(
       test_name=DAG_ID,
       workload_function=run_nemo_workload,
+      owner=test_owner.BRYAN_W,
       hypercomputer=HYPERCOMPUTER,
       model_id=MODEL_ID,
       framework=FRAMEWORK,

dags/map_reproducibility/a3mega/llama_3_70b_nemo.py

@@ -19,6 +19,7 @@
 
 from airflow import models
 from dags import composer_env
+from dags.common import test_owner
 from dags.map_reproducibility.utils.common_utils import get_scheduled_time, run_workload_with_quarantine
 from dags.map_reproducibility.utils.common_utils import run_nemo_workload
 
@@ -67,6 +68,7 @@
   run_workload_with_quarantine(
       test_name=DAG_ID,
       workload_function=run_nemo_workload,
+      owner=test_owner.BRYAN_W,
       hypercomputer=HYPERCOMPUTER,
       model_id=MODEL_ID,
       framework=FRAMEWORK,

dags/map_reproducibility/a3mega/mixtral_8_7b_nemo.py

@@ -18,6 +18,7 @@
 
 from airflow import models
 from dags import composer_env
+from dags.common import test_owner
 from dags.map_reproducibility.utils.common_utils import get_scheduled_time, run_workload_with_quarantine
 from dags.map_reproducibility.utils.common_utils import run_nemo_workload
 
@@ -55,6 +56,7 @@
   run_workload_with_quarantine(
       test_name=DAG_ID,
       workload_function=run_nemo_workload,
+      owner=test_owner.BRYAN_W,
       hypercomputer=HYPERCOMPUTER,
       model_id=MODEL_ID,
       framework=FRAMEWORK,