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/vm_resource.py b/dags/common/vm_resource.py index 274c06024..f056575b2 100644 --- a/dags/common/vm_resource.py +++ b/dags/common/vm_resource.py @@ -124,6 +124,8 @@ class Zone(enum.Enum): US_EAST5_C = "us-east5-c" # reserved v5e in tpu-prod-env-multipod US_WEST4_B = "us-west4-b" + US_WEST4_A = "us-west4-a" + US_EAST4_B = "us-east4-b" # reserved v5e in cloud-tpu-inference-test US_WEST1_C = "us-west1-c" # reserved a3+ cluster in supercomputer-testing diff --git a/dags/inference/configs/jetstream_benchmark_serving_gce_config.py b/dags/inference/configs/jetstream_benchmark_serving_gce_config.py index ec5112356..fe8a0d735 100644 --- a/dags/inference/configs/jetstream_benchmark_serving_gce_config.py +++ b/dags/inference/configs/jetstream_benchmark_serving_gce_config.py @@ -25,6 +25,7 @@ PROJECT_NAME = Project.CLOUD_ML_AUTO_SOLUTIONS.value RUNTIME_IMAGE = RuntimeVersion.TPU_UBUNTU2204_BASE.value GCS_SUBFOLDER_PREFIX = test_owner.Team.INFERENCE.value +VENV_DIR = "venv-312" def get_config( @@ -70,11 +71,14 @@ def get_config( # Make the PATH change permanent for subsequent sessions (optional, but good practice) "echo 'export PATH=\"$HOME/.local/bin:$PATH\"' >> ~/.bashrc", "source ~/.bashrc", - "uv venv --python 3.12 venv-312 --seed", - "source venv-312/bin/activate", + f"uv venv --python 3.12 {VENV_DIR} --seed --clear", + f"source {VENV_DIR}/bin/activate", "pip install uv", "uv pip install maxtext --resolution=lowest", "install_maxtext_github_deps", + "uv pip install rouge-score", + "sudo apt-get install -y git-lfs", + "git lfs install", ) set_up_cmds += setup_maxtext_cmds @@ -113,15 +117,19 @@ def get_config( # Let gcs path be directly used, else use maxtext/assets dir if not model_configs["tokenizer"].startswith("gs://"): - tokenizer_path = f"assets/{model_configs['tokenizer']}" - full_tokenizer_path = f"maxtext/assets/{model_configs['tokenizer']}" + tokenizer_path = ( + f"src/maxtext/assets/tokenizers/{model_configs['tokenizer']}" + ) + full_tokenizer_path = ( + f"maxtext/src/maxtext/assets/tokenizers/{model_configs['tokenizer']}" + ) else: tokenizer_path = model_configs["tokenizer"] full_tokenizer_path = model_configs["tokenizer"] run_model_cmds = ( # Start virtual environment - "source .env/bin/activate", + f"source {VENV_DIR}/bin/activate", "wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json > /dev/null 2>&1", # Get commit hash of the maxtext and jetstream repos f"export METADATA_DICT='{json.dumps(additional_metadata_dict)}'", @@ -131,58 +139,37 @@ def get_config( 'export METADATA_DICT=$(jq -c \'. + { "jetstream_commit_hash": $newVal}\' --arg newVal ${JETSTREAM_COMMIT_HASH} <<<"$METADATA_DICT")', ### Benchmark "cd maxtext", - # Configure flags - f"export MODEL_NAME={model_configs['model_name']}", - f"export TOKENIZER_PATH={tokenizer_path}", - f"export WEIGHT_DTYPE={model_configs['weight_dtype']}", - f"export SCAN_LAYERS={model_configs['scan_layers']}", - f"export MAX_PREFILL_PREDICT_LENGTH={model_configs['max_prefill_predict_length']}", - f"export MAX_TARGET_LENGTH={model_configs['max_target_length']}", - f"export ATTENTION={model_configs['attention']}", - f"export ICI_FSDP_PARALLELISM={model_configs['ici_fsdp_parallelism']}", - f"export ICI_AUTOREGRESSIVE_PARALLELISM={model_configs['ici_autoregressive_parallelism']}", - f"export ICI_TENSOR_PARALLELISM={model_configs['ici_tensor_parallelism']}", - f"export UNSCANNED_CKPT_PATH={model_configs['checkpoint']}", - "export LOAD_PARAMETERS_PATH=${UNSCANNED_CKPT_PATH}", - f"export QUANTIZATION={model_configs['quantization']}", - f"export QUANTIZE_KVCACHE={model_configs['quantize_kvcache']}", - f"export KV_QUANT_DTYPE={model_configs['kv_quant_dtype']}", - f"export PER_DEVICE_BATCH_SIZE={model_configs['per_device_batch_size']}", - f"export PREFILL_CACHE_AXIS_ORDER={model_configs['prefill_cache_axis_order']}", - f"export AR_CACHE_AXIS_ORDER={model_configs['ar_cache_axis_order']}", - f"export COMPUTE_AXIS_ORDER={model_configs['compute_axis_order']}", - f"export RESHAPE_Q={model_configs['reshape_q']}", - f"export KV_QUANT_AXIS={model_configs['kv_quant_axis']}", # Start JetStream MaxText server in the background - """python3 -m MaxText.maxengine_server \ - src/maxtext/configs/inference/inference_jetstream.yml \ - model_name=${MODEL_NAME} \ - tokenizer_path=${TOKENIZER_PATH} \ - weight_dtype=${WEIGHT_DTYPE} \ - scan_layers=${SCAN_LAYERS} \ - max_prefill_predict_length=${MAX_PREFILL_PREDICT_LENGTH} \ - max_target_length=${MAX_TARGET_LENGTH} \ - attention=${ATTENTION} \ - ici_fsdp_parallelism=${ICI_FSDP_PARALLELISM} \ - ici_autoregressive_parallelism=${ICI_AUTOREGRESSIVE_PARALLELISM} \ - ici_tensor_parallelism=${ICI_TENSOR_PARALLELISM} \ - load_parameters_path=${LOAD_PARAMETERS_PATH} \ - quantization=${QUANTIZATION} \ - quantize_kvcache=${QUANTIZE_KVCACHE} \\""" - + ( - """kv_quant_dtype=${KV_QUANT_DTYPE} \\""" - if model_configs["kv_quant_dtype"] - else "" - ) - + """per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \ - prefill_cache_axis_order=${PREFILL_CACHE_AXIS_ORDER} \ - ar_cache_axis_order=${AR_CACHE_AXIS_ORDER} \ - compute_axis_order=${COMPUTE_AXIS_ORDER} \ - reshape_q=${RESHAPE_Q} \ - kv_quant_axis=${KV_QUANT_AXIS} &""", + f"""python3 -m MaxText.maxengine_server \\ + src/maxtext/configs/inference/inference_jetstream.yml \\ + model_name='{model_configs['model_name']}' \\ + tokenizer_path='{tokenizer_path}' \\ + weight_dtype='{model_configs['weight_dtype']}' \\ + scan_layers='{model_configs['scan_layers']}' \\ + max_prefill_predict_length='{model_configs['max_prefill_predict_length']}' \\ + max_target_length='{model_configs['max_target_length']}' \\ + attention='{model_configs['attention']}' \\ + ici_fsdp_parallelism='{model_configs['ici_fsdp_parallelism']}' \\ + ici_autoregressive_parallelism='{model_configs['ici_autoregressive_parallelism']}' \\ + ici_tensor_parallelism='{model_configs['ici_tensor_parallelism']}' \\ + load_parameters_path='{model_configs['checkpoint']}' \\ + quantization='"{model_configs.get('quantization') or ""}"' \\ + quantize_kvcache='{model_configs['quantize_kvcache']}' \\ + kv_quant_dtype='"{model_configs.get('kv_quant_dtype') or ""}"' \\ + per_device_batch_size='{model_configs['per_device_batch_size']}' \\ + prefill_cache_axis_order='{model_configs['prefill_cache_axis_order']}' \\ + ar_cache_axis_order='{model_configs['ar_cache_axis_order']}' \\ + compute_axis_order='{model_configs['compute_axis_order']}' \\ + reshape_q='{model_configs['reshape_q']}' \\ + kv_quant_axis='"{model_configs.get('kv_quant_axis') or ""}"' &""", "cd ..", # Give server time to start f"sleep {model_configs['sleep_time']}", + "cd JetStream", + # Since we change to the Jetstream dir as root, we need to change the ownership of the dir to the user running the script + "sudo chown -R $USER:$USER $HOME/JetStream", + "git lfs pull", + "cd ..", # Run benchmark, run eval, save benchmark and eval results, and save predictions to /tmp/request-outputs.json f"""python JetStream/benchmarks/benchmark_serving.py \ --tokenizer {full_tokenizer_path} \ @@ -233,7 +220,6 @@ def get_config( json_lines=metric_config.JSONLinesConfig("metric_report.jsonl"), use_runtime_generated_gcs_folder=True, ) - return task.run_queued_resource_test( task_test_config=job_test_config, task_gcp_config=job_gcp_config, diff --git a/dags/inference/jetstream_inference_e2e.py b/dags/inference/jetstream_inference_e2e.py index 1b8c6a012..01ef07789 100644 --- a/dags/inference/jetstream_inference_e2e.py +++ b/dags/inference/jetstream_inference_e2e.py @@ -16,8 +16,7 @@ import datetime from airflow import models -from dags.common.vm_resource import TpuVersion -from dags.inference.maxtext_model_config_generator import generate_model_configs +from dags.inference.maxtext_model_config_generator import generate_model_configs, TpuConfig """A JetStream inference E2E test (JAX nightly, no schedule) DAG. @@ -72,7 +71,7 @@ "jetstream_branch": "", "sleep_time": 360, "time_out_in_min": 60, - "tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)], + "tpu_version_cores": [(TpuConfig.V5E, 8), (TpuConfig.TRILLIUM, 8)], "model_name": LLAMA2_7B, "tokenizer": "tokenizer.llama2", "weight_dtype": "bfloat16", @@ -100,7 +99,7 @@ "jetstream_branch": "", "sleep_time": 360, "time_out_in_min": 60, - "tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)], + "tpu_version_cores": [(TpuConfig.V5E, 8), (TpuConfig.TRILLIUM, 8)], "model_name": GEMMA_7B, "tokenizer": "tokenizer.gemma", "weight_dtype": "bfloat16", diff --git a/dags/inference/maxtext_inference.py b/dags/inference/maxtext_inference.py index 5e01ded8b..ce575fca6 100644 --- a/dags/inference/maxtext_inference.py +++ b/dags/inference/maxtext_inference.py @@ -18,8 +18,8 @@ import numpy as np from airflow import models from dags import composer_env -from dags.common.vm_resource import TpuVersion -from dags.inference.maxtext_model_config_generator import generate_model_configs + +from dags.inference.maxtext_model_config_generator import generate_model_configs, TpuConfig USER_PREFIX = "" @@ -117,7 +117,10 @@ "jetstream_branch": jetstream_branch, "sleep_time": 360, "time_out_in_min": 120, - "tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)], + "tpu_version_cores": [ + (TpuConfig.temp_V5E, 8), + (TpuConfig.temp_TRILLIUM, 8), + ], "model_name": LLAMA2_7B, "tokenizer": "tokenizer.llama2", "weight_dtype": "bfloat16", @@ -155,7 +158,10 @@ "jetstream_branch": jetstream_branch, "sleep_time": 360, "time_out_in_min": 120, - "tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)], + "tpu_version_cores": [ + (TpuConfig.temp_V5E, 8), + (TpuConfig.temp_TRILLIUM, 8), + ], "model_name": LLAMA2_13B, "tokenizer": "tokenizer.llama2", "weight_dtype": "bfloat16", @@ -193,7 +199,10 @@ "jetstream_branch": jetstream_branch, "sleep_time": 360, "time_out_in_min": 240, - "tpu_version_cores": [(TpuVersion.V5P, 8), (TpuVersion.TRILLIUM, 8)], + "tpu_version_cores": [ + (TpuConfig.V5P, 8), + (TpuConfig.temp_TRILLIUM, 8), + ], "model_name": LLAMA2_70B, "tokenizer": "tokenizer.llama2", "weight_dtype": "bfloat16", @@ -230,7 +239,10 @@ "jetstream_branch": jetstream_branch, "sleep_time": 360, "time_out_in_min": 120, - "tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)], + "tpu_version_cores": [ + (TpuConfig.temp_V5E, 8), + (TpuConfig.temp_TRILLIUM, 8), + ], "model_name": GEMMA_7B, "tokenizer": "tokenizer.gemma", "weight_dtype": "bfloat16", @@ -269,7 +281,10 @@ "jetstream_branch": jetstream_branch, "sleep_time": 240, "time_out_in_min": 240, - "tpu_version_cores": [(TpuVersion.V5P, 8), (TpuVersion.TRILLIUM, 8)], + "tpu_version_cores": [ + (TpuConfig.V5P, 8), + (TpuConfig.temp_TRILLIUM, 8), + ], "model_name": MIXTRAL_8_7B, "tokenizer": "gs://maxtext-external/mixtral-8x7B-v0.1-Instruct/tokenizer.mistral", "weight_dtype": "bfloat16", @@ -576,7 +591,6 @@ }, } ) - # run_configs = [ # f"{LLAMA2_7B}-{BASE_MODE}-{W_BF16_KV_BF16}", # f"{LLAMA2_7B}-{BASE_MODE}-{W_INT8_KV_INT8}", diff --git a/dags/inference/maxtext_inference_microbenchmark.py b/dags/inference/maxtext_inference_microbenchmark.py index 4be520d2c..ebec12aa8 100644 --- a/dags/inference/maxtext_inference_microbenchmark.py +++ b/dags/inference/maxtext_inference_microbenchmark.py @@ -19,8 +19,9 @@ import itertools import numpy from airflow import models -from dags.common.vm_resource import TpuVersion, Zone, Project, V5_NETWORKS, V5E_SUBNETWORKS, V5P_SUBNETWORKS, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK + from dags.inference.configs import maxtext_inference_microbenchmark_gce_config +from dags.inference.maxtext_model_config_generator import TpuConfig from dags.multipod.configs.common import SetupMode USER_PREFIX = "" @@ -110,6 +111,8 @@ def generate_model_configs( tpu_version, tpu_cores, ): + config = tpu_version.value + model_configs = {} model_configs["model_config_name"] = model_config_name @@ -177,32 +180,18 @@ def generate_model_configs( test_name = f"{test_name_prefix}-{test_run_tag}" model_configs["run_name"] = test_run_tag - if tpu_version == TpuVersion.V5E: - # v5e benchmarks - project_name = Project.TPU_PROD_ENV_AUTOMATED.value - zone = Zone.US_EAST1_C.value - network = V5_NETWORKS - subnetwork = V5E_SUBNETWORKS - runtime_version = RuntimeVersion.V2_ALPHA_TPUV5_LITE.value - if tpu_version == TpuVersion.TRILLIUM: - project_name = Project.TPU_PROD_ENV_AUTOMATED.value - zone = Zone.US_EAST5_A.value - network = V6E_GCE_NETWORK - subnetwork = V6E_GCE_SUBNETWORK - runtime_version = RuntimeVersion.V2_ALPHA_TPUV6.value - maxtext_kv_cache_layout_optimization = ( maxtext_inference_microbenchmark_gce_config.config( - tpu_version=tpu_version, + tpu_version=config.tpu_version_name, tpu_cores=tpu_cores, - tpu_zone=zone, + tpu_zone=config.zone, time_out_in_min=sweep_model_configs["time_out_in_min"], test_name=test_name, test_mode=SetupMode.STABLE, - project_name=project_name, - runtime_version=runtime_version, - network=network, - subnetwork=subnetwork, + project_name=config.project_name, + runtime_version=config.runtime_version, + network=config.network, + subnetwork=config.subnetwork, is_tpu_reserved=True, model_configs=model_configs, maxtext_branch=model_configs["maxtext_branch"], @@ -247,7 +236,7 @@ def generate_model_configs( if not MAXTEXT_BRANCH else f"-b {MAXTEXT_BRANCH}", "sleep_time": 60, - "tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)], + "tpu_version_cores": [(TpuConfig.V5E, 8), (TpuConfig.TRILLIUM, 8)], "model_name": LLAMA2_7B, "tokenizer": "tokenizer.llama2", "weight_dtype": "bfloat16", diff --git a/dags/inference/maxtext_model_config_generator.py b/dags/inference/maxtext_model_config_generator.py index 7f597f1f3..239bb56ba 100644 --- a/dags/inference/maxtext_model_config_generator.py +++ b/dags/inference/maxtext_model_config_generator.py @@ -14,11 +14,66 @@ """A helper to generate maxtext model configs.""" +from dataclasses import dataclass +from enum import Enum from dags.common.vm_resource import TpuVersion, Zone, Project, V5_NETWORKS, V5E_SUBNETWORKS, V5P_SUBNETWORKS, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK from dags.inference.configs import jetstream_benchmark_serving_gce_config from dags.multipod.configs.common import SetupMode +@dataclass +class TpuResource: + tpu_version_name: TpuVersion + runtime_version: str + project_name: str + zone: str + network: str + subnetwork: str + + +class TpuConfig(Enum): + V5E = TpuResource( + tpu_version_name=TpuVersion.V5E, + runtime_version=RuntimeVersion.V2_ALPHA_TPUV5_LITE.value, + project_name=Project.TPU_PROD_ENV_AUTOMATED.value, + zone=Zone.US_EAST1_C.value, + network=V5_NETWORKS, + subnetwork=V5E_SUBNETWORKS, + ) + V5P = TpuResource( + tpu_version_name=TpuVersion.V5P, + runtime_version=RuntimeVersion.V2_ALPHA_TPUV5.value, + project_name=Project.TPU_PROD_ENV_AUTOMATED.value, + zone=Zone.US_EAST5_A.value, + network=V5_NETWORKS, + subnetwork=V5P_SUBNETWORKS, + ) + TRILLIUM = TpuResource( + tpu_version_name=TpuVersion.TRILLIUM, + runtime_version=RuntimeVersion.V2_ALPHA_TPUV6.value, + project_name=Project.TPU_PROD_ENV_AUTOMATED.value, + zone=Zone.US_EAST5_A.value, + network=V6E_GCE_NETWORK, + subnetwork=V6E_GCE_SUBNETWORK, + ) + temp_V5E = TpuResource( + tpu_version_name=TpuVersion.V5E, + runtime_version=RuntimeVersion.V2_ALPHA_TPUV5_LITE.value, + project_name="cienet-cmcs", + zone=Zone.US_WEST4_A.value, + network="projects/cienet-cmcs/global/networks/mas-test", + subnetwork="projects/cienet-cmcs/regions/us-west4/subnetworks/mas-test", + ) + temp_TRILLIUM = TpuResource( + tpu_version_name=TpuVersion.TRILLIUM, + runtime_version=RuntimeVersion.V2_ALPHA_TPUV6.value, + project_name=Project.TPU_PROD_ENV_AUTOMATED.value, + zone=Zone.SOUTHAMERICA_WEST1_A.value, + network=V5_NETWORKS, + subnetwork="projects/tpu-prod-env-automated/regions/southamerica-west1-a/subnetworks/mas-test", + ) + + def generate_model_configs( test_name_prefix, model_config_name, @@ -99,37 +154,20 @@ def generate_model_configs( test_name = f"{test_name_prefix}-{test_run_tag}" - if tpu_version == TpuVersion.V5E: - # v5e benchmarks - project_name = Project.TPU_PROD_ENV_AUTOMATED.value - zone = Zone.US_EAST1_C.value - network = V5_NETWORKS - subnetwork = V5E_SUBNETWORKS - runtime_version = RuntimeVersion.V2_ALPHA_TPUV5_LITE.value - elif tpu_version == TpuVersion.V5P: - zone = Zone.US_EAST5_A.value - runtime_version = RuntimeVersion.V2_ALPHA_TPUV5.value - project_name = Project.TPU_PROD_ENV_AUTOMATED.value - network = V5_NETWORKS - subnetwork = V5P_SUBNETWORKS - elif tpu_version == TpuVersion.TRILLIUM: - zone = Zone.US_EAST5_A.value - runtime_version = RuntimeVersion.V2_ALPHA_TPUV6.value - project_name = Project.TPU_PROD_ENV_AUTOMATED.value - network = V6E_GCE_NETWORK - subnetwork = V6E_GCE_SUBNETWORK + config = tpu_version.value + jetstream_benchmark_serving = ( jetstream_benchmark_serving_gce_config.get_config( - tpu_version=tpu_version, + tpu_version=config.tpu_version_name, tpu_cores=tpu_cores, - tpu_zone=zone, + tpu_zone=config.zone, time_out_in_min=sweep_model_configs["time_out_in_min"], test_name=test_name, test_mode=SetupMode.STABLE, - project_name=project_name, - runtime_version=runtime_version, - network=network, - subnetwork=subnetwork, + project_name=config.project_name, + runtime_version=config.runtime_version, + network=config.network, + subnetwork=config.subnetwork, is_tpu_reserved=True, model_configs=model_configs, maxtext_branch=model_configs["maxtext_branch"], 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." diff --git a/xlml/utils/startup_script.py b/xlml/utils/startup_script.py index 00c3acba4..20d720fa9 100644 --- a/xlml/utils/startup_script.py +++ b/xlml/utils/startup_script.py @@ -15,59 +15,73 @@ """Utility for startup scripts.""" import shlex +import textwrap + + +block_background_upgrade_services_script = textwrap.dedent( + """\ + sudo systemctl stop apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service unattended-upgrades.service + sudo systemctl mask apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service unattended-upgrades.service + sudo fuser -k -KILL /usr/bin/apt /usr/bin/apt-get /usr/bin/dpkg || true + sudo fuser -k -KILL /var/lib/dpkg/lock /var/lib/dpkg/lock-frontend /var/cache/apt/archives/lock || true + sudo dpkg --configure -a +""" +) def generate_startup_script(main_command: str) -> str: escaped_command = shlex.quote(main_command) - return f""" -set -o pipefail -bash -c {escaped_command} 2>&1 | tee /tmp/logs & -pid=$! -echo $pid > /tmp/main_process_id.txt -wait $pid -exit_status=$? -echo $exit_status > /tmp/process_exit_status.txt -""" + return block_background_upgrade_services_script + textwrap.dedent( + f"""\ + set -o pipefail + bash -c {escaped_command} 2>&1 | tee /tmp/logs & + pid=$! + echo $pid > /tmp/main_process_id.txt + wait $pid + exit_status=$? + echo $exit_status > /tmp/process_exit_status.txt + """ + ) def monitor_startup_script() -> str: return """ -# File paths -pid_file="/tmp/main_process_id.txt" -status_file="/tmp/process_exit_status.txt" -log_file="/tmp/logs" - -echo "LOGGER: Waiting for the workload to show up in $pid_file" - -# Wait until the PID file exists -while [ ! -f "$pid_file" ]; do - sleep 1 -done - -# Extract PID from pid_file -pid=$(cat "$pid_file") - -echo "LOGGER: Streaming worker 0 logs." -# Tail the log file and terminate when the process with $pid exits -tail -f --pid=$pid --retry $log_file - -echo "LOGGER: Process $pid has finished." - -# Check if status_file contain any number -if ! grep -q '[0-9]' "$status_file"; then - echo "LOGGER: The status file contains no exit status." - exit 1 -fi - -# Read and output the exit status -exit_status=$(cat "$status_file") - -# Check the exit_status -if [ "$exit_status" -eq 0 ]; then - echo "LOGGER: The process exited successfully." - exit 0 -else - echo "LOGGER: The process failed with exit status $exit_status." - exit 1 -fi + # File paths + pid_file="/tmp/main_process_id.txt" + status_file="/tmp/process_exit_status.txt" + log_file="/tmp/logs" + + echo "LOGGER: Waiting for the workload to show up in $pid_file" + + # Wait until the PID file exists + while [ ! -f "$pid_file" ]; do + sleep 1 + done + + # Extract PID from pid_file + pid=$(cat "$pid_file") + + echo "LOGGER: Streaming worker 0 logs." + # Tail the log file and terminate when the process with $pid exits + tail -f --pid=$pid --retry $log_file + + echo "LOGGER: Process $pid has finished." + + # Check if status_file contain any number + if ! grep -q '[0-9]' "$status_file"; then + echo "LOGGER: The status file contains no exit status." + exit 1 + fi + + # Read and output the exit status + exit_status=$(cat "$status_file") + + # Check the exit_status + if [ "$exit_status" -eq 0 ]; then + echo "LOGGER: The process exited successfully." + exit 0 + else + echo "LOGGER: The process failed with exit status $exit_status." + exit 1 + fi """ diff --git a/xlml/utils/tpu.py b/xlml/utils/tpu.py index 7e5e1d40f..8b312ec3d 100644 --- a/xlml/utils/tpu.py +++ b/xlml/utils/tpu.py @@ -119,15 +119,16 @@ def create_queued_resource_request( ) ) - startup_script_command = '' - + startup_cmds = [] if use_startup_script: - main_command = '\n'.join( + startup_cmds.extend( task_test_config.set_up_cmds + task_test_config.run_model_cmds ) - startup_script_command = startup_script.generate_startup_script( - main_command - ) + + main_command = '\n'.join(startup_cmds) + startup_script_command = startup_script.generate_startup_script( + main_command + ) metadata = { 'ssh-keys': f'{ssh_keys.user}:{ssh_keys.public}',