Relaxing the right constraint

Author: parambole

docker_build_dependency_image.sh

@@ -30,7 +30,7 @@ set -e
 export LOCAL_IMAGE_NAME=maxdiffusion_base_image
 
 # Use Docker BuildKit so we can cache pip packages.
-export DOCKER_BUILDKIT=1
+#export DOCKER_BUILDKIT=1
 
 echo "Starting to build your docker image. This will take a few minutes but the image can be reused as you iterate."
 
@@ -88,4 +88,4 @@ else
       -t ${LOCAL_IMAGE_NAME} \
       -f maxdiffusion_dependencies.Dockerfile .
   fi
-fi
+fi

maxdiffusion_dependencies.Dockerfile

@@ -52,7 +52,7 @@ COPY . .
 
 RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION"
 
-RUN --mount=type=cache,target=/root/.cache/pip bash setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION}
+RUN bash setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION}
 
 # Cleanup
-RUN rm -rf /root/.cache/pip
+RUN rm -rf /root/.cache/pip

requirements.txt

@@ -19,12 +19,12 @@ Pillow
 pylint
 pyink
 pytest==8.2.2
-tensorflow==2.17.0
+tensorflow>=2.17.0
 tensorflow-datasets>=4.9.6
 ruff>=0.1.5,<=0.2
 git+https://github.com/mlperf/logging.git
 opencv-python-headless==4.10.0.84
-orbax-checkpoint>=0.10.3
+orbax-checkpoint==0.10.3
 tokenizers==0.21.0
 huggingface_hub==0.30.2
 transformers==4.48.1

requirements_with_jax_stable_stack.txt

@@ -14,7 +14,7 @@ jaxlib>=0.4.30
 Jinja2
 opencv-python-headless==4.10.0.84
 optax>=0.2.3
-orbax-checkpoint>=0.10.3
+orbax-checkpoint==0.10.3
 parameterized
 Pillow
 pyink