264 lines
7.6 KiB
Python
264 lines
7.6 KiB
Python
import asyncio
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
import pytest
|
|
from flask import Flask
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT / "api"))
|
|
|
|
|
|
class _Pipeline:
|
|
def incrby(self, *_args):
|
|
return self
|
|
|
|
def execute(self):
|
|
return []
|
|
|
|
|
|
class _FakeRedis:
|
|
def __init__(self, acquire_result=None):
|
|
self.acquire_result = acquire_result or [1, b"ok"]
|
|
self.eval_calls = []
|
|
|
|
def zremrangebyscore(self, *_args):
|
|
return 0
|
|
|
|
def zcard(self, *_args):
|
|
return 0
|
|
|
|
def eval(self, script, key_count, *args):
|
|
self.eval_calls.append((script, key_count, args))
|
|
if "endpoint_concurrency" in script:
|
|
return self.acquire_result
|
|
return 1
|
|
|
|
def pipeline(self):
|
|
return _Pipeline()
|
|
|
|
def delete(self, *_args):
|
|
return 1
|
|
|
|
def incr(self, *_args):
|
|
return 1
|
|
|
|
def expire(self, *_args):
|
|
return True
|
|
|
|
|
|
@pytest.fixture()
|
|
def pool_app(monkeypatch):
|
|
from insurance.db.database import db
|
|
from insurance.models.model_pool import ModelEndpoint, ModelPool, ModelQuotaGroup
|
|
from insurance.utils.credential_crypto import encrypt_secret
|
|
|
|
app = Flask(__name__)
|
|
app.config.update(
|
|
SQLALCHEMY_DATABASE_URI="sqlite:///:memory:",
|
|
SQLALCHEMY_TRACK_MODIFICATIONS=False,
|
|
)
|
|
db.init_app(app)
|
|
monkeypatch.setenv("INSURANCE_CREDENTIAL_MASTER_KEY", "11" * 32)
|
|
monkeypatch.setenv("INSURANCE_CREDENTIAL_KEY_VERSION", "1")
|
|
with app.app_context():
|
|
ModelPool.__table__.create(bind=db.engine, checkfirst=True)
|
|
ModelQuotaGroup.__table__.create(bind=db.engine, checkfirst=True)
|
|
ModelEndpoint.__table__.create(bind=db.engine, checkfirst=True)
|
|
pool = ModelPool(
|
|
id=1,
|
|
name="ppt-default",
|
|
purpose="ppt_extract",
|
|
acquire_timeout_ms=1,
|
|
max_attempts=2,
|
|
)
|
|
group = ModelQuotaGroup(
|
|
id=1,
|
|
name="vendor-account",
|
|
provider="openai",
|
|
max_concurrency=2,
|
|
rpm_limit=10,
|
|
tpm_limit=10000,
|
|
)
|
|
endpoint = ModelEndpoint(
|
|
id=1,
|
|
pool_id=1,
|
|
quota_group_id=1,
|
|
name="endpoint-a",
|
|
provider="openai",
|
|
base_url="https://api.example.com/v1",
|
|
model_name="model-a",
|
|
encrypted_api_key=encrypt_secret(
|
|
"secret-a", context="model_endpoint:1:api_key"
|
|
),
|
|
max_concurrency=1,
|
|
rpm_limit=5,
|
|
capability_json='{"json": true}',
|
|
)
|
|
db.session.add_all([pool, group, endpoint])
|
|
db.session.commit()
|
|
return app
|
|
|
|
|
|
def test_model_pool_acquires_and_releases_cross_process_lease(
|
|
pool_app, monkeypatch
|
|
):
|
|
from insurance.model_pool import service
|
|
|
|
redis = _FakeRedis()
|
|
monkeypatch.setattr(service, "redis_client", redis)
|
|
with pool_app.app_context():
|
|
lease = service.acquire_endpoint(
|
|
"ppt_extract",
|
|
required_capabilities={"json": True},
|
|
reserved_tokens=100,
|
|
)
|
|
assert lease.api_key == "secret-a"
|
|
assert lease.endpoint.id == 1
|
|
lease.release(actual_tokens=80)
|
|
|
|
assert len(redis.eval_calls) == 2
|
|
acquire_args = redis.eval_calls[0][2]
|
|
assert "insurance:model:leases:endpoint:1" in acquire_args
|
|
assert "insurance:model:leases:quota:1" in acquire_args
|
|
|
|
|
|
def test_model_pool_fails_closed_when_capacity_is_busy(pool_app, monkeypatch):
|
|
from insurance.model_pool import service
|
|
|
|
monkeypatch.setattr(
|
|
service,
|
|
"redis_client",
|
|
_FakeRedis([0, b"quota_concurrency"]),
|
|
)
|
|
with pool_app.app_context(), pytest.raises(
|
|
service.ModelPoolBusy, match="quota_concurrency"
|
|
):
|
|
service.acquire_endpoint("ppt_extract")
|
|
|
|
|
|
def test_model_pool_auth_failure_disables_endpoint(pool_app, monkeypatch):
|
|
from insurance.model_pool import service
|
|
from insurance.models.model_pool import ModelEndpoint
|
|
|
|
class _Response:
|
|
status_code = 401
|
|
headers = {}
|
|
|
|
class _AuthError(RuntimeError):
|
|
response = _Response()
|
|
|
|
monkeypatch.setattr(service, "redis_client", _FakeRedis())
|
|
with pool_app.app_context():
|
|
service.record_failure(1, _AuthError("invalid credential"))
|
|
endpoint = service.db.session.get(ModelEndpoint, 1)
|
|
assert endpoint.enabled is False
|
|
assert endpoint.status == "disabled_auth"
|
|
payload = endpoint.to_dict()
|
|
|
|
assert payload["apiKeyMasked"] == "********"
|
|
assert "secret-a" not in str(payload)
|
|
|
|
|
|
def test_model_pool_429_cools_down_endpoint(pool_app, monkeypatch):
|
|
from insurance.model_pool import service
|
|
from insurance.models.model_pool import ModelEndpoint
|
|
|
|
class _Response:
|
|
status_code = 429
|
|
headers = {"Retry-After": "45"}
|
|
|
|
class _RateLimitError(RuntimeError):
|
|
response = _Response()
|
|
|
|
monkeypatch.setattr(service, "redis_client", _FakeRedis())
|
|
with pool_app.app_context():
|
|
service.record_failure(1, _RateLimitError("rate limited"))
|
|
endpoint = service.db.session.get(ModelEndpoint, 1)
|
|
assert endpoint.enabled is True
|
|
assert endpoint.status == "cooldown"
|
|
assert endpoint.cooldown_until is not None
|
|
|
|
|
|
def test_llm_pool_switches_endpoint_after_429(monkeypatch):
|
|
from insurance.model_pool import service
|
|
from insurance.ppt import llm_client as module
|
|
|
|
released = []
|
|
endpoints = [
|
|
types.SimpleNamespace(
|
|
id=1,
|
|
name="first",
|
|
provider="openai",
|
|
base_url="https://one.example/v1",
|
|
model_name="model-a",
|
|
timeout_seconds=10,
|
|
pool=types.SimpleNamespace(max_attempts=2),
|
|
),
|
|
types.SimpleNamespace(
|
|
id=2,
|
|
name="second",
|
|
provider="openai",
|
|
base_url="https://two.example/v1",
|
|
model_name="model-a",
|
|
timeout_seconds=10,
|
|
pool=types.SimpleNamespace(max_attempts=2),
|
|
),
|
|
]
|
|
|
|
class _Lease:
|
|
def __init__(self, endpoint):
|
|
self.endpoint = endpoint
|
|
self.api_key = "secret"
|
|
|
|
def release(self, actual_tokens=None):
|
|
released.append((self.endpoint.id, actual_tokens))
|
|
|
|
monkeypatch.setattr(
|
|
service,
|
|
"acquire_endpoint",
|
|
lambda *_args, **_kwargs: _Lease(endpoints.pop(0)),
|
|
)
|
|
failures = []
|
|
successes = []
|
|
monkeypatch.setattr(
|
|
service, "record_failure", lambda endpoint_id, _exc: failures.append(endpoint_id)
|
|
)
|
|
monkeypatch.setattr(
|
|
service, "record_success", lambda endpoint_id: successes.append(endpoint_id)
|
|
)
|
|
calls = []
|
|
|
|
async def _fake_call(config, *_args, **_kwargs):
|
|
calls.append(config.base_url)
|
|
if len(calls) == 1:
|
|
request = httpx.Request("POST", config.base_url)
|
|
response = httpx.Response(429, request=request)
|
|
raise httpx.HTTPStatusError(
|
|
"rate limited", request=request, response=response
|
|
)
|
|
return module.LLMResponse(
|
|
content="ok",
|
|
provider=config.name,
|
|
tokens={"input": 10, "output": 2},
|
|
)
|
|
|
|
monkeypatch.setattr(module, "_call_provider", _fake_call)
|
|
client = module.LLMClient("ppt")
|
|
response = asyncio.run(
|
|
client._call_model_pool(
|
|
"ppt_extract",
|
|
[{"role": "user", "content": "hello"}],
|
|
json_mode=False,
|
|
temperature=0.2,
|
|
)
|
|
)
|
|
|
|
assert response.content == "ok"
|
|
assert failures == [1]
|
|
assert successes == [2]
|
|
assert released == [(1, None), (2, 12)]
|