feat: automate image builds with different libtpu version daily and enhance JobSet DAG validations

.github/workflows/dag-check.yml

@@ -3,7 +3,6 @@ name: DAG Check
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:

.github/workflows/pyink-check.yml

@@ -2,11 +2,12 @@ name: Formatter
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
   push:
     branches: [master]
 
+  workflow_dispatch: {}
+
 jobs:
   format_check:
     runs-on: ubuntu-latest

.github/workflows/pylint-check.yml

@@ -2,12 +2,13 @@ name: Linter
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:
     branches: [master]
 
+  workflow_dispatch: {}
+
 jobs:
   linting_check:
     runs-on: ubuntu-latest

.github/workflows/require-checklist.yml

@@ -2,10 +2,13 @@ name: Require Checklist
 on:
   pull_request:
     types: [opened, edited, synchronize]
+
+  workflow_dispatch: {}
+
 jobs:
   check_pr_body:
     runs-on: ubuntu-latest
     steps:
       - uses: mheap/require-checklist-action@v2
         with:
-          requireChecklist: false # If this is true and there are no checklists detected, the action will fail
+          requireChecklist: false # If this is true and there are no checklists detected, the action will fail

.github/workflows/unit-test.yml

@@ -3,7 +3,6 @@ name: Unit Test
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:

dags/common/scheduling_helper/scheduling_helper.py

@@ -58,12 +58,12 @@ class DayOfWeek(enum.Enum):
         "multi_host_nodepool_rollback": DefaultTimeout,
         "node_pool_ttr_disk_size": dt.timedelta(minutes=90),
         "node_pool_ttr_update_label": dt.timedelta(minutes=90),
-        "tpu_info_format_validation_dag": DefaultTimeout,
+        "tpu_info_format_validation_dag": dt.timedelta(minutes=45),
         "tpu_sdk_monitoring_validation": DefaultTimeout,
         "jobset_ttr_kill_process": dt.timedelta(minutes=90),
         "jobset_uptime_validation": dt.timedelta(minutes=90),
         "jobset_ttr_drain_restart": DefaultTimeout,
-        "tpu_info_metrics_verification": DefaultTimeout,
+        "tpu_info_metrics_verification": dt.timedelta(minutes=45),
         "jobset_ttr_node_reboot": dt.timedelta(minutes=90),
     },
     TPU_INTERRUPTION_MOCK_CLUSTER.name: {

dags/common/vm_resource.py

@@ -398,3 +398,11 @@ class DockerImage(enum.Enum):
       "gcr.io/tpu-prod-env-multipod/maxtext_post_training_nightly:"
       f"{datetime.datetime.today().strftime('%Y-%m-%d')}"
   )
+  LIBTPU_NIGHTLY = (
+      "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-nightly:"
+      "latest"
+  )
+  LIBTPU_STABLE = (
+      "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-stable:"
+      "latest"
+  )

dags/tpu_observability/dockerfiles/libtpu_nightly.Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.11
+
+ARG LIBTPU_VERSION
+
+RUN apt-get update && apt-get install -y vim nano && nano --version && apt-get install -y curl gpg
+
+RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && apt-get update -y && apt-get install google-cloud-cli -y
+
+RUN pip install jax
+
+RUN pip install --no-cache-dir libtpu==${LIBTPU_VERSION} -f https://storage.googleapis.com/libtpu-wheels/index.html
+
+RUN pip install tpu-info

dags/tpu_observability/dockerfiles/libtpu_stable.Dockerfile

@@ -0,0 +1,9 @@
+FROM python:3.11
+
+RUN apt-get update && apt-get install -y vim nano && nano --version && apt-get install -y curl gpg
+
+RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && apt-get update -y && apt-get install google-cloud-cli -y
+
+RUN pip install jax libtpu
+
+RUN pip install tpu-info

dags/tpu_observability/jobset_ttr_drain_restart.py

@@ -24,6 +24,7 @@
 from airflow.decorators import task
 
 from dags import composer_env
+from dags.common.vm_resource import DockerImage
 from dags.tpu_observability.utils import jobset_util as jobset
 from dags.tpu_observability.utils import subprocess_util as subprocess
 from dags.tpu_observability.utils import node_pool_util as node_pool
@@ -130,6 +131,7 @@ def check_nodes_number(
           gcs_path=GCS_JOBSET_CONFIG_PATH,
           dag_name=DAG_ID,
           node_pool_selector=selector,
+          image=DockerImage.LIBTPU_STABLE.value,
       )
 
       cluster_info = node_pool.build_node_pool_info_from_gcs_yaml(

dags/tpu_observability/jobset_ttr_kill_process.py

@@ -30,6 +30,7 @@
 
 
 from dags import composer_env
+from dags.common.vm_resource import DockerImage
 from dags.tpu_observability.utils import jobset_util as jobset
 from dags.tpu_observability.utils import node_pool_util as node_pool
 from dags.tpu_observability.utils import subprocess_util as subprocess
@@ -133,6 +134,7 @@ def kill_tpu_pod_workload(info: node_pool.Info, pod_name: str) -> None:
           gcs_path=GCS_JOBSET_CONFIG_PATH,
           dag_name=DAG_ID,
           node_pool_selector=selector,
+          image=DockerImage.LIBTPU_STABLE.value,
       )
 
       cluster_info = node_pool.build_node_pool_info_from_gcs_yaml(

dags/tpu_observability/jobset_ttr_node_pool_resize.py

@@ -22,6 +22,7 @@
 from airflow.utils.task_group import TaskGroup
 
 from dags import composer_env
+from dags.common.vm_resource import DockerImage
 from dags.tpu_observability.utils import jobset_util as jobset
 from dags.tpu_observability.utils import node_pool_util as node_pool
 from dags.tpu_observability.utils.jobset_util import JobSet, Workload
@@ -97,6 +98,7 @@
           gcs_path=GCS_JOBSET_CONFIG_PATH,
           dag_name=DAG_ID,
           node_pool_selector=selector,
+          image=DockerImage.LIBTPU_STABLE.value,
       )
 
       cluster_info = node_pool.build_node_pool_info_from_gcs_yaml(

dags/tpu_observability/jobset_ttr_node_reboot.py

@@ -22,6 +22,7 @@
 from airflow.utils.trigger_rule import TriggerRule
 
 from dags import composer_env
+from dags.common.vm_resource import DockerImage
 from dags.tpu_observability.utils import jobset_util as jobset
 from dags.tpu_observability.utils import node_pool_util as node_pool
 from dags.tpu_observability.utils.jobset_util import Workload
@@ -102,6 +103,7 @@
           dag_name=DAG_ID,
           node_pool_selector=selector,
           privileged=True,
+          image=DockerImage.LIBTPU_STABLE.value,
       )
 
       cluster_info = node_pool.build_node_pool_info_from_gcs_yaml(

dags/tpu_observability/jobset_ttr_pod_delete.py

@@ -22,6 +22,7 @@
 from airflow.utils.task_group import TaskGroup
 
 from dags import composer_env
+from dags.common.vm_resource import DockerImage
 from dags.tpu_observability.utils import jobset_util as jobset
 from dags.tpu_observability.utils import node_pool_util as node_pool
 from dags.tpu_observability.utils.jobset_util import Workload
@@ -93,6 +94,7 @@
           gcs_path=GCS_JOBSET_CONFIG_PATH,
           dag_name=DAG_ID,
           node_pool_selector=selector,
+          image=DockerImage.LIBTPU_STABLE.value,
       )
 
       cluster_info = node_pool.build_node_pool_info_from_gcs_yaml(

dags/tpu_observability/jobset_ttr_rollback.py

@@ -22,6 +22,7 @@
 from airflow.utils.task_group import TaskGroup
 
 from dags import composer_env
+from dags.common.vm_resource import DockerImage
 from dags.tpu_observability.utils import jobset_util as jobset
 from dags.tpu_observability.utils import node_pool_util as node_pool
 from dags.tpu_observability.utils.jobset_util import Workload
@@ -94,6 +95,7 @@
           gcs_path=GCS_JOBSET_CONFIG_PATH,
           dag_name=DAG_ID,
           node_pool_selector=selector,
+          image=DockerImage.LIBTPU_STABLE.value,
       )
 
       cluster_info = node_pool.build_node_pool_info_from_gcs_yaml(

dags/tpu_observability/jobset_uptime_validation.py

@@ -23,6 +23,7 @@
 from airflow.utils.trigger_rule import TriggerRule
 
 from dags import composer_env
+from dags.common.vm_resource import DockerImage
 from dags.tpu_observability.configs.common import (
     MachineConfigMap,
     GCS_CONFIG_PATH,
@@ -99,12 +100,13 @@ def get_current_time() -> TimeUtil:
     with TaskGroup(  # pylint: disable=unexpected-keyword-arg
         group_id=f"v{config.tpu_version.value}"
     ):
-      selector = jobset.generate_node_pool_selector("jobset-rollback-ttr")
+      selector = jobset.generate_node_pool_selector("uptime-val")
 
       jobset_config = jobset.build_jobset_from_gcs_yaml(
           gcs_path=GCS_JOBSET_CONFIG_PATH,
           dag_name=DAG_ID,
           node_pool_selector=selector,
+          image=DockerImage.LIBTPU_STABLE.value,
       )
 
       cluster_info = node_pool.build_node_pool_info_from_gcs_yaml(