feat: Redis atomic balance reservation — prevent concurrent overspend
- balance.py: reserve_balance/finalize_balance/refund_balance with Lua - utils.py: get_model_max_cost from llmusage history, update_model_max_cost - init.py: register reserve_balance, finalize_balance, refund_balance on env - chat/completions/index.dspy: reserve before inference, refund on exception - DB migration: llm.max_cost DECIMAL(10,4) for historical max charge
This commit is contained in:
parent
514b674f04
commit
6351ee9b56
250
llmage/balance.py
Normal file
250
llmage/balance.py
Normal file
@ -0,0 +1,250 @@
|
||||
"""
|
||||
Redis atomic balance reservation for llmage.
|
||||
Prevents concurrent overspend: pre-deduct before inference, settle after accounting.
|
||||
|
||||
Key patterns:
|
||||
balance:{userorgid} — current pre-deducted balance (int, cents)
|
||||
reserve:{luid} — {userorgid}|{llmid}|{max_cost} (TTL 600s)
|
||||
model:max_cost:{llmid} — historical max actual customer charge (int, cents)
|
||||
"""
|
||||
import asyncio
|
||||
from appPublic.log import debug, exception
|
||||
|
||||
# ── Lua scripts ──────────────────────────────────────────────
|
||||
|
||||
RESERVE_LUA = """
|
||||
local bal_key = KEYS[1]
|
||||
local cost_key = KEYS[2]
|
||||
local reserve_key = KEYS[3]
|
||||
local max_cost = tonumber(ARGV[1])
|
||||
local db_balance = tonumber(ARGV[2]) -- 0 means no DB fallback
|
||||
local reserve_val = ARGV[3]
|
||||
|
||||
-- Get or init max_cost: use stored if higher, else set from arg
|
||||
local stored_max = redis.call('GET', cost_key)
|
||||
if stored_max then
|
||||
max_cost = math.max(max_cost, tonumber(stored_max))
|
||||
elseif max_cost > 0 then
|
||||
redis.call('SET', cost_key, max_cost)
|
||||
end
|
||||
if max_cost <= 0 then
|
||||
return {0, 'max_cost is zero'}
|
||||
end
|
||||
|
||||
-- Load or init balance
|
||||
local balance = redis.call('GET', bal_key)
|
||||
if not balance then
|
||||
if db_balance > 0 then
|
||||
balance = db_balance
|
||||
redis.call('SET', bal_key, db_balance)
|
||||
else
|
||||
return {0, 'balance not initialized'}
|
||||
end
|
||||
end
|
||||
balance = tonumber(balance)
|
||||
|
||||
-- Check and deduct
|
||||
if balance < max_cost then
|
||||
return {0, 'insufficient balance'}
|
||||
end
|
||||
redis.call('DECRBY', bal_key, max_cost)
|
||||
redis.call('SETEX', reserve_key, 600, reserve_val)
|
||||
return {1, max_cost}
|
||||
"""
|
||||
|
||||
FINALIZE_LUA = """
|
||||
local cost_key = KEYS[1]
|
||||
local reserve_key = KEYS[2]
|
||||
local actual_cost = tonumber(ARGV[1])
|
||||
|
||||
-- Atomically pop reserve
|
||||
local reserve_val = redis.call('GETDEL', reserve_key)
|
||||
if not reserve_val then
|
||||
return {0, 'no reserve'}
|
||||
end
|
||||
|
||||
-- Parse: userorgid|llmid|max_cost
|
||||
local i = 1
|
||||
local userorgid, llmid, max_cost
|
||||
for part in string.gmatch(reserve_val, '([^|]+)') do
|
||||
if i == 1 then userorgid = part
|
||||
elseif i == 2 then llmid = part
|
||||
elseif i == 3 then max_cost = tonumber(part) end
|
||||
i = i + 1
|
||||
end
|
||||
|
||||
-- Update max_cost if actual exceeds
|
||||
if actual_cost > max_cost then
|
||||
redis.call('SET', cost_key, actual_cost)
|
||||
max_cost = actual_cost
|
||||
end
|
||||
|
||||
-- Adjust balance
|
||||
local bal_key = 'balance:' .. userorgid
|
||||
local diff = max_cost - actual_cost
|
||||
if diff > 0 then
|
||||
redis.call('INCRBY', bal_key, diff)
|
||||
elseif diff < 0 then
|
||||
redis.call('DECRBY', bal_key, -diff)
|
||||
end
|
||||
|
||||
return {1, userorgid, max_cost, actual_cost, diff}
|
||||
"""
|
||||
|
||||
REFUND_LUA = """
|
||||
local reserve_key = KEYS[1]
|
||||
|
||||
local reserve_val = redis.call('GETDEL', reserve_key)
|
||||
if not reserve_val then
|
||||
return {0, 'no reserve'}
|
||||
end
|
||||
|
||||
local i = 1
|
||||
local userorgid, llmid, max_cost
|
||||
for part in string.gmatch(reserve_val, '([^|]+)') do
|
||||
if i == 1 then userorgid = part
|
||||
elseif i == 2 then llmid = part
|
||||
elseif i == 3 then max_cost = tonumber(part) end
|
||||
i = i + 1
|
||||
end
|
||||
|
||||
if max_cost > 0 then
|
||||
redis.call('INCRBY', 'balance:' .. userorgid, max_cost)
|
||||
end
|
||||
return {1, max_cost, userorgid}
|
||||
"""
|
||||
|
||||
|
||||
def _cents(f):
|
||||
return int(round(float(f) * 100))
|
||||
|
||||
|
||||
def _from_cents(c):
|
||||
return round(int(c) / 100, 2)
|
||||
|
||||
|
||||
async def _redis_eval(env, script, nkeys, *keys_and_args):
|
||||
"""Eval Lua script via env's Redis connection."""
|
||||
if not hasattr(env, 'redis') or env.redis is None:
|
||||
raise RuntimeError('Redis unavailable')
|
||||
return await asyncio.to_thread(env.redis.eval, script, nkeys, *keys_and_args)
|
||||
|
||||
|
||||
async def reserve_balance(env, llmid, userorgid, luid):
|
||||
"""Atomically check balance and deduct max_cost via Redis Lua."""
|
||||
try:
|
||||
bal_key = f'balance:{userorgid}'
|
||||
cost_key = f'model:max_cost:{llmid}'
|
||||
reserve_key = f'reserve:{luid}'
|
||||
|
||||
# Load max_cost from DB if not cached
|
||||
redis = env.redis
|
||||
max_cost = 0
|
||||
stored = await asyncio.to_thread(redis.get, cost_key)
|
||||
if stored:
|
||||
max_cost = int(stored)
|
||||
else:
|
||||
try:
|
||||
from .utils import get_model_max_cost
|
||||
mc = await get_model_max_cost(llmid)
|
||||
if mc and mc > 0:
|
||||
max_cost = _cents(mc)
|
||||
except Exception as e:
|
||||
debug(f'reserve_balance: get_model_max_cost failed: {e}')
|
||||
|
||||
if max_cost <= 0:
|
||||
return {'ok': True, 'max_cost': 0, 'no_history': True}
|
||||
|
||||
# Load balance from DB if not cached
|
||||
db_balance = 0
|
||||
stored_bal = await asyncio.to_thread(redis.get, bal_key)
|
||||
if not stored_bal:
|
||||
try:
|
||||
from accounting.getaccount import getCustomerBalance
|
||||
from sqlor.dbpools import get_sor_context
|
||||
async with get_sor_context(env, 'accounting') as sor:
|
||||
bal = await getCustomerBalance(sor, userorgid)
|
||||
if bal is not None:
|
||||
db_balance = _cents(float(bal))
|
||||
except Exception as e:
|
||||
debug(f'reserve_balance: getCustomerBalance failed: {e}')
|
||||
|
||||
reserve_val = f'{userorgid}|{llmid}|{max_cost}'
|
||||
result = await _redis_eval(env, RESERVE_LUA, 3,
|
||||
bal_key, cost_key, reserve_key,
|
||||
str(max_cost), str(db_balance), reserve_val)
|
||||
|
||||
if result[0] == 0:
|
||||
return {'ok': False, 'reason': result[1]}
|
||||
return {'ok': True, 'max_cost': _from_cents(result[1])}
|
||||
|
||||
except RuntimeError:
|
||||
debug('reserve_balance: Redis unavailable, skip')
|
||||
return {'ok': True, 'max_cost': 0, 'no_redis': True}
|
||||
except Exception as e:
|
||||
exception(f'reserve_balance error: {e}')
|
||||
return {'ok': False, 'reason': str(e)}
|
||||
|
||||
|
||||
async def finalize_balance(env, luid, actual_cost):
|
||||
"""After accounting: adjust balance, update max_cost."""
|
||||
try:
|
||||
cost_key = 'model:max_cost:dummy'
|
||||
reserve_key = f'reserve:{luid}'
|
||||
|
||||
# We need the llmid to get the correct cost_key, but we won't know it
|
||||
# until parsing the reserve. So use a placeholder and let Lua handle it.
|
||||
# Actually FINALIZE_LUA only uses cost_key for updating, so pass a dummy.
|
||||
# The real cost_key needs the llmid which comes from reserve parsing.
|
||||
|
||||
redis = env.redis
|
||||
reserve_val = await asyncio.to_thread(redis.get, reserve_key)
|
||||
if not reserve_val:
|
||||
debug(f'finalize_balance: no reserve for luid={luid}')
|
||||
return
|
||||
|
||||
# Parse to get llmid
|
||||
parts = reserve_val.split('|')
|
||||
if len(parts) < 3:
|
||||
return
|
||||
llmid = parts[1]
|
||||
cost_key = f'model:max_cost:{llmid}'
|
||||
|
||||
result = await _redis_eval(env, FINALIZE_LUA, 2,
|
||||
cost_key, reserve_key, str(_cents(actual_cost)))
|
||||
|
||||
if result[0] == 0:
|
||||
debug(f'finalize_balance: {result[1]}')
|
||||
return
|
||||
|
||||
max_c = _from_cents(result[2])
|
||||
diff = _from_cents(result[4])
|
||||
debug(f'finalize_balance: luid={luid} max={max_c} actual={actual_cost} diff={diff}')
|
||||
|
||||
# Persist max_cost to DB if actual exceeded previous max
|
||||
if actual_cost > max_c:
|
||||
try:
|
||||
from .utils import update_model_max_cost
|
||||
await update_model_max_cost(llmid, actual_cost)
|
||||
except Exception as e:
|
||||
debug(f'finalize_balance: DB persist failed: {e}')
|
||||
|
||||
except RuntimeError:
|
||||
debug('finalize_balance: Redis unavailable, skip')
|
||||
except Exception as e:
|
||||
exception(f'finalize_balance error: {e}')
|
||||
|
||||
|
||||
async def refund_balance(env, luid):
|
||||
"""Full refund on API failure."""
|
||||
try:
|
||||
reserve_key = f'reserve:{luid}'
|
||||
result = await _redis_eval(env, REFUND_LUA, 1, reserve_key)
|
||||
if result[0] == 1:
|
||||
debug(f'refund_balance: luid={luid} refund={_from_cents(result[1])}')
|
||||
else:
|
||||
debug(f'refund_balance: luid={luid} skipped ({result[1]})')
|
||||
except RuntimeError:
|
||||
pass
|
||||
except Exception as e:
|
||||
exception(f'refund_balance error: {e}')
|
||||
@ -28,6 +28,10 @@ from .utils import (
|
||||
get_llmid_cached,
|
||||
invalidate_llmid_cache,
|
||||
get_plaza_models,
|
||||
# Redis balance reservation
|
||||
reserve_balance,
|
||||
finalize_balance,
|
||||
refund_balance,
|
||||
)
|
||||
|
||||
from .llmclient import (
|
||||
@ -166,6 +170,12 @@ def load_llmage():
|
||||
env.llm_query_price = llm_query_price
|
||||
env.get_llms_by_catelog_to_customer = get_llms_by_catelog_to_customer
|
||||
env.get_plaza_models = get_plaza_models
|
||||
env._reserve_balance = reserve_balance
|
||||
env._finalize_balance = finalize_balance
|
||||
env._refund_balance = refund_balance
|
||||
env.reserve_balance = lambda llmid, userorgid, luid: reserve_balance(env, llmid, userorgid, luid)
|
||||
env.finalize_balance = lambda luid, actual_cost: finalize_balance(env, luid, actual_cost)
|
||||
env.refund_balance = lambda luid: refund_balance(env, luid)
|
||||
env.backup_accounted_llmusage = backup_accounted_llmusage
|
||||
env.read_ioinfo_content = read_ioinfo_content
|
||||
env.get_llmusage_by_id = get_llmusage_by_id
|
||||
|
||||
@ -617,3 +617,24 @@ ServerEnv().os = _os
|
||||
ServerEnv().file_to_b64 = _file_to_b64
|
||||
ServerEnv().file_mime = _file_mime
|
||||
|
||||
async def get_model_max_cost(llmid):
|
||||
"""Get historical max actual customer charge for a model from llmusage."""
|
||||
env = ServerEnv()
|
||||
async with get_sor_context(env, 'llmage') as sor:
|
||||
sql = """SELECT MAX(amount) as max_cost
|
||||
FROM llmusage
|
||||
WHERE llmid = ${llmid}$
|
||||
AND status = 'SUCCEEDED'
|
||||
AND amount IS NOT NULL
|
||||
AND amount > 0"""
|
||||
recs = await sor.sqlExe(sql, {'llmid': llmid})
|
||||
if recs and recs[0].max_cost:
|
||||
return float(recs[0].max_cost)
|
||||
return 0
|
||||
|
||||
async def update_model_max_cost(llmid, new_max):
|
||||
"""Persist updated max_cost to llm table."""
|
||||
env = ServerEnv()
|
||||
async with get_sor_context(env, 'llmage') as sor:
|
||||
await sor.U('llm', {'id': llmid, 'max_cost': float(new_max)})
|
||||
|
||||
|
||||
@ -1,13 +1,19 @@
|
||||
async def gen():
|
||||
env = request._run_ns.copy()
|
||||
f = partial(inference_generator, request, params_kw=params_kw)
|
||||
if params_kw.stream:
|
||||
async for l in f():
|
||||
yield f'data: {l}\n'
|
||||
yield 'data: [DONE]\n\n'
|
||||
else:
|
||||
async for l in f():
|
||||
yield l
|
||||
try:
|
||||
if params_kw.stream:
|
||||
async for l in f():
|
||||
yield f'data: {l}\n'
|
||||
yield 'data: [DONE]\n\n'
|
||||
else:
|
||||
async for l in f():
|
||||
yield l
|
||||
except Exception as e:
|
||||
luid = params_kw.get('_luid')
|
||||
if luid:
|
||||
await env.refund_balance(luid)
|
||||
raise
|
||||
|
||||
debug_params('params_kw', params_kw)
|
||||
catelogid = params_kw.catelogid or 't2t'
|
||||
@ -37,9 +43,16 @@ params_kw.llmid = llmid
|
||||
params_kw.llmcatelogid = catelogid
|
||||
|
||||
debug(f'{params_kw.llmid=}')
|
||||
luid = getID()
|
||||
reserved = await env.reserve_balance(params_kw.llmid, userorgid, luid)
|
||||
if not reserved.get('ok'):
|
||||
debug(f'{userid=} balance not enough: {reserved}')
|
||||
return openai_429()
|
||||
params_kw._luid = luid
|
||||
params_kw._reserved = reserved
|
||||
f = await checkCustomerBalance(params_kw.llmid, userid, userorgid)
|
||||
if not f:
|
||||
debug(f'{userid=} balance not enough')
|
||||
if not f and reserved.get('no_redis'):
|
||||
debug(f'{userid=} balance not enough (DB fallback)')
|
||||
return openai_429()
|
||||
# debug(f'{tools=}, {request._run_ns.tools=}')
|
||||
return await env.stream_response(request, gen, content_type='application/json')
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user