Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docker-bake.hcl
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ variable "BASE_TAG_PYTHON" {

# Pin for nmp-automodel-base.
variable "BASE_TAG_AUTOMODEL" {
default = "65632527b258367fb8855ef431f596abb8538577"
default = "907446243676db5406e2a0421d1ac060ad1f1e14"
Comment thread
soluwalana marked this conversation as resolved.
}

# The tag for base images if needed
Expand Down
31 changes: 30 additions & 1 deletion docker/Dockerfile.nmp-automodel-base
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
"pillow>=12.2.0,<13" \
"pyasn1>=0.6.3,<1" \
"tornado>=6.5.5,<7" \
"urllib3>=2.7.0,<3"
"urllib3>=2.7.0,<3" \
"mlflow-skinny>=3.11.1,<3.12.0" \
"grpcio>=1.81.1,<2"

# Published base image (same filesystem as builder).
FROM ${PYTORCH_BASE} AS nmp-automodel-base
Expand All @@ -113,6 +115,33 @@ COPY --from=nmp-automodel-base-builder /opt/Automodel /opt/Automodel
# Builder pins uv 0.9.14 but does not ship it in the venv layer; PyTorch base may ship 0.10.x.
COPY --from=ghcr.io/astral-sh/uv:0.9.14 /uv /bin/uv

# Stale copies under NGC system site-packages (verified on nvcr.io/nvidia/pytorch:26.05-py3).
# CVE scanners read dist-packages even when --system-site-packages resolves imports from /opt/venv.
# rm -rf (not pip uninstall) avoids corrupted dist-info RECORDs on the NGC base image.
RUN rm -rf \
/usr/local/lib/python3.12/dist-packages/aiohttp \
/usr/local/lib/python3.12/dist-packages/aiohttp-*.dist-info \
/usr/local/lib/python3.12/dist-packages/black \
/usr/local/lib/python3.12/dist-packages/_black* \
/usr/local/lib/python3.12/dist-packages/black-*.dist-info \
/usr/local/lib/python3.12/dist-packages/blackd \
/usr/local/lib/python3.12/dist-packages/tornado \
/usr/local/lib/python3.12/dist-packages/tornado-*.dist-info \
/usr/local/lib/python3.12/dist-packages/onnx \
/usr/local/lib/python3.12/dist-packages/onnx-*.dist-info \
/usr/local/lib/python3.12/dist-packages/pyasn1 \
/usr/local/lib/python3.12/dist-packages/pyasn1-*.dist-info \
/usr/local/lib/python3.12/dist-packages/urllib3 \
/usr/local/lib/python3.12/dist-packages/urllib3-*.dist-info \
/usr/local/lib/python3.12/dist-packages/grpc \
/usr/local/lib/python3.12/dist-packages/grpcio \
/usr/local/lib/python3.12/dist-packages/grpcio-*.dist-info \
Comment thread
coderabbitai[bot] marked this conversation as resolved.
/usr/local/lib/python3.12/dist-packages/wandb \
/usr/local/lib/python3.12/dist-packages/wandb-*.dist-info \
/usr/local/lib/python3.12/dist-packages/mlflow \
/usr/local/lib/python3.12/dist-packages/mlflow-*.dist-info \
/usr/local/lib/python3.12/dist-packages/mlflow_skinny-*.dist-info

ENV VIRTUAL_ENV=/opt/venv \
UV_PROJECT_ENVIRONMENT=/opt/venv \
HF_HUB_ENABLE_HF_TRANSFER=1 \
Expand Down
70 changes: 46 additions & 24 deletions docker/Dockerfile.nmp-unsloth-training
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,11 @@
# bitsandbytes, xformers, etc.). `--overrides preserve_base_torch.txt`
# blocks uv from installing/upgrading torch into the venv so the NGC
# base's PyTorch + CUDA remain the runtime stack.
# 1b. bitsandbytes — compiled from source against NGC CUDA 13.1 (same pattern
# as docker/Dockerfile.nmp-automodel-base). PyPI wheels
# only ship through cuda130; source build replaces the wheel from step 1.
# 1c. mamba-ssm + causal-conv1d — installed from the prebuilt cu13.1.1 / cp312
# wheels shared with docker/Dockerfile.nmp-automodel-base (the
# causal-conv1d-wheel / mamba-ssm-wheel bake contexts). Required by hybrid
# Mamba/SSM models (e.g. NVIDIA Nemotron-H *-A3B) whose remote code imports
# mamba_ssm at load time.
# 1b. bitsandbytes — compiled from source against NGC CUDA 13.2 (same pattern
# as docker/Dockerfile.nmp-automodel-base). PyPI wheels only ship through
# cuda130; source build replaces the wheel from step 1.
# 1c. mamba-ssm + causal-conv1d — prebuilt cu13.2 / cp312 wheels (shared with
# docker/Dockerfile.nmp-automodel-base).
# 1d. flash-attn — optional for unsloth and currently NOT installed (see the
# commented TODO below). Without it Unsloth falls back when xformers is
# also missing (common on newer CUDA stacks), logging "FA2 = False /
Expand All @@ -30,9 +27,10 @@
# Publish target: nmp-unsloth-training
# Default tag: `local` (override via BAKE_TAG at build time).

# NGC PyTorch base. 26.02-py3 ships PyTorch 2.11 + CUDA 13.1 + Python 3.12.
# Override at build time: --set nmp-unsloth-training.args.PYTORCH_BASE=...
ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.02-py3
# NGC PyTorch base. 26.05-py3 ships PyTorch 2.12 + CUDA 13.2 + Python 3.12 (aligned with
# docker/Dockerfile.nmp-automodel-base). Override at build time:
# --set nmp-unsloth-training.args.PYTORCH_BASE=...
ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.05-py3

# Prebuilt CUDA-extension wheels (mamba-ssm + causal-conv1d), shared with
# docker/Dockerfile.nmp-automodel-base. The bake `nmp-unsloth-training` target
Expand Down Expand Up @@ -74,7 +72,7 @@ ARG USER_GID=1000
ARG UNSLOTH_VERSION=2026.6.1
ARG TRANSFORMERS_VERSION=4.57.6
ARG HF_HUB_VERSION=0.36.2
ARG BITSANDBYTES_VERSION=0.49.1
ARG BITSANDBYTES_VERSION=0.49.2
ARG BNB_MAX_JOBS=10

WORKDIR /app
Expand Down Expand Up @@ -106,7 +104,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
transformers==${TRANSFORMERS_VERSION} \
huggingface-hub==${HF_HUB_VERSION}

# Step 1b: bitsandbytes from source — matches automodel base (CUDA 13.1 nvcc).
# Step 1b: bitsandbytes from source — matches automodel base (CUDA 13.2 nvcc).
RUN --mount=type=cache,target=/root/.cache/uv \
git clone https://github.com/bitsandbytes-foundation/bitsandbytes.git /tmp/bitsandbytes && \
cd /tmp/bitsandbytes && \
Expand All @@ -118,20 +116,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
uv pip uninstall --python ${VIRTUAL_ENV}/bin/python scikit-build-core && \
rm -rf /tmp/bitsandbytes

# Step 1c: mamba-ssm + causal-conv1d — required by hybrid Mamba/SSM models such
# as NVIDIA Nemotron-H (e.g. *-Nano-*-A3B), whose remote modeling code imports
# mamba_ssm at load time. Installed from the SAME prebuilt cu13.1.1 / cp312
# wheels as docker/Dockerfile.nmp-automodel-base — built once by the
# causal-conv1d-wheel / mamba-ssm-wheel bake targets (docker/base/Dockerfile.mamba-wheel)
# and shared across both images, so the versions are single-sourced and nothing
# recompiles here. The wheels are built from this same NGC 26.02 base (torch 2.11
# + CUDA 13.1), so they're ABI-matched. --no-deps so the HF stack isn't
# re-resolved (triton is already provided by the NGC base).
# Step 1c: mamba-ssm + causal-conv1d — cu13.2 / cp312 wheels shared with automodel base.
RUN --mount=from=causal-conv1d-wheel-src,target=/tmp/causal-conv1d-wheel-src,readonly \
--mount=from=mamba-ssm-wheel-src,target=/tmp/mamba-ssm-wheel-src,readonly \
uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache --no-deps \
/tmp/causal-conv1d-wheel-src/wheels/cu13.1.1/causal_conv1d-*cp312*.whl \
/tmp/mamba-ssm-wheel-src/wheels/cu13.1.1/mamba_ssm-2.3.0-cp312*.whl
/tmp/causal-conv1d-wheel-src/wheels/cu13.2/causal_conv1d-*cp312*.whl \
/tmp/mamba-ssm-wheel-src/wheels/cu13.2/mamba_ssm-2.3.0-cp312*.whl

# TODO: Step 1d: Flash Attention 2 — compiled from source against the NGC 26.02 torch.
# /usr/local/cuda symlinks to an older toolkit; use /usr/local/cuda-13.1 instead.
Expand Down Expand Up @@ -179,6 +169,38 @@ RUN --mount=type=cache,target=/root/.cache/uv \
huggingface-hub==${HF_HUB_VERSION} \
transformers==${TRANSFORMERS_VERSION}

# Pin nspect-tracked deps into /opt/venv before removing stale NGC system copies.
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache \
--overrides /opt/docker/preserve_base_torch.txt \
"aiohttp>=3.13.4,<4" \
"urllib3>=2.7.0,<3" \
"grpcio>=1.81.1,<2" \
"mlflow-skinny>=3.11.1,<3.12.0"

# Stale NGC system site-packages (verified on nvcr.io/nvidia/pytorch:26.05-py3).
RUN rm -rf \
/usr/local/lib/python3.12/dist-packages/aiohttp \
/usr/local/lib/python3.12/dist-packages/aiohttp-*.dist-info \
/usr/local/lib/python3.12/dist-packages/black \
/usr/local/lib/python3.12/dist-packages/_black* \
/usr/local/lib/python3.12/dist-packages/black-*.dist-info \
/usr/local/lib/python3.12/dist-packages/blackd \
/usr/local/lib/python3.12/dist-packages/tornado \
/usr/local/lib/python3.12/dist-packages/tornado-*.dist-info \
/usr/local/lib/python3.12/dist-packages/onnx \
/usr/local/lib/python3.12/dist-packages/onnx-*.dist-info \
/usr/local/lib/python3.12/dist-packages/urllib3 \
/usr/local/lib/python3.12/dist-packages/urllib3-*.dist-info \
/usr/local/lib/python3.12/dist-packages/grpc \
/usr/local/lib/python3.12/dist-packages/grpcio \
/usr/local/lib/python3.12/dist-packages/grpcio-*.dist-info \
/usr/local/lib/python3.12/dist-packages/wandb \
/usr/local/lib/python3.12/dist-packages/wandb-*.dist-info \
/usr/local/lib/python3.12/dist-packages/mlflow \
/usr/local/lib/python3.12/dist-packages/mlflow-*.dist-info \
/usr/local/lib/python3.12/dist-packages/mlflow_skinny-*.dist-info

ENTRYPOINT ["/opt/venv/bin/python"]
CMD ["-m", "nmp.unsloth.tasks.training", "--help"]

Expand Down
70 changes: 35 additions & 35 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -85,12 +85,12 @@ dev = [
"ty==0.0.17",
"hatchling>=1.26.3",
"nmp-testing",
"mlflow-skinny", # Optional dep for fine-tuning (shipped in training images); skinny variant to test API surface without bloating uv.lock
"nemo-platform-plugin", # Optional runtime dep; included here for type checking
"nvidia-nat-atif>=1.7.0,<1.8", # ATIF schema models for agentic-use trajectory validation
"nemo-platform-sdk-tools", # Include SDK/license tool migration target for testing
"nmp-build-tools", # Build hook helpers imported by package-local hatch_build.py files
"nmp-dev-mcp", # Include MCP dev tools
"mlflow-skinny>=3.11.1,<3.12.0", # Optional dep for fine-tuning (shipped in training images); skinny variant to test API surface without bloating uv.lock
"nemo-platform-plugin", # Optional runtime dep; included here for type checking
"nvidia-nat-atif>=1.7.0,<1.8", # ATIF schema models for agentic-use trajectory validation
"nemo-platform-sdk-tools", # Include SDK/license tool migration target for testing
"nmp-build-tools", # Build hook helpers imported by package-local hatch_build.py files
"nmp-dev-mcp", # Include MCP dev tools
# These are the dev dependencies from `sdk/python/nemo-platform`.
# Including here until we find a better way to include them.
"pyright==1.1.399",
Expand Down Expand Up @@ -258,21 +258,21 @@ constraint-dependencies = [
"GitPython>=3.1.49",
"Mako>=1.3.12",
"Pygments>=2.20.0",
"aiohttp>=3.13.4", # High/Medium/Low OOS – auditor-tasks + customizer
"aiohttp>=3.13.4", # High/Medium/Low OOS – auditor-tasks + customizer
"authlib>=1.6.11",
"black>=26.3.1",
"cryptography>=46.0.7",
"diffusers>=0.38.0",
"filelock>=3.20.3",
"jupyter-server>=2.18.0",
"jupyterlab>=4.5.7",
"langchain-community>=0.3.31,<0.4", # 0.4.x removes the vertexai import path used by ragas 0.4.3
"langchain-community>=0.3.31,<0.4", # 0.4.x removes the vertexai import path used by ragas 0.4.3
"langchain-core>=1.4.8",
"langchain-nvidia-ai-endpoints>=1.4.1",
"langchain-openai>=1.3.3",
"langgraph>=1.2.6",
"langsmith==0.8.2", # osv-scanner UNKNOWN for 0.8.9+; PyPI metadata is MIT
"json-repair==0.58.7", # osv-scanner UNKNOWN for 0.60.1+; PyPI metadata is MIT
"langsmith==0.8.2", # osv-scanner UNKNOWN for 0.8.9+; PyPI metadata is MIT
"json-repair==0.58.7", # osv-scanner UNKNOWN for 0.60.1+; PyPI metadata is MIT
"litellm>=1.83.10",
"lxml>=6.1.0",
"mistune>=3.2.1",
Expand All @@ -283,7 +283,7 @@ constraint-dependencies = [
"pyasn1>=0.6.3",
"python-multipart>=0.0.27",
"regex>=2025.10.22",
"safetensors>=0.8.0rc0", # explicit prerelease so uv allows the rc
"safetensors>=0.8.0rc0", # explicit prerelease so uv allows the rc
"sqlfluff==4.1.0",
"starlette>=1.3.1",
"urllib3>=2.7.0",
Expand All @@ -296,37 +296,37 @@ override-dependencies = [
"jsonpath-ng>=1.6.1",
"grpcio>=1.71.0",
"huggingface-hub>=1.0.1,<2.0.0",
"pydantic[email]>=2.9.2", # openai-harmony depends on a later pydantic, but that causes the stainless gen to hit an infinite loop
"ray; sys_platform == 'never'", # ray has an unfixed critical CVE, as for why we're removing it this way see https://github.com/astral-sh/uv/issues/9174
"fschat; sys_platform == 'never'", # fschat is unmaintained with unfixed High/Medium CVEs; transitive via garak but never imported
"pydantic[email]>=2.9.2", # openai-harmony depends on a later pydantic, but that causes the stainless gen to hit an infinite loop
"ray; sys_platform == 'never'", # ray has an unfixed critical CVE, as for why we're removing it this way see https://github.com/astral-sh/uv/issues/9174
"fschat; sys_platform == 'never'", # fschat is unmaintained with unfixed High/Medium CVEs; transitive via garak but never imported
"diskcache; sys_platform == 'never'", # unfixed CVE-2025-69872; transitive via ragas but not imported
"botocore>=1.40.46,<1.40.62", # ngcsdk requires botocore>=1.37; upper bound required by aiobotocore 2.25.1 (aioboto3 15.5.0)
"langchain>=1.3.11", # LangChain CVEs through CVE-2026-45134 and CVE-2025-68664
"langchain-core>=1.4.8", # LangChain Core CVEs through GHSA-qh6h-p6c9-ff54 and CVE-2026-44843
"nltk>=3.9.4", # NLTK CVEs with available patches; remaining NLTK alerts have no upstream fix
"Pillow>=12.2.0", # Pillow CVEs through CVE-2026-48624; override: fastembed (via nemoguardrails) pins Pillow<12
"botocore>=1.40.46,<1.40.62", # ngcsdk requires botocore>=1.37; upper bound required by aiobotocore 2.25.1 (aioboto3 15.5.0)
"langchain>=1.3.11", # LangChain CVEs through CVE-2026-45134 and CVE-2025-68664
"langchain-core>=1.4.8", # LangChain Core CVEs through GHSA-qh6h-p6c9-ff54 and CVE-2026-44843
"nltk>=3.9.4", # NLTK CVEs with available patches; remaining NLTK alerts have no upstream fix
"Pillow>=12.2.0", # Pillow CVEs through CVE-2026-48624; override: fastembed (via nemoguardrails) pins Pillow<12
# todo(dn-v2): unpin these. this was originally causing tests to fail.
"aiodns==3.5.0",
"pycares==4.11.0",
"authlib>=1.6.9", # Critical – nmp-api/core/cpu-tasks/gpu-tasks
"aiohttp>=3.13.4", # High/Medium/Low OOS – auditor-tasks + customizer
"jaraco-context>=6.1.0", # High OOS – auditor-tasks + customizer
"pyasn1>=0.6.3", # High OOS – auditor-tasks + customizer
"setuptools>=78.1.1", # High OOS – nmp-cpu-tasks/gpu-tasks/customizer-tasks
"urllib3>=2.7.0", # High OOS – auditor-tasks + customizer
"wheel>=0.46.2", # High OOS – widespread
"cbor2>=5.9.0", # High – customizer + nmp-gpu-tasks
"simpleeval>=1.0.5", # High – nmp-api + nmp-gpu-tasks
"ujson>=5.12.0", # High – nmp-api/cpu-tasks/gpu-tasks
"xgrammar>=0.1.32", # High – customizer + nmp-gpu-tasks
"fastmcp>=3.2.0", # GHSA-vv7q-7jx5-f767 (Critical) + GHSA-rww4-4w9c-7733 (High); overrides vendored sdk/python/nemo-platform <3 constraint
"wandb>=0.25.1", # CVE-2026-33186
"click>=8.2.0", # Below CVEs are based on constraints that garak has
"authlib>=1.6.9", # Critical – nmp-api/core/cpu-tasks/gpu-tasks
"aiohttp>=3.13.4", # High/Medium/Low OOS – auditor-tasks + customizer
"jaraco-context>=6.1.0", # High OOS – auditor-tasks + customizer
"pyasn1>=0.6.3", # High OOS – auditor-tasks + customizer
"setuptools>=78.1.1", # High OOS – nmp-cpu-tasks/gpu-tasks/customizer-tasks
"urllib3>=2.7.0", # High OOS – auditor-tasks + customizer
"wheel>=0.46.2", # High OOS – widespread
"cbor2>=5.9.0", # High – customizer + nmp-gpu-tasks
"simpleeval>=1.0.5", # High – nmp-api + nmp-gpu-tasks
"ujson>=5.12.0", # High – nmp-api/cpu-tasks/gpu-tasks
"xgrammar>=0.1.32", # High – customizer + nmp-gpu-tasks
"fastmcp>=3.2.0", # GHSA-vv7q-7jx5-f767 (Critical) + GHSA-rww4-4w9c-7733 (High); overrides vendored sdk/python/nemo-platform <3 constraint
"wandb>=0.25.1", # CVE-2026-33186
"click>=8.2.0", # Below CVEs are based on constraints that garak has
"langchain-openai>=1.3.3",
"litellm>=1.83.10", # CVE-2026-42208 (Critical SQL injection), CVE-2026-40217, and 5 other High CVEs
"litellm>=1.83.10", # CVE-2026-42208 (Critical SQL injection), CVE-2026-40217, and 5 other High CVEs
"python-dotenv>=1.2.2",
"openai>=2.26.0",
"sqlfluff>=4.1.0", # CVE-2026-46373
"sqlfluff>=4.1.0", # CVE-2026-46373

]

Expand Down
Loading
Loading