diff --git a/cmd/ormb-storage-initializer/cmd/option.go b/cmd/ormb-storage-initializer/cmd/option.go index 0caf309a..002e88f3 100644 --- a/cmd/ormb-storage-initializer/cmd/option.go +++ b/cmd/ormb-storage-initializer/cmd/option.go @@ -18,4 +18,5 @@ package cmd var ( usernameOpt, passwordOpt string passwordFromStdinOpt, insecureOpt, plainHTTPOpt bool + reLayoutOpt bool ) diff --git a/cmd/ormb-storage-initializer/cmd/prerun.go b/cmd/ormb-storage-initializer/cmd/prerun.go index 663aa438..2882713a 100644 --- a/cmd/ormb-storage-initializer/cmd/prerun.go +++ b/cmd/ormb-storage-initializer/cmd/prerun.go @@ -1,3 +1,18 @@ +/* +Copyright © 2020 NAME HERE + +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 ( diff --git a/cmd/ormb-storage-initializer/cmd/pull-and-export.go b/cmd/ormb-storage-initializer/cmd/pull-and-export.go index dab1a171..c4861f15 100644 --- a/cmd/ormb-storage-initializer/cmd/pull-and-export.go +++ b/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") }