feat(ormb-storage-initializer): Relayout model file for serving (#110)

cmd/ormb-storage-initializer/cmd/option.go

@@ -18,4 +18,5 @@ package cmd
 var (
 	usernameOpt, passwordOpt                        string
 	passwordFromStdinOpt, insecureOpt, plainHTTPOpt bool
+	reLayoutOpt                                     bool
 )

cmd/ormb-storage-initializer/cmd/prerun.go

@@ -1,3 +1,18 @@
+/*
+Copyright © 2020 NAME HERE <EMAIL ADDRESS>
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
 package cmd
 
 import (

cmd/ormb-storage-initializer/cmd/pull-and-export.go

@@ -19,13 +19,16 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"path"
 	"path/filepath"
 	"strings"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
+	"gopkg.in/yaml.v2"
 
 	"github.com/kleveross/ormb/pkg/consts"
+	"github.com/kleveross/ormb/pkg/model"
 	"github.com/kleveross/ormb/pkg/oras"
 	"github.com/kleveross/ormb/pkg/ormb"
 )
@@ -83,37 +86,73 @@ var pullExportCmd = &cobra.Command{
 			return err
 		}
 
-		// Move the files in model directory to the upper directory.
-		// e.g. Move /mnt/models/model to /mnt/models (dstDir).
-		// Seldon core will run `--model_base_path=dstDir` directly.
-		originalDir, err := filepath.Abs(
-			filepath.Join(dstDir, consts.ORMBModelDirectory))
-		if err != nil {
-			return err
+		// For model registry, can not move any exported data, but for model
+		// serving, it must relayout model file so that it will work.
+		if !reLayoutOpt {
+			return nil
 		}
-		destinationDir, err := filepath.Abs(dstDir)
-		if err != nil {
+
+		if err := relayoutModel(dstDir); err != nil {
 			return err
 		}
-		files, err := ioutil.ReadDir(originalDir)
-		if err != nil {
+
+		return nil
+	},
+}
+
+func relayoutModel(modelDir string) error {
+	// Move the files in model directory to the upper directory.
+	// e.g. Move /mnt/models/model to /mnt/models (dstDir).
+	// but for tensorflow serving, MUST move /mnt/models/model to /mnt/models/1 (dstDir).
+	// Seldon core will run `--model_base_path=dstDir` directly.
+	originalDir, err := filepath.Abs(
+		filepath.Join(modelDir, consts.ORMBModelDirectory))
+	if err != nil {
+		return err
+	}
+	destinationDir, err := filepath.Abs(modelDir)
+	if err != nil {
+		return err
+	}
+
+	// For TensorFlow serving, MUST move /mnt/models/model to /mnt/models/1 (dstDir).
+	ormbfileBytes, err := ioutil.ReadFile(path.Join(modelDir, "ormbfile.yaml"))
+	if err != nil {
+		return err
+	}
+	var metadata model.Metadata
+	err = yaml.Unmarshal(ormbfileBytes, &metadata)
+	if err != nil {
+		return err
+	}
+	if metadata.Format == string(model.FormatSavedModel) {
+		if err := os.Rename(originalDir, path.Join(modelDir, "1")); err != nil {
 			return err
 		}
-		for _, f := range files {
-			oldPath := filepath.Join(originalDir, f.Name())
-			newPath := filepath.Join(destinationDir, f.Name())
-			fmt.Printf("Moving %s to %s\n", oldPath, newPath)
-			if err := os.Rename(oldPath, newPath); err != nil {
-				return err
-			}
-		}
+
 		return nil
-	},
+	}
+
+	// For other format model.
+	files, err := ioutil.ReadDir(originalDir)
+	if err != nil {
+		return err
+	}
+	for _, f := range files {
+		oldPath := filepath.Join(originalDir, f.Name())
+		newPath := filepath.Join(destinationDir, f.Name())
+		fmt.Printf("Moving %s to %s\n", oldPath, newPath)
+		if err := os.Rename(oldPath, newPath); err != nil {
+			return err
+		}
+	}
+	return nil
 }
 
 func init() {
 	rootCmd.AddCommand(pullExportCmd)
 
-	pullExportCmd.Flags().BoolVarP(&plainHTTPOpt, "plain-http", "", false, "use plain http and not https")
+	pullExportCmd.Flags().BoolVarP(&reLayoutOpt, "relayout", "", true, "relayout data for model serving")
+	pullExportCmd.Flags().BoolVarP(&plainHTTPOpt, "plain-http", "", true, "use plain http and not https")
 	pullExportCmd.Flags().BoolVarP(&insecureOpt, "insecure", "", true, "allow connections to TLS registry without certs")
 }