forked from luopingyi/xc_agent_platform_demo
66 lines
3.0 KiB
Python
66 lines
3.0 KiB
Python
|
"""Fresh authenticated queue snapshots. Never expose response bodies or credentials."""
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import urllib.request
|
||
|
|
from .task_add_client import resolve_token
|
||
|
|
|
||
|
|
TERMINAL = {'success', 'failed', 'fail', 'error', 'cancelled', 'canceled'}
|
||
|
|
|
||
|
|
|
||
|
|
def model_key(value):
|
||
|
|
value = str(value or '').strip().lower().rstrip('/')
|
||
|
|
for prefix in ('https://huggingface.co/', 'http://huggingface.co/', 'https://modelscope.cn/models/'):
|
||
|
|
if value.startswith(prefix):
|
||
|
|
value = value[len(prefix):]
|
||
|
|
return value
|
||
|
|
|
||
|
|
|
||
|
|
def read_queue():
|
||
|
|
"""Return (count, reason, model identities); count None means fail closed."""
|
||
|
|
try:
|
||
|
|
token = resolve_token()
|
||
|
|
except (ValueError, OSError, UnicodeError):
|
||
|
|
return None, 'missing/invalid_token', set()
|
||
|
|
count, models, seen, total = 0, set(), 0, None
|
||
|
|
for page in range(1, 21):
|
||
|
|
try:
|
||
|
|
request = urllib.request.Request(
|
||
|
|
os.environ.get('MODELHUB_API_BASE', 'https://modelhub.org.cn').rstrip('/')
|
||
|
|
+ f'/api/adapt/task/page?onlyMine=true&pageSize=100¤t={page}',
|
||
|
|
headers={'Xc-Token': token, 'Accept': 'application/json'}, method='GET')
|
||
|
|
with urllib.request.urlopen(request, timeout=30) as response:
|
||
|
|
if not 200 <= response.status < 300:
|
||
|
|
return None, 'http_error', set()
|
||
|
|
body = json.loads(response.read())
|
||
|
|
except (ValueError, UnicodeError):
|
||
|
|
return None, 'malformed_json', set()
|
||
|
|
except OSError:
|
||
|
|
return None, 'transport/http_error', set()
|
||
|
|
if not isinstance(body, dict) or type(body.get('code')) not in (int, str) or str(body['code']) != '0':
|
||
|
|
return None, 'platform_error', set()
|
||
|
|
data = body.get('data')
|
||
|
|
if not isinstance(data, dict) or type(data.get('total')) is not int or data['total'] < 0:
|
||
|
|
return None, 'malformed_data', set()
|
||
|
|
records = data.get('records')
|
||
|
|
if not isinstance(records, list) or any(not isinstance(r, dict) for r in records):
|
||
|
|
return None, 'malformed_records', set()
|
||
|
|
if total is not None and total != data['total']:
|
||
|
|
return None, 'pagination_total_changed', set()
|
||
|
|
total = data['total']
|
||
|
|
seen += len(records)
|
||
|
|
if seen > total:
|
||
|
|
return None, 'pagination_inconsistent', set()
|
||
|
|
for record in records:
|
||
|
|
count += str(record.get('status', '')).strip().lower() not in TERMINAL
|
||
|
|
for field in ('modelAddress', 'model_address'):
|
||
|
|
if record.get(field):
|
||
|
|
models.add(model_key(record[field]))
|
||
|
|
if record.get('modelAuthorName') and record.get('modelId'):
|
||
|
|
models.add(model_key(f"{record['modelAuthorName']}/{record['modelId']}"))
|
||
|
|
if seen == total:
|
||
|
|
# Do not retain arbitrary response fields, which may echo credentials.
|
||
|
|
return count, 'ok', {m for m in models if token.lower() not in m}
|
||
|
|
if not records:
|
||
|
|
break
|
||
|
|
return None, 'incomplete_pagination', set()
|