feature: add model_version_policy to model config

Author: laurenyu

docker/build_artifacts/sagemaker/serve.py

@@ -99,9 +99,17 @@ def _create_tfs_config(self):
         config = "model_config_list: {\n"
         for m in models:
             config += "  config: {\n"
-            config += "    name: '{}',\n".format(os.path.basename(m))
-            config += "    base_path: '{}',\n".format(m)
+            config += "    name: '{}'\n".format(os.path.basename(m))
+            config += "    base_path: '{}'\n".format(m)
             config += "    model_platform: 'tensorflow'\n"
+
+            config += "    model_version_policy: {\n"
+            config += "      specific: {\n"
+            for version in tfs_utils.find_model_versions(m):
+                config += "        versions: {}\n".format(version)
+            config += "      }\n"
+            config += "    }\n"
+
             config += "  }\n"
         config += "}\n"
 

docker/build_artifacts/sagemaker/tfs_utils.py

@@ -142,6 +142,10 @@ def find_models():
         return models
 
 
+def find_model_versions(model_path):
+    return [version.lstrip("0") for version in os.listdir(model_path)]
+
+
 def _find_saved_model_files(path):
     for e in os.scandir(path):
         if e.is_dir():

test/integration/local/test_container.py

@@ -71,11 +71,14 @@ def container(request, docker_base_name, tag, runtime_config):
         subprocess.check_call("docker rm -f sagemaker-tensorflow-serving-test".split())
 
 
-def make_request(data, content_type="application/json", method="predict"):
+def make_request(data, content_type="application/json", method="predict", version=None):
+    custom_attributes = "tfs-model-name=half_plus_three,tfs-method={}".format(method)
+    if version:
+        custom_attributes += ",tfs-model-version={}".format(version)
+
     headers = {
         "Content-Type": content_type,
-        "X-Amzn-SageMaker-Custom-Attributes":
-            "tfs-model-name=half_plus_three,tfs-method=%s" % method
+        "X-Amzn-SageMaker-Custom-Attributes": custom_attributes,
     }
     response = requests.post(BASE_URL, data=data, headers=headers)
     return json.loads(response.content.decode("utf-8"))
@@ -101,6 +104,18 @@ def test_predict_twice():
     assert z == {"predictions": [3.5, 4.0, 5.5]}
 
 
+def test_predict_specific_versions():
+    x = {
+        "instances": [1.0, 2.0, 5.0]
+    }
+
+    y = make_request(json.dumps(x), version=123)
+    assert y == {"predictions": [3.5, 4.0, 5.5]}
+
+    y = make_request(json.dumps(x), version=124)
+    assert y == {"predictions": [3.5, 4.0, 5.5]}
+
+
 def test_predict_two_instances():
     x = {
         "instances": [[1.0, 2.0, 5.0], [1.0, 2.0, 5.0]]

test/resources/models/half_plus_three/00000124/assets/foo.txt

@@ -0,0 +1 @@
+asset-file-contents