[REFACTOR] Enable validation logic in GenerationConfig (#2411)

This PR enables a centralized validation logic in GenerationConfig.

Author: tqchen

cpp/json_ffi/json_ffi_engine.cc

@@ -97,9 +97,14 @@ bool JSONFFIEngine::AddRequest(std::string request_json_str, std::string request
   gen_cfg->stop_token_ids = conv_template_.stop_token_ids;
   gen_cfg->debug_config = request.debug_config.value_or(DebugConfig());
 
-  Request engine_request(request_id, inputs, GenerationConfig(gen_cfg));
-  this->engine_->AddRequest(engine_request);
+  Result<GenerationConfig> res_gen_config = GenerationConfig::Validate(GenerationConfig(gen_cfg));
+  if (res_gen_config.IsErr()) {
+    err_ = res_gen_config.UnwrapErr();
+    return false;
+  }
 
+  Request engine_request(request_id, inputs, res_gen_config.Unwrap());
+  this->engine_->AddRequest(engine_request);
   return true;
 }
 

cpp/serve/config.cc

@@ -23,33 +23,57 @@ namespace serve {
 
 TVM_REGISTER_OBJECT_TYPE(GenerationConfigNode);
 
-GenerationConfig::GenerationConfig(String config_json_str, const GenerationConfig& default_config) {
+Result<GenerationConfig> GenerationConfig::Validate(GenerationConfig cfg) {
+  using TResult = Result<GenerationConfig>;
+  if (cfg->n <= 0) {
+    return TResult::Error("\"n\" should be at least 1");
+  }
+  if (cfg->temperature < 0) {
+    return TResult::Error("\"temperature\" should be non-negative");
+  }
+  if (cfg->top_p < 0 || cfg->top_p > 1) {
+    return TResult::Error("\"top_p\" should be in range [0, 1]");
+  }
+  if (std::fabs(cfg->frequency_penalty) > 2.0) {
+    return TResult::Error("frequency_penalty must be in [-2, 2]!");
+  }
+  if (cfg->repetition_penalty <= 0) {
+    return TResult::Error("\"repetition_penalty\" must be positive");
+  }
+  if (cfg->top_logprobs < 0 || cfg->top_logprobs > 5) {
+    return TResult::Error("At most 5 top logprob tokens are supported");
+  }
+  if (cfg->top_logprobs != 0 && !(cfg->logprobs)) {
+    return TResult::Error("\"logprobs\" must be true to support \"top_logprobs\"");
+  }
+  for (const auto& item : cfg->logit_bias) {
+    double bias_value = item.second;
+    if (std::fabs(bias_value) > 100.0) {
+      return TResult::Error("Logit bias value should be in range [-100, 100].");
+    }
+  }
+  return TResult::Ok(cfg);
+}
+
+Result<GenerationConfig> GenerationConfig::FromJSON(String config_json_str,
+                                                    const GenerationConfig& default_config) {
+  using TResult = Result<GenerationConfig>;
   picojson::object config = json::ParseToJSONObject(config_json_str);
   ObjectPtr<GenerationConfigNode> n = make_object<GenerationConfigNode>();
 
   n->n = json::LookupOrDefault<int64_t>(config, "n", default_config->n);
-  CHECK_GT(n->n, 0) << "\"n\" should be at least 1";
   n->temperature =
       json::LookupOrDefault<double>(config, "temperature", default_config->temperature);
-  CHECK_GE(n->temperature, 0) << "\"temperature\" should be non-negative";
   n->top_p = json::LookupOrDefault<double>(config, "top_p", default_config->top_p);
-  CHECK(n->top_p >= 0 && n->top_p <= 1) << "\"top_p\" should be in range [0, 1]";
   n->frequency_penalty =
       json::LookupOrDefault<double>(config, "frequency_penalty", default_config->frequency_penalty);
-  CHECK(std::fabs(n->frequency_penalty) <= 2.0) << "Frequency penalty must be in [-2, 2]!";
   n->presence_penalty =
       json::LookupOrDefault<double>(config, "presence_penalty", default_config->presence_penalty);
-  CHECK(std::fabs(n->presence_penalty) <= 2.0) << "Presence penalty must be in [-2, 2]!";
   n->repetition_penalty = json::LookupOrDefault<double>(config, "repetition_penalty",
                                                         default_config->repetition_penalty);
-  CHECK(n->repetition_penalty > 0) << "Repetition penalty must be a positive number!";
   n->logprobs = json::LookupOrDefault<bool>(config, "logprobs", default_config->logprobs);
   n->top_logprobs =
       json::LookupOrDefault<int64_t>(config, "top_logprobs", default_config->top_logprobs);
-  CHECK(n->top_logprobs >= 0 && n->top_logprobs <= 5)
-      << "At most 5 top logprob tokens are supported";
-  CHECK(n->top_logprobs == 0 || n->logprobs)
-      << "\"logprobs\" must be true to support \"top_logprobs\"";
 
   std::optional<picojson::object> logit_bias_obj =
       json::LookupOptional<picojson::object>(config, "logit_bias");
@@ -59,7 +83,6 @@ GenerationConfig::GenerationConfig(String config_json_str, const GenerationConfi
     for (auto [token_id_str, bias] : logit_bias_obj.value()) {
       CHECK(bias.is<double>());
       double bias_value = bias.get<double>();
-      CHECK_LE(std::fabs(bias_value), 100.0) << "Logit bias value should be in range [-100, 100].";
       logit_bias.emplace_back(std::stoi(token_id_str), bias_value);
     }
     n->logit_bias = std::move(logit_bias);
@@ -78,7 +101,9 @@ GenerationConfig::GenerationConfig(String config_json_str, const GenerationConfi
     Array<String> stop_strs;
     stop_strs.reserve(stop_strs_arr.value().size());
     for (const picojson::value& v : stop_strs_arr.value()) {
-      CHECK(v.is<std::string>()) << "Invalid stop string in stop_strs";
+      if (!v.is<std::string>()) {
+        return TResult::Error("Invalid stop string in stop_strs");
+      }
       stop_strs.push_back(v.get<std::string>());
     }
     n->stop_strs = std::move(stop_strs);
@@ -91,7 +116,9 @@ GenerationConfig::GenerationConfig(String config_json_str, const GenerationConfi
     std::vector<int> stop_token_ids;
     stop_token_ids.reserve(stop_token_ids_arr.value().size());
     for (const picojson::value& v : stop_token_ids_arr.value()) {
-      CHECK(v.is<int64_t>()) << "Invalid stop token in stop_token_ids";
+      if (!v.is<int64_t>()) {
+        return TResult::Error("Invalid stop token in stop_token_ids");
+      }
       stop_token_ids.push_back(v.get<int64_t>());
     }
     n->stop_token_ids = std::move(stop_token_ids);
@@ -123,8 +150,7 @@ GenerationConfig::GenerationConfig(String config_json_str, const GenerationConfi
     n->debug_config.ignore_eos =
         json::LookupOrDefault<bool>(debug_config_obj.value(), "ignore_eos", false);
   }
-
-  data_ = std::move(n);
+  return Validate(GenerationConfig(n));
 }
 
 GenerationConfig GenerationConfig::GetDefaultFromModelConfig(

cpp/serve/config.h

@@ -68,7 +68,19 @@ class GenerationConfigNode : public Object {
 
 class GenerationConfig : public ObjectRef {
  public:
-  explicit GenerationConfig(String config_json_str, const GenerationConfig& default_config);
+  /*!
+   * \brief Run validation of generation config and ensure values are in bound.
+   * \return The validtaed Generation config or error.
+   */
+  static Result<GenerationConfig> Validate(GenerationConfig cfg);
+
+  /*!
+   * \brief Create generation config from JSON.
+   * \param config_json_str The json string for generation config
+   * \param default_config The default config
+   */
+  static Result<GenerationConfig> FromJSON(String config_json_str,
+                                           const GenerationConfig& default_config);
 
   /*! \brief Get the default generation config from the model config. */
   static GenerationConfig GetDefaultFromModelConfig(const picojson::object& json);
@@ -192,7 +204,7 @@ class EngineConfigNode : public Object {
   /*************** Debug ***************/
   bool verbose = false;
 
-  TVM_DLL String AsJSONString() const;
+  String AsJSONString() const;
 
   static constexpr const char* _type_key = "mlc.serve.EngineConfig";
   static constexpr const bool _type_has_method_sequal_reduce = false;
@@ -203,14 +215,14 @@ class EngineConfigNode : public Object {
 class EngineConfig : public ObjectRef {
  public:
   /*! \brief Create EngineConfig from JSON object and inferred config. */
-  TVM_DLL static EngineConfig FromJSONAndInferredConfig(
-      const picojson::object& json, const InferrableEngineConfig& inferred_config);
+  static EngineConfig FromJSONAndInferredConfig(const picojson::object& json,
+                                                const InferrableEngineConfig& inferred_config);
 
   /*!
    * \brief Get all the models and model libs from the JSON string for engine initialization.
    * \return The parsed models/model libs from config or error message.
    */
-  TVM_DLL static Result<std::vector<std::pair<std::string, std::string>>>
+  static Result<std::vector<std::pair<std::string, std::string>>>
   GetModelsAndModelLibsFromJSONString(const std::string& json_str);
 
   TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(EngineConfig, ObjectRef, EngineConfigNode);
@@ -225,13 +237,13 @@ struct InferrableEngineConfig {
   std::optional<int64_t> max_history_size;
 
   /*! \brief Infer the config for KV cache from a given initial config. */
-  TVM_DLL static Result<InferrableEngineConfig> InferForKVCache(
+  static Result<InferrableEngineConfig> InferForKVCache(
       EngineMode mode, Device device, double gpu_memory_utilization,
       const std::vector<picojson::object>& model_configs,
       const std::vector<ModelMetadata>& model_metadata, InferrableEngineConfig init_config,
       bool verbose);
   /*! \brief Infer the config for RNN state from a given initial config. */
-  TVM_DLL static Result<InferrableEngineConfig> InferForRNNState(
+  static Result<InferrableEngineConfig> InferForRNNState(
       EngineMode mode, Device device, double gpu_memory_utilization,
       const std::vector<picojson::object>& model_configs,
       const std::vector<ModelMetadata>& model_metadata, InferrableEngineConfig init_config,

cpp/serve/engine.cc

@@ -601,9 +601,10 @@ class EngineModule : public ModuleNode {
   void Abort(const String& request_id) { return GetEngine()->AbortRequest(request_id); }
 
   Request CreateRequest(String id, Array<Data> inputs, String generation_cfg_json_str) {
-    return Request(
-        std::move(id), std::move(inputs),
-        GenerationConfig(std::move(generation_cfg_json_str), default_generation_config_));
+    auto gen_config =
+        GenerationConfig::FromJSON(std::move(generation_cfg_json_str), default_generation_config_);
+    CHECK(gen_config.IsOk()) << gen_config.UnwrapErr();
+    return Request(std::move(id), std::move(inputs), gen_config.Unwrap());
   }
   /*! \brief Redirection to `Engine::Step`. */
   void Step() { return GetEngine()->Step(); }

cpp/serve/model.h

@@ -355,16 +355,16 @@ class Model : public ObjectRef {
    * \param trace_enabled A boolean indicating whether tracing is enabled.
    * \return The created runtime module.
    */
-  TVM_DLL static Model Create(String reload_lib_path, String model_path,
-                              const picojson::object& model_config, DLDevice device,
-                              const Optional<Session>& session, bool trace_enabled);
+  static Model Create(String reload_lib_path, String model_path,
+                      const picojson::object& model_config, DLDevice device,
+                      const Optional<Session>& session, bool trace_enabled);
 
   /*!
    * Load the model config from the given model path.
    * \param model_path The path to the model weight parameters.
    * \return The model config json object.
    */
-  TVM_DLL static Result<picojson::object> LoadModelConfig(const String& model_path);
+  static Result<picojson::object> LoadModelConfig(const String& model_path);
 
   TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Model, ObjectRef, ModelObj);
 };

cpp/serve/sampler/sampler.h

@@ -126,7 +126,7 @@ class SamplerObj : public Object {
 class Sampler : public ObjectRef {
  public:
   /*! * \brief Create a CPU sampler. */
-  TVM_DLL static Sampler CreateCPUSampler(Optional<EventTraceRecorder> trace_recorder);
+  static Sampler CreateCPUSampler(Optional<EventTraceRecorder> trace_recorder);
   /*!
    * \brief Create a GPU sampler.
    * \param max_num_sample The max number of samples to sample at a time.
@@ -135,9 +135,8 @@ class Sampler : public ObjectRef {
    * \param device The device that the model runs on.
    * \param trace_recorder The event trace recorder.
    */
-  TVM_DLL static Sampler CreateGPUSampler(int max_num_sample, int vocab_size, FunctionTable* ft,
-                                          DLDevice device,
-                                          Optional<EventTraceRecorder> trace_recorder);
+  static Sampler CreateGPUSampler(int max_num_sample, int vocab_size, FunctionTable* ft,
+                                  DLDevice device, Optional<EventTraceRecorder> trace_recorder);
 
   /*! \brief Check if the given device supports GPU sampling. */
   static bool SupportGPUSampler(Device device) {

cpp/serve/threaded_engine.cc

@@ -220,9 +220,10 @@ class ThreadedEngineImpl : public ThreadedEngine {
   }
 
   Request CreateRequest(String id, Array<Data> inputs, String generation_cfg_json_str) const {
-    return Request(
-        std::move(id), std::move(inputs),
-        GenerationConfig(std::move(generation_cfg_json_str), GetDefaultGenerationConfig()));
+    auto gen_config = GenerationConfig::FromJSON(std::move(generation_cfg_json_str),
+                                                 GetDefaultGenerationConfig());
+    CHECK(gen_config.IsOk()) << gen_config.UnwrapErr();
+    return Request(std::move(id), std::move(inputs), gen_config.Unwrap());
   }
 
   EngineConfig GetCompleteEngineConfig() const final {