System:
Hardware: Mac Studio (Apple Silicon)
Environment: Python 3.12, mlx (latest), Uvicorn/FastAPI
Task: Running a Kokoro TTS server using mlx-audio.
The Problem:
I am receiving a RuntimeError: There is no Stream(gpu, 0) in current thread whenever the server processes an inference request. The server starts fine, but as soon as model.generate() is called within the ASGI worker thread, it crashes.
Error Trace:
Plaintext
File ".../mlx_audio/tts/models/interpolate.py", line 34, in interpolate
curr_size = max(1, int(math.ceil(float(input_np.shape[i + 2]) * float(scale_factor[i]))))
RuntimeError: There is no Stream(gpu, 0) in current thread.
Diagnostic History & Attempts:
I have attempted to isolate the issue with the following interventions, all of which failed to resolve the thread-context crash:
Thread Initialization: Injected mx.set_default_device(mx.gpu) at the start of the worker _run() method to try and force a stream context creation in the thread.
CPU Offloading: Patched interpolate.py to move the input array to numpy (CPU) to perform shape calculations and interpolation, thinking this would bypass the MLX GPU stream requirement.
Variable Type Casting: Explicitly cast shapes to float and int to ensure no MLX array types were being passed into math operations that might trigger a GPU stream check.
Dependencies: Resolved multiple ModuleNotFoundError (espeakng-loader) and AttributeError (EspeakWrapper path) issues, and confirmed the espeak-ng path is correctly accessible via ESPEAK_DATA.
Questions for the Maintainers:
How can I properly propagate or initialize the MLX GPU stream context within a FastAPI/Uvicorn background thread?
Is there a recommended way to perform inference on MLX models within an ASGI/async environment that circumvents thread-local stream limitations?
Are there known compatibility issues between MLX stream management and starlette/anyio task groups?
Any guidance on how to properly handle the GPU context in this architecture would be appreciated.
System:
Hardware: Mac Studio (Apple Silicon)
Environment: Python 3.12, mlx (latest), Uvicorn/FastAPI
Task: Running a Kokoro TTS server using mlx-audio.
The Problem:
I am receiving a RuntimeError: There is no Stream(gpu, 0) in current thread whenever the server processes an inference request. The server starts fine, but as soon as model.generate() is called within the ASGI worker thread, it crashes.
Error Trace:
Plaintext
File ".../mlx_audio/tts/models/interpolate.py", line 34, in interpolate
curr_size = max(1, int(math.ceil(float(input_np.shape[i + 2]) * float(scale_factor[i]))))
RuntimeError: There is no Stream(gpu, 0) in current thread.
Diagnostic History & Attempts:
I have attempted to isolate the issue with the following interventions, all of which failed to resolve the thread-context crash:
Thread Initialization: Injected mx.set_default_device(mx.gpu) at the start of the worker _run() method to try and force a stream context creation in the thread.
CPU Offloading: Patched interpolate.py to move the input array to numpy (CPU) to perform shape calculations and interpolation, thinking this would bypass the MLX GPU stream requirement.
Variable Type Casting: Explicitly cast shapes to float and int to ensure no MLX array types were being passed into math operations that might trigger a GPU stream check.
Dependencies: Resolved multiple ModuleNotFoundError (espeakng-loader) and AttributeError (EspeakWrapper path) issues, and confirmed the espeak-ng path is correctly accessible via ESPEAK_DATA.
Questions for the Maintainers:
How can I properly propagate or initialize the MLX GPU stream context within a FastAPI/Uvicorn background thread?
Is there a recommended way to perform inference on MLX models within an ASGI/async environment that circumvents thread-local stream limitations?
Are there known compatibility issues between MLX stream management and starlette/anyio task groups?
Any guidance on how to properly handle the GPU context in this architecture would be appreciated.