yumoqing/llmage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(balance): extend atomic balance reserve to async+sync+product paths

Browse Source 
- balance.py: module-level redis.asyncio singleton (works in all
  processes without env.redis injection), centralized tpac/self-org
  skip policy, userid param, extend_reserve, invalidate_balance_cache,
  finalize updates max_cost on cold start (no reserve)
- llmclient.py: unified reserve injection in _inference_generator
  (covers all 22 dspy entries); fix luid mismatch (reserve key ==
  llmusage id == finalize key); stream exception refund
- syncinference.py / asyncinference.py: luid reuse + failure refunds
  (submit failure, outer exception, query_task_status FAILED)
- product_interface.py: reserve + refund on product stream/non-stream
- accounting.py: finalize after llm_accounting (was missing for
  direct llmage accounting path -> reserve leaked until TTL)
- init.py: lambda signatures synced
- v1/chat/completions/index.dspy: pass userid to reserve
This commit is contained in:
yumoqing 2026-08-04 18:35:53 +08:00
parent 89e6d7cb25
commit e3f61f446d
8 changed files with 264 additions and 66 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
llmage/accounting.py
View File

@ -195,6 +195,13 @@ async def llm_accounting(llmusage):
'accounting_status': 'accounted'
}
await sor.U('llmusage', ns)
# Finalize Redis balance reservation (non-critical)
try:
from .balance import finalize_balance
await finalize_balance(env, llmusage.id, trans_amount,
llmid=llmusage.llmid)
except Exception as e:
exception(f'finalize_balance failed (non-critical): {e}')
async def get_accounting_llmusages(luid=None):
env = ServerEnv()

21
llmage/asyncinference.py
View File

@ -13,6 +13,7 @@ from appPublic.base64_to_file import base64_to_file, getFilenameFromBase64
from ahserver.serverenv import get_serverenv, ServerEnv
from ahserver.filestorage import FileStorage
from .accounting import llm_accounting, llm_charging
from .balance import refund_balance
from .utils import *
# Global set to keep references to background tasks
@ -67,7 +68,7 @@ async def async_uapi_request(request, llm,
uapi = env.UpAppApi(request)
userid = await env.uapi_data.get_calluserid(llm.upappid, orgid=llm.ownerid)
b = None
luid = getID()
luid = params_kw.get('_luid') or getID()
try:
start_timestamp = time.time()
if llm.callbackurl:
@ -77,6 +78,12 @@ async def async_uapi_request(request, llm,
try:
b = await uapi.call(llm.upappid, llm.apiname, userid, params=params_kw)
except Exception as e:
# Refund balance reservation on submission failure
try:
if luid:
await refund_balance(ServerEnv(), luid)
except Exception:
pass
estr = erase_apikey(e)
ed = {"error": f"ERROR:{estr}", "status": "FAILED"}
exception(f'{ed}')
@ -123,6 +130,12 @@ async def async_uapi_request(request, llm,
task.add_done_callback(_background_tasks.discard)
except Exception as e:
# Refund balance reservation on outer failure
try:
if luid:
await refund_balance(ServerEnv(), luid)
except Exception:
pass
ed = {"error": f"ERROR:{e}", "status": "FAILED"}
s = json.dumps(ed, ensure_ascii=False)
s = ''.join(s.split('\n'))
@ -229,6 +242,12 @@ async def query_task_status(request, luid, onetime=False):
ns['usages'] = json.dumps(new_output['usage'])
await append_new_llmoutput(llmusage.ioinfo, new_output)
await modify_llmusage(ns)
if llmusage.status == 'FAILED':
# Async task failed — refund the balance reservation
try:
await refund_balance(ServerEnv(), luid)
except Exception:
pass
if llmusage.status in ['UNKNOWN', 'FAILED', 'SUCCEEDED']:
critical(f'finished .. {llmusage.status=}')
return

214
llmage/balance.py
View File

@ -2,14 +2,54 @@
Redis atomic balance reservation for llmage.
Prevents concurrent overspend: pre-deduct before inference, settle after accounting.
Uses a module-level redis.asyncio singleton so the atomic reserve works in
EVERY process (web workers + backend_accounting) without needing env.redis
to be injected. This fixes the previous fatal issue where env.redis was never
set, so reserve always silently fell back to the non-atomic DB check.
Key patterns:
balance:{userorgid} — current pre-deducted balance (int, cents)
reserve:{luid} — {userorgid}|{llmid}|{max_cost} (TTL 600s)
reserve:{luid} — {userorgid}|{llmid}|{max_cost} (TTL)
model:max_cost:{llmid} — historical max actual customer charge (int, cents)
Reserve TTL: 600s for fast (stream/sync) calls, 3600s for async tasks that
may run for minutes before finalize/refund.
"""
import asyncio
from appPublic.log import debug, exception
# ── Module-level async Redis singleton ───────────────────────
_redis = None
_redis_lock = asyncio.Lock()
def _redis_url():
try:
from appPublic.jsonConfig import getConfig
config = getConfig()
url = getattr(config.website, 'session_redis', None)
if url:
u = getattr(url, 'url', None)
if u:
return u
except Exception:
pass
return "redis://127.0.0.1:6379"
async def _get_redis():
"""Return a shared redis.asyncio client (lazy singleton)."""
global _redis
if _redis is not None:
return _redis
async with _redis_lock:
if _redis is None:
import redis.asyncio as aioredis
_redis = await aioredis.from_url(
_redis_url(), decode_responses=True)
return _redis
# ── Lua scripts ──────────────────────────────────────────────
RESERVE_LUA = """
@ -19,6 +59,7 @@ 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]
local ttl = tonumber(ARGV[4])
-- Get or init max_cost: use stored if higher, else set from arg
local stored_max = redis.call('GET', cost_key)
@ -48,7 +89,7 @@ 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)
redis.call('SETEX', reserve_key, ttl, reserve_val)
return {1, max_cost}
"""
@ -123,24 +164,69 @@ 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 _update_max_cost(llmid, actual_cost):
"""Bump model:max_cost if actual exceeds stored value; persist to DB.
async def reserve_balance(env, llmid, userorgid, luid):
"""Atomically check balance and deduct max_cost via Redis Lua."""
Keeps the pre-deduct baseline growing even when the original reserve was
skipped (no_history cold start) — otherwise reserve would never activate.
"""
try:
redis = await _get_redis()
cost_key = f'model:max_cost:{llmid}'
cost_cents = _cents(actual_cost)
stored = await redis.get(cost_key)
if stored and int(stored) >= cost_cents:
return
await redis.set(cost_key, cost_cents)
from .utils import update_model_max_cost
await update_model_max_cost(llmid, actual_cost)
except Exception as e:
debug(f'_update_max_cost failed: {e}')
async def reserve_balance(env, llmid, userorgid, luid, ttl=600, userid=None):
"""Atomically check balance and deduct max_cost via Redis Lua.
Policy (centralized so every entry behaves the same):
- Self-owned org (llm.ownerid == userorgid): skip reserve.
- tpac user (external balance system): skip reserve.
- No pricing (ppid empty): skip reserve (availability handled elsewhere).
Returns:
{'ok': True, 'max_cost': X} — reserved X
{'ok': True, 'max_cost': 0, 'skip': ...} — reserve skipped by policy
{'ok': True, 'max_cost': 0, 'no_history': True} — no max_cost data yet
{'ok': False, 'reason': ...} — insufficient balance
{'ok': True, 'max_cost': 0, 'no_redis': True} — Redis down, DB fallback
"""
try:
# ── Policy checks (need llm info) ──
try:
from .utils import get_llmage_llm, get_user_tpac
llm = await get_llmage_llm(llmid)
if llm and llm.ownerid == userorgid:
return {'ok': True, 'max_cost': 0, 'skip': 'self_org'}
if not llm or not llm.ppid:
return {'ok': True, 'max_cost': 0, 'skip': 'no_ppid'}
# tpac user: balance lives in external system, skip redis reserve
if userid:
try:
tpac = await get_user_tpac(userid)
if tpac:
return {'ok': True, 'max_cost': 0, 'skip': 'tpac'}
except Exception as e:
debug(f'reserve_balance: tpac check failed: {e}')
except Exception as e:
debug(f'reserve_balance: policy check failed: {e}')
redis = await _get_redis()
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
# Load max_cost from Redis if cached, else from DB history
max_cost = 0
stored = await asyncio.to_thread(redis.get, cost_key)
stored = await redis.get(cost_key)
if stored:
max_cost = int(stored)
else:
@ -151,13 +237,13 @@ async def reserve_balance(env, llmid, userorgid, luid):
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
# Load balance from DB if not cached in Redis
db_balance = 0
stored_bal = await asyncio.to_thread(redis.get, bal_key)
stored_bal = await redis.get(bal_key)
if not stored_bal:
try:
from accounting.getaccount import getCustomerBalance
@ -168,70 +254,62 @@ async def reserve_balance(env, llmid, userorgid, luid):
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,
result = await redis.eval(RESERVE_LUA, 3,
bal_key, cost_key, reserve_key,
str(max_cost), str(db_balance), reserve_val)
str(max_cost), str(db_balance), reserve_val, str(ttl))
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 not configured, skip')
return {'ok': True, 'max_cost': 0, 'no_redis': True}
except Exception as e:
# Redis down or unreachable — fall back to DB check, don't block
debug(f'reserve_balance: Redis error, falling back to DB: {e}')
return {'ok': True, 'max_cost': 0, 'no_redis': True}
async def finalize_balance(env, luid, actual_cost):
"""After accounting: adjust balance, update max_cost."""
async def finalize_balance(env, luid, actual_cost, llmid=None):
"""After accounting: adjust balance, update max_cost.
llmid: needed for cold-start max_cost updates when no reserve existed
(reserve skipped due to no_history) — keeps the pre-deduct baseline
growing so future calls can reserve.
"""
try:
cost_key = 'model:max_cost:dummy'
redis = await _get_redis()
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)
reserve_val = await redis.get(reserve_key)
if not reserve_val:
debug(f'finalize_balance: no reserve for luid={luid}')
if llmid:
await _update_max_cost(llmid, actual_cost)
return
# Parse to get llmid
# Parse to get llmid (reserve_val: userorgid|llmid|max_cost)
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,
result = await redis.eval(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')
await _update_max_cost(llmid, actual_cost)
except Exception as e:
exception(f'finalize_balance error: {e}')
@ -239,13 +317,39 @@ async def finalize_balance(env, luid, actual_cost):
async def refund_balance(env, luid):
"""Full refund on API failure."""
try:
redis = await _get_redis()
reserve_key = f'reserve:{luid}'
result = await _redis_eval(env, REFUND_LUA, 1, reserve_key)
result = await redis.eval(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}')
async def extend_reserve(env, luid, ttl):
"""Extend TTL of an existing reserve.
Used when an entry pre-reserved with the short (600s) TTL but the
request is dispatched to async mode, where the task may run longer.
"""
try:
redis = await _get_redis()
await redis.expire(f'reserve:{luid}', ttl)
except Exception as e:
debug(f'extend_reserve failed: {e}')
async def invalidate_balance_cache(userorgid):
"""Delete the cached Redis balance so the next reserve reloads from DB.
Must be called after recharge / recharge reversal, otherwise the
pre-deduct baseline stays stale and valid requests get rejected.
"""
try:
redis = await _get_redis()
await redis.delete(f'balance:{userorgid}')
debug(f'invalidate_balance_cache: balance:{userorgid} deleted')
except Exception as e:
debug(f'invalidate_balance_cache failed: {e}')

4
llmage/init.py
View File

@ -173,8 +173,8 @@ def load_llmage():
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.reserve_balance = lambda llmid, userorgid, luid, ttl=600, userid=None: reserve_balance(env, llmid, userorgid, luid, ttl=ttl, userid=userid)
env.finalize_balance = lambda luid, actual_cost, llmid=None: finalize_balance(env, luid, actual_cost, llmid=llmid)
env.refund_balance = lambda luid: refund_balance(env, luid)
env.backup_accounted_llmusage = backup_accounted_llmusage
env.read_ioinfo_content = read_ioinfo_content

27
llmage/llmclient.py
View File

@ -14,7 +14,7 @@ from ahserver.filestorage import FileStorage
from .asyncinference import async_uapi_request
from .syncinference import sync_uapi_request
from .accounting import llm_accounting, llm_charging
from .balance import refund_balance
from .balance import refund_balance, reserve_balance, extend_reserve
from .utils import *
async def uapi_request(request, llm, callerid, callerorgid, params_kw=None):
@ -27,7 +27,7 @@ async def uapi_request(request, llm, callerid, callerorgid, params_kw=None):
userid = await env.uapi_data.get_calluserid(llm.upappid, orgid=llm.ownerid)
outlines = []
txt = ''
luid = getID()
luid = params_kw.get('_luid') or getID()
llmusage = None
try:
start_timestamp = time.time()
@ -161,6 +161,29 @@ async def _inference_generator(request, callerid, callerorgid,
yield errmsg
return
params_kw.model = llm.model
# ── Unified balance reserve ────────────────────────────────────────
# Covers every entry that did not pre-reserve. v1/chat/completions
# reserves at dspy level and passes _luid; everyone else reserves here.
# Reusing _luid downstream also fixes the old mismatch where the
# reserve key could never equal the llmusage id used by finalize.
if params_kw.get('_luid'):
if llm.stream == 'async':
# Entry reserved with the short TTL; async tasks run longer
await extend_reserve(env, params_kw._luid, 3600)
else:
_luid = getID()
_ttl = 3600 if llm.stream == 'async' else 600
reserved = await reserve_balance(env, llm.id, callerorgid, _luid,
ttl=_ttl, userid=callerid)
if not reserved.get('ok'):
debug(f'balance reserve rejected: {reserved}')
errmsg = json.dumps({'status': 'FAILED',
'error': f'余额不足(balance reserve rejected): {reserved.get("reason")}'},
ensure_ascii=False) + '\n'
yield errmsg
return
params_kw._luid = _luid
params_kw._reserved = reserved
if llm.stream == 'async':
if llm.callbackurl:
cb_url = env.entire_url(llm.callbackurl)

46
llmage/product_interface.py
View File

@ -22,6 +22,7 @@ from .utils import (
write_llmusage,
)
from .accounting import llm_charging
from .balance import reserve_balance, refund_balance
async def _resolve_llm(resource_ref_id):
@ -219,24 +220,46 @@ async def execute_product_service(resource_ref_id, user_id, user_org_id, request
if not params_kw.get('transno'):
params_kw.transno = luid
# ── Balance reservation (product path) ─────────────────────────
# Same atomic pre-deduct as the direct dspy path; policy checks
# (self-org / tpac / no-ppid) live inside reserve_balance.
ttl = 3600 if full_llm.stream == 'async' else 600
reserved = await reserve_balance(env, full_llm.id, user_org_id, luid,
ttl=ttl, userid=user_id)
if not reserved.get('ok'):
return {'success': False,
'message': f'余额不足(balance reserve rejected): {reserved.get("reason")}',
'status': 'FAILED', 'task_id': luid}
params_kw._luid = luid
try:
if full_llm.stream == 'async':
from .asyncinference import async_uapi_request_product
result = await async_uapi_request_product(
full_llm, userid, user_id, user_org_id, params_kw, luid)
return result
elif not full_llm.stream:
from .syncinference import sync_uapi_request_product
result = await sync_uapi_request_product(
full_llm, userid, user_id, user_org_id, params_kw, luid)
return result
else:
result = await _collect_stream(full_llm, userid, user_id,
user_org_id, params_kw, luid)
return result
# Sub-callers swallow exceptions and return success=False —
# refund here too (refund is idempotent via GETDEL)
if not result.get('success'):
try:
await refund_balance(env, luid)
except Exception:
pass
return result
except Exception as e:
# Refund balance reservation on failure
try:
await refund_balance(env, luid)
except Exception:
pass
exception(f'execute_product_service error: {e}')
return {'success': False, 'message': str(e), 'status': 'FAILED',
'task_id': luid}
@ -349,6 +372,16 @@ async def execute_product_service_stream(resource_ref_id, user_id, user_org_id,
if not params_kw.get('transno'):
params_kw.transno = luid
# ── Balance reservation (product stream path) ──────────────────
reserved = await reserve_balance(env, full_llm.id, user_org_id, luid,
ttl=600, userid=user_id)
if not reserved.get('ok'):
yield {'chunk': None, 'usage_data': None, 'done': True,
'error': f'余额不足(balance reserve rejected): {reserved.get("reason")}',
'task_id': luid, 'status': 'FAILED'}
return
params_kw._luid = luid
outlines = []
txt = ''
usage = None
@ -416,6 +449,11 @@ async def execute_product_service_stream(resource_ref_id, user_id, user_org_id,
}
except Exception as e:
# Refund balance reservation on failure
try:
await refund_balance(env, luid)
except Exception:
pass
exception(f'stream error: {e}')
yield {'chunk': None, 'usage_data': None, 'done': True,
'error': str(e), 'task_id': luid, 'status': 'FAILED'}

9
llmage/syncinference.py
View File

@ -14,6 +14,7 @@ from appPublic.base64_to_file import base64_to_file, getFilenameFromBase64
from ahserver.serverenv import get_serverenv, ServerEnv
from ahserver.filestorage import FileStorage
from .accounting import llm_accounting, llm_charging
from .balance import refund_balance
from .utils import *
async def sync_uapi_request(request, llm, callerid, callerorgid, params_kw=None):
@ -28,7 +29,7 @@ async def sync_uapi_request(request, llm, callerid, callerorgid, params_kw=None)
outlines = []
b = None
d = None
luid = getID()
luid = params_kw.get('_luid') or getID()
try:
start_timestamp = time.time()
responsed_seconds = None
@ -69,6 +70,12 @@ async def sync_uapi_request(request, llm, callerid, callerorgid, params_kw=None)
yield b
# await write_llmusage(llmusage)
except Exception as e:
# Refund balance reservation on failure
try:
if luid:
await refund_balance(ServerEnv(), luid)
except Exception:
pass
exception(f'{e=},{format_exc()}, {b=}')
estr = erase_apikey(e)
ed = {"error": f"ERROR:{estr}", "status": "FAILED" ,"llmusageid": luid}

2
wwwroot/v1/chat/completions/index.dspy
View File

@ -49,7 +49,7 @@ params_kw.llmcatelogid = catelogid
debug(f'{params_kw.llmid=}')
luid = getID()
reserved = await env.reserve_balance(params_kw.llmid, userorgid, luid)
reserved = await env.reserve_balance(params_kw.llmid, userorgid, luid, userid=userid)
if not reserved.get('ok'):
debug(f'{userid=} balance not enough: {reserved}')
return openai_429()