fix: fix loading user custom script

aws · Aug 30, 2022 · b53de53 · b53de53
1 parent 0cc05b2
commit b53de53
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/src/sagemaker_inference/default_handler_service.py b/src/sagemaker_inference/default_handler_service.py
@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 import os
+import sys
 
 from sagemaker_inference.transformer import Transformer
 
@@ -57,10 +58,7 @@ def initialize(self, context):
         model_dir = properties.get("model_dir")
 
         # add model_dir/code to python path
-        code_dir_path = "{}:".format(model_dir + "/code")
-        if PYTHON_PATH_ENV in os.environ:
-            os.environ[PYTHON_PATH_ENV] = code_dir_path + os.environ[PYTHON_PATH_ENV]
-        else:
-            os.environ[PYTHON_PATH_ENV] = code_dir_path
+        code_dir_path = model_dir + "/code"
+        sys.path.insert(1, code_dir_path)
 
         self._service.validate_and_initialize(model_dir=model_dir, context=context)