Plugins: Allow exchanging analysis providers and inference providers alike - #782
Conversation
Adds two provider-agnostic seams so an accelerator (GPU, NPU, ...) can be shipped as an out-of-tree plugin instead of forking core: * register_onnx_provider() gains only_models/exclude_models so a plugin can offer an ONNX execution provider for just the graphs it can compile, scoped by the session label passed to create_onnx_session (musicnn, clap, whisper_encoder, ...). resolve_providers() honors the scope and leaves unmatched models on their default chain. * register_analysis_provider(component, factory) is wired end to end (it was a logging-only stub). A plugin can replace a whole analysis component; core consults the registered provider first and falls back to the built-in. The lyrics pipeline resolves its ASR/Whisper backend through lyrics/_asr_backend so the built-in ONNX Whisper can be swapped without touching call sites. clap_analyzer pins the CLAP audio model's symbolic time axis to a static shape for any provider whose graph compiler needs it, generalizing the existing CoreML-only path.
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
Which test did you run on this PR, can you share more details ? My expectation is that on container image (normal image, Nvidia, novx) and native image (Windows, Linux and Macos) on both Intel and Arm without plugin installed this is a NO-OP. With plugin installed, the plugin should be able to:
Can you share more about the above points? Thanks. |
|
I have run the changed code locally with and without the plugin. Without the plugin, no additional ExecutionProvider is added. The only real change is the
This is correct. Without the plugin, nothing changes. Default is always CPUExecutionProvider or CUDAExcecutionProvider if Nvidia image is used (this I cannot test, I don't have Nvidia GPUs, sorry). Here's an example log file of analysis running purely on CPU:
This is currently not in. Do you want me to implement that? The mechanism is currently inverted, the core calls
Yes and no. I did not change this mechanism, so just like you implemented before, you can provide a .json to install additional dependencies via PyPi. What is not possible is to install system dependencies. This would be insecure imo. This is a similar problem to Nvidia. You need system dependencies and a different base image to get everything working perfectly. I have CI running to publish base images for all supported AMD GPUs here: https://github.com/Schaka/rocm-migraphx-ort-builder One last point:
Currently, I only did this for whisper. It only made sense for me to do it for whisper, because there are so many speech to text models that all work completely different. For something like CLAP or even musicnn, I don't think there is any alternative so I'm not sure we could benefit from switching out the model. Does this all make sense to you? Let me know if you need anything more from me. |
There was a problem hiding this comment.
Remember that this PR MUST allow to everyone to use different provider for any of the actual model by a plugin. So the core idea should be change this to allow multiple use, not only this specific one.
Unit test is failing, please check:
=========================== short test summary info ============================
FAILED test/unit/test_plugin_analysis_seams.py::TestGetAsrBackend::test_falls_back_to_builtin_when_no_override - AssertionError: assert <module 'lyrics.whisper_onnx' from '/home/runner/work/AudioMuse-AI/AudioMuse-AI/lyrics/whisper_onnx.py'> is <module 'lyrics.whisper_onnx'>
+ where <module 'lyrics.whisper_onnx' from '/home/runner/work/AudioMuse-AI/AudioMuse-AI/lyrics/whisper_onnx.py'> = <function get_asr_backend at 0x7febd5fc7a60>()
+ where <function get_asr_backend at 0x7febd5fc7a60> = <module 'lyrics._asr_backend' from '/home/runner/work/AudioMuse-AI/AudioMuse-AI/lyrics/_asr_backend.py'>.get_asr_backend
FAILED test/unit/test_plugin_analysis_seams.py::TestGetAsrBackend::test_manager_error_falls_back_to_builtin - AssertionError: assert <module 'lyrics.whisper_onnx' from '/home/runner/work/AudioMuse-AI/AudioMuse-AI/lyrics/whisper_onnx.py'> is <module 'lyrics.whisper_onnx'>
+ where <module 'lyrics.whisper_onnx' from '/home/runner/work/AudioMuse-AI/AudioMuse-AI/lyrics/whisper_onnx.py'> = <function get_asr_backend at 0x7febd5fc7a60>()
+ where <function get_asr_backend at 0x7febd5fc7a60> = <module 'lyrics._asr_backend' from '/home/runner/work/AudioMuse-AI/AudioMuse-AI/lyrics/_asr_backend.py'>.get_asr_backend
====== 2 failed, 2667 passed, 1 skipped, 22 warnings in 154.71s (0:02:34) ======
Code Review point of attention
1. The per-track fallback never gets a label
In tasks/analysis/song.py the MusiCNN models can be loaded in two different ways.
The normal way: at the start of every album the worker loads the two models one time and reuses them for all the tracks of that album. This is ensure_musicnn_sessions calling load_musicnn_sessions, line 244. Your PR gives this one label='musicnn', so here it works.
The fallback way: if that album load fails, load_musicnn_sessions catches the error and returns None. Then analyze_track receives onnx_sessions=None, and _sessions_for_track (line 428) builds a new session for every single track. This one calls resolve_providers() with no argument, so label stays None, and it never gets a label.
With label=None your new filter drops the plugin provider:
if only and label not in only: # None not in ['musicnn'] -> skip
continueSo a plugin with only_models=['musicnn'] uses the GPU in the normal way, and goes back to CPU in the fallback way, without any message in the log. What do you think about adding label='musicnn' on line 428 too?
2. Today the ASR model is loaded one time per album, not one time per song
How it works now: the built-in whisper is a module singleton. It is loaded on the first track that needs lyrics, kept in memory for all the tracks of the album, and freed at the end of the album by cleanup_optional_models.
With the new seam this is not guaranteed anymore. get_asr_backend() has no cache, and the manager runs factory() on every call (once per song in _transcribe, and again in is_lyrics_loaded and unload_lyrics_models). So if a plugin gives a callable that builds a new object, the model is loaded again for every song, and the old one is never freed, because is_loaded() on a new object returns False.
What do you think about keeping the same behaviour as today, so the backend is resolved one time and cached? Or, if you prefer, give the choice to the plugin developer, for example a flag on register_analysis_provider, so the plugin can decide if it wants a cached backend or a new one every time.
3. Only 4 labels really work, and the docs miss one
The labels that really arrive to your new filter are only musicnn, clap, whisper_encoder and whisper_decoder. The docs list the first three and miss whisper_decoder, which is exactly the graph you say MIGraphX cannot compile. So someone who writes exclude_models=['whisper_encoder'] still gets the provider on the decoder, and there is no warning.
The other labels (embedding, prediction, gte, silero_vad) never arrive to the filter, because those sessions already receive a ready provider list, so scoping them does nothing.
What do you think about writing the 4 real labels in the docs, and logging a warning when a plugin uses a label that does not exist?
4. Other things, can you check these too?
_PREPARED_MODEL_PROVIDERSintasks/clap_analyzer.pyhasMIGraphXExecutionProviderwritten inside core. The idea of the PR is that a new accelerator does not need a core change, but the next provider that needs static shapes will need one. What do you think about a flag at registration time, for exampleneeds_static_shapes=True, so core does not need to know the provider names?only_models='musicnn'written as a plain string becomes['m','u','s','i','c','n','n'], so it matches nothing and there is no error. Can you accept a string and convert it to a list, like_plugin_targetsalready does fortargets?- Nothing checks that a plugin backend really has the 4 methods. If
is_loadedis missing,lyrics/__init__.pycatches the error andis_lyrics_loaded()returns True forever, so we run a full memory cleanup after every album for nothing. And iftranscribeis broken, every song fails and we never go back to the built-in whisper. Can we check the 4 methods one time when we resolve, log a warning and use the built-in if something is missing? register_analysis_provideraccepts anycomponentname. If a plugin writes'ASR'or'musicnn', the value is stored, the plugin saysok, and nothing happens, with no log. The old stub at least logged a message. Same thing if two plugins register'asr': the winner depends on dict order, with no warning.positionis accepted and it is also in the doc row, but it is never read, so a plugin provider is always added after CUDA. This is true also before your PR, but since you are touching this API, maybe it is a good moment to honour it or to remove it.
5. Half of the models have no label, so a plugin can never reach them
Only 4 sessions ask resolve_providers, so only 4 models can be scoped today. These other sessions pass their own provider list, so they never ask resolve_providers and a plugin provider can never arrive there:
gte, the lyrics text embedding,lyrics/gte_onnx.py:103, CPU onlysilero_vad,lyrics/silero_onnx.py:73, CPU only- the CLAP text model,
tasks/clap_analyzer.py:209, it builds its own CUDA or CPU list - MusiCNN in the per-track path,
tasks/analysis/song.py:428, this is my point 1
The reason is that create_onnx_session does provider_options or resolve_providers(...), so when the call site already passes a list, resolve_providers is not called and the label is used only for the log line.
So today the feature works only for the models your own plugin needs. For example someone with OpenVINO would like to say "use it for MusiCNN and for gte, but not for CLAP and not for Whisper", and the gte part is simply not possible.
What do you think about making this general for all the models? An idea: keep CPU as the default for these sessions, but let the plugin opt-in, for example resolve_providers(label='gte', cpu_only_default=True). In this way nothing changes for the normal user, and a plugin can use only_models=['gte'] if it wants. For the CLAP text model please keep the Flask process on CPU, because that one is for thread safety, and open it only on the worker.
|
Thank you for the extensive review.
They were running locally. As far as I understand, it had something to do with my local environment missing a dependency. This came down to my lack of experience with Python. I initially tried to be minimally invasive with my changes, but I agree the that we should cover all models and cache instantiation of the asr backend. Good catch, honestly. I didn't really think about that because I come from other languages and framework. I now gated the asr backend and validate that methods exist. Do you think this is enough or is there a way to validate that methods also take the correct parameter maybe? Hopefully I covered everything now (in a proper Python way) |
NeptuneHub
left a comment
There was a problem hiding this comment.
Hi,
thanks for your additional update, there is still some small update that I need (the first is probably the most important one):
1. Document what transcribe() must return
Background: after a song is transcribed, core reads a confidence value (avg_logprob) from the result to decide if the transcript is good enough. If the confidence is too low, the transcript is thrown away and the song is treated as instrumental.
The problem: the docs only list the four methods a plugin must provide, they never say what transcribe() must return. If a plugin returns only text and language (a very natural thing to do) and forgets avg_logprob, core assumes the worst confidence and drops every transcript. No crash, no error, the plugin looks like it works, but the user gets zero lyrics for the whole library. This is exactly the trap the faster-whisper plugin this PR is made for could fall into.
Fix asked: document the required return keys (text, language, avg_logprob). If possible, also skip the quality check when avg_logprob is missing instead of assuming the worst, so this failure becomes impossible.
2. Make plugin failures visible in the log
Background: when a plugin cannot provide its component, core falls back to the built-in behavior. This is the right choice, but the user must be able to see it happened.
The problem: two cases are hard to diagnose today. When a plugin factory fails, the error message says which component failed but not which plugin, so the user cannot tell which plugin to fix or remove. And when a factory returns None, nothing is logged at all: the user installs the plugin, core silently keeps using the built-in Whisper, and there is no trace anywhere of why.
Fix asked: add the plugin id to the factory error message, and log a message when a factory returns None.
3. Small doc additions
Background: these are the first questions every plugin author will ask.
The problem: three things are true in the code but not written anywhere: the execution provider must already be present in the onnxruntime build of the image (a plugin cannot add one via pip into its own _lib); needs_static_shapes only affects the CLAP audio model; and if a plugin registers a provider core already manages (like CUDA), core wins and the plugin entry is dropped.
Fix asked: one sentence for each in docs/PLUGIN.md.
4. Drop the .gitignore change
Background: the PR touches .gitignore.
The problem: the added entries look like local development stuff not related to this PR (.claude/ is already ignored, .tokensave is not used anywhere, and .venv-test is added without a final newline).
Fix asked: remove this hunk from the PR.
Thanks again!
|
Again, thank you so much for the expansive review. Hopefully the documentation is enough now. I asked Claude to summarize my changes, so it's not initially written by me, but I went over to correct/shorten things where possible. I'm not great at writing docs, so hopefully it covers everything you'd want it to now.
This was a really good catch. In my comment I only wondered about whether we should check method signatures for Let me know if you need anything more and sorry for taking so long to cover all the bases. |
…fidence, name the owning plugin in fallback warnings, fix PLUGIN.md static-shape claim
…arn on impossible confidence, None-init and stale-owner cleanup
|
|
Hi, Commit 1 (58a0b83)
Commit 2 (822d847)
I personally did some test with no plugins installed, and for a normal users nothing should change, so I'm going to merge this PR. PLEASE:
Your plugin will become "the case study" for future developer, so really thanks for your effort. Important: Support for plugin that add execution provider will be released in the next release, that will be the |



This solves issue #755
It was not feasible to expand the plugin installer scripts to be able to install system dependencies into the image.
Even if it was, it would create a messy image and probably take absolutely forever, because we also need to compile dependencies and set up
PATHanyway.Once you merge these changes, I will be able to provide my own plugin to enable MiGraphX AMD GPU (plugin already written and tested against these changes). I have a base image ready in the form of https://github.com/Schaka/rocm-migraphx-ort-builder and will instruct users to either use my image (containing your Audiomuse-AI code) to provide system dependencies for MiGraphX to work, or I will provide a more clear cut base image and create another PR to your CI that will build a tag for AMD, similar to Nvidia (just
BASE_IMAGEswapped).Let me know if you need any more changes made to this.
Looking forward to having this upstream so I can publish the plugin.
Edit: I kept whisper as being named whisper. Let me know if you prefer renaming to ASR entirely