diff --git a/dockerfiles/pytorch/Dockerfile b/dockerfiles/pytorch/Dockerfile index 4ca98a86..12cb541e 100644 --- a/dockerfiles/pytorch/Dockerfile +++ b/dockerfiles/pytorch/Dockerfile @@ -16,13 +16,13 @@ RUN apt-get update && \ apt-get install -y \ build-essential \ bzip2 \ + cmake \ curl \ git \ git-lfs \ tar \ gcc \ g++ \ - cmake \ libprotobuf-dev \ protobuf-compiler \ python3.11 \ @@ -44,14 +44,17 @@ RUN curl -O https://bootstrap.pypa.io/get-pip.py && \ python get-pip.py && \ rm get-pip.py -# install wheel and setuptools +# Upgrade pip +RUN pip install --no-cache-dir --upgrade pip + +# Install wheel and setuptools RUN pip install --no-cache-dir --upgrade pip ".[torch,st,diffusers]" -# copy application +# Copy application 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 entrypoint and change permissions COPY --chmod=0755 scripts/entrypoint.sh entrypoint.sh ENTRYPOINT ["bash", "-c", "./entrypoint.sh"] diff --git a/scripts/entrypoint.sh b/scripts/entrypoint.sh old mode 100644 new mode 100755 index f89d83c7..68969353 --- a/scripts/entrypoint.sh +++ b/scripts/entrypoint.sh @@ -1,6 +1,6 @@ #!/bin/bash -# Define the default port +# Set the default port PORT=5000 # Check if AIP_MODE is set and adjust the port for Vertex AI @@ -8,7 +8,48 @@ if [[ ! -z "${AIP_MODE}" ]]; then PORT=${AIP_HTTP_PORT} fi -# Check if HF_MODEL_DIR is set and if not skip installing custom dependencies +# Check that only one of HF_MODEL_ID or HF_MODEL_DIR is provided +if [[ ! -z "${HF_MODEL_ID}" && ! -z "${HF_MODEL_DIR}" ]]; then + echo "Error: Both HF_MODEL_ID and HF_MODEL_DIR are set. Please provide only one." + exit 1 +elif [[ -z "${HF_MODEL_ID}" && -z "${HF_MODEL_DIR}" ]]; then + echo "Error: Neither HF_MODEL_ID nor HF_MODEL_DIR is set. Please provide one of them." + exit 1 +fi + +# If HF_MODEL_ID is provided, download handler.py and requirements.txt if available +if [[ ! -z "${HF_MODEL_ID}" ]]; then + filename=${HF_DEFAULT_PIPELINE_NAME:-handler.py} + revision=${HF_REVISION:-main} + + echo "Downloading $filename for model ${HF_MODEL_ID}" + huggingface-cli download ${HF_MODEL_ID} "$filename" --revision "$revision" --local-dir /tmp + + # Check if handler.py was downloaded successfully + if [ -f "/tmp/$filename" ]; then + echo "$filename downloaded successfully, checking if there's a requirements.txt file..." + rm /tmp/$filename + + # Attempt to download requirements.txt + echo "Downloading requirements.txt for model ${HF_MODEL_ID}" + huggingface-cli download "${HF_MODEL_ID}" requirements.txt --revision "$revision" --local-dir /tmp + + # Check if requirements.txt was downloaded successfully + if [ -f "/tmp/requirements.txt" ]; then + echo "requirements.txt downloaded successfully, now installing the dependencies..." + + # Install dependencies + pip install -r /tmp/requirements.txt --no-cache-dir + rm /tmp/requirements.txt + else + echo "${HF_MODEL_ID} with revision $revision contains a custom handler at $filename but doesn't contain a requirements.txt file, so skipping downloading and installing extra requirements from it." + fi + else + echo "${HF_MODEL_ID} with revision $revision doesn't contain a $filename file, so skipping download." + fi +fi + +# If HF_MODEL_DIR is provided, check for requirements.txt and install dependencies if available if [[ ! -z "${HF_MODEL_DIR}" ]]; then # Check if requirements.txt exists and if so install dependencies if [ -f "${HF_MODEL_DIR}/requirements.txt" ]; then diff --git a/setup.py b/setup.py index d5b3da68..658d67ab 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ # We don't declare our dependency on transformers here because we build with # different packages for different variants -VERSION = "0.5.4" +VERSION = "0.5.5" # Ubuntu packages # libsndfile1-dev: torchaudio requires the development version of the libsndfile package which can be installed via a system package manager. On Ubuntu it can be installed as follows: apt install libsndfile1-dev @@ -13,6 +13,9 @@ # libavcodec-extra : libavcodec-extra includes additional codecs for ffmpeg install_requires = [ + # Due to an error affecting kenlm and cmake (see https://github.com/kpu/kenlm/pull/464) + # Also see the transformers patch for it https://github.com/huggingface/transformers/pull/37091 + "kenlm@git+https://github.com/kpu/kenlm@ba83eafdce6553addd885ed3da461bb0d60f8df7", "transformers[sklearn,sentencepiece,audio,vision]==4.48.0", "huggingface_hub[hf_transfer]==0.27.1", # vision