fix: stabilize background operations and large embedding results
This commit is contained in:
@@ -0,0 +1,11 @@
|
||||
"""Bound embedding result frames so large notes do not exceed pipe line limits."""
|
||||
import json
|
||||
|
||||
|
||||
def response_lines(response, operation):
|
||||
if operation == 'embedding' and 'result' in response and 'error_code' not in response:
|
||||
vectors = response['result']
|
||||
for offset in range(0, len(vectors), 128):
|
||||
yield json.dumps({'embedding_offset': offset, 'embedding_chunk': vectors[offset:offset + 128]}, allow_nan=False) + '\n'
|
||||
response = {**response, 'result': [], 'embedding_count': len(vectors)}
|
||||
yield json.dumps(response, ensure_ascii=False, allow_nan=False) + '\n'
|
||||
@@ -189,15 +189,30 @@ class Runtime:
|
||||
await process.stdin.drain()
|
||||
process.stdin.close()
|
||||
final = None
|
||||
vectors = []
|
||||
while line := await process.stdout.readline():
|
||||
message = json.loads(line)
|
||||
if "progress" in message:
|
||||
if "embedding_chunk" in message:
|
||||
chunk = message['embedding_chunk']
|
||||
if (operation != 'embedding' or not isinstance(chunk, list)
|
||||
or message.get('embedding_offset') != len(vectors)
|
||||
or len(vectors) + len(chunk) > len(payload.get('texts', []))):
|
||||
raise ProviderError('LOCAL_MODEL_INVALID_RESPONSE', '本地向量传输顺序或数量无效。')
|
||||
vectors.extend(chunk)
|
||||
elif "progress" in message:
|
||||
callback = runtime_progress.get()
|
||||
if callback:
|
||||
callback(message)
|
||||
else:
|
||||
final = message
|
||||
await process.wait()
|
||||
if isinstance(final, dict) and 'embedding_count' in final:
|
||||
if (final['embedding_count'] != len(vectors)
|
||||
or len(vectors) != len(payload.get('texts', []))):
|
||||
raise ProviderError('LOCAL_MODEL_INVALID_RESPONSE', '本地向量传输不完整。')
|
||||
final['result'] = vectors
|
||||
elif vectors:
|
||||
raise ProviderError('LOCAL_MODEL_INVALID_RESPONSE', '本地向量传输缺少结束标记。')
|
||||
return final
|
||||
try:
|
||||
result = await asyncio.wait_for(receive(), config.timeout_seconds)
|
||||
|
||||
@@ -216,4 +216,6 @@ if __name__ == "__main__":
|
||||
response = {"error_code": "LOCAL_INFERENCE_FAILED", "message": "本地推理失败,请检查媒体格式、模型和设备配置。"}
|
||||
if "error_code" in response:
|
||||
response["diagnostics"] = {"requested_device": request["config"]["device"], "actual_device": request.get("_actual_device", "unknown")}
|
||||
sys.stdout.buffer.write((json.dumps(response, ensure_ascii=False, allow_nan=False) + "\n").encode("utf-8"))
|
||||
from protocol import response_lines
|
||||
for line in response_lines(response, request['operation']):
|
||||
sys.stdout.buffer.write(line.encode('utf-8'))
|
||||
|
||||
Reference in New Issue
Block a user