Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ROCm support for Dockerfile and build_image.sh #3381

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

ROCm support for Dockerfile and build_image.sh #3381

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ pip install torchserve-nightly torch-model-archiver-nightly torch-workflow-archi
# Install dependencies
python ./ts_scripts/install_dependencies.py

# Include depeendencies for accelerator support with the relevant optional flags
# Include dependencies for accelerator support with the relevant optional flags
python ./ts_scripts/install_dependencies.py --rocm=rocm61
python ./ts_scripts/install_dependencies.py --cuda=cu121

Expand Down
67 changes: 58 additions & 9 deletions docker/Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,12 +37,12 @@ ARG BRANCH_NAME
ARG REPO_URL=https://github.com/pytorch/serve.git
ENV PYTHONUNBUFFERED TRUE

RUN --mount=type=cache,id=apt-dev,target=/var/cache/apt \
RUN --mount=type=cache,sharing=locked,id=apt-dev,target=/var/cache/apt \
apt-get update && \
apt-get upgrade -y && \
apt-get install software-properties-common -y && \
add-apt-repository -y ppa:deadsnakes/ppa && \
apt remove python-pip python3-pip && \
apt remove -y python-pip python3-pip && \
DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \
ca-certificates \
g++ \
Expand All @@ -55,6 +55,13 @@ RUN --mount=type=cache,id=apt-dev,target=/var/cache/apt \
git \
&& rm -rf /var/lib/apt/lists/*

RUN --mount=type=cache,sharing=locked,id=apt-dev,target=/var/cache/apt \
if [ "$USE_ROCM_VERSION" ]; then \
apt-get update && \
DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y rocm-dev amd-smi-lib \
&& rm -rf /var/lib/apt/lists/* ; \
fi

# Make the virtual environment and "activating" it by adding it first to the path.
# From here on the python$PYTHON_VERSION interpreter is used and the packages
# are installed in /home/venv which is what we need for the "runtime-image"
Expand All @@ -67,6 +74,7 @@ RUN python -m pip install -U pip setuptools
RUN export USE_CUDA=1

ARG USE_CUDA_VERSION=""
ARG USE_ROCM_VERSION=""

COPY ./ serve

Expand All @@ -76,7 +84,6 @@ RUN \
git clone --recursive $REPO_URL -b $BRANCH_NAME serve; \
fi


WORKDIR "serve"

RUN cp docker/dockerd-entrypoint.sh /usr/local/bin/dockerd-entrypoint.sh
Expand All @@ -90,6 +97,14 @@ RUN \
else \
python ./ts_scripts/install_dependencies.py;\
fi; \
elif echo "${BASE_IMAGE}" | grep -q "rocm/"; then \
# Install ROCm version specific binary when ROCm version is specified as a build arg
if [ "$USE_ROCM_VERSION" ]; then \
python ./ts_scripts/install_dependencies.py --rocm $USE_ROCM_VERSION;\
# Install the binary with the latest CPU image on a ROCm base image
else \
python ./ts_scripts/install_dependencies.py; \
fi; \
# Install the CPU binary
else \
python ./ts_scripts/install_dependencies.py; \
Expand All @@ -111,13 +126,14 @@ FROM ${BASE_IMAGE} AS production-image
# Re-state ARG PYTHON_VERSION to make it active in this build-stage (uses default define at the top)
ARG PYTHON_VERSION
ENV PYTHONUNBUFFERED TRUE
ARG USE_ROCM_VERSION

RUN --mount=type=cache,target=/var/cache/apt \
RUN --mount=type=cache,sharing=locked,target=/var/cache/apt \
apt-get update && \
apt-get upgrade -y && \
apt-get install software-properties-common -y && \
add-apt-repository ppa:deadsnakes/ppa -y && \
apt remove python-pip python3-pip && \
apt remove -y python-pip python3-pip && \
DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \
python$PYTHON_VERSION \
python3-distutils \
Expand All @@ -130,13 +146,25 @@ RUN --mount=type=cache,target=/var/cache/apt \
&& rm -rf /var/lib/apt/lists/* \
&& cd /tmp

RUN --mount=type=cache,sharing=locked,id=apt-dev,target=/var/cache/apt \
if [ "$USE_ROCM_VERSION" ]; then \
apt-get update && \
DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y rocm-dev amd-smi-lib \
&& rm -rf /var/lib/apt/lists/* ; \
fi

RUN useradd -m model-server \
&& mkdir -p /home/model-server/tmp

COPY --chown=model-server --from=compile-image /home/venv /home/venv
COPY --from=compile-image /usr/local/bin/dockerd-entrypoint.sh /usr/local/bin/dockerd-entrypoint.sh
ENV PATH="/home/venv/bin:$PATH"

RUN \
if [ "$USE_ROCM_VERSION" ]; then \
python -m pip install /opt/rocm/share/amd_smi; \
fi

RUN chmod +x /usr/local/bin/dockerd-entrypoint.sh \
&& chown -R model-server /home/model-server

Expand All @@ -157,13 +185,14 @@ FROM ${BASE_IMAGE} AS ci-image
ARG PYTHON_VERSION
ARG BRANCH_NAME
ENV PYTHONUNBUFFERED TRUE
ARG USE_ROCM_VERSION

RUN --mount=type=cache,target=/var/cache/apt \
RUN --mount=type=cache,sharing=locked,target=/var/cache/apt \
apt-get update && \
apt-get upgrade -y && \
apt-get install software-properties-common -y && \
add-apt-repository -y ppa:deadsnakes/ppa && \
apt remove python-pip python3-pip && \
apt remove -y python-pip python3-pip && \
DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \
python$PYTHON_VERSION \
python3-distutils \
Expand All @@ -183,13 +212,24 @@ RUN --mount=type=cache,target=/var/cache/apt \
&& rm -rf /var/lib/apt/lists/* \
&& cd /tmp

RUN --mount=type=cache,sharing=locked,id=apt-dev,target=/var/cache/apt \
if [ "$USE_ROCM_VERSION" ]; then \
apt-get update && \
DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y rocm-dev amd-smi-lib \
&& rm -rf /var/lib/apt/lists/* ; \
fi

COPY --from=compile-image /home/venv /home/venv

ENV PATH="/home/venv/bin:$PATH"

RUN python -m pip install --no-cache-dir -r https://raw.githubusercontent.com/pytorch/serve/$BRANCH_NAME/requirements/developer.txt

RUN \
if [ "$USE_ROCM_VERSION" ]; then \
python -m pip install /opt/rocm/share/amd_smi; \
fi

RUN mkdir /home/serve
ENV TS_RUN_IN_DOCKER True

Expand All @@ -203,11 +243,13 @@ ARG PYTHON_VERSION
ARG BRANCH_NAME
ARG BUILD_FROM_SRC
ARG LOCAL_CHANGES
ARG USE_ROCM_VERSION
ARG BUILD_WITH_IPEX
ARG IPEX_VERSION=1.11.0
ARG IPEX_URL=https://software.intel.com/ipex-whl-stable
ENV PYTHONUNBUFFERED TRUE
RUN --mount=type=cache,target=/var/cache/apt \

RUN --mount=type=cache,sharing=locked,target=/var/cache/apt \
apt-get update && \
apt-get upgrade -y && \
apt-get install software-properties-common -y && \
Expand Down Expand Up @@ -243,10 +285,17 @@ RUN \

COPY --from=compile-image /home/venv /home/venv
ENV PATH="/home/venv/bin:$PATH"

RUN \
if [ "$USE_ROCM_VERSION" ]; then \
python -m pip install /opt/rocm/share/amd_smi; \
fi

WORKDIR "serve"

RUN python -m pip install -U pip setuptools \
&& python -m pip install --no-cache-dir -r requirements/developer.txt \
&& python ts_scripts/install_from_src.py \
&& python ts_scripts/install_from_src.py --environment=dev\
&& useradd -m model-server \
&& mkdir -p /home/model-server/tmp \
&& cp docker/dockerd-entrypoint.sh /usr/local/bin/dockerd-entrypoint.sh \
Expand Down
Loading