working colocated example

Author: lkolluru05

Dockerfile.colocated

@@ -0,0 +1,60 @@
+# Use the JAX image with the custom-built sidecar as the base.
+
+FROM gcr.io/cloud-tpu-multipod-dev/sujinesh_sidecar_debug@sha256:9c9bea3836db6ef736abcdf1838074d8ba18e55e6c59272f19a1d15b1f9819aa
+# Defines a build argument for the requirements file. This contains the user's custom
+
+# Set the working directory (this is already inherited)
+WORKDIR /app
+
+# Copy the user's requirements file into the image.
+COPY . .
+
+# Install the additional user-provided dependencies, strictly enforcing the rules
+# from the base image's constraints file.
+# First, generate the requirements and a custom constraints file in its own layer for better caching.
+RUN \
+    # Install toml, generate the requirements files, and then clean up.
+    pip install toml && \
+    python3 colocated/extract_deps.py && \
+    pip uninstall -y toml && \
+    \
+    # Create a new constraints file that includes the correct jaxlib version.
+    JAXLIB_VERSION=$(pip show jaxlib | grep Version | awk '{print $2}') && \
+    cat /opt/venv/server_constraints.txt > /tmp/constraints.txt && \
+    echo "jaxlib==$JAXLIB_VERSION" >> /tmp/constraints.txt
+
+# Now, install the dependencies and the axlearn package.
+RUN \
+    # Capture the initial jax and jaxlib versions from the base image.
+    JAX_VERSION_BEFORE=$(pip show jax | grep Version | awk '{print $2}') && \
+    JAXLIB_VERSION_BEFORE=$(pip show jaxlib | grep Version | awk '{print $2}') && \
+    \
+    # Install the main dependencies, using our custom constraints.
+    uv pip install -r /tmp/requirements.txt -c /tmp/constraints.txt &&  \
+    \
+    # Install the JAX-dependent packages without resolving their dependencies.
+    uv pip install --no-deps -r /tmp/requirements_nodeps.txt && \
+    \
+    # Install the axlearn package itself, without its dependencies.
+    uv pip install --no-deps . && \
+    \
+    # Clean up the temporary files and the cache.
+    rm /tmp/requirements.txt && \
+    rm /tmp/requirements_nodeps.txt && \
+    rm /tmp/constraints.txt && \
+    uv cache clean && \
+    \
+    # Capture the final jax and jaxlib versions.
+    JAX_VERSION_AFTER=$(pip show jax | grep Version | awk '{print $2}') && \
+    JAXLIB_VERSION_AFTER=$(pip show jaxlib | grep Version | awk '{print $2}') && \
+    \
+    # Verify that the versions have not changed.
+    if [ "$JAX_VERSION_BEFORE" != "$JAX_VERSION_AFTER" ] || [ "$JAXLIB_VERSION_BEFORE" != "$JAXLIB_VERSION_AFTER" ]; then \
+        echo "ERROR: jax or jaxlib version changed!" >&2; \
+        echo "jax version before: $JAX_VERSION_BEFORE, after: $JAX_VERSION_AFTER" >&2; \
+        echo "jaxlib version before: $JAXLIB_VERSION_BEFORE, after: $JAXLIB_VERSION_AFTER" >&2; \
+        exit 1; \
+    fi
+
+# Note: The ENTRYPOINT and CMD are inherited from the base image, so they do not
+# need to be redefined here. I.e. the sidecar will be launched automatically.

axlearn/cloud/gcp/bundler.py

@@ -48,6 +48,7 @@
 
 import os
 import subprocess
+import time
 from typing import Optional
 
 from absl import app, flags, logging
@@ -58,10 +59,10 @@
 from axlearn.cloud.common.bundler import register_bundler
 from axlearn.cloud.common.docker import registry_from_repo
 from axlearn.cloud.common.utils import canonicalize_to_list, to_bool
-from axlearn.cloud.gcp.cloud_build import wait_for_cloud_build
+from axlearn.cloud.gcp.cloud_build import get_cloud_build_status
 from axlearn.cloud.gcp.config import gcp_settings
 from axlearn.cloud.gcp.utils import common_flags
-from axlearn.common.config import REQUIRED, Required, config_class, maybe_set_config
+from axlearn.common.config import REQUIRED, Required, config_class, maybe_set_config, config_for_class
 
 FLAGS = flags.FLAGS
 
@@ -98,19 +99,77 @@ class ArtifactRegistryBundler(DockerBundler):
 
     TYPE = "artifactregistry"
 
+    @config_class
+    class Config(DockerBundler.Config):
+        """Configures CloudBuildBundler.
+
+        Attributes:
+            colocated_image_required: Bool to build a colocated image
+        """
+        # Build image asynchronously.
+        colocated_image_required: bool = False
+        colocated_image_name: str = None
+        colocated_dockerfile: str = None
+        
+
     @classmethod
     def from_spec(cls, spec: list[str], *, fv: Optional[flags.FlagValues]) -> DockerBundler.Config:
-        cfg = super().from_spec(spec, fv=fv)
+        cfg: ArtifactRegistryBundler.Config = super().from_spec(spec, fv=fv)
         cfg.repo = cfg.repo or gcp_settings("docker_repo", required=False, fv=fv)
         cfg.dockerfile = cfg.dockerfile or gcp_settings("default_dockerfile", required=False, fv=fv)
+        cfg.colocated_image_required = cfg.colocated_image_required or gcp_settings("colocated_image_required", required=False, fv=fv)
+        cfg.colocated_image_name = cfg.colocated_image_name or gcp_settings("colocated_image_name", required=False, fv=fv)
+        cfg.colocated_dockerfile = cfg.colocated_dockerfile or gcp_settings("colocated_dockerfile", required=False, fv=fv)
+        return cfg
+
+    def _build_and_push(self, *args, **kwargs):
+        cfg = self.config
+        subprocess.run(
+            ["gcloud", "auth", "configure-docker", registry_from_repo(cfg.repo)],
+            check=True,
+        )
+
+        print("actual",cfg)
+        actual_name = cfg.image
+        actual_dockerfile=cfg.dockerfile
+        actual_target=cfg.target
+        if bool(cfg.colocated_image_required):
+            
+            cfg.dockerfile=cfg.colocated_dockerfile
+            cfg.image=cfg.colocated_image_name
+            cfg.target=None
+            print("updated config: ",cfg)
+            colocated_bundler_class = ColocatedArtifactRegistryBundler(cfg=cfg)
+            colocated_image_name = colocated_bundler_class.bundle(tag="latest")
+            print(colocated_image_name)
+        
+            cfg.dockerfile=actual_dockerfile
+            cfg.image=actual_name
+            cfg.target=actual_target
+
+
+            
+        return super()._build_and_push(*args, **kwargs)
+
+
+class ColocatedArtifactRegistryBundler(DockerBundler):
+    """A DockerBundler that reads configs from gcp_settings, and auths to Artifact Registry."""
+
+    @classmethod
+    def from_spec(cls, spec: list[str], *, fv: Optional[flags.FlagValues]) -> DockerBundler.Config:
+        cfg: ColocatedArtifactRegistryBundler.Config = super().from_spec(spec, fv=fv)
+        cfg.repo = cfg.repo or gcp_settings("docker_repo", required=False, fv=fv)
+        cfg.dockerfile = cfg.colocated_dockerfile or gcp_settings("colocated_dockerfile", required=False, fv=fv)
         return cfg
 
     def _build_and_push(self, *args, **kwargs):
         cfg = self.config
+        print("colocated",cfg)
         subprocess.run(
             ["gcloud", "auth", "configure-docker", registry_from_repo(cfg.repo)],
             check=True,
         )
+        
         return super()._build_and_push(*args, **kwargs)
 
 
@@ -237,14 +296,36 @@ def wait_until_finished(self, name: str, wait_timeout=3600):
             TimeoutError: If the build does not complete within the overall timeout.
             ValueError: If the async build fails.
         """
+        start_time = time.perf_counter()
         cfg: CloudBuildBundler.Config = self.config
-        if cfg.is_async:
-            wait_for_cloud_build(
-                project_id=cfg.project,
-                image_id=self.id(name),
-                tags=[name],
-                wait_timeout=wait_timeout,
-            )
+        while cfg.is_async:
+            elapsed_time = time.perf_counter() - start_time
+            if elapsed_time > wait_timeout:
+                timeout_msg = (
+                    f"Timed out waiting for CloudBuild to finish for more than "
+                    f"{wait_timeout} seconds."
+                )
+                logging.error(timeout_msg)
+                raise TimeoutError(timeout_msg)
+            try:
+                build_status = get_cloud_build_status(
+                    project_id=cfg.project, image_name=self.id(name), tags=[name]
+                )
+            except Exception as e:  # pylint: disable=broad-except
+                # TODO(liang-he,markblee): Distinguish transient from non-transient errors.
+                logging.warning("Failed to get the CloudBuild status, will retry: %s", e)
+            else:
+                if not build_status:
+                    logging.warning("CloudBuild for %s does not exist yet.", name)
+                elif build_status.is_pending():
+                    logging.info("CloudBuild for %s is pending: %s.", name, build_status)
+                elif build_status.is_success():
+                    logging.info("CloudBuild for %s is successful: %s.", name, build_status)
+                    return
+                else:
+                    # Unknown status is also considered a failure.
+                    raise RuntimeError(f"CloudBuild for {name} failed: {build_status}.")
+            time.sleep(30)
 
 
 def with_tpu_extras(bundler: Bundler.Config) -> Bundler.Config:
@@ -263,4 +344,4 @@ def with_tpu_extras(bundler: Bundler.Config) -> Bundler.Config:
 if __name__ == "__main__":
     common_flags()
     bundler_main_flags()
-    app.run(bundler_main)
+    app.run(bundler_main)