Merge branch 'apple:main' into jax_py3.12_nightly

andersensam · web-flow · commit 72de49b0e010 · 2025-10-28T10:25:49.000-07:00
diff --git a/axlearn/cloud/common/bundler.py b/axlearn/cloud/common/bundler.py
@@ -629,12 +629,12 @@ def install_command(self, bundle_id: str) -> str:
         )
         pip_install_cmd = (
             f"if [[ -f {config.CONFIG_DIR}/requirements.txt ]]; then "
-            f"python3 -m pip install -r {config.CONFIG_DIR}/requirements.txt; "
-            "else python3 -m pip install .; fi"
+            f"python3 -m uv pip install -r {config.CONFIG_DIR}/requirements.txt; "
+            "else python3 -m uv pip install .; fi"
         )
         return (
             f"{copy_cmd} && tar -xzf axlearn.tar.gz && "
-            f"python3 -m pip install --upgrade pip && {pip_install_cmd}"
+            f"python3 -m uv pip install --upgrade pip && {pip_install_cmd}"
         )
 
 
diff --git a/axlearn/cloud/gcp/jobset_utils.py b/axlearn/cloud/gcp/jobset_utils.py
@@ -587,6 +587,56 @@ def _build_shared_memory_volumes(self, shared_memory: str) -> Nested[Any]:
             "emptyDir": {"medium": "Memory", "sizeLimit": shared_memory},
         }
 
+    def set_up_gcsfuse(self, cfg: "TPUJobBuilder.Config", volumes: list, annotations: dict) -> None:
+        """Sets up GCSFuse volumes and annotations.
+
+        Args:
+            cfg: The TPUJobBuilder configuration.
+            volumes: The list of volumes to append GCSFuse volumes to.
+            annotations: The dictionary of annotations to update with GCSFuse settings.
+        """
+        # Increases the shared memory volumes when enabled gcsfuse. This is useful when grain
+        # prefetch is enabled.
+        volumes.append(self._build_shared_memory_volumes(cfg.gcsfuse_mount.shared_memory))
+        # Mount a GCS bucket as a volume.
+        annotations.update(
+            {
+                "gke-gcsfuse/volumes": "true",
+                "gke-gcsfuse/cpu-request": cfg.gcsfuse_mount.cpu,
+                "gke-gcsfuse/memory-request": cfg.gcsfuse_mount.memory,
+                "gke-gcsfuse/ephemeral-storage-request": cfg.gcsfuse_mount.ephemeral_gb,
+                # GCSFuse will set limits=request if we only set requests:
+                # https://github.com/GoogleCloudPlatform/gcs-fuse-csi-driver/blob/main/pkg/webhook/config.go#L110
+                "gke-gcsfuse/cpu-limit": "0",
+                "gke-gcsfuse/memory-limit": "0",
+                "gke-gcsfuse/ephemeral-storage-limit": "0",
+            }
+        )
+        # Parse GCSFuseMount path into bucket, prefix.
+        parsed = urlparse(cfg.gcsfuse_mount.gcs_path)
+        # https://cloud.google.com/kubernetes-engine/docs/how-to/persistent-volumes/cloud-storage-fuse-csi-driver#consume-ephemeral-volume-pod
+        # Caveat: --implicit-dirs might have negative impacts on i/o performance. See
+        # https://github.com/googlecloudplatform/gcsfuse/blob/master/docs/semantics.md .
+        # See https://cloud.google.com/storage/docs/cloud-storage-fuse/config-file for more
+        # details about mountOptions.
+        # The mountOptions are following https://github.com/AI-Hypercomputer/maxtext/pull/1070.
+        volumes.append(
+            dict(
+                name=cfg.gcsfuse_mount.name,
+                csi=dict(
+                    driver="gcsfuse.csi.storage.gke.io",
+                    readOnly=cfg.gcsfuse_mount.read_only,
+                    volumeAttributes=dict(
+                        bucketName=parsed.netloc,
+                        # pylint: disable=line-too-long
+                        mountOptions=f"only-dir={parsed.path.lstrip('/')},implicit-dirs,metadata-cache:ttl-secs:-1,metadata-cache:stat-cache-max-size-mb:-1,metadata-cache:type-cache-max-size-mb:-1,kernel-list-cache-ttl-secs=-1,gcs-connection:http-client-timeout:{cfg.gcsfuse_mount.http_client_timeout}",
+                        gcsfuseMetadataPrefetchOnMount="false",  # Improves first-time read.
+                        disableMetrics="false",  # Enables GCSFuse metrics by default.
+                    ),
+                ),
+            )
+        )
+
     def _build_pod(self) -> Nested[Any]:
         """Builds a config for a single Pod, which is a set of containers.
 
@@ -601,47 +651,7 @@ def _build_pod(self) -> Nested[Any]:
 
         volumes.append(dict(name="shared-output", emptyDir={}))
         if cfg.gcsfuse_mount:
-            # Increases the shared memory volumes when enabled gcsfuse. This is useful when grain
-            # prefetch is enabled.
-            volumes.append(self._build_shared_memory_volumes(cfg.gcsfuse_mount.shared_memory))
-            # Mount a GCS bucket as a volume.
-            annotations.update(
-                {
-                    "gke-gcsfuse/volumes": "true",
-                    "gke-gcsfuse/cpu-request": cfg.gcsfuse_mount.cpu,
-                    "gke-gcsfuse/memory-request": cfg.gcsfuse_mount.memory,
-                    "gke-gcsfuse/ephemeral-storage-request": cfg.gcsfuse_mount.ephemeral_gb,
-                    # GCSFuse will set limits=request if we only set requests:
-                    # https://github.com/GoogleCloudPlatform/gcs-fuse-csi-driver/blob/main/pkg/webhook/config.go#L110
-                    "gke-gcsfuse/cpu-limit": "0",
-                    "gke-gcsfuse/memory-limit": "0",
-                    "gke-gcsfuse/ephemeral-storage-limit": "0",
-                }
-            )
-            # Parse GCSFuseMount path into bucket, prefix.
-            parsed = urlparse(cfg.gcsfuse_mount.gcs_path)
-            # https://cloud.google.com/kubernetes-engine/docs/how-to/persistent-volumes/cloud-storage-fuse-csi-driver#consume-ephemeral-volume-pod
-            # Caveat: --implicit-dirs might have negative impacts on i/o performance. See
-            # https://github.com/googlecloudplatform/gcsfuse/blob/master/docs/semantics.md .
-            # See https://cloud.google.com/storage/docs/cloud-storage-fuse/config-file for more
-            # details about mountOptions.
-            # The mountOptions are following https://github.com/AI-Hypercomputer/maxtext/pull/1070.
-            volumes.append(
-                dict(
-                    name=cfg.gcsfuse_mount.name,
-                    csi=dict(
-                        driver="gcsfuse.csi.storage.gke.io",
-                        readOnly=cfg.gcsfuse_mount.read_only,
-                        volumeAttributes=dict(
-                            bucketName=parsed.netloc,
-                            # pylint: disable=line-too-long
-                            mountOptions=f"only-dir={parsed.path.lstrip('/')},implicit-dirs,metadata-cache:ttl-secs:-1,metadata-cache:stat-cache-max-size-mb:-1,metadata-cache:type-cache-max-size-mb:-1,kernel-list-cache-ttl-secs=-1,gcs-connection:http-client-timeout:{cfg.gcsfuse_mount.http_client_timeout}",
-                            gcsfuseMetadataPrefetchOnMount="false",  # Improves first-time read.
-                            disableMetrics="false",  # Enables GCSFuse metrics by default.
-                        ),
-                    ),
-                )
-            )
+            self.set_up_gcsfuse(cfg, volumes, annotations)
         if cfg.host_mounts:
             for mount in cfg.host_mounts:
                 volumes.append(
diff --git a/axlearn/cloud/gcp/pathways_utils.py b/axlearn/cloud/gcp/pathways_utils.py
@@ -63,9 +63,9 @@
 # The container name of pathways proxy.
 _PATHWAYS_PROXY_CONTAINER_NAME = "pathways-proxy"
 # The k8s replicatedJob name for pathways-head pods.
-_PATHWAYS_HEAD_REPLICATED_JOB_NAME = "pathways-head"
+_PATHWAYS_HEAD_REPLICATED_JOB_NAME = "pwhd"
 # The k8s replicatedJob name for pathways-worker pods.
-_PATHWAYS_WORKER_REPLICATED_JOB_NAME = "pathways-worker"
+_PATHWAYS_WORKER_REPLICATED_JOB_NAME = "pwwk"
 
 # Add node-selector for cpu workload to avoid sharing nodes with system services.
 _PATHWAYS_HEAD_NODE_POOL_SELECTOR_KEY = "axlearn/nodepool_type"
@@ -439,17 +439,12 @@ def _build_pathways_head_pod(self) -> Nested[Any]:
             labels.update({BASTION_JOB_VERSION_LABEL: os.environ.get(BASTION_JOB_VERSION_ENV_VAR)})
 
         volumes.append(dict(name="shared-output", emptyDir={}))
-        volumes.append(dict(name="shared-memory", emptyDir=dict(medium="Memory")))
-
         if cfg.gcsfuse_mount:
-            annotations.update(
-                {
-                    "gke-gcsfuse/volumes": "true",
-                    "gke-gcsfuse/cpu-limit": cfg.gcsfuse_mount.cpu,
-                    "gke-gcsfuse/memory-limit": cfg.gcsfuse_mount.memory,
-                    "gke-gcsfuse/ephemeral-storage-limit": cfg.gcsfuse_mount.ephemeral_gb,
-                }
-            )
+            self._inner.set_up_gcsfuse(cfg, volumes, annotations)
+        else:
+            # gcsfuse mounts shared-memory. To avoid double mounting, we only mount
+            # shared-memory explicitly when gcsfuse_mount is not enabled.
+            volumes.append(dict(name="shared-memory", emptyDir=dict(medium="Memory")))
 
         node_selector = {
             _PATHWAYS_HEAD_NODE_POOL_SELECTOR_KEY: _PATHWAYS_HEAD_NODE_POOL_SELECTOR_VALUE,
diff --git a/axlearn/cloud/gcp/pathways_utils_test.py b/axlearn/cloud/gcp/pathways_utils_test.py
@@ -299,30 +299,65 @@ def test_validate_head_name(self):
         with self._job_config(CloudBuildBundler) as (cfg, bundler_cfg):
             cfg.inner.set(
                 project="test-project",
-                name="a" * 40,
+                name="a" * 49,
                 command="test_command",
                 output_dir="FAKE",
-            ).instantiate(bundler=bundler_cfg.instantiate())
+            )
 
             with self.assertRaisesRegex(
-                ValueError, r"pathways-head-1-1-abcde exceeds max \(63\) by 1 chars."
+                ValueError, r"pwhd-1-1-abcde exceeds max \(63\) by 1 chars."
             ):
                 _ = cfg.instantiate(bundler=bundler_cfg.instantiate())
 
     def test_validate_worker_name(self):
         with self._job_config(CloudBuildBundler) as (cfg, bundler_cfg):
             cfg.inner.set(
                 project="test-project",
-                name="a" * 38,
+                name="a" * 49,
                 command="test_command",
                 output_dir="FAKE",
-            ).instantiate(bundler=bundler_cfg.instantiate())
+            )
 
+            # Both head and worker have the same name length, so head validation runs first
             with self.assertRaisesRegex(
-                ValueError, r"pathways-worker-1-2-abcde exceeds max \(63\) by 1 chars."
+                ValueError, r"pwhd-1-1-abcde exceeds max \(63\) by 1 chars."
             ):
                 _ = cfg.instantiate(bundler=bundler_cfg.instantiate())
 
+    def test_build_pathways_head_pod_with_gcsfuse(self):
+        with (
+            self._job_config(
+                CloudBuildBundler,
+                gcsfuse_mount_spec=["mount_path=/tmp/gcsfuse", "gcs_path=gs://test-bucket/path"],
+            ) as (cfg, bundler_cfg),
+        ):
+            cfg.inner.set(
+                project="test-project",
+                name="test",
+                command="test_command",
+                output_dir="FAKE",
+            ).instantiate(bundler=bundler_cfg.instantiate())
+
+            builder = cfg.instantiate(bundler=bundler_cfg.instantiate())
+            # pylint: disable-next=protected-access
+            pod = builder._build_pathways_head_pod()
+            pod_spec = pod["spec"]
+            annotations = pod["metadata"]["annotations"]
+
+            # Verify gcsfuse annotations are set correctly
+            self.assertEqual(annotations.get("gke-gcsfuse/volumes"), "true")
+            self.assertIn("gke-gcsfuse/cpu-request", annotations)
+            self.assertIn("gke-gcsfuse/memory-request", annotations)
+            self.assertIn("gke-gcsfuse/ephemeral-storage-request", annotations)
+            # Verify limit annotations are set to "0"
+            self.assertEqual(annotations.get("gke-gcsfuse/cpu-limit"), "0")
+            self.assertEqual(annotations.get("gke-gcsfuse/memory-limit"), "0")
+            self.assertEqual(annotations.get("gke-gcsfuse/ephemeral-storage-limit"), "0")
+
+            # Verify shared memory volume is present
+            volume_names = [v["name"] for v in pod_spec["volumes"]]
+            self.assertIn("shared-memory", volume_names)
+
 
 class PathwaysMultiheadReplicatedJobTest(TestCase):
     """Tests PathwaysMultiheadReplicatedJob."""
@@ -383,37 +418,33 @@ def test_replicated_job(self, num_replicas):
                     annotations.get("axlearn/replicatedjob-load-balancer-port", {}),
                 )
 
-                if replicated_job_name.startswith("pathways-head"):
+                if replicated_job_name.startswith("pwhd"):
                     self.assertEqual(replicated_job["replicas"], num_replicas)
-                elif replicated_job_name.startswith("pathways-worker"):
+                elif replicated_job_name.startswith("pwwk"):
                     self.assertEqual(replicated_job["replicas"], 1)
 
     def test_validate_head_name(self):
         with self._job_config(CloudBuildBundler, 2) as (cfg, bundler_cfg):
             cfg.inner.set(
                 project="test-project",
-                name="a" * 40,
+                name="a" * 49,
                 command="test_command",
                 output_dir="FAKE",
-            ).instantiate(bundler=bundler_cfg.instantiate())
+            )
 
-        with self.assertRaisesRegex(
-            ValueError, r"pathways-head-1-1-abcde exceeds max \(63\) by 1 chars."
-        ):
+        with self.assertRaisesRegex(ValueError, r"pwhd-1-1-abcde exceeds max \(63\) by 1 chars."):
             _ = cfg.instantiate(bundler=bundler_cfg.instantiate())
 
     def test_validate_worker_name(self):
         with self._job_config(CloudBuildBundler, 2) as (cfg, bundler_cfg):
             cfg.inner.set(
                 project="test-project",
-                name="a" * 36,
+                name="a" * 47,
                 command="test_command",
                 output_dir="FAKE",
-            ).instantiate(bundler=bundler_cfg.instantiate())
+            )
 
-        with self.assertRaisesRegex(
-            ValueError, r"pathways-worker-2-0-2-abcde exceeds max \(63\) by 1 chars."
-        ):
+        with self.assertRaisesRegex(ValueError, r"pwwk-2-0-2-abcde exceeds max \(63\) by 1 chars."):
             _ = cfg.instantiate(bundler=bundler_cfg.instantiate())
 
 
diff --git a/axlearn/cloud/gcp/scripts/start_vm.sh b/axlearn/cloud/gcp/scripts/start_vm.sh
@@ -6,6 +6,8 @@
 #
 # Downloads install-bundle from metadata:bucket, installs, and sets label:boot_status = 'done'.
 
+set -xv
+
 echo "=== AXLearn start_vm.sh ==="
 
 # Setup logs.
@@ -94,6 +96,8 @@ if [[ " ${tar_bundlers[*]} " =~ " ${BUNDLER_TYPE} " ]]; then
     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
   }
@@ -105,7 +109,8 @@ if [[ " ${tar_bundlers[*]} " =~ " ${BUNDLER_TYPE} " ]]; then
     # 'Until' seemingly necessary as of 08/06/23 to avoid background setup process damaging partial
     # pip install state.
     # TODO(markblee,tom_gunter): Revisit this and understand why and how to wait until ready.
-    until python3 -m pip install .[gcp]; do
+    export UV_HTTP_TIMEOUT=300
+    until uv pip install .[gcp]; do
       echo "Attempting to install bundle again..."
       sleep 1
     done
diff --git a/axlearn/common/array_serialization_test.py b/axlearn/common/array_serialization_test.py
@@ -383,11 +383,10 @@ async def mock_ts_open(spec_arg, *args, **kwargs):
         deserialize_specs = captured_specs[-len(data) :]
         self.assertGreater(len(deserialize_specs), 0)
         for spec in deserialize_specs:
-            # gcs_grpc driver should be used when either:
-            # - JAX_PLATFORMS is "proxy" (running on GCP/Pathways), OR
+            # gcs_grpc driver should be used when:
             # - ENABLE_GCS_GRPC is "true" (explicit opt-in)
             # Otherwise, gcs driver should be used (no change)
-            should_use_grpc = jax_platforms == "proxy" or enable_gcs_grpc == "true"
+            should_use_grpc = enable_gcs_grpc == "true"
             expected_driver = "gcs_grpc" if should_use_grpc else "gcs"
             self.assertEqual(spec["kvstore"]["driver"], expected_driver)
 
diff --git a/axlearn/common/compiler_options.py b/axlearn/common/compiler_options.py
@@ -87,8 +87,6 @@ def default_xla_options(
             xla_latency_hiding_scheduler_rerun=2,
             # Improved performance for v6e.
             xla_tpu_scoped_vmem_limit_kib=98304,
-            # For megascale performance.
-            xla_jf_crs_combiner_threshold_count=10,
             # Disable collective matmul. Collective matmul could negatively affect performance in
             # some cases. Even in cases where collective matmul provides gains, the gains are
             # marginal on v6e due to the high arithmetic intensity.
@@ -106,8 +104,9 @@ def default_xla_options(
         # fusion and allreduce SC offloading by default.
         options.update(
             xla_tpu_enable_async_collective_fusion_fuse_all_gather="true",
-            # Always enable SparseCore offloading for allreduce.
-            xla_tpu_enable_sparse_core_collective_offload_all_reduce="true",
+            # Allreduce SparseCore offloading leads to quality discrepancy on v6e in JAX 0.6.2.
+            # TODO(changlan): Review and enable it later.
+            # xla_tpu_enable_sparse_core_collective_offload_all_reduce="true",
         )
 
         options.update(
@@ -455,7 +454,9 @@ def infer_xla_performance_flags(
                 xla_tpu_enable_async_collective_fusion_fuse_reduce_scatter="false",
                 xla_tpu_enable_sparse_core_collective_offload_all_gather="true",
                 xla_tpu_enable_sparse_core_collective_offload_reduce_scatter="true",
-                xla_tpu_enable_sparse_core_collective_offload_all_reduce="true",
+                # Allreduce SparseCore offloading leads to quality discrepancy on v6e in JAX 0.6.2.
+                # TODO(changlan): Review and enable it later.
+                # xla_tpu_enable_sparse_core_collective_offload_all_reduce="true",
                 xla_tpu_enable_all_gather_offload_tracing="true",
                 xla_tpu_enable_reduce_scatter_offload_tracing="true",
                 xla_tpu_enable_all_reduce_offload_tracing="true",

Original file line number	Diff line number	Diff line change
`@@ -629,12 +629,12 @@ def install_command(self, bundle_id: str) -> str:`
`629`	`629`	`)`
`630`	`630`	`pip_install_cmd = (`
`631`	`631`	`f"if [[ -f {config.CONFIG_DIR}/requirements.txt ]]; then "`
`632`		`- f"python3 -m pip install -r {config.CONFIG_DIR}/requirements.txt; "`
`633`		`- "else python3 -m pip install .; fi"`
	`632`	`+ f"python3 -m uv pip install -r {config.CONFIG_DIR}/requirements.txt; "`
	`633`	`+ "else python3 -m uv pip install .; fi"`
`634`	`634`	`)`
`635`	`635`	`return (`
`636`	`636`	`f"{copy_cmd} && tar -xzf axlearn.tar.gz && "`
`637`		`- f"python3 -m pip install --upgrade pip && {pip_install_cmd}"`
	`637`	`+ f"python3 -m uv pip install --upgrade pip && {pip_install_cmd}"`
`638`	`638`	`)`
`639`	`639`
`640`	`640`