[0.28.0 dlc][cherry-pick] AOT fix and pin numpy (#2086)

Co-authored-by: Andrew Song <[email protected]> Co-authored-by: Siddharth Venkatesan <[email protected]>
deepjavalibrary · Jun 18, 2024 · 73a782d · 73a782d
1 parent bdee553
commit 73a782d
Show file tree

Hide file tree

Showing 14 changed files with 74 additions and 35 deletions.
diff --git a/.github/workflows/continuous.yml b/.github/workflows/continuous.yml
@@ -32,7 +32,7 @@ jobs:
         with:
           python-version: '3.11'
       - name: install Python Dependencies
-        run: pip3 install numpy yapf
+        run: pip3 install "numpy<2" yapf
       # Enable gradle cache: https://github.com/actions/cache/blob/master/examples.md#java---gradle
       - uses: actions/cache@v4
         with:

diff --git a/.github/workflows/llm_inf2_integration.yml b/.github/workflows/llm_inf2_integration.yml
@@ -54,7 +54,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy pillow wheel
+        run: pip3 install requests "numpy<2" pillow wheel
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh pytorch-inf2 ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers
@@ -109,7 +109,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy pillow
+        run: pip3 install requests "numpy<2" pillow
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh pytorch-inf2 ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers
@@ -267,7 +267,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy pillow
+        run: pip3 install requests "numpy<2" pillow
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh pytorch-inf2 ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers
@@ -365,7 +365,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy pillow
+        run: pip3 install requests "numpy<2" pillow
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh pytorch-inf2 ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers

diff --git a/.github/workflows/llm_integration.yml b/.github/workflows/llm_integration.yml
@@ -72,7 +72,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy huggingface_hub
+        run: pip3 install requests "numpy<2" huggingface_hub
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh ${{ matrix.arch }} ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers
@@ -176,7 +176,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy huggingface_hub
+        run: pip3 install requests "numpy<2" huggingface_hub
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh tensorrt-llm ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers
@@ -290,7 +290,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy huggingface_hub
+        run: pip3 install requests "numpy<2" huggingface_hub
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh tensorrt-llm ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers
@@ -515,7 +515,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy huggingface_hub
+        run: pip3 install requests "numpy<2" huggingface_hub
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh lmi ${{ github.event.inputs.djl-version }}
       - name: Download docker
@@ -605,7 +605,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy
+        run: pip3 install requests "numpy<2"
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh lmi ${{ github.event.inputs.djl-version }}
       - name: Download docker
@@ -723,7 +723,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy huggingface_hub
+        run: pip3 install requests "numpy<2" huggingface_hub
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh lmi ${{ github.event.inputs.djl-version }}
       - name: Download docker
@@ -833,7 +833,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy huggingface_hub
+        run: pip3 install requests "numpy<2" huggingface_hub
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh lmi ${{ github.event.inputs.djl-version }}
       - name: Download docker
@@ -915,7 +915,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy huggingface_hub
+        run: pip3 install requests "numpy<2" huggingface_hub
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh lmi ${{ github.event.inputs.djl-version }}
       - name: Download docker

diff --git a/.github/workflows/llm_integration_p4d.yml b/.github/workflows/llm_integration_p4d.yml
@@ -49,7 +49,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy
+        run: pip3 install requests "numpy<2"
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh lmi ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers
@@ -126,7 +126,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy
+        run: pip3 install requests "numpy<2"
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh tensorrt-llm ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers
@@ -185,7 +185,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy
+        run: pip3 install requests "numpy<2"
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh lmi ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers

diff --git a/.github/workflows/lmi-dist-deps-build.yml b/.github/workflows/lmi-dist-deps-build.yml
@@ -45,7 +45,7 @@ jobs:
           python -m venv venv
           . ./venv/bin/activate
           python -m pip install --upgrade pip
-          python -m pip install numpy cmake awscli packaging wheel setuptools ninja git-remote-codecommit \
+          python -m pip install "numpy<2" cmake awscli packaging wheel setuptools ninja git-remote-codecommit \
                                 torch==2.3.0 --extra-index-url https://download.pytorch.org/whl/cu121
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v4

diff --git a/.github/workflows/lmi-no-code.yml b/.github/workflows/lmi-no-code.yml
@@ -94,7 +94,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy
+        run: pip3 install requests "numpy<2"
       - name: Install s5cmd
         working-directory: serving/docker
         run: sudo scripts/install_s5cmd.sh x64
@@ -188,7 +188,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy huggingface_hub
+        run: pip3 install requests "numpy<2" huggingface_hub
       - name: Install s5cmd
         working-directory: serving/docker
         run: sudo scripts/install_s5cmd.sh x64

diff --git a/.github/workflows/lmic_performance.yml b/.github/workflows/lmic_performance.yml
@@ -53,7 +53,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy datetime
+        run: pip3 install requests "numpy<2" datetime
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh lmi ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers
@@ -89,7 +89,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy datetime
+        run: pip3 install requests "numpy<2" datetime
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh lmi ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers
@@ -125,7 +125,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy datetime
+        run: pip3 install requests "numpy<2" datetime
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh lmi ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers
@@ -161,7 +161,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy datetime
+        run: pip3 install requests "numpy<2" datetime
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh lmi ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers
@@ -197,7 +197,7 @@ jobs:
         with:
           python-version: '3.10.x'
       - name: Install pip dependencies
-        run: pip3 install requests numpy datetime
+        run: pip3 install requests "numpy<2" datetime
       - name: Build container name
         run: ./serving/docker/scripts/docker_name_builder.sh lmi ${{ github.event.inputs.djl-version }}
       - name: Download models and dockers

diff --git a/engines/python/setup/setup.py b/engines/python/setup/setup.py
@@ -56,7 +56,7 @@ def run(self):
     requirements = ['psutil', 'packaging', 'wheel']
 
     test_requirements = [
-        'numpy', 'requests', 'Pillow', 'transformers', 'torch', 'einops',
+        'numpy<2', 'requests', 'Pillow', 'transformers', 'torch', 'einops',
         'accelerate', 'sentencepiece', 'protobuf', "peft", 'yapf',
         'pydantic>=2.0', "objgraph"
     ]

diff --git a/engines/python/src/test/resources/accumulate/requirements.txt b/engines/python/src/test/resources/accumulate/requirements.txt
@@ -1 +1 @@
-numpy
+numpy<2
diff --git a/serving/docker/partition/sm_neo_quantize.py b/serving/docker/partition/sm_neo_quantize.py
@@ -20,7 +20,7 @@
 
 from sm_neo_utils import (CompilationFatalError, write_error_to_file,
                           get_neo_env_vars)
-from utils import extract_python_jar
+from utils import extract_python_jar, load_properties
 from properties_manager import PropertiesManager
 from partition import PartitionService
 
@@ -38,13 +38,14 @@ def __init__(self):
         self.INPUT_MODEL_DIRECTORY: Final[str] = env[1]
         self.OUTPUT_MODEL_DIRECTORY: Final[str] = env[2]
         self.COMPILATION_ERROR_FILE: Final[str] = env[3]
-        self.COMPILER_CACHE_LOCATION: Final[str] = env[4]
+        self.HF_CACHE_LOCATION: Final[str] = env[5]
+        self.TARGET_INSTANCE_TYPE: Final[str] = env[6]
 
     def update_dataset_cache_location(self):
         logging.info(
-            f"Updating HuggingFace Datasets cache directory to: {self.COMPILER_CACHE_LOCATION}"
+            f"Updating HuggingFace Datasets cache directory to: {self.HF_CACHE_LOCATION}"
         )
-        os.environ['HF_DATASETS_CACHE'] = self.COMPILER_CACHE_LOCATION
+        os.environ['HF_DATASETS_CACHE'] = self.HF_CACHE_LOCATION
         #os.environ['HF_DATASETS_OFFLINE'] = "1"
 
     def initialize_partition_args_namespace(self):
@@ -72,6 +73,8 @@ def construct_properties_manager(self):
         given serving.properties
         """
         # Default to awq quantization
+        # TODO: update this when new quantization methods are added,
+        # since envvar overrides customer serving.properties
         os.environ['OPTION_QUANTIZE'] = 'awq'
         logging.debug("Constructing PropertiesManager from "
                       f"serving.properties\nargs:{self.args}\n")
@@ -89,11 +92,42 @@ def run_quantization(self) -> str:
             raise CompilationFatalError(
                 f"Encountered an error during quantization: {exc}")
 
+    def write_properties(self):
+        """
+        Updates outputted serving.properties.
+        If a user passes in tensor_parallel_degree, it is passed through to the output.
+        Otherwise, tensor_parallel_degree is not outputted so that it can be defined
+        during serving.
+        """
+        customer_properties = load_properties(self.INPUT_MODEL_DIRECTORY)
+        user_tensor_parallel_degree = customer_properties.get(
+            "option.tensor_parallel_degree")
+        if os.environ.get("OPTION_TENSOR_PARALLEL_DEGREE"):
+            user_tensor_parallel_degree = os.environ.get(
+                "OPTION_TENSOR_PARALLEL_DEGREE")
+
+        output_properties = self.properties_manager.properties
+        if user_tensor_parallel_degree:
+            logging.info(
+                f"User passed tensor_parallel_degree={user_tensor_parallel_degree}"
+            )
+            output_properties[
+                "option.tensor_parallel_degree"] = user_tensor_parallel_degree
+        else:
+            logging.info(
+                "User did not passs tensor_parallel_degree. Outputted serving.properties"
+                "will not include this field.")
+            del output_properties["option.tensor_parallel_degree"]
+
+        self.properties_manager.properties = output_properties
+        self.properties_manager.generate_properties_file()
+
     def neo_quantize(self):
         self.update_dataset_cache_location()
         self.initialize_partition_args_namespace()
         self.construct_properties_manager()
         self.run_quantization()
+        self.write_properties()
 
 
 def main():

diff --git a/serving/docker/partition/sm_neo_trt_llm_partition.py b/serving/docker/partition/sm_neo_trt_llm_partition.py
@@ -62,7 +62,8 @@ def main():
     compilation_error_file = None
     try:
         (compiler_options, input_model_directory, compiled_model_directory,
-         compilation_error_file, neo_cache_dir) = get_neo_env_vars()
+         compilation_error_file, neo_cache_dir,
+         neo_hf_cache_dir) = get_neo_env_vars()
 
         # Neo requires that serving.properties is in same dir as model files
         serving_properties = load_properties(input_model_directory)

diff --git a/serving/docker/partition/sm_neo_utils.py b/serving/docker/partition/sm_neo_utils.py
@@ -30,16 +30,19 @@ def write_error_to_file(error_message, error_file):
 def get_neo_env_vars():
     """
     Get environment variables required by the SageMaker Neo interface
+
+    TODO: Update the return type to a dictionary to allow for easier changes
     """
     try:
         compiler_options = os.environ.get("COMPILER_OPTIONS")
         input_model_directory = os.environ["SM_NEO_INPUT_MODEL_DIR"]
         compiled_model_directory = os.environ["SM_NEO_COMPILED_MODEL_DIR"]
         compilation_error_file = os.environ["SM_NEO_COMPILATION_ERROR_FILE"]
-        neo_cache_dir = os.environ["SM_NEO_CACHE_DIR"]
+        neo_cache_dir = os.environ.get("SM_NEO_CACHE_DIR")
+        neo_hf_cache_dir = os.environ.get("SM_NEO_HF_CACHE_DIR")
         return (compiler_options, input_model_directory,
                 compiled_model_directory, compilation_error_file,
-                neo_cache_dir)
+                neo_cache_dir, neo_hf_cache_dir)
     except KeyError as exc:
         raise InputConfiguration(
             f"SageMaker Neo environment variable '{exc.args[0]}' expected but not found"

diff --git a/serving/docker/pytorch-gpu.Dockerfile b/serving/docker/pytorch-gpu.Dockerfile
@@ -19,6 +19,7 @@ ARG torch_version=2.2.2
 ARG torch_vision_version=0.17.2
 ARG onnx_version=1.17.1
 ARG python_version=3.10
+ARG numpy_version=1.26.4
 
 RUN mkdir -p /opt/djl/conf && \
     mkdir -p /opt/ml/model
@@ -56,7 +57,7 @@ RUN chmod +x /usr/local/bin/dockerd-entrypoint.sh && \
     curl -o $(ls -d /usr/local/djl-serving-*/)lib/onnxruntime_gpu-$onnx_version.jar https://publish.djl.ai/onnxruntime/$onnx_version/onnxruntime_gpu-$onnx_version.jar && \
     scripts/install_python.sh ${python_version} && \
     scripts/install_s5cmd.sh x64 && \
-    pip3 install numpy && pip3 install torch==${torch_version} torchvision==${torch_vision_version} --extra-index-url https://download.pytorch.org/whl/cu121 && \
+    pip3 install numpy==${numpy_version} && pip3 install torch==${torch_version} torchvision==${torch_vision_version} --extra-index-url https://download.pytorch.org/whl/cu121 && \
     scripts/patch_oss_dlc.sh python && \
     scripts/security_patch.sh pytorch-gpu && \
     useradd -m -d /home/djl djl && \

diff --git a/serving/docker/scripts/install_python.sh b/serving/docker/scripts/install_python.sh
@@ -20,5 +20,5 @@ else
   rm -rf get-pip.py
 fi
 python3 -m pip --no-cache-dir install -U pip
-python3 -m pip --no-cache-dir install -U numpy awscli
+python3 -m pip --no-cache-dir install -U "numpy<2" awscli
 ln -sf /usr/bin/python3 /usr/bin/python