streamline python packaging with uv (#135)

* Revert "add nemo-run as a git submodule (#186)"

This reverts commit 80f2c48.

* add uv based dockerfile

* add back version

* add some dockerfile comments

* keep subpackage sources in release images

* fix submodule install path

* docker image

* updating dockerfiles

* update uv docker

* update lockfile

* move uv dockerfile to ci, restructure devcontainer on ngc

* devcontainer attempts

* fix devcontainer

.devcontainer/devcontainer.json

@@ -7,8 +7,8 @@
 	"build": {
 		"context": "${localWorkspaceFolder}",
 		"dockerfile": "${localWorkspaceFolder}/Dockerfile",
-		"cacheFrom": "nvcr.io/nvidian/cvai_bnmo_trng/bionemo:bionemo2-latest",
-		"cacheFrom": "gitlab-master.nvidia.com:5005/clara-discovery/bionemo/bionemo-inter:bionemo2-latest-cache",
+		"cacheFrom": "type=registry,ref=nvcr.io/nvidian/cvai_bnmo_trng/bionemo:bionemo2-devcontainer-cache",
+		"cacheTo": "type=registry,ref=nvcr.io/nvidian/cvai_bnmo_trng/bionemo:bionemo2-devcontainer-cache,mode=max",
 		"target": "dev"
 	},
 	"mounts": [
@@ -34,7 +34,8 @@
 				"ms-azuretools.vscode-docker",
 				"charliermarsh.ruff",
 				"njpwerner.autodocstring",
-				"ms-toolsai.jupyter"
+				"ms-toolsai.jupyter",
+				"tamasfe.even-better-toml"
 			],
 			"settings": {
 				"python.analysis.extraPaths": [
@@ -46,10 +47,9 @@
 					"./sub-packages/bionemo-testing/src",
 					"./sub-packages/bionemo-example_model/src",
 					"./3rdparty/NeMo",
-					"./3rdparty/Megatron-LM",
-					"./3rdparty/NeMo-Run"
+					"./3rdparty/Megatron-LM"
 				],
-				"python.pythonPath": "/bin/python3",
+				"python.defaultInterpreterPath": "/usr/bin/python",
 				"python.testing.pytestEnabled": true,
 				"python.testing.pytestArgs": [
 					"sub-packages/",

.devcontainer/postCreateCommand.sh

@@ -1,10 +1,3 @@
 #!/bin/bash
 
-pip install --disable-pip-version-check --no-cache-dir -e 3rdparty/Megatron-LM
-pip install --disable-pip-version-check --no-cache-dir -e 3rdparty/NeMo[all]
-pip install --disable-pip-version-check --no-cache-dir -e 3rdparty/NeMo-Run
-
-for SUB_PKG in sub-packages/bionemo-*;
-do
-    pip install --disable-pip-version-check --no-cache-dir --no-deps -e $SUB_PKG
-done
+sudo pip install --no-deps --editable ./3rdparty/* ./sub-packages/bionemo-*

.dockerignore

@@ -4,3 +4,19 @@ MNISTCustom
 lightning_logs
 models
 test_data
+Dockerfile*
+.dockerignore
+.gitignore
+__pycache__
+*.pyc
+*.pyo
+*.pyd
+.Python
+env/
+venv/
+*.so
+.vscode/
+.idea/
+*.swp
+*.swo
+*.egg-info

.github/dependabot.yml

@@ -6,7 +6,7 @@ updates:
         interval: "weekly"
     directory: "/"
     target-branch: "main"
-    open-pull-requests-limit: 3
+    open-pull-requests-limit: 2
     reviewers:
       - "pstjohn"
       - "jstjohn"

.github/workflows/pre-commit.yml

@@ -14,6 +14,7 @@ jobs:
     - uses: actions/checkout@v4
       with:
         fetch-depth: 0
+        submodules: 'recursive'
     - uses: actions/setup-python@v5
       with:
         python-version: '3.10'

.gitignore

@@ -187,3 +187,5 @@ coverage.xml
 # System files
 .DS_Store
 Thumbs.db
+
+.python_history

.gitmodules

@@ -4,6 +4,3 @@
 [submodule "3rdparty/NeMo"]
 	path = 3rdparty/NeMo
 	url = https://github.com/NVIDIA/NeMo.git
-[submodule "3rdparty/NeMo-Run"]
-	path = 3rdparty/NeMo-Run
-	url = https://github.com/NVIDIA/NeMo-Run.git

.secrets.baseline

@@ -135,7 +135,16 @@
         "is_verified": false,
         "line_number": 83
       }
+    ],
+    "pyproject.toml": [
+      {
+        "type": "Hex High Entropy String",
+        "filename": "pyproject.toml",
+        "hashed_secret": "79670e9c9d1c7ea5b81a96a2053d81437712c78e",
+        "is_verified": false,
+        "line_number": 104
+      }
     ]
   },
-  "generated_at": "2024-08-22T16:20:49Z"
+  "generated_at": "2024-09-19T16:05:13Z"
 }

Dockerfile

@@ -47,43 +47,11 @@ RUN CAUSAL_CONV1D_FORCE_BUILD=TRUE pip --disable-pip-version-check --no-cache-di
 RUN pip --disable-pip-version-check --no-cache-dir install \
   git+https://github.com/state-spaces/mamba.git@v2.0.3
 
-FROM bionemo2-base AS pip-requirements
-
-# Copy and install pypi depedencies.
-RUN mkdir /tmp/pip-tmp
-WORKDIR /tmp/pip-tmp
-
-COPY requirements-dev.txt requirements-test.txt requirements-cve.txt /tmp/pip-tmp/
-
-# We want to only copy the requirements.txt, setup.py, and pyproject.toml files for *ALL* sub-packages
-# but we **can't** do COPY sub-packages/**/{requirements.txt,...} /<destination> because this will overwrite!
-# So....we copy everything into a temporary image and remove everything else!
-# Later, we can copy the result from the temporary image and get what we want
-# **WITHOUT** invalidating the cache for successive layers!
-COPY sub-packages/ /tmp/pip-tmp/sub-packages
-# remove all directories that aren't the top-level sub-packages/bionemo-{xyz}
-RUN find sub-packages/ -type d | grep "bionemo-[a-zA-Z0-9\-]*/" | xargs rm -rf && \
-    # only keep the requirements-related files
-    find sub-packages/ -type f | grep -v -E "requirements.txt|pyproject.toml|setup.py" | xargs rm
-
-FROM bionemo2-base AS dev
+RUN pip install hatchling   # needed to install nemo-run
+ARG NEMU_RUN_TAG=34259bd3e752fef94045a9a019e4aaf62bd11ce2
+RUN pip install nemo_run@git+https://github.com/NVIDIA/NeMo-Run.git@${NEMU_RUN_TAG}
 
 RUN mkdir -p /workspace/bionemo2/
-WORKDIR /workspace/bionemo2
-
-# We get the sub-packcages/ top-level structure + requirements.txt files
-COPY --from=pip-requirements /tmp/pip-tmp/ /workspace/bionemo2/
-
-RUN pip install -r requirements-dev.txt -r requirements-test.txt -r requirements-cve.txt
-
-# We calculate paths to each requirements.txt file and dynamically construct the pip install command.
-# This command will expand to something like:
-#   pip install --disable-pip-version-check --no-cache-dir \
-#      -r bionemo-core/requirements.txt \
-#      -r bionemo-pytorch/requirements.txt \
-#      -r bionemo-lmm/requirements.txt \
-#      (etc.)
-RUN X=""; for sub in $(echo sub-packages/bionemo-*); do X="-r ${sub}/requirements.txt ${X}"; done; eval "pip install --disable-pip-version-check --no-cache-dir ${X}"
 
 # Delete the temporary /build directory.
 WORKDIR /workspace
@@ -92,16 +60,54 @@ RUN rm -rf /build
 # Addressing Security Scan Vulnerabilities
 RUN rm -rf /opt/pytorch/pytorch/third_party/onnx
 RUN apt-get update  && \
-    apt-get install -y openssh-client=1:8.9p1-3ubuntu0.10 && \
-    rm -rf /var/lib/apt/lists/*
+  apt-get install -y openssh-client=1:8.9p1-3ubuntu0.10 && \
+  rm -rf /var/lib/apt/lists/*
 RUN apt purge -y libslurm37 libpmi2-0 && \
-    apt autoremove -y
+  apt autoremove -y
 RUN source /usr/local/nvm/nvm.sh && \
-    NODE_VER=$(nvm current) && \
-    nvm deactivate && \
-    nvm uninstall $NODE_VER && \
-    sed -i "/NVM/d" /root/.bashrc && \
-    sed -i "/nvm.sh/d" /etc/bash.bashrc
+  NODE_VER=$(nvm current) && \
+  nvm deactivate && \
+  nvm uninstall $NODE_VER && \
+  sed -i "/NVM/d" /root/.bashrc && \
+  sed -i "/nvm.sh/d" /etc/bash.bashrc
+
+# Use UV to install python packages from the workspace. This just installs packages into the system's python
+# environment, and does not use the current uv.lock file.
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
+ENV UV_LINK_MODE=copy \
+  UV_COMPILE_BYTECODE=1 \
+  UV_PYTHON_DOWNLOADS=never \
+  UV_SYSTEM_PYTHON=true
+
+WORKDIR /workspace/bionemo2
+
+# Install 3rd-party deps and bionemo submodules.
+COPY ./3rdparty /workspace/bionemo2/3rdparty
+COPY ./sub-packages /workspace/bionemo2/sub-packages
+
+# Note, we need to mount the .git folder here so that setuptools-scm is able to fetch git tag for version.
+RUN --mount=type=bind,source=./.git,target=./.git \
+  --mount=type=cache,id=uv-cache,target=/root/.cache,sharing=locked \
+  <<EOT
+uv pip install --no-build-isolation -v ./3rdparty/* ./sub-packages/bionemo-*
+rm -rf ./3rdparty
+rm -rf /tmp/*
+EOT
+
+# In the devcontainer image, we just copy over the finished `dist-packages` folder from the build image back into the
+# base pytorch container. We can then set up a non-root user and uninstall the bionemo and 3rd-party packages, so that
+# they can be installed in an editable fashion from the workspace directory. This lets us install all the package
+# dependencies in a cached fashion, so they don't have to be built from scratch every time the devcontainer is rebuilt.
+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 \
+  <<EOT
+apt-get update -qy
+apt-get install -qyy \
+  sudo
+rm -rf /tmp/* /var/tmp/*
+EOT
 
 # Create a non-root user to use inside a devcontainer.
 ARG USERNAME=bionemo
@@ -112,35 +118,20 @@ RUN groupadd --gid $USER_GID $USERNAME \
   && echo $USERNAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME \
   && chmod 0440 /etc/sudoers.d/$USERNAME
 
-RUN find /usr/local/lib/python3.10/dist-packages/ -type f -print0 | xargs -0 -P 0 -n 10000 chown $USERNAME:$USER_GID
-
-ENV PATH="/home/bionemo/.local/bin:${PATH}"
-
-
-# Create a release image with bionemo2 installed.
-FROM dev AS release
-
-# Install 3rd-party deps
-COPY ./3rdparty /build
-WORKDIR /build/Megatron-LM
-RUN pip install --disable-pip-version-check --no-cache-dir .
-
-WORKDIR /build/NeMo
-RUN pip install --disable-pip-version-check --no-cache-dir .[all]
+RUN rm -rf /usr/local/lib/python3.10/dist-packages
+COPY --from=bionemo2-base --chown=$USERNAME:$USERNAME --chmod=777 \
+  /usr/local/lib/python3.10/dist-packages /usr/local/lib/python3.10/dist-packages
+RUN <<EOT
+  rm -rf /usr/local/lib/python3.10/dist-packages/bionemo*
+  pip uninstall -y nemo_toolkit megatron_core
+EOT
 
-WORKDIR /build/NeMo-Run
-RUN SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0.dev1 pip install --disable-pip-version-check --no-cache-dir .
+# The 'release' target needs to be last so that it's the default build target. In the future, we could consider a setup
+# similar to the devcontainer above, where we copy the dist-packages folder from the build image into the release image.
+# This would reduce the overall image size by reducing the number of intermediate layers. In the meantime, we match the
+# existing release image build by copying over remaining files from the repo into the container.
+FROM bionemo2-base AS release
 
-WORKDIR /workspace
-RUN rm -rf /build
-
-# Install bionemo2 submodules
-WORKDIR /workspace/bionemo2/
 COPY VERSION .
-COPY ./sub-packages /workspace/bionemo2/sub-packages
-# Dynamically install the code for each bionemo namespace package.
-RUN for sub in sub-packages/bionemo-*; do pushd ${sub} && pip install --no-build-isolation --no-cache-dir --disable-pip-version-check --no-deps -e . && popd; done
-
-WORKDIR /workspace/bionemo2/
 COPY ./scripts ./scripts
 COPY ./README.md ./