diff --git a/.github/workflows/dag-check.yml b/.github/workflows/dag-check.yml index 96fcb76bf..ac3982d0e 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/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index c57e99460..a7dbb5c13 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -13,9 +13,18 @@ # limitations under the License. """ -A DAG orchestrates the process of verifying TensorCore utilization metrics. - -This is done by comparing data from Cloud Logging and Cloud Monitoring. +tpu_info_validation_dag: +A comprehensive DAG that orchestrates the end-to-end validation of the +`tpu-info` observability tool. It performs two primary types of +verification: + +1. Format Validation: Parses the tool's raw output into structured tables + (TPU Chips, Runtime Utilization, TensorCore Utilization, and Latency) + and validates row counts and data integrity using regex. + +2. CLI Validation: Verifies command-line options (`-help`, `--version`, + `--process`) to ensure metadata, help documentation, and + process-to-chip mapping are functional inside live TPU worker pods. """ import datetime @@ -52,34 +61,6 @@ SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) -@task -def get_tpu_info_from_pod(info: node_pool.Info, pod_name: str) -> str: - """ - Executes the `tpu-info` command in a specified pod and returns its output. - - This task uses kubectl to run the 'tpu-info' command inside the given pod - in the 'default' namespace. The output of the command is captured and - returned. - - Args: - kubeconfig: The path to the kubeconfig file. - pod_name: The name of the pod to execute the command in. - - Returns: - The standard output from the 'tpu-info' command. - """ - with tempfile.NamedTemporaryFile() as temp_config_file: - env = os.environ.copy() - env["KUBECONFIG"] = temp_config_file.name - - cmd = " && ".join([ - jobset.Command.get_credentials_command(info), - f"kubectl exec {pod_name} -n default -- tpu-info", - ]) - - return subprocess.run_exec(cmd, env=env) - - @task def verify_table_amount(tpu_info_output: list[tpu_info.Table]): """ @@ -292,6 +273,27 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]): ) +@task +def validate_tpu_info_patterns(output: str, cmd_name: str): + """Matches output against patterns defined in a local spec.""" + patterns_map = { + tpu_info.TpuInfoCmd.HELP.value: ["--streaming", "--rate RATE"], + tpu_info.TpuInfoCmd.VERSION.value: [ + "tpu-info version:", + "libtpu version:", + ], + tpu_info.TpuInfoCmd.PROCESS.value: [ + "TPU Process Info", + "/dev/vfio/", + "python", + ], + } + patterns = patterns_map.get(cmd_name, []) + for pattern in patterns: + if pattern not in output: + raise AssertionError(f"Cmd '{cmd_name}' missing pattern: {pattern}") + + # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with models.DAG( # pylint: disable=unexpected-keyword-arg @@ -303,37 +305,31 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]): catchup=False, tags=["gke", "tpu-observability", "tpu-info", "TPU", "v6e-16"], description=( - "This DAG verifies the format of the tables in the tpu-info output " - "using tpu-info CLI tool. It includes 4 tables: TPU Chips, TPU " - "Runtime Utilization, TensorCore Utilization, and TPU Buffer Transfer " - "Latency." + "Validates the tpu-info CLI tool by verifying its command-line options. " + "It ensures that help documentation, version metadata, and process " + "mapping are functional across TPU worker pods." ), doc_md=""" - # Format Validation DAG - # This DAG verifies the format of the tables in the tpu-info output. + # TPU Info CLI Validation DAG + This DAG automates the end-to-end validation of the `tpu-info` observability tool. ### Description - This DAG automates the validation of the tpu-info command-line tool's - output format.It verifies the structure and content of key metric tables, - including "TPU Chips", "TPU Runtime Utilization", "TensorCore - Utilization", and "TPU Buffer Transfer Latency", by running the tool on a - live GKE cluster with TPU node pools. - - ### Prerequisites - This test requires an existing GKE cluster. - A pre-built Docker image containing the necessary jax, libtpu, and - tpu-info packages must also be available in a repository accessible - by the GKE cluster. + The test verifies that the `tpu-info` CLI correctly interprets various **command-line options** (e.g., `-help`, `--version`, `--process`) and returns the expected metadata or + metric structures. This ensures the tool is correctly installed and compatible + with the current TPU runtime environment. + + ### Validation Scope + The DAG executes the following command options inside the TPU pods: + * **Help Option (`-help`)**: Verifies the presence of required flags like `--streaming` and `--rate`. + * **Version Option (`--version`)**: Validates that the tool reports the correct version and `libtpu` metadata. + * **Process Option (`--process`)**: Ensures that active PIDs can be mapped to specific TPU chips. ### Procedures - The DAG begins by creating temporary GKE TPU node pools for the test. - Once the node pools are running, it schedules a Kubernetes JobSet and - waits for the pods to become active. It then executes the tpu-info - command within these pods to capture the raw text output. This output is - parsed into structured tables, and a series of validation tasks check - each table for the correct structure, row counts, and data formats. - Finally, regardless of the test outcome, the DAG cleans up all created - resources, including the JobSet and the temporary node pools. + 1. **Environment Setup**: Dynamically creates GKE TPU node pools. + 2. **Workload Deployment**: Runs a JAX benchmark JobSet to generate active TPU processes. + 3. **CLI Execution**: Iterates through defined `TpuInfoCmd` options within the worker pods. + 4. **Pattern Matching**: Performs regex and string validation on the CLI output to ensure accuracy. + 5. **Teardown**: Automatically cleans up all Kubernetes resources and node pools to minimize costs. """, ) as dag: for machine in MachineConfigMap: @@ -421,9 +417,11 @@ def generate_second_node_pool_name( job_apply_time=apply_time, ) - outputs_of_tpu_info = ( - get_tpu_info_from_pod.override(task_id="get_tpu_info") - .partial(info=cluster_info) + raw_metric_data = ( + tpu_info.get_tpu_info_from_pod.override(task_id="get_tpu_info") + .partial( + info=cluster_info, cmd_str=tpu_info.TpuInfoCmd.TPU_INFO.value + ) .expand(pod_name=running_pods) ) @@ -432,7 +430,20 @@ def generate_second_node_pool_name( task_id="get_each_metric_table" ) .partial() - .expand(output=outputs_of_tpu_info) + .expand(output=raw_metric_data.map(lambda x: x["output"])) + ) + + cli_raw_results = ( + tpu_info.get_tpu_info_from_pod.override(task_id="get_cli_output") + .partial(info=cluster_info) + .expand( + pod_name=running_pods, + cmd_str=[ + tpu_info.TpuInfoCmd.HELP.value, + tpu_info.TpuInfoCmd.VERSION.value, + tpu_info.TpuInfoCmd.PROCESS.value, + ], + ) ) # Keyword arguments are generated dynamically at runtime (pylint does not @@ -472,6 +483,10 @@ def generate_second_node_pool_name( .expand(tpu_info_output=output_of_tpu_info) ) + cli_validation = validate_tpu_info_patterns.override( + task_id="validate_cli_output" + ).expand_kwargs(cli_raw_results) + clean_up_workload = jobset.end_workload.override( task_id="clean_up_workload", trigger_rule=TriggerRule.ALL_DONE )( @@ -509,10 +524,11 @@ def generate_second_node_pool_name( validate_runtime_metric, validate_tensorcore_metric, validate_latency_metric, + cli_validation, ], ) - 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) @@ -525,8 +541,9 @@ def generate_second_node_pool_name( apply_time, running_pods, wait_for_job_start, - outputs_of_tpu_info, + raw_metric_data, output_of_tpu_info, + cli_raw_results, verification_group, clean_up_workload, cleanup_node_pool, diff --git a/dags/tpu_observability/utils/tpu_info_util.py b/dags/tpu_observability/utils/tpu_info_util.py index c408463d4..f8a3a2b0c 100644 --- a/dags/tpu_observability/utils/tpu_info_util.py +++ b/dags/tpu_observability/utils/tpu_info_util.py @@ -1,11 +1,25 @@ """Utility for parsing the output of the 'tpu-info' command.""" -from dataclasses import dataclass -from enum import auto -from enum import IntEnum +import os import re +import tempfile +from dataclasses import dataclass +from enum import auto, Enum, IntEnum from airflow.decorators import task +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 + + +class TpuInfoCmd(Enum): + """Defines the available tpu-info CLI commands.""" + + TPU_INFO = "tpu-info" + HELP = "tpu-info -help" + VERSION = "tpu-info --version" + PROCESS = "tpu-info --process" + # A type alias for a parsed row, mapping column headers to their values. _TableRow = dict[str, str] @@ -23,7 +37,8 @@ def parse_body(self): """Parses the raw_body string to populate the structured body attribute.""" class TableLineIndex(IntEnum): - """Below is an example of the text returned by tpu-info, formatted as a table. + """Below is an example of the text returned by tpu-info, + formatted as a table. TPU Chips ┏━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━┓ @@ -96,6 +111,23 @@ def parse_tpu_info_output(output: str) -> list[Table]: return parsed_tables +@task +def get_tpu_info_from_pod( + info: node_pool.Info, pod_name: str, cmd_str: str +) -> dict: + """Executes command and returns a dict containing metadata.""" + with tempfile.NamedTemporaryFile() as temp_config_file: + env = os.environ.copy() + env["KUBECONFIG"] = temp_config_file.name + cmd = " && ".join([ + jobset.Command.get_credentials_command(info), + f"kubectl exec {pod_name} -n default -- {cmd_str}", + ]) + output = subprocess.run_exec(cmd, env=env) + + return {"output": output, "cmd_name": cmd_str} + + if __name__ == "__main__": full_output = """ Libtpu version: 0.0.23 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."