Multistage build with baked-in models needed for the enabled guardrails

Dockerfile.server

@@ -1,44 +1,78 @@
-FROM registry.access.redhat.com/ubi9/python-312 as build
+FROM registry.access.redhat.com/ubi9/python-312:latest as build
 
 USER 0
 WORKDIR /app
 
-RUN dnf install -y gcc gcc-c++ git && \
-    pip install --no-cache-dir poetry==1.8.2 pyyaml==6.0.2 && \
+ARG POETRY_VERSION=2.2.1
+ARG PYYAML_VERSION=6.0.3
+ARG GUARDRAILS_PROFILE=opensource
+
+RUN dnf install -y --nodocs --setopt=install_weak_deps=False \
+        gcc \
+        gcc-c++ \
+        git && \
+    pip install --no-cache-dir \
+        poetry==${POETRY_VERSION} \
+        pyyaml==${PYYAML_VERSION} && \
     dnf clean all && \
-    rm -rf /var/cache/dnf
+    rm -rf /var/cache/dnf /tmp/* /var/tmp/*
 
-COPY pyproject.toml poetry.lock* README.md ./
+COPY pyproject.toml poetry.lock README.md ./
 COPY nemoguardrails/ ./nemoguardrails/
-COPY examples/ ./examples/
-COPY chat-ui/ ./chat-ui/
-COPY scripts/provider-list.yaml ./scripts/
-COPY scripts/filter_guardrails.py ./scripts/
-COPY scripts/entrypoint.sh ./scripts/
-RUN chmod +x ./scripts/entrypoint.sh
+COPY scripts/ ./scripts/
+COPY examples/bots/ ./examples/bots/
 
-ARG GUARDRAILS_PROFILE=opensource
+RUN chmod +x ./scripts/entrypoint.sh
 RUN python3 ./scripts/filter_guardrails.py ./scripts/provider-list.yaml $GUARDRAILS_PROFILE
 
 ENV POETRY_VIRTUALENVS_IN_PROJECT=1 \
-    POETRY_NO_INTERACTION=1
+    POETRY_NO_INTERACTION=1 \
+    POETRY_VIRTUALENVS_CREATE=true \
+    POETRY_NO_CACHE=1 \
+    HF_HOME=/app/.cache/huggingface \
+    TRANSFORMERS_CACHE=/app/.cache/transformers \
+    SENTENCE_TRANSFORMERS_HOME=/app/.cache/sentence_transformers \
+    NLTK_DATA=/app/.cache/nltk_data \
+    FASTEMBED_CACHE_PATH=/app/.cache/fastembed
 
-RUN poetry install --no-ansi --extras="sdd jailbreak openai nvidia tracing" && \
-    poetry run pip install "spacy>=3.4.4,<4.0.0" && \
-    poetry run python -m spacy download en_core_web_lg
+RUN poetry install --no-ansi --only main --extras="sdd jailbreak openai nvidia tracing models" && \
+    poetry cache clear --all pypi && \
+    rm -rf /root/.cache/pip /tmp/* /var/tmp/* /opt/app-root/src/.cache/pypoetry
 
-FROM registry.access.redhat.com/ubi9/python-312
+RUN set -ex && \
+    GUARDRAILS_PROFILE=$GUARDRAILS_PROFILE poetry run python ./scripts/pre_download_required_models.py && \
+    find /app/.cache -type f -name "*.pyc" -delete && \
+    find /app/.cache -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true && \
+    find /app/.cache -type d -name ".git" -exec rm -rf {} + 2>/dev/null || true && \
+    poetry cache clear --all pypi && \
+    rm -rf /root/.cache/pip /tmp/* /var/tmp/*
+
+RUN dnf remove -y gcc gcc-c++ git && \
+    dnf clean all && \
+    rm -rf /var/cache/dnf
+
+FROM registry.access.redhat.com/ubi9/python-312:latest as runtime
 
 USER 0
 WORKDIR /app
 
 COPY --from=build /app /app
+
 RUN rm -f /etc/security/namespace.conf /usr/lib64/security/pam_namespace.so || true && \
     chgrp -R 0 /app && \
-    chmod -R g+rwX /app
+    chmod -R g=u /app && \
+    chmod +x /app/scripts/entrypoint.sh
 
 USER 1001
 
-ENV PATH="/app/.venv/bin:$PATH"
+ENV HF_HOME=/app/.cache/huggingface \
+    TRANSFORMERS_CACHE=/app/.cache/transformers \
+    SENTENCE_TRANSFORMERS_HOME=/app/.cache/sentence_transformers \
+    NLTK_DATA=/app/.cache/nltk_data \
+    FASTEMBED_CACHE_PATH=/app/.cache/fastembed \
+    PATH="/app/.venv/bin:$PATH" \
+    PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1
+
 EXPOSE 8000
-ENTRYPOINT ["./scripts/entrypoint.sh"]
+ENTRYPOINT ["./scripts/entrypoint.sh"]

pyproject.toml

@@ -105,6 +105,13 @@ google-cloud-language = { version = ">=2.14.0", optional = true }
 # jailbreak injection
 yara-python = { version = "^4.5.1", optional = true }
 
+# models - optional dependencies for various embedding and NLP models
+torch = { version = ">=2.0.0", optional = true, source = "pytorch-cpu" }
+spacy = { version = ">=3.4.4,<4.0.0,!=3.7.0", python = "<3.14", optional = true }
+nltk = { version = ">=3.8", optional = true }
+sentence-transformers = { version = ">=2.2.0", optional = true }
+einops = { version = ">=0.6.0", optional = true }
+
 [tool.poetry.extras]
 sdd = ["presidio-analyzer", "presidio-anonymizer"]
 eval = ["tqdm", "numpy", "streamlit", "tornado"]
@@ -113,6 +120,7 @@ gcp = ["google-cloud-language"]
 tracing = ["opentelemetry-api", "aiofiles"]
 nvidia = ["langchain-nvidia-ai-endpoints"]
 jailbreak = ["yara-python"]
+models = ["torch", "spacy", "nltk", "sentence-transformers", "einops"]
 # Poetry does not support recursive dependencies, so we need to add all the dependencies here.
 # I also support their decision. There is no PEP for recursive dependencies, but it has been supported in pip since version 21.2.
 # It is here for backward compatibility.
@@ -128,6 +136,11 @@ all = [
   "aiofiles",
   "langchain-nvidia-ai-endpoints",
   "yara-python",
+  "torch",
+  "spacy",
+  "nltk",
+  "sentence-transformers",
+  "einops",
 ]
 
 [tool.poetry.group.dev]
@@ -193,6 +206,11 @@ log-level = "DEBUG"
 # phase, which will cause tests to fail or "magically" ignored.
 log_cli = "False"
 
+[[tool.poetry.source]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+priority = "explicit"
+
 [build-system]
 requires = ["poetry-core>=1.0.0,<2.0.0"]
 build-backend = "poetry.core.masonry.api"

scripts/discover_required_models.py

@@ -0,0 +1,207 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Discover all models required by NeMo Guardrails based on guardrails profile.
+"""
+
+import ast
+import logging
+import os
+import re
+import sys
+from pathlib import Path
+from typing import Dict, List, Set
+
+import yaml
+
+logging.basicConfig(level=logging.INFO, format="%(message)s")
+
+
+class ModelDiscoverer:
+    MODEL_KEYS = ("spacy", "sentence_transformers", "huggingface", "nltk", "custom")
+
+    PATTERNS = {
+        "spacy": [
+            r"spacy\.load\(['\"]([^'\"]+)['\"]",
+            r"spacy\.util\.is_package\(['\"]([^'\"]+)['\"]",
+            r"\"model_name\":\s*['\"]([^'\"]+)['\"]",  # For dict configs
+        ],
+        "sentence_transformers": [
+            r"SentenceTransformer\(['\"]([^'\"]+)['\"]",
+            r"sentence-transformers/([a-zA-Z0-9_\-]+)",
+        ],
+        "huggingface": [
+            r"from_pretrained\(['\"]([^'\"]+)['\"]",
+            r"model_name\s*=\s*['\"]([^'\"]+)['\"]",
+        ],
+        "nltk": [
+            r"nltk\.download\(['\"]([^'\"]+)['\"]",
+        ],
+    }
+
+    def __init__(self, profile: str = "opensource"):
+        self.profile = profile
+        self.models: Dict[str, Set[str]] = {k: set() for k in self.MODEL_KEYS}
+
+    def get_active_guardrails(self) -> List[str]:
+        config_path = Path("scripts/provider-list.yaml")
+        if not config_path.exists():
+            logging.error(f"Missing config: {config_path}")
+            sys.exit(1)
+        try:
+            with open(config_path) as f:
+                config = yaml.safe_load(f)
+            profile_config = config["profiles"].get(self.profile, {})
+            closed_source = config.get("closed_source_guardrails", [])
+            include_closed = profile_config.get("include_closed_source", False)
+        except Exception as e:
+            logging.error(f"Error loading provider-list.yaml: {e}")
+            sys.exit(1)
+
+        library_path = Path("nemoguardrails/library")
+        if not library_path.exists():
+            logging.error(f"Missing directory: {library_path}")
+            sys.exit(1)
+
+        available = [
+            item.name
+            for item in library_path.iterdir()
+            if item.is_dir() and not item.name.startswith("_")
+        ]
+        return (
+            available
+            if include_closed
+            else [gr for gr in available if gr not in closed_source]
+        )
+
+    @staticmethod
+    def _extract_from_ast(tree: ast.AST) -> Dict[str, Set[str]]:
+        models = {k: set() for k in ModelDiscoverer.MODEL_KEYS}
+        for node in ast.walk(tree):
+            if (
+                isinstance(node, ast.Call)
+                and getattr(getattr(node.func, "attr", None), "lower", lambda: "")()
+                == "load"
+                and getattr(getattr(node.func, "value", None), "id", None) == "spacy"
+            ):
+                if (
+                    node.args
+                    and isinstance(node.args[0], ast.Constant)
+                    and isinstance(node.args[0].value, str)
+                ):
+                    models["spacy"].add(node.args[0].value)
+            if (
+                isinstance(node, ast.Call)
+                and getattr(node.func, "id", None) == "SentenceTransformer"
+            ):
+                if (
+                    node.args
+                    and isinstance(node.args[0], ast.Constant)
+                    and isinstance(node.args[0].value, str)
+                ):
+                    name = node.args[0].value
+                    if not name.startswith("sentence-transformers/"):
+                        name = f"sentence-transformers/{name}"
+                    models["sentence_transformers"].add(name)
+            if (
+                isinstance(node, ast.Call)
+                and getattr(node.func, "attr", None) == "from_pretrained"
+                and node.args
+                and isinstance(node.args[0], ast.Constant)
+                and isinstance(node.args[0].value, str)
+            ):
+                models["huggingface"].add(node.args[0].value)
+            if (
+                isinstance(node, ast.Call)
+                and getattr(node.func, "attr", None) == "download"
+                and getattr(getattr(node.func, "value", None), "id", None) == "nltk"
+            ):
+                if (
+                    node.args
+                    and isinstance(node.args[0], ast.Constant)
+                    and isinstance(node.args[0].value, str)
+                ):
+                    models["nltk"].add(node.args[0].value)
+        return models
+
+    def extract_models_from_ast(self, file_path: Path) -> Dict[str, Set[str]]:
+        try:
+            with open(file_path, "r", encoding="utf-8") as f:
+                tree = ast.parse(f.read(), filename=str(file_path))
+            return self._extract_from_ast(tree)
+        except Exception as e:
+            logging.warning(f"Error parsing {file_path}: {e}")
+            return {k: set() for k in self.MODEL_KEYS}
+
+    def scan_file(self, file_path: Path):
+        ast_models = self.extract_models_from_ast(file_path)
+        for key in self.models:
+            self.models[key].update(ast_models.get(key, set()))
+
+        try:
+            content = file_path.read_text(encoding="utf-8")
+            for model_type, patterns in self.PATTERNS.items():
+                for pattern in patterns:
+                    for match in re.findall(pattern, content, re.IGNORECASE):
+                        if model_type == "sentence_transformers":
+                            if not match.startswith("sentence-transformers/"):
+                                match = f"sentence-transformers/{match}"
+                        self.models[model_type].add(match)
+        except Exception as e:
+            logging.warning(f"Error scanning {file_path}: {e}")
+
+    def discover(self) -> Dict[str, Set[str]]:
+        for guardrail in self.get_active_guardrails():
+            guardrail_path = Path(f"nemoguardrails/library/{guardrail}")
+            if guardrail_path.exists():
+                for py_file in guardrail_path.rglob("*.py"):
+                    self.scan_file(py_file)
+                if (guardrail_path / "Dockerfile").exists():
+                    self.models["custom"].add(f"{guardrail}_custom_models")
+
+        for py_file in Path("nemoguardrails").rglob("*.py"):
+            if "library" not in str(py_file):
+                self.scan_file(py_file)
+
+        return self.models
+
+    def print_summary(self):
+        active_guardrails = self.get_active_guardrails()
+        print(f"Discovering models for profile: {self.profile}")
+        print(
+            f"Active guardrails ({len(active_guardrails)}): {', '.join(active_guardrails)}"
+        )
+        for category in self.MODEL_KEYS:
+            models = self.models[category]
+            if models:
+                print(f"\n{category.upper()}:")
+                for model in sorted(models):
+                    print(f"  - {model}")
+
+
+def main():
+    profile = os.environ.get("GUARDRAILS_PROFILE", "opensource")
+    if len(sys.argv) > 1:
+        profile = sys.argv[1]
+    discoverer = ModelDiscoverer(profile)
+    discoverer.discover()
+    if "--quiet" not in sys.argv:
+        discoverer.print_summary()
+
+
+if __name__ == "__main__":
+    main()

scripts/entrypoint.sh

@@ -24,4 +24,4 @@ exec /app/.venv/bin/nemoguardrails server \
   --config "/app/config" \
   --port "$PORT" \
   --default-config-id "$CONFIG_ID" \
-  --disable-chat-ui
+  --disable-chat-ui

scripts/filter_guardrails.py

@@ -1,11 +1,27 @@
 #!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
 import os
-import sys
-import yaml
 import shutil
-import logging
+import sys
 from pathlib import Path
 
+import yaml
+
 logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
 logger = logging.getLogger(__name__)