Use tensorflow/serving Docker images

cortexlabs · Aug 14, 2019 · 2e7a4ac · 2e7a4ac
1 parent 483fb13
commit 2e7a4ac
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 23 deletions.
diff --git a/images/tf-serve-gpu/Dockerfile b/images/tf-serve-gpu/Dockerfile
@@ -1,8 +1 @@
-FROM cortexlabs/tf-base-gpu
-
-ARG TF_SERV_VERSION="1.14.0"
-
-RUN curl -o tensorflow-model-server.deb http://storage.googleapis.com/tensorflow-serving-apt/pool/tensorflow-model-server-${TF_SERV_VERSION}/t/tensorflow-model-server/tensorflow-model-server_${TF_SERV_VERSION}_all.deb
-RUN dpkg -i tensorflow-model-server.deb
-
-ENTRYPOINT ["tensorflow_model_server"]
+FROM tensorflow/serving:1.14.0-gpu
diff --git a/images/tf-serve/Dockerfile b/images/tf-serve/Dockerfile
@@ -1,12 +1 @@
-FROM cortexlabs/tf-base
-
-ARG TF_SERV_VERSION="1.14.0"
-
-RUN apt-get update -qq && apt-get install -y -q \
-        curl \
-    && apt-get clean -qq && rm -rf /var/lib/apt/lists/*
-
-RUN curl -o tensorflow-model-server.deb http://storage.googleapis.com/tensorflow-serving-apt/pool/tensorflow-model-server-${TF_SERV_VERSION}/t/tensorflow-model-server/tensorflow-model-server_${TF_SERV_VERSION}_all.deb
-RUN dpkg -i tensorflow-model-server.deb
-
-ENTRYPOINT ["tensorflow_model_server"]
+FROM tensorflow/serving:1.14.0
diff --git a/pkg/workloads/cortex/tf_api/api.py b/pkg/workloads/cortex/tf_api/api.py
@@ -127,7 +127,7 @@ def create_prediction_request(transformed_sample):
     signature_def = local_cache["metadata"]["signatureDef"]
     signature_key = list(signature_def.keys())[0]
     prediction_request = predict_pb2.PredictRequest()
-    prediction_request.model_spec.name = "default"
+    prediction_request.model_spec.name = "model"
     prediction_request.model_spec.signature_name = signature_key
 
     for column_name, value in transformed_sample.items():
@@ -150,7 +150,7 @@ def create_raw_prediction_request(sample):
     signature_def = local_cache["metadata"]["signatureDef"]
     signature_key = list(signature_def.keys())[0]
     prediction_request = predict_pb2.PredictRequest()
-    prediction_request.model_spec.name = "default"
+    prediction_request.model_spec.name = "model"
     prediction_request.model_spec.signature_name = signature_key
 
     for column_name, value in sample.items():
@@ -252,7 +252,7 @@ def parse_response_proto(response_proto):
 
 def create_get_model_metadata_request():
     get_model_metadata_request = get_model_metadata_pb2.GetModelMetadataRequest()
-    get_model_metadata_request.model_spec.name = "default"
+    get_model_metadata_request.model_spec.name = "model"
     get_model_metadata_request.metadata_field.append("signature_def")
     return get_model_metadata_request