diff --git a/cmd/ormb-storage-initializer/cmd/pull-and-export.go b/cmd/ormb-storage-initializer/cmd/pull-and-export.go index 964aaba4..efd40653 100644 --- a/cmd/ormb-storage-initializer/cmd/pull-and-export.go +++ b/cmd/ormb-storage-initializer/cmd/pull-and-export.go @@ -18,12 +18,14 @@ package cmd import ( "fmt" "os" + "path" "path/filepath" "strings" "github.com/spf13/cobra" "github.com/spf13/viper" + "github.com/kleveross/ormb/pkg/consts" "github.com/kleveross/ormb/pkg/oras" "github.com/kleveross/ormb/pkg/ormb" ) @@ -81,13 +83,41 @@ var pullExportCmd = &cobra.Command{ return err } + // For model serving, need relayout option. + if !reLayoutOpt { + return nil + } + + if err := relayoutModel(dstDir); err != nil { + return err + } + return nil }, } +func relayoutModel(modelDir string) error { + // Rename the `model` directory as `1`. + // e.g. move /mnt/models/model to /mnt/models/1 (dstDir). + // so that, the trtserver will serving, the refenence as + // https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/model_repository.html#section-onnx-models. + originalDir, err := filepath.Abs( + filepath.Join(modelDir, consts.ORMBModelDirectory)) + if err != nil { + return err + } + + if err := os.Rename(originalDir, path.Join(modelDir, "1")); err != nil { + return err + } + + return nil +} + func init() { rootCmd.AddCommand(pullExportCmd) + 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") } diff --git a/examples/Caffe-model/model/1/mobilenet.caffemodel b/examples/Caffe-model/model/mobilenet.caffemodel similarity index 100% rename from examples/Caffe-model/model/1/mobilenet.caffemodel rename to examples/Caffe-model/model/mobilenet.caffemodel diff --git a/examples/Caffe-model/model/1/mobilenet_deploy.prototxt b/examples/Caffe-model/model/mobilenet_deploy.prototxt similarity index 100% rename from examples/Caffe-model/model/1/mobilenet_deploy.prototxt rename to examples/Caffe-model/model/mobilenet_deploy.prototxt diff --git a/examples/Caffe2-model/model/1/init_net.pb b/examples/Caffe2-model/model/init_net.pb similarity index 100% rename from examples/Caffe2-model/model/1/init_net.pb rename to examples/Caffe2-model/model/init_net.pb diff --git a/examples/Caffe2-model/model/1/predict_net.pb b/examples/Caffe2-model/model/predict_net.pb similarity index 100% rename from examples/Caffe2-model/model/1/predict_net.pb rename to examples/Caffe2-model/model/predict_net.pb diff --git a/examples/GraphDef-model/model/1/graphdef.pb b/examples/GraphDef-model/model/graphdef.pb similarity index 100% rename from examples/GraphDef-model/model/1/graphdef.pb rename to examples/GraphDef-model/model/graphdef.pb diff --git a/examples/H5-model/model/1/model.h5 b/examples/H5-model/model/model.h5 similarity index 100% rename from examples/H5-model/model/1/model.h5 rename to examples/H5-model/model/model.h5 diff --git a/examples/MXNETParams-model/model/1/mobilenet_v2-0000.params b/examples/MXNETParams-model/model/mobilenet_v2-0000.params similarity index 100% rename from examples/MXNETParams-model/model/1/mobilenet_v2-0000.params rename to examples/MXNETParams-model/model/mobilenet_v2-0000.params diff --git a/examples/MXNETParams-model/model/1/mobilenet_v2-symbol.json b/examples/MXNETParams-model/model/mobilenet_v2-symbol.json similarity index 100% rename from examples/MXNETParams-model/model/1/mobilenet_v2-symbol.json rename to examples/MXNETParams-model/model/mobilenet_v2-symbol.json diff --git a/examples/ONNX-model/model/1/densenet121-1.2.onnx b/examples/ONNX-model/model/densenet121-1.2.onnx similarity index 100% rename from examples/ONNX-model/model/1/densenet121-1.2.onnx rename to examples/ONNX-model/model/densenet121-1.2.onnx diff --git a/examples/ONNX-model/model/1/densenet_labels.txt b/examples/ONNX-model/model/densenet_labels.txt similarity index 100% rename from examples/ONNX-model/model/1/densenet_labels.txt rename to examples/ONNX-model/model/densenet_labels.txt diff --git a/examples/PMML-model/model/1/single_audit_mlp.pmml b/examples/PMML-model/model/single_audit_mlp.pmml similarity index 100% rename from examples/PMML-model/model/1/single_audit_mlp.pmml rename to examples/PMML-model/model/single_audit_mlp.pmml diff --git a/examples/SavedModel-fashion/model/1/saved_model.pb b/examples/SavedModel-fashion/model/saved_model.pb similarity index 100% rename from examples/SavedModel-fashion/model/1/saved_model.pb rename to examples/SavedModel-fashion/model/saved_model.pb diff --git a/examples/SavedModel-fashion/model/1/variables/variables.data-00000-of-00001 b/examples/SavedModel-fashion/model/variables/variables.data-00000-of-00001 similarity index 100% rename from examples/SavedModel-fashion/model/1/variables/variables.data-00000-of-00001 rename to examples/SavedModel-fashion/model/variables/variables.data-00000-of-00001 diff --git a/examples/SavedModel-fashion/model/1/variables/variables.index b/examples/SavedModel-fashion/model/variables/variables.index similarity index 100% rename from examples/SavedModel-fashion/model/1/variables/variables.index rename to examples/SavedModel-fashion/model/variables/variables.index diff --git a/examples/TensorRT-model/model/1/model.plan b/examples/TensorRT-model/model/model.plan similarity index 100% rename from examples/TensorRT-model/model/1/model.plan rename to examples/TensorRT-model/model/model.plan diff --git a/examples/TorchScript-model/model/1/model.pt b/examples/TorchScript-model/model/model.pt similarity index 100% rename from examples/TorchScript-model/model/1/model.pt rename to examples/TorchScript-model/model/model.pt diff --git a/pkg/model/format.go b/pkg/model/format.go index f55d3ed5..223016a8 100644 --- a/pkg/model/format.go +++ b/pkg/model/format.go @@ -34,13 +34,6 @@ func (f Format) ValidateDirectory(rootPath string) error { if err != nil { return err } - if len(fileList) != 1 || !fileList[0].IsDir() { - return fmt.Errorf("the model directory structure is error") - } - fileList, err = ioutil.ReadDir(path.Join(modelFilePath, fileList[0].Name())) - if err != nil { - return err - } switch f { case FormatSavedModel: