294 lines
11 KiB
Python
294 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
from fastapi.testclient import TestClient
|
|
|
|
|
|
def _write_config(tmp_path: Path) -> Path:
|
|
cfg = {
|
|
"ray": {
|
|
"address": "http://127.0.0.1:8265",
|
|
"shared_root": "/private",
|
|
"entrypoint_resources": {"worker_node": 1},
|
|
"runtime_env": {"env_vars": {}},
|
|
},
|
|
"data": {
|
|
"user_root": str(tmp_path / "users"),
|
|
},
|
|
"service": {
|
|
"api": {"host": "127.0.0.1", "port": 0},
|
|
"auth": {"token_env": "MVP_INTERNAL_TOKEN"},
|
|
"sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
|
|
"scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
|
|
},
|
|
"serving": {
|
|
"serve": {"http_port": 8000, "proxy_location": "HeadOnly"},
|
|
"llm": {"accelerator_type": "H20"},
|
|
},
|
|
}
|
|
p = tmp_path / "cfg.yaml"
|
|
p.write_text(yaml.safe_dump(cfg), encoding="utf-8")
|
|
return p
|
|
|
|
|
|
def test_serving_api_crud_flow(tmp_path: Path, monkeypatch):
|
|
from argus.service import app as app_mod
|
|
|
|
cfg_path = _write_config(tmp_path)
|
|
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
|
|
|
class _Scheduler:
|
|
def __init__(self, **kwargs):
|
|
self.tool = object()
|
|
|
|
def run_forever(self, stop_flag):
|
|
return None
|
|
|
|
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
|
monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: f"mk-{user_id}")
|
|
|
|
app = app_mod.create_app(str(cfg_path))
|
|
|
|
admin_headers = {"authorization": "Bearer admin-token"}
|
|
with TestClient(app) as c:
|
|
r = c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
|
assert r.status_code == 200
|
|
r2 = c.post("/api/v2/users/alice/tokens", headers=admin_headers)
|
|
assert r2.status_code == 200
|
|
user_token = r2.json()["token"]
|
|
|
|
headers = {"authorization": f"Bearer {user_token}"}
|
|
|
|
spec_yaml = (
|
|
"model_id: qwen-0.5b\n"
|
|
"model_source: $HOME/common/hf/models--Qwen--Qwen2.5-0.5B-Instruct/snapshots/sha\n"
|
|
"num_replicas: 1\n"
|
|
"gpus_per_replica: 1\n"
|
|
)
|
|
r3 = c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
|
|
assert r3.status_code == 200
|
|
assert r3.json()["model_key"] == "mk-alice"
|
|
assert r3.json()["state"] == "QUEUED"
|
|
|
|
r4 = c.get("/api/v2/serve/models?limit=10&offset=0", headers=headers)
|
|
assert r4.status_code == 200
|
|
obj = r4.json()
|
|
assert obj["openai_base_url"] is None
|
|
assert len(obj["items"]) == 1
|
|
assert obj["items"][0]["model_key"] == "mk-alice"
|
|
assert obj["items"][0]["endpoint"]["openai_base_url"] == "http://testserver:8000/serve/mk-alice/v1"
|
|
|
|
r5 = c.get("/api/v2/serve/models/mk-alice", headers=headers)
|
|
assert r5.status_code == 200
|
|
detail = r5.json()
|
|
assert detail["model"]["model_key"] == "mk-alice"
|
|
assert "model_id_prefix" in detail["model"]
|
|
assert "resolved_spec_yaml" in detail
|
|
assert isinstance(detail.get("events"), list)
|
|
|
|
r6 = c.patch("/api/v2/serve/models/mk-alice", headers=headers, json={"num_replicas": 2})
|
|
assert r6.status_code == 200
|
|
assert r6.json()["state"] == "QUEUED"
|
|
|
|
r7 = c.delete("/api/v2/serve/models/mk-alice", headers=headers)
|
|
assert r7.status_code == 200
|
|
assert r7.json()["state"] == "DELETING"
|
|
|
|
# Admin status endpoint
|
|
r8 = c.get("/api/v2/serve/status", headers=admin_headers)
|
|
assert r8.status_code == 200
|
|
assert r8.json()["http_port"] == 8000
|
|
assert r8.json()["openai_base_url"] is None
|
|
|
|
|
|
def test_serving_api_rejects_path_outside_user_and_hf(tmp_path: Path, monkeypatch):
|
|
from argus.service import app as app_mod
|
|
|
|
cfg_path = _write_config(tmp_path)
|
|
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
|
|
|
class _Scheduler:
|
|
def __init__(self, **kwargs):
|
|
self.tool = object()
|
|
|
|
def run_forever(self, stop_flag):
|
|
return None
|
|
|
|
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
|
monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: f"mk-{user_id}")
|
|
|
|
app = app_mod.create_app(str(cfg_path))
|
|
|
|
admin_headers = {"authorization": "Bearer admin-token"}
|
|
with TestClient(app) as c:
|
|
c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
|
r2 = c.post("/api/v2/users/alice/tokens", headers=admin_headers)
|
|
user_token = r2.json()["token"]
|
|
headers = {"authorization": f"Bearer {user_token}"}
|
|
|
|
spec_yaml = (
|
|
"model_id: x\n"
|
|
"model_source: /private/users/bob/models/evil\n"
|
|
"num_replicas: 1\n"
|
|
"gpus_per_replica: 1\n"
|
|
)
|
|
r3 = c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
|
|
assert r3.status_code == 403
|
|
|
|
|
|
def test_serving_api_invalid_yaml_and_non_mapping(tmp_path: Path, monkeypatch):
|
|
from argus.service import app as app_mod
|
|
|
|
cfg_path = _write_config(tmp_path)
|
|
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
|
|
|
class _Scheduler:
|
|
def __init__(self, **kwargs):
|
|
self.tool = object()
|
|
|
|
def run_forever(self, stop_flag):
|
|
return None
|
|
|
|
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
|
monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: f"mk-{user_id}")
|
|
app = app_mod.create_app(str(cfg_path))
|
|
|
|
with TestClient(app) as c:
|
|
# Create a user token
|
|
admin_headers = {"authorization": "Bearer admin-token"}
|
|
c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
|
token = c.post("/api/v2/users/alice/tokens", headers=admin_headers).json()["token"]
|
|
headers = {"authorization": f"Bearer {token}"}
|
|
|
|
r = c.post("/api/v2/serve/models", headers=headers, data=": bad\n")
|
|
assert r.status_code == 400
|
|
|
|
r2 = c.post("/api/v2/serve/models", headers=headers, data="- 1\n- 2\n")
|
|
assert r2.status_code == 400
|
|
|
|
|
|
def test_serving_api_engine_kwargs_binary_rejected(tmp_path: Path, monkeypatch):
|
|
"""
|
|
yaml !!binary is parsed as bytes, which is not JSON-serializable.
|
|
"""
|
|
from argus.service import app as app_mod
|
|
|
|
cfg_path = _write_config(tmp_path)
|
|
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
|
|
|
class _Scheduler:
|
|
def __init__(self, **kwargs):
|
|
self.tool = object()
|
|
|
|
def run_forever(self, stop_flag):
|
|
return None
|
|
|
|
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
|
monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: f"mk-{user_id}")
|
|
app = app_mod.create_app(str(cfg_path))
|
|
|
|
admin_headers = {"authorization": "Bearer admin-token"}
|
|
with TestClient(app) as c:
|
|
c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
|
token = c.post("/api/v2/users/alice/tokens", headers=admin_headers).json()["token"]
|
|
headers = {"authorization": f"Bearer {token}"}
|
|
|
|
spec_yaml = (
|
|
"model_id: x\n"
|
|
"model_source: $HOME/common/hf/x\n"
|
|
"engine_kwargs:\n"
|
|
" blob: !!binary \"AQID\"\n"
|
|
)
|
|
r = c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
|
|
assert r.status_code == 400
|
|
|
|
|
|
def test_serving_api_list_include_deleted_and_forwarded_base_url(tmp_path: Path, monkeypatch):
|
|
from argus.service import app as app_mod
|
|
from argus.service.config import V2Config
|
|
from argus.service.db import Db
|
|
|
|
cfg_path = _write_config(tmp_path)
|
|
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
|
|
|
class _Scheduler:
|
|
def __init__(self, **kwargs):
|
|
self.tool = object()
|
|
|
|
def run_forever(self, stop_flag):
|
|
return None
|
|
|
|
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
|
keys = iter(["mk-alice-1", "mk-alice-2"])
|
|
monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: next(keys))
|
|
|
|
app = app_mod.create_app(str(cfg_path))
|
|
|
|
admin_headers = {"authorization": "Bearer admin-token"}
|
|
with TestClient(app) as c:
|
|
c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
|
token = c.post("/api/v2/users/alice/tokens", headers=admin_headers).json()["token"]
|
|
headers = {"authorization": f"Bearer {token}"}
|
|
|
|
spec_yaml = "model_id: x\nmodel_source: $HOME/common/hf/x\n"
|
|
c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
|
|
c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
|
|
|
|
# Mark one model as DELETED directly in DB (sets deleted_at).
|
|
root = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
|
|
v2_cfg = V2Config.from_root_dict(root)
|
|
db = Db(v2_cfg.sqlite.db_path)
|
|
db.set_serve_model_state(model_key="mk-alice-2", state="DELETED")
|
|
|
|
r1 = c.get(
|
|
"/api/v2/serve/models?limit=10&offset=0&include_deleted=0",
|
|
headers={**headers, "x-forwarded-host": "example.com:8080", "x-forwarded-proto": "https"},
|
|
)
|
|
assert r1.status_code == 200
|
|
assert r1.json()["openai_base_url"] is None
|
|
assert {m["model_key"] for m in r1.json()["items"]} == {"mk-alice-1"}
|
|
assert r1.json()["items"][0]["endpoint"]["openai_base_url"] == "https://example.com:8000/serve/mk-alice-1/v1"
|
|
|
|
# IPv6 bracketed host should be preserved.
|
|
r1b = c.get(
|
|
"/api/v2/serve/models?limit=10&offset=0&include_deleted=0",
|
|
headers={**headers, "x-forwarded-host": "[::1]:8080", "x-forwarded-proto": "https"},
|
|
)
|
|
assert r1b.status_code == 200
|
|
assert r1b.json()["items"][0]["endpoint"]["openai_base_url"] == "https://[::1]:8000/serve/mk-alice-1/v1"
|
|
|
|
r2 = c.get("/api/v2/serve/models?include_deleted=1", headers=headers)
|
|
assert r2.status_code == 200
|
|
assert {m["model_key"] for m in r2.json()["items"]} == {"mk-alice-1", "mk-alice-2"}
|
|
|
|
|
|
def test_serving_api_patch_invalid_num_replicas(tmp_path: Path, monkeypatch):
|
|
from argus.service import app as app_mod
|
|
|
|
cfg_path = _write_config(tmp_path)
|
|
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
|
|
|
class _Scheduler:
|
|
def __init__(self, **kwargs):
|
|
self.tool = object()
|
|
|
|
def run_forever(self, stop_flag):
|
|
return None
|
|
|
|
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
|
monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: "mk-alice")
|
|
|
|
app = app_mod.create_app(str(cfg_path))
|
|
|
|
admin_headers = {"authorization": "Bearer admin-token"}
|
|
with TestClient(app) as c:
|
|
c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
|
token = c.post("/api/v2/users/alice/tokens", headers=admin_headers).json()["token"]
|
|
headers = {"authorization": f"Bearer {token}"}
|
|
|
|
c.post("/api/v2/serve/models", headers=headers, data="model_id: x\nmodel_source: $HOME/common/hf/x\n")
|
|
r = c.patch("/api/v2/serve/models/mk-alice", headers=headers, json={"num_replicas": 0})
|
|
assert r.status_code == 422
|