chore: add debug script for volcengine sync
This commit is contained in:
@@ -0,0 +1,203 @@
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import json
|
||||
import hmac
|
||||
import hashlib
|
||||
import datetime
|
||||
import requests
|
||||
from urllib.parse import urlparse
|
||||
|
||||
# ==========================================
|
||||
# 🛠️ Configuration
|
||||
# ==========================================
|
||||
# Volcengine API Config
|
||||
ACCESS_KEY = os.getenv(
|
||||
"VOLC_ACCESS_KEY", "AKLTNGYxZmUyNzNmZjk3NGMyZmEzNjVjZmY4NGM5ZjcxMjU"
|
||||
)
|
||||
SECRET_KEY = os.getenv(
|
||||
"VOLC_SECRET_KEY", "Tm1VNE5XSmlORGs1WXpGbU5EVmtabUUyT1RZeVltTTBNalkxTWpjeE1tTQ=="
|
||||
)
|
||||
HOST = "open.volcengineapi.com"
|
||||
REGION = "cn-beijing"
|
||||
BUFFER = 300000
|
||||
|
||||
# Suffixes to clean from Volcengine response
|
||||
SUFFIXES_TO_REMOVE = [
|
||||
"_data_collaboration_resource_pack",
|
||||
"协作奖励计划资源包",
|
||||
"免费在线推理资源包",
|
||||
"免费在线推理包",
|
||||
"资源包",
|
||||
]
|
||||
|
||||
# Volcengine Name -> LiteLLM Model Name Map
|
||||
VOLC_MODEL_MAP = {
|
||||
"doubao-seed-1-6": ["doubao-seed-1-6-250615", "doubao-seed-1-6-251015"],
|
||||
"doubao-seed-1-8": "doubao-seed-1-8-251228",
|
||||
"doubao-seed-code": "doubao-seed-code-preview-251028",
|
||||
"doubao-seed-translation": "doubao-seed-translation-250915",
|
||||
"deepseek-v3": "deepseek-v3-250324",
|
||||
"deepseek-v3-1": "deepseek-v3-1-terminus",
|
||||
"deepseek-v3-2": "deepseek-v3-2-251201",
|
||||
"deepseek-r1": "deepseek-r1-250528",
|
||||
"kimi-k2": "kimi-k2-thinking-251104",
|
||||
"glm-4_7": "glm-4-7-251222",
|
||||
}
|
||||
|
||||
|
||||
# ==========================================
|
||||
# 🧮 Utility Functions
|
||||
# ==========================================
|
||||
def normalize_volc_name(raw_name):
|
||||
"""Volcengine API response cleaning."""
|
||||
if not raw_name:
|
||||
return "unknown"
|
||||
clean_name = raw_name
|
||||
for suffix in sorted(SUFFIXES_TO_REMOVE, key=len, reverse=True):
|
||||
if clean_name.endswith(suffix):
|
||||
clean_name = clean_name.replace(suffix, "")
|
||||
break
|
||||
clean_name = clean_name.replace("_", "-").replace(".", "-")
|
||||
clean_name = clean_name.lower().strip()
|
||||
return VOLC_MODEL_MAP.get(clean_name, clean_name)
|
||||
|
||||
|
||||
def sign(key, msg):
|
||||
return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest()
|
||||
|
||||
|
||||
def get_signature_key(key, date_stamp, region_name, service_name):
|
||||
k_date = sign(key.encode("utf-8"), date_stamp)
|
||||
k_region = sign(k_date, region_name)
|
||||
k_service = sign(k_region, service_name)
|
||||
k_signing = sign(k_service, "request")
|
||||
return k_signing
|
||||
|
||||
|
||||
def get_signed_headers(action, version, body_dict, service="billing"):
|
||||
t = datetime.datetime.utcnow()
|
||||
amz_date = t.strftime("%Y%m%dT%H%M%SZ")
|
||||
date_stamp = t.strftime("%Y%m%d")
|
||||
request_body = json.dumps(body_dict)
|
||||
payload_hash = hashlib.sha256(request_body.encode("utf-8")).hexdigest()
|
||||
|
||||
canonical_headers = f"content-type:application/json\nhost:{HOST}\nx-content-sha256:{payload_hash}\nx-date:{amz_date}\n"
|
||||
signed_headers = "content-type;host;x-content-sha256;x-date"
|
||||
canonical_request = f"POST\n/\nAction={action}&Version={version}\n{canonical_headers}\n{signed_headers}\n{payload_hash}"
|
||||
|
||||
credential_scope = f"{date_stamp}/{REGION}/{service}/request"
|
||||
string_to_sign = f"HMAC-SHA256\n{amz_date}\n{credential_scope}\n{hashlib.sha256(canonical_request.encode('utf-8')).hexdigest()}"
|
||||
|
||||
signing_key = get_signature_key(SECRET_KEY, date_stamp, REGION, service)
|
||||
signature = hmac.new(
|
||||
signing_key, string_to_sign.encode("utf-8"), hashlib.sha256
|
||||
).hexdigest()
|
||||
|
||||
authorization_header = f"HMAC-SHA256 Credential={ACCESS_KEY}/{credential_scope}, SignedHeaders={signed_headers}, Signature={signature}"
|
||||
|
||||
return {
|
||||
"Content-Type": "application/json",
|
||||
"Host": HOST,
|
||||
"X-Date": amz_date,
|
||||
"X-Content-Sha256": payload_hash,
|
||||
"Authorization": authorization_header,
|
||||
}, request_body
|
||||
|
||||
|
||||
# ==========================================
|
||||
# 🔄 Fetch & Process Logic
|
||||
# ==========================================
|
||||
def fetch_all_packages():
|
||||
print(f"Fetching packages from {HOST}...")
|
||||
action = "ListResourcePackages"
|
||||
version = "2022-01-01"
|
||||
all_packages = []
|
||||
current_token = ""
|
||||
|
||||
while True:
|
||||
try:
|
||||
payload = {
|
||||
"ResourceType": "Package",
|
||||
"MaxResults": "20",
|
||||
"NextToken": current_token,
|
||||
}
|
||||
headers, body = get_signed_headers(action, version, payload)
|
||||
url = f"https://{HOST}/?Action={action}&Version={version}"
|
||||
response = requests.post(url, data=body, headers=headers, timeout=60)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
items = data.get("Result", {}).get("List", [])
|
||||
print(f"Fetched {len(items)} items in this page.")
|
||||
if items:
|
||||
all_packages.extend(items)
|
||||
next_token = data.get("Result", {}).get("NextToken", "")
|
||||
if not next_token or next_token == current_token:
|
||||
break
|
||||
current_token = next_token
|
||||
time.sleep(0.2)
|
||||
except Exception as e:
|
||||
print(f"Error fetching packages: {e}")
|
||||
break
|
||||
return all_packages
|
||||
|
||||
|
||||
def process_packages_breakdown(packages):
|
||||
model_counts = {}
|
||||
print("\n--- Raw Package Processing ---")
|
||||
for pkg in packages:
|
||||
raw_code = pkg.get("ConfigurationCode", "")
|
||||
raw_name = pkg.get("ConfigurationName", "Unknown")
|
||||
target_raw = raw_code if raw_code else raw_name
|
||||
model_name = normalize_volc_name(target_raw)
|
||||
|
||||
available = pkg.get("AvailableAmount", 0)
|
||||
unit = pkg.get("Unit", "token")
|
||||
|
||||
print(
|
||||
f"Pkg: Code='{raw_code}', Name='{raw_name}' -> Norm='{model_name}', Avail={available} {unit}"
|
||||
)
|
||||
|
||||
try:
|
||||
if unit.lower() == "token":
|
||||
model_counts[model_name] = model_counts.get(model_name, 0) + int(
|
||||
available
|
||||
)
|
||||
except ValueError:
|
||||
continue
|
||||
return model_counts
|
||||
|
||||
|
||||
# ==========================================
|
||||
# ▶️ Main Execution
|
||||
# ==========================================
|
||||
if __name__ == "__main__":
|
||||
print(f"Debug Script for Volcengine Sync")
|
||||
print(f"Access Key: {ACCESS_KEY[:4]}...{ACCESS_KEY[-4:]}")
|
||||
|
||||
packages = fetch_all_packages()
|
||||
print(f"\nTotal Packages Fetched: {len(packages)}")
|
||||
|
||||
model_balances = process_packages_breakdown(packages)
|
||||
|
||||
print("\n--- Final Model Balances ---")
|
||||
for model, balance in model_balances.items():
|
||||
remaining = balance - BUFFER
|
||||
print(f"Model: {model}")
|
||||
print(f" Balance : {balance:,}")
|
||||
print(f" Buffer : {BUFFER:,}")
|
||||
print(f" Remaining: {remaining:,}")
|
||||
|
||||
# Simulate Redis Keys
|
||||
targets = []
|
||||
if model in VOLC_MODEL_MAP:
|
||||
val = VOLC_MODEL_MAP[model]
|
||||
targets = val if isinstance(val, list) else [val]
|
||||
if not targets:
|
||||
targets = [model]
|
||||
else:
|
||||
targets = [model]
|
||||
|
||||
print(f" Redis Keys that would be set:")
|
||||
for t in targets:
|
||||
print(f" - limit:dynamic:volcengine:remaining:{t} = {int(remaining)}")
|
||||
Reference in New Issue
Block a user