Open
Conversation
Contributor
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughESM2 model architecture refactored to rename the base model namespace from "esm" to "model" across core components, model classes, and related tests. A new Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Collaborator
|
@gagank1 : Are you working on a readme? |
Collaborator
Author
|
/ok to test |
@gagank1, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/ |
gagank1
requested review from
cspades,
dorotat-nv,
jomitchellnv,
jstjohn,
jwilber,
pstjohn,
savitha-eng and
trvachov
as code owners
Collaborator
Author
|
/ok to test |
@gagank1, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/ |
Collaborator
Author
|
/ok to test c34c09b |
Collaborator
Author
|
@coderabbitai review |
Contributor
✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (3)
bionemo-recipes/recipes/esm2_peft_te/example_8m_checkpoint/esm_nv.py (1)
405-407: Same_tied_weights_keysdict type concern.See comment on
modeling_esm_te.py.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@bionemo-recipes/recipes/esm2_peft_te/example_8m_checkpoint/esm_nv.py` around lines 405 - 407, The class variable _tied_weights_keys uses the modern built-in generic dict[...] typing which may be incompatible with the rest of the codebase; change its annotation to use typing.Dict[str, str] (and ensure Dict is imported) or use typing.Mapping if immutability is desired, mirroring the fix applied in modeling_esm_te.py so the declaration becomes ClassVar[Dict[str, str]] with the same key/value entries retained.bionemo-recipes/recipes/esm2_native_te/example_8m_checkpoint/esm_nv.py (1)
405-407: Same_tied_weights_keysdict type concern as inmodeling_esm_te.py.See comment on the canonical file — HF expects
list[str], notdict[str, str].🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@bionemo-recipes/recipes/esm2_native_te/example_8m_checkpoint/esm_nv.py` around lines 405 - 407, The _tied_weights_keys typed as ClassVar[dict[str, str]] should be changed to ClassVar[list[str]] to match HF expectations (same fix as in modeling_esm_te.py); replace the dict literal with a list of the relevant parameter names (e.g. ["lm_head.decoder.weight", "model.embeddings.word_embeddings.weight"]) and update any usages that assume dict semantics to use the list order or explicit pairing where needed.bionemo-recipes/recipes/esm2_accelerate_te/example_8m_checkpoint/esm_nv.py (1)
405-407: Same_tied_weights_keysdict type concern.See comment on
modeling_esm_te.py.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@bionemo-recipes/recipes/esm2_accelerate_te/example_8m_checkpoint/esm_nv.py` around lines 405 - 407, _tied_weights_keys is annotated as a mutable dict ClassVar which repeats the same typing concern as in modeling_esm_te.py; change the annotation to an immutable mapping type (e.g., ClassVar[Mapping[str, str]] from typing) and, to avoid accidental mutation, assign a read-only view (e.g., types.MappingProxyType({"lm_head.decoder.weight": "model.embeddings.word_embeddings.weight"})); update the import list to include typing.Mapping and types if not present and mirror the same pattern used/fixed in modeling_esm_te.py.
🧹 Nitpick comments (4)
bionemo-recipes/vllm/launch.sh (1)
50-50:exec $DOCKER_CMDis unquoted — word splitting will break paths with spaces.If
PROJECT_ROOTcontains spaces (e.g.,/home/user/my projects/...), the-vargument will be incorrectly split into multiple tokens. Use a Bash array to avoid this:🔧 Proposed fix (array-based approach)
Replace the string-based
DOCKER_CMDwith an array throughout the script:-DOCKER_CMD="docker run -itd ..." +DOCKER_CMD=("docker" "run" "-itd" "--gpus" "all" "--network" "host" "--ipc=host" "-e" "HF_TOKEN" "--rm" "--name" "${CONTAINER}_dev") if [ "$MOUNT_DIR" = true ]; then PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" - DOCKER_CMD="$DOCKER_CMD -v ${PROJECT_ROOT}:/workspace/bionemo-framework" + DOCKER_CMD+=("-v" "${PROJECT_ROOT}:/workspace/bionemo-framework") fi -DOCKER_CMD="$DOCKER_CMD $CONTAINER /bin/bash" +DOCKER_CMD+=("$CONTAINER" "/bin/bash") -exec $DOCKER_CMD +exec "${DOCKER_CMD[@]}"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@bionemo-recipes/vllm/launch.sh` at line 50, The exec $DOCKER_CMD call uses an unquoted string which allows word-splitting and breaks mount paths with spaces (e.g., PROJECT_ROOT used in -v). Convert the string DOCKER_CMD into a Bash array (e.g., DOCKER_CMD=(docker run ...)) and update all places that build/append to DOCKER_CMD so they push elements into the array, then replace exec $DOCKER_CMD with exec "${DOCKER_CMD[@]}" so each argument (including the -v PROJECT_ROOT value) is preserved; update any helper code that concatenates DOCKER_CMD to use array operations instead.bionemo-recipes/vllm/Dockerfile (2)
2-3: Base image is hosted on an internal NVIDIA GitLab registry — not pullable outside NVIDIA.
gitlab-master.nvidia.com:5005/dl/dgx/vllm:main-py3.43005406-develrequires internal network/credentials access. The commented-outnvcr.io/nvidia/vllm:26.01-py3alternative on line 1 is the publicly accessible equivalent. Once an NGC release with vLLM ≥ 0.14 becomes available, switching to the public image will make this recipe usable by external contributors without additional setup.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@bionemo-recipes/vllm/Dockerfile` around lines 2 - 3, The Dockerfile's FROM line currently uses an internal image tag "gitlab-master.nvidia.com:5005/dl/dgx/vllm:main-py3.43005406-devel" which is not pullable externally; replace that base image with the public equivalent "nvcr.io/nvidia/vllm:26.01-py3" (or parametrize the base via a build ARG) in the FROM instruction so external contributors can build without internal credentials, and retain a short comment noting the vLLM>=0.14 requirement and to switch back when an official public image with the needed version is available.
30-30: Pin thetransformer_engineversion for reproducible builds.
pip install --no-build-isolation transformer_engine[pytorch]with no version specifier will install whichever version is latest at build time. TE releases frequently and has had breaking API changes between major versions (e.g., 1.x → 2.x). A silent version bump can break the integration without any change to this file.🔧 Proposed fix
-RUN pip install --no-build-isolation transformer_engine[pytorch] +RUN pip install --no-build-isolation "transformer_engine[pytorch]==<verified_version>"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@bionemo-recipes/vllm/Dockerfile` at line 30, The Dockerfile currently installs transformer_engine without a version pin (the RUN pip install --no-build-isolation transformer_engine[pytorch] line); change that to install a specific, tested TE release by updating that RUN to include an exact version specifier (for example: RUN pip install --no-build-isolation transformer_engine==<MAJOR.MINOR.PATCH>[pytorch]) so builds are reproducible and won’t break on upstream major/minor bumps; optionally add a short comment noting the chosen compatible version.bionemo-recipes/vllm/test_esm2_golden_values.py (1)
46-63:sys.path.insert+os.chdiris fragile for test infrastructure.
sys.path.insert(0, ...)at module level (line 46) andos.chdirinsidefresh_exportmake this script sensitive to working directory and import order. This is acceptable for a standalone validation script run manually inside a container, but consider adding a note that this is not designed to run as part of a pytest suite.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@bionemo-recipes/vllm/test_esm2_golden_values.py` around lines 46 - 63, The test module mutates import paths and working directory (sys.path.insert(0, ...), os.chdir(...) inside fresh_export) which is fragile for pytest; update the file to document this by adding a clear module-level comment or docstring near sys.path.insert and a brief note on fresh_export explaining it intentionally changes cwd for export_hf_checkpoint and that the script is not intended to be run under pytest/parallel test runners (referencing sys.path.insert, ESM2_MODEL_DIR, and fresh_export by name), or alternatively guard execution with a main-check so pytest won't import/run it implicitly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@bionemo-recipes/vllm/Dockerfile`:
- Around line 1-36: Add a non-root user and switch to it in the Dockerfile so
the container does not run as root; create a user/group (e.g., bionemo), chown
the application directory (/workspace/bionemo) and any cache/venv paths to that
user, and add a final USER bionemo line. Locate the Dockerfile sections around
WORKDIR /workspace/bionemo and COPY . . and insert user creation
(useradd/addgroup or groupadd) and chown before switching context, then add USER
bionemo at the end so subsequent runtime processes run unprivileged.
In `@bionemo-recipes/vllm/launch.sh`:
- Around line 36-40: The script hardcodes "--name vllm_dev" in DOCKER_CMD which
prevents multiple instances and misrepresents the positional $CONTAINER
argument; update the usage comment to clarify that the positional argument is
the image name (or image_name) and change the DOCKER_CMD assignments in the
HEADLESS branch to use a container name derived from $CONTAINER (for example
--name "$CONTAINER" or --name "${CONTAINER}_dev") instead of "vllm_dev" so the
image argument is also used as the container name and avoids name collisions.
In `@bionemo-recipes/vllm/README.md`:
- Around line 21-23: Change the fenced code block that currently uses the wrong
language tag; locate the block containing the shell command "python
test_esm2_golden_values.py" in the README and replace the opening fence language
identifier from ```python to ```bash so the command is treated as a shell
snippet and highlighted correctly.
- Line 17: The README's "or use `launch.sh`" is misleading because `launch.sh
--mount_dir` mounts the repo to `/workspace/bionemo-framework` while the manual
`docker run` mounts to `/workspace/bionemo`, causing `python
test_esm2_golden_values.py` to fail when run from WORKDIR; fix by either (A)
update `launch.sh` to mount the project root into `/workspace/bionemo` instead
of `/workspace/bionemo-framework` (adjust mount target and any downstream path
references in launch.sh), or (B) update README.md to explicitly document that
`launch.sh --mount_dir` mounts at `/workspace/bionemo-framework` and instruct
users to cd into the correct subdirectory (where `test_esm2_golden_values.py`
lives) before running the test; mention the affected files `launch.sh`,
`README.md`, and `test_esm2_golden_values.py` so reviewers can locate the
changes.
In `@bionemo-recipes/vllm/test_esm2_golden_values.py`:
- Around line 147-210: The script currently only prints comparisons and never
fails; add assertions that enforce the golden-value tolerances using RTOL and
ATOL: for each pair in pairs (refer to the pairs list and variables a, b),
assert np.allclose(a, b, rtol=RTOL, atol=ATOL) (or equivalently assert
(np.abs(a-b) <= ATOL + RTOL * np.abs(b)).all()) and fail the test if not, and
also assert cosine_sim(a, b) exceeds a sensible threshold or that exact is True
when ATOL/RTOL are zero; add per-sequence assertions inside the per-sequence
loop to ensure each sequence max-diff <= ATOL + RTOL * max(|b_i|) so the test
fails on unacceptable drift.
---
Duplicate comments:
In `@bionemo-recipes/recipes/esm2_accelerate_te/example_8m_checkpoint/esm_nv.py`:
- Around line 405-407: _tied_weights_keys is annotated as a mutable dict
ClassVar which repeats the same typing concern as in modeling_esm_te.py; change
the annotation to an immutable mapping type (e.g., ClassVar[Mapping[str, str]]
from typing) and, to avoid accidental mutation, assign a read-only view (e.g.,
types.MappingProxyType({"lm_head.decoder.weight":
"model.embeddings.word_embeddings.weight"})); update the import list to include
typing.Mapping and types if not present and mirror the same pattern used/fixed
in modeling_esm_te.py.
In `@bionemo-recipes/recipes/esm2_native_te/example_8m_checkpoint/esm_nv.py`:
- Around line 405-407: The _tied_weights_keys typed as ClassVar[dict[str, str]]
should be changed to ClassVar[list[str]] to match HF expectations (same fix as
in modeling_esm_te.py); replace the dict literal with a list of the relevant
parameter names (e.g. ["lm_head.decoder.weight",
"model.embeddings.word_embeddings.weight"]) and update any usages that assume
dict semantics to use the list order or explicit pairing where needed.
In `@bionemo-recipes/recipes/esm2_peft_te/example_8m_checkpoint/esm_nv.py`:
- Around line 405-407: The class variable _tied_weights_keys uses the modern
built-in generic dict[...] typing which may be incompatible with the rest of the
codebase; change its annotation to use typing.Dict[str, str] (and ensure Dict is
imported) or use typing.Mapping if immutability is desired, mirroring the fix
applied in modeling_esm_te.py so the declaration becomes ClassVar[Dict[str,
str]] with the same key/value entries retained.
---
Nitpick comments:
In `@bionemo-recipes/vllm/Dockerfile`:
- Around line 2-3: The Dockerfile's FROM line currently uses an internal image
tag "gitlab-master.nvidia.com:5005/dl/dgx/vllm:main-py3.43005406-devel" which is
not pullable externally; replace that base image with the public equivalent
"nvcr.io/nvidia/vllm:26.01-py3" (or parametrize the base via a build ARG) in the
FROM instruction so external contributors can build without internal
credentials, and retain a short comment noting the vLLM>=0.14 requirement and to
switch back when an official public image with the needed version is available.
- Line 30: The Dockerfile currently installs transformer_engine without a
version pin (the RUN pip install --no-build-isolation
transformer_engine[pytorch] line); change that to install a specific, tested TE
release by updating that RUN to include an exact version specifier (for example:
RUN pip install --no-build-isolation
transformer_engine==<MAJOR.MINOR.PATCH>[pytorch]) so builds are reproducible and
won’t break on upstream major/minor bumps; optionally add a short comment noting
the chosen compatible version.
In `@bionemo-recipes/vllm/launch.sh`:
- Line 50: The exec $DOCKER_CMD call uses an unquoted string which allows
word-splitting and breaks mount paths with spaces (e.g., PROJECT_ROOT used in
-v). Convert the string DOCKER_CMD into a Bash array (e.g., DOCKER_CMD=(docker
run ...)) and update all places that build/append to DOCKER_CMD so they push
elements into the array, then replace exec $DOCKER_CMD with exec
"${DOCKER_CMD[@]}" so each argument (including the -v PROJECT_ROOT value) is
preserved; update any helper code that concatenates DOCKER_CMD to use array
operations instead.
In `@bionemo-recipes/vllm/test_esm2_golden_values.py`:
- Around line 46-63: The test module mutates import paths and working directory
(sys.path.insert(0, ...), os.chdir(...) inside fresh_export) which is fragile
for pytest; update the file to document this by adding a clear module-level
comment or docstring near sys.path.insert and a brief note on fresh_export
explaining it intentionally changes cwd for export_hf_checkpoint and that the
script is not intended to be run under pytest/parallel test runners (referencing
sys.path.insert, ESM2_MODEL_DIR, and fresh_export by name), or alternatively
guard execution with a main-check so pytest won't import/run it implicitly.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (20)
bionemo-recipes/models/esm2/convert.pybionemo-recipes/models/esm2/export.pybionemo-recipes/models/esm2/modeling_esm_te.pybionemo-recipes/models/esm2/tests/test_cp_bshd.pybionemo-recipes/models/esm2/tests/test_cp_thd.pybionemo-recipes/models/esm2/tests/test_distributed_fp8.pybionemo-recipes/models/esm2/tests/test_distributed_strategies.pybionemo-recipes/models/esm2/tests/test_modeling_esm_te.pybionemo-recipes/recipes/esm2_accelerate_te/example_8m_checkpoint/esm_nv.pybionemo-recipes/recipes/esm2_native_te/example_8m_checkpoint/esm_nv.pybionemo-recipes/recipes/esm2_native_te/tests/test_stop_and_go.pybionemo-recipes/recipes/esm2_native_te/train_ddp.pybionemo-recipes/recipes/esm2_native_te/train_ddp_cp.pybionemo-recipes/recipes/esm2_native_te/train_fsdp2.pybionemo-recipes/recipes/esm2_native_te/train_fsdp2_cp.pybionemo-recipes/recipes/esm2_peft_te/example_8m_checkpoint/esm_nv.pybionemo-recipes/vllm/Dockerfilebionemo-recipes/vllm/README.mdbionemo-recipes/vllm/launch.shbionemo-recipes/vllm/test_esm2_golden_values.py
bionemo-recipes/vllm/Dockerfile
Outdated
Comment on lines
+1
to
+36
| # FROM nvcr.io/nvidia/vllm:26.01-py3 | ||
| FROM gitlab-master.nvidia.com:5005/dl/dgx/vllm:main-py3.43005406-devel | ||
| # using this because we need vllm >= 0.14 to work with Transformers v5. no released nvidia version with this yet. | ||
|
|
||
| # The vLLM image has CUDA 13.1 runtime and nvcc, but missing dev headers (cusparse.h, nvtx, etc.) | ||
| # Install cuda-keyring to add NVIDIA's apt repo, then install the dev headers for transformer_engine | ||
| RUN apt-get update && apt-get install -y --no-install-recommends wget && \ | ||
| wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb && \ | ||
| dpkg -i cuda-keyring_1.1-1_all.deb && \ | ||
| rm cuda-keyring_1.1-1_all.deb && \ | ||
| apt-get update && apt-get install -y --no-install-recommends \ | ||
| cuda-nvtx-13-1 \ | ||
| cuda-cupti-dev-13-1 \ | ||
| cuda-nvml-dev-13-1 \ | ||
| libcusparse-dev-13-1 \ | ||
| libcusolver-dev-13-1 \ | ||
| libcufft-dev-13-1 \ | ||
| libnvjitlink-dev-13-1 \ | ||
| libnvjpeg-dev-13-1 \ | ||
| libcublasmp0-dev-cuda-13 \ | ||
| libcudnn9-cuda-13 \ | ||
| && rm -rf /var/lib/apt/lists/* | ||
|
|
||
| # Install remaining dependencies | ||
| RUN --mount=type=cache,target=/root/.cache/pip \ | ||
| --mount=type=bind,source=requirements.txt,target=/requirements.txt \ | ||
| pip install -r /requirements.txt | ||
|
|
||
| # Install transformer_engine from source (force build for CUDA 13.1, not pre-built cu12 wheel) | ||
| RUN pip install --no-build-isolation transformer_engine[pytorch] | ||
|
|
||
| RUN pip install transformers[torch]==5.0.0 | ||
|
|
||
|
|
||
| WORKDIR /workspace/bionemo | ||
| COPY . . |
Contributor
There was a problem hiding this comment.
Container runs as root — add a USER directive.
No USER command is present, so all processes run as root inside the container. This is flagged by Trivy (DS-0002) and violates the principle of least privilege.
🔒 Proposed fix
WORKDIR /workspace/bionemo
COPY . .
+
+RUN useradd -m -u 1000 appuser && chown -R appuser /workspace/bionemo
+USER appuser📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| # FROM nvcr.io/nvidia/vllm:26.01-py3 | |
| FROM gitlab-master.nvidia.com:5005/dl/dgx/vllm:main-py3.43005406-devel | |
| # using this because we need vllm >= 0.14 to work with Transformers v5. no released nvidia version with this yet. | |
| # The vLLM image has CUDA 13.1 runtime and nvcc, but missing dev headers (cusparse.h, nvtx, etc.) | |
| # Install cuda-keyring to add NVIDIA's apt repo, then install the dev headers for transformer_engine | |
| RUN apt-get update && apt-get install -y --no-install-recommends wget && \ | |
| wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb && \ | |
| dpkg -i cuda-keyring_1.1-1_all.deb && \ | |
| rm cuda-keyring_1.1-1_all.deb && \ | |
| apt-get update && apt-get install -y --no-install-recommends \ | |
| cuda-nvtx-13-1 \ | |
| cuda-cupti-dev-13-1 \ | |
| cuda-nvml-dev-13-1 \ | |
| libcusparse-dev-13-1 \ | |
| libcusolver-dev-13-1 \ | |
| libcufft-dev-13-1 \ | |
| libnvjitlink-dev-13-1 \ | |
| libnvjpeg-dev-13-1 \ | |
| libcublasmp0-dev-cuda-13 \ | |
| libcudnn9-cuda-13 \ | |
| && rm -rf /var/lib/apt/lists/* | |
| # Install remaining dependencies | |
| RUN --mount=type=cache,target=/root/.cache/pip \ | |
| --mount=type=bind,source=requirements.txt,target=/requirements.txt \ | |
| pip install -r /requirements.txt | |
| # Install transformer_engine from source (force build for CUDA 13.1, not pre-built cu12 wheel) | |
| RUN pip install --no-build-isolation transformer_engine[pytorch] | |
| RUN pip install transformers[torch]==5.0.0 | |
| WORKDIR /workspace/bionemo | |
| COPY . . | |
| # FROM nvcr.io/nvidia/vllm:26.01-py3 | |
| FROM gitlab-master.nvidia.com:5005/dl/dgx/vllm:main-py3.43005406-devel | |
| # using this because we need vllm >= 0.14 to work with Transformers v5. no released nvidia version with this yet. | |
| # The vLLM image has CUDA 13.1 runtime and nvcc, but missing dev headers (cusparse.h, nvtx, etc.) | |
| # Install cuda-keyring to add NVIDIA's apt repo, then install the dev headers for transformer_engine | |
| RUN apt-get update && apt-get install -y --no-install-recommends wget && \ | |
| wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb && \ | |
| dpkg -i cuda-keyring_1.1-1_all.deb && \ | |
| rm cuda-keyring_1.1-1_all.deb && \ | |
| apt-get update && apt-get install -y --no-install-recommends \ | |
| cuda-nvtx-13-1 \ | |
| cuda-cupti-dev-13-1 \ | |
| cuda-nvml-dev-13-1 \ | |
| libcusparse-dev-13-1 \ | |
| libcusolver-dev-13-1 \ | |
| libcufft-dev-13-1 \ | |
| libnvjitlink-dev-13-1 \ | |
| libnvjpeg-dev-13-1 \ | |
| libcublasmp0-dev-cuda-13 \ | |
| libcudnn9-cuda-13 \ | |
| && rm -rf /var/lib/apt/lists/* | |
| # Install remaining dependencies | |
| RUN --mount=type=cache,target=/root/.cache/pip \ | |
| --mount=type=bind,source=requirements.txt,target=/requirements.txt \ | |
| pip install -r /requirements.txt | |
| # Install transformer_engine from source (force build for CUDA 13.1, not pre-built cu12 wheel) | |
| RUN pip install --no-build-isolation transformer_engine[pytorch] | |
| RUN pip install transformers[torch]==5.0.0 | |
| WORKDIR /workspace/bionemo | |
| COPY . . | |
| RUN useradd -m -u 1000 appuser && chown -R appuser /workspace/bionemo | |
| USER appuser |
🧰 Tools
🪛 Trivy (0.69.1)
[error] 1-1: Image user should not be 'root'
Specify at least 1 USER command in Dockerfile with non-root user as argument
Rule: DS-0002
(IaC/Dockerfile)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@bionemo-recipes/vllm/Dockerfile` around lines 1 - 36, Add a non-root user and
switch to it in the Dockerfile so the container does not run as root; create a
user/group (e.g., bionemo), chown the application directory (/workspace/bionemo)
and any cache/venv paths to that user, and add a final USER bionemo line. Locate
the Dockerfile sections around WORKDIR /workspace/bionemo and COPY . . and
insert user creation (useradd/addgroup or groupadd) and chown before switching
context, then add USER bionemo at the end so subsequent runtime processes run
unprivileged.
bionemo-recipes/vllm/launch.sh
Outdated
Comment on lines
+36
to
+40
| if [ "$HEADLESS" = true ]; then | ||
| DOCKER_CMD="docker run -itd --gpus all --network host --ipc=host -e HF_TOKEN --rm --name vllm_dev" | ||
| else | ||
| DOCKER_CMD="docker run -it --gpus all --network host --ipc=host -e HF_TOKEN --rm --name vllm_dev" | ||
| fi |
Contributor
There was a problem hiding this comment.
Hardcoded --name vllm_dev prevents running multiple containers and shadows the $CONTAINER argument's meaning.
The script's positional argument ($CONTAINER) is the Docker image name, but the running container name is always vllm_dev. This has two problems:
- Starting a second instance (e.g., a second headless container) will fail with
docker: Error response from daemon: Conflict. The container name "/vllm_dev" is already in use. - The usage comment says
<container_name>but it's really<image_name>, which is confusing.
Consider using $CONTAINER as both the image name and the container name (or derive the container name from it):
🔧 Proposed fix
-DOCKER_CMD="docker run -itd --gpus all --network host --ipc=host -e HF_TOKEN --rm --name vllm_dev"
+DOCKER_CMD="docker run -itd --gpus all --network host --ipc=host -e HF_TOKEN --rm --name ${CONTAINER}_dev"-DOCKER_CMD="docker run -it --gpus all --network host --ipc=host -e HF_TOKEN --rm --name vllm_dev"
+DOCKER_CMD="docker run -it --gpus all --network host --ipc=host -e HF_TOKEN --rm --name ${CONTAINER}_dev"Also update the usage comment on line 4:
-# Usage: ./launch.sh <container_name> [--mount_dir] [--headless]
+# Usage: ./launch.sh <image_name> [--mount_dir] [--headless]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if [ "$HEADLESS" = true ]; then | |
| DOCKER_CMD="docker run -itd --gpus all --network host --ipc=host -e HF_TOKEN --rm --name vllm_dev" | |
| else | |
| DOCKER_CMD="docker run -it --gpus all --network host --ipc=host -e HF_TOKEN --rm --name vllm_dev" | |
| fi | |
| if [ "$HEADLESS" = true ]; then | |
| DOCKER_CMD="docker run -itd --gpus all --network host --ipc=host -e HF_TOKEN --rm --name ${CONTAINER}_dev" | |
| else | |
| DOCKER_CMD="docker run -it --gpus all --network host --ipc=host -e HF_TOKEN --rm --name ${CONTAINER}_dev" | |
| fi |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@bionemo-recipes/vllm/launch.sh` around lines 36 - 40, The script hardcodes
"--name vllm_dev" in DOCKER_CMD which prevents multiple instances and
misrepresents the positional $CONTAINER argument; update the usage comment to
clarify that the positional argument is the image name (or image_name) and
change the DOCKER_CMD assignments in the HEADLESS branch to use a container name
derived from $CONTAINER (for example --name "$CONTAINER" or --name
"${CONTAINER}_dev") instead of "vllm_dev" so the image argument is also used as
the container name and avoids name collisions.
bionemo-recipes/vllm/README.md
Outdated
| docker run -it --gpus all --network host --ipc=host -e HF_TOKEN --rm -v ${PWD}:/workspace/bionemo vllm /bin/bash | ||
| ``` | ||
|
|
||
| or use `launch.sh`. |
Contributor
There was a problem hiding this comment.
launch.sh mounts to a different container path than the manual docker run command — the README's test instruction will break after using launch.sh --mount_dir.
The manual command (line 14) mounts ${PWD} → /workspace/bionemo, which is the container's WORKDIR, so python test_esm2_golden_values.py works directly. However, launch.sh --mount_dir mounts the project root (two levels up) to /workspace/bionemo-framework, leaving /workspace/bionemo populated only by the image's COPY . . layer. Running python test_esm2_golden_values.py from WORKDIR after using --mount_dir will fail unless the user navigates to the correct subdirectory. The "or use launch.sh" phrasing implies equivalence — either clarify the path difference or update launch.sh to mount to /workspace/bionemo instead of /workspace/bionemo-framework.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@bionemo-recipes/vllm/README.md` at line 17, The README's "or use `launch.sh`"
is misleading because `launch.sh --mount_dir` mounts the repo to
`/workspace/bionemo-framework` while the manual `docker run` mounts to
`/workspace/bionemo`, causing `python test_esm2_golden_values.py` to fail when
run from WORKDIR; fix by either (A) update `launch.sh` to mount the project root
into `/workspace/bionemo` instead of `/workspace/bionemo-framework` (adjust
mount target and any downstream path references in launch.sh), or (B) update
README.md to explicitly document that `launch.sh --mount_dir` mounts at
`/workspace/bionemo-framework` and instruct users to cd into the correct
subdirectory (where `test_esm2_golden_values.py` lives) before running the test;
mention the affected files `launch.sh`, `README.md`, and
`test_esm2_golden_values.py` so reviewers can locate the changes.
bionemo-recipes/vllm/README.md
Outdated
Comment on lines
+21
to
+23
| ```python | ||
| python test_esm2_golden_values.py | ||
| ``` |
Contributor
There was a problem hiding this comment.
Wrong code fence language identifier — should be bash, not python.
The python test_esm2_golden_values.py command is a shell invocation, not Python source code. Using ```python causes syntax highlighters to misparse it.
📝 Proposed fix
-```python
+```bash
python test_esm2_golden_values.py</details>
<details>
<summary>🤖 Prompt for AI Agents</summary>
Verify each finding against the current code and only fix it if needed.
In @bionemo-recipes/vllm/README.md around lines 21 - 23, Change the fenced code
block that currently uses the wrong language tag; locate the block containing
the shell command "python test_esm2_golden_values.py" in the README and replace
the opening fence language identifier from python to bash so the command
is treated as a shell snippet and highlighted correctly.
</details>
<!-- fingerprinting:phantom:poseidon:churro -->
<!-- This is an auto-generated comment by CodeRabbit -->
Comment on lines
+147
to
+210
| if __name__ == "__main__": | ||
| print(f"GPUs: {torch.cuda.device_count()}") | ||
|
|
||
| # Step 0: fresh export (facebook HF -> our TE format) | ||
| print("\n[0/3] Exporting checkpoint ...") | ||
| MODEL_ID = fresh_export() | ||
|
|
||
| print(f"MODEL_ID: {MODEL_ID}") | ||
| print(f"REFERENCE_MODEL_ID: {REFERENCE_MODEL_ID}") | ||
| print(f"Sequences: {len(SEQUENCES)}") | ||
|
|
||
| # 1) vLLM on exported model | ||
| print("\n[1/3] vLLM inference (exported model) ...") | ||
| emb_vllm = vllm_embed(MODEL_ID, SEQUENCES) | ||
|
|
||
| # 2) HuggingFace on exported model | ||
| print("\n[2/3] HuggingFace inference (exported model) ...") | ||
| emb_hf_exported = hf_embed(MODEL_ID, SEQUENCES) | ||
|
|
||
| # 3) HuggingFace on reference Hub model | ||
| print("\n[3/3] HuggingFace inference (reference model) ...") | ||
| emb_hf_reference = hf_embed(REFERENCE_MODEL_ID, SEQUENCES) | ||
|
|
||
| # ---- Pairwise comparisons ---- | ||
| pairs = [ | ||
| ("vLLM (exported)", "HF (exported)", emb_vllm, emb_hf_exported), | ||
| ("vLLM (exported)", "HF (reference)", emb_vllm, emb_hf_reference), | ||
| ("HF (exported)", "HF (reference)", emb_hf_exported, emb_hf_reference), | ||
| ] | ||
|
|
||
| # ---- Summary table ---- | ||
| header = f"{'Pair':<35} {'max |diff|':>14} {'mean |diff|':>14} {'cos sim':>12} {'exact':>7}" | ||
| sep = "-" * len(header) | ||
| print(f"\n{sep}") | ||
| print(header) | ||
| print(sep) | ||
|
|
||
| for name_a, name_b, a, b in pairs: | ||
| diffs = np.abs(a.astype(np.float64) - b.astype(np.float64)) | ||
| label = f"{name_a} vs {name_b}" | ||
| exact = np.array_equal(a, b) | ||
| print( | ||
| f"{label:<35} {diffs.max():>14.8e} {diffs.mean():>14.8e} " | ||
| f"{cosine_sim(a, b):>12.10f} {'YES' if exact else 'NO':>7}" | ||
| ) | ||
|
|
||
| print(sep) | ||
| print(f"Tolerance: rtol={RTOL}, atol={ATOL} (0 = exact match required)") | ||
|
|
||
| # Per-sequence breakdown | ||
| short = {"vLLM (exported)": "vllm", "HF (exported)": "hf_exp", "HF (reference)": "hf_ref"} | ||
| print("\nPer-sequence max |diff|:") | ||
| for i in range(len(SEQUENCES)): | ||
| row = f" seq {i}:" | ||
| for name_a, name_b, a, b in pairs: | ||
| d = float(np.abs(a[i].astype(np.float64) - b[i].astype(np.float64)).max()) | ||
| row += f" {short[name_a]}_vs_{short[name_b]}={d:.8e}" | ||
| print(row) | ||
|
|
||
| print(sep) | ||
|
|
||
| # Cleanup | ||
| if torch.distributed.is_initialized(): | ||
| torch.distributed.destroy_process_group() |
Contributor
There was a problem hiding this comment.
Test script has no assertions — it will never fail.
This "test" prints a comparison table but never asserts that results are within tolerance. RTOL and ATOL (line 74) are defined but unused. A silent pass regardless of output drift defeats the purpose of a golden-value test.
Add assertions after the comparison loop, for example:
Suggested assertion block
+ all_passed = True
for name_a, name_b, a, b in pairs:
diffs = np.abs(a.astype(np.float64) - b.astype(np.float64))
label = f"{name_a} vs {name_b}"
exact = np.array_equal(a, b)
print(
f"{label:<35} {diffs.max():>14.8e} {diffs.mean():>14.8e} "
f"{cosine_sim(a, b):>12.10f} {'YES' if exact else 'NO':>7}"
)
+ if not np.allclose(a, b, rtol=RTOL, atol=ATOL):
+ all_passed = False
print(sep)
- print(f"Tolerance: rtol={RTOL}, atol={ATOL} (0 = exact match required)")
+ print(f"Tolerance: rtol={RTOL}, atol={ATOL}")
+
+ assert all_passed, "Golden value comparison failed — see table above for details."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@bionemo-recipes/vllm/test_esm2_golden_values.py` around lines 147 - 210, The
script currently only prints comparisons and never fails; add assertions that
enforce the golden-value tolerances using RTOL and ATOL: for each pair in pairs
(refer to the pairs list and variables a, b), assert np.allclose(a, b,
rtol=RTOL, atol=ATOL) (or equivalently assert (np.abs(a-b) <= ATOL + RTOL *
np.abs(b)).all()) and fail the test if not, and also assert cosine_sim(a, b)
exceeds a sensible threshold or that exact is True when ATOL/RTOL are zero; add
per-sequence assertions inside the per-sequence loop to ensure each sequence
max-diff <= ATOL + RTOL * max(|b_i|) so the test fails on unacceptable drift.
pstjohn
reviewed
Feb 23, 2026
Comment on lines
+77
to
+78
| # only needed at runtime for FP8 training efficiency; users who train with FP8 pass | ||
| # padded_vocab_size explicitly. Keeping vocab_size-sized weights in the checkpoint |
Collaborator
There was a problem hiding this comment.
users who train with FP8 pass padded_vocab_size explicitly
i'm not sure that's a great assumption 😆. CC @jomitchellnv
Collaborator
There was a problem hiding this comment.
yea that's not true, users who train with FP8 don't pass this
Collaborator
Author
|
/ok to test 36cdbb2 |
gagank1
added
the
ciflow:all-recipes
Collaborator
Author
|
/ok to test 36cdbb2 |
broland-hat
requested changes
Mar 2, 2026
Collaborator
broland-hat
left a comment
broland-hat
left a comment
There was a problem hiding this comment.
Appears to be changes due to (1) layer name changes, (2) toggle the inclusion of a 'pooling layer'. Let me know when the README is updated with a docker build-and-test workflow that I can use to successfully verify the tests added.
| # Export without vocab padding so the checkpoint stores embeddings at the real | ||
| # vocab_size. This avoids shape-mismatch errors in vLLM's VocabParallelEmbedding, | ||
| # which expects vocab_size-shaped weights. | ||
| model_te = convert_esm_hf_to_te(model_hf_masked_lm, padded_vocab_size=None) |
Collaborator
There was a problem hiding this comment.
[not blocking]: Okay made some changes to convert_esm_hf_to_te
| ## Runtime Requirements | ||
|
|
||
| We recommend using the latest [NVIDIA PyTorch container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) | ||
| for optimal performance and compatibility. Refer to the provided Dockerfile for details. |
Collaborator
There was a problem hiding this comment.
[not blocking]: okay build is fine
build_image_r20260302T163713_bionemo-recipes_models_esm2-br_review_gkaushik_esm2-vllm-20260302T163713-36cdbb27.log
| output = model(**inputs) | ||
| ``` | ||
|
|
||
| ### Running inference with vLLM |
Collaborator
There was a problem hiding this comment.
README.md
- HW requirements.
- In the 'Runtime requirements' section please put the HW requirements at a course level of detail, or link/refer to elsewhere in the project where to find the info.
- On which HW has the code been tested (pytests) so far
- Eg are both amd64/x86_64 and arm64 expected to work.
- which generations of NVIDIA devices are expected to work?
- In the 'Runtime requirements' section please put the HW requirements at a course level of detail, or link/refer to elsewhere in the project where to find the info.
- Build and test, user to verify environment
- Recommend: Immediately after given HW and software prerequistes, give the exact docker build command, and command to run the python pytest/unittest suite. If these are the same as https://github.com/NVIDIA/bionemo-framework/tree/main/bionemo-recipes/models#cicd-integration, you can refer to this link.
-top-level Dockerfile WORKDIR is /workspace/bionemo2, recommend the same for this Dockerfile for consistency
| torchao!=0.14.0 | ||
| transformer_engine[pytorch] | ||
| transformers | ||
| vllm |
Collaborator
There was a problem hiding this comment.
[not blocking]: Better to pin, this dep may update
Collaborator
There was a problem hiding this comment.
Actually, I think it is a recipes design choice to avoid pinning unless absolutely necessary to maximize compatibility when installing into external environments. @pstjohn thoughts?
Collaborator
There was a problem hiding this comment.
yes, please don't pin. When we get uv working, we can have a uv.lock file with a last-known good configuration and let a dependabot script update them. But we want to know if a vllm update breaks our stuff -- we'll pin (temporarily), fix it here or upstream, and then un-pin
Signed-off-by: Gagan Kaushik <gkaushik@nvidia.com>
Signed-off-by: Gagan Kaushik <gkaushik@nvidia.com>
Collaborator
Author
|
/ok to test 6779320 |
Signed-off-by: Gagan Kaushik <gkaushik@nvidia.com>
gagank1
removed
the
ciflow:all-recipes
Collaborator
Author
|
/ok to test 6280742 |
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
Signed-off-by: Gagan Kaushik <gkaushik@nvidia.com>
Signed-off-by: Gagan Kaushik <gkaushik@nvidia.com>
Collaborator
Author
|
/ok to test c31c3e8 |
Signed-off-by: Gagan Kaushik <gkaushik@nvidia.com>
Collaborator
Author
|
/ok to test ef612a3 |
pstjohn
reviewed
Mar 9, 2026
Collaborator
There was a problem hiding this comment.
wait, what? why doesn't this work by just adding vllm to requirements.txt?
Collaborator
Author
There was a problem hiding this comment.
To link it to an existing pytorch install (ie the one in our base image) you have to build it: https://docs.vllm.ai/en/latest/getting_started/installation/gpu/#use-an-existing-pytorch-installation
Collaborator
Author
There was a problem hiding this comment.
Btw that script is optional, the way i thought about it was someone might want to quickly build the image for training without waiting for vllm to install and then maybe decide to install it in which case they can use that script inside the container. or they can use the build args with the Dockerfile to install it on first build
Signed-off-by: Gagan Kaushik <gkaushik@nvidia.com>
gagank1
requested review from
polinabinder1,
skothenhill-nv and
yzhang123
as code owners
Collaborator
Author
|
/ok to test 60043cd |
Signed-off-by: Gagan Kaushik <gkaushik@nvidia.com>
Collaborator
Author
|
/ok to test c2d2351 |
❌ 1 Tests Failed:
View the top 1 failed test(s) by shortest run time
To view more test analytics, go to the Test Analytics Dashboard |
pstjohn
reviewed
Mar 10, 2026
Collaborator
There was a problem hiding this comment.
can we flatten this directory? e.g. no esm2/ folder. we can put modeling_esm_te.py and modeling_llama_te.py right next to each other.
Collaborator
Author
There was a problem hiding this comment.
we could but i have a number of files for llama3 as well so it might get pretty messy: https://github.com/NVIDIA/bionemo-framework/pull/1501/changes
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR makes the ESM2 model compatible with vLLM. Primary issues were a naming incompatibility (vLLM expects
model.prefix and ESM2 usesesm.) andNVEsmModeldefaults toadd_pooling_layer=Truewhen loading the checkpoint even though it's exported without pooler weights.Usage
python test_esm2_golden_values.pyfrom inside the container, instructions to build and run it are provided.Type of changes
Triggering Code Rabbit AI Review
To trigger a code review from code rabbit, comment on a pull request with one of these commands:
See https://docs.coderabbit.ai/reference/review-commands for a full list of commands.
Pre-submit Checklist