feat: handle oom

assets/huggingface-vit-template/huggingface-infer/config.pbtxt

@@ -1,6 +1,6 @@
 name: "huggingface-infer"
 platform: "onnxruntime_onnx"
-max_batch_size: 32
+max_batch_size: 16
 input [
   {
     name: "pixel_values"

assets/huggingface-vit-template/huggingface/config.pbtxt

@@ -1,6 +1,6 @@
 name: "huggingface"
 platform: "ensemble"
-max_batch_size: 32
+max_batch_size: 16
 input [
   {
     name: "input"

assets/huggingface-vit-template/pre/config.pbtxt

@@ -1,6 +1,6 @@
 name: "pre"
 backend: "python"
-max_batch_size: 32
+max_batch_size: 16
 input [
   {
     name: "input"

cmd/main/misc.go

@@ -120,6 +120,8 @@ func errorHandler(ctx context.Context, mux *runtime.ServeMux, marshaler runtime.
 				switch v.Violations[0].Type {
 				case "UPDATE", "DELETE", "STATE", "RENAME":
 					httpStatus = http.StatusUnprocessableEntity
+				case "MAX BATCH SIZE LIMITATION":
+					httpStatus = http.StatusBadRequest
 				}
 			default:
 				httpStatus = runtime.HTTPStatusFromCode(s.Code())

config/config.go

@@ -80,15 +80,24 @@ type PipelineBackendConfig struct {
 	}
 }
 
+type MaxBatchSizeConfig struct {
+	Unspecified    int `koanf:"unspecified"`
+	Classification int `koanf:"classification"`
+	Detection      int `koanf:"detection"`
+	Keypoint       int `koanf:"keypoint"`
+	Ocr            int `koanf:"ocr"`
+}
+
 // AppConfig defines
 type AppConfig struct {
-	Server          ServerConfig          `koanf:"server"`
-	Database        DatabaseConfig        `koanf:"database"`
-	TritonServer    TritonServerConfig    `koanf:"tritonserver"`
-	MgmtBackend     MgmtBackendConfig     `koanf:"mgmtbackend"`
-	Cache           CacheConfig           `koanf:"cache"`
-	UsageBackend    UsageBackendConfig    `koanf:"usagebackend"`
-	PipelineBackend PipelineBackendConfig `koanf:"pipelinebackend"`
+	Server                 ServerConfig          `koanf:"server"`
+	Database               DatabaseConfig        `koanf:"database"`
+	TritonServer           TritonServerConfig    `koanf:"tritonserver"`
+	MgmtBackend            MgmtBackendConfig     `koanf:"mgmtbackend"`
+	Cache                  CacheConfig           `koanf:"cache"`
+	UsageBackend           UsageBackendConfig    `koanf:"usagebackend"`
+	PipelineBackend        PipelineBackendConfig `koanf:"pipelinebackend"`
+	MaxBatchSizeLimitation MaxBatchSizeConfig    `koanf:"maxbatchsizelimitation"`
 }
 
 // Config - Global variable to export

config/config.yaml

@@ -45,3 +45,9 @@ pipelinebackend:
   https:
     cert:   # /ssl/tls.crt
     key:    # /ssl/tls.key
+maxbatchsizelimitation:
+  unspecified: 2
+  classification: 16 
+  detection: 8
+  keypoint: 8
+  ocr: 2

integration-test/rest.js

@@ -44,7 +44,7 @@ export default function (data) {
     });
   }
 
-  // Create Model API
+  // // Create Model API
   createModel.CreateModelFromLocal()
   createModel.CreateModelFromGitHub()
 

integration-test/rest_create_model.js

@@ -13,6 +13,10 @@ const cls_model = open(`${__ENV.TEST_FOLDER_ABS_PATH}/integration-test/data/dumm
 const det_model = open(`${__ENV.TEST_FOLDER_ABS_PATH}/integration-test/data/dummy-det-model.zip`, "b");
 const keypoint_model = open(`${__ENV.TEST_FOLDER_ABS_PATH}/integration-test/data/dummy-keypoint-model.zip`, "b");
 const unspecified_model = open(`${__ENV.TEST_FOLDER_ABS_PATH}/integration-test/data/dummy-unspecified-model.zip`, "b");
+const cls_model_bz17 = open(`${__ENV.TEST_FOLDER_ABS_PATH}/integration-test/data/dummy-cls-model-bz17.zip`, "b");
+const det_model_bz9 = open(`${__ENV.TEST_FOLDER_ABS_PATH}/integration-test/data/dummy-det-model-bz9.zip`, "b");
+const keypoint_model_bz9 = open(`${__ENV.TEST_FOLDER_ABS_PATH}/integration-test/data/dummy-keypoint-model-bz9.zip`, "b");
+const unspecified_model_bz3 = open(`${__ENV.TEST_FOLDER_ABS_PATH}/integration-test/data/dummy-unspecified-model-bz3.zip`, "b");
 
 
 export function CreateModelFromLocal() {
@@ -188,6 +192,90 @@ export function CreateModelFromLocal() {
           r.status === 204
       });
     });
+
+    group("Model Backend API: Upload a model which exceed max batch size limitation", function () {
+      let fd_cls = new FormData();
+      let model_id_cls = randomString(10)
+      let model_description = randomString(20)
+      fd_cls.append("id", model_id_cls);
+      fd_cls.append("description", model_description);
+      fd_cls.append("model_definition", "model-definitions/local");
+      fd_cls.append("content", http.file(cls_model_bz17, "dummy-cls-model-bz17.zip"));
+      check(http.request("POST", `${apiHost}/v1alpha/models:multipart`, fd_cls.body(), {
+        headers: genHeader(`multipart/form-data; boundary=${fd_cls.boundary}`),
+      }), {
+        "POST /v1alpha/models:multipart task cls response status": (r) =>
+          r.status === 400,
+      });
+
+      let fd_det = new FormData();
+      let model_id_det = randomString(10)
+      model_description = randomString(20)
+      fd_det.append("id", model_id_det);
+      fd_det.append("description", model_description);
+      fd_det.append("model_definition", "model-definitions/local");
+      fd_det.append("content", http.file(det_model_bz9, "dummy-det-model-bz9.zip"));
+      check(http.request("POST", `${apiHost}/v1alpha/models:multipart`, fd_det.body(), {
+        headers: genHeader(`multipart/form-data; boundary=${fd_det.boundary}`),
+      }), {
+        "POST /v1alpha/models:multipart task det response status": (r) =>
+          r.status === 400,
+      });
+
+      let fd_keypoint = new FormData();
+      let model_id_keypoint = randomString(10)
+      model_description = randomString(20)
+      fd_keypoint.append("id", model_id_keypoint);
+      fd_keypoint.append("description", model_description);
+      fd_keypoint.append("model_definition", "model-definitions/local");
+      fd_keypoint.append("content", http.file(keypoint_model_bz9, "dummy-keypoint-model-bz9.zip"));
+      check(http.request("POST", `${apiHost}/v1alpha/models:multipart`, fd_keypoint.body(), {
+        headers: genHeader(`multipart/form-data; boundary=${fd_keypoint.boundary}`),
+      }), {
+        "POST /v1alpha/models:multipart task keypoint response status": (r) =>
+          r.status === 400,
+      });
+
+      let fd_unspecified = new FormData();
+      let model_id_unspecified = randomString(10)
+      model_description = randomString(20)
+      fd_unspecified.append("id", model_id_unspecified);
+      fd_unspecified.append("description", model_description);
+      fd_unspecified.append("model_definition", "model-definitions/local");
+      fd_unspecified.append("content", http.file(unspecified_model_bz3, "dummy-unspecified-model-bz3.zip"));
+      check(http.request("POST", `${apiHost}/v1alpha/models:multipart`, fd_unspecified.body(), {
+        headers: genHeader(`multipart/form-data; boundary=${fd_unspecified.boundary}`),
+      }), {
+        "POST /v1alpha/models:multipart task unspecified response status": (r) =>
+          r.status === 400,
+      });
+
+      // clean up
+      check(http.request("DELETE", `${apiHost}/v1alpha/models/${model_id_cls}`, null, {
+        headers: genHeader(`application/json`),
+      }), {
+        "DELETE clean up response status": (r) =>
+          r.status === 404
+      });
+      check(http.request("DELETE", `${apiHost}/v1alpha/models/${model_id_det}`, null, {
+        headers: genHeader(`application/json`),
+      }), {
+        "DELETE clean up response status": (r) =>
+          r.status === 404
+      });
+      check(http.request("DELETE", `${apiHost}/v1alpha/models/${model_id_keypoint}`, null, {
+        headers: genHeader(`application/json`),
+      }), {
+        "DELETE clean up response status": (r) =>
+          r.status === 404
+      });
+      check(http.request("DELETE", `${apiHost}/v1alpha/models/${model_id_unspecified}`, null, {
+        headers: genHeader(`application/json`),
+      }), {
+        "DELETE clean up response status": (r) =>
+          r.status === 404
+      });
+    });    
   }
 }
 

internal/triton/triton.go

@@ -497,8 +497,8 @@ func (ts *triton) PostProcess(inferResponse *inferenceserver.ModelInferResponse,
 }
 
 func (ts *triton) LoadModelRequest(modelName string) (*inferenceserver.RepositoryModelLoadResponse, error) {
-	// Create context for our request with 60 second timeout
-	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	// Create context for our request with 600 second timeout. The time for warmup model inference
+	ctx, cancel := context.WithTimeout(context.Background(), 600*time.Second)
 	defer cancel()
 
 	// Create status request for a given model

internal/util/util.go

@@ -313,6 +313,35 @@ func ConvertAllJSONEnumValueToProtoStyle(enumRegistry map[string]map[string]int3
 	}
 }
 
+func GetMaxBatchSize(configFilePath string) (int, error) {
+	if _, err := os.Stat(configFilePath); errors.Is(err, os.ErrNotExist) {
+		return -1, err
+	}
+	file, err := os.Open(configFilePath)
+	if err != nil {
+		return -1, err
+	}
+	defer file.Close()
+
+	scanner := bufio.NewScanner(file)
+	r, err := regexp.Compile(`max_batch_size:`)
+	if err != nil {
+		return -1, err
+	}
+
+	for scanner.Scan() {
+		if r.MatchString(scanner.Text()) {
+			maxBatchSize := scanner.Text()
+			maxBatchSize = strings.Trim(maxBatchSize, "max_batch_size:")
+			maxBatchSize = strings.Trim(maxBatchSize, " ")
+			intMaxBatchSize, err := strconv.Atoi(maxBatchSize)
+			return intMaxBatchSize, err
+		}
+	}
+
+	return -1, fmt.Errorf("not found")
+}
+
 func DoSupportBatch(configFilePath string) (bool, error) {
 	if _, err := os.Stat(configFilePath); errors.Is(err, os.ErrNotExist) {
 		return false, err