feature: add edge packaging job support (#507)

mnuyens · ChoiByungWook · commit fe853569ffb7 · 2020-12-08T14:46:06.000-08:00
diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py
@@ -101,6 +101,8 @@ def __init__(
         self.sagemaker_session = sagemaker_session
         self.endpoint_name = None
         self._is_compiled_model = False
+        self._compilation_job_name = None
+        self._is_edge_packaged_model = False
         self._enable_network_isolation = enable_network_isolation
         self.model_kms_key = model_kms_key
 
@@ -336,6 +338,50 @@ def _get_framework_version(self):
         """Placeholder docstring"""
         return getattr(self, "framework_version", None)
 
+    def _edge_packaging_job_config(
+        self,
+        output_path,
+        role,
+        model_name,
+        model_version,
+        packaging_job_name,
+        compilation_job_name,
+        resource_key,
+        s3_kms_key,
+        tags,
+    ):
+        """Creates a request object for a packaging job.
+
+        Args:
+            output_path (str): where in S3 to store the output of the job
+            role (str): what role to use when executing the job
+            packaging_job_name (str): what to name the packaging job
+            compilation_job_name (str): what compilation job to source the model from
+            resource_key (str): the kms key to encrypt the disk with
+            s3_kms_key (str): the kms key to encrypt the output with
+            tags (list[dict]): List of tags for labeling an edge packaging job. For
+                more, see
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
+        Returns:
+            dict: the request object to use when creating a packaging job
+        """
+        output_model_config = {
+            "S3OutputLocation": output_path,
+        }
+        if s3_kms_key is not None:
+            output_model_config["KmsKeyId"] = s3_kms_key
+
+        return {
+            "output_model_config": output_model_config,
+            "role": role,
+            "tags": tags,
+            "model_name": model_name,
+            "model_version": model_version,
+            "job_name": packaging_job_name,
+            "compilation_job_name": compilation_job_name,
+            "resource_key": resource_key,
+        }
+
     def _compilation_job_config(
         self,
         target_instance_type,
@@ -438,6 +484,64 @@ def _compilation_image_uri(self, region, target_instance_type, framework, framew
             version=framework_version,
         )
 
+    def package_for_edge(
+        self,
+        output_path,
+        model_name,
+        model_version,
+        role=None,
+        job_name=None,
+        resource_key=None,
+        s3_kms_key=None,
+        tags=None,
+    ):
+        """Package this ``Model`` with SageMaker Edge.
+
+        Creates a new EdgePackagingJob and wait for it to finish.
+        model_data will now point to the packaged artifacts.
+
+        Args:
+            output_path (str): Specifies where to store the packaged model
+            role (str): Execution role
+            model_name (str): the name to attach to the model metadata
+            model_version (str): the version to attach to the model metadata
+            job_name (str): The name of the edge packaging job
+            resource_key (str): the kms key to encrypt the disk with
+            s3_kms_key (str): the kms key to encrypt the output with
+            tags (list[dict]): List of tags for labeling an edge packaging job. For
+                more, see
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
+
+        Returns:
+            sagemaker.model.Model: A SageMaker ``Model`` object. See
+            :func:`~sagemaker.model.Model` for full details.
+        """
+        if self._compilation_job_name is None:
+            raise ValueError("You must first compile this model")
+        if job_name is None:
+            job_name = f"packaging{self._compilation_job_name[11:]}"
+        if role is None:
+            role = self.sagemaker_session.expand_role(role)
+
+        self._init_sagemaker_session_if_does_not_exist(None)
+        config = self._edge_packaging_job_config(
+            output_path,
+            role,
+            model_name,
+            model_version,
+            job_name,
+            self._compilation_job_name,
+            resource_key,
+            s3_kms_key,
+            tags,
+        )
+        self.sagemaker_session.package_model_for_edge(**config)
+        job_status = self.sagemaker_session.wait_for_edge_packaging_job(job_name)
+        self.model_data = job_status["ModelArtifact"]
+        self._is_edge_packaged_model = True
+
+        return self
+
     def compile(
         self,
         target_instance_family,
@@ -557,6 +661,8 @@ def compile(
                 "supported for deployment via SageMaker. Please deploy the model manually."
             )
 
+        self._compilation_job_name = job_name
+
         return self
 
     def deploy(
diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py
@@ -1802,6 +1802,48 @@ def compile_model(
         LOGGER.info("Creating compilation-job with name: %s", job_name)
         self.sagemaker_client.create_compilation_job(**compilation_job_request)
 
+    def package_model_for_edge(
+        self,
+        output_model_config,
+        role,
+        job_name,
+        compilation_job_name,
+        model_name,
+        model_version,
+        resource_key,
+        tags,
+    ):
+        """Create an Amazon SageMaker Edge packaging job.
+
+        Args:
+            output_model_config (dict): Identifies the Amazon S3 location where you want Amazon
+                SageMaker Edge to save the results of edge packaging job
+            role (str): An AWS IAM role (either name or full ARN). The Amazon SageMaker Edge
+                edge packaging jobs use this role to access model artifacts. You must grant
+                sufficient permissions to this role.
+            job_name (str): Name of the edge packaging job being created.
+            compilation_job_name (str): Name of the compilation job being created.
+            resource_key (str): KMS key to encrypt the disk used to package the job
+            tags (list[dict]): List of tags for labeling a compile model job. For more, see
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
+        """
+        edge_packaging_job_request = {
+            "OutputConfig": output_model_config,
+            "RoleArn": role,
+            "ModelName": model_name,
+            "ModelVersion": model_version,
+            "EdgePackagingJobName": job_name,
+            "CompilationJobName": compilation_job_name,
+        }
+
+        if tags is not None:
+            edge_packaging_job_request["Tags"] = tags
+        if resource_key is not None:
+            edge_packaging_job_request["ResourceKey"] = (resource_key,)
+
+        LOGGER.info("Creating edge-packaging-job with name: %s", job_name)
+        self.sagemaker_client.create_edge_packaging_job(**edge_packaging_job_request)
+
     def tune(  # noqa: C901
         self,
         job_name,
@@ -3108,6 +3150,23 @@ def wait_for_compilation_job(self, job, poll=5):
         self._check_job_status(job, desc, "CompilationJobStatus")
         return desc
 
+    def wait_for_edge_packaging_job(self, job, poll=5):
+        """Wait for an Amazon SageMaker Edge packaging job to complete.
+
+        Args:
+            job (str): Name of the edge packaging job to wait for.
+            poll (int): Polling interval in seconds (default: 5).
+
+        Returns:
+            (dict): Return value from the ``DescribeEdgePackagingJob`` API.
+
+        Raises:
+            exceptions.UnexpectedStatusException: If the compilation job fails.
+        """
+        desc = _wait_until(lambda: _edge_packaging_job_status(self.sagemaker_client, job), poll)
+        self._check_job_status(job, desc, "EdgePackagingJobStatus")
+        return desc
+
     def wait_for_tuning_job(self, job, poll=5):
         """Wait for an Amazon SageMaker hyperparameter tuning job to complete.
 
@@ -4186,6 +4245,38 @@ def _processing_job_status(sagemaker_client, job_name):
     return desc
 
 
+def _edge_packaging_job_status(sagemaker_client, job_name):
+    """Process the current status of a packaging job
+
+    Args:
+        sagemaker_client (boto3.client.sagemaker): a sagemaker client
+        job_name (str): the name of the job to inspec
+
+    Returns:
+        Dict: the status of the edge packaging job
+    """
+    package_status_codes = {
+        "Completed": "!",
+        "InProgress": ".",
+        "Failed": "*",
+        "Stopped": "s",
+        "Stopping": "_",
+    }
+    in_progress_statuses = ["InProgress", "Stopping", "Starting"]
+
+    desc = sagemaker_client.describe_edge_packaging_job(EdgePackagingJobName=job_name)
+    status = desc["EdgePackagingJobStatus"]
+
+    status = _STATUS_CODE_TABLE.get(status, status)
+    print(package_status_codes.get(status, "?"), end="")
+    sys.stdout.flush()
+
+    if status in in_progress_statuses:
+        return None
+
+    return desc
+
+
 def _compilation_job_status(sagemaker_client, job_name):
     """Placeholder docstring"""
     compile_status_codes = {
diff --git a/tests/integ/test_edge.py b/tests/integ/test_edge.py
@@ -0,0 +1,69 @@
+# Copyright 2020-2020 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.
+from __future__ import absolute_import
+
+import os
+
+import pytest
+
+from sagemaker.mxnet.estimator import MXNet
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
+from tests.integ.timeout import timeout
+
+
+@pytest.fixture(scope="module")
+def mxnet_training_job(
+    sagemaker_session,
+    cpu_instance_type,
+    mxnet_training_latest_version,
+    mxnet_training_latest_py_version,
+):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+        script_path = os.path.join(DATA_DIR, "mxnet_mnist", "mnist_neo.py")
+        data_path = os.path.join(DATA_DIR, "mxnet_mnist")
+
+        mx = MXNet(
+            entry_point=script_path,
+            role="SageMakerRole",
+            framework_version=mxnet_training_latest_version,
+            py_version=mxnet_training_latest_py_version,
+            instance_count=1,
+            instance_type=cpu_instance_type,
+            sagemaker_session=sagemaker_session,
+        )
+
+        train_input = mx.sagemaker_session.upload_data(
+            path=os.path.join(data_path, "train"), key_prefix="integ-test-data/mxnet_mnist/train"
+        )
+        test_input = mx.sagemaker_session.upload_data(
+            path=os.path.join(data_path, "test"), key_prefix="integ-test-data/mxnet_mnist/test"
+        )
+
+        mx.fit({"train": train_input, "test": test_input})
+        return mx.latest_training_job.name
+
+
+def test_edge_packaging_job(mxnet_training_job, sagemaker_session):
+    estimator = MXNet.attach(mxnet_training_job, sagemaker_session=sagemaker_session)
+    model = estimator.compile_model(
+        target_instance_family="rasp3b",
+        input_shape={"data": [1, 1, 28, 28], "softmax_label": [1]},
+        output_path=estimator.output_path,
+    )
+
+    model.package_for_edge(
+        output_path=estimator.output_path,
+        role=estimator.role,
+        model_name="sdk-test-model",
+        model_version="1.0",
+    )
diff --git a/tests/unit/sagemaker/model/test_edge.py b/tests/unit/sagemaker/model/test_edge.py
@@ -0,0 +1,73 @@
+# Copyright 2020-2020 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.
+from __future__ import absolute_import
+
+import pytest
+from mock import Mock
+
+from sagemaker.model import Model
+
+MODEL_DATA = "s3://bucket/model.tar.gz"
+MODEL_IMAGE = "mi"
+
+REGION = "us-west-2"
+
+DESCRIBE_EDGE_PACKAGING_JOB_RESPONSE = {
+    "EdgePackagingJobStatus": "Completed",
+    "ModelArtifact": "s3://output-path/package-model.tar.gz",
+}
+
+
+@pytest.fixture
+def sagemaker_session():
+    return Mock(boto_region_name=REGION)
+
+
+def _create_model(sagemaker_session=None):
+    model = Model(MODEL_IMAGE, MODEL_DATA, role="role", sagemaker_session=sagemaker_session)
+    model._compilation_job_name = "compilation-test-name"
+    model._is_compiled_model = True
+    return model
+
+
+def test_package_model(sagemaker_session):
+    sagemaker_session.wait_for_edge_packaging_job = Mock(
+        return_value=DESCRIBE_EDGE_PACKAGING_JOB_RESPONSE
+    )
+    model = _create_model(sagemaker_session)
+    model.package_for_edge(
+        output_path="s3://output",
+        role="role",
+        model_name="model_name",
+        model_version="1.0",
+    )
+    assert model._is_edge_packaged_model is True
+
+
+def test_package_validates_compiled():
+    sagemaker_session.wait_for_edge_packaging_job = Mock(
+        return_value=DESCRIBE_EDGE_PACKAGING_JOB_RESPONSE
+    )
+    sagemaker_session.package_model_for_edge = Mock()
+    model = _create_model()
+    model._compilation_job_name = None
+
+    with pytest.raises(ValueError) as e:
+        model.package_for_edge(
+            output_path="s3://output",
+            role="role",
+            model_name="model_name",
+            model_version="1.0",
+        )
+
+    assert "You must first compile this model" in str(e)
diff --git a/tests/unit/test_mxnet.py b/tests/unit/test_mxnet.py
@@ -42,6 +42,7 @@
 ACCELERATOR_TYPE = "ml.eia.medium"
 IMAGE = "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.4.0-cpu-py3"
 COMPILATION_JOB_NAME = "{}-{}".format("compilation-sagemaker-mxnet", TIMESTAMP)
+EDGE_PACKAGING_JOB_NAME = "{}-{}".format("compilation-sagemaker-mxnet", TIMESTAMP)
 FRAMEWORK = "mxnet"
 ROLE = "Dummy"
 REGION = "us-west-2"
