Merge pull request #10 from aws/master

Author: mufaddal-rohawala

CHANGELOG.md

@@ -1,5 +1,21 @@
 # Changelog
 
+## v2.63.0 (2021-10-13)
+
+### Features
+
+ * support configurable retry for pipeline steps
+
+## v2.62.0 (2021-10-12)
+
+### Features
+
+ * Hugging Face Transformers 4.10 for Pt1.8/TF2.4 & Transformers 4.11 for PT1.9&TF2.5
+
+### Bug Fixes and Other Changes
+
+ * repack_model script used in pipelines to support source_dir and dependencies
+
 ## v2.61.0 (2021-10-11)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.61.1.dev0
+2.63.1.dev0

src/sagemaker/workflow/_repack_model.py

@@ -34,15 +34,19 @@
 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):
+    """Repack custom dependencies and code into an existing model TAR archive
+
+    Args:
+        inference_script (str): The path to the custom entry point.
+        model_archive (str): The name of the model TAR archive.
+        dependencies (str): A space-delimited string of paths to custom dependencies.
+        source_dir (str): The path to a custom source directory.
+    """
 
     # 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 +55,53 @@
         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:
+            if os.path.exists(code_dir):
+                shutil.rmtree(code_dir)
+                shutil.copytree(source_dir, code_dir)
+
+        # 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")
+
+
+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

@@ -28,13 +28,14 @@
 from sagemaker.sklearn.estimator import SKLearn
 from sagemaker.workflow.entities import RequestType
 from sagemaker.workflow.properties import Properties
-from sagemaker.session import get_create_model_package_request
-from sagemaker.session import get_model_package_args
+from sagemaker.session import get_create_model_package_request, get_model_package_args
 from sagemaker.workflow.steps import (
     StepTypeEnum,
     TrainingStep,
     Step,
+    ConfigurableRetryStep,
 )
+from sagemaker.workflow.retry import RetryPolicy
 
 FRAMEWORK_VERSION = "0.23-1"
 INSTANCE_TYPE = "ml.m5.large"
@@ -60,6 +61,7 @@ def __init__(
         source_dir: str = None,
         dependencies: List = None,
         depends_on: Union[List[str], List[Step]] = None,
+        retry_policies: List[RetryPolicy] = None,
         subnets=None,
         security_group_ids=None,
         **kwargs,
@@ -126,6 +128,7 @@ def __init__(
                     This is not supported with "local code" in Local Mode.
             depends_on (List[str] or List[Step]): A list of step names or instances
                     this step depends on
+            retry_policies (List[RetryPolicy]): The list of retry policies for the current step
             subnets (list[str]): List of subnet ids. If not specified, the re-packing
                     job will be created without VPC config.
             security_group_ids (list[str]): List of security group ids. If not
@@ -145,6 +148,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 +165,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,
@@ -171,6 +181,7 @@ def __init__(
             display_name=display_name,
             description=description,
             depends_on=depends_on,
+            retry_policies=retry_policies,
             estimator=repacker,
             inputs=inputs,
         )
@@ -252,7 +263,7 @@ def properties(self):
         return self._properties
 
 
-class _RegisterModelStep(Step):
+class _RegisterModelStep(ConfigurableRetryStep):
     """Register model step in workflow that creates a model package.
 
     Attributes:
@@ -295,6 +306,7 @@ def __init__(
         display_name: str = None,
         description=None,
         depends_on: Union[List[str], List[Step]] = None,
+        retry_policies: List[RetryPolicy] = None,
         tags=None,
         container_def_list=None,
         **kwargs,
@@ -332,10 +344,11 @@ def __init__(
             description (str): Model Package description (default: None).
             depends_on (List[str] or List[Step]): A list of step names or instances
                 this step depends on
+            retry_policies (List[RetryPolicy]): The list of retry policies for the current step
             **kwargs: additional arguments to `create_model`.
         """
         super(_RegisterModelStep, self).__init__(
-            name, display_name, description, StepTypeEnum.REGISTER_MODEL, depends_on
+            name, StepTypeEnum.REGISTER_MODEL, display_name, description, depends_on, retry_policies
         )
         self.estimator = estimator
         self.model_data = model_data