Skip to content

Commit

Permalink
Add pytorch/inference/{cpu,gpu}/2.3.1/transformers/4.48.0/py311 (#135)
Browse files Browse the repository at this point in the history
* Add `pytorch/inference/gpu/2.3.1/transformers/4.47.0/py311`

* Add `pytorch/inference/cpu/2.3.1/transformers/4.47.0/py311`

* Bump `transformers` version based on `huggingface-inference-toolkit` (WIP)

* Set `huggingface-inference-toolkit` released version

* Update `entrypoint.sh` for PyTorch Inference DLC

* Fix code-comment formatting

* Fix `PORT` handling and remove `-u` from `entrypoint.sh`

* Add `flash-attn` installation to PyTorch Training on GPU

* Fix `entrypoint.sh` on env var check and `HF_MODEL_DIR` handling
  • Loading branch information
alvarobartt authored Jan 28, 2025
1 parent 04b2ab7 commit 60b829f
Show file tree
Hide file tree
Showing 4 changed files with 284 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
FROM ubuntu:22.04
SHELL ["/bin/bash", "-c"]

LABEL maintainer="Hugging Face"

ENV DEBIAN_FRONTEND=noninteractive

WORKDIR /app

# Install required dependencies
RUN apt-get update && \
apt-get install software-properties-common -y && \
add-apt-repository ppa:deadsnakes/ppa && \
apt-get -y upgrade --only-upgrade systemd openssl cryptsetup && \
apt-get install -y \
build-essential \
bzip2 \
curl \
git \
git-lfs \
tar \
gcc \
g++ \
cmake \
libprotobuf-dev \
protobuf-compiler \
python3.11 \
python3.11-dev \
libsndfile1-dev \
ffmpeg && \
rm -rf /var/lib/apt/lists/*

# Set Python 3.11 as the default python version
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1 && \
ln -sf /usr/bin/python3.11 /usr/bin/python

# Install pip from source
RUN curl -O https://bootstrap.pypa.io/get-pip.py && \
python get-pip.py && \
rm get-pip.py

# Hugging Face Inference Toolkit
ARG HF_INFERENCE_TOOLKIT_VERSION=0.5.4
ARG HF_INFERENCE_TOOLKIT_URL=git+https://github.com/huggingface/huggingface-inference-toolkit.git@${HF_INFERENCE_TOOLKIT_VERSION}
RUN pip install --upgrade "huggingface-inference-toolkit[torch,diffusers,st,google] @ ${HF_INFERENCE_TOOLKIT_URL}" --no-cache-dir

ENV HF_HUB_ENABLE_HF_TRANSFER="1"

# Install Google CLI single command
RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" \
| tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg \
| apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
apt-get update -y && \
apt-get install google-cloud-sdk -y && \
apt-get clean autoremove --yes && \
rm -rf /var/lib/{apt,dpkg,cache,log}

# Copy entrypoint and change permissions
COPY --chmod=0755 containers/pytorch/inference/cpu/2.3.1/transformers/4.48.0/py311/entrypoint.sh entrypoint.sh
ENTRYPOINT ["bash", "-c", "./entrypoint.sh"]
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
#!/bin/bash

set -eo pipefail

# Define the default port
readonly DEFAULT_PORT=5000

# Check if `AIP_MODE` is set and adjust the port for Vertex AI
if [[ ! -z "${AIP_MODE}" ]]; then
PORT="${AIP_HTTP_PORT:-$DEFAULT_PORT}"
else
PORT="$DEFAULT_PORT"
fi

# Check if exactly one of HF_MODEL_ID, HF_MODEL_DIR, or AIP_STORAGE_URI is set, by
# concatenating thosa and counting non-empty ones using pattern substitution
if [[ $(printf "%s\n" "${HF_MODEL_ID:+x}" "${HF_MODEL_DIR:+x}" "${AIP_STORAGE_URI:+x}" | grep -c .) -ne 1 ]]; then
echo "ERROR: Exactly one of HF_MODEL_ID, HF_MODEL_DIR, or AIP_STORAGE_URI must be provided."
exit 1
fi

# If `HF_MODEL_ID` is a path instead of a Hub ID, then clear its value and assign it
# to the `HF_MODEL_DIR` instead, including a user warning
if [[ -d "${HF_MODEL_ID:-}" ]]; then
echo "WARNING: HF_MODEL_ID is a path, please use HF_MODEL_DIR for paths instead."
HF_MODEL_DIR="${HF_MODEL_ID}"
HF_MODEL_ID=""
fi

# Check if `MODEL_ID` starts with "gcs://"
if [[ "${AIP_STORAGE_URI:-}" == gs://* ]]; then
echo "INFO: AIP_STORAGE_URI set and starts with 'gs://', proceeding to download from GCS."
echo "INFO: AIP_STORAGE_URI: $AIP_STORAGE_URI"

# Define the target directory
TARGET_DIR="/opt/huggingface/model"
mkdir -p "$TARGET_DIR"

# Check if `gsutil` is available
if ! command -v gsutil &> /dev/null; then
echo "ERROR: gsutil command not found. Please install Google Cloud SDK." >&2
exit 1
fi

# Use `gsutil` to copy the content from GCS to the target directory
echo "INFO: Running: gsutil -m cp -e -r \"$AIP_STORAGE_URI/*\" \"$TARGET_DIR\""
if ! gsutil -m cp -e -r "$AIP_STORAGE_URI/*" "$TARGET_DIR"; then
echo "ERROR: Failed to download model from GCS." >&2
exit 1
fi

echo "INFO: Model downloaded successfully to ${TARGET_DIR}."
echo "INFO: Updating HF_MODEL_DIR to point to the local directory."
HF_MODEL_DIR="$TARGET_DIR"
AIP_STORAGE_URI=""
fi

# If `HF_MODEL_DIR` is set is a valid directory
if [[ -n "${HF_MODEL_DIR:-}" ]]; then
if [[ ! -d "${HF_MODEL_DIR}" ]]; then
echo "ERROR: Provided HF_MODEL_DIR is not a valid directory" >&2
exit 1
fi

# Check if `requirements.txt` exists and if so install dependencies
if [[ -f "${HF_MODEL_DIR}/requirements.txt" ]]; then
echo "INFO: Installing custom dependencies from ${HF_MODEL_DIR}/requirements.txt"
pip install -r "${HF_MODEL_DIR}/requirements.txt" --no-cache-dir

# Check if `handler.py` is missing when `requirements.txt` is present
if [[ ! -f "${HF_MODEL_DIR}/handler.py" ]]; then
echo "WARNING: requirements.txt is present, but handler.py is missing in ${HF_MODEL_DIR}."
echo "WARNING: If you intend to run custom code, make sure to include handler.py."
fi
fi
fi

# Start the server
exec uvicorn huggingface_inference_toolkit.webservice_starlette:app --host 0.0.0.0 --port "${PORT}"
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
FROM nvidia/cuda:12.1.0-devel-ubuntu22.04
SHELL ["/bin/bash", "-c"]

LABEL maintainer="Hugging Face"

ENV DEBIAN_FRONTEND=noninteractive

WORKDIR /app

# Install required dependencies
RUN apt-get update && \
apt-get install software-properties-common -y && \
add-apt-repository ppa:deadsnakes/ppa && \
apt-get -y upgrade --only-upgrade systemd openssl cryptsetup && \
apt-get install -y \
build-essential \
bzip2 \
curl \
git \
git-lfs \
tar \
gcc \
g++ \
cmake \
libprotobuf-dev \
protobuf-compiler \
python3.11 \
python3.11-dev \
libsndfile1-dev \
ffmpeg && \
rm -rf /var/lib/apt/lists/*

# Set Python 3.11 as the default python version
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1 && \
ln -sf /usr/bin/python3.11 /usr/bin/python

# Install pip from source
RUN curl -O https://bootstrap.pypa.io/get-pip.py && \
python get-pip.py && \
rm get-pip.py

# Hugging Face Inference Toolkit
ARG HF_INFERENCE_TOOLKIT_VERSION=0.5.4
ARG HF_INFERENCE_TOOLKIT_URL=git+https://github.com/huggingface/huggingface-inference-toolkit.git@${HF_INFERENCE_TOOLKIT_VERSION}
RUN pip install --upgrade "huggingface-inference-toolkit[torch,diffusers,st,google] @ ${HF_INFERENCE_TOOLKIT_URL}" --no-cache-dir

ENV HF_HUB_ENABLE_HF_TRANSFER="1"

# Install Flash Attention 2
RUN pip install --no-cache-dir packaging ninja
RUN MAX_JOBS=4 pip install --no-build-isolation "flash-attn==2.6.3"

# Install Google CLI single command
RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" \
| tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg \
| apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
apt-get update -y && \
apt-get install google-cloud-sdk -y && \
apt-get clean autoremove --yes && \
rm -rf /var/lib/{apt,dpkg,cache,log}

# Copy entrypoint and change permissions
COPY --chmod=0755 containers/pytorch/inference/gpu/2.3.1/transformers/4.48.0/py311/entrypoint.sh entrypoint.sh
ENTRYPOINT ["bash", "-c", "./entrypoint.sh"]
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
#!/bin/bash

set -eo pipefail

# Define the default port
readonly DEFAULT_PORT=5000

# Check if `AIP_MODE` is set and adjust the port for Vertex AI
if [[ ! -z "${AIP_MODE}" ]]; then
PORT="${AIP_HTTP_PORT:-$DEFAULT_PORT}"
else
PORT="$DEFAULT_PORT"
fi

# Check if exactly one of HF_MODEL_ID, HF_MODEL_DIR, or AIP_STORAGE_URI is set, by
# concatenating thosa and counting non-empty ones using pattern substitution
if [[ $(printf "%s\n" "${HF_MODEL_ID:+x}" "${HF_MODEL_DIR:+x}" "${AIP_STORAGE_URI:+x}" | grep -c .) -ne 1 ]]; then
echo "ERROR: Exactly one of HF_MODEL_ID, HF_MODEL_DIR, or AIP_STORAGE_URI must be provided."
exit 1
fi

# If `HF_MODEL_ID` is a path instead of a Hub ID, then clear its value and assign it
# to the `HF_MODEL_DIR` instead, including a user warning
if [[ -d "${HF_MODEL_ID:-}" ]]; then
echo "WARNING: HF_MODEL_ID is a path, please use HF_MODEL_DIR for paths instead."
HF_MODEL_DIR="${HF_MODEL_ID}"
HF_MODEL_ID=""
fi

# Check if `MODEL_ID` starts with "gcs://"
if [[ "${AIP_STORAGE_URI:-}" == gs://* ]]; then
echo "INFO: AIP_STORAGE_URI set and starts with 'gs://', proceeding to download from GCS."
echo "INFO: AIP_STORAGE_URI: $AIP_STORAGE_URI"

# Define the target directory
TARGET_DIR="/opt/huggingface/model"
mkdir -p "$TARGET_DIR"

# Check if `gsutil` is available
if ! command -v gsutil &> /dev/null; then
echo "ERROR: gsutil command not found. Please install Google Cloud SDK." >&2
exit 1
fi

# Use `gsutil` to copy the content from GCS to the target directory
echo "INFO: Running: gsutil -m cp -e -r \"$AIP_STORAGE_URI/*\" \"$TARGET_DIR\""
if ! gsutil -m cp -e -r "$AIP_STORAGE_URI/*" "$TARGET_DIR"; then
echo "ERROR: Failed to download model from GCS." >&2
exit 1
fi

echo "INFO: Model downloaded successfully to ${TARGET_DIR}."
echo "INFO: Updating HF_MODEL_DIR to point to the local directory."
HF_MODEL_DIR="$TARGET_DIR"
AIP_STORAGE_URI=""
fi

# If `HF_MODEL_DIR` is set is a valid directory
if [[ -n "${HF_MODEL_DIR:-}" ]]; then
if [[ ! -d "${HF_MODEL_DIR}" ]]; then
echo "ERROR: Provided HF_MODEL_DIR is not a valid directory" >&2
exit 1
fi

# Check if `requirements.txt` exists and if so install dependencies
if [[ -f "${HF_MODEL_DIR}/requirements.txt" ]]; then
echo "INFO: Installing custom dependencies from ${HF_MODEL_DIR}/requirements.txt"
pip install -r "${HF_MODEL_DIR}/requirements.txt" --no-cache-dir

# Check if `handler.py` is missing when `requirements.txt` is present
if [[ ! -f "${HF_MODEL_DIR}/handler.py" ]]; then
echo "WARNING: requirements.txt is present, but handler.py is missing in ${HF_MODEL_DIR}."
echo "WARNING: If you intend to run custom code, make sure to include handler.py."
fi
fi
fi

# Start the server
exec uvicorn huggingface_inference_toolkit.webservice_starlette:app --host 0.0.0.0 --port "${PORT}"

0 comments on commit 60b829f

Please sign in to comment.