From 8f6887280c44f14d6b546d79d79db50c8dc5eda1 Mon Sep 17 00:00:00 2001
From: Ta-Ching Chen <hello@tachingchen.com>
Date: Sat, 2 Nov 2019 23:51:31 +0800
Subject: [PATCH] Allow to set API type for tensorflow serving environment
 (#1371)

Tensorflow Serving supports three kinds of API: predict, classify, regress. In order to fully support possible API types, this PR adds code to split the entry point with separator ":" to get model name and the API type separately. And append the API type to proxy URL when receiving the requests.

$ fission fn create --name t1 --pkg <pkg name> \
    --env tensorflow --entrypoint "half_plus_two:regress"
---
 environments/tensorflow-serving/server.go | 24 +++++++++++++++++++++--
 examples/tensorflow-serving/README.md     |  5 +++--
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/environments/tensorflow-serving/server.go b/environments/tensorflow-serving/server.go
index d935476f..060f711d 100644
--- a/environments/tensorflow-serving/server.go
+++ b/environments/tensorflow-serving/server.go
@@ -28,6 +28,7 @@ var (
 
 	// for tensorflow serving to use
 	MODEL_NAME = ""
+	API_TYPE   = ""
 )
 
 type (
@@ -77,8 +78,26 @@ func specializeHandlerV2(logger *zap.Logger) func(http.ResponseWriter, *http.Req
 			return
 		}
 
+		// Tensorflow-serving supports three types of API: classify, regress, predict
+		// To get the API type, we need to split the entry point with separator ":"
+		// POST http://host:port/v1/models/${MODEL_NAME}:(classify|regress:predict)
+		entrypoint := strings.Split(loadreq.FunctionName, ":")
+		modelDir, apiType := "", ""
+
+		if len(entrypoint) == 0 {
+			logger.Error("unable to load model due to empty entrypoint")
+			w.WriteHeader(http.StatusBadRequest)
+			return
+		} else if len(entrypoint) == 1 {
+			modelDir = entrypoint[0]
+			apiType = "predict" // assign default API type
+		} else {
+			modelDir = entrypoint[0]
+			apiType = entrypoint[1]
+		}
+
 		// To ensure we load model from the expected path
-		basePath := fmt.Sprintf("%v/%v", loadreq.FilePath, loadreq.FunctionName)
+		basePath := fmt.Sprintf("%v/%v", loadreq.FilePath, modelDir)
 		basePath, err = filepath.Abs(basePath)
 		if err != nil {
 			msg := "error getting absolute path of model"
@@ -104,6 +123,7 @@ func specializeHandlerV2(logger *zap.Logger) func(http.ResponseWriter, *http.Req
 
 		// get directory name that holds model
 		MODEL_NAME = filepath.Base(basePath)
+		API_TYPE = apiType
 
 		argModelBasePath := fmt.Sprintf("--model_base_path=%v", basePath)
 		argModelName := fmt.Sprintf("--model_name=%v", MODEL_NAME)
@@ -191,7 +211,7 @@ func main() {
 		director := func(req *http.Request) {
 			req.URL.Scheme = "http"
 			req.URL.Host = "localhost:8501"
-			req.URL.Path = fmt.Sprintf("/v1/models/%v:predict", MODEL_NAME)
+			req.URL.Path = fmt.Sprintf("/v1/models/%v:%v", MODEL_NAME, API_TYPE)
 		}
 
 		proxy := &httputil.ReverseProxy{
diff --git a/examples/tensorflow-serving/README.md b/examples/tensorflow-serving/README.md
index 43bd1722..f4c2e7de 100644
--- a/examples/tensorflow-serving/README.md
+++ b/examples/tensorflow-serving/README.md
@@ -15,9 +15,10 @@ $ fission pkg create --env tensorflow --deploy half_plus_two.zip
 
 ## Create Function
 
-Here, the `--entrypoint` represents the name of top directory contains trained model.
+Here, the `--entrypoint` represents the name of the top directory contains the trained model and what kind of API the model supports.
+Currently, three API are supported: `predict`, `classify`, `regress`. `predict` API will be used if no API kind was given.
 
 ```bash
-$ fission fn create --name t1 --pkg <pkg name> --env tensorflow --entrypoint "half_plus_two"
+$ fission fn create --name t1 --pkg <pkg name> --env tensorflow --entrypoint "half_plus_two:predict"
 $ fission fn test --name t1 --body '{"instances": [1.0, 2.0, 0.0]}' --method POST
 ```