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

@@ -18,6 +18,12 @@
 import enum
 from xlml.apis.xpk_cluster_config import XpkClusterConfig
 
+# Cienet Reservation for V5e and V6e
+CIENET_NETWORK_PREFIX = f"projects/cienet-cmcs"
+CIENET_NETWORKS = f"{CIENET_NETWORK_PREFIX}/global/networks/mas-test"
+CIENET_V5E_SUBNETWORKS = (
+    f"{CIENET_NETWORK_PREFIX}/regions/us-west4/subnetworks/mas-test"
+)
 
 V5_NETWORKS_PREFIX = "projects/tpu-prod-env-automated"
 V5_NETWORKS = f"{V5_NETWORKS_PREFIX}/global/networks/mas-test"
@@ -27,7 +33,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 +78,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 +131,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,10 @@ 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",
+      "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 +73,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 +119,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 +141,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 +222,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

@@ -576,7 +576,6 @@
           },
       }
   )
-
   # run_configs = [
   #     f"{LLAMA2_7B}-{BASE_MODE}-{W_BF16_KV_BF16}",
   #     f"{LLAMA2_7B}-{BASE_MODE}-{W_INT8_KV_INT8}",
@@ -616,6 +615,7 @@
                     request_rate=request_rate,
                     tpu_version=tpu_version,
                     tpu_cores=tpu_cores,
+                    backup_tpu_resources=True,
                 )
             )
             dags.append(jetstream_benchmark_serving_kv_cache_layout)

dags/inference/maxtext_model_config_generator.py

@@ -14,7 +14,7 @@
 
 """A helper to generate maxtext model configs."""
 
-from dags.common.vm_resource import TpuVersion, Zone, Project, V5_NETWORKS, V5E_SUBNETWORKS, V5P_SUBNETWORKS, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK
+from dags.common.vm_resource import TpuVersion, Zone, Project, V5_NETWORKS, V5E_SUBNETWORKS, V5P_SUBNETWORKS, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK, V6E_SUBNETWORKS, CIENET_NETWORKS, CIENET_V5E_SUBNETWORKS
 from dags.inference.configs import jetstream_benchmark_serving_gce_config
 from dags.multipod.configs.common import SetupMode
 
@@ -28,6 +28,7 @@ def generate_model_configs(
     request_rate,
     tpu_version,
     tpu_cores,
+    backup_tpu_resources=False,
 ):
   model_configs = {}
   model_configs["model_config_name"] = model_config_name
@@ -99,25 +100,55 @@ 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
+  # Default configurations for each TPU version
+  tpu_configs = {
+      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,
+      },
+      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,
+      },
+      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,
+      },
+  }
+
+  # Overrides when backup TPU resources are requested
+  backup_overrides = {
+      TpuVersion.V5E: {
+          "project_name": Project.CIENET_CMCS.value,
+          "zone": Zone.US_WEST4_A.value,
+          "network": CIENET_NETWORKS,
+          "subnetwork": CIENET_V5E_SUBNETWORKS,
+      },
+      TpuVersion.TRILLIUM: {
+          "zone": Zone.SOUTHAMERICA_WEST1_A.value,
+          "network": V5_NETWORKS,
+          "subnetwork": V6E_SUBNETWORKS,
+      },
+  }
+
+  config = tpu_configs[tpu_version]
+  if backup_tpu_resources and tpu_version in backup_overrides:
+    config.update(backup_overrides[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,