From 77b4a0129581c82ebe097ecd0240c49f745a15e2 Mon Sep 17 00:00:00 2001
From: rvalkenaers <rien.valkenaers@gmail.com>
Date: Wed, 13 Sep 2023 12:34:29 +0000
Subject: [PATCH] Fix displaying MLFlow predictors

Signed-off-by: rvalkenaers <rien.valkenaers@gmail.com>
---
 frontend/src/app/shared/utils.ts            | 3 +++
 frontend/src/app/types/kfserving/v1beta1.ts | 2 ++
 2 files changed, 5 insertions(+)

diff --git a/frontend/src/app/shared/utils.ts b/frontend/src/app/shared/utils.ts
index 6b78ce8..1b90378 100644
--- a/frontend/src/app/shared/utils.ts
+++ b/frontend/src/app/shared/utils.ts
@@ -227,6 +227,9 @@ export function getPredictorRuntime(predictor: PredictorSpec): string {
   if (predictorType === PredictorType.Lightgbm) {
     return 'LightGBM ModelServer';
   }
+  if (predictorType === PredictorType.MLFlow) {
+    return 'MLFlow ModelServer';
+  }
   if (predictorType === PredictorType.Custom) {
     return 'Custom ModelServer';
   }
diff --git a/frontend/src/app/types/kfserving/v1beta1.ts b/frontend/src/app/types/kfserving/v1beta1.ts
index 5e175fd..42a993d 100644
--- a/frontend/src/app/types/kfserving/v1beta1.ts
+++ b/frontend/src/app/types/kfserving/v1beta1.ts
@@ -47,6 +47,7 @@ export enum PredictorType {
   Xgboost = 'xgboost',
   Pmml = 'pmml',
   Lightgbm = 'lightgbm',
+  MLFlow = 'mlflow',
   Custom = 'custom',
 }
 
@@ -59,6 +60,7 @@ export interface PredictorSpec extends V1PodSpec, ComponentExtensionSpec {
   onnx?: PredictorExtensionSpec;
   pmml?: PredictorExtensionSpec;
   lightgbm?: PredictorExtensionSpec;
+  mlflow?: PredictorExtensionSpec;
   model?: ModelSpec;
 }