feature: upgrade Neo MxNet to 1.7

Author: Rui Wang Napieralski

src/sagemaker/image_uri_config/neo-mxnet.json

@@ -2,20 +2,20 @@
     "processors": ["cpu", "gpu"],
     "scope": ["inference"],
     "version_aliases": {
-        "0.12.1": "1.5",
-        "1.0.0": "1.5",
-        "1.1.0": "1.5",
-        "1.2": "1.5",
-        "1.2.0": "1.5",
-        "1.2.1": "1.5",
-        "1.3": "1.5",
-        "1.3.0": "1.5",
-        "1.4": "1.5",
-        "1.4.0": "1.5",
-        "1.4.1": "1.5"
+        "0.12.1": "1.7",
+        "1.0.0": "1.7",
+        "1.1.0": "1.7",
+        "1.2": "1.7",
+        "1.2.0": "1.7",
+        "1.2.1": "1.7",
+        "1.3": "1.7",
+        "1.3.0": "1.7",
+        "1.4": "1.7",
+        "1.4.0": "1.7",
+        "1.4.1": "1.7"
     },
     "versions": {
-        "1.5": {
+        "1.7": {
             "py_versions": ["py3"],
             "registries": {
                 "af-south-1": "774647643957",
@@ -42,7 +42,7 @@
                 "us-west-1": "710691900526",
                 "us-west-2": "301217895009"
             },
-            "repository": "sagemaker-neo-mxnet"
+            "repository": "sagemaker-inference-mxnet"
         }
     }
 }

tests/data/mxnet_mnist/mnist_neo.py

@@ -19,6 +19,7 @@
 
 import mxnet as mx
 import numpy as np
+import neomxnet  # noqa: F401
 
 
 def load_data(path):
@@ -104,27 +105,28 @@ def train(
         save(model_dir, mlp_model)
 
 
-def neo_preprocess(payload, content_type):
-    logging.info("Invoking user-defined pre-processing function")
-
-    if content_type != "application/vnd+python.numpy+binary":
-        raise RuntimeError("Content type must be application/vnd+python.numpy+binary")
-
-    return np.asarray(json.loads(payload.decode("utf-8")))
+def model_fn(path_to_model_files):
+    ctx = mx.cpu()
+    sym, arg_params, aux_params = mx.model.load_checkpoint(
+        os.path.join(path_to_model_files, "compiled"), 0
+    )
+    mod = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
+    mod.bind(
+        for_training=False, data_shapes=[("data", (1, 1, 28, 28))], label_shapes=mod._label_shapes
+    )
+    mod.set_params(arg_params, aux_params, allow_missing=True)
+    return mod
 
 
-# NOTE: this function cannot use MXNet
-def neo_postprocess(result):
-    logging.info("Invoking user-defined post-processing function")
+def transform_fn(mod, payload, input_content_type, requested_output_content_type):
+    if input_content_type != "application/vnd+python.numpy+binary":
+        raise RuntimeError("Input content type must be application/vnd+python.numpy+binary")
 
-    # Softmax (assumes batch size 1)
+    inference_payload = np.asarray(json.loads(payload.decode("utf-8")))
+    result = mod.predict(inference_payload)
     result = np.squeeze(result)
-    result_exp = np.exp(result - np.max(result))
-    result = result_exp / np.sum(result_exp)
-
-    response_body = json.dumps(result.tolist())
+    response_body = json.dumps(result.asnumpy().tolist())
     content_type = "application/json"
-
     return response_body, content_type
 
 
@@ -135,7 +137,7 @@ def neo_postprocess(result):
     parser = argparse.ArgumentParser()
 
     parser.add_argument("--batch-size", type=int, default=100)
-    parser.add_argument("--epochs", type=int, default=10)
+    parser.add_argument("--epochs", type=int, default=1)
     parser.add_argument("--learning-rate", type=float, default=0.1)
 
     parser.add_argument("--model-dir", type=str, default=os.environ["SM_MODEL_DIR"])

tests/integ/test_neo_mxnet.py

@@ -58,9 +58,6 @@ def mxnet_training_job(
 
 
 @pytest.mark.canary_quick
-@pytest.mark.skip(
-    reason="This test is failing because the image uri and the training script format has changed."
-)
 def test_attach_deploy(
     mxnet_training_job, sagemaker_session, cpu_instance_type, cpu_instance_family
 ):
@@ -71,7 +68,7 @@ def test_attach_deploy(
 
         estimator.compile_model(
             target_instance_family=cpu_instance_family,
-            input_shape={"data": [1, 1, 28, 28]},
+            input_shape={"data": [1, 1, 28, 28], "softmax_label": [1]},
             output_path=estimator.output_path,
         )
 
@@ -89,9 +86,6 @@ def test_attach_deploy(
         predictor.predict(data)
 
 
-@pytest.mark.skip(
-    reason="This test is failing because the image uri and the training script format has changed."
-)
 def test_deploy_model(
     mxnet_training_job,
     sagemaker_session,
@@ -123,7 +117,7 @@ def test_deploy_model(
 
         model.compile(
             target_instance_family=cpu_instance_family,
-            input_shape={"data": [1, 1, 28, 28]},
+            input_shape={"data": [1, 1, 28, 28], "softmax_label": [1]},
             role=role,
             job_name=unique_name_from_base("test-deploy-model-compilation-job"),
             output_path="/".join(model_data.split("/")[:-1]),
@@ -165,7 +159,7 @@ def test_inferentia_deploy_model(
 
         model.compile(
             target_instance_family=inf_instance_family,
-            input_shape={"data": [1, 1, 28, 28]},
+            input_shape={"data": [1, 1, 28, 28], "softmax_label": [1]},
             role=role,
             job_name=unique_name_from_base("test-deploy-model-compilation-job"),
             output_path="/".join(model_data.split("/")[:-1]),