Merge branch 'main' into jax_py3.12_nightly

Author: camiloCienet

Dockerfile

@@ -1,7 +1,8 @@
 # syntax=docker/dockerfile:1
 
 ARG TARGET=base
-ARG BASE_IMAGE=ubuntu:24.04
+ARG BASE_IMAGE=ubuntu:22.04
+ARG BASE_IMAGE_COLOCATED=us-docker.pkg.dev/cloud-tpu-v2-images/pathways-colocated-python/sidecar:2025_10_29-python_3.10-jax_0.6.2
 
 FROM ${BASE_IMAGE} AS base
 
@@ -18,7 +19,7 @@ RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.
     curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && \
     apt-get update -y -qq && \
     apt-get install -y -qq apt-transport-https ca-certificates gcc g++ \
-    git screen ca-certificates google-perftools google-cloud-cli python3.12-venv && \
+    git screen ca-certificates google-perftools google-cloud-cli python3.10-venv python3.10-dev && \
     apt clean -y -qq
 
 # Setup.
@@ -102,12 +103,40 @@ COPY pyproject.toml README.md /root/
 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
 RUN if [ "$INSTALL_PATHWAYS_JAXLIB" = "true" ]; then \
-      uv pip install --prerelease=allow "jaxlib==0.5.3.dev20250918" \
+      uv pip install --prerelease=allow "jaxlib==0.6.2.dev20251021" \
         --find-links https://storage.googleapis.com/axlearn-wheels/wheels.html; \
     fi
 RUN uv pip install -qq --no-deps libtpu==0.0.28.dev20251104+nightly && uv cache clean
 COPY . .
 
+################################################################################
+# Colocated Python container spec.                                             #
+################################################################################
+
+FROM ${BASE_IMAGE_COLOCATED} AS colocated-python
+
+WORKDIR /app
+COPY . .
+
+# Install the additional user-provided dependencies, strictly enforcing the rules
+# from the base image's constraints file.
+RUN \
+    # 1. Install user-provided dependencies with modified constraints
+    grep -v "^numpy" /opt/venv/server_constraints.txt | grep -v "^scipy" > /tmp/modified_constraints.txt && \
+    echo "--> Installing user-provided dependencies..." && \
+    uv pip install ".[core,gcp]" -c /tmp/modified_constraints.txt && \
+    \
+    # 2. Override numpy and scipy with specific versions
+    uv pip install numpy==2.1.1 scipy==1.15.3 && \
+    \
+    # 3. Verify that the colocated_python_cpu_client is present.
+    echo "--> Verifying JAX patch integrity..." && \
+    python -c "from jax._src.lib import _jax; _jax.colocated_python_cpu_client" && \
+    echo "--> JAX patch verification successful." && \
+    \
+    # 4. Clean the cache to keep the image slim.
+    uv cache clean
+
 ################################################################################
 # GPU container spec.                                                          #
 ################################################################################

axlearn/audio/aligner/ctc_aligner.py

@@ -488,8 +488,11 @@ def test_base(self, loop):
         expected_search_lattice = np.where(
             np.isinf(self.expected_search_lattices), 0.0, self.expected_search_lattices
         )
+        # Transpose expected to match new [T+1, B, 2*(L+1)] layout
+        expected_search_lattice = np.transpose(expected_search_lattice, (1, 0, 2))
+        expected_backtrace = np.transpose(self.expected_backtrace, (1, 0, 2))
         self.assertNestedAllClose(search_lattice, expected_search_lattice)
-        self.assertNestedEqual(state.backtrace, self.expected_backtrace)
+        self.assertNestedEqual(state.backtrace, expected_backtrace)
 
     @parameterized.parameters(("python",), ("lax",))
     def test_to_alignment(self, loop):
@@ -557,6 +560,7 @@ def setUp(self):
         # pylint: disable=invalid-name
         T = 4
         L = 3
+        # Create in [B, T+1, 2*(L+1)] format first
         search_lattices = np.full((3, (T + 1), 2 * (L + 1)), fill_value=-np.inf)
         frame_lengths = np.array([3, 4, 3], dtype=np.int32)
         label_lengths = np.array([2, 3, 2], dtype=np.int32)
@@ -571,6 +575,9 @@ def setUp(self):
         search_lattices[2, frame_lengths[2], 2 * label_lengths[2]] = 10.0
         search_lattices[2, frame_lengths[2], 2 * label_lengths[2] + 1] = -np.inf
 
+        # Transpose to [T+1, B, 2*(L+1)] to match the new layout
+        search_lattices = np.transpose(search_lattices, (1, 0, 2))
+
         self.search_lattices = jnp.array(search_lattices)
         self.frame_lengths = jnp.array(frame_lengths)
         self.label_lengths = jnp.array(label_lengths)

axlearn/audio/aligner/ctc_aligner_test.py

@@ -150,8 +150,7 @@ def default_config(cls):
     class Scorer(Protocol):
         """A function that computes word errors."""
 
-        def __call__(self, hypotheses: list[str], reference: list[str]) -> WordErrors:
-            ...
+        def __call__(self, hypotheses: list[str], reference: list[str]) -> WordErrors: ...
 
     def __init__(
         self,

axlearn/audio/evaler_asr.py

@@ -133,13 +133,13 @@ def _compute_metrics(
         if brevity_penalty:
             decode_kwargs["brevity_penalty"] = brevity_penalty
 
-        cfg: WordErrorRateMetricCalculator.Config = (
-            WordErrorRateMetricCalculator.default_config().set(
-                vocab=config_for_class(seqio.SentencePieceVocabulary).set(
-                    sentencepiece_model_file=vocab_file,
-                ),
-                model_method_kwargs=decode_kwargs,
-            )
+        cfg: (
+            WordErrorRateMetricCalculator.Config
+        ) = WordErrorRateMetricCalculator.default_config().set(
+            vocab=config_for_class(seqio.SentencePieceVocabulary).set(
+                sentencepiece_model_file=vocab_file,
+            ),
+            model_method_kwargs=decode_kwargs,
         )
         calculator: WordErrorRateMetricCalculator = cfg.set(name="test-metric").instantiate(
             parent=None, model=model, model_param_partition_specs={}

axlearn/audio/evaler_asr_test.py

@@ -1690,10 +1690,9 @@ def test_sync_jobs_for_valid_pending_to_sudden_invalid_jobs(self):
         mock_validator_cfg = MockStatefulJobValidator.default_config()
         mock_append_to_job_history = mock.MagicMock()
 
-        with self._patch_bastion(
-            validator_cfg=mock_validator_cfg
-        ) as mock_bastion, mock.patch.object(
-            mock_bastion, "_append_to_job_history", mock_append_to_job_history
+        with (
+            self._patch_bastion(validator_cfg=mock_validator_cfg) as mock_bastion,
+            mock.patch.object(mock_bastion, "_append_to_job_history", mock_append_to_job_history),
         ):
             os.makedirs(mock_bastion._active_dir, exist_ok=True)
             os.makedirs(_JOB_DIR, exist_ok=True)
@@ -1802,10 +1801,9 @@ def test_sync_jobs_for_immediate_invalid_pending_jobs(self):
         mock_validator_cfg = MockAlwaysInvalidValidator.default_config()
         mock_append_to_job_history = mock.MagicMock()
 
-        with self._patch_bastion(
-            validator_cfg=mock_validator_cfg
-        ) as mock_bastion, mock.patch.object(
-            mock_bastion, "_append_to_job_history", mock_append_to_job_history
+        with (
+            self._patch_bastion(validator_cfg=mock_validator_cfg) as mock_bastion,
+            mock.patch.object(mock_bastion, "_append_to_job_history", mock_append_to_job_history),
         ):
             os.makedirs(mock_bastion._active_dir, exist_ok=True)
             os.makedirs(_JOB_DIR, exist_ok=True)

axlearn/cloud/common/bastion_test.py

@@ -301,6 +301,8 @@ class Config(Bundler.Config):
         cache_from: Optional[Sequence[str]] = None
         # Skip the build + push step (e.g., using a pre-built image).
         skip_bundle: bool = False
+        # Sidecar names to build images for.
+        sidecars: list[str] = []
 
     def __init__(self, cfg: Config):
         super().__init__(cfg)
@@ -337,6 +339,7 @@ def from_spec(cls, spec: list[str], *, fv: Optional[flags.FlagValues]) -> Config
         - platform: The image target platform.
         - allow_dirty: Whether to ignore dirty git status.
         - cache_from: A comma-separated list of cache sources.
+        - sidecars: A comma-separated list of sidecar names.
         - skip_bundle: Whether to skip the build + push. This option is intended to be used when an
             image has already been pre-built offline, in which case we may still want to leverage
             the install commands implemented by the bundler.
@@ -346,13 +349,15 @@ def from_spec(cls, spec: list[str], *, fv: Optional[flags.FlagValues]) -> Config
         cfg: BaseDockerBundler.Config = super().from_spec(spec, fv=fv)
         kwargs = parse_kv_flags(spec, delimiter="=")
         cache_from = canonicalize_to_list(kwargs.pop("cache_from", None))
+        sidecars = canonicalize_to_list(kwargs.pop("sidecars", None))
         skip_bundle = to_bool(kwargs.pop("skip_bundle", False))
         allow_dirty = to_bool(kwargs.pop("allow_dirty", False))
         # Non-config specs are treated as build args.
         build_args = {k: kwargs.pop(k) for k in list(kwargs.keys()) if k not in cfg}
         return cfg.set(
             build_args=build_args,
             cache_from=cache_from,
+            sidecars=sidecars,
             skip_bundle=skip_bundle,
             allow_dirty=allow_dirty,
             **kwargs,
@@ -485,6 +490,16 @@ def _build_and_push(
         labels: dict[str, str],
     ) -> str:
         cfg: DockerBundler.Config = self.config
+
+        _, tag = image.rsplit(":", maxsplit=1)
+        for sidecar in cfg.sidecars:
+            sidecar_bundler = cfg.set(
+                image=sidecar,
+                target=sidecar,
+                sidecars=[],
+            ).instantiate()
+            sidecar_bundler.bundle(tag=tag)
+
         return docker_push(
             docker_build(
                 dockerfile=dockerfile,

axlearn/cloud/common/bundler.py

@@ -1,7 +1,7 @@
 # Copyright © 2023 Apple Inc.
 
 """Utilities to clean resources."""
-
+import logging
 from collections.abc import Sequence
 from enum import Enum
 
@@ -89,6 +89,12 @@ def sweep(self, jobs: dict[str, JobSpec]) -> Sequence[str]:
             schedule_result = scheduler.schedule(
                 dict(my_job=job_spec.metadata),
             )
-            if schedule_result.job_verdicts["my_job"].over_limits:
+            job_verdict = schedule_result.job_verdicts.get("my_job")
+            # In certain corner cases, schedule_result may not contain the job being passed in
+            # TODO(yangwwei): revisit and fix it in a proper way
+            if job_verdict is None:
+                logging.warning("job %s does not exist in the verdicts, skipping", job_name)
+                continue
+            if job_verdict.over_limits:
                 result.append(job_name)
         return result

axlearn/cloud/common/cleaner.py

@@ -4,6 +4,7 @@
 
 import datetime
 from collections.abc import Sequence
+from unittest import mock
 
 from absl.testing import parameterized
 
@@ -15,7 +16,7 @@
     UnschedulableCleaner,
 )
 from axlearn.cloud.common.quota import QuotaInfo
-from axlearn.cloud.common.scheduler import JobMetadata, JobScheduler
+from axlearn.cloud.common.scheduler import BaseScheduler, JobMetadata, JobScheduler, JobVerdict
 from axlearn.cloud.common.types import JobSpec
 from axlearn.common.config import REQUIRED, Required, config_class, config_for_function
 
@@ -116,3 +117,99 @@ def mock_jobs(*, prefix: str, resource_count: int):
         cleaner = cfg.instantiate()
         result = cleaner.sweep(jobs)
         self.assertSequenceEqual(result, list(unschedulable_jobs.keys()))
+
+    def test_sweep_missing_job_in_verdicts(self):
+        """Test that sweep handles the case when job is not in job_verdicts."""
+        # Create a mock schedule result without the job being scheduled
+        mock_schedule_result = BaseScheduler.ScheduleResults(
+            project_limits={},
+            project_usages={},
+            job_verdicts={},  # empty job verdicts
+            unused_limits=None,
+        )
+
+        # Create a simple job
+        job = new_jobspec(
+            name="test_job",
+            command="echo hello",
+            metadata=JobMetadata(
+                user_id="user_1",
+                project_id="default",
+                creation_time=datetime.datetime(1900, 1, 1, 0, 0, 0),
+                resources={"aws_4:p5.48xlarge": 1},
+            ),
+        )
+        jobs = {"test_job": job}
+
+        # Create cleaner
+        quota = QuotaInfo(
+            total_resources=[{"aws_4:p5.48xlarge": 10}],
+            project_resources={},
+            project_membership={},
+        )
+        cfg = UnschedulableCleaner.default_config().set(
+            scheduler=JobScheduler.default_config().set(
+                quota=config_for_function(lambda: lambda: quota)
+            )
+        )
+        cleaner = cfg.instantiate()
+
+        # Patch the scheduler.schedule method to return our mock result
+        # and verify the logging
+        with self.assertLogs(level="INFO") as log_context:
+            with mock.patch.object(JobScheduler, "schedule", return_value=mock_schedule_result):
+                result = cleaner.sweep(jobs)
+
+            # Should return empty list since "my_job" is not in job_verdicts
+            self.assertSequenceEqual(result, [])
+
+        # Verify the log message was emitted
+        self.assertTrue(
+            any(
+                "job test_job does not exist in the verdicts, skipping" in log
+                for log in log_context.output
+            )
+        )
+
+    def test_sweep_job_with_over_limits(self):
+        """Test that sweep correctly identifies jobs with over_limits."""
+        # Create a mock schedule result with a verdict that has over_limits
+        mock_schedule_result = BaseScheduler.ScheduleResults(
+            project_limits={},
+            project_usages={},
+            job_verdicts={"my_job": JobVerdict(over_limits={"aws_4:p5.48xlarge"})},
+            unused_limits=None,
+        )
+
+        # Create a simple job
+        job = new_jobspec(
+            name="test_job",
+            command="echo hello",
+            metadata=JobMetadata(
+                user_id="user_1",
+                project_id="default",
+                creation_time=datetime.datetime(1900, 1, 1, 0, 0, 0),
+                resources={"aws_4:p5.48xlarge": 100},
+            ),
+        )
+        jobs = {"test_job": job}
+
+        # Create cleaner
+        quota = QuotaInfo(
+            total_resources=[{"aws_4:p5.48xlarge": 10}],
+            project_resources={},
+            project_membership={},
+        )
+        cfg = UnschedulableCleaner.default_config().set(
+            scheduler=JobScheduler.default_config().set(
+                quota=config_for_function(lambda: lambda: quota)
+            )
+        )
+        cleaner = cfg.instantiate()
+
+        # Patch the scheduler.schedule method to return our mock result
+        with mock.patch.object(JobScheduler, "schedule", return_value=mock_schedule_result):
+            result = cleaner.sweep(jobs)
+
+        # Should return the job since it has over_limits
+        self.assertSequenceEqual(result, ["test_job"])

axlearn/cloud/common/cleaner_test.py

@@ -0,0 +1,22 @@
+# Copyright © 2025 Apple Inc.
+
+"""Command patcher interface for job processing."""
+
+from typing import Any, Dict
+
+from axlearn.common.config import Configurable
+
+
+class UserCommandPatcher(Configurable):
+    """A command patcher interface for modifying a job's user command."""
+
+    def patch(self, command: str, **kwargs: Dict[str, Any]) -> str:
+        """Patches the command string.
+
+        Args:
+            command: The original user command string.
+
+        Returns:
+            The patched command string.
+        """
+        raise NotImplementedError(type(self))