fix: fix _repack_model script used in pipelines not supporting dependencies or source_dir arguments, which are used for script-mode in models

Author: Payton Staub

src/sagemaker/workflow/_repack_model.py

@@ -33,16 +33,11 @@
 # repacking is some short-lived hackery, right??
 from distutils.dir_util import copy_tree
 
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--inference_script", type=str, default="inference.py")
-    parser.add_argument("--model_archive", type=str, default="model.tar.gz")
-    args = parser.parse_args()
+def repack(inference_script, model_archive, dependencies=None, source_dir=None):
 
     # the data directory contains a model archive generated by a previous training job
     data_directory = "/opt/ml/input/data/training"
-    model_path = os.path.join(data_directory, args.model_archive)
+    model_path = os.path.join(data_directory, model_archive)
 
     # create a temporary directory
     with tempfile.TemporaryDirectory() as tmp:
@@ -51,17 +46,55 @@
         shutil.copy2(model_path, local_path)
         src_dir = os.path.join(tmp, "src")
         # create the "code" directory which will contain the inference script
-        os.makedirs(os.path.join(src_dir, "code"))
+        code_dir = os.path.join(src_dir, "code")
+        os.makedirs(code_dir)
         # extract the contents of the previous training job's model archive to the "src"
         # directory of this training job
         with tarfile.open(name=local_path, mode="r:gz") as tf:
             tf.extractall(path=src_dir)
 
-        # generate a path to the custom inference script
-        entry_point = os.path.join("/opt/ml/code", args.inference_script)
-        # copy the custom inference script to the "src" dir
-        shutil.copy2(entry_point, os.path.join(src_dir, "code", args.inference_script))
+        # copy the custom inference script to code/ 
+        entry_point = os.path.join("/opt/ml/code", inference_script)
+        shutil.copy2(entry_point, os.path.join(src_dir, "code", inference_script))
+
+        # copy source_dir to code/
+        if source_dir:
+            actual_source_dir_path = os.path.join("/opt/ml/code", source_dir)
+            if os.path.exists(actual_source_dir_path):
+                for item in os.listdir(actual_source_dir_path):
+                    s = os.path.join(actual_source_dir_path, item)
+                    d = os.path.join(os.path.join(code_dir, source_dir), item)
+                    if os.path.isdir(s):
+                        shutil.copytree(s, d)
+                    else:
+                        shutil.copy2(s, d)
+
+        # copy any dependencies to code/lib/
+        if dependencies:
+            for dependency in dependencies.split(' '):
+                actual_dependency_path = os.path.join("/opt/ml/code", dependency)
+                lib_dir = os.path.join(code_dir, "lib")
+                if not os.path.exists(lib_dir):
+                        os.mkdir(lib_dir)
+                if os.path.isdir(actual_dependency_path):
+                    shutil.copytree(actual_dependency_path, os.path.join(lib_dir, os.path.basename(actual_dependency_path)))
+                else:
+                    shutil.copy2(actual_dependency_path, lib_dir)
 
         # copy the "src" dir, which includes the previous training job's model and the
         # custom inference script, to the output of this training job
-        copy_tree(src_dir, "/opt/ml/model")
+        copy_tree(src_dir, "/opt/ml/model")  
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--inference_script", type=str, default="inference.py")
+    parser.add_argument("--dependencies", type=str, default=None)
+    parser.add_argument("--source_dir", type=str, default=None)
+    parser.add_argument("--model_archive", type=str, default="model.tar.gz")
+    args, extra = parser.parse_known_args()
+    repack(
+        inference_script=args.inference_script,
+        dependencies=args.dependencies,
+        source_dir=args.source_dir,
+        model_archive=args.model_archive
+        )

src/sagemaker/workflow/_utils.py

@@ -145,6 +145,11 @@ def __init__(
         self._source_dir = source_dir
         self._dependencies = dependencies
 
+        # convert dependencies array into space-delimited string
+        dependencies_hyperparameter = None
+        if self._dependencies:
+            dependencies_hyperparameter = ' '.join(self._dependencies)
+
         # the real estimator and inputs
         repacker = SKLearn(
             framework_version=FRAMEWORK_VERSION,
@@ -157,6 +162,8 @@ def __init__(
             hyperparameters={
                 "inference_script": self._entry_point_basename,
                 "model_archive": self._model_archive,
+                "dependencies": dependencies_hyperparameter,
+                "source_dir": self._source_dir
             },
             subnets=subnets,
             security_group_ids=security_group_ids,

tests/unit/sagemaker/workflow/test_repack_model_script.py

@@ -0,0 +1,173 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+from sagemaker.workflow import _repack_model
+
+from pathlib import Path
+import shutil
+import tarfile
+import os
+import pytest
+import time
+
+@pytest.mark.skip(reason="""This test operates on the root file system 
+                            and will likely fail due to permission errors.
+                            Temporarily remove this skip decorator and run
+                            the test after making changes to _repack_model.py""")
+def test_repack_entry_point_only(tmp):
+    model_name = "xg-boost-model"
+    fake_model_path = os.path.join(tmp, model_name)
+
+    # create a fake model
+    open(fake_model_path, "w")
+
+    # create model.tar.gz
+    model_tar_name = "model-%s.tar.gz" % time.time()
+    model_tar_location = os.path.join(tmp, model_tar_name)
+    with tarfile.open(model_tar_location, mode="w:gz") as t:
+        t.add(fake_model_path, arcname=model_name)
+
+    # move model.tar.gz to /opt/ml/input/data/training
+    Path("/opt/ml/input/data/training").mkdir(parents=True, exist_ok=True)
+    shutil.move(model_tar_location, os.path.join("/opt/ml/input/data/training", model_tar_name))
+
+    # create files that will be added to model.tar.gz
+    create_file_tree(
+        "/opt/ml/code",
+        [
+            "inference.py",
+        ],
+    )
+
+    # repack
+    _repack_model.repack(
+        inference_script="inference.py",
+        model_archive=model_tar_name
+    )
+
+    # /opt/ml/model should now have the original model and the inference script
+    assert os.path.exists(os.path.join("/opt/ml/model", model_name))
+    assert os.path.exists(os.path.join("/opt/ml/model/code", "inference.py"))
+
+@pytest.mark.skip(reason="""This test operates on the root file system 
+                            and will likely fail due to permission errors.
+                            Temporarily remove this skip decorator and run
+                            the test after making changes to _repack_model.py""")
+def test_repack_with_dependencies(tmp):
+    model_name = "xg-boost-model"
+    fake_model_path = os.path.join(tmp, model_name)
+
+    # create a fake model
+    open(fake_model_path, "w")
+
+    # create model.tar.gz
+    model_tar_name = "model-%s.tar.gz" % time.time()
+    model_tar_location = os.path.join(tmp, model_tar_name)
+    with tarfile.open(model_tar_location, mode="w:gz") as t:
+        t.add(fake_model_path, arcname=model_name)
+
+    # move model.tar.gz to /opt/ml/input/data/training
+    Path("/opt/ml/input/data/training").mkdir(parents=True, exist_ok=True)
+    shutil.move(model_tar_location, os.path.join("/opt/ml/input/data/training", model_tar_name))
+
+    # create files that will be added to model.tar.gz
+    create_file_tree(
+        "/opt/ml/code",
+        [
+            "inference.py",
+            "dependencies/a",
+            "bb",
+            "dependencies/some/dir/b"
+        ],
+    )
+
+    # repack
+    _repack_model.repack(
+        inference_script="inference.py",
+        model_archive=model_tar_name,
+        dependencies=["dependencies/a", "bb", "dependencies/some/dir"]
+    )
+
+    # /opt/ml/model should now have the original model and the inference script
+    assert os.path.exists(os.path.join("/opt/ml/model", model_name))
+    assert os.path.exists(os.path.join("/opt/ml/model/code", "inference.py"))
+    assert os.path.exists(os.path.join("/opt/ml/model/code/lib", "a"))
+    assert os.path.exists(os.path.join("/opt/ml/model/code/lib", "bb"))
+    assert os.path.exists(os.path.join("/opt/ml/model/code/lib/dir", "b"))
+
+@pytest.mark.skip(reason="""This test operates on the root file system 
+                            and will likely fail due to permission errors.
+                            Temporarily remove this skip decorator and run
+                            the test after making changes to _repack_model.py""")
+def test_repack_with_source_dir_and_dependencies(tmp):
+    model_name = "xg-boost-model"
+    fake_model_path = os.path.join(tmp, model_name)
+
+    # create a fake model
+    open(fake_model_path, "w")
+
+    # create model.tar.gz
+    model_tar_name = "model-%s.tar.gz" % time.time()
+    model_tar_location = os.path.join(tmp, model_tar_name)
+    with tarfile.open(model_tar_location, mode="w:gz") as t:
+        t.add(fake_model_path, arcname=model_name)
+
+    # move model.tar.gz to /opt/ml/input/data/training
+    Path("/opt/ml/input/data/training").mkdir(parents=True, exist_ok=True)
+    shutil.move(model_tar_location, os.path.join("/opt/ml/input/data/training", model_tar_name))
+
+    # create files that will be added to model.tar.gz
+    create_file_tree(
+        "/opt/ml/code",
+        [
+            "inference.py",
+            "dependencies/a",
+            "bb",
+            "dependencies/some/dir/b",
+            "sourcedir/foo.py",
+            "sourcedir/some/dir/a"
+        ],
+    )
+
+    # repack
+    _repack_model.repack(
+        inference_script="inference.py",
+        model_archive=model_tar_name,
+        dependencies=["dependencies/a", "bb", "dependencies/some/dir"],
+        source_dir="sourcedir"
+    )
+
+    # /opt/ml/model should now have the original model and the inference script
+    assert os.path.exists(os.path.join("/opt/ml/model", model_name))
+    assert os.path.exists(os.path.join("/opt/ml/model/code", "inference.py"))
+    assert os.path.exists(os.path.join("/opt/ml/model/code/lib", "a"))
+    assert os.path.exists(os.path.join("/opt/ml/model/code/lib", "bb"))
+    assert os.path.exists(os.path.join("/opt/ml/model/code/lib/dir", "b"))
+    assert os.path.exists(os.path.join("/opt/ml/model/code/sourcedir", "foo.py"))
+    assert os.path.exists(os.path.join("/opt/ml/model/code/sourcedir/some/dir", "a"))
+
+
+def create_file_tree(root, tree):
+    for file in tree:
+        try:
+            os.makedirs(os.path.join(root, os.path.dirname(file)))
+        except:  # noqa: E722 Using bare except because p2/3 incompatibility issues.
+            pass
+        with open(os.path.join(root, file), "a") as f:
+            f.write(file)
+
+
+@pytest.fixture()
+def tmp(tmpdir):
+    yield str(tmpdir)