343 lines
10 KiB
Python
343 lines
10 KiB
Python
import logging
|
|
import os
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from app.faster_qwen_worker import (
|
|
DEFAULT_CPU_AFFINITY_COUNT,
|
|
DEFAULT_CPU_THREADS,
|
|
MAX_NEW_TOKENS,
|
|
FasterQwenWorkerClient,
|
|
FasterQwenWorkerTimeout,
|
|
faster_qwen_worker_main,
|
|
_generation_kwargs,
|
|
_configure_worker_environment,
|
|
)
|
|
|
|
|
|
class _FakeConnection:
|
|
def __init__(self, responses):
|
|
self.responses = list(responses)
|
|
self.sent = []
|
|
self.closed = False
|
|
|
|
def poll(self, _timeout):
|
|
return bool(self.responses)
|
|
|
|
def recv(self):
|
|
return self.responses.pop(0)
|
|
|
|
def send(self, message):
|
|
self.sent.append(message)
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
|
|
class _FakeChildConnection:
|
|
def close(self):
|
|
return None
|
|
|
|
|
|
class _FakeProcess:
|
|
_next_pid = 4000
|
|
|
|
def __init__(self):
|
|
type(self)._next_pid += 1
|
|
self.pid = type(self)._next_pid
|
|
self.alive = False
|
|
self.terminated = False
|
|
self.killed = False
|
|
|
|
def start(self):
|
|
self.alive = True
|
|
|
|
def is_alive(self):
|
|
return self.alive
|
|
|
|
def join(self, timeout=None):
|
|
return None
|
|
|
|
def terminate(self):
|
|
self.terminated = True
|
|
self.alive = False
|
|
|
|
def kill(self):
|
|
self.killed = True
|
|
self.alive = False
|
|
|
|
|
|
class _FakeContext:
|
|
def __init__(self, connections):
|
|
self.connections = list(connections)
|
|
self.processes = []
|
|
|
|
def Pipe(self, duplex=True):
|
|
self.assert_duplex = duplex
|
|
return self.connections.pop(0), _FakeChildConnection()
|
|
|
|
def Process(self, **_kwargs):
|
|
process = _FakeProcess()
|
|
self.processes.append(process)
|
|
return process
|
|
|
|
|
|
def _ready(pid):
|
|
return {
|
|
"type": "ready",
|
|
"pid": pid,
|
|
"load_ms": 100,
|
|
"warmup_ms": 50,
|
|
"max_new_tokens": MAX_NEW_TOKENS,
|
|
}
|
|
|
|
|
|
def _progress(stage):
|
|
return {
|
|
"type": "progress",
|
|
"stage": stage,
|
|
"pid": 4000,
|
|
}
|
|
|
|
|
|
def _probe_worker_main(connection, _settings):
|
|
connection.send({
|
|
"type": "ready",
|
|
"pid": os.getpid(),
|
|
"load_ms": 1,
|
|
"warmup_ms": 1,
|
|
"max_new_tokens": MAX_NEW_TOKENS,
|
|
})
|
|
while True:
|
|
message = connection.recv()
|
|
if message.get("command") == "stop":
|
|
break
|
|
if message.get("command") == "synthesize":
|
|
audio = b"RIFF-spawn-probe"
|
|
connection.send({
|
|
"type": "result",
|
|
"request_id": message["request_id"],
|
|
"audio": audio,
|
|
"duration_ms": 2,
|
|
"bytes": len(audio),
|
|
})
|
|
connection.close()
|
|
|
|
|
|
class FasterQwenWorkerTests(unittest.TestCase):
|
|
def test_worker_environment_limits_native_thread_pools(self):
|
|
names = (
|
|
"OMP_NUM_THREADS",
|
|
"MKL_NUM_THREADS",
|
|
"OPENBLAS_NUM_THREADS",
|
|
"NUMEXPR_NUM_THREADS",
|
|
"VECLIB_MAXIMUM_THREADS",
|
|
"BLIS_NUM_THREADS",
|
|
)
|
|
with patch.dict(os.environ, {}, clear=False):
|
|
threads = _configure_worker_environment({})
|
|
|
|
self.assertEqual(threads, DEFAULT_CPU_THREADS)
|
|
for name in names:
|
|
self.assertEqual(os.environ[name], str(DEFAULT_CPU_THREADS))
|
|
self.assertEqual(os.environ["TOKENIZERS_PARALLELISM"], "false")
|
|
|
|
def test_default_worker_resource_limits_are_four_threads_and_eight_cores(self):
|
|
self.assertEqual(DEFAULT_CPU_THREADS, 4)
|
|
self.assertEqual(DEFAULT_CPU_AFFINITY_COUNT, 8)
|
|
|
|
def test_worker_environment_accepts_bounded_override(self):
|
|
with patch.dict(os.environ, {}, clear=False):
|
|
self.assertEqual(_configure_worker_environment({"cpu_threads": 4}), 4)
|
|
self.assertEqual(_configure_worker_environment({"cpu_threads": 99}), 8)
|
|
|
|
def test_generation_is_hard_limited_to_384_tokens(self):
|
|
kwargs = _generation_kwargs(
|
|
{"language": "Chinese", "non_streaming_mode": True},
|
|
"测试",
|
|
)
|
|
|
|
self.assertEqual(MAX_NEW_TOKENS, 384)
|
|
self.assertEqual(kwargs["max_new_tokens"], 384)
|
|
|
|
def test_client_returns_worker_audio(self):
|
|
connection = _FakeConnection([
|
|
_progress("model_loading"),
|
|
_progress("warmup_started"),
|
|
_ready(4100),
|
|
{
|
|
"type": "result",
|
|
"request_id": "placeholder",
|
|
"audio": b"RIFF-audio",
|
|
"duration_ms": 1234,
|
|
"bytes": 10,
|
|
},
|
|
])
|
|
context = _FakeContext([connection])
|
|
client = FasterQwenWorkerClient(
|
|
{},
|
|
logging.getLogger("test-faster-qwen-worker"),
|
|
context=context,
|
|
)
|
|
|
|
client.ensure_ready()
|
|
request_id = "fixed-request-id"
|
|
connection.responses[0]["request_id"] = request_id
|
|
with patch("app.faster_qwen_worker.uuid.uuid4") as make_uuid:
|
|
make_uuid.return_value.hex = request_id
|
|
audio, metadata = client.synthesize("测试")
|
|
|
|
self.assertEqual(audio, b"RIFF-audio")
|
|
self.assertEqual(metadata["duration_ms"], 1234)
|
|
self.assertEqual(len(context.processes), 1)
|
|
client.close()
|
|
|
|
def test_client_default_startup_timeout_allows_slow_model_load(self):
|
|
client = FasterQwenWorkerClient({})
|
|
|
|
self.assertEqual(client.startup_timeout_seconds, 900.0)
|
|
client.close()
|
|
|
|
def test_worker_applies_process_limits_after_model_warmup(self):
|
|
order = []
|
|
|
|
class FakeModel:
|
|
def warmup(self):
|
|
order.append("warmup")
|
|
|
|
class WorkerConnection:
|
|
def __init__(self):
|
|
self.sent = []
|
|
self.closed = False
|
|
|
|
def send(self, message):
|
|
self.sent.append(message)
|
|
|
|
def recv(self):
|
|
return {"command": "stop"}
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
connection = WorkerConnection()
|
|
with (
|
|
patch("app.faster_qwen_worker._configure_worker_environment", return_value=3),
|
|
patch(
|
|
"app.faster_qwen_worker._load_runtime",
|
|
side_effect=lambda _settings, progress=None: (
|
|
progress("from_pretrained_started") if progress else None,
|
|
order.append("load"),
|
|
{"model": FakeModel()},
|
|
)[-1],
|
|
),
|
|
patch(
|
|
"app.faster_qwen_worker._apply_worker_process_limits",
|
|
side_effect=lambda _settings: order.append("limits") or {
|
|
"cpu_affinity_count": 6,
|
|
"process_priority": "below_normal",
|
|
},
|
|
),
|
|
patch(
|
|
"app.faster_qwen_worker._synthesize_wav",
|
|
side_effect=lambda _runtime, _text: order.append("voice_clone") or b"wav",
|
|
),
|
|
):
|
|
faster_qwen_worker_main(connection, {})
|
|
|
|
self.assertEqual(order, ["load", "warmup", "voice_clone", "limits"])
|
|
progress_stages = [
|
|
item.get("stage")
|
|
for item in connection.sent
|
|
if item.get("type") == "progress"
|
|
]
|
|
self.assertEqual(
|
|
progress_stages,
|
|
[
|
|
"process_started",
|
|
"environment_ready",
|
|
"model_loading",
|
|
"from_pretrained_started",
|
|
"model_loaded",
|
|
"warmup_started",
|
|
"graph_warmup_complete",
|
|
"voice_clone_warmup_started",
|
|
"voice_clone_warmup_complete",
|
|
"warmup_complete",
|
|
],
|
|
)
|
|
self.assertEqual(connection.sent[-1]["type"], "ready")
|
|
|
|
def test_timeout_terminates_worker_and_starts_a_prewarmed_replacement(self):
|
|
first_connection = _FakeConnection([_ready(4200)])
|
|
replacement_connection = _FakeConnection([_ready(4300)])
|
|
context = _FakeContext([first_connection, replacement_connection])
|
|
client = FasterQwenWorkerClient(
|
|
{},
|
|
logging.getLogger("test-faster-qwen-timeout"),
|
|
synthesis_timeout_seconds=120,
|
|
context=context,
|
|
)
|
|
|
|
with self.assertRaises(FasterQwenWorkerTimeout):
|
|
client.synthesize("会超时的播报")
|
|
|
|
self.assertEqual(len(context.processes), 2)
|
|
self.assertTrue(context.processes[0].terminated)
|
|
self.assertTrue(context.processes[1].is_alive())
|
|
self.assertEqual(client.worker_pid, context.processes[1].pid)
|
|
client.close()
|
|
|
|
def test_real_spawned_worker_round_trip(self):
|
|
client = FasterQwenWorkerClient(
|
|
{},
|
|
logging.getLogger("test-faster-qwen-spawn"),
|
|
startup_timeout_seconds=20,
|
|
process_target=_probe_worker_main,
|
|
)
|
|
|
|
audio, metadata = client.synthesize("测试子进程")
|
|
|
|
self.assertEqual(audio, b"RIFF-spawn-probe")
|
|
self.assertEqual(metadata["duration_ms"], 2)
|
|
self.assertGreater(client.worker_pid, 0)
|
|
client.close()
|
|
|
|
def test_startup_failure_backoff_blocks_immediate_retry(self):
|
|
from app.faster_qwen_worker import (
|
|
STARTUP_FAILURE_BACKOFF_SECONDS,
|
|
FasterQwenWorkerError,
|
|
)
|
|
|
|
# 启动永远超时（poll 返回 False），触发启动失败
|
|
timeout_connection = _FakeConnection([])
|
|
timeout_connection.poll = lambda _timeout: False
|
|
context = _FakeContext([timeout_connection, _FakeConnection([_ready(4400)])])
|
|
client = FasterQwenWorkerClient(
|
|
{},
|
|
logging.getLogger("test-faster-qwen-backoff"),
|
|
startup_timeout_seconds=0.01,
|
|
context=context,
|
|
)
|
|
|
|
with self.assertRaises(FasterQwenWorkerError):
|
|
client.ensure_ready()
|
|
self.assertGreater(client._next_start_after, 0.0)
|
|
|
|
# 退避期内立即重试应直接报退避错误，且不创建新进程
|
|
with self.assertRaises(FasterQwenWorkerError) as ctx:
|
|
client.ensure_ready()
|
|
self.assertIn("退避", str(ctx.exception))
|
|
self.assertEqual(len(context.processes), 1)
|
|
|
|
# 退避期过后允许重新启动
|
|
client._next_start_after = 0.0
|
|
client.ensure_ready()
|
|
self.assertEqual(len(context.processes), 2)
|
|
self.assertEqual(client._next_start_after, 0.0)
|
|
self.assertGreater(STARTUP_FAILURE_BACKOFF_SECONDS, 0)
|
|
client.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|