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/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: { diff --git a/dags/common/vm_resource.py b/dags/common/vm_resource.py index 217371c21..ed653e576 100644 --- a/dags/common/vm_resource.py +++ b/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" + ) diff --git a/dags/tpu_observability/dockerfiles/libtpu_nightly.Dockerfile b/dags/tpu_observability/dockerfiles/libtpu_nightly.Dockerfile new file mode 100644 index 000000000..10983910a --- /dev/null +++ b/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 diff --git a/dags/tpu_observability/dockerfiles/libtpu_stable.Dockerfile b/dags/tpu_observability/dockerfiles/libtpu_stable.Dockerfile new file mode 100644 index 000000000..79498423a --- /dev/null +++ b/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 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 bbf45257a..9ecc204d0 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.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..193f95284 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.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 5745281bd..33c4227f7 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.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..2f3cce4ab 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.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..79f9e0216 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, @@ -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( diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index d38fab39c..383e1a377 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, @@ -91,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} @@ -293,6 +299,14 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]): ) +@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). with models.DAG( # pylint: disable=unexpected-keyword-arg @@ -340,23 +354,16 @@ 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.LIBTPU_STABLE.value, + "nightly": DockerImage.LIBTPU_NIGHTLY.value, + } + for machine in MachineConfigMap: config = machine.value - selector = jobset.generate_node_pool_selector( - "tpu-info-format-validation-dag" - ) - - # 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, - ) + 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, @@ -377,83 +384,121 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]): ) 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( - node_pool=cluster_info, - jobset_config=jobset_config, - workload_type=Workload.JAX_TPU_BENCHMARK, + create_second_node_pool = node_pool.create.override( + task_id="node_pool_2", + retries=2, + )( + node_pool=cluster_info_2, ) - 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) - ) + 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, + ) - output_of_tpu_info = ( - tpu_info.parse_tpu_info_output.override( - task_id="get_each_metric_table" + startup = jobset.create_jobset_startup_tasks( + node_pool=cluster_info, + jobset_config=jobset_config, + workload_type=Workload.JAX_TPU_BENCHMARK, ) - .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") - .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_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) - ) + output_of_tpu_info = ( + tpu_info.parse_tpu_info_output.override( + task_id="get_each_metric_table" + ) + .partial() + .expand(output=outputs_of_tpu_info) + ) - validate_runtime_metric = ( - validate_runtime_table.override(task_id="validate_runtime_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_tensorcore_metric = ( - validate_tensorcore_table.override( - task_id="validate_tensorcore_metric" + 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) ) - .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_runtime_metric = ( + validate_runtime_table.override( + task_id="validate_runtime_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_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) + ) + + 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 + ) + + chain( + *startup.tasks, + outputs_of_tpu_info, + output_of_tpu_info, + verification_group, + clean_up_workload, + ) + image_task_groups.append(image_tg) # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). @@ -476,27 +521,14 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]): setups=create_node_pool, ) - chain( - verify_table_amount_task, - [ - validate_tpu_chips_metric, - validate_runtime_metric, - validate_tensorcore_metric, - validate_latency_metric, - ], - ) + chain(cleanup_first_node_pool, cleanup_second_node_pool) - chain(create_first_node_pool, create_second_node_pool) + _ = [create_first_node_pool, create_second_node_pool] + chain(selector, cluster_info, cluster_info_2, create_node_pool) - chain(cleanup_first_node_pool, cleanup_second_node_pool) + current_node = create_node_pool + for tg in image_task_groups: + chain(current_node, tg) + current_node = tg - chain( - selector, - create_node_pool, - *startup.tasks, - outputs_of_tpu_info, - output_of_tpu_info, - verification_group, - clean_up_workload, - cleanup_node_pool, - ) + chain(current_node, cleanup_node_pool) diff --git a/dags/tpu_observability/tpu_info_metric.py b/dags/tpu_observability/tpu_info_metric.py index e1d98e59f..6e2a2ccdb 100644 --- a/dags/tpu_observability/tpu_info_metric.py +++ b/dags/tpu_observability/tpu_info_metric.py @@ -81,6 +81,15 @@ 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. @@ -305,15 +314,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( @@ -327,6 +334,29 @@ 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.""" @@ -443,7 +473,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 +485,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 +499,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 +513,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" @@ -540,8 +570,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/tpu_info_metrics_dag.py b/dags/tpu_observability/tpu_info_metrics_dag.py index 120c0de25..cf9da8539 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, @@ -60,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): @@ -75,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)): @@ -92,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, @@ -171,12 +174,14 @@ 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 @@ -275,6 +280,11 @@ def summarize_results( pools. """, ) as dag: + docker_images = { + "stable": DockerImage.LIBTPU_STABLE.value, + "nightly": DockerImage.LIBTPU_NIGHTLY.value, + } + for machine in MachineConfigMap: config = machine.value @@ -316,72 +326,84 @@ 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: + startup = jobset.create_jobset_startup_tasks( + node_pool=cluster_info, + jobset_config=jobset_config, + workload_type=Workload.JAX_TPU_BENCHMARK, + ) - verification_results = {} - all_verification_groups = [] + verification_results = {} + all_verification_groups = [] - for strategy in ALL_METRIC_STRATEGIES: - group_id = f"verify_{strategy.dag_id_suffix}" + for strategy in ALL_METRIC_STRATEGIES: + group_id = f"verify_{strategy.dag_id_suffix}" - 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, + 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) ) - .expand(pod_name=startup.running_pods) - ) - tpu_info_metric_output = ( - tpu_info.parse_tpu_info_output.override( - task_id="get_each_metric_table" + tpu_info_metric_output = ( + tpu_info.parse_tpu_info_output.override( + task_id="get_each_metric_table" + ) + .partial() + .expand(output=tpu_info_metric_outputs) ) - .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, - ) - .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 +422,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..12767779d 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,86 @@ def validate_monitoring_sdk(info: node_pool.Info, pod_name: str) -> None: and hardware devices inside the container. """, ) as dag: + docker_images = { + "stable": DockerImage.LIBTPU_STABLE.value, + "nightly": DockerImage.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) diff --git a/dags/tpu_observability/utils/tpu_info_util.py b/dags/tpu_observability/utils/tpu_info_util.py index c408463d4..c552447a4 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,95 @@ 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) 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."