use uv cache in docker

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Dockerfile

@@ -91,12 +91,12 @@ ENV UV_LINK_MODE=copy \
   UV_COMPILE_BYTECODE=1 \
   UV_PYTHON_DOWNLOADS=never \
   UV_SYSTEM_PYTHON=true \
-  UV_NO_CACHE=1 \
   UV_BREAK_SYSTEM_PACKAGES=1
 
 # Install the bionemo-geometric requirements ahead of copying over the rest of the repo, so that we can cache their
 # installation. These involve building some torch extensions, so they can take a while to install.
 RUN --mount=type=bind,source=./sub-packages/bionemo-geometric/requirements.txt,target=/requirements-pyg.txt \
+    --mount=type=cache,target=/root/.cache \
   uv pip install --no-build-isolation -r /requirements-pyg.txt
 
 COPY --from=rust-env /usr/local/cargo /usr/local/cargo
@@ -105,7 +105,7 @@ COPY --from=rust-env /usr/local/rustup /usr/local/rustup
 ENV PATH="/usr/local/cargo/bin:/usr/local/rustup/bin:${PATH}"
 ENV RUSTUP_HOME="/usr/local/rustup"
 
-RUN <<EOF
+RUN --mount=type=cache,target=/root/.cache <<EOF
 set -eo pipefail
 uv pip install maturin --no-build-isolation
 
@@ -115,7 +115,6 @@ sed -i 's/^Version: 0\.0\.0$/Version: 0.1.45/' \
   /usr/local/lib/python3.12/dist-packages/tensorstore-0.0.0.dist-info/METADATA
 mv /usr/local/lib/python3.12/dist-packages/tensorstore-0.0.0.dist-info \
 /usr/local/lib/python3.12/dist-packages/tensorstore-0.1.45.dist-info
-rm -rf /root/.cache/*
 EOF
 
 WORKDIR /workspace/bionemo2
@@ -129,9 +128,9 @@ COPY ./sub-packages /workspace/bionemo2/sub-packages
 # Includes a hack to install tensorstore 0.1.45, which doesn't distribute a pypi wheel for python 3.12, and the metadata
 # in the source distribution doesn't match the expected pypi version.
 RUN --mount=type=bind,source=./.git,target=./.git \
-  --mount=type=bind,source=./requirements-test.txt,target=/requirements-test.txt \
-  --mount=type=bind,source=./requirements-cve.txt,target=/requirements-cve.txt \
-  <<EOF
+    --mount=type=bind,source=./requirements-test.txt,target=/requirements-test.txt \
+    --mount=type=bind,source=./requirements-cve.txt,target=/requirements-cve.txt \
+    --mount=type=cache,target=/root/.cache <<EOF
 set -eo pipefail
 
 uv pip install --no-build-isolation \
@@ -152,8 +151,8 @@ EOF
 FROM ${BASE_IMAGE} AS dev
 
 RUN --mount=type=cache,id=apt-cache,target=/var/cache/apt,sharing=locked \
-  --mount=type=cache,id=apt-lib,target=/var/lib/apt,sharing=locked \
-  <<EOF
+    --mount=type=cache,id=apt-lib,target=/var/lib/apt,sharing=locked \
+    <<EOF
 set -eo pipefail
 apt-get update -qy
 apt-get install -qyy \
@@ -196,7 +195,7 @@ ENV PATH="/usr/local/cargo/bin:/usr/local/rustup/bin:${PATH}"
 ENV RUSTUP_HOME="/usr/local/rustup"
 
 RUN --mount=type=bind,source=./requirements-dev.txt,target=/workspace/bionemo2/requirements-dev.txt \
-  --mount=type=cache,id=uv-cache,target=/root/.cache,sharing=locked <<EOF
+    --mount=type=cache,target=/root/.cache <<EOF
   set -eo pipefail
   uv pip install -r /workspace/bionemo2/requirements-dev.txt
   rm -rf /tmp/*

pyproject.toml

@@ -76,6 +76,7 @@ dev-dependencies = [
     "tenacity",
 ]
 no-build-isolation-package = ["flash-attn"]
+cache-keys = [{ git = { commit = true } }]
 
 [tool.black]
 line-length = 119