Merge branch 'GoogleCloudPlatform:master' into sav01-save-local

Author: camiloCienet

.github/CODEOWNERS

@@ -1,24 +1,28 @@
 # Default owners for everything in the repo, unless a later match takes precedence.
-* @richardsliu @yixinshi @RissyRan @hyeygit
+* @richardsliu @yixinshi @RissyRan @hyeygit @jiangjy1982
 
-dags/common @richardsliu @yixinshi @RissyRan @hyeygit @zpcore @ManfeiBai @gobbleturk @shralex @jiangjy1982 @ortibazar @parambole @vipannalla @crankshaw-google @polydier1
+dags/common @richardsliu @yixinshi @RissyRan @hyeygit @bhavya01 @pgmoka @qihqi @zhanyong-wan @gobbleturk @shralex @jiangjy1982 @ortibazar @parambole @vipannalla @crankshaw-google @polydier1
 
 dags/solutions_team/configs/tensorflow @chandrasekhard2 @ZhaoyueCheng @richardsliu
 dags/solutions_team/solutionsteam_tf* @chandrasekhard2 @ZhaoyueCheng @richardsliu
 
-dags/pytorch_xla @vanbasten23 @zpcore @ManfeiBai
-dags/legacy_test/tests/pytorch @vanbasten23 @zpcore @ManfeiBai
+dags/pytorch_xla @bhavya01 @pgmoka @qihqi @zhanyong-wan
+dags/legacy_test/tests/pytorch @bhavya01 @pgmoka @qihqi @zhanyong-wan
 
-dags/mantaray @ManfeiBai @tengyifei @zpcore
+dags/mantaray @qihqi @bhavya01 @pgmoka @zhanyong-wan
 
-dags/multipod @tonyjohnchen @raymondzouu @gobbleturk @shralex @RissyRan @jiangjy1982
+dags/multipod @tonyjohnchen @raymondzouu @gobbleturk @shralex @RissyRan @jiangjy1982 @notabee @xibinliu
 
 dags/mlcompass @ortibazar @sganeshb @brajiang @wlzhg
 
-dags/sparsity_diffusion_devx @RissyRan @parambole @jiangjy1982 @aireenmei @michelle-yooh
-dags/sparsity_diffusion_devx/project_bite* @RissyRan @parambole @jiangjy1982 @aireenmei @michelle-yooh @jiya-zhang
-dags/sparsity_diffusion_devx/configs/project_bite* @RissyRan @parambole @jiangjy1982 @aireenmei @michelle-yooh @jiya-zhang
+dags/sparsity_diffusion_devx @RissyRan @parambole @jiangjy1982 @aireenmei @michelle-yooh @entrpn @coolkp @rohan-bierneni @shuningjin
+dags/sparsity_diffusion_devx/project_bite* @RissyRan @parambole @jiangjy1982 @aireenmei @michelle-yooh @entrpn @coolkp @rohan-bierneni @shuningjin @jiya-zhang
+dags/sparsity_diffusion_devx/configs/project_bite* @RissyRan @parambole @jiangjy1982 @aireenmei @michelle-yooh @entrpn @coolkp @rohan-bierneni @shuningjin @jiya-zhang
 
-dags/inference @vipannalla @mailvijayasingh @sixiang-google @joezijunzhou @singh-mitali
+dags/inference @vipannalla @mailvijayasingh @sixiang-google @joezijunzhou
 
 dags/map_reproducibility @crankshaw-google @polydier1 @gunjanj007 @stingram @utkarshsharma1 @cneel8986 @Doris26 @wang2yn84 @zshu188
+
+dags/orbax @abhinavclemson @xuefgu @xibinliu @andrewyct @severus-ho 
+
+dags/tpu_observability @xibinliu @andrewyct @severus-ho

.github/requirements.txt

@@ -1,10 +1,11 @@
+sqlalchemy-spanner==1.11.0
 apache-airflow==2.10.2
 apache-airflow-providers-sendgrid
 fabric
-google-cloud-bigquery
-google-cloud-compute
-google-cloud-storage
-google-cloud-container
+google-cloud-bigquery>=3.29.0
+google-cloud-compute>=1.26.0
+google-cloud-storage>=2.18.2
+google-cloud-container>=2.54.0
 google-cloud-tpu>=1.16.0
 jsonlines
 tensorflow-cpu

dags/common/model_configs.py

@@ -25,9 +25,18 @@ class MaxTextV5eModelConfigs(enum.Enum):
   DEFAULT_128B = "default_128b_v5e_256"
   GPT3_175B = "gpt_3_175b_v5e_256"
   LLAMA2_7B = "llama2_7b_v5e_256"
+  LLAMA2_13B = "llama2_13b_v5e_256"
   LLAMA2_70B = "llama2_70b_v5e_256"
 
 
+class MaxTextV5pModelConfigs(enum.Enum):
+  # Refers to model configs in https://github.com/AI-Hypercomputer/maxtext/blob/main/benchmarks/maxtext_v5p_model_configs.py
+  GPT3_175B = "gpt_3_175b_v5p_128"
+  GPT3_175B_SC = "gpt_3_175b_v5p_128_sc"
+  LLAMA2_7B = "llama2_7b_v5p_128"
+  LLAMA2_70B = "llama2_70b_v5p_128"
+
+
 class MaxTextTrilliumModelConfigs(enum.Enum):
   # Refers to model configs in https://github.com/AI-Hypercomputer/maxtext/blob/main/benchmarks/maxtext_trillium_model_configs.py
   GPT3_175B = "gpt_3_175b"

dags/common/test_owner.py

@@ -50,6 +50,7 @@ class Team(enum.Enum):
 MOHIT_K = "Mohit K."
 ANISHA_M = "Anisha M."
 YUWEI_Y = "Yuwei Y."
+RISHABH_B = "Rishabh B."
 
 # Multi-tier Checkpointing
 ABHINAV_S = "ABHINAV S."

dags/common/vm_resource.py

@@ -263,6 +263,13 @@ class XpkClusters:
       project=Project.TPU_PROD_ENV_MULTIPOD.value,
       zone=Zone.EUROPE_WEST4_B.value,
   )
+  TPU_V6E_8_CLUSTER = XpkClusterConfig(
+      name="bodaborg-v6e-8-yucmhab-c",
+      device_version=TpuVersion.TRILLIUM,
+      core_count=8,
+      project=Project.TPU_PROD_ENV_ONE_VM.value,
+      zone=Zone.US_EAST5_B.value,
+  )
   TPU_V6E_256_CLUSTER = XpkClusterConfig(
       name="bodaborg-v6e-256",
       device_version=TpuVersion.TRILLIUM,

dags/examples/always_fail.py

@@ -0,0 +1,57 @@
+import logging
+
+from airflow import models
+
+from airflow.decorators import task
+from airflow.exceptions import AirflowException
+from airflow.utils.task_group import TaskGroup
+
+"""
+A DAG that always fails when triggered. This DAG is an example DAG used to
+trigger the on_failure_actions.py/DagRunListener plugin and post a Github
+Issue.
+"""
+
+
+@task
+def task_a():
+  logging.info("task A")
+
+
+# Add or override the owner of the task, in order to assign issue assignees.
+@task(owner="airflow")
+def task_b():
+  logging.info("task B")
+  raise AirflowException("task B failed")
+
+
+@task
+def task_c():
+  logging.info("task C")
+
+
+with models.DAG(
+    dag_id="on_failure_actions_trigger",
+    schedule=None,
+    tags=[
+        "test_dag",
+        "on_failure_alert",  # Add this to enable DagRunListener
+    ],
+    catchup=False,
+    default_args={
+        "retries": 0,  # Set to 0 for throwing exception immediately
+    },
+) as dag:
+  with TaskGroup(group_id="Test1") as group_a:
+    with TaskGroup(group_id="Subgroup1") as group_b:
+      taskA = task_a()
+
+  with TaskGroup(group_id="Test2") as group_c:
+    with TaskGroup(group_id="Subgroup2") as group_d:
+      taskB = task_b()
+
+  with TaskGroup(group_id="Test3") as group_e:
+    with TaskGroup(group_id="Subgroup3") as group_f:
+      taskC = task_c()
+
+  group_a >> group_c >> group_e

dags/legacy_test/tests/pytorch/nightly/common.libsonnet

@@ -105,7 +105,7 @@ local volumes = import 'templates/volumes.libsonnet';
         sudo apt install -y libsndfile-dev
         pip3 install --user --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
         pip install --user \
-          'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev-cp310-cp310-linux_x86_64.whl' \
+          'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.9.0.dev-cp310-cp310-linux_x86_64.whl' \
           -f https://storage.googleapis.com/libtpu-releases/index.html \
           -f https://storage.googleapis.com/libtpu-wheels/index.html
         pip3 install pillow

…/pytorch/r2.7/accelerate-smoke.libsonnet …/pytorch/r2.8/accelerate-smoke.libsonnetdags/legacy_test/tests/pytorch/r2.7/accelerate-smoke.libsonnet renamed to dags/legacy_test/tests/pytorch/r2.8/accelerate-smoke.libsonnet dags/legacy_test/tests/pytorch/r2.7/accelerate-smoke.libsonnet renamed to dags/legacy_test/tests/pytorch/r2.8/accelerate-smoke.libsonnet

@@ -18,22 +18,22 @@ local mixins = import 'templates/mixins.libsonnet';
 local utils = import 'templates/utils.libsonnet';
 local volumes = import 'templates/volumes.libsonnet';
 
-local rcVersion = 'rc5';
+local rcVersion = 'rc1';
 
-// make sure the vision commit aligns with upstream. E.g., for 2.7 release:
-// https://github.com/pytorch/pytorch/blob/release/2.7/.github/ci_commit_pins/vision.txt.
+// make sure the vision commit aligns with upstream. E.g., for 2.8 release:
+// https://github.com/pytorch/pytorch/blob/release/2.8/.github/ci_commit_pins/vision.txt.
 local vision_commit = 'd23a6e1664d20707c11781299611436e1f0c104f';
 
 
 {
-  local r2_7 = {
-    frameworkPrefix: 'pt-2-7',
+  local r2_8 = {
+    frameworkPrefix: 'pt-2-8',
     tpuSettings+: {
       softwareVersion: 'tpu-ubuntu2204-base',
     },
-    imageTag: 'r2.7.0-%(rc)s_3.10' % {rc: rcVersion},
+    imageTag: 'r2.8.0-%(rc)s_3.12' % {rc: rcVersion},
   },
-  PyTorchTest:: common.PyTorchTest + r2_7 {
+  PyTorchTest:: common.PyTorchTest + r2_8 {
     local config = self,
 
     podTemplate+:: {
@@ -74,7 +74,7 @@ local vision_commit = 'd23a6e1664d20707c11781299611436e1f0c104f';
 
                 ctc = cloud_tpu_client.Client(tpu=os.path.basename('$(TPU_NAME)'), zone=os.path.dirname('$(TPU_NAME)'))
                 ctc.wait_for_healthy()
-                ctc.configure_tpu_version(f'pytorch-2.7-dev{libtpu_date}', restart_type='always')
+                ctc.configure_tpu_version(f'pytorch-2.8-dev{libtpu_date}', restart_type='always')
                 ctc.wait_for_healthy()
               |||,
             ],
@@ -110,16 +110,16 @@ local vision_commit = 'd23a6e1664d20707c11781299611436e1f0c104f';
         sudo apt install -y libopenblas-base
         # for huggingface tests
         sudo apt install -y libsndfile-dev
-        # Install torchvision by pinned commit in PyTorch 2.7 release branch.
-        pip install torch==2.7.0 --index-url https://download.pytorch.org/whl/test/cpu
-        # torchvision commit reference: https://github.com/pytorch/pytorch/blob/release/2.7/.github/ci_commit_pins/vision.txt
+        # Install torchvision by pinned commit in PyTorch 2.8 release branch.
+        pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/test/cpu
+        # torchvision commit reference: https://github.com/pytorch/pytorch/blob/release/2.8/.github/ci_commit_pins/vision.txt
         pip install --user --no-use-pep517 "git+https://github.com/pytorch/vision.git@%(vision_commit)s"
-        pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0%(rc)s-cp310-cp310-linux_x86_64.whl
+        pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0%(rc)s-cp310-cp310-linux_x86_64.whl
         pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
         pip install pillow
         git clone --depth=1 https://github.com/pytorch/pytorch.git
         cd pytorch
-        git clone -b v2.7.0-%(rc)s https://github.com/pytorch/xla.git
+        git clone -b v2.8.0-%(rc)s https://github.com/pytorch/xla.git
       ||| % {rc: rcVersion, vision_commit: vision_commit},
     },
     podTemplate+:: {
@@ -155,12 +155,12 @@ local vision_commit = 'd23a6e1664d20707c11781299611436e1f0c104f';
 
         nvidia-smi
         pip uninstall -y torch torchvision
-        pip install torch==2.7.0 --index-url https://download.pytorch.org/whl/test/cpu
+        pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/test/cpu
         pip install --user --no-use-pep517 "git+https://github.com/pytorch/vision.git@%(vision_commit)s"
-        pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.6/torch_xla-2.7.0%(rc)s-cp310-cp310-linux_x86_64.whl
+        pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.6/torch_xla-2.8.0%(rc)s-cp310-cp310-linux_x86_64.whl
 
         mkdir -p pytorch/xla
-        git clone -b v2.7.0-%(rc)s https://github.com/pytorch/xla.git pytorch/xla
+        git clone -b v2.8.0-%(rc)s https://github.com/pytorch/xla.git pytorch/xla
 
         %(cmd)s
 
@@ -238,5 +238,5 @@ local vision_commit = 'd23a6e1664d20707c11781299611436e1f0c104f';
   HuggingfacePipVersionConstraints:: common.HuggingfacePipVersionConstraints,
 
   // DEPRECATED: Use PyTorchTpuVmMixin instead
-  tpu_vm_r2_7_install: self.PyTorchTpuVmMixin.tpuSettings.tpuVmPytorchSetup,
+  tpu_vm_r2_8_install: self.PyTorchTpuVmMixin.tpuSettings.tpuVmPytorchSetup,
 }