Add colocated python support

Co-authored-by: lkolluru05 <lkolluru@google.com>
ORIGINAL_AUTHOR=rohitc33 <70339497+rohitc33@users.noreply.github.com>
COPYBARA_INTEGRATE_REVIEW=#1353 from rohitc33:rohit-colocated-2 e3ac888

Author: 2 people

Dockerfile

@@ -2,6 +2,7 @@
 
 ARG TARGET=base
 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
 
@@ -102,11 +103,39 @@ 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
 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/cloud/common/bundler.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/gcp/bundler.py

@@ -129,14 +129,14 @@ class Config(BaseDockerBundler.Config):
                 from flags.
             is_async: Whether to build asynchronously. If True, callers should invoke
                 `wait_until_finished()` to wait for bundling to complete.
+            private_worker_pool: If provided, should be the identifier of a private worker pool.
+                See: https://cloud.google.com/build/docs/private-pools/private-pools-overview
         """
 
         # GCP project.
         project: Required[str] = REQUIRED
         # Build image asynchronously.
         is_async: bool = True
-        # If provided, should be the identifier of a private worker pool.
-        # See: https://cloud.google.com/build/docs/private-pools/private-pools-overview
         private_worker_pool: Optional[str] = None
 
     @classmethod
@@ -175,9 +175,14 @@ def _build_and_push(
         )
         image_path, image_tag = image.rsplit(":", maxsplit=1)
         latest_tag = f"{image_path}:latest"
-        cloudbuild_yaml = f"""
-steps:
-- name: "gcr.io/cloud-builders/docker"
+
+        # Build steps - start with main image
+        build_steps = []
+        images_list = [f'"{image}"', f'"{latest_tag}"']
+
+        # Main image build step
+        build_steps.append(
+            f"""- name: "gcr.io/cloud-builders/docker"
   args: [
     "build",
     "-f", "{os.path.relpath(dockerfile, context)}",
@@ -193,11 +198,44 @@ def _build_and_push(
     "."
   ]
   env:
-  - "DOCKER_BUILDKIT=1"
+  - "DOCKER_BUILDKIT=1\""""
+        )
+
+        # Add sidecar image build steps
+        for sidecar in cfg.sidecars:
+            sidecar_target = sidecar
+            sidecar_image_path = f"{cfg.repo}/{sidecar}"
+            sidecar_image = f"{sidecar_image_path}:{image_tag}"
+            sidecar_latest_image = f"{sidecar_image_path}:latest"
+
+            build_steps.append(
+                f"""- name: "gcr.io/cloud-builders/docker"
+  args: [
+    "build",
+    "-f", "{os.path.relpath(dockerfile, context)}",
+    "-t", "{sidecar_image}",
+    "-t", "{sidecar_latest_image}",
+    "--target", "{sidecar_target}",
+    "--cache-from", "{sidecar_image}",
+    "--cache-from", "{sidecar_latest_image}",
+    {cache_from}
+    {build_platform}
+    {build_args}
+    {labels}
+    "."
+  ]
+  env:
+  - "DOCKER_BUILDKIT=1\""""
+            )
+
+            images_list.extend([f'"{sidecar_image}"', f'"{sidecar_latest_image}"'])
+
+        cloudbuild_yaml = f"""
+steps:
+{chr(10).join(build_steps)}
 timeout: 3600s
 images:
-- "{image}"
-- "{latest_tag}"
+{chr(10).join([f"- {img}" for img in images_list])}
 tags: [{image_tag}]
 options:
   logging: CLOUD_LOGGING_ONLY