Turn off OpenMP multi-threading for python processes (#277)

* multiprocessing optimization

* ultiprocessing optimization with inference

* fix format issues

* add default env values in serving

* resolving comments

* add unit test

* Add negative test cases for set_default_env

* fix some typo

* split unit test into two testcases

* fixing Nvidia key error

Author: haixiw

docker/1.5-1/base/Dockerfile.cpu

@@ -23,7 +23,13 @@ ENV PYTHONDONTWRITEBYTECODE=1
 ENV PYTHONUNBUFFERED=1
 ENV PYTHONIOENCODING='utf-8'
 
-RUN apt-get update && \
+RUN  rm /etc/apt/sources.list.d/cuda.list && \
+        rm /etc/apt/sources.list.d/nvidia-ml.list && \
+        apt-key del 7fa2af80 && \
+        apt-get update && apt-get install -y --no-install-recommends wget && \
+        wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-keyring_1.0-1_all.deb && \
+        dpkg -i cuda-keyring_1.0-1_all.deb && \
+        apt-get update && \
     apt-get -y upgrade && \
     apt-get -y install --no-install-recommends \
         build-essential \

src/sagemaker_xgboost_container/constants/sm_env_constants.py

@@ -32,3 +32,6 @@
 SAGEMAKER_INFERENCE_OUTPUT = 'SAGEMAKER_INFERENCE_OUTPUT'
 SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT = 'SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT'
 SAGEMAKER_BATCH = 'SAGEMAKER_BATCH'
+
+# Multiprocessing related constants
+ONE_THREAD_PER_PROCESS = '1'

src/sagemaker_xgboost_container/serving.py

@@ -27,6 +27,7 @@
 from sagemaker_xgboost_container import encoder as xgb_encoders
 from sagemaker_xgboost_container.algorithm_mode import serve
 from sagemaker_xgboost_container.serving_mms import start_mxnet_model_server
+from sagemaker_xgboost_container.constants import sm_env_constants
 
 logging.basicConfig(
     format="%(asctime)s %(levelname)s - %(name)s - %(message)s", level=logging.INFO
@@ -43,6 +44,23 @@ def is_multi_model():
     return os.environ.get("SAGEMAKER_MULTI_MODEL")
 
 
+def set_default_serving_env_if_unspecified():
+    """Set default values for environment variables if they aren't already specified.
+
+    set "OMP_NUM_THREADS" = sm_env_constants.ONE_THREAD_PER_PROCESS
+    Single-thread processes by default. Multithreading can introduce significant
+    performance overhead due to task switching.
+    """
+    env_default_dict = {"OMP_NUM_THREADS": sm_env_constants.ONE_THREAD_PER_PROCESS}
+    for always_specified_key, default_value in env_default_dict.items():
+        try:
+            # If this does not throw, the user has specified a non-default value.
+            os.environ[always_specified_key]
+        except KeyError:
+            #  Key that is always specified is not set in the environment. Set default value.
+            os.environ[always_specified_key] = default_value
+
+
 def default_model_fn(model_dir):
     """Load a model. For XGBoost Framework, a default function to load a model is not provided.
     Users should provide customized model_fn() in script.
@@ -148,6 +166,8 @@ def serving_entrypoint():
     NOTE: If the inference server is multi-model, MxNet Model Server will be used as the base server. Otherwise,
         GUnicorn is used as the base server.
     """
+    set_default_serving_env_if_unspecified()
+
     if is_multi_model():
         start_mxnet_model_server()
     else:

test/unit/test_serving.py

@@ -20,6 +20,7 @@
 from sagemaker_algorithm_toolkit.exceptions import UserError
 from sagemaker_containers.beta.framework import (content_types, encoders, errors)
 from sagemaker_xgboost_container import serving
+from sagemaker_xgboost_container.constants import sm_env_constants
 
 TEST_CONFIG_FILE = "test_dir"
 
@@ -102,6 +103,23 @@ def test_serving_entrypoint_start_gunicorn(mock_server):
     mock_server.start.assert_called_once()
 
 
+@patch('sagemaker_xgboost_container.serving.server')
+@patch('sagemaker_xgboost_container.serving.set_default_serving_env_if_unspecified')
+def test_serving_entrypoint_set_default_env_positive(mock_set_default_serving_env_if_unspecified, mock_server):
+    serving.serving_entrypoint()
+    mock_set_default_serving_env_if_unspecified.assert_called_once()
+    assert os.getenv('OMP_NUM_THREADS') == sm_env_constants.ONE_THREAD_PER_PROCESS
+
+
+@patch('sagemaker_xgboost_container.serving.server')
+@patch('sagemaker_xgboost_container.serving.set_default_serving_env_if_unspecified')
+def test_serving_entrypoint_set_default_env_negative(mock_set_default_serving_env_if_unspecified, mock_server):
+    with patch.dict(os.environ, {"OMP_NUM_THREADS": "USER_SPECIFIED_VALUE"}, clear=True):
+        serving.serving_entrypoint()
+        mock_set_default_serving_env_if_unspecified.assert_called_once()
+        assert os.getenv('OMP_NUM_THREADS') == "USER_SPECIFIED_VALUE"
+
+
 @patch.dict(os.environ, {'SAGEMAKER_MULTI_MODEL': 'True', })
 @patch('sagemaker_xgboost_container.serving.start_mxnet_model_server')
 def test_serving_entrypoint_start_mms(mock_start_mxnet_model_server):