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()