Fix Dockerfile and some minor updates

.dockerignore

@@ -1,10 +1,11 @@
+.git
 .github
 .pytest_cache
 .ruff_cache
 .tox
 .venv
 .gitignore
-makefile
+Makefile
 __pycache__
 tests
 .vscode

MANIFEST.in

@@ -1,6 +1,5 @@
-
 include LICENSE
 include README.md
 
 recursive-exclude * __pycache__
-recursive-exclude * *.py[co]
+recursive-exclude * *.py[co]

makefile → Makefile

@@ -1,39 +1,35 @@
-.PHONY: quality style unit-test integ-test
+.PHONY: quality style unit-test integ-test inference-pytorch-gpu inference-pytorch-cpu inference-pytorch-inf2 stop-all
 
-check_dirs := src tests 
+check_dirs := src tests
 
-# run tests
+# Check that source code meets quality standards
+quality:
+	ruff check $(check_dirs)
+
+# Format source code automatically
+style:
+	ruff check $(check_dirs) --fix
 
+# Run unit tests
 unit-test:
 	RUN_SLOW=True python3 -m pytest -s -v tests/unit -n 10 --log-cli-level='ERROR'
 
+# Run integration tests
 integ-test:
 	python3 -m pytest -s -v tests/integ/
 
-# Check that source code meets quality standards
-
-quality:
-	ruff check $(check_dirs) 
-
-# Format source code automatically
-
-style: 
-	ruff check $(check_dirs) --fix
-
+# Build Docker image for PyTorch on GPU
 inference-pytorch-gpu:
 	docker build -f dockerfiles/pytorch/Dockerfile -t integration-test-pytorch:gpu .
 
+# Build Docker image for PyTorch on CPU
 inference-pytorch-cpu:
 	docker build --build-arg="BASE_IMAGE=ubuntu:22.04" -f dockerfiles/pytorch/Dockerfile -t integration-test-pytorch:cpu .
 
+# Build Docker image for PyTorch on AWS Inferentia2
 inference-pytorch-inf2:
 	docker build -f dockerfiles/pytorch/Dockerfile.inf2 -t integration-test-pytorch:inf2 .
 
-vertex-pytorch-gpu:
-	docker build -t vertex -f dockerfiles/pytorch/Dockerfile -t integration-test-pytorch:gpu .
-
-vertex-pytorch-cpu:
-	docker build  -t vertex --build-arg="BASE_IMAGE=ubuntu:22.04" -f dockerfiles/pytorch/Dockerfile -t integration-test-pytorch:cpu .
-
+# Stop all and prune/clean the Docker Containers 
 stop-all:
-	docker stop $$(docker ps -a -q) && docker container prune --force
+	docker stop $$(docker ps -a -q) && docker container prune --force

dockerfiles/pytorch/Dockerfile

@@ -1,6 +1,6 @@
 ARG BASE_IMAGE=nvidia/cuda:12.1.0-devel-ubuntu22.04
 
-FROM $BASE_IMAGE as base 
+FROM $BASE_IMAGE AS base
 SHELL ["/bin/bash", "-c"]
 
 LABEL maintainer="Hugging Face"
@@ -25,9 +25,8 @@ RUN apt-get update && \
     cmake \
     libprotobuf-dev \
     protobuf-compiler \
-    python3-dev \
-    python3-pip \
     python3.11 \
+    python3.11-dev \
     libsndfile1-dev \
     ffmpeg \
     && apt-get clean autoremove --yes \
@@ -36,6 +35,15 @@ RUN apt-get update && \
 # Copying only necessary files as filtered by .dockerignore
 COPY . .
 
+# 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
+
 # install wheel and setuptools
 RUN pip install --no-cache-dir --upgrade pip ".[torch,st,diffusers]"
 
@@ -44,11 +52,6 @@ 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 --chmod=0755  scripts/entrypoint.sh entrypoint.sh
+COPY --chmod=0755 scripts/entrypoint.sh entrypoint.sh
 
 ENTRYPOINT ["bash", "-c", "./entrypoint.sh"]
-
-FROM base AS vertex
-
-# Install `google` extra for Vertex AI compatibility
-RUN pip install --no-cache-dir --upgrade ".[google]"

pyproject.toml

@@ -5,25 +5,25 @@ scripts_are_modules = true
 
 [tool.ruff]
 select = [
-  "E", # pycodestyle errors
-  "W", # pycodestyle warnings
-  "F", # pyflakes
-  "I", # isort
-  "C", # flake8-comprehensions
-  "B", # flake8-bugbear
+    "E", # pycodestyle errors
+    "W", # pycodestyle warnings
+    "F", # pyflakes
+    "I", # isort
+    "C", # flake8-comprehensions
+    "B", # flake8-bugbear
 ]
 ignore = [
-  "E501", # Line length (handled by ruff-format)
-  "B008", # do not perform function calls in argument defaults
-  "C901", # too complex
+    "E501", # Line length (handled by ruff-format)
+    "B008", # do not perform function calls in argument defaults
+    "C901", # too complex
 ]
 # Same as Black.
 line-length = 119
 
 # Allow unused variables when underscore-prefixed.
 dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
 
-# Assume Python 3.11.
+# Assume Python 3.11
 target-version = "py311"
 
 per-file-ignores = { "__init__.py" = ["F401"] }

scripts/entrypoint.sh

@@ -5,16 +5,16 @@ PORT=5000
 
 # Check if AIP_MODE is set and adjust the port for Vertex AI
 if [[ ! -z "${AIP_MODE}" ]]; then
-  PORT=${AIP_HTTP_PORT}
+    PORT=${AIP_HTTP_PORT}
 fi
 
 # 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
+    # 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

setup.py

@@ -62,9 +62,7 @@
     packages=find_packages(where="src"),
     install_requires=install_requires,
     extras_require=extras,
-    entry_points={
-        "console_scripts": "serve=sagemaker_huggingface_inference_toolkit.serving:main"
-    },
+    entry_points={"console_scripts": "serve=sagemaker_huggingface_inference_toolkit.serving:main"},
     python_requires=">=3.8",
     license="Apache License 2.0",
     classifiers=[