From 6d5a20c258531e86a440786d4731cc81ac3dd16e Mon Sep 17 00:00:00 2001 From: Erick Galinkin Date: Tue, 18 Feb 2025 16:56:21 -0500 Subject: [PATCH] fix: set workdir to models and specify entrypoint explicitly. * Fix bug with git to fetch models, use wget instead * Fix bug where workdir is kept as /models. Specify entrypoint explicitly. --- .../library/jailbreak_detection/Dockerfile | 12 +++++------ .../jailbreak_detection/Dockerfile-GPU | 21 ++++++++----------- 2 files changed, 15 insertions(+), 18 deletions(-) diff --git a/nemoguardrails/library/jailbreak_detection/Dockerfile b/nemoguardrails/library/jailbreak_detection/Dockerfile index 5cf76dee4..209baafc0 100644 --- a/nemoguardrails/library/jailbreak_detection/Dockerfile +++ b/nemoguardrails/library/jailbreak_detection/Dockerfile @@ -4,6 +4,11 @@ FROM python:3.10-slim # Install git and wget RUN apt-get update && apt-get install -y git gcc g++ python3-dev wget && apt-get clean +# Predownload embedding-based jailbreak detection models, set environment variable for path +WORKDIR /models +RUN wget https://huggingface.co/nvidia/NemoGuard-JailbreakDetect/resolve/main/snowflake.pkl +ENV EMBEDDING_CLASSIFIER_PATH=/models + # Set working directory WORKDIR /app @@ -21,14 +26,9 @@ ENV JAILBREAK_CHECK_DEVICE=cpu # Predownload the GPT2 model. RUN python -c "from transformers import GPT2LMHeadModel, GPT2TokenizerFast; GPT2LMHeadModel.from_pretrained('gpt2-large'); GPT2TokenizerFast.from_pretrained('gpt2-large');" -# Predownload embedding-based jailbreak detection models, set environment variable for path -WORKDIR /models -RUN wget https://huggingface.co/nvidia/NemoGuard-JailbreakDetect/resolve/main/snowflake.pkl -ENV EMBEDDING_CLASSIFIER_PATH=/models - # Expose a port for the server EXPOSE 1337 # Start the server as the default command -ENTRYPOINT ["/usr/local/bin/python", "server.py"] +ENTRYPOINT ["/usr/local/bin/python", "/app/server.py"] CMD ["--port=1337"] diff --git a/nemoguardrails/library/jailbreak_detection/Dockerfile-GPU b/nemoguardrails/library/jailbreak_detection/Dockerfile-GPU index 418e4026a..5180617d8 100644 --- a/nemoguardrails/library/jailbreak_detection/Dockerfile-GPU +++ b/nemoguardrails/library/jailbreak_detection/Dockerfile-GPU @@ -10,6 +10,14 @@ RUN apt-get install python3 python3-pip python3-dev -y # Upgrade pip RUN pip install --no-cache-dir --upgrade pip +# Set the device on which the model should load e.g., "cpu", "cuda:0", etc. +ENV JAILBREAK_CHECK_DEVICE=cuda:0 + +# Predownload embedding-based jailbreak detection models, set environment variable for path +WORKDIR /models +RUN wget https://huggingface.co/nvidia/NemoGuard-JailbreakDetect/resolve/main/snowflake.pkl +ENV EMBEDDING_CLASSIFIER_PATH=/models + # Set working directory WORKDIR /app @@ -24,20 +32,9 @@ COPY . . # Predownload the GPT2 model. RUN python3 -c "from transformers import GPT2LMHeadModel, GPT2TokenizerFast; GPT2LMHeadModel.from_pretrained('gpt2-large'); GPT2TokenizerFast.from_pretrained('gpt2-large');" -# Set the device on which the model should load e.g., "cpu", "cuda:0", etc. -ENV JAILBREAK_CHECK_DEVICE=cuda:0 - -# Predownload embedding-based jailbreak detection models, set environment variable for path -WORKDIR /models -RUN wget https://huggingface.co/nvidia/NemoGuard-JailbreakDetect/resolve/main/snowflake.pkl -ENV EMBEDDING_CLASSIFIER_PATH=/models - -# To use nvidia/nv-embedqa-e5-v5 model, uncomment the line below and set your API key -# ENV NVIDIA_API_KEY= - # Expose a port for the server EXPOSE 1337 # Start the server as the default command -ENTRYPOINT ["/usr/local/bin/python", "server.py"] +ENTRYPOINT ["/usr/local/bin/python", "/app/server.py"] CMD ["--port=1337"]