diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index d661d99841..91911c75fd 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -734,21 +734,25 @@ void server_models::load_models() { std::set hidden_models; { std::set preset_paths; + auto add_hf_path = [&preset_paths](const common_preset & preset, const char * repo_key, const char * file_key) { + std::string hf_repo; + if (!preset.get_option(repo_key, hf_repo) || hf_repo.empty()) { + return; + } + std::string hf_file; + preset.get_option(file_key, hf_file); + std::string path = common_download_resolve_path(hf_repo, hf_file); + if (!path.empty()) { + preset_paths.insert(path); + } + }; for (const auto & [name, preset] : custom_presets) { std::string val; if (!preset.get_option(COMMON_ARG_PRESET_DEDUP_CACHE_MODELS, val) || !common_arg_utils::is_truthy(val)) { continue; } - std::string hf_repo; - if (!preset.get_option("LLAMA_ARG_HF_REPO", hf_repo) || hf_repo.empty()) { - continue; - } - std::string hf_file; - preset.get_option("LLAMA_ARG_HF_FILE", hf_file); - std::string path = common_download_resolve_path(hf_repo, hf_file); - if (!path.empty()) { - preset_paths.insert(path); - } + add_hf_path(preset, "LLAMA_ARG_HF_REPO", "LLAMA_ARG_HF_FILE"); + add_hf_path(preset, "LLAMA_ARG_SPEC_DRAFT_HF_REPO", "LLAMA_ARG_SPEC_DRAFT_MODEL"); } if (!preset_paths.empty()) { for (const auto & [name, preset] : cached_models) { diff --git a/tools/server/tests/unit/test_router.py b/tools/server/tests/unit/test_router.py index bae1565177..5355ab7367 100644 --- a/tools/server/tests/unit/test_router.py +++ b/tools/server/tests/unit/test_router.py @@ -433,12 +433,14 @@ def test_router_dedup_cache_models(): global server preset_path = os.path.join(TMP_DIR, "test_dedup.ini") - cache_id = "ggml-org/test-model-stories260K:F32" + main_cache_id = "ggml-org/test-model-stories260K:F32" + draft_cache_id = "ggml-org/test-model-stories260K-infill:F32" with open(preset_path, "w") as f: f.write( "[model-dedup]\n" "hf-repo = ggml-org/test-model-stories260K\n" + "spec-draft-hf = ggml-org/test-model-stories260K-infill\n" "dedup-cache-models = 1\n" ) @@ -448,12 +450,13 @@ def test_router_dedup_cache_models(): try: ids = _get_model_ids(is_reload=False) assert "model-dedup" in ids - assert cache_id not in ids, "cache model should be hidden by dedup" + assert main_cache_id not in ids, "main cache model should be hidden by dedup" + assert draft_cache_id not in ids, "draft cache model should be hidden by dedup" # other cache models are unaffected assert "ggml-org/tinygemma3-GGUF:Q8_0" in ids # the hidden model is only hidden from the listing, it can still be used - res = server.make_request("POST", "/tokenize", data={"model": cache_id, "content": "hello"}) + res = server.make_request("POST", "/tokenize", data={"model": main_cache_id, "content": "hello"}) assert res.status_code == 200 # disabling the flag brings the cache entry back on reload @@ -461,9 +464,11 @@ def test_router_dedup_cache_models(): f.write( "[model-dedup]\n" "hf-repo = ggml-org/test-model-stories260K\n" + "spec-draft-hf = ggml-org/test-model-stories260K-infill\n" ) ids = _get_model_ids(is_reload=True) - assert cache_id in ids + assert main_cache_id in ids + assert draft_cache_id in ids # the flag also works from the global section with open(preset_path, "w") as f: @@ -473,10 +478,12 @@ def test_router_dedup_cache_models(): "\n" "[model-dedup]\n" "hf-repo = ggml-org/test-model-stories260K\n" + "spec-draft-hf = ggml-org/test-model-stories260K-infill\n" ) ids = _get_model_ids(is_reload=True) assert "model-dedup" in ids - assert cache_id not in ids, "cache model should be hidden by global dedup" + assert main_cache_id not in ids, "main cache model should be hidden by global dedup" + assert draft_cache_id not in ids, "draft cache model should be hidden by global dedup" finally: os.remove(preset_path)