Add integration test for float with cog-runtime (#2486)

* Add integration test for float with cog-runtime

* Check if cog-runtime and cog combined can
receive a float input from the CLI

* Handle number in openapi schema explicitly

* Parse a number in the JSON to the inputs
* This fixes the CLI issue with cog-runtime

* Add support for integers

* Explicitly support integer

* Add more test cases

* Remove test for int with float

Author: 8W9aG

pkg/predict/input.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"strconv"
 	"strings"
 
 	"github.com/getkin/kin-openapi/openapi3"
@@ -20,6 +21,8 @@ type Input struct {
 	File   *string
 	Array  *[]any
 	Json   *json.RawMessage
+	Float  *float32
+	Int    *int32
 }
 
 type Inputs map[string]Input
@@ -51,11 +54,12 @@ func NewInputs(keyVals map[string][]string, schema *openapi3.T) (Inputs, error)
 					property, err := propertiesSchemas.JSONLookup(key)
 					if err == nil {
 						propertySchema := property.(*openapi3.Schema)
-						if propertySchema.Type.Is("object") {
+						switch {
+						case propertySchema.Type.Is("object"):
 							encodedVal := json.RawMessage(val)
 							input[key] = Input{Json: &encodedVal}
 							continue
-						} else if propertySchema.Type.Is("array") {
+						case propertySchema.Type.Is("array"):
 							var parsed any
 							err := json.Unmarshal([]byte(val), &parsed)
 							if err == nil {
@@ -69,6 +73,29 @@ func NewInputs(keyVals map[string][]string, schema *openapi3.T) (Inputs, error)
 							var arr = []any{val}
 							input[key] = Input{Array: &arr}
 							continue
+						case propertySchema.Type.Is("number"):
+							value, err := strconv.ParseInt(val, 10, 32)
+							if err == nil {
+								valueInt := int32(value)
+								input[key] = Input{Int: &valueInt}
+								continue
+							} else {
+								value, err := strconv.ParseFloat(val, 32)
+								if err != nil {
+									return input, err
+								}
+								float := float32(value)
+								input[key] = Input{Float: &float}
+								continue
+							}
+						case propertySchema.Type.Is("integer"):
+							value, err := strconv.ParseInt(val, 10, 32)
+							if err != nil {
+								return input, err
+							}
+							valueInt := int32(value)
+							input[key] = Input{Int: &valueInt}
+							continue
 						}
 					}
 				}
@@ -131,6 +158,10 @@ func (inputs *Inputs) toMap() (map[string]any, error) {
 			keyVals[key] = dataURLs
 		case input.Json != nil:
 			keyVals[key] = *input.Json
+		case input.Float != nil:
+			keyVals[key] = *input.Float
+		case input.Int != nil:
+			keyVals[key] = *input.Int
 		}
 	}
 	return keyVals, nil

test-integration/test_integration/fixtures/cog-runtime-float/cog.yaml

@@ -0,0 +1,4 @@
+build:
+  python_version: "3.11"
+  cog_runtime: true
+predict: "predict.py:Predictor"

test-integration/test_integration/fixtures/cog-runtime-float/predict.py

@@ -0,0 +1,8 @@
+from cog import BasePredictor, Input
+
+
+class Predictor(BasePredictor):
+    def predict(
+        self, num: float = Input(description="Number of things")
+    ) -> float:
+        return num * 2.0

test-integration/test_integration/fixtures/cog-runtime-int/cog.yaml

@@ -0,0 +1,4 @@
+build:
+  python_version: "3.11"
+  cog_runtime: true
+predict: "predict.py:Predictor"

test-integration/test_integration/fixtures/cog-runtime-int/predict.py

@@ -0,0 +1,8 @@
+from cog import BasePredictor, Input
+
+
+class Predictor(BasePredictor):
+    def predict(
+        self, num: int = Input(description="Number of things")
+    ) -> int:
+        return num * 2

test-integration/test_integration/test_predict.py

@@ -814,3 +814,55 @@ def test_predict_pipeline(cog_binary):
     )
     assert result.returncode == 0
     assert result.stdout == "HELLO TEST\n"
+
+
+def test_predict_cog_runtime_float(cog_binary):
+    project_dir = Path(__file__).parent / "fixtures/cog-runtime-float"
+    result = subprocess.run(
+        [cog_binary, "predict", "--debug", "-i", "num=10"],
+        cwd=project_dir,
+        capture_output=True,
+        text=True,
+        timeout=120.0,
+    )
+    assert result.returncode == 0
+    assert result.stdout == "20\n"
+
+
+def test_predict_cog_runtime_float_negative(cog_binary):
+    project_dir = Path(__file__).parent / "fixtures/cog-runtime-float"
+    result = subprocess.run(
+        [cog_binary, "predict", "--debug", "-i", "num=-10"],
+        cwd=project_dir,
+        capture_output=True,
+        text=True,
+        timeout=120.0,
+    )
+    assert result.returncode == 0
+    assert result.stdout == "-20\n"
+
+
+def test_predict_cog_runtime_int(cog_binary):
+    project_dir = Path(__file__).parent / "fixtures/cog-runtime-int"
+    result = subprocess.run(
+        [cog_binary, "predict", "--debug", "-i", "num=10"],
+        cwd=project_dir,
+        capture_output=True,
+        text=True,
+        timeout=120.0,
+    )
+    assert result.returncode == 0
+    assert result.stdout == "20\n"
+
+
+def test_predict_cog_runtime_int_negative(cog_binary):
+    project_dir = Path(__file__).parent / "fixtures/cog-runtime-int"
+    result = subprocess.run(
+        [cog_binary, "predict", "--debug", "-i", "num=-10"],
+        cwd=project_dir,
+        capture_output=True,
+        text=True,
+        timeout=120.0,
+    )
+    assert result.returncode == 0
+    assert result.stdout == "-20\n"