55 lines
2.7 KiB
Python
55 lines
2.7 KiB
Python
|
"""One bounded, read-only metadata request to a registered loopback model service."""
|
||
|
|
import ipaddress
|
||
|
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
import re
|
||
|
|
import shlex
|
||
|
|
import urllib.error
|
||
|
|
import urllib.request
|
||
|
|
from urllib.parse import urlsplit
|
||
|
|
|
||
|
|
|
||
|
|
class NoRedirect(urllib.request.HTTPRedirectHandler):
|
||
|
|
def redirect_request(self, *args):
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def check_model(target, credential_files):
|
||
|
|
result = {'name': '模型身份', 'ok': False, 'message': '模型目录检查未完成'}
|
||
|
|
try:
|
||
|
|
url = urlsplit(target['url'])
|
||
|
|
if (url.scheme != 'http' or not ipaddress.ip_address(url.hostname).is_loopback
|
||
|
|
or url.username or url.password or url.query or url.fragment or url.path not in ('', '/')):
|
||
|
|
raise ValueError('仅允许登记的本机模型服务')
|
||
|
|
key = ''
|
||
|
|
for filename in credential_files:
|
||
|
|
path = Path(filename)
|
||
|
|
if path.stat().st_size > 262144:
|
||
|
|
raise ValueError('认证引用文件过大')
|
||
|
|
for line in path.read_text().splitlines():
|
||
|
|
match = re.match(r'^(?:export\s+)?VLLM_API_KEY=(.*)$', line.strip())
|
||
|
|
if not match:
|
||
|
|
continue
|
||
|
|
values = shlex.split(match[1], comments=True)
|
||
|
|
if len(values) > 1:
|
||
|
|
raise ValueError('认证引用格式不支持')
|
||
|
|
key = values[0] if values else ''
|
||
|
|
headers = {'Authorization': 'Bearer ' + key} if key else {}
|
||
|
|
request = urllib.request.Request(target['url'].rstrip('/') + '/v1/models', headers=headers)
|
||
|
|
opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect())
|
||
|
|
with opener.open(request, timeout=3) as response:
|
||
|
|
data = response.read(262145)
|
||
|
|
if len(data) > 262144:
|
||
|
|
raise ValueError('模型目录超过检查上限')
|
||
|
|
payload = json.loads(data)
|
||
|
|
models = payload.get('data') if isinstance(payload, dict) else None
|
||
|
|
if not isinstance(models, list) or any(not isinstance(m, dict) or not isinstance(m.get('id'), str) for m in models):
|
||
|
|
raise ValueError('模型目录格式不支持')
|
||
|
|
result['ok'] = any(m['id'] == target['model'] for m in models)
|
||
|
|
result['message'] = ('已核对模型名称：' if result['ok'] else '当前服务未提供登记的模型名称：') + target['model']
|
||
|
|
except urllib.error.HTTPError as exc:
|
||
|
|
result['message'] = f'模型目录校验失败（HTTP {exc.code}），请核对服务与认证配置；未发送推理'
|
||
|
|
except (OSError, ValueError, TypeError, KeyError):
|
||
|
|
result['message'] = '模型目录或认证引用不可用，未发送推理；请核对原服务配置'
|
||
|
|
return result
|