Fix: Fixed Jetstream_benchmarking_servig DAG by camiloCienet · Pull Request #276 · CIeNET-International/ml-auto-solutions-3

camiloCienet · 2026-06-12T09:14:04Z

Description

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

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run one-shot tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed.

alfredyu-cienet · 2026-06-15T00:36:20Z

      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 apt-get -o Dpkg::Lock::Timeout=300 -y update",


this means the apt-get was locked right?

but why does it get locked? are there more than one process trying to executing command on a same TPU VM?

I believe is because of this -->

(From GEMINI): When a fresh Google Cloud VM or TPU boots up for the very first time, the operating system executes several background tasks immediately. These often include: Cloud-init: Running default Google Cloud initialization scripts. Unattended Upgrades: The OS automatically checking for and installing critical security patches in the background. Startup Scripts: Any default infrastructure scripts setting up the TPU drivers. These background processes take up the dpkg lock.

Note: While running the DAGs I noticed Unattended Upgrades errors that keeps the Lock.

After running for a while a few tests are failing because of this. Seems that the Dpkg::Lock::Timeout=300 is not effective in all test. So I decided to disable this backgorund process that is holding the lock unattended-upgrades :

First command stop any unattended-upgrades second command totally disable it.

"sudo systemctl stop unattended-upgrades", "sudo systemctl disable unattended-upgrades",

let's do Dpkg::Lock::Timeout=300 (not terminating anything), but do something like this

BASH_HEADER = r""" apt_safe() { local PIDS=$(sudo fuser /var/lib/dpkg/lock-frontend /var/lib/dpkg/lock 2>/dev/null | awk '{print $1}') if [ -n "$PIDS" ]; then echo "Warning: apt is currently locked. Analyzing occupying processes..." echo "----------------------------------------" for PID in $PIDS; do local BIN_FILE=$(sudo readlink -f "/proc/$PID/exe" 2>/dev/null) local CMD_ARGS=$(ps -p "$PID" -o args= 2>/dev/null) # print pid, bin-file, cmd done echo "Applying 300 seconds timeout..." sudo apt-get -o Dpkg::Lock::Timeout=300 "$@" else sudo apt-get "$@" fi } """

xxx_command = ( ..., apt_safe -y install jq, ..., )

job_test_config = test_config.TpuVmTest( ..., set_up_cmds=BASH_HEADER + "; && \n".join(set_up_cmds), run_model_cmds=BASH_HEADER + "; && \n".join(run_model_cmds),

now we can assume the apt-get won't be blocked and no timeout required right?

alfredyu-cienet · 2026-06-15T01:09:38Z

      # Install uv using the standalone installer and add to PATH permanently in the current session
      "curl -LsSf https://astral.sh/uv/install.sh | sh",
      'export PATH="$HOME/.local/bin:$PATH"',
+      'export XLMLTEST_SSH_EXTERNAL_IPS="1"',


TBD and/or memo(@alfredyu-cienet):

all TPU-VM enables external IP, so using that is okay, the problem is how to configure that (XLMLTEST_SSH_EXTERNAL_IPS is likely a testing purpose config for local tests)

tpu.py:

queued_resource = tpu_api.QueuedResource( # TODO(ranran): enable configuration via `AcceleratorConfig` tpu=tpu_api.QueuedResource.Tpu( node_spec=[ tpu_api.QueuedResource.Tpu.NodeSpec( node_id=node_id, multi_node_params=multi_node_params, parent=parent, node=tpu_api.Node( accelerator_type=accelerator.name, description='noteardown', runtime_version=accelerator.runtime_version, network_config=tpu_api.NetworkConfig( network=accelerator.network, subnetwork=accelerator.subnetwork, enable_external_ips=True, ), metadata=metadata, labels=labels, scheduling_config=tpu_api.SchedulingConfig( preemptible=accelerator.preemptible, reserved=accelerator.reserved, ), ), ) ], ),

You are totally right. Just test it without that export and it works. Not sure why I did this before. I fixed it

alfredyu-cienet · 2026-06-15T01:10:22Z

  run_model_cmds = (
      # Start virtual environment
-      "source .env/bin/activate",
+      "source venv-312/bin/activate",


declare the name of venv and reuse that definition

alfredyu-cienet · 2026-06-15T01:12:28Z

-      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']}",
+      "pwd >&2",


does it a debugging log or sort? if it is, please remove them

Yes. Sorry remove it already

alfredyu-cienet · 2026-06-15T01:13:43Z

+      "sudo apt-get -o Dpkg::Lock::Timeout=300 install -y git-lfs",
+      "git lfs install",


what does lfs for? does it have a release version so that we don't need an addition command?

Ok this Git Large File Storage (Git LFS) was necessary because seems that there are some files in the Jetstream repo such as jetstream/tests/engine/external_tokenizers/llama2/tokenizer.model which contain the weights of the model, that git struggle to load(To big). (I tried without this command and it fails.)

From gemini:

Standard Git Can't Handle Large Files Well
Git is designed to track code (text files). If you put large binaries (like machine learning model weights, compiled binaries, or large benchmark datasets) directly into a standard Git repository, the repository becomes bloated, slow, and nearly impossible to clone.

Git LFS solves this by replacing the large files in the repository with tiny "pointer files." The actual large files are stored on a remote server.

The command git lfs pull looks at the repository, finds all those pointer files, and downloads the actual, heavy binaries from the remote LFS server.

alfredyu-cienet · 2026-06-15T01:16:42Z

+      "cd JetStream",
+      "pwd >&2",
+      "sudo chown -R $USER:$USER /home/sa_112632397993248756658/JetStream",
+      "git lfs pull",
+      "cd ..",
+      "pwd >&2",


what does it for?

So since we change to a new repo to execute Jetstream inference we dont have correct permissions. This command allow us to execute the previous explained git lfs pull to load all those necessary binaries. Tokenizer , datasets , etc.

alfredyu-cienet · 2026-06-15T01:17:51Z

    endpoints = [nodes[0].network_endpoints[0]]

-  use_external_ips = os.getenv('XLMLTEST_SSH_EXTERNAL_IPS', '0') == '1'
+  use_external_ips = os.environ.get('XLMLTEST_SSH_EXTERNAL_IPS') == '1'


question, does this change is based on pylint or something else?

I revert this change. And test it it is working without any issue. We dont need this change anymore.

alfredyu-cienet · 2026-06-15T01:18:58Z

TBD and/or memo(@alfredyu-cienet):

consider configure in somewhere else, as this would effect other DAGs

Refactor and test it. It worked. Let me know what do you think of the change. Thanks

alfredyu-cienet · 2026-06-15T08:19:35Z

-    project_name = Project.TPU_PROD_ENV_AUTOMATED.value
-    network = V6E_GCE_NETWORK
-    subnetwork = V6E_GCE_SUBNETWORK
+  # Default configurations for each TPU version


class TpuResource: zone: str runtime_version: str project_name: str network: str subnetwork: str class TpuConfig(Enum): V5E = TpuResource( 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, ) V5P = ... TRILLIUM = ...

define the temp config under the same enum, and put text literal in this enum (remove all new definitions related to cienet-cmcs, like V6E_SUBNETWORKS, CIENET_NETWORKS, CIENET_V5E_SUBNETWORKS etc.)

temp_V5E = ... temp_TRILLIUM = ...

and then DAG specify either

from dags.inference.maxtext_model_config_generator import TpuConfig

TpuConfig.V5E
TpuConfig.V5P
TpuConfig.TRILLIUM
TpuConfig.temp_V5E
TpuConfig.temp_TRILLIUM

Done and tested it. Please let me know what do you think.

alfredyu-cienet · 2026-06-16T03:16:05Z

      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 apt-get -o Dpkg::Lock::Timeout=300 -y update",


let's do Dpkg::Lock::Timeout=300 (not terminating anything), but do something like this

BASH_HEADER = r""" apt_safe() { local PIDS=$(sudo fuser /var/lib/dpkg/lock-frontend /var/lib/dpkg/lock 2>/dev/null | awk '{print $1}') if [ -n "$PIDS" ]; then echo "Warning: apt is currently locked. Analyzing occupying processes..." echo "----------------------------------------" for PID in $PIDS; do local BIN_FILE=$(sudo readlink -f "/proc/$PID/exe" 2>/dev/null) local CMD_ARGS=$(ps -p "$PID" -o args= 2>/dev/null) # print pid, bin-file, cmd done echo "Applying 300 seconds timeout..." sudo apt-get -o Dpkg::Lock::Timeout=300 "$@" else sudo apt-get "$@" fi } """

xxx_command = ( ..., apt_safe -y install jq, ..., )

job_test_config = test_config.TpuVmTest( ..., set_up_cmds=BASH_HEADER + "; && \n".join(set_up_cmds), run_model_cmds=BASH_HEADER + "; && \n".join(run_model_cmds),

alfredyu-cienet · 2026-06-16T03:18:04Z

 )
 V6E_SUBNETWORKS = (
-    f"{V5_NETWORKS_PREFIX}/regions/us-central2/subnetworks/mas-test"
+    f"{V5_NETWORKS_PREFIX}/regions/southamerica-west1-a/subnetworks/mas-test"


Timeout Lock

There is an issue with the Dpkg::Lock::Timeout=300. In my testing even with the timeout =300 some tests will fail with the error message

Waiting for cache lock: Could not get lock /var/lib/dpkg/lock-frontend. It is held by process 6581 (unattended-upgr)...

Killing background apt dpkg process

So maybe we might need to kill the dpkg process. I already tests this way and seems to be working without any issue and since the TPU vm are ephimeral I think would be ok. --> https://screenshot.googleplex.com/bKjebzNNrDegcwp

set_up_cmds = ( "pip install --upgrade pip", # 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 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 ..", ) Let me know what do you think.

@alfredyu-cienet For context. If use_startup_script=True a new task will be created in Airflow --> https://screenshot.googleplex.com/ATSFfAyuhLiN4A6 . In that task this commands that clean the lock files will be executed.

alfredyu-cienet · 2026-06-16T03:18:08Z

  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"


remove this definition

alfredyu-cienet · 2026-06-16T03:18:39Z

+  # reserved v5e in cienet-cmcs
+  US_WEST4_A = "us-west4-a"
+  # reserved v6e in cienet-cmcs
+  US_EAST4_B = "us-east4-b"


not mentioning cienet-cmcs

alfredyu-cienet · 2026-06-16T03:19:56Z

      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",


try not to hard code the sa_112632397993248756658 part

alfredyu-cienet · 2026-06-16T03:20:40Z

+  runtime_version = config.runtime_version
+  project_name = config.project_name
+  zone = config.zone
+  network = config.network
+  subnetwork = config.subnetwork


remove and inline to code below

alfredyu-cienet · 2026-06-16T03:21:29Z

-    project_name = Project.TPU_PROD_ENV_AUTOMATED.value
-    network = V6E_GCE_NETWORK
-    subnetwork = V6E_GCE_SUBNETWORK
+  # Handle both TpuConfig and TpuVersion for backwards compatibility


remove this comment (straightforward)

alfredyu-cienet · 2026-06-16T03:22:33Z

+  # 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_", "")]


remove, instead, add another attr for the actual tpu-version in TpuResource

I make some changes to remove this check. I update DAG's files

Just to be clear the function generate_model_configs in dags/inference/maxtext_model_config_generator.py only affects the following files:

dags/inference/maxtext_inference_microbenchmark.py (This has a function with same name in same file. I refactored since TpuConfig already contains all the neccessary info)

dags/inference/jetstream_inference_e2e.py

dags/inference/maxtext_inference.py.

Now all of them are using the class TpuConfig(Enum). Already tested it worked. Let me know what you think..

alfredyu-cienet · 2026-06-16T23:53:08Z

      )

-    startup_script_command = ''
+    dpkg_lock_fix_cmds = [


please rename, block_background_upgrade_services_script

please add comment(s), e.g. the reason we want to block them, whether it'd effect automation or not etc.

please move to the module startup_script

let's make it as a single string (concatenated)

alfredyu-cienet · 2026-06-16T23:54:41Z


-    startup_script_command = ''
+    dpkg_lock_fix_cmds = [
+        'set -x',


do we need this line?

(since the script will be processed through startup_script.generate_startup_script)

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

[cienet-private] Enable workflow checks for the primary working branches

d4c530a

camiloCienet requested review from alfredyu-cienet, andrewyct, severus-ho and xibinliu as code owners June 12, 2026 09:14

alfredyu-cienet reviewed Jun 15, 2026

View reviewed changes

camiloCienet force-pushed the maxtext/user/camilo/fix-jetstream-benchmark-serving branch 6 times, most recently from 647aaac to 4344b39 Compare June 15, 2026 07:51

alfredyu-cienet reviewed Jun 15, 2026

View reviewed changes

camiloCienet force-pushed the maxtext/user/camilo/fix-jetstream-benchmark-serving branch 4 times, most recently from 33d331e to 11bec15 Compare June 16, 2026 02:29

alfredyu-cienet reviewed Jun 16, 2026

View reviewed changes

camiloCienet force-pushed the maxtext/user/camilo/fix-jetstream-benchmark-serving branch 3 times, most recently from e6cea3a to 7102b6b Compare June 16, 2026 12:53

alfredyu-cienet reviewed Jun 16, 2026

View reviewed changes

camiloCienet force-pushed the maxtext/user/camilo/fix-jetstream-benchmark-serving branch 2 times, most recently from 9eaea64 to 78d0a44 Compare June 17, 2026 07:41

alfredyu-cienet force-pushed the dev branch from f8824db to 09c9e42 Compare June 17, 2026 08:08

camiloCienet force-pushed the maxtext/user/camilo/fix-jetstream-benchmark-serving branch from 78d0a44 to e459999 Compare June 17, 2026 09:15

Fix: Fixed Jetstream_benchmarking_servig DAG

4c8e672

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

camiloCienet force-pushed the maxtext/user/camilo/fix-jetstream-benchmark-serving branch from e459999 to 4c8e672 Compare June 17, 2026 09:16

alfredyu-cienet force-pushed the dev branch from 00a4856 to fdcbae9 Compare June 18, 2026 00:43

alfredyu-cienet force-pushed the dev branch from 50d7f44 to 6663e92 Compare June 28, 2026 23:55

alfredyu-cienet force-pushed the dev branch 4 times, most recently from d242c8d to f76e189 Compare July 3, 2026 08:25

alfredyu-cienet force-pushed the dev branch 4 times, most recently from a577785 to 71b8d8a Compare July 14, 2026 03:12

alfredyu-cienet force-pushed the dev branch 4 times, most recently from b588728 to 0e938b1 Compare July 23, 2026 01:43

alfredyu-cienet force-pushed the dev branch 2 times, most recently from 54f9a78 to 8e4ea82 Compare July 27, 2026 07:46

		"sudo apt-get -o Dpkg::Lock::Timeout=300 install -y git-lfs",
		"git lfs install",

Uh oh!

Conversation

camiloCienet commented Jun 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

camiloCienet Jun 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

camiloCienet Jun 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

camiloCienet Jun 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

camiloCienet Jun 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Timeout Lock

Killing background apt dpkg process

Uh oh!

camiloCienet Jun 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

camiloCienet commented Jun 12, 2026 •

edited

Loading

camiloCienet Jun 15, 2026 •

edited

Loading

camiloCienet Jun 15, 2026 •

edited

Loading

camiloCienet Jun 15, 2026 •

edited

Loading

camiloCienet Jun 16, 2026 •

edited

Loading

camiloCienet Jun 16, 2026 •

edited

Loading

camiloCienet Jun 16, 2026 •

edited

Loading