Fix: Don't load default model in MME mode (#130)

* Fix: Don't load default model in MME mode

* Remove whitespace from blank line

* Fix flake8 issue in other file

* Remove --model specification from MME test

* Trigger build

* Upgrade sagemaker and torch package

* Unpin sagemaker in tox

* Pin protobuf version

* Update protobuf version for python3.6

* Trigger build

* Revert dependency changes, keep only test dep changes

* Trigger build

* Revert version change when setting env variables

* Trigger build

* Trigger build

Co-authored-by: Nikhil Kulkarni <[email protected]>

Author: nikhil-sk

src/sagemaker_pytorch_serving_container/torchserve.py

@@ -90,10 +90,13 @@ def start_torchserve(handler_service=DEFAULT_HANDLER_SERVICE):
         TS_CONFIG_FILE,
         "--log-config",
         DEFAULT_TS_LOG_FILE,
-        "--models",
-        DEFAULT_TS_MODEL_NAME + "=" + environment.model_dir
     ]
 
+    default_model_path_args = ["--models", DEFAULT_TS_MODEL_NAME + "=" + environment.model_dir]
+
+    if not ENABLE_MULTI_MODEL:
+        ts_torchserve_cmd += default_model_path_args
+
     print(ts_torchserve_cmd)
 
     logger.info(ts_torchserve_cmd)

test/unit/test_model_server.py

@@ -104,8 +104,6 @@ def test_start_torchserve_default_service_handler_multi_model(
         torchserve.TS_CONFIG_FILE,
         "--log-config",
         torchserve.DEFAULT_TS_LOG_FILE,
-        "--models",
-        "model=/opt/ml/model"
     ]
 
     subprocess_popen.assert_called_once_with(ts_model_server_cmd)

test/utils/file_utils.py

@@ -19,7 +19,7 @@
 def make_tarfile(script, model, output_path, filename="model.tar.gz", script_path=None):
     output_filename = os.path.join(output_path, filename)
     with tarfile.open(output_filename, "w:gz") as tar:
-        if(script_path):
+        if (script_path):
             tar.add(script, arcname=os.path.join(script_path, os.path.basename(script)))
         else:
             tar.add(script, arcname=os.path.basename(script))

tox.ini

@@ -70,6 +70,7 @@ deps =
     six
     future
     pyyaml
+    protobuf == 3.19.6
 
 [testenv:flake8]
 basepython = python3