change: add tflite to Neo-supported frameworks (#1364)

Author: ashishgupta023

src/sagemaker/model.py

@@ -24,7 +24,9 @@
 
 LOGGER = logging.getLogger("sagemaker")
 
-NEO_ALLOWED_FRAMEWORKS = set(["mxnet", "tensorflow", "keras", "pytorch", "onnx", "xgboost"])
+NEO_ALLOWED_FRAMEWORKS = set(
+    ["mxnet", "tensorflow", "keras", "pytorch", "onnx", "xgboost", "tflite"]
+)
 
 NEO_IMAGE_ACCOUNT = {
     "us-west-1": "710691900526",

tests/unit/test_model.py

@@ -560,6 +560,22 @@ def test_compile_model_for_edge_device(sagemaker_session, tmpdir):
     assert model._is_compiled_model is False
 
 
+def test_compile_model_for_edge_device_tflite(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="tflite",
+        job_name="tflite-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
@@ -576,6 +592,22 @@ def test_compile_model_for_cloud(sagemaker_session, tmpdir):
     assert model._is_compiled_model is True
 
 
+def test_compile_model_for_cloud_tflite(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="tflite",
+        job_name="tflite-compile-model",
+    )
+    assert model._is_compiled_model is True
+
+
 @patch("sagemaker.session.Session")
 @patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
 def test_compile_creates_session(session):