Add warning message if user tries to compile a model for edge device (#659)

Author: wweic

src/sagemaker/model.py

@@ -19,6 +19,9 @@
 from sagemaker import fw_utils, local, session, utils
 from sagemaker.transformer import Transformer
 
+logging.basicConfig()
+LOGGER = logging.getLogger('sagemaker')
+LOGGER.setLevel(logging.INFO)
 
 NEO_ALLOWED_TARGET_INSTANCE_FAMILY = set(['ml_c5', 'ml_m5', 'ml_c4', 'ml_m4', 'jetson_tx1', 'jetson_tx2', 'ml_p2',
                                           'ml_p3', 'deeplens', 'rasp3b'])
@@ -190,9 +193,13 @@ def compile(self, target_instance_family, input_shape, output_path, role,
         self.sagemaker_session.compile_model(**config)
         job_status = self.sagemaker_session.wait_for_compilation_job(job_name)
         self.model_data = job_status['ModelArtifacts']['S3ModelArtifacts']
-        self.image = self._neo_image(self.sagemaker_session.boto_region_name, target_instance_family, framework,
-                                     framework_version)
-        self._is_compiled_model = True
+        if target_instance_family.startswith('ml_'):
+            self.image = self._neo_image(self.sagemaker_session.boto_region_name, target_instance_family, framework,
+                                         framework_version)
+            self._is_compiled_model = True
+        else:
+            LOGGER.warning("The intance type {} is not supported to deploy via SageMaker,"
+                           "please deploy the model on the device by yourself.".format(target_instance_family))
         return self
 
     def deploy(self, initial_instance_count, instance_type, accelerator_type=None, endpoint_name=None,

tests/unit/test_model.py

@@ -88,6 +88,13 @@
     'CertifyForMarketplace': False
 }
 
+DESCRIBE_COMPILATION_JOB_RESPONSE = {
+    'CompilationJobStatus': "Completed",
+    'ModelArtifacts': {
+        'S3ModelArtifacts': 's3://output-path/model.tar.gz'
+    }
+}
+
 
 class DummyFrameworkModel(FrameworkModel):
 
@@ -351,3 +358,21 @@ def test_delete_non_deployed_model(sagemaker_session):
     model = DummyFrameworkModel(sagemaker_session)
     with pytest.raises(ValueError, match='The SageMaker model must be created first before attempting to delete.'):
         model.delete_model()
+
+
+def test_compile_model_for_edge_device(sagemaker_session, tmpdir):
+    sagemaker_session.wait_for_compilation_job = Mock(
+        return_value=DESCRIBE_COMPILATION_JOB_RESPONSE)
+    model = DummyFrameworkModel(sagemaker_session, source_dir=str(tmpdir))
+    model.compile(target_instance_family='deeplens', input_shape={'data': [1, 3, 1024, 1024]},
+                  output_path='s3://output', role='role', framework='tensorflow', job_name="compile-model")
+    assert model._is_compiled_model is False
+
+
+def test_compile_model_for_cloud(sagemaker_session, tmpdir):
+    sagemaker_session.wait_for_compilation_job = Mock(
+        return_value=DESCRIBE_COMPILATION_JOB_RESPONSE)
+    model = DummyFrameworkModel(sagemaker_session, source_dir=str(tmpdir))
+    model.compile(target_instance_family='ml_c4', input_shape={'data': [1, 3, 1024, 1024]},
+                  output_path='s3://output', role='role', framework='tensorflow', job_name="compile-model")
+    assert model._is_compiled_model is True