Fix tracker key bug (#378)

Author: vishalbollu

docs/deployments/apis.md

@@ -11,7 +11,7 @@ Serve models at scale.
   model_format: <string>  # model format, must be "tensorflow" or "onnx" (default: "onnx" if model path ends with .onnx, "tensorflow" if model path ends with .zip)
   request_handler: <string>  # path to the request handler implementation file, relative to the cortex root
   tracker:
-    key: <string> # json key to track in the response payload
+    key: <string> # json key to track if the response payload is a dictionary
     model_type: <string> # model type, must be "classification" or "regression"
   compute:
     min_replicas: <int>  # minimum number of replicas (default: 1)

pkg/operator/api/userconfig/apis.go

@@ -42,7 +42,7 @@ type API struct {
 }
 
 type Tracker struct {
-	Key       string    `json:"key" yaml:"key"`
+	Key       *string   `json:"key" yaml:"key"`
 	ModelType ModelType `json:"model_type" yaml:"model_type"`
 }
 
@@ -72,10 +72,8 @@ var apiValidation = &cr.StructValidation{
 				DefaultNil: true,
 				StructFieldValidations: []*cr.StructFieldValidation{
 					{
-						StructField: "Key",
-						StringValidation: &cr.StringValidation{
-							Required: true,
-						},
+						StructField:         "Key",
+						StringPtrValidation: &cr.StringPtrValidation{},
 					},
 					{
 						StructField: "ModelType",

pkg/workloads/cortex/lib/api_utils.py

@@ -73,23 +73,34 @@ def extract_predicted_values(api, predictions):
 
     tracker = api.get("tracker")
     for prediction in predictions:
-        predicted_value = prediction.get(tracker["key"])
-        if predicted_value is None:
-            raise ValueError(
-                "failed to track key '{}': not found in response payload".format(tracker["key"])
-            )
+        if tracker.get("key") is not None:
+            key = tracker["key"]
+            if type(prediction) != dict:
+                raise ValueError(
+                    "failed to track key '{}': expected prediction to be of type dict but found '{}'".format(
+                        key, type(prediction)
+                    )
+                )
+            if prediction.get(key) is None:
+                raise ValueError(
+                    "failed to track key '{}': not found in prediction".format(tracker["key"])
+                )
+            predicted_value = prediction[key]
+        else:
+            predicted_value = prediction
+
         if tracker["model_type"] == "classification":
             if type(predicted_value) != str and type(predicted_value) != int:
                 raise ValueError(
-                    "failed to track key '{}': expected type 'str' or 'int' but encountered '{}'".format(
-                        tracker["key"], type(predicted_value)
+                    "failed to track classification prediction: expected type 'str' or 'int' but encountered '{}'".format(
+                        type(predicted_value)
                     )
                 )
         else:
             if type(predicted_value) != float and type(predicted_value) != int:  # allow ints
                 raise ValueError(
-                    "failed to track key '{}': expected type 'float' or 'int' but encountered '{}'".format(
-                        tracker["key"], type(predicted_value)
+                    "failed to track regression prediction: expected type 'float' or 'int' but encountered '{}'".format(
+                        type(predicted_value)
                     )
                 )
         predicted_values.append(predicted_value)