Fixed vllm and vllm-omni tasks to correct routing logic for HF DLCs

[dwarez]

Previously, the set of vLLM and vLLM-omni tasks used in routing for Hugging Face DLCs was wrong, which resulted in the wrong DLCs selection for models like google/gemma-4-E2B-it and other image-text-to-text models.

[pagezyhf]

is it always true?
many any-to-any run on vllm, no? https://recipes.vllm.ai/Google/gemma-4-12B-it

[dwarez]

iiuc, not all models on the hub flagged as any-to-any would work on vllm. In fact, also gemma-4-12B-it has a dedicated vllm image to make it work vllm/vllm-openai:gemma4-unified

[rsareddy0329]

Seems like a bug fix, Can we add more details in the description on the issue faced currently and how it is being addressed.

[dwarez]

@rsareddy0329 in the current routing logic, tasks that should be handled with vllm (like image-text-to-text) are instead routed to vllm-omni, which results in the wrong resolution of the image uri when using the ModelBuilder abstraction. We'd like to leave only the any-to-any task to be run in vllm-omni for the moment, and even with this logic there are edge cases (like google/gemma-4-12B-it which is flagged as any-to-any but should be served with vllm)

[github-actions]

Claude finished @dwarez's task in 1m 9s —— View job

---

Code Review

• Read the PR diff and changed files
• Verify routing logic consistency across call sites
• Check test coverage

---

I reviewed the routing change. Overall this is a clean, well-scoped fix. The core insight is correct: image-text-to-text and audio-text-to-text models produce text output and are served by vLLM, so folding them into OMNI_TASKS (multimodal/vLLM-omni) was mis-routing them and yielding the wrong image URI. Moving them into a new VLLM_TASKS set and leaving only any-to-any in OMNI_TASKS matches how these pipeline tags map to servers.

Correctness / consistency ✅

• Both auto-selection call sites are updated consistently — model_builder.py:3561 and model_builder_utils.py:684 — so the ModelBuilder.build() path and the _detect_huggingface_image() path agree. Good; a mismatch here would have been a subtle bug.
• No stale references to the old OMNI_TASKS membership remain (grep confirms only the two updated call sites use it).

Tests ✅

• test_constants.py, test_model_builder_utils_extended_coverage.py, and test_model_builder_workflows.py are all updated to reflect the new routing, including a positive case (google/gemma-4-E2B-it → huggingface-vllm). Coverage is adequate for the behavior change.

Points worth confirming (not blockers):

1. any-to-any → vLLM-omni is a heuristic, not a guarantee — echoing @pagezyhf's earlier question. As you noted in your own comment, google/gemma-4-12B-it is tagged any-to-any yet is served by vLLM. So this PR fixes the larger class of mis-routings but knowingly leaves an any-to-any edge case that will still route to omni. That's a reasonable incremental improvement, but it may be worth a short code comment above OMNI_TASKS documenting the known limitation (and that model tags aren't a perfect signal), so the next reader doesn't assume it's exhaustive. Users can still override via model_server= explicitly, which mitigates it.

2. Pre-existing TEI-branch inconsistency (out of scope, flagging only): model_builder.py:3567 routes ["sentence-similarity", "feature-extraction", "text-ranking"] to TEI, while model_builder_utils.py:688 only lists ["sentence-similarity", "feature-extraction"] (no text-ranking). Not introduced by this PR, but since you're touching this exact routing block in both files, it may be a good opportunity to align them.

No security or performance concerns — this is a constant/set change plus routing dispatch. LGTM pending the any-to-any documentation note.

Note: the inline-comment tool isn't available in this run, so the file/line references are inlined above.
· fix-hf-vllm-routing