Merge pull request #39 from huggingface/add-pip-install-entrypoint

added entrypoint for pip install if exists

Author: philschmid

dockerfiles/starlette/pytorch/Dockerfile.cpu

@@ -5,7 +5,7 @@ COPY starlette_requirements.txt /tmp/requirements.txt
 RUN pip install  --no-cache-dir -r /tmp/requirements.txt && rm /tmp/requirements.txt
 
 # Think about a better solution -> base contaienr has pt 1.13. thats why need below 0.14
-RUN pip install  --no-cache-dir sentence_transformers torchvision~="0.14.0" diffusers=="0.9.0" accelerate=="0.14.0"
+RUN pip install  --no-cache-dir sentence_transformers=="2.2.2" torchvision=="0.14.0" diffusers=="0.9.0" accelerate=="0.14.0"
 
 # Add upgrade due to issue in base container upgrade https://github.com/mamba-org/mamba/issues/2170
 RUN pip install transformers==4.25.1  --no-cache-dir --upgrade 
@@ -14,7 +14,11 @@ RUN pip install transformers==4.25.1  --no-cache-dir --upgrade
 COPY src/huggingface_inference_toolkit huggingface_inference_toolkit
 COPY src/huggingface_inference_toolkit/webservice_starlette.py webservice_starlette.py
 
+# copy entrypoint and change permissions
+COPY scripts/entrypoint.sh entrypoint.sh
+RUN chmod +x entrypoint.sh
+
 # run app
-ENTRYPOINT ["uvicorn", "webservice_starlette:app", "--host", "0.0.0.0", "--port", "5000"]
+ENTRYPOINT ["/bin/bash", "entrypoint.sh"]
 
 

dockerfiles/starlette/pytorch/Dockerfile.gpu

@@ -5,7 +5,7 @@ COPY starlette_requirements.txt /tmp/requirements.txt
 RUN pip install  --no-cache-dir -r /tmp/requirements.txt && rm /tmp/requirements.txt
 
 # Think about a better solution -> base contaienr has pt 1.13. thats why need below 0.14
-RUN pip install  --no-cache-dir sentence_transformers torchvision=="0.14.0" diffusers=="0.9.0" accelerate=="0.14.0" --extra-index-url https://download.pytorch.org/whl/cu116
+RUN pip install  --no-cache-dir sentence_transformers=="2.2.2" torchvision=="0.14.0" diffusers=="0.9.0" accelerate=="0.14.0" --extra-index-url https://download.pytorch.org/whl/cu116
 
 # Add upgrade due to issue in base container upgrade https://github.com/mamba-org/mamba/issues/2170
 RUN pip install transformers==4.25.1  --no-cache-dir --upgrade 
@@ -14,6 +14,9 @@ RUN pip install transformers==4.25.1  --no-cache-dir --upgrade
 COPY src/huggingface_inference_toolkit huggingface_inference_toolkit
 COPY src/huggingface_inference_toolkit/webservice_starlette.py webservice_starlette.py
 
-# run app
-ENTRYPOINT ["uvicorn", "webservice_starlette:app", "--host", "0.0.0.0", "--port", "5000"]
+# copy entrypoint and change permissions
+COPY scripts/entrypoint.sh entrypoint.sh
+RUN chmod +x entrypoint.sh
 
+# run app
+ENTRYPOINT ["/bin/bash", "entrypoint.sh"]

dockerfiles/starlette/tensorflow/Dockerfile.cpu

@@ -8,7 +8,9 @@ RUN pip install  --no-cache-dir -r /tmp/requirements.txt && rm /tmp/requirements
 COPY src/huggingface_inference_toolkit huggingface_inference_toolkit
 COPY src/huggingface_inference_toolkit/webservice_starlette.py webservice_starlette.py
 
-# run app
-ENTRYPOINT ["uvicorn", "webservice_starlette:app", "--host", "0.0.0.0", "--port", "5000"]
-
+# copy entrypoint and change permissions
+COPY scripts/entrypoint.sh entrypoint.sh
+RUN chmod +x entrypoint.sh
 
+# run app
+ENTRYPOINT ["/bin/bash", "entrypoint.sh"]

dockerfiles/starlette/tensorflow/Dockerfile.gpu

@@ -8,5 +8,9 @@ RUN pip install  --no-cache-dir -r /tmp/requirements.txt && rm /tmp/requirements
 COPY src/huggingface_inference_toolkit huggingface_inference_toolkit
 COPY src/huggingface_inference_toolkit/webservice_starlette.py webservice_starlette.py
 
+# copy entrypoint and change permissions
+COPY scripts/entrypoint.sh entrypoint.sh
+RUN chmod +x entrypoint.sh
+
 # run app
-ENTRYPOINT ["uvicorn", "webservice_starlette:app","--host", "0.0.0.0", "--port", "5000"]
+ENTRYPOINT ["/bin/bash", "entrypoint.sh"]

scripts/entrypoint.sh

@@ -0,0 +1,13 @@
+# /bin/bash
+
+# check if HF_MODEL_DIR is set and if not skip installing custom dependencies
+if [ -z "${HF_MODEL_DIR}" ]; then 
+  # check if requirements.txt exists and if so install dependencies
+  if [ -f "${HF_MODEL_DIR}/requirements.txt" ]; then
+    echo "Installing custom dependencies from ${HF_MODEL_DIR}/requirements.txt"
+    pip install -r ${HF_MODEL_DIR}/requirements.txt --no-cache-dir;
+  fi
+fi
+
+# start the server
+uvicorn webservice_starlette:app --host 0.0.0.0 --port 5000

src/huggingface_inference_toolkit/diffusers_utils.py

@@ -30,7 +30,6 @@ def check_supported_pipeline(model_dir):
 
 class DiffusersPipelineImageToText:
     def __init__(self, model_dir: str, device: str = None):  # needs "cuda" for GPU
-
         self.pipeline = StableDiffusionPipeline.from_pretrained(model_dir, torch_dtype=torch.float16)
         # try to use DPMSolverMultistepScheduler
         try: