feat(MLflow)[KLV-30]: Support mlflow

examples/MLflow-model/model/MLmodel

@@ -0,0 +1,13 @@
+artifact_path: model
+flavors:
+  python_function:
+    env: conda.yaml
+    loader_module: mlflow.sklearn
+    model_path: model.pkl
+    python_version: 3.6.12
+  sklearn:
+    pickled_model: model.pkl
+    serialization_format: cloudpickle
+    sklearn_version: 0.19.1
+run_id: 850277c51810451fab9b06976cf6cf84
+utc_time_created: '2020-10-19 13:08:18.484317'

examples/MLflow-model/model/conda.yaml

@@ -0,0 +1,11 @@
+channels:
+- defaults
+- conda-forge
+dependencies:
+- python=3.6.12
+- scikit-learn=0.19.1
+- pip
+- pip:
+  - mlflow
+  - cloudpickle==1.6.0
+name: mlflow-env

examples/MLflow-model/ormbfile.yaml

@@ -0,0 +1,2 @@
+author: "c1 <wangsiyuan.wangsy@bytedance.com>"
+format: "MLflow"

pkg/model/format.go

@@ -1,6 +1,7 @@
 package model
 
 import (
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -23,6 +24,7 @@ const (
 	FormatTensorRT    Format = "TensorRT"
 	FormatSKLearn     Format = "SKLearn"
 	FormatXGBoost     Format = "XGBoost"
+	FormatMLflow      Format = "MLflow"
 	FormatOthers      Format = "Others"
 )
 
@@ -62,7 +64,12 @@ func (f Format) ValidateDirectory(rootPath string) error {
 		err = f.validateForSKLearn(modelFilePath, fileList)
 	case FormatXGBoost:
 		err = f.validateForXGBoost(modelFilePath, fileList)
+	case FormatMLflow:
+		err = f.validateForMLflow(modelFilePath, fileList)
+	default:
+		err = errors.New("unrecognized model format, please check the ormbfile.yaml")
 	}
+
 	if err != nil {
 		return err
 	}
@@ -252,3 +259,17 @@ func (f Format) validateForXGBoost(modelPath string, files []os.FileInfo) error
 	}
 	return nil
 }
+
+func (f Format) validateForMLflow(modelPath string, files []os.FileInfo) error {
+	var isMLflowFile bool
+	for _, file := range files {
+		if file.Name() == "MLmodel" {
+			// assuming that user would not fool the tool
+			isMLflowFile = true
+		}
+	}
+	if !isMLflowFile {
+		return fmt.Errorf("there are no MLmodel file in %v, directory", modelPath)
+	}
+	return nil
+}

pkg/model/format_test.go

@@ -79,4 +79,16 @@ var _ = Describe("Format", func() {
 		err := xgboostFormat.ValidateDirectory("../../examples/XGBoost-model")
 		Expect(err).To(BeNil())
 	})
+
+	It("Should validate MLmodel format successfully", func() {
+		mlmodelFormat := model.FormatMLflow
+		err := mlmodelFormat.ValidateDirectory("../../examples/MLflow-model")
+		Expect(err).To(BeNil())
+	})
+
+	It("Should validate fail for corruptted format", func() {
+		others := model.FormatOthers
+		err := others.ValidateDirectory("../../test/Corrupt")
+		Expect(err).Should(HaveOccurred())
+	})
 })

test/Corrupt/model/corrupt

@@ -0,0 +1,8 @@
+#!/usr/bin/env -S scheme --script
+
+(define ode-to-a-nightingale "\'Tis not through envy of thy happy lot, but being too happy in thine happiness.")
+
+(define (main)
+  (display ode-to-a-nightingale))
+
+(main)

test/Corrupt/ormbfile.yaml

@@ -0,0 +1,2 @@
+author: "c1 <wangsiyuan.wangsy@bytedance.com>"
+format: "corruptted :)"