Fix: Fixed Jetstream_benchmarking_servig DAG

- Point v5e tests to new reservation in CIENET-CMCS project
- Point v6e tests to new reservation in tpu-prod-env-automated

Author: camiloCienet

dags/common/vm_resource.py

@@ -27,7 +27,7 @@
     f"{V5_NETWORKS_PREFIX}/regions/europe-west4/subnetworks/mas-test-v2"
 )
 V6E_SUBNETWORKS = (
-    f"{V5_NETWORKS_PREFIX}/regions/us-central2/subnetworks/mas-test"
+    f"{V5_NETWORKS_PREFIX}/regions/southamerica-west1-a/subnetworks/mas-test"
 )
 # TODO: Figure V6E_GCE_NETWORK and V6E_GCE_SUBNETWORK
 V6E_GCE_NETWORK = "default"
@@ -72,6 +72,7 @@ class Project(enum.Enum):
   TPU_PROD_ENV_LARGE_ADHOC = "tpu-prod-env-large-adhoc"
   TPU_PROD_ENV_ONE_VM = "tpu-prod-env-one-vm"
   TPU_PROD_ENV_LARGE_CONT = "tpu-prod-env-large-cont"
+  CIENET_CMCS = "cienet-cmcs"
 
 
 class ImageProject(enum.Enum):
@@ -124,6 +125,10 @@ class Zone(enum.Enum):
   US_EAST5_C = "us-east5-c"
   # reserved v5e in tpu-prod-env-multipod
   US_WEST4_B = "us-west4-b"
+  # reserved v5e in cienet-cmcs
+  US_WEST4_A = "us-west4-a"
+  # reserved v6e in cienet-cmcs
+  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

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(
@@ -55,8 +56,20 @@ def get_config(
       # Download jetstream and maxtext
       f"if [ ! -d maxtext ]; then git clone {maxtext_branch} https://github.com/google/maxtext.git; fi",
       f"if [ ! -d JetStream ]; then git clone {jetstream_branch} https://github.com/google/JetStream.git; fi",
-      "sudo apt-get -y update",
-      "sudo apt-get -y install jq",
+      "sudo systemctl stop unattended-upgrades",
+      "sudo systemctl disable unattended-upgrades",
+      # 2. Aggressively kill any lingering apt, dpkg, or unattended-upgrades processes
+      "sudo pkill -9 unattended-upgr || true",
+      "sudo pkill -9 apt || true",
+      "sudo pkill -9 dpkg || true",
+      # 3. Remove any remaining lock files
+      "sudo rm -f /var/lib/apt/lists/lock",
+      "sudo rm -f /var/cache/apt/archives/lock",
+      "sudo rm -f /var/lib/dpkg/lock*",
+      # 4. Repair any interrupted package configurations
+      "sudo dpkg --configure -a",
+      "sudo apt-get -o Dpkg::Lock::Timeout=300 -y update",
+      "sudo apt-get -o Dpkg::Lock::Timeout=300 -y install jq",
       "cd JetStream && pip install -e . && cd benchmarks && pip install -r requirements.in",
       "pip install torch --index-url https://download.pytorch.org/whl/cpu",
       "cd ..",
@@ -70,11 +83,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 -o Dpkg::Lock::Timeout=300 install -y git-lfs",
+      "git lfs install",
   )
 
   set_up_cmds += setup_maxtext_cmds
@@ -113,15 +129,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 +151,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/sa_112632397993248756658/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 +232,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,

dags/inference/maxtext_inference.py

@@ -19,7 +19,7 @@
 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}",

dags/inference/maxtext_model_config_generator.py

@@ -14,11 +14,60 @@
 
 """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:
+  runtime_version: str
+  project_name: str
+  zone: str
+  network: str
+  subnetwork: str
+
+
+class TpuConfig(Enum):
+  V5E = TpuResource(
+      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(
+      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(
+      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(
+      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(
+      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,28 +148,22 @@ 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
+  # Handle both TpuConfig and TpuVersion for backwards compatibility
+  if isinstance(tpu_version, TpuConfig):
+    config = tpu_version.value
+    tpu_version_enum = TpuVersion[tpu_version.name.replace("temp_", "")]
+  else:
+    config = TpuConfig[tpu_version.name].value
+    tpu_version_enum = tpu_version
+
+  runtime_version = config.runtime_version
+  project_name = config.project_name
+  zone = config.zone
+  network = config.network
+  subnetwork = config.subnetwork
   jetstream_benchmark_serving = (
       jetstream_benchmark_serving_gce_config.get_config(
-          tpu_version=tpu_version,
+          tpu_version=tpu_version_enum,
           tpu_cores=tpu_cores,
           tpu_zone=zone,
           time_out_in_min=sweep_model_configs["time_out_in_min"],