test

Author: lkolluru05

Dockerfile

@@ -107,15 +107,20 @@ ARG EXTRAS=
 # Install a custom jaxlib that includes backport of Pathways shared memory feature.
 # PR: https://github.com/openxla/xla/pull/31417
 # Needed until Jax is upgraded to 0.8.0 or newer.
-ARG INSTALL_PATHWAYS_JAXLIB=false
+ARG INSTALL_PATHWAYS_JAXLIB=true
 
 # Ensure we install the TPU version, even if building locally.
 # Jax will fallback to CPU when run on a machine without TPU.
 RUN uv pip install -qq --prerelease=allow .[core,tpu] && uv cache clean
 RUN if [ -n "$EXTRAS" ]; then uv pip install -qq .[$EXTRAS] && uv cache clean; fi
+
+COPY jaxlib-0.6.2.dev20251021-cp310-cp310-manylinux2014_x86_64.whl .
+
+# 2. RUN the pip install command using the new, simple path *inside* the container
 RUN if [ "$INSTALL_PATHWAYS_JAXLIB" = "true" ]; then \
-      uv pip install --prerelease=allow "jaxlib==0.5.3.dev20250918" \
-        --find-links https://storage.googleapis.com/axlearn-wheels/wheels.html; \
+    #   uv pip install --prerelease=allow "jaxlib==0.6.2.dev20251020" \
+    #     --find-links https://storage.googleapis.com/axlearn-wheels/wheels.html; \
+        uv pip install jaxlib-0.6.2.dev20251021-cp310-cp310-manylinux2014_x86_64.whl; \
     fi
 COPY . .
 

axlearn/cloud/gcp/pathways_utils.py

@@ -334,29 +334,29 @@ def _build_pathways_head_container(self) -> dict:
             }
         )
 
-        # pylint: disable=line-too-long
-        env_list.append(
-            {
-                "name": "NUM_REPLICAS",
-                "valueFrom": {
-                    "fieldRef": {
-                        "fieldPath": "metadata.annotations['jobset.sigs.k8s.io/replicatedjob-replicas']"
-                    }
-                },
-            }
-        )
+        # # pylint: disable=line-too-long
+        # env_list.append(
+        #     {
+        #         "name": "NUM_REPLICAS",
+        #         "valueFrom": {
+        #             "fieldRef": {
+        #                 "fieldPath": "metadata.annotations['jobset.sigs.k8s.io/replicatedjob-replicas']"
+        #             }
+        #         },
+        #     }
+        # )
         # pylint: enable=line-too-long
 
-        env_list.append(
-            {
-                "name": "REPLICA_ID",
-                "valueFrom": {
-                    "fieldRef": {
-                        "fieldPath": "metadata.annotations['jobset.sigs.k8s.io/job-index']"
-                    }
-                },
-            }
-        )
+        # env_list.append(
+        #     {
+        #         "name": "REPLICA_ID",
+        #         "valueFrom": {
+        #             "fieldRef": {
+        #                 "fieldPath": "metadata.annotations['jobset.sigs.k8s.io/job-index']"
+        #             }
+        #         },
+        #     }
+        # )
 
         head_container["env"] = env_list
 

jaxlib-0.6.2.dev20251021-cp310-cp310-manylinux2014_x86_64.whl

@@ -43,7 +43,7 @@ core = [
     "tensorflow-metadata==1.17.2", # Otherwise Seqio will report no core for tfds
     "tensorflow-io>=0.37.1",  # for tensorflow-2.16. Note that 0.37.0 results in "pure virtual method called".
     "tensorflow-metadata>=1.0.0", # Otherwise Seqio will report no core for tfds
-    "tensorflow_text==2.17.0", # implied by seqio, but also used directly for text processing
+    "tensorflow_text==2.19.0", # implied by seqio, but also used directly for text processing
     "tensorstore>=0.1.63",  # used for supporting GDA checkpoints
     "toml",  # for config management
     "typing-extensions==4.12.2",
@@ -71,15 +71,14 @@ dev = [
     "pylint==2.17.7",
     "pytest",  # test runner
     "pytest-xdist",  # pytest plugin for test parallelism
-    "pytest-timeout",  # pytest plugin for forcing timeout of tests
     "pytype==2022.4.22",  # type checking
     "scikit-learn==1.5.2",  # test-only
     # Fix AttributeError: module 'scipy.linalg' has no attribute 'tril' and related scipy import errors.
     "scipy==1.12.0",
     "sentencepiece != 0.1.92",
     "tqdm",  # test-only
     "timm==0.6.12",  # DiT Dependency test-only
-    "torch>=2.1.1",  # test-only
+    "torch>=1.12.1",  # test-only
     "torchvision==0.16.1",  # test-only
     "safetensors<=0.5.3",  # TODO: Remove once torch dependency is >=2.3.0
     "transformers==4.51.3",  # test-only
@@ -199,7 +198,6 @@ line-length = 100
 target-version = 'py39'
 
 [tool.pytest.ini_options]
-timeout = 300
 addopts = "-rs -s -p no:warnings --junitxml=test-results/testing.xml"
 markers = [
     "gs_login: tests needing GS login.",
@@ -219,13 +217,9 @@ junit_family="xunit2"
 line_length = 100
 profile = "black"
 
-[tool.uv.pip]
+[tool.uv]
 find-links = [
+    "https://storage.googleapis.com/axlearn-wheels/wheels.html",
     "https://storage.googleapis.com/jax-releases/libtpu_releases.html",
     "https://storage.googleapis.com/jax-releases/jax_cuda_releases.html",
-]
-
-[tool.uv]
-override-dependencies = [
-    "ml-dtypes>=0.5,<0.6",
-]
+]