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

tools/server/tests/unit/test_compat_gcp.py

@@ -0,0 +1,60 @@
+import pytest
+from utils import *
+
+server: ServerProcess
+
+
+@pytest.fixture(autouse=True)
+def create_server():
+    global server
+    server = ServerPreset.tinyllama2()
+    server.gcp_compat = True
+
+
+def test_gcp_predict_camel_case():
+    global server
+    server.start()
+    res = server.make_request("POST", "/predict", data={
+        "instances": [
+            {
+                "@requestFormat": "chatCompletions",
+                "max_tokens": 8,
+                "messages": [
+                    {"role": "user", "content": "What is the meaning of life?"},
+                ],
+            }
+        ],
+    })
+    assert res.status_code == 200
+    assert "predictions" in res.body
+    assert len(res.body["predictions"]) == 1
+    prediction = res.body["predictions"][0]
+    assert "choices" in prediction
+    assert len(prediction["choices"]) == 1
+    assert prediction["choices"][0]["message"]["role"] == "assistant"
+    assert len(prediction["choices"][0]["message"]["content"]) > 0
+
+
+def test_gcp_predict_multiple_instances():
+    global server
+    server.n_slots = 2
+    server.start()
+    res = server.make_request("POST", "/predict", data={
+        "instances": [
+            {
+                "@requestFormat": "chatCompletions",
+                "max_tokens": 8,
+                "messages": [{"role": "user", "content": "Say hello"}],
+            },
+            {
+                "@requestFormat": "chatCompletions",
+                "max_tokens": 8,
+                "messages": [{"role": "user", "content": "Say world"}],
+            },
+        ],
+    })
+    assert res.status_code == 200
+    assert len(res.body["predictions"]) == 2
+    for prediction in res.body["predictions"]:
+        assert "choices" in prediction
+        assert len(prediction["choices"][0]["message"]["content"]) > 0

tools/server/tests/utils.py

@@ -108,6 +108,7 @@ class ServerProcess:
     no_cache_idle_slots: bool = False
     log_path: str | None = None
     webui_mcp_proxy: bool = False
+    gcp_compat: bool = False
 
     # session variables
     process: subprocess.Popen | None = None
@@ -122,6 +123,9 @@ class ServerProcess:
         self.external_server = "DEBUG_EXTERNAL" in os.environ
 
     def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None:
+        env = {**os.environ}
+        if "LLAMA_CACHE" not in os.environ:
+            env["LLAMA_CACHE"] = "tmp"
         if self.external_server:
             print(f"[external_server]: Assuming external server running on {self.server_host}:{self.server_port}")
             return
@@ -248,6 +252,8 @@ class ServerProcess:
             server_args.append("--no-cache-idle-slots")
         if self.webui_mcp_proxy:
             server_args.append("--webui-mcp-proxy")
+        if self.gcp_compat:
+            env["AIP_MODE"] = "PREDICTION"
 
         args = [str(arg) for arg in [server_path, *server_args]]
         print(f"tests: starting server with: {' '.join(args)}")
@@ -268,7 +274,7 @@ class ServerProcess:
             creationflags=flags,
             stdout=self._log,
             stderr=self._log if self._log != sys.stdout else sys.stdout,
-            env={**os.environ, "LLAMA_CACHE": "tmp"} if "LLAMA_CACHE" not in os.environ else None,
+            env=env,
         )
         server_instances.add(self)