Fix env module usage from sagemaker-inference-toolkit

Author: dk19y

src/sagemaker_pytorch_serving_container/torchserve.py

@@ -155,8 +155,8 @@ def _generate_ts_config_properties():
     env = environment.Environment()
 
     user_defined_configuration = {
-        "default_response_timeout": env.torchserve_timeout,
-        "default_workers_per_model": env.torchserve_workers,
+        "default_response_timeout": env.model_server_timeout,
+        "default_workers_per_model": env.model_server_workers,
         "inference_address": "http://0.0.0.0:{}".format(env.inference_http_port),
         "management_address": "http://0.0.0.0:{}".format(env.management_http_port),
     }

test/unit/test_model_server.py

@@ -205,18 +205,18 @@ def test_create_torchserve_config_file(write_file, generate_ts_config_props):
 @patch("sagemaker_inference.utils.read_file", return_value=DEFAULT_CONFIGURATION)
 @patch("sagemaker_inference.environment.Environment")
 def test_generate_ts_config_properties(env, read_file):
-    torchserve_timeout = "torchserve_timeout"
+    model_sever_timeout = "torchserve_timeout"
     torchserve_workers = "torchserve_workers"
     http_port = "http_port"
 
-    env.return_value.torchserve_timeout = torchserve_timeout
+    env.return_value.model_sever_timeout = torchserve_timeout
     env.return_value.torchserve_workers = torchserve_workers
     env.return_value.inference_http_port = http_port
 
     ts_config_properties = torchserve._generate_ts_config_properties()
 
     inference_address = "inference_address=http://0.0.0.0:{}\n".format(http_port)
-    server_timeout = "default_response_timeout={}\n".format(torchserve_timeout)
+    server_timeout = "default_response_timeout={}\n".format(model_sever_timeout)
     workers = "default_workers_per_model={}\n".format(torchserve_workers)
 
     read_file.assert_called_once_with(torchserve.DEFAULT_TS_CONFIG_FILE)