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

Co-authored-by: Andrew Song <40076917+a-ys@users.noreply.github.com> Co-authored-by: Siddharth Venkatesan <siddvenk@umich.edu>
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