From ea29ef3d4ae5c59f09e021fc685cc0dda54030c2 Mon Sep 17 00:00:00 2001 From: Alfred Yu Date: Tue, 12 Aug 2025 10:03:26 +0800 Subject: [PATCH 01/19] [cienet-private] Enable workflow checks for the primary working branches --- .github/workflows/dag-check.yml | 1 - .github/workflows/pyink-check.yml | 3 +- .github/workflows/pylint-check.yml | 3 +- .github/workflows/require-checklist.yml | 5 ++- .github/workflows/unit-test.yml | 1 - scripts/code-style.sh | 41 +++++++++++++++++++------ 6 files changed, 40 insertions(+), 14 deletions(-) diff --git a/.github/workflows/dag-check.yml b/.github/workflows/dag-check.yml index b85625224..c1d3e7cf3 100644 --- a/.github/workflows/dag-check.yml +++ b/.github/workflows/dag-check.yml @@ -3,7 +3,6 @@ name: DAG Check on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: diff --git a/.github/workflows/pyink-check.yml b/.github/workflows/pyink-check.yml index ea7a8f250..c5cb1cd83 100644 --- a/.github/workflows/pyink-check.yml +++ b/.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 diff --git a/.github/workflows/pylint-check.yml b/.github/workflows/pylint-check.yml index 5e23d2812..02f4cb31e 100644 --- a/.github/workflows/pylint-check.yml +++ b/.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 diff --git a/.github/workflows/require-checklist.yml b/.github/workflows/require-checklist.yml index d15d19d99..4da288575 100644 --- a/.github/workflows/require-checklist.yml +++ b/.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 \ No newline at end of file + requireChecklist: false # If this is true and there are no checklists detected, the action will fail diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 4f1723cb2..b771ca7e9 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -3,7 +3,6 @@ name: Unit Test on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: diff --git a/scripts/code-style.sh b/scripts/code-style.sh index 36cfa13e9..ea50709ac 100755 --- a/scripts/code-style.sh +++ b/scripts/code-style.sh @@ -19,14 +19,37 @@ set -e FOLDERS_TO_FORMAT=("dags" "xlml") -for folder in "${FOLDERS_TO_FORMAT[@]}" -do - pyink "$folder" --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff -done - -for folder in "${FOLDERS_TO_FORMAT[@]}" -do - pylint "./$folder" --fail-under=9.6 -done +HEAD_SHA="$(git rev-parse HEAD)" +BASE_BRANCH="dev" + +if ! git rev-parse --verify "$BASE_BRANCH" >/dev/null 2>&1; then + git fetch origin "$BASE_BRANCH":"$BASE_BRANCH" || { + echo "[code-style] base branch '$BASE_BRANCH' not found, skip diff-based check." + exit 0 + } +fi + +CHANGED_PY_FILES="$( + git diff --name-only --diff-filter=ACM "${BASE_BRANCH}" "${HEAD_SHA}" \ + | grep '\.py$' \ + | while read -r f; do + for folder in "${FOLDERS_TO_FORMAT[@]}"; do + if [[ "$f" == "$folder/"* ]]; then + echo "$f" + break + fi + done + done \ + | sort -u +)" + +if [[ -z "${CHANGED_PY_FILES}" ]]; then + echo "[pre-push hook] no changed files detected between ${HEAD_SHA} and ${BASE_BRANCH}" + exit 1 +fi + +pyink ${CHANGED_PY_FILES} --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff + +pylint ${CHANGED_PY_FILES} --fail-under=9.6 --disable=E1123 echo "Successfully clean up all codes." From fae3f4ba6a68bc61a4bbb872f744a437cc51646e Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Mon, 13 Apr 2026 08:16:17 +0000 Subject: [PATCH 02/19] Add nightly Dockerfile --- dags/tpu_observability/configs/Dockerfile | 11 +++++++++++ 1 file changed, 11 insertions(+) create mode 100644 dags/tpu_observability/configs/Dockerfile diff --git a/dags/tpu_observability/configs/Dockerfile b/dags/tpu_observability/configs/Dockerfile new file mode 100644 index 000000000..d8a766219 --- /dev/null +++ b/dags/tpu_observability/configs/Dockerfile @@ -0,0 +1,11 @@ +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 --no-cache-dir libtpu==0.0.20.dev20250722+nightly -f https://storage.googleapis.com/libtpu-wheels/index.html + +RUN pip install tpu-info From 836da2b518f70fa69459f806dd564550b43834cc Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Fri, 17 Apr 2026 03:02:23 +0000 Subject: [PATCH 03/19] Add nightly and stable libtpu dockerfile --- dags/tpu_observability/configs/Dockerfile-nightly | 11 +++++++++++ .../configs/{Dockerfile => Dockerfile-stable} | 2 -- 2 files changed, 11 insertions(+), 2 deletions(-) create mode 100644 dags/tpu_observability/configs/Dockerfile-nightly rename dags/tpu_observability/configs/{Dockerfile => Dockerfile-stable} (80%) diff --git a/dags/tpu_observability/configs/Dockerfile-nightly b/dags/tpu_observability/configs/Dockerfile-nightly new file mode 100644 index 000000000..1ed89774f --- /dev/null +++ b/dags/tpu_observability/configs/Dockerfile-nightly @@ -0,0 +1,11 @@ +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 + +RUN pip install --no-cache-dir libtpu==${LIBTPU_VERSION} -f https://storage.googleapis.com/libtpu-wheels/index.html + +RUN pip install tpu-info diff --git a/dags/tpu_observability/configs/Dockerfile b/dags/tpu_observability/configs/Dockerfile-stable similarity index 80% rename from dags/tpu_observability/configs/Dockerfile rename to dags/tpu_observability/configs/Dockerfile-stable index d8a766219..79498423a 100644 --- a/dags/tpu_observability/configs/Dockerfile +++ b/dags/tpu_observability/configs/Dockerfile-stable @@ -6,6 +6,4 @@ RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages. RUN pip install jax libtpu -RUN pip install --no-cache-dir libtpu==0.0.20.dev20250722+nightly -f https://storage.googleapis.com/libtpu-wheels/index.html - RUN pip install tpu-info From 13470ee10c091af4780328aed0847a6cefdf4c36 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Fri, 17 Apr 2026 03:13:28 +0000 Subject: [PATCH 04/19] Add argument to Dockerfile-nigtly --- dags/tpu_observability/configs/Dockerfile-nightly | 2 ++ 1 file changed, 2 insertions(+) diff --git a/dags/tpu_observability/configs/Dockerfile-nightly b/dags/tpu_observability/configs/Dockerfile-nightly index 1ed89774f..10983910a 100644 --- a/dags/tpu_observability/configs/Dockerfile-nightly +++ b/dags/tpu_observability/configs/Dockerfile-nightly @@ -1,5 +1,7 @@ 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 From b5b166380f5c1d8a14597fa1bf1f426ca315683f Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Tue, 21 Apr 2026 04:11:52 +0000 Subject: [PATCH 05/19] feat: update tpu-info parsing and add TPU_OBS image for DAG validation - Add TPU_OBS_LIBTPU_ image to vm_resource for environment testing. - Update tpu_info_utils logic to support parsing the latest tpu-info output format. - Add a new DAG to validate compatibility across different TPU images. --- dags/common/vm_resource.py | 6 + .../tpu_info_format_validation_dags.py | 233 ++++++++++-------- dags/tpu_observability/tpu_info_metric.py | 8 +- dags/tpu_observability/utils/tpu_info_util.py | 150 ++++++----- 4 files changed, 219 insertions(+), 178 deletions(-) diff --git a/dags/common/vm_resource.py b/dags/common/vm_resource.py index 217371c21..a08e66d44 100644 --- a/dags/common/vm_resource.py +++ b/dags/common/vm_resource.py @@ -398,3 +398,9 @@ class DockerImage(enum.Enum): "gcr.io/tpu-prod-env-multipod/maxtext_post_training_nightly:" f"{datetime.datetime.today().strftime('%Y-%m-%d')}" ) + TPU_OBS_LIBTPU_NIGHTLY = ( + "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-nightly:latest" + ) + TPU_OBS_LIBTPU_STABLE = ( + "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-stable:latest" + ) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index d38fab39c..45c16ab9c 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -34,6 +34,7 @@ from dags import composer_env from dags.common import test_owner +from dags.common.vm_resource import DockerImage from dags.tpu_observability.configs.common import ( MachineConfigMap, TpuConfig, @@ -292,6 +293,12 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]): f" output:\n{content.raw_body}" ) +@task +def generate_second_node_pool_name( + node_pool_info: node_pool.Info, +) -> str: + """Generates a second node pool name.""" + return f"{node_pool_info.node_pool_name}-2" # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). @@ -340,111 +347,119 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]): resources, including the JobSet and the temporary node pools. """, ) as dag: + docker_images = { + "stable": DockerImage.TPU_OBS_LIBTPU_STABLE.value, + "nightly": DockerImage.TPU_OBS_LIBTPU_NIGHTLY.value, + } + for machine in MachineConfigMap: config = machine.value - selector = jobset.generate_node_pool_selector( - "tpu-info-format-validation-dag" + DAG_ID + ) + + cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( + gcs_path=GCS_CONFIG_PATH, + dag_name=DAG_ID, + is_prod=composer_env.is_prod_env(), + machine_type=config.machine_version.value, + tpu_topology=config.tpu_topology, + node_pool_selector=selector, ) + cluster_info_2 = copy.deepcopy(cluster_info) + cluster_info_2.node_pool_name = f"{cluster_info.node_pool_name}-2" + # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id=f"v{config.tpu_version.value}" - ): - jobset_config = jobset.build_jobset_from_gcs_yaml( - gcs_path=GCS_JOBSET_CONFIG_PATH, - dag_name=DAG_ID, - node_pool_selector=selector, - ) - - cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( - gcs_path=GCS_CONFIG_PATH, - dag_name=DAG_ID, - is_prod=composer_env.is_prod_env(), - machine_type=config.machine_version.value, - tpu_topology=config.tpu_topology, - node_pool_selector=selector, - ) - - cluster_info_2 = copy.deepcopy(cluster_info) - cluster_info_2.node_pool_name = f"{cluster_info.node_pool_name}-2" - - # Keyword arguments are generated dynamically at runtime (pylint does not - # know this signature). - with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id="create_node_pool" - ) as create_node_pool: - create_first_node_pool = node_pool.create.override( - task_id="node_pool_1", - retries=2, - )( - node_pool=cluster_info, - ) - - create_second_node_pool = node_pool.create.override( - task_id="node_pool_2", - retries=2, - )( - node_pool=cluster_info_2, - ) - - startup = jobset.create_jobset_startup_tasks( + group_id="create_node_pool" + ) as create_node_pool: + create_first_node_pool = node_pool.create.override( + task_id="node_pool_1", + )( node_pool=cluster_info, - jobset_config=jobset_config, - workload_type=Workload.JAX_TPU_BENCHMARK, - ) - - outputs_of_tpu_info = ( - get_tpu_info_from_pod.override(task_id="get_tpu_info") - .partial(info=cluster_info) - .expand(pod_name=startup.running_pods) ) - output_of_tpu_info = ( - tpu_info.parse_tpu_info_output.override( - task_id="get_each_metric_table" - ) - .partial() - .expand(output=outputs_of_tpu_info) + create_second_node_pool = node_pool.create.override( + task_id="node_pool_2", + retries=2, + )( + node_pool=cluster_info_2, ) + image_task_groups = [] + for type_name, image_url in docker_images.items(): # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id="verification_group" - ) as verification_group: - verify_table_amount_task = ( - verify_table_amount.override(task_id="verify_table_amount_task") - .partial() - .expand(tpu_info_output=output_of_tpu_info) + group_id=f"v{config.tpu_version.value}_{type_name}" + ) as image_tg: + image_task_groups.append(image_tg) + + jobset_config = jobset.build_jobset_from_gcs_yaml( + gcs_path=GCS_JOBSET_CONFIG_PATH, + dag_name=DAG_ID, + node_pool_selector=selector, + image=image_url, ) - validate_tpu_chips_metric = ( - validate_chips_table.override(task_id="validate_tpu_chips_metric") - .partial(tpu_config=config) - .expand(tpu_info_output=output_of_tpu_info) + startup = jobset.create_jobset_startup_tasks( + node_pool=cluster_info, + jobset_config=jobset_config, + workload_type=Workload.JAX_TPU_BENCHMARK, ) - validate_runtime_metric = ( - validate_runtime_table.override(task_id="validate_runtime_metric") - .partial() - .expand(tpu_info_output=output_of_tpu_info) + outputs_of_tpu_info = ( + get_tpu_info_from_pod.override(task_id="get_tpu_info") + .partial(info=cluster_info) + .expand(pod_name=startup.running_pods) ) - validate_tensorcore_metric = ( - validate_tensorcore_table.override( - task_id="validate_tensorcore_metric" + output_of_tpu_info = ( + tpu_info.parse_tpu_info_output.override( + task_id="get_each_metric_table" ) .partial() - .expand(tpu_info_output=output_of_tpu_info) + .expand(output=outputs_of_tpu_info) ) - validate_latency_metric = ( - validate_latency_table.override(task_id="validate_latency_metric") - .partial() - .expand(tpu_info_output=output_of_tpu_info) - ) + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id="verification_group" + ) as verification_group: + verify_table_amount_task = ( + verify_table_amount.override(task_id="verify_table_amount_task") + .partial() + .expand(tpu_info_output=output_of_tpu_info) + ) + + validate_tpu_chips_metric = ( + validate_chips_table.override(task_id="validate_tpu_chips_metric") + .partial(tpu_config=config) + .expand(tpu_info_output=output_of_tpu_info) + ) + + validate_runtime_metric = ( + validate_runtime_table.override(task_id="validate_runtime_metric") + .partial() + .expand(tpu_info_output=output_of_tpu_info) + ) + + validate_tensorcore_metric = ( + validate_tensorcore_table.override( + task_id="validate_tensorcore_metric" + ) + .partial() + .expand(tpu_info_output=output_of_tpu_info) + ) + + validate_latency_metric = ( + validate_latency_table.override(task_id="validate_latency_metric") + .partial() + .expand(tpu_info_output=output_of_tpu_info) + ) clean_up_workload = jobset.end_workload.override( task_id="clean_up_workload", trigger_rule=TriggerRule.ALL_DONE @@ -455,40 +470,40 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]): setups=startup.jobset_start_time ) - # Keyword arguments are generated dynamically at runtime (pylint does not - # know this signature). - with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id="cleanup_node_pool" - ) as cleanup_node_pool: - cleanup_first_node_pool = node_pool.delete.override( - task_id="cleanup_node_pool_1", - trigger_rule=TriggerRule.ALL_DONE, - retries=2, - )(node_pool=cluster_info).as_teardown( - setups=create_node_pool, - ) + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id="cleanup_node_pool" + ) as cleanup_node_pool: + cleanup_first_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool_1", + trigger_rule=TriggerRule.ALL_DONE, + retries=2, + )(node_pool=cluster_info).as_teardown( + setups=create_node_pool, + ) - cleanup_second_node_pool = node_pool.delete.override( - task_id="cleanup_node_pool_2", - trigger_rule=TriggerRule.ALL_DONE, - retries=2, - )(node_pool=cluster_info_2).as_teardown( - setups=create_node_pool, - ) + cleanup_second_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool_2", + trigger_rule=TriggerRule.ALL_DONE, + retries=2, + )(node_pool=cluster_info_2).as_teardown( + setups=create_node_pool, + ) - chain( - verify_table_amount_task, - [ - validate_tpu_chips_metric, - validate_runtime_metric, - validate_tensorcore_metric, - validate_latency_metric, - ], - ) + chain( + verify_table_amount_task, + [ + validate_tpu_chips_metric, + validate_runtime_metric, + validate_tensorcore_metric, + validate_latency_metric, + ], + ) - chain(create_first_node_pool, create_second_node_pool) + _ = [create_first_node_pool, create_second_node_pool] - chain(cleanup_first_node_pool, cleanup_second_node_pool) + chain(cleanup_first_node_pool, cleanup_second_node_pool) chain( selector, diff --git a/dags/tpu_observability/tpu_info_metric.py b/dags/tpu_observability/tpu_info_metric.py index e1d98e59f..9b0d9198c 100644 --- a/dags/tpu_observability/tpu_info_metric.py +++ b/dags/tpu_observability/tpu_info_metric.py @@ -540,8 +540,8 @@ def parse_from_tpu_info( # Current Monitoring API aggregation differs from tpu-info, making it # unsuitable as a Source of Truth. Investigation for a valid verification # method is ongoing. - # BufferTransferLatencyStrategy(), - # HostToDeviceTransferLatenciesStrategy(), - # DeviceToHostTransferLatenciesStrategy(), - # CollectiveEndToEndLatencyLatenciesStrategy(), + BufferTransferLatencyStrategy(), + HostToDeviceTransferLatenciesStrategy(), + DeviceToHostTransferLatenciesStrategy(), + CollectiveEndToEndLatencyLatenciesStrategy(), ] diff --git a/dags/tpu_observability/utils/tpu_info_util.py b/dags/tpu_observability/utils/tpu_info_util.py index c408463d4..49c8027a8 100644 --- a/dags/tpu_observability/utils/tpu_info_util.py +++ b/dags/tpu_observability/utils/tpu_info_util.py @@ -7,6 +7,7 @@ from airflow.decorators import task + # A type alias for a parsed row, mapping column headers to their values. _TableRow = dict[str, str] @@ -17,7 +18,7 @@ class Table: name: str raw_body: str - body: list[_TableRow] + body: list[_TableRow] | None = None def parse_body(self): """Parses the raw_body string to populate the structured body attribute.""" @@ -25,35 +26,32 @@ def parse_body(self): class TableLineIndex(IntEnum): """Below is an example of the text returned by tpu-info, formatted as a table. - TPU Chips - ┏━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━┓ - ┃ Chip ┃ Type ┃ Devices ┃ PID ┃ - ┡━━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━┩ - │ /dev/vfio/0 │ TPU v6e chip │ 1 │ 1016 │ - │ /dev/vfio/1 │ TPU v6e chip │ 1 │ 1016 │ - │ /dev/vfio/2 │ TPU v6e chip │ 1 │ 1016 │ - │ /dev/vfio/3 │ TPU v6e chip │ 1 │ 1016 │ - └─────────────┴──────────────┴─────────┴──────┘ + | Chip | Type | Devices | PID | + |-------------|--------------|---------|-----| + | /dev/vfio/0 | TPU v6e chip | 1 | 24 | + | /dev/vfio/1 | TPU v6e chip | 1 | 24 | + | /dev/vfio/2 | TPU v6e chip | 1 | 24 | + | /dev/vfio/3 | TPU v6e chip | 1 | 24 | """ - UPPER_BORDER = 0 - HEADER = auto() - SEPARATOR = auto() - DATA = auto() - LOWER_BORDER = -1 + HEADER = 0 + SEPARATOR = 1 + DATA = 2 lines = self.raw_body.strip().split("\n") - if len(lines) < max(TableLineIndex): + + if len(lines) < TableLineIndex.DATA: self.body = [] return header_line = lines[TableLineIndex.HEADER] - headers = [h.strip() for h in header_line.split("┃") if h.strip()] - data_lines = lines[TableLineIndex.DATA : TableLineIndex.LOWER_BORDER] + headers = [h.strip() for h in header_line.split("|")[1:-1]] + + data_lines = lines[TableLineIndex.DATA :] parsed_body = [] for line in data_lines: - columns = line.split("│")[1:-1] + columns = line.split("|")[1:-1] if len(columns) != len(headers): continue @@ -67,75 +65,97 @@ class TableLineIndex(IntEnum): @task def parse_tpu_info_output(output: str) -> list[Table]: - """Splits a multi-table string from tpu-info into a structured TpuInfo object. + """Splits a multi-table Markdown string from tpu-info into a structured TpuInfo object. Args: output: The raw string output from the 'tpu-info' command. Returns: - A TpuInfo object with attributes populated for each found table. + A list of Table objects with attributes populated for each found table. """ - title_pattern = re.compile(r"(^[^\n].*)\n┏", re.MULTILINE) - table_block_pattern = re.compile(r"(^┏[\s\S]*?┘)", re.MULTILINE) + pattern = re.compile( + r"^([A-Za-z][^\n]*?)\s*\n+" + r"(^\|[^\n]*(?:\n\|[^\n]*)*)", + re.MULTILINE + ) - titles = [s.strip() for s in title_pattern.findall(output)] - blocks = table_block_pattern.findall(output) + parsed_tables = [] - if len(titles) != len(blocks): - raise ValueError( - "Mismatch between found table titles and table blocks. " - f"Found {len(titles)} titles and {len(blocks)} blocks." - ) + for match in pattern.finditer(output): + name = match.group(1).strip() + raw_body = match.group(2).strip() - parsed_tables = [] - for name, raw_body in zip(titles, blocks): table = Table(name=name, raw_body=raw_body, body=None) + table.parse_body() parsed_tables.append(table) + if not parsed_tables: + raise ValueError("Failed to parse any tables from the tpu-info output.") + return parsed_tables if __name__ == "__main__": full_output = """ - Libtpu version: 0.0.23 - Accelerator type: v6e - TPU Chips -┏━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━┓ -┃ Chip ┃ Type ┃ Devices ┃ PID ┃ -┡━━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━┩ -│ /dev/vfio/0 │ TPU v6e chip │ 1 │ 1098 │ -│ /dev/vfio/1 │ TPU v6e chip │ 1 │ 1098 │ -│ /dev/vfio/2 │ TPU v6e chip │ 1 │ 1098 │ -│ /dev/vfio/3 │ TPU v6e chip │ 1 │ 1098 │ -└─────────────┴──────────────┴─────────┴──────┘ + +| Chip | Type | Devices | PID | +|-------------|--------------|---------|-----| +| /dev/vfio/0 | TPU v6e chip | 1 | 24 | +| /dev/vfio/1 | TPU v6e chip | 1 | 24 | +| /dev/vfio/2 | TPU v6e chip | 1 | 24 | +| /dev/vfio/3 | TPU v6e chip | 1 | 24 | + TPU Runtime Utilization -┏━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓ -┃ Device ┃ HBM Usage (GiB) ┃ Duty cycle ┃ -┡━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩ -│ 8 │ 18.45 GiB / 31.25 GiB │ 100.00% │ -│ 9 │ 10.40 GiB / 31.25 GiB │ 100.00% │ -│ 12 │ 10.40 GiB / 31.25 GiB │ 100.00% │ -│ 13 │ 10.40 GiB / 31.25 GiB │ 100.00% │ -└────────┴───────────────────────┴────────────┘ + +| Chip | HBM Usage (GiB) | Duty cycle | +|------|----------------------|------------| +| 0 | 8.44 GiB / 31.25 GiB | 100.00% | +| 1 | 8.44 GiB / 31.25 GiB | 100.00% | +| 4 | 8.44 GiB / 31.25 GiB | 100.00% | +| 5 | 8.44 GiB / 31.25 GiB | 100.00% | + TensorCore Utilization -┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┓ -┃ Chip ID ┃ TensorCore Utilization ┃ -┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━┩ -│ 0 │ 15.14% │ -│ 1 │ 14.56% │ -│ 2 │ 15.53% │ -│ 3 │ 14.97% │ -└─────────┴────────────────────────┘ + +| Core ID | TensorCore Utilization | +|---------|------------------------| +| 0 | 8.71% | +| 1 | 8.56% | +| 2 | 8.52% | +| 3 | 8.20% | + TPU Buffer Transfer Latency -┏━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┓ -┃ Buffer Size ┃ P50 ┃ P90 ┃ P95 ┃ P999 ┃ -┡━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━┩ -│ 8MB+ │ 30154.32 us │ 65472.43 us │ 73990.97 us │ 103220.65 us │ -│ 4MB+ │ 16622.62 us │ 33210.22 us │ 36404.47 us │ 50954.72 us │ -└─────────────┴─────────────┴─────────────┴─────────────┴──────────────┘ + +| Buffer Size | P50 | P90 | P95 | P999 | +|-------------|-------------|-------------|--------------|--------------| +| 8MB+ | 54031.13 us | 96579.79 us | 103413.24 us | 131976.57 us | + +TPU Inbound Buffer Transfer Latency + +| Buffer Size | P50 | P90 | P95 | P999 | +|-------------|-------------|-------------|--------------|--------------| +| 8MB+ | 53198.63 us | 93329.45 us | 101954.57 us | 132554.97 us | + +╭──────────────────────── Host Compute Latency Status ─────────────────────────╮ +│ WARNING: Host Compute Latency metrics unavailable. Did you start a │ +│ MULTI_SLICE workload with `TPU_RUNTIME_METRICS_PORTS=8431,8432,8433,8434`? │ +╰──────────────────────────────────────────────────────────────────────────────╯ +TPU gRPC TCP Minimum RTT + +| P50 | P90 | P95 | P999 | +|----------|----------|----------|----------| +| 66.79 us | 82.26 us | 84.50 us | 86.18 us | + +TPU gRPC TCP Delivery Rate + +| P50 | P90 | P95 | P999 | +|---------------|---------------|---------------|---------------| +| 12408.00 Mbps | 27133.21 Mbps | 30492.25 Mbps | 35873.24 Mbps | """ tpu_info_output = parse_tpu_info_output(full_output) print(tpu_info_output) + for i in tpu_info_output: + print(i.name) + print(i.body) From 8d1bc010810a443b9eb0f95c095a3f3936c7a961 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Wed, 22 Apr 2026 03:38:36 +0000 Subject: [PATCH 06/19] refactor(tpu-info): expand validation metrics and enable cross-image testing Updated the tpu-info format validation DAG to include new network and latency metrics, and refactored the workflow to test multiple libtpu images. Changes: - Added new expected tables to `verify_table_amount`: - "TPU Inbound Buffer Transfer Latency" - "TPU gRPC TCP Minimum RTT" - "TPU gRPC TCP Delivery Rate" - Note: "Host Compute Latency Status" is currently commented out as the metric is broken even in manual testing. - Refactored DAG to run sequential validations for both 'stable' and 'nightly' libtpu Docker images daily. - Optimized node pool usage by creating the infrastructure once and running multiple workloads before cleanup. --- .../tpu_info_format_validation_dags.py | 241 ++++++++++-------- 1 file changed, 131 insertions(+), 110 deletions(-) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 45c16ab9c..287821be1 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -92,6 +92,11 @@ def verify_table_amount(tpu_info_output: list[tpu_info.Table]): "TPU Runtime Utilization", "TensorCore Utilization", "TPU Buffer Transfer Latency", + "TPU Inbound Buffer Transfer Latency", + # This metric is not available in currently tpu-info version + # "Host Compute Latency Status", + "TPU gRPC TCP Minimum RTT", + "TPU gRPC TCP Delivery Rate", } found_names = {table.name for table in tpu_info_output} @@ -354,9 +359,12 @@ def generate_second_node_pool_name( for machine in MachineConfigMap: config = machine.value - selector = jobset.generate_node_pool_selector( - DAG_ID - ) + + with TaskGroup(group_id=f"v{config.tpu_version.value}"): + + selector = jobset.generate_node_pool_selector( + DAG_ID + ) cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( gcs_path=GCS_CONFIG_PATH, @@ -370,16 +378,16 @@ def generate_second_node_pool_name( cluster_info_2 = copy.deepcopy(cluster_info) cluster_info_2.node_pool_name = f"{cluster_info.node_pool_name}-2" - # Keyword arguments are generated dynamically at runtime (pylint does not - # know this signature). - with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id="create_node_pool" - ) as create_node_pool: - create_first_node_pool = node_pool.create.override( - task_id="node_pool_1", - )( - node_pool=cluster_info, - ) + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id="create_node_pool" + ) as create_node_pool: + create_first_node_pool = node_pool.create.override( + task_id="node_pool_1", + )( + node_pool=cluster_info, + ) create_second_node_pool = node_pool.create.override( task_id="node_pool_2", @@ -388,21 +396,20 @@ def generate_second_node_pool_name( node_pool=cluster_info_2, ) - image_task_groups = [] - for type_name, image_url in docker_images.items(): - # Keyword arguments are generated dynamically at runtime (pylint does not - # know this signature). - with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id=f"v{config.tpu_version.value}_{type_name}" - ) as image_tg: - image_task_groups.append(image_tg) - - jobset_config = jobset.build_jobset_from_gcs_yaml( - gcs_path=GCS_JOBSET_CONFIG_PATH, - dag_name=DAG_ID, - node_pool_selector=selector, - image=image_url, - ) + image_task_groups = [] + for type_name, image_url in docker_images.items(): + # Keyword arguments are generated dynamically at runtime (pylint does + # not know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id=f"v{config.tpu_version.value}_{type_name}" + ) as image_tg: + + jobset_config = jobset.build_jobset_from_gcs_yaml( + gcs_path=GCS_JOBSET_CONFIG_PATH, + dag_name=DAG_ID, + node_pool_selector=selector, + image=image_url, + ) startup = jobset.create_jobset_startup_tasks( node_pool=cluster_info, @@ -416,102 +423,116 @@ def generate_second_node_pool_name( .expand(pod_name=startup.running_pods) ) - output_of_tpu_info = ( - tpu_info.parse_tpu_info_output.override( - task_id="get_each_metric_table" - ) - .partial() - .expand(output=outputs_of_tpu_info) - ) - - # Keyword arguments are generated dynamically at runtime (pylint does not - # know this signature). - with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id="verification_group" - ) as verification_group: - verify_table_amount_task = ( - verify_table_amount.override(task_id="verify_table_amount_task") + output_of_tpu_info = ( + tpu_info.parse_tpu_info_output.override( + task_id="get_each_metric_table" + ) .partial() - .expand(tpu_info_output=output_of_tpu_info) + .expand(output=outputs_of_tpu_info) ) - validate_tpu_chips_metric = ( - validate_chips_table.override(task_id="validate_tpu_chips_metric") - .partial(tpu_config=config) - .expand(tpu_info_output=output_of_tpu_info) - ) + # Keyword arguments are generated dynamically at runtime (pylint does + # not know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id="verification_group" + ) as verification_group: + verify_table_amount_task = ( + verify_table_amount.override( + task_id="verify_table_amount_task" + ) + .partial() + .expand(tpu_info_output=output_of_tpu_info) + ) - validate_runtime_metric = ( - validate_runtime_table.override(task_id="validate_runtime_metric") - .partial() - .expand(tpu_info_output=output_of_tpu_info) - ) + validate_tpu_chips_metric = ( + validate_chips_table.override( + task_id="validate_tpu_chips_metric" + ) + .partial(tpu_config=config) + .expand(tpu_info_output=output_of_tpu_info) + ) - validate_tensorcore_metric = ( - validate_tensorcore_table.override( - task_id="validate_tensorcore_metric" - ) - .partial() - .expand(tpu_info_output=output_of_tpu_info) - ) + validate_runtime_metric = ( + validate_runtime_table.override( + task_id="validate_runtime_metric" + ) + .partial() + .expand(tpu_info_output=output_of_tpu_info) + ) - validate_latency_metric = ( - validate_latency_table.override(task_id="validate_latency_metric") - .partial() - .expand(tpu_info_output=output_of_tpu_info) - ) + validate_tensorcore_metric = ( + validate_tensorcore_table.override( + task_id="validate_tensorcore_metric" + ) + .partial() + .expand(tpu_info_output=output_of_tpu_info) + ) - clean_up_workload = jobset.end_workload.override( - task_id="clean_up_workload", trigger_rule=TriggerRule.ALL_DONE - )( - node_pool=cluster_info, - jobset_config=jobset_config, - ).as_teardown( - setups=startup.jobset_start_time - ) + validate_latency_metric = ( + validate_latency_table.override( + task_id="validate_latency_metric" + ) + .partial() + .expand(tpu_info_output=output_of_tpu_info) + ) - # Keyword arguments are generated dynamically at runtime (pylint does not - # know this signature). - with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id="cleanup_node_pool" - ) as cleanup_node_pool: - cleanup_first_node_pool = node_pool.delete.override( - task_id="cleanup_node_pool_1", - trigger_rule=TriggerRule.ALL_DONE, - retries=2, - )(node_pool=cluster_info).as_teardown( - setups=create_node_pool, + chain( + verify_table_amount_task, + [ + validate_tpu_chips_metric, + validate_runtime_metric, + validate_tensorcore_metric, + validate_latency_metric, + ], + ) + + clean_up_workload = jobset.end_workload.override( + task_id="clean_up_workload", trigger_rule=TriggerRule.ALL_DONE + )( + node_pool=cluster_info, + jobset_config=jobset_config, + ).as_teardown( + setups=startup.jobset_start_time ) - cleanup_second_node_pool = node_pool.delete.override( - task_id="cleanup_node_pool_2", - trigger_rule=TriggerRule.ALL_DONE, - retries=2, - )(node_pool=cluster_info_2).as_teardown( - setups=create_node_pool, + chain( + startup.task_group, + outputs_of_tpu_info, + output_of_tpu_info, + verification_group, + clean_up_workload, ) + image_task_groups.append(image_tg) - chain( - verify_table_amount_task, - [ - validate_tpu_chips_metric, - validate_runtime_metric, - validate_tensorcore_metric, - validate_latency_metric, - ], + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id="cleanup_node_pool" + ) as cleanup_node_pool: + cleanup_first_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool_1", + trigger_rule=TriggerRule.ALL_DONE, + retries=2, + )(node_pool=cluster_info).as_teardown( + setups=create_node_pool, ) - _ = [create_first_node_pool, create_second_node_pool] + cleanup_second_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool_2", + trigger_rule=TriggerRule.ALL_DONE, + retries=2, + )(node_pool=cluster_info_2).as_teardown( + setups=create_node_pool, + ) chain(cleanup_first_node_pool, cleanup_second_node_pool) - chain( - selector, - create_node_pool, - *startup.tasks, - outputs_of_tpu_info, - output_of_tpu_info, - verification_group, - clean_up_workload, - cleanup_node_pool, - ) + _ = [create_first_node_pool, create_second_node_pool] + chain(selector, cluster_info, cluster_info_2, create_node_pool) + + current_node = create_node_pool + for tg in image_task_groups: + chain(current_node, tg) + current_node = tg + + chain(current_node, cleanup_node_pool) From 18d0e8ba6d5591960b3715c25bd99c76cbc8b9b1 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Wed, 22 Apr 2026 07:16:18 +0000 Subject: [PATCH 07/19] chore(tpu-observability): configure stable docker image for jobset DAGs Explicitly specify the stable Docker image when building jobsets in: - jobset_ttr_kill_process.py - jobset_ttr_node_pool_resize.py - jobset_ttr_pod_delete_exist.py - jobset_ttr_pod_delete.py - jobset_ttr_rollback.py - jobset_uptime_validation.py This ensures we are using a known stable image for these validation tests. --- .../jobset_ttr_kill_process.py | 2 + .../jobset_ttr_node_pool_resize.py | 2 + .../jobset_ttr_pod_delete.py | 2 + dags/tpu_observability/jobset_ttr_rollback.py | 2 + .../jobset_uptime_validation.py | 2 + .../tpu_observability/tpu_info_metrics_dag.py | 138 +++++++++-------- .../tpu_sdk_monitoring_validation_dag.py | 145 ++++++++++-------- 7 files changed, 168 insertions(+), 125 deletions(-) diff --git a/dags/tpu_observability/jobset_ttr_kill_process.py b/dags/tpu_observability/jobset_ttr_kill_process.py index bbf45257a..ae3793d00 100644 --- a/dags/tpu_observability/jobset_ttr_kill_process.py +++ b/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.TPU_OBS_LIBTPU_STABLE.value, ) cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( diff --git a/dags/tpu_observability/jobset_ttr_node_pool_resize.py b/dags/tpu_observability/jobset_ttr_node_pool_resize.py index 9fabae6a3..f1f7479be 100644 --- a/dags/tpu_observability/jobset_ttr_node_pool_resize.py +++ b/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.TPU_OBS_LIBTPU_STABLE.value, ) cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( diff --git a/dags/tpu_observability/jobset_ttr_pod_delete.py b/dags/tpu_observability/jobset_ttr_pod_delete.py index 5745281bd..e84252303 100644 --- a/dags/tpu_observability/jobset_ttr_pod_delete.py +++ b/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.TPU_OBS_LIBTPU_STABLE.value, ) cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( diff --git a/dags/tpu_observability/jobset_ttr_rollback.py b/dags/tpu_observability/jobset_ttr_rollback.py index b85d23c7d..e282af9b5 100644 --- a/dags/tpu_observability/jobset_ttr_rollback.py +++ b/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.TPU_OBS_LIBTPU_STABLE.value, ) cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( diff --git a/dags/tpu_observability/jobset_uptime_validation.py b/dags/tpu_observability/jobset_uptime_validation.py index 0558e3e8d..a5cc011cd 100644 --- a/dags/tpu_observability/jobset_uptime_validation.py +++ b/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, @@ -105,6 +106,7 @@ def get_current_time() -> TimeUtil: gcs_path=GCS_JOBSET_CONFIG_PATH, dag_name=DAG_ID, node_pool_selector=selector, + image=DockerImage.TPU_OBS_LIBTPU_STABLE.value, ) cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( diff --git a/dags/tpu_observability/tpu_info_metrics_dag.py b/dags/tpu_observability/tpu_info_metrics_dag.py index 120c0de25..f468aba81 100644 --- a/dags/tpu_observability/tpu_info_metrics_dag.py +++ b/dags/tpu_observability/tpu_info_metrics_dag.py @@ -31,6 +31,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, @@ -275,6 +276,11 @@ def summarize_results( pools. """, ) as dag: + docker_images = { + "stable": DockerImage.TPU_OBS_LIBTPU_STABLE.value, + "nightly": DockerImage.TPU_OBS_LIBTPU_NIGHTLY.value, + } + for machine in MachineConfigMap: config = machine.value @@ -316,72 +322,81 @@ def summarize_results( _ = [create_first_node_pool, create_second_node_pool] - startup = jobset.create_jobset_startup_tasks( - node_pool=cluster_info, - jobset_config=jobset_config, - workload_type=Workload.JAX_TPU_BENCHMARK, - ) + image_task_groups = [] + for type_name, image_url in docker_images.items(): + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id=f"v{config.tpu_version.value}_{type_name}" + ) as image_tg: - verification_results = {} - all_verification_groups = [] + startup = jobset.create_jobset_startup_tasks( + node_pool=cluster_info, + jobset_config=jobset_config, + workload_type=Workload.JAX_TPU_BENCHMARK, + ) - for strategy in ALL_METRIC_STRATEGIES: - group_id = f"verify_{strategy.dag_id_suffix}" + verification_results = {} + all_verification_groups = [] - with TaskGroup(group_id=group_id) as verification_group: - tpu_info_metric_outputs = ( - get_tpu_info_metric_from_pod.override( - task_id="get_tpu_info_metric_table" - ) - .partial( - node_pool=cluster_info, - jobset_config=jobset_config, - metric_name=strategy.tpu_info_metric_name, - ) - .expand(pod_name=startup.running_pods) - ) + for strategy in ALL_METRIC_STRATEGIES: + group_id = f"verify_{strategy.dag_id_suffix}" - tpu_info_metric_output = ( - tpu_info.parse_tpu_info_output.override( - task_id="get_each_metric_table" + with TaskGroup(group_id=group_id) as verification_group: + tpu_info_metric_outputs = ( + get_tpu_info_metric_from_pod.override( + task_id="get_tpu_info_metric_table" + ) + .partial( + node_pool=cluster_info, + jobset_config=jobset_config, + metric_name=strategy.tpu_info_metric_name, + ) + .expand(pod_name=startup.running_pods) ) - .partial() - .expand(output=tpu_info_metric_outputs) - ) - verify_metric = ( - run_metric_verification.override(task_id="run_verification") - .partial( - node_pool=cluster_info, - job_apply_time=startup.jobset_start_time, - metric_strategy=strategy, + tpu_info_metric_output = ( + tpu_info.parse_tpu_info_output.override( + task_id="get_each_metric_table" + ) + .partial() + .expand(output=tpu_info_metric_outputs) ) - .expand( - comparison_data=startup.running_pods.zip( - tpu_info_metric_output + + verify_metric = ( + run_metric_verification.override(task_id="run_verification") + .partial( + node_pool=cluster_info, + job_apply_time=startup.jobset_start_time, + metric_strategy=strategy, ) + .expand(comparison_data=startup.running_pods.zip(tpu_info_metric_output)) ) - ) - all_verification_groups.append(verification_group) + all_verification_groups.append(verification_group) + verification_results[strategy.dag_id_suffix] = verify_metric - verification_results[strategy.dag_id_suffix] = verify_metric + summary = summarize_results.override( + task_id="summarize_results", trigger_rule=TriggerRule.ALL_DONE + )( + verification_results_dict=verification_results, + active_pods=startup.running_pods, + ) - summary = summarize_results.override( - task_id="summarize_results", trigger_rule=TriggerRule.ALL_DONE - )( - verification_results_dict=verification_results, - active_pods=startup.running_pods, - ) + clean_up_workload = jobset.end_workload.override( + task_id="clean_up_workload", trigger_rule=TriggerRule.ALL_DONE + )( + node_pool=cluster_info, + jobset_config=jobset_config, + ).as_teardown( + setups=startup.jobset_start_time + ) - clean_up_workload = jobset.end_workload.override( - task_id="clean_up_workload", trigger_rule=TriggerRule.ALL_DONE - )( - node_pool=cluster_info, - jobset_config=jobset_config, - ).as_teardown( - setups=startup.jobset_start_time - ) + chain( + *startup.tasks, + all_verification_groups, + summary, + clean_up_workload, + ) + image_task_groups.append(image_tg) with TaskGroup(group_id="cleanup_node_pool") as cleanup_node_pool: cleanup_first_node_pool = node_pool.delete.override( @@ -400,12 +415,11 @@ def summarize_results( chain(cleanup_first_node_pool, cleanup_second_node_pool) - chain( - selector, - create_node_pool, - *startup.tasks, - all_verification_groups, - summary, - clean_up_workload, - cleanup_node_pool, - ) + chain(selector, cluster_info, cluster_info_2, create_node_pool) + + current_node = create_node_pool + for tg in image_task_groups: + chain(current_node, tg) + current_node = tg + + chain(current_node, cleanup_node_pool) diff --git a/dags/tpu_observability/tpu_sdk_monitoring_validation_dag.py b/dags/tpu_observability/tpu_sdk_monitoring_validation_dag.py index 8f7b07523..4d62a0724 100644 --- a/dags/tpu_observability/tpu_sdk_monitoring_validation_dag.py +++ b/dags/tpu_observability/tpu_sdk_monitoring_validation_dag.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 tpu_monitoring_sdk_util as sdk from dags.tpu_observability.utils import node_pool_util as node_pool @@ -123,69 +124,87 @@ def validate_monitoring_sdk(info: node_pool.Info, pod_name: str) -> None: and hardware devices inside the container. """, ) as dag: + docker_images = { + "stable": DockerImage.TPU_OBS_LIBTPU_STABLE.value, + "nightly": DockerImage.TPU_OBS_LIBTPU_NIGHTLY.value, + } + for machine in MachineConfigMap: config = machine.value - # Keyword arguments are generated dynamically at runtime (pylint does not - # know this signature). - with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id=f"v{config.tpu_version.value}" - ): - selector = jobset.generate_node_pool_selector( - "tpu-sdk-monitoring-validation" - ) - - jobset_config = jobset.build_jobset_from_gcs_yaml( - gcs_path=GCS_JOBSET_CONFIG_PATH, - dag_name=DAG_ID, - node_pool_selector=selector, - ) - - cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( - gcs_path=GCS_CONFIG_PATH, - dag_name=DAG_ID, - is_prod=composer_env.is_prod_env(), - machine_type=config.machine_version.value, - tpu_topology=config.tpu_topology, - node_pool_selector=selector, - ) - - create_node_pool = node_pool.create.override(task_id="create_node_pool")( - node_pool=cluster_info, - ) - - startup = jobset.create_jobset_startup_tasks( - node_pool=cluster_info, - jobset_config=jobset_config, - workload_type=Workload.JAX_TPU_BENCHMARK, - ) - - sdk_validation = ( - validate_monitoring_sdk.override(task_id="sdk_validation") - .partial(info=cluster_info) - .expand(pod_name=startup.running_pods) - ) - - cleanup_workload = jobset.end_workload.override( - task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE - )( - node_pool=cluster_info, - jobset_config=jobset_config, - ).as_teardown( - setups=startup.jobset_start_time - ) - - cleanup_node_pool = node_pool.delete.override( - task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE - )(node_pool=cluster_info).as_teardown( - setups=create_node_pool, - ) - - chain( - selector, - create_node_pool, - *startup.tasks, - sdk_validation, - cleanup_workload, - cleanup_node_pool, - ) + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id=f"v{config.tpu_version.value}" + ): + selector = jobset.generate_node_pool_selector( + "tpu-sdk-monitoring-validation" + ) + + jobset_config = jobset.build_jobset_from_gcs_yaml( + gcs_path=GCS_JOBSET_CONFIG_PATH, + dag_name=DAG_ID, + node_pool_selector=selector, + ) + + cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( + gcs_path=GCS_CONFIG_PATH, + dag_name=DAG_ID, + is_prod=composer_env.is_prod_env(), + machine_type=config.machine_version.value, + tpu_topology=config.tpu_topology, + node_pool_selector=selector, + ) + + create_node_pool = node_pool.create.override(task_id="create_node_pool")( + node_pool=cluster_info, + ) + + image_task_groups = [] + for type_name, image_url in docker_images.items(): + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id=f"v{config.tpu_version.value}_{type_name}" + ) as image_tg: + + startup = jobset.create_jobset_startup_tasks( + node_pool=cluster_info, + jobset_config=jobset_config, + workload_type=Workload.JAX_TPU_BENCHMARK, + ) + + sdk_validation = ( + validate_monitoring_sdk.override(task_id="sdk_validation") + .partial(info=cluster_info) + .expand(pod_name=startup.running_pods) + ) + + cleanup_workload = jobset.end_workload.override( + task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE + )( + node_pool=cluster_info, + jobset_config=jobset_config, + ).as_teardown( + setups=startup.jobset_start_time + ) + + chain( + *startup.tasks, + sdk_validation, + cleanup_workload, + ) + image_task_groups.append(image_tg) + + cleanup_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE + )(node_pool=cluster_info).as_teardown( + setups=create_node_pool, + ) + + chain(selector, jobset_config, cluster_info, create_node_pool) + + current_node = create_node_pool + for tg in image_task_groups: + chain(current_node, tg) + current_node = tg + + chain(current_node, cleanup_node_pool) From 3d7b51c528fa3aff2ccb20d2ced3458a969c0732 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Thu, 23 Apr 2026 07:20:09 +0000 Subject: [PATCH 08/19] chore: Change "Latencies" related metric tolerance_percent to 10.0 --- dags/tpu_observability/tpu_info_metric.py | 8 ++++---- dags/tpu_observability/tpu_info_metrics_dag.py | 6 +++++- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/dags/tpu_observability/tpu_info_metric.py b/dags/tpu_observability/tpu_info_metric.py index 9b0d9198c..80c43db63 100644 --- a/dags/tpu_observability/tpu_info_metric.py +++ b/dags/tpu_observability/tpu_info_metric.py @@ -443,7 +443,7 @@ class BufferTransferLatencyStrategy(_BaseDistributionStrategy): ) tpu_info_metric_name = "buffer_transfer_latency" dag_id_suffix = "buffer_transfer_latency" - tolerance_percent = 3.0 + tolerance_percent = 10.0 _monitoring_group_by_label = "buffer_size" _tpu_info_table_name = "TPU Buffer Transfer Latency" _tpu_info_group_by_key = "Buffer Size" @@ -455,7 +455,7 @@ class HostToDeviceTransferLatenciesStrategy(_BaseDistributionStrategy): metric_name = "kubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencies" tpu_info_metric_name = "host_to_device_transfer_latency" dag_id_suffix = "host_to_device_transfer_latency" - tolerance_percent = 3.0 + tolerance_percent = 10.0 _monitoring_group_by_label = "buffer_size" _tpu_info_table_name = "TPU Host to Device Transfer Latency" _tpu_info_group_by_key = "Buffer Size" @@ -469,7 +469,7 @@ class DeviceToHostTransferLatenciesStrategy(_BaseDistributionStrategy): metric_name = "kubernetes.io/container/multislice/accelerator/device_to_host_transfer_latencies" tpu_info_metric_name = "device_to_host_transfer_latency" dag_id_suffix = "device_to_host_transfer_latency" - tolerance_percent = 3.0 + tolerance_percent = 10.0 _monitoring_group_by_label = "buffer_size" _tpu_info_table_name = "TPU Device to Host Transfer Latency" _tpu_info_group_by_key = "Buffer Size" @@ -483,7 +483,7 @@ class CollectiveEndToEndLatencyLatenciesStrategy(_BaseDistributionStrategy): metric_name = "kubernetes.io/container/multislice/network/collective_end_to_end_latencies" tpu_info_metric_name = "collective_e2e_latency" dag_id_suffix = "collective_e2e_latency" - tolerance_percent = 3.0 + tolerance_percent = 10.0 _monitoring_group_by_label = "collective_type" _tpu_info_table_name = "TPU Collective End to End Latency" _tpu_info_group_by_key = "Buffer Size" diff --git a/dags/tpu_observability/tpu_info_metrics_dag.py b/dags/tpu_observability/tpu_info_metrics_dag.py index f468aba81..0de3b8966 100644 --- a/dags/tpu_observability/tpu_info_metrics_dag.py +++ b/dags/tpu_observability/tpu_info_metrics_dag.py @@ -368,7 +368,11 @@ def summarize_results( job_apply_time=startup.jobset_start_time, metric_strategy=strategy, ) - .expand(comparison_data=startup.running_pods.zip(tpu_info_metric_output)) + .expand( + comparison_data=startup.running_pods.zip( + tpu_info_metric_output + ) + ) ) all_verification_groups.append(verification_group) From 0d06dbc6b8619b8733645d272702c1e35286c371 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Thu, 23 Apr 2026 08:57:05 +0000 Subject: [PATCH 09/19] format fix --- dags/common/vm_resource.py | 8 ++------ dags/tpu_observability/utils/tpu_info_util.py | 4 +--- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/dags/common/vm_resource.py b/dags/common/vm_resource.py index a08e66d44..2d9a57474 100644 --- a/dags/common/vm_resource.py +++ b/dags/common/vm_resource.py @@ -398,9 +398,5 @@ class DockerImage(enum.Enum): "gcr.io/tpu-prod-env-multipod/maxtext_post_training_nightly:" f"{datetime.datetime.today().strftime('%Y-%m-%d')}" ) - TPU_OBS_LIBTPU_NIGHTLY = ( - "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-nightly:latest" - ) - TPU_OBS_LIBTPU_STABLE = ( - "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-stable:latest" - ) + TPU_OBS_LIBTPU_NIGHTLY = "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-nightly:latest" + TPU_OBS_LIBTPU_STABLE = "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-stable:latest" diff --git a/dags/tpu_observability/utils/tpu_info_util.py b/dags/tpu_observability/utils/tpu_info_util.py index 49c8027a8..1112b653a 100644 --- a/dags/tpu_observability/utils/tpu_info_util.py +++ b/dags/tpu_observability/utils/tpu_info_util.py @@ -74,9 +74,7 @@ def parse_tpu_info_output(output: str) -> list[Table]: A list of Table objects with attributes populated for each found table. """ pattern = re.compile( - r"^([A-Za-z][^\n]*?)\s*\n+" - r"(^\|[^\n]*(?:\n\|[^\n]*)*)", - re.MULTILINE + r"^([A-Za-z][^\n]*?)\s*\n+" r"(^\|[^\n]*(?:\n\|[^\n]*)*)", re.MULTILINE ) parsed_tables = [] From fe2113c9ef410e17918e39c53e77cff44b2d9bde Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Thu, 23 Apr 2026 09:02:32 +0000 Subject: [PATCH 10/19] fix --- .../tpu_info_format_validation_dags.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 287821be1..77935b288 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -298,6 +298,7 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]): f" output:\n{content.raw_body}" ) + @task def generate_second_node_pool_name( node_pool_info: node_pool.Info, @@ -305,6 +306,7 @@ def generate_second_node_pool_name( """Generates a second node pool name.""" return f"{node_pool_info.node_pool_name}-2" + # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with models.DAG( # pylint: disable=unexpected-keyword-arg @@ -361,10 +363,7 @@ def generate_second_node_pool_name( config = machine.value with TaskGroup(group_id=f"v{config.tpu_version.value}"): - - selector = jobset.generate_node_pool_selector( - DAG_ID - ) + selector = jobset.generate_node_pool_selector(DAG_ID) cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( gcs_path=GCS_CONFIG_PATH, @@ -403,7 +402,6 @@ def generate_second_node_pool_name( with TaskGroup( # pylint: disable=unexpected-keyword-arg group_id=f"v{config.tpu_version.value}_{type_name}" ) as image_tg: - jobset_config = jobset.build_jobset_from_gcs_yaml( gcs_path=GCS_JOBSET_CONFIG_PATH, dag_name=DAG_ID, @@ -437,9 +435,7 @@ def generate_second_node_pool_name( group_id="verification_group" ) as verification_group: verify_table_amount_task = ( - verify_table_amount.override( - task_id="verify_table_amount_task" - ) + verify_table_amount.override(task_id="verify_table_amount_task") .partial() .expand(tpu_info_output=output_of_tpu_info) ) From 965f83c33d60893d3396b79fbf6d5c69a1724feb Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Thu, 23 Apr 2026 10:01:02 +0000 Subject: [PATCH 11/19] fix: change parse_from_tpu_info logic since some metric using DistributionStrategy don't have group_key in tpu-info table --- dags/tpu_observability/tpu_info_metric.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/dags/tpu_observability/tpu_info_metric.py b/dags/tpu_observability/tpu_info_metric.py index 80c43db63..cd4edd99e 100644 --- a/dags/tpu_observability/tpu_info_metric.py +++ b/dags/tpu_observability/tpu_info_metric.py @@ -305,15 +305,13 @@ def parse_from_tpu_info( for metric_table in tpu_info_metric_output: if metric_table.name == table_name: for row_dict in metric_table.body: - group_value = row_dict.get(group_key) - if not group_value: - continue + group_value = row_dict.get(group_key, "summary") - parsed_values_by_group[group_value] = {} + if group_value not in parsed_values_by_group: + parsed_values_by_group[group_value] = {} for p in self.percentiles_to_check: # Use Enum name as key (e.g. "P999") value_str = row_dict.get(p.name, "") - match = re.search(r"([\d\.]+)", value_str) if match: parsed_values_by_group[group_value][p.name] = float( From 7a65ad7ab7a6bb47c4fa4cf7e0ec3e7e6a2f27d8 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Fri, 24 Apr 2026 02:55:44 +0000 Subject: [PATCH 12/19] fix --- dags/tpu_observability/utils/tpu_info_util.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dags/tpu_observability/utils/tpu_info_util.py b/dags/tpu_observability/utils/tpu_info_util.py index 1112b653a..c552447a4 100644 --- a/dags/tpu_observability/utils/tpu_info_util.py +++ b/dags/tpu_observability/utils/tpu_info_util.py @@ -154,6 +154,6 @@ def parse_tpu_info_output(output: str) -> list[Table]: tpu_info_output = parse_tpu_info_output(full_output) print(tpu_info_output) - for i in tpu_info_output: - print(i.name) - print(i.body) + # for i in tpu_info_output: + # print(i.name) + # print(i.body) From 4c8cb11547064af59125573c5b01ae59d0adc3a8 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Fri, 24 Apr 2026 03:47:35 +0000 Subject: [PATCH 13/19] feat(tpu_observability): use custom labels for distribution metrics in tpu_info verification --- dags/tpu_observability/tpu_info_metric.py | 34 +++++++++++++++++++ .../tpu_observability/tpu_info_metrics_dag.py | 16 ++++++--- 2 files changed, 45 insertions(+), 5 deletions(-) diff --git a/dags/tpu_observability/tpu_info_metric.py b/dags/tpu_observability/tpu_info_metric.py index cd4edd99e..a777abfc4 100644 --- a/dags/tpu_observability/tpu_info_metric.py +++ b/dags/tpu_observability/tpu_info_metric.py @@ -81,6 +81,16 @@ def parse_from_tpu_info( """ pass + def get_labels( + self, tpu_info_metric_output: list[tpu_info.Table] + ) -> list[str] | None: + """Gets labels for the metric values. + + Default returns None, which implies using default device labels. + """ + return None + + class _BaseSimplePointStrategy(BaseMetricStrategy): """Base strategy for parsing single numeric data points from Cloud Monitoring. @@ -325,6 +335,30 @@ def parse_from_tpu_info( return tpu_info_data_values + def get_labels( + self, tpu_info_metric_output: list[tpu_info.Table] + ) -> list[str] | None: + """Parses labels for percentiles from `tpu-info` output tables.""" + parsed_values_by_group: dict[str, dict[str, float]] = {} + table_name = self._tpu_info_table_name + group_key = self._tpu_info_group_by_key + + for metric_table in tpu_info_metric_output: + if metric_table.name == table_name: + for row_dict in metric_table.body: + group_value = row_dict.get(group_key, "summary") + + if group_value not in parsed_values_by_group: + parsed_values_by_group[group_value] = {} + + labels = [] + for group_value in sorted(parsed_values_by_group.keys()): + for p in self.percentiles_to_check: + labels.append(f"{group_value} {p.name}") + + return labels + + class MemoryUsedStrategy(_BaseSimplePointStrategy): """Strategy for verifying Used HBM Memory.""" diff --git a/dags/tpu_observability/tpu_info_metrics_dag.py b/dags/tpu_observability/tpu_info_metrics_dag.py index 0de3b8966..ffac67cd3 100644 --- a/dags/tpu_observability/tpu_info_metrics_dag.py +++ b/dags/tpu_observability/tpu_info_metrics_dag.py @@ -61,6 +61,7 @@ def compare_metric_values( pod_name: str, metric_display_name: str, tolerance_percent: float, + labels: list[str] | None = None, ): """Compares two lists of metric values and checks if they are within a tolerance range.""" if len(cmd_values) != len(monitoring_values): @@ -76,15 +77,15 @@ def compare_metric_values( metric_display_name, ) logging.info( - "%-12s%-15s%-17s%-12s%-15s%-10s", - "Device", + "%-20s%-15s%-17s%-12s%-15s%-10s", + "Label" if labels else "Device", "TPU-Info Val", "Monitoring Val", "Difference", "Allowed Diff", "Result", ) - logging.info("-" * 85) + logging.info("-" * 95) all_passed = True for i, (log_val, mon_val) in enumerate(zip(cmd_values, monitoring_values)): @@ -93,9 +94,10 @@ def compare_metric_values( passed = diff <= allowed_diff if not passed: all_passed = False + label = labels[i] if labels else f"Device {i}" logging.info( - "%-12s%-15.2f%-17.2f%-12.2f%-15.2f%-10s", - f"Device {i}", + "%-20s%-15.2f%-17.2f%-12.2f%-15.2f%-10s", + label, log_val, mon_val, diff, @@ -114,6 +116,7 @@ def compare_metric_values( ) + @task def get_tpu_info_metric_from_pod( node_pool: node_pool.Info, @@ -172,14 +175,17 @@ def run_metric_verification( metric_strategy.dag_id_suffix, ) + labels = metric_strategy.get_labels(tpu_info_output) compare_metric_values( cmd_values, monitoring_values, pod_name, metric_display_name=metric_strategy.dag_id_suffix, tolerance_percent=tolerance_for_metric, + labels=labels, ) + return True From 6b7899491ee8aac8ffec58231af4c1e61c8f75a5 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Fri, 24 Apr 2026 08:30:52 +0000 Subject: [PATCH 14/19] Change Dockerfile folder and name --- .../Dockerfile-nightly => dockerfiles/libtpu_nightly.Dockerfile} | 0 .../Dockerfile-stable => dockerfiles/libtpu_stable.Dockerfile} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename dags/tpu_observability/{configs/Dockerfile-nightly => dockerfiles/libtpu_nightly.Dockerfile} (100%) rename dags/tpu_observability/{configs/Dockerfile-stable => dockerfiles/libtpu_stable.Dockerfile} (100%) diff --git a/dags/tpu_observability/configs/Dockerfile-nightly b/dags/tpu_observability/dockerfiles/libtpu_nightly.Dockerfile similarity index 100% rename from dags/tpu_observability/configs/Dockerfile-nightly rename to dags/tpu_observability/dockerfiles/libtpu_nightly.Dockerfile diff --git a/dags/tpu_observability/configs/Dockerfile-stable b/dags/tpu_observability/dockerfiles/libtpu_stable.Dockerfile similarity index 100% rename from dags/tpu_observability/configs/Dockerfile-stable rename to dags/tpu_observability/dockerfiles/libtpu_stable.Dockerfile From aab432a6c08005839ea83cee86d635f8ed669cd2 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Fri, 15 May 2026 03:37:06 +0000 Subject: [PATCH 15/19] fix format --- dags/tpu_observability/tpu_info_metric.py | 2 -- dags/tpu_observability/tpu_info_metrics_dag.py | 2 -- 2 files changed, 4 deletions(-) diff --git a/dags/tpu_observability/tpu_info_metric.py b/dags/tpu_observability/tpu_info_metric.py index a777abfc4..6e2a2ccdb 100644 --- a/dags/tpu_observability/tpu_info_metric.py +++ b/dags/tpu_observability/tpu_info_metric.py @@ -91,7 +91,6 @@ def get_labels( return None - class _BaseSimplePointStrategy(BaseMetricStrategy): """Base strategy for parsing single numeric data points from Cloud Monitoring. @@ -359,7 +358,6 @@ def get_labels( return labels - class MemoryUsedStrategy(_BaseSimplePointStrategy): """Strategy for verifying Used HBM Memory.""" diff --git a/dags/tpu_observability/tpu_info_metrics_dag.py b/dags/tpu_observability/tpu_info_metrics_dag.py index ffac67cd3..77dfc0f6f 100644 --- a/dags/tpu_observability/tpu_info_metrics_dag.py +++ b/dags/tpu_observability/tpu_info_metrics_dag.py @@ -116,7 +116,6 @@ def compare_metric_values( ) - @task def get_tpu_info_metric_from_pod( node_pool: node_pool.Info, @@ -185,7 +184,6 @@ def run_metric_verification( labels=labels, ) - return True From 976fbe83feb05753c9594bc4178ff8813ace6a05 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Wed, 20 May 2026 02:05:44 +0000 Subject: [PATCH 16/19] update tpu_info_format_validation_dags and tpu_info_metrics_dag DAG timeout to 45 --- dags/common/scheduling_helper/scheduling_helper.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dags/common/scheduling_helper/scheduling_helper.py b/dags/common/scheduling_helper/scheduling_helper.py index f9a874ef0..2c8650909 100644 --- a/dags/common/scheduling_helper/scheduling_helper.py +++ b/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: { From 8a83f0602f8acdb94d3f0f076368d20fcaf37599 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Wed, 20 May 2026 02:26:20 +0000 Subject: [PATCH 17/19] Update libtpu image global parameter in `vm_resource` --- dags/common/vm_resource.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/dags/common/vm_resource.py b/dags/common/vm_resource.py index 2d9a57474..ed653e576 100644 --- a/dags/common/vm_resource.py +++ b/dags/common/vm_resource.py @@ -398,5 +398,11 @@ class DockerImage(enum.Enum): "gcr.io/tpu-prod-env-multipod/maxtext_post_training_nightly:" f"{datetime.datetime.today().strftime('%Y-%m-%d')}" ) - TPU_OBS_LIBTPU_NIGHTLY = "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-nightly:latest" - TPU_OBS_LIBTPU_STABLE = "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-stable:latest" + 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" + ) From eef9d279e649e2731633d6507dfdcd7403882df8 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Thu, 28 May 2026 09:04:21 +0000 Subject: [PATCH 18/19] image parameter fix --- .../jobset_ttr_drain_restart.py | 2 + .../jobset_ttr_kill_process.py | 2 +- .../jobset_ttr_node_pool_resize.py | 2 +- .../jobset_ttr_node_reboot.py | 2 + .../jobset_ttr_pod_delete.py | 2 +- dags/tpu_observability/jobset_ttr_rollback.py | 2 +- .../jobset_uptime_validation.py | 4 +- .../tpu_info_format_validation_dags.py | 44 +++++++++---------- .../tpu_observability/tpu_info_metrics_dag.py | 5 +-- .../tpu_sdk_monitoring_validation_dag.py | 5 +-- 10 files changed, 36 insertions(+), 34 deletions(-) diff --git a/dags/tpu_observability/jobset_ttr_drain_restart.py b/dags/tpu_observability/jobset_ttr_drain_restart.py index 311fbc36e..08ad20f42 100644 --- a/dags/tpu_observability/jobset_ttr_drain_restart.py +++ b/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( diff --git a/dags/tpu_observability/jobset_ttr_kill_process.py b/dags/tpu_observability/jobset_ttr_kill_process.py index ae3793d00..9ecc204d0 100644 --- a/dags/tpu_observability/jobset_ttr_kill_process.py +++ b/dags/tpu_observability/jobset_ttr_kill_process.py @@ -134,7 +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.TPU_OBS_LIBTPU_STABLE.value, + image=DockerImage.LIBTPU_STABLE.value, ) cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( diff --git a/dags/tpu_observability/jobset_ttr_node_pool_resize.py b/dags/tpu_observability/jobset_ttr_node_pool_resize.py index f1f7479be..193f95284 100644 --- a/dags/tpu_observability/jobset_ttr_node_pool_resize.py +++ b/dags/tpu_observability/jobset_ttr_node_pool_resize.py @@ -98,7 +98,7 @@ gcs_path=GCS_JOBSET_CONFIG_PATH, dag_name=DAG_ID, node_pool_selector=selector, - image=DockerImage.TPU_OBS_LIBTPU_STABLE.value, + image=DockerImage.LIBTPU_STABLE.value, ) cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( diff --git a/dags/tpu_observability/jobset_ttr_node_reboot.py b/dags/tpu_observability/jobset_ttr_node_reboot.py index 9e970f026..de883c295 100644 --- a/dags/tpu_observability/jobset_ttr_node_reboot.py +++ b/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( diff --git a/dags/tpu_observability/jobset_ttr_pod_delete.py b/dags/tpu_observability/jobset_ttr_pod_delete.py index e84252303..33c4227f7 100644 --- a/dags/tpu_observability/jobset_ttr_pod_delete.py +++ b/dags/tpu_observability/jobset_ttr_pod_delete.py @@ -94,7 +94,7 @@ gcs_path=GCS_JOBSET_CONFIG_PATH, dag_name=DAG_ID, node_pool_selector=selector, - image=DockerImage.TPU_OBS_LIBTPU_STABLE.value, + image=DockerImage.LIBTPU_STABLE.value, ) cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( diff --git a/dags/tpu_observability/jobset_ttr_rollback.py b/dags/tpu_observability/jobset_ttr_rollback.py index e282af9b5..2f3cce4ab 100644 --- a/dags/tpu_observability/jobset_ttr_rollback.py +++ b/dags/tpu_observability/jobset_ttr_rollback.py @@ -95,7 +95,7 @@ gcs_path=GCS_JOBSET_CONFIG_PATH, dag_name=DAG_ID, node_pool_selector=selector, - image=DockerImage.TPU_OBS_LIBTPU_STABLE.value, + image=DockerImage.LIBTPU_STABLE.value, ) cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( diff --git a/dags/tpu_observability/jobset_uptime_validation.py b/dags/tpu_observability/jobset_uptime_validation.py index a5cc011cd..79f9e0216 100644 --- a/dags/tpu_observability/jobset_uptime_validation.py +++ b/dags/tpu_observability/jobset_uptime_validation.py @@ -100,13 +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.TPU_OBS_LIBTPU_STABLE.value, + image=DockerImage.LIBTPU_STABLE.value, ) cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 77935b288..a5c125866 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -355,8 +355,8 @@ def generate_second_node_pool_name( """, ) as dag: docker_images = { - "stable": DockerImage.TPU_OBS_LIBTPU_STABLE.value, - "nightly": DockerImage.TPU_OBS_LIBTPU_NIGHTLY.value, + "stable": DockerImage.LIBTPU_STABLE.value, + "nightly": DockerImage.LIBTPU_NIGHTLY.value, } for machine in MachineConfigMap: @@ -365,17 +365,17 @@ def generate_second_node_pool_name( with TaskGroup(group_id=f"v{config.tpu_version.value}"): selector = jobset.generate_node_pool_selector(DAG_ID) - cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( - gcs_path=GCS_CONFIG_PATH, - dag_name=DAG_ID, - is_prod=composer_env.is_prod_env(), - machine_type=config.machine_version.value, - tpu_topology=config.tpu_topology, - node_pool_selector=selector, - ) + cluster_info = node_pool.build_node_pool_info_from_gcs_yaml( + gcs_path=GCS_CONFIG_PATH, + dag_name=DAG_ID, + is_prod=composer_env.is_prod_env(), + machine_type=config.machine_version.value, + tpu_topology=config.tpu_topology, + node_pool_selector=selector, + ) - cluster_info_2 = copy.deepcopy(cluster_info) - cluster_info_2.node_pool_name = f"{cluster_info.node_pool_name}-2" + cluster_info_2 = copy.deepcopy(cluster_info) + cluster_info_2.node_pool_name = f"{cluster_info.node_pool_name}-2" # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). @@ -409,17 +409,17 @@ def generate_second_node_pool_name( image=image_url, ) - startup = jobset.create_jobset_startup_tasks( - node_pool=cluster_info, - jobset_config=jobset_config, - workload_type=Workload.JAX_TPU_BENCHMARK, - ) + startup = jobset.create_jobset_startup_tasks( + node_pool=cluster_info, + jobset_config=jobset_config, + workload_type=Workload.JAX_TPU_BENCHMARK, + ) - outputs_of_tpu_info = ( - get_tpu_info_from_pod.override(task_id="get_tpu_info") - .partial(info=cluster_info) - .expand(pod_name=startup.running_pods) - ) + outputs_of_tpu_info = ( + get_tpu_info_from_pod.override(task_id="get_tpu_info") + .partial(info=cluster_info) + .expand(pod_name=startup.running_pods) + ) output_of_tpu_info = ( tpu_info.parse_tpu_info_output.override( diff --git a/dags/tpu_observability/tpu_info_metrics_dag.py b/dags/tpu_observability/tpu_info_metrics_dag.py index 77dfc0f6f..cf9da8539 100644 --- a/dags/tpu_observability/tpu_info_metrics_dag.py +++ b/dags/tpu_observability/tpu_info_metrics_dag.py @@ -281,8 +281,8 @@ def summarize_results( """, ) as dag: docker_images = { - "stable": DockerImage.TPU_OBS_LIBTPU_STABLE.value, - "nightly": DockerImage.TPU_OBS_LIBTPU_NIGHTLY.value, + "stable": DockerImage.LIBTPU_STABLE.value, + "nightly": DockerImage.LIBTPU_NIGHTLY.value, } for machine in MachineConfigMap: @@ -331,7 +331,6 @@ def summarize_results( with TaskGroup( # pylint: disable=unexpected-keyword-arg group_id=f"v{config.tpu_version.value}_{type_name}" ) as image_tg: - startup = jobset.create_jobset_startup_tasks( node_pool=cluster_info, jobset_config=jobset_config, diff --git a/dags/tpu_observability/tpu_sdk_monitoring_validation_dag.py b/dags/tpu_observability/tpu_sdk_monitoring_validation_dag.py index 4d62a0724..12767779d 100644 --- a/dags/tpu_observability/tpu_sdk_monitoring_validation_dag.py +++ b/dags/tpu_observability/tpu_sdk_monitoring_validation_dag.py @@ -125,8 +125,8 @@ def validate_monitoring_sdk(info: node_pool.Info, pod_name: str) -> None: """, ) as dag: docker_images = { - "stable": DockerImage.TPU_OBS_LIBTPU_STABLE.value, - "nightly": DockerImage.TPU_OBS_LIBTPU_NIGHTLY.value, + "stable": DockerImage.LIBTPU_STABLE.value, + "nightly": DockerImage.LIBTPU_NIGHTLY.value, } for machine in MachineConfigMap: @@ -165,7 +165,6 @@ def validate_monitoring_sdk(info: node_pool.Info, pod_name: str) -> None: with TaskGroup( # pylint: disable=unexpected-keyword-arg group_id=f"v{config.tpu_version.value}_{type_name}" ) as image_tg: - startup = jobset.create_jobset_startup_tasks( node_pool=cluster_info, jobset_config=jobset_config, From 964712338fc54ddf02b95979eab7448927061f17 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Thu, 9 Jul 2026 08:53:46 +0000 Subject: [PATCH 19/19] fix --- dags/tpu_observability/tpu_info_format_validation_dags.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index a5c125866..383e1a377 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -492,7 +492,7 @@ def generate_second_node_pool_name( ) chain( - startup.task_group, + *startup.tasks, outputs_of_tpu_info, output_of_tpu_info, verification_group,