Merge branch 'main' into jax_0.8.0_py3.12_v7x

Author: andersensam

axlearn/cloud/gcp/pathways_utils.py

@@ -95,6 +95,8 @@ def get_pathways_tpu_version(gke_machine_type: str) -> str:
     https://github.com/google/pathways-job/blob/4417de7aa23d3c2316e400a3a327512834374475/internal/controller/pathwaysjob_controller.go#L70-L82
     """
     pathways_tpu_devices = {
+        # 7x
+        "tpu7x-standard-4t": "tpu7x",
         # v6e
         "ct6e-standard-4t": "tpuv6e",
         # v5p

axlearn/cloud/gcp/scripts/start_vm.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 #
-# Copyright © 2023 Apple Inc.
+# Copyright © 2025 Apple Inc.
 #
 # Boot script for GCP VM.
 #
@@ -28,6 +28,10 @@ trap 'echo "Bootstrap failed." >> ${SETUP_LOG_PATH}; \
       gcloud compute instances add-labels ${JOB_NAME} --labels=boot_status=failed --zone ${ZONE}; \
       exit 1' ERR
 
+# Ensure gcloud is not installed via snap.
+# This causes a permissions issue where the snap data directory is not writeable when using gcloud.
+snap remove --purge google-cloud-cli || echo "Ignoring snap removal failure..."
+
 # Ensure that gcloud is installed.
 while [[ ! -x $(which gcloud) ]]
 do
@@ -88,18 +92,16 @@ if [[ " ${tar_bundlers[*]} " =~ " ${BUNDLER_TYPE} " ]]; then
 
   # Install + activate Python 3.10.
   install_py310() {
-    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh
-    /bin/bash miniconda.sh -b -p /opt/conda
-    source /opt/conda/etc/profile.d/conda.sh
-    # Accept conda ToS to allow conda env creation
-    conda tos accept
-    conda create -y -n py310 python=3.10
-    conda activate py310
-    conda info --envs
-    # update pip and pre-install needed utilities
-    pip install -U pip swig flit uv
-    # Add conda to .profile file, and use login shell to source.
-    echo 'source /opt/conda/etc/profile.d/conda.sh && conda activate py310' >> ~/.profile
+    # Install Python 3.10 with the required venv module
+    apt-get update && \
+      DEBIAN_FRONTEND=noninteractive apt-get install -y python3.10-venv python3.10-dev && \
+      apt clean -y
+    mkdir -p /opt/venv
+    # Create a new virtual environment
+    python3.10 -m venv /opt/venv
+    /opt/venv/bin/python3.10 -m pip install -U pip swig flit uv
+    echo 'source /opt/venv/bin/activate' >> ~/.profile
+    source /opt/venv/bin/activate
   }
   install_py310 >> ${SETUP_LOG_PATH} 2>&1
   log "Using python3: $(which python3)"

axlearn/cloud/gcp/system_characteristics.py

@@ -416,38 +416,38 @@ class _SystemCharacteristics:
     "v6e-256": _SystemCharacteristics(
         "16x16", 64, "tpu-v6e-slice", "ct6e-standard-4t", 4, AcceleratorType["TPU"], "v6e-256"
     ),
-    # v7x
-    "v7x-8": _SystemCharacteristics(
+    # 7x
+    "7x-8": _SystemCharacteristics(
         "2x2x1", 1, "tpu7x", "tpu7x-standard-4t", 4, AcceleratorType["TPU"], "v7x-8"
     ),
-    "v7x-16": _SystemCharacteristics(
+    "7x-16": _SystemCharacteristics(
         "2x2x2", 2, "tpu7x", "tpu7x-standard-4t", 4, AcceleratorType["TPU"], "v7x-16"
     ),
-    "v7x-32": _SystemCharacteristics(
+    "7x-32": _SystemCharacteristics(
         "2x2x4", 4, "tpu7x", "tpu7x-standard-4t", 4, AcceleratorType["TPU"], "v7x-32"
     ),
-    "v7x-64": _SystemCharacteristics(
+    "7x-64": _SystemCharacteristics(
         "2x4x4", 8, "tpu7x", "tpu7x-standard-4t", 4, AcceleratorType["TPU"], "v7x-64"
     ),
-    "v7x-128": _SystemCharacteristics(
+    "7x-128": _SystemCharacteristics(
         "4x4x4", 16, "tpu7x", "tpu7x-standard-4t", 4, AcceleratorType["TPU"], "v7x-128"
     ),
-    "v7x-256": _SystemCharacteristics(
+    "7x-256": _SystemCharacteristics(
         "4x4x8", 32, "tpu7x", "tpu7x-standard-4t", 4, AcceleratorType["TPU"], "v7x-256"
     ),
-    "v7x-512": _SystemCharacteristics(
+    "7x-512": _SystemCharacteristics(
         "4x8x8", 64, "tpu7x", "tpu7x-standard-4t", 4, AcceleratorType["TPU"], "v7x-512"
     ),
-    "v7x-1024": _SystemCharacteristics(
+    "7x-1024": _SystemCharacteristics(
         "8x8x8", 128, "tpu7x", "tpu7x-standard-4t", 4, AcceleratorType["TPU"], "v7x-1024"
     ),
-    "v7x-2048": _SystemCharacteristics(
+    "7x-2048": _SystemCharacteristics(
         "8x8x16", 256, "tpu7x", "tpu7x-standard-4t", 4, AcceleratorType["TPU"], "v7x-2048"
     ),
-    "v7x-4096": _SystemCharacteristics(
+    "7x-4096": _SystemCharacteristics(
         "8x16x16", 512, "tpu7x", "tpu7x-standard-4t", 4, AcceleratorType["TPU"], "v7x-4096"
     ),
-    "v7x-8192": _SystemCharacteristics(
+    "7x-8192": _SystemCharacteristics(
         "16x16x16", 1024, "tpu7x", "tpu7x-standard-4t", 4, AcceleratorType["TPU"], "v7x-8192"
     ),
     # v4
@@ -530,13 +530,15 @@ class _SystemCharacteristics:
     "ct4p-hightpu-4t": 407,
     "ct5lp-hightpu-4t": 192,
     "ct6e-standard-4t": 720,
+    "tpu7x-standard-4t": 960,
 }
 
 # Reference doc https://cloud.google.com/kubernetes-engine/docs/concepts/plan-tpus
 GCE_MACHINE_TYPE_TO_CPU_CHARACTERISTICS = {
     "ct5p-hightpu-4t": 208,
     "ct4p-hightpu-4t": 240,
     "ct6e-standard-4t": 180,
+    "tpu7x-standard-4t": 224,
 }
 
 

axlearn/cloud/gcp/tpu.py

@@ -63,7 +63,7 @@ def infer_tpu_workers(tpu_type: str) -> int:
     try:
         if match is not None:
             tpu_version, tpu_cores = match.groups()
-            if tpu_version in {"v3", "v4", "v5p", "v7x"}:
+            if tpu_version in {"v3", "v4", "v5p", "7x"}:
                 return int(tpu_cores) // 8
             if tpu_version in {"v5litepod", "v6e"}:
                 return int(tpu_cores) // 4

axlearn/cloud/gcp/vertexai_tensorboard.py

@@ -2,6 +2,7 @@
 
 """Tools to upload model summaries to VertexAI Tensorboard."""
 
+import hashlib
 import multiprocessing
 import re
 import resource
@@ -34,10 +35,15 @@ def _vertexai_experiment_name_from_output_dir(output_dir: str) -> str:
     # Vertex AI Tensorboard requires experiment_name to match "[a-z0-9][a-z0-9-]+".
     experiment_name = match.group(1).lower().replace("/", "-").replace("_", "-").replace(".", "-")
     if len(experiment_name) >= _VERTEXAI_EXP_NAME_MAX_LEN:  # Vertex AI length limit.
-        raise ValueError(
-            rf"Experiment name must be less than {_VERTEXAI_EXP_NAME_MAX_LEN} chars long."
-            rf"{experiment_name} is {len(experiment_name)} chars."
+        logging.warning(
+            "Experiment name '%s' is %d chars (exceeds %d limit). Truncating to fit.",
+            experiment_name,
+            len(experiment_name),
+            _VERTEXAI_EXP_NAME_MAX_LEN,
         )
+        hash_suffix = hashlib.md5(experiment_name.encode()).hexdigest()[:8]
+        max_name_len = _VERTEXAI_EXP_NAME_MAX_LEN - 1 - len(hash_suffix)
+        experiment_name = f"{experiment_name[:max_name_len]}-{hash_suffix}"
     return experiment_name
 
 

axlearn/cloud/gcp/vertexai_tensorboard_test.py

@@ -12,6 +12,7 @@
 from axlearn.cloud.gcp import config as gcp_config
 from axlearn.cloud.gcp import test_utils
 from axlearn.cloud.gcp.vertexai_tensorboard import (
+    _VERTEXAI_EXP_NAME_MAX_LEN,
     VertexAITensorboardUploader,
     _vertexai_experiment_name_from_output_dir,
 )
@@ -79,5 +80,10 @@ def test_uploader_calls(
 
 class ExperimentNameTest(absltest.TestCase):
     def test_exp_name_len(self):
-        with self.assertRaises(ValueError):
-            _vertexai_experiment_name_from_output_dir("gs://abc/" + "a" * 128)
+        # Test that long experiment names are properly truncated
+        long_input = "gs://abc/" + "a" * _VERTEXAI_EXP_NAME_MAX_LEN
+        result = _vertexai_experiment_name_from_output_dir(long_input)
+        self.assertLessEqual(len(result), _VERTEXAI_EXP_NAME_MAX_LEN)
+        self.assertGreater(len(result), 0)
+        # The result should end with a hash suffix when truncated
+        self.assertRegex(result, r"-[a-f0-9]{8}$")

axlearn/common/compiler_options.py

@@ -132,7 +132,7 @@ def default_xla_options(
             xla_should_add_loop_invariant_op_in_chain="true",
             xla_tpu_use_enhanced_launch_barrier="true",
         )
-    if version == "v7x":
+    if version == "7x":
         options.update(
             # Many of the v7x flags are similar to v6e
             xla_tpu_scoped_vmem_limit_kib=65536,
@@ -307,7 +307,7 @@ class NotTpuError(ValueError):
 def infer_tpu_type(instance_type: str) -> str:
     """Infers tpu type (e.g. v4-8 or v6e-8-1) from instance type
     (e.g. tpu-v4-8, v4-8, tpu-v6e-8-1 or v6e-8-1)."""
-    if not (instance_type and re.fullmatch(r"(tpu-)?v.+-\d+", instance_type)):
+    if not (instance_type and re.fullmatch(r"(tpu-)?v?.+-\d+", instance_type)):
         raise NotTpuError(f"Invalid TPU instance: {instance_type}")
     return instance_type.replace("tpu-", "")
 
@@ -399,7 +399,7 @@ def infer_xsc_compiler_options(
 
 
 _TPU_VERSION_ALIASES = {"v5e": "v5litepod"}
-_TPU_VERSIONS = ("v3", "v4", "v5litepod", "v5p", "v6e", "v7x")
+_TPU_VERSIONS = ("v3", "v4", "v5litepod", "v5p", "v6e", "7x")
 
 
 def infer_xla_performance_flags(