Validate create graph parameters

Author: michalkulakowski

src/config.cpp

@@ -92,6 +92,40 @@ bool Config::validate() {
             std::cerr << "Error: --task parameter not set." << std::endl;
             return false;
         }
+        std::vector allowedPipelineTypes = {"LM", "LM_CB", "VLM", "VLM_CB", "AUTO"};
+        if (serverSettings.hfSettings.graphSettings.pipelineType.has_value() && std::find(allowedPipelineTypes.begin(), allowedPipelineTypes.end(), serverSettings.hfSettings.graphSettings.pipelineType) == allowedPipelineTypes.end()) {
+            std::cerr << "pipeline_type: " << serverSettings.hfSettings.graphSettings.pipelineType.value() << " is not allowed. Supported types: LM, LM_CB, VLM, VLM_CB, AUTO" << std::endl;
+            return false;
+        }
+
+        std::vector allowedTargeDevices = {"CPU", "GPU", "NPU", "HETERO"};
+        if (std::find(allowedTargeDevices.begin(), allowedTargeDevices.end(), serverSettings.hfSettings.graphSettings.targetDevice) == allowedTargeDevices.end()) {
+            std::cerr << "target_device: " << serverSettings.hfSettings.graphSettings.targetDevice << " is not allowed. Supported devices: CPU, GPU, NPU, HETERO" << std::endl;
+            return false;
+        }
+
+        std::vector allowedBoolValues = {"false", "true"};
+        if (std::find(allowedBoolValues.begin(), allowedBoolValues.end(), serverSettings.hfSettings.graphSettings.enablePrefixCaching) == allowedBoolValues.end()) {
+            std::cerr << "enable_prefix_caching: " << serverSettings.hfSettings.graphSettings.enablePrefixCaching << " is not allowed. Supported values: true, false" << std::endl;
+            return false;
+        }
+
+        if (std::find(allowedBoolValues.begin(), allowedBoolValues.end(), serverSettings.hfSettings.graphSettings.dynamicSplitFuse) == allowedBoolValues.end()) {
+            std::cerr << "dynamic_split_fuse: " << serverSettings.hfSettings.graphSettings.dynamicSplitFuse << " is not allowed. Supported values: true, false" << std::endl;
+            return false;
+        }
+
+        if (serverSettings.hfSettings.graphSettings.targetDevice != "NPU") {
+            if (serverSettings.hfSettings.graphSettings.pluginConfig.maxPromptLength.has_value()) {
+                std::cerr << "max_prompt_len is only supported for NPU target device";
+                return false;
+            }
+        }
+
+        if (serverSettings.hfSettings.sourceModel.rfind("OpenVINO/", 0) != 0) {
+            std::cerr << "For now only OpenVINO models are supported";
+            return false;
+        }
         return true;
     }
     if (this->serverSettings.listServables) {

src/graph_export/BUILD

@@ -27,6 +27,9 @@ cc_library(
         "@ovms//src:libovms_module",
         "@ovms//src:libovmsfilesystem",
         "@com_github_tencent_rapidjson//:rapidjson",
+        "@mediapipe//mediapipe/framework/port:parse_text_proto",
+        "@mediapipe//mediapipe/framework:calculator_graph",
+        "@ovms//src:libovmsschema",
     ],
     visibility = ["//visibility:public"],
 )

src/graph_export/graph_export.cpp

@@ -25,6 +25,8 @@
 #include <rapidjson/document.h>
 #include <rapidjson/stringbuffer.h>
 #include <rapidjson/writer.h>
+#include <rapidjson/istreamwrapper.h>
+#include <rapidjson/prettywriter.h>
 #pragma warning(pop)
 
 #include "../capi_frontend/server_settings.hpp"
@@ -33,8 +35,12 @@
 #include "../logging.hpp"
 #include "../status.hpp"
 #include "../stringutils.hpp"
+#include "../schema.hpp"
 #include "graph_export_types.hpp"
 
+#include "mediapipe/framework/port/parse_text_proto.h"
+#include "mediapipe/framework/calculator_graph.h"
+
 namespace ovms {
 
 Status createFile(const std::string& filePath, const std::string& contents) {
@@ -116,12 +122,31 @@ static Status createTextGenerationGraphTemplate(const std::string& directoryPath
         }
     }
     })";
-
+    ::mediapipe::CalculatorGraphConfig config;
+    bool success = ::google::protobuf::TextFormat::ParseFromString(oss.str(), &config);
+    if (!success) {
+        SPDLOG_ERROR("Created graph config couldn't be parsed.");
+        return StatusCode::MEDIAPIPE_GRAPH_CONFIG_FILE_INVALID;
+    }
     // clang-format on
     std::string fullPath = FileSystem::joinPath({directoryPath, "graph.pbtxt"});
     return createFile(fullPath, oss.str());
 }
 
+static Status validateSubconfigSchema(const std::string& subconfig, const std::string& type) {
+    rapidjson::Document subconfigJson;
+    rapidjson::ParseResult parseResult = subconfigJson.Parse(subconfig.c_str());
+    if (parseResult.Code()) {
+        SPDLOG_LOGGER_ERROR(modelmanager_logger, "Created {} subconfig file is not a valid JSON file. Error: {}", type, rapidjson::GetParseError_En(parseResult.Code()));
+        return StatusCode::JSON_INVALID;
+    }
+    if (validateJsonAgainstSchema(subconfigJson, MEDIAPIPE_SUBCONFIG_SCHEMA.c_str()) != StatusCode::OK) {
+        SPDLOG_ERROR("Created {} subconfig file is not in valid configuration format", type);
+        return StatusCode::JSON_INVALID;
+    }
+    return StatusCode::OK;
+}
+
 static Status createRerankSubconfigTemplate(const std::string& directoryPath, const RerankGraphSettingsImpl& graphSettings) {
     std::ostringstream oss;
     // clang-format off
@@ -144,6 +169,10 @@ static Status createRerankSubconfigTemplate(const std::string& directoryPath, co
             }
         ]
     })";
+    auto status = validateSubconfigSchema(oss.str(), "rerank");
+    if (!status.ok()){
+        return status;
+    }
     // clang-format on
     std::string fullPath = FileSystem::joinPath({directoryPath, "subconfig.json"});
     return createFile(fullPath, oss.str());
@@ -171,6 +200,10 @@ static Status createEmbeddingsSubconfigTemplate(const std::string& directoryPath
             }
         ]
     })";
+    auto status = validateSubconfigSchema(oss.str(), "embeddings");
+    if (!status.ok()){
+        return status;
+    }
     // clang-format on
     std::string fullPath = FileSystem::joinPath({directoryPath, "subconfig.json"});
     return createFile(fullPath, oss.str());
@@ -210,6 +243,12 @@ static Status createRerankGraphTemplate(const std::string& directoryPath, const
         output_stream: "RESPONSE_PAYLOAD:output"
     })";
 
+    ::mediapipe::CalculatorGraphConfig config;
+    bool success = ::google::protobuf::TextFormat::ParseFromString(oss.str(), &config);
+    if (!success) {
+        SPDLOG_ERROR("Created rerank graph config couldn't be parsed.");
+        return StatusCode::MEDIAPIPE_GRAPH_CONFIG_FILE_INVALID;
+    }
     // clang-format on
     std::string fullPath = FileSystem::joinPath({directoryPath, "graph.pbtxt"});
     auto status = createFile(fullPath, oss.str());
@@ -259,6 +298,12 @@ static Status createEmbeddingsGraphTemplate(const std::string& directoryPath, co
     }
     })";
 
+    ::mediapipe::CalculatorGraphConfig config;
+    bool success = ::google::protobuf::TextFormat::ParseFromString(oss.str(), &config);
+    if (!success) {
+        SPDLOG_ERROR("Created embeddings graph config couldn't be parsed.");
+        return StatusCode::MEDIAPIPE_GRAPH_CONFIG_FILE_INVALID;
+    }
     // clang-format on
     std::string fullPath = FileSystem::joinPath({directoryPath, "graph.pbtxt"});
     auto status = createFile(fullPath, oss.str());

src/test/graph_export_test.cpp

@@ -288,6 +288,36 @@ TEST_F(GraphCreationTest, rerankPositiveDefault) {
     ASSERT_EQ(expectedRerankJsonContents, jsonContents) << jsonContents;
 }
 
+TEST_F(GraphCreationTest, rerankCreatedJsonInvalid) {
+    ovms::HFSettingsImpl hfSettings;
+    hfSettings.task = ovms::rerank;
+    hfSettings.rerankGraphSettings.targetDevice = "GPU";
+    hfSettings.rerankGraphSettings.modelName = "myModel\t";
+    hfSettings.rerankGraphSettings.numStreams = 2;
+    hfSettings.rerankGraphSettings.maxDocLength = 18;
+    hfSettings.rerankGraphSettings.version = 2;
+    std::string graphPath = ovms::FileSystem::appendSlash(this->directoryPath) + "graph.pbtxt";
+    std::string subconfigPath = ovms::FileSystem::appendSlash(this->directoryPath) + "subconfig.json";
+    std::unique_ptr<ovms::GraphExport> graphExporter = std::make_unique<ovms::GraphExport>();
+    auto status = graphExporter->createServableConfig(this->directoryPath, hfSettings);
+    ASSERT_EQ(status, ovms::StatusCode::JSON_INVALID);
+}
+
+TEST_F(GraphCreationTest, rerankCreatedPbtxtInvalid) {
+    ovms::HFSettingsImpl hfSettings;
+    hfSettings.task = ovms::rerank;
+    hfSettings.rerankGraphSettings.targetDevice = "GPU";
+    hfSettings.rerankGraphSettings.modelName = "myModel\"";
+    hfSettings.rerankGraphSettings.numStreams = 2;
+    hfSettings.rerankGraphSettings.maxDocLength = 18;
+    hfSettings.rerankGraphSettings.version = 2;
+    std::string graphPath = ovms::FileSystem::appendSlash(this->directoryPath) + "graph.pbtxt";
+    std::string subconfigPath = ovms::FileSystem::appendSlash(this->directoryPath) + "subconfig.json";
+    std::unique_ptr<ovms::GraphExport> graphExporter = std::make_unique<ovms::GraphExport>();
+    auto status = graphExporter->createServableConfig(this->directoryPath, hfSettings);
+    ASSERT_EQ(status, ovms::StatusCode::MEDIAPIPE_GRAPH_CONFIG_FILE_INVALID);
+}
+
 TEST_F(GraphCreationTest, embeddingsPositiveDefault) {
     ovms::HFSettingsImpl hfSettings;
     hfSettings.task = ovms::embeddings;
@@ -312,6 +342,38 @@ TEST_F(GraphCreationTest, embeddingsPositiveDefault) {
     ASSERT_EQ(expectedEmbeddingsJsonContents, jsonContents) << jsonContents;
 }
 
+TEST_F(GraphCreationTest, embeddingsCreatedJsonInvalid) {
+    ovms::HFSettingsImpl hfSettings;
+    hfSettings.task = ovms::embeddings;
+    hfSettings.embeddingsGraphSettings.targetDevice = "GPU";
+    hfSettings.embeddingsGraphSettings.modelName = "myModel\t";
+    hfSettings.embeddingsGraphSettings.numStreams = 2;
+    hfSettings.embeddingsGraphSettings.truncate = "true";
+    hfSettings.embeddingsGraphSettings.normalize = "true";
+    hfSettings.embeddingsGraphSettings.version = 2;
+    std::string graphPath = ovms::FileSystem::appendSlash(this->directoryPath) + "graph.pbtxt";
+    std::string subconfigPath = ovms::FileSystem::appendSlash(this->directoryPath) + "subconfig.json";
+    std::unique_ptr<ovms::GraphExport> graphExporter = std::make_unique<ovms::GraphExport>();
+    auto status = graphExporter->createServableConfig(this->directoryPath, hfSettings);
+    ASSERT_EQ(status, ovms::StatusCode::JSON_INVALID);
+}
+
+TEST_F(GraphCreationTest, embeddingsCreatedPbtxtInvalid) {
+    ovms::HFSettingsImpl hfSettings;
+    hfSettings.task = ovms::embeddings;
+    hfSettings.embeddingsGraphSettings.targetDevice = "GPU";
+    hfSettings.embeddingsGraphSettings.modelName = "myModel\"";
+    hfSettings.embeddingsGraphSettings.numStreams = 2;
+    hfSettings.embeddingsGraphSettings.truncate = "true";
+    hfSettings.embeddingsGraphSettings.normalize = "true";
+    hfSettings.embeddingsGraphSettings.version = 2;
+    std::string graphPath = ovms::FileSystem::appendSlash(this->directoryPath) + "graph.pbtxt";
+    std::string subconfigPath = ovms::FileSystem::appendSlash(this->directoryPath) + "subconfig.json";
+    std::unique_ptr<ovms::GraphExport> graphExporter = std::make_unique<ovms::GraphExport>();
+    auto status = graphExporter->createServableConfig(this->directoryPath, hfSettings);
+    ASSERT_EQ(status, ovms::StatusCode::MEDIAPIPE_GRAPH_CONFIG_FILE_INVALID);
+}
+
 TEST_F(GraphCreationTest, positivePluginConfigAll) {
     ovms::HFSettingsImpl hfSettings;
     ovms::TextGenGraphSettingsImpl graphSettings;
@@ -377,3 +439,13 @@ TEST_F(GraphCreationTest, negativeGraphOptionsNotInitialized) {
     status = graphExporter->createServableConfig(this->directoryPath, hfSettings);
     ASSERT_EQ(status, ovms::StatusCode::OK) << status.string();
 }
+
+TEST_F(GraphCreationTest, negativeCreatedPbtxtInvalid) {
+    ovms::HFSettingsImpl hfSettings;
+    hfSettings.sourceModel = "\"";
+
+    std::string graphPath = ovms::FileSystem::appendSlash(this->directoryPath) + "graph.pbtxt";
+    std::unique_ptr<ovms::GraphExport> graphExporter = std::make_unique<ovms::GraphExport>();
+    auto status = graphExporter->createServableConfig(this->directoryPath, hfSettings);
+    ASSERT_EQ(status, ovms::StatusCode::MEDIAPIPE_GRAPH_CONFIG_FILE_INVALID);
+}

src/test/ovmsconfig_test.cpp

@@ -739,6 +739,118 @@ TEST(OvmsGraphConfigTest, positiveDefaultStart) {
     ASSERT_EQ(graphSettings.draftModelDirName.has_value(), false);
 }
 
+TEST(OvmsGraphConfigTest, negativePipelineType) {
+    std::string modelName = "OpenVINO/Phi-3-mini-FastDraft-50M-int8-ov";
+    std::string downloadPath = "test/repository";
+    char* n_argv[] = {
+        (char*)"ovms",
+        (char*)"--pull",
+        (char*)"--source_model",
+        (char*)modelName.c_str(),
+        (char*)"--model_repository_path",
+        (char*)downloadPath.c_str(),
+        (char*)"--pipeline_type",
+        (char*)"INVALID",
+    };
+
+    int arg_count = 8;
+    ConstructorEnabledConfig config;
+    EXPECT_EXIT(ovms::Config::instance().parse(arg_count, n_argv), ::testing::ExitedWithCode(OVMS_EX_USAGE), "pipeline_type: INVALID is not allowed. Supported types: LM, LM_CB, VLM, VLM_CB, AUTO");
+}
+
+TEST(OvmsGraphConfigTest, negativeTargetDevice) {
+    std::string modelName = "OpenVINO/Phi-3-mini-FastDraft-50M-int8-ov";
+    std::string downloadPath = "test/repository";
+    char* n_argv[] = {
+        (char*)"ovms",
+        (char*)"--pull",
+        (char*)"--source_model",
+        (char*)modelName.c_str(),
+        (char*)"--model_repository_path",
+        (char*)downloadPath.c_str(),
+        (char*)"--graph_target_device",
+        (char*)"INVALID",
+    };
+
+    int arg_count = 8;
+    ConstructorEnabledConfig config;
+    EXPECT_EXIT(ovms::Config::instance().parse(arg_count, n_argv), ::testing::ExitedWithCode(OVMS_EX_USAGE), "target_device: INVALID is not allowed. Supported devices: CPU, GPU, NPU, HETERO");
+}
+
+TEST(OvmsGraphConfigTest, negativeEnablePrefixCaching) {
+    std::string modelName = "OpenVINO/Phi-3-mini-FastDraft-50M-int8-ov";
+    std::string downloadPath = "test/repository";
+    char* n_argv[] = {
+        (char*)"ovms",
+        (char*)"--pull",
+        (char*)"--source_model",
+        (char*)modelName.c_str(),
+        (char*)"--model_repository_path",
+        (char*)downloadPath.c_str(),
+        (char*)"--enable_prefix_caching",
+        (char*)"INVALID",
+    };
+
+    int arg_count = 8;
+    ConstructorEnabledConfig config;
+    EXPECT_EXIT(ovms::Config::instance().parse(arg_count, n_argv), ::testing::ExitedWithCode(OVMS_EX_USAGE), "enable_prefix_caching: INVALID is not allowed. Supported values: true, false");
+}
+
+TEST(OvmsGraphConfigTest, negativeDynamicSplitFuse) {
+    std::string modelName = "OpenVINO/Phi-3-mini-FastDraft-50M-int8-ov";
+    std::string downloadPath = "test/repository";
+    char* n_argv[] = {
+        (char*)"ovms",
+        (char*)"--pull",
+        (char*)"--source_model",
+        (char*)modelName.c_str(),
+        (char*)"--model_repository_path",
+        (char*)downloadPath.c_str(),
+        (char*)"--dynamic_split_fuse",
+        (char*)"INVALID",
+    };
+
+    int arg_count = 8;
+    ConstructorEnabledConfig config;
+    EXPECT_EXIT(ovms::Config::instance().parse(arg_count, n_argv), ::testing::ExitedWithCode(OVMS_EX_USAGE), "dynamic_split_fuse: INVALID is not allowed. Supported values: true, false");
+}
+
+TEST(OvmsGraphConfigTest, negativeMaxPromptLength) {
+    std::string modelName = "OpenVINO/Phi-3-mini-FastDraft-50M-int8-ov";
+    std::string downloadPath = "test/repository";
+    char* n_argv[] = {
+        (char*)"ovms",
+        (char*)"--pull",
+        (char*)"--source_model",
+        (char*)modelName.c_str(),
+        (char*)"--model_repository_path",
+        (char*)downloadPath.c_str(),
+        (char*)"--max_prompt_len",
+        (char*)"10",
+    };
+
+    int arg_count = 8;
+    ConstructorEnabledConfig config;
+    EXPECT_EXIT(ovms::Config::instance().parse(arg_count, n_argv), ::testing::ExitedWithCode(OVMS_EX_USAGE), "max_prompt_len is only supported for NPU target device");
+}
+
+TEST(OvmsGraphConfigTest, negativeSourceModel) {
+    std::string modelName = "NonOpenVINO/Phi-3-mini-FastDraft-50M-int8-ov";
+    std::string downloadPath = "test/repository";
+    char* n_argv[] = {
+        (char*)"ovms",
+        (char*)"--pull",
+        (char*)"--source_model",
+        (char*)modelName.c_str(),
+        (char*)"--model_repository_path",
+        (char*)downloadPath.c_str(),
+    };
+
+    int arg_count = 6;
+    ConstructorEnabledConfig config;
+    EXPECT_EXIT(ovms::Config::instance().parse(arg_count, n_argv), ::testing::ExitedWithCode(OVMS_EX_USAGE), "For now only OpenVINO models are supported");
+}
+
 TEST(OvmsGraphConfigTest, positiveAllChangedRerank) {
     std::string modelName = "OpenVINO/Phi-3-mini-FastDraft-50M-int8-ov";
     std::string downloadPath = "test/repository";