server: support Vertex AI compatible API (#22545)

* server: support Vertex AI compatible API * a bit safer * support other AIP_* env var * various fixes * if AIP_MODE is unset, do nothing * fix test case * fix windows build
2026-05-08 15:23:04 +02:00
parent 9dcf835528
commit 29debb3a6a
5 changed files with 296 additions and 4 deletions
@@ -4,7 +4,9 @@

 #include <cpp-httplib/httplib.h>

+#include <cstdlib>
 #include <functional>
+#include <future>
 #include <string>
 #include <thread>

@@ -51,11 +53,51 @@ static void log_server_request(const httplib::Request & req, const httplib::Resp
    SRV_DBG("response: %s\n", res.body.c_str());
 }

+// For Google Cloud Platform deployment compatibility
+struct gcp_params {
+    bool enabled;
+    std::string path_health;
+    std::string path_predict;
+    int port;
+
+    // Ref: https://docs.cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables
+    gcp_params() {
+        enabled = getenv("AIP_MODE", "") == "PREDICTION";
+        path_health = getenv("AIP_HEALTH_ROUTE", "", true); // default: using the route defined in server.cpp
+        path_predict = getenv("AIP_PREDICT_ROUTE", "/predict", true);
+        port = std::stoi(getenv("AIP_HTTP_PORT", "8080"));
+    }
+
+    static std::string getenv(const char * name, const std::string & default_value, bool ensure_leading_slash = false) {
+        const char * value = std::getenv(name);
+        if (value == nullptr || value[0] == '\0') {
+            return default_value;
+        }
+        std::string val = value;
+        if (ensure_leading_slash && !val.empty() && val[0] != '/') {
+            val.insert(val.begin(), '/');
+        }
+        return val;
+    }
+};
+
 bool server_http_context::init(const common_params & params) {
+    const gcp_params gcp;
+
    path_prefix = params.api_prefix;
    port = params.port;
    hostname = params.hostname;

+    if (gcp.enabled) {
+        LOG_INF("%s: Google Cloud Platform compat: health route = %s, predict route = %s, port = %d\n", __func__, gcp.path_health.c_str(), gcp.path_predict.c_str(), gcp.port);
+
+        if (port != gcp.port) {
+            LOG_WRN("%s: Google Cloud Platform compat: overriding server port %d with AIP_HTTP_PORT %d\n", __func__, port, gcp.port);
+        }
+
+        port = gcp.port;
+    }
+
    auto & srv = pimpl->srv;

 #ifdef CPPHTTPLIB_OPENSSL_SUPPORT
@@ -420,6 +462,7 @@ static void process_handler_response(server_http_req_ptr && request, server_http
 }

 void server_http_context::get(const std::string & path, const server_http_context::handler_t & handler) const {
+    handlers.emplace(path, handler);
    pimpl->srv->Get(path_prefix + path, [handler](const httplib::Request & req, httplib::Response & res) {
        server_http_req_ptr request = std::make_unique<server_http_req>(server_http_req{
            get_params(req),
@@ -436,6 +479,7 @@ void server_http_context::get(const std::string & path, const server_http_contex
 }

 void server_http_context::post(const std::string & path, const server_http_context::handler_t & handler) const {
+    handlers.emplace(path, handler);
    pimpl->srv->Post(path_prefix + path, [handler](const httplib::Request & req, httplib::Response & res) {
        std::string body = req.body;
        std::map<std::string, uploaded_file> files;
@@ -481,3 +525,176 @@ void server_http_context::post(const std::string & path, const server_http_conte
    });
 }

+//
+// Vertex AI Prediction protocol (AIP_PREDICT_ROUTE)
+// https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements
+//
+
+// Derives the camelCase @requestFormat alias for a registered path.
+// e.g. "/v1/chat/completions" -> "chatCompletions", "/apply-template" -> "applyTemplate"
+static std::string path_to_gcp_format(const std::string & path) {
+    std::string s = path;
+    if (s.size() > 3 && s[0] == '/' && s[1] == 'v' && s[2] == '1') {
+        s = s.substr(3);
+    }
+    if (!s.empty() && s[0] == '/') {
+        s = s.substr(1);
+    }
+    std::string result;
+    bool cap = false;
+    for (unsigned char c : s) {
+        if (c == ':') break; // stop before path parameters
+        if (c == '/' || c == '-' || c == '_') {
+            cap = true;
+        } else {
+            result += cap ? (char)std::toupper(c) : (char)c;
+            cap = false;
+        }
+    }
+    return result;
+}
+
+static json parse_gcp_predict_response(const server_http_res_ptr & res) {
+    if (res == nullptr) {
+        throw std::runtime_error("empty response from internal handler");
+    }
+    if (res->is_stream()) {
+        throw std::invalid_argument("predict route does not support streaming responses");
+    }
+    if (res->data.empty()) {
+        return nullptr;
+    }
+    try {
+        return json::parse(res->data);
+    } catch (...) {
+        return res->data;
+    }
+}
+
+void server_http_context::register_gcp_compat() {
+    const gcp_params gcp;
+
+    if (!gcp.enabled) {
+        // do nothing
+        return;
+    }
+
+    if (handlers.count(gcp.path_predict)) {
+        LOG_ERR("%s: AIP_PREDICT_ROUTE=%s conflicts with an existing llama-server route\n", __func__, gcp.path_predict.c_str());
+        exit(1);
+    }
+
+    // camelCase alias -> canonical path (first registration wins on collision)
+    // e.g. "chatCompletions" -> "/v1/chat/completions"
+    std::unordered_map<std::string, std::string> alias_to_path;
+    for (const auto & [path, _] : handlers) {
+        alias_to_path.emplace(path_to_gcp_format(path), path);
+    }
+
+    if (!gcp.path_health.empty()) {
+        auto health_handler = handlers.find("/health");
+        GGML_ASSERT(health_handler != handlers.end());
+        get(gcp.path_health, health_handler->second);
+    }
+
+    post(gcp.path_predict, [this, alias_to_path = std::move(alias_to_path)](const server_http_req & req) -> server_http_res_ptr {
+        static const auto build_error = [](const std::string & message, error_type type) -> json {
+            return json {{"error", format_error_response(message, type)}};
+        };
+
+        json data;
+        try {
+            data = json::parse(req.body);
+        } catch (const std::exception & e) {
+            auto res = std::make_unique<server_http_res>();
+            res->status = 400;
+            res->data = safe_json_to_str({{"error", format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST)}});
+            return res;
+        }
+        if (!data.is_object()) {
+            auto res = std::make_unique<server_http_res>();
+            res->status = 400;
+            res->data = safe_json_to_str({{"error", format_error_response("request body must be a JSON object", ERROR_TYPE_INVALID_REQUEST)}});
+            return res;
+        }
+        if (!data.contains("instances") || !data.at("instances").is_array()) {
+            auto res = std::make_unique<server_http_res>();
+            res->status = 400;
+            res->data = safe_json_to_str({{"error", format_error_response("request body must include an array field named instances", ERROR_TYPE_INVALID_REQUEST)}});
+            return res;
+        }
+
+        const json & instances = data.at("instances");
+        static const size_t MAX_INSTANCES = 128;
+        if (instances.size() > MAX_INSTANCES) {
+            auto res = std::make_unique<server_http_res>();
+            res->status = 400;
+            res->data = safe_json_to_str({{"error", format_error_response("instances array exceeds maximum size of " + std::to_string(MAX_INSTANCES), ERROR_TYPE_INVALID_REQUEST)}});
+            return res;
+        }
+
+        std::vector<std::future<json>> futures;
+        futures.reserve(instances.size());
+
+        for (const auto & instance : instances) {
+            futures.push_back(std::async(std::launch::async, [this, &req, &alias_to_path, instance]() -> json {
+                if (!instance.is_object()) {
+                    return build_error("each instance must be a JSON object", ERROR_TYPE_INVALID_REQUEST);
+                }
+                if (!instance.contains("@requestFormat") || !instance.at("@requestFormat").is_string()) {
+                    return build_error("each instance must include a string @requestFormat", ERROR_TYPE_INVALID_REQUEST);
+                }
+
+                try {
+                    json payload = instance;
+                    const std::string format = payload.at("@requestFormat").get<std::string>();
+                    payload.erase("@requestFormat");
+
+                    if (payload.contains("stream")) {
+                        LOG_WRN("%s: ignoring client-provided stream field in instance, streaming is not supported in predict route\n", __func__);
+                        payload["stream"] = false;
+                    }
+
+                    // accept both camelCase aliases (e.g. "chatCompletions") and direct paths
+                    std::string dispatch_path;
+                    auto it_alias = alias_to_path.find(format);
+                    if (it_alias != alias_to_path.end()) {
+                        dispatch_path = it_alias->second;
+                    } else if (handlers.count(format)) {
+                        dispatch_path = format;
+                    } else {
+                        return build_error("no handler registered for @requestFormat: " + format, ERROR_TYPE_INVALID_REQUEST);
+                    }
+
+                    const server_http_req internal_req {
+                        req.params,
+                        req.headers,
+                        path_prefix + dispatch_path,
+                        req.query_string,
+                        payload.dump(),
+                        {},
+                        req.should_stop,
+                    };
+
+                    server_http_res_ptr internal_res = handlers.at(dispatch_path)(internal_req);
+                    return parse_gcp_predict_response(internal_res);
+                } catch (const std::invalid_argument & e) {
+                    return build_error(e.what(), ERROR_TYPE_INVALID_REQUEST);
+                } catch (const std::exception & e) {
+                    return build_error(e.what(), ERROR_TYPE_SERVER);
+                } catch (...) {
+                    return build_error("unknown error", ERROR_TYPE_SERVER);
+                }
+            }));
+        }
+
+        json predictions = json::array();
+        for (auto & future : futures) {
+            predictions.push_back(future.get());
+        }
+
+        auto res = std::make_unique<server_http_res>();
+        res->data = safe_json_to_str({{"predictions", predictions}});
+        return res;
+    });
+}