fix: preserve inference script in model repack. (#1432)

Author: chuyang-deng

src/sagemaker/utils.py

@@ -534,10 +534,6 @@ def _create_or_update_code_dir(
         tmp:
     """
     code_dir = os.path.join(model_dir, "code")
-    if os.path.exists(code_dir):
-        for filename in os.listdir(code_dir):
-            if filename.endswith(".py"):
-                os.remove(os.path.join(code_dir, filename))
     if source_directory and source_directory.lower().startswith("s3://"):
         local_code_path = os.path.join(tmp, "local_code.tar.gz")
         download_file_from_url(source_directory, local_code_path, sagemaker_session)
@@ -552,7 +548,13 @@ def _create_or_update_code_dir(
     else:
         if not os.path.exists(code_dir):
             os.mkdir(code_dir)
-        shutil.copy2(inference_script, code_dir)
+        try:
+            shutil.copy2(inference_script, code_dir)
+        except FileNotFoundError:
+            if os.path.exists(os.path.join(code_dir, inference_script)):
+                pass
+            else:
+                raise
 
     for dependency in dependencies:
         lib_dir = os.path.join(code_dir, "lib")

tests/data/pytorch_mnist/packed_model.tar.gz

@@ -31,6 +31,7 @@
 
 MNIST_DIR = os.path.join(DATA_DIR, "pytorch_mnist")
 MNIST_SCRIPT = os.path.join(MNIST_DIR, "mnist.py")
+PACKED_MODEL = os.path.join(MNIST_DIR, "packed_model.tar.gz")
 
 EIA_DIR = os.path.join(DATA_DIR, "pytorch_eia")
 EIA_MODEL = os.path.join(EIA_DIR, "model_mnist.tar.gz")
@@ -124,6 +125,31 @@ def test_deploy_model(pytorch_training_job, sagemaker_session, cpu_instance_type
         assert output.shape == (batch_size, 10)
 
 
+@pytest.mark.skipif(
+    PYTHON_VERSION == "py2",
+    reason="Python 2 is supported by PyTorch {} and lower versions.".format(LATEST_PY2_VERSION),
+)
+def test_deploy_packed_model_with_entry_point_name(sagemaker_session, cpu_instance_type):
+    endpoint_name = "test-pytorch-deploy-model-{}".format(sagemaker_timestamp())
+
+    with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+        model_data = sagemaker_session.upload_data(path=PACKED_MODEL)
+        model = PyTorchModel(
+            model_data,
+            "SageMakerRole",
+            entry_point="mnist.py",
+            framework_version="1.4.0",
+            sagemaker_session=sagemaker_session,
+        )
+        predictor = model.deploy(1, cpu_instance_type, endpoint_name=endpoint_name)
+
+        batch_size = 100
+        data = numpy.random.rand(batch_size, 1, 28, 28).astype(numpy.float32)
+        output = predictor.predict(data)
+
+        assert output.shape == (batch_size, 10)
+
+
 @pytest.mark.skipif(PYTHON_VERSION == "py2", reason="PyTorch EIA does not support Python 2.")
 @pytest.mark.skipif(
     test_region() not in EI_SUPPORTED_REGIONS, reason="EI isn't supported in that specific region."

tests/integ/test_pytorch_train.py

@@ -643,6 +643,36 @@ def test_repack_model_with_inference_code_and_requirements(tmp, fake_s3):
     assert list_tar_files(fake_s3.fake_upload_path, tmp) == {
         "/code/requirements.txt",
         "/code/new-inference.py",
+        "/code/old-inference.py",
+        "/model",
+    }
+
+
+def test_repack_model_with_same_inference_file_name(tmp, fake_s3):
+    create_file_tree(
+        tmp,
+        [
+            "inference.py",
+            "model-dir/model",
+            "model-dir/code/inference.py",
+            "model-dir/code/requirements.txt",
+        ],
+    )
+
+    fake_s3.tar_and_upload("model-dir", "s3://fake/location")
+
+    sagemaker.utils.repack_model(
+        os.path.join(tmp, "inference.py"),
+        None,
+        None,
+        "s3://fake/location",
+        "s3://destination-bucket/repacked-model",
+        fake_s3.sagemaker_session,
+    )
+
+    assert list_tar_files(fake_s3.fake_upload_path, tmp) == {
+        "/code/requirements.txt",
+        "/code/inference.py",
         "/model",
     }