forked from phoenix-oss/llama-stack-mirror
88a796ca5a
# What does this PR do? fix a bug where models registered at runtime could not be used. ``` $ llama-stack-client models register test-model --provider-id nvidia --provider-model-id meta/llama-3.1-70b-instruct $ curl http://localhost:8321/v1/openai/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{ "model": "test-model", "messages": [{"role": "user", "content": "What is the weather like in Boston today?"}] }' =(client)=> {"detail":"Internal server error: An unexpected error occurred."} =(server)=> TypeError: Missing required arguments; Expected either ('messages' and 'model') or ('messages', 'model' and 'stream') arguments to be given ``` *root cause:* test-model is not added to ModelRegistryHelper's alias_to_provider_id_map. as part of the fix, this adds tests for ModelRegistryHelper and defines its expected behavior. user visible behavior changes - | action | existing behavior | new behavior | | -- | -- | -- | | double register | success (but no change) | error | | register unknown | success (fail when used) | error | existing behavior for register unknown model and double register - ``` $ llama-stack-client models register test-model --provider-id nvidia --provider-model-id meta/llama-3.1-70b-instruct-unknown Successfully registered model test-model $ llama-stack-client models list | grep test-model │ llm │ test-model │ meta/llama-3.1-70b-instruct-unknown │ │ nv… │ $ llama-stack-client models register test-model --provider-id nvidia --provider-model-id meta/llama-3.1-70b-instruct Successfully registered model test-model $ llama-stack-client models list | grep test-model │ llm │ test-model │ meta/llama-3.1-70b-instruct-unknown │ │ nv… │ ``` new behavior for register unknown - ``` $ llama-stack-client models register test-model --provider-id nvidia --provider-model-id meta/llama-3.1-70b-instruct-unknown ╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ │ Failed to register model │ │ │ │ Error Type: BadRequestError │ │ Details: Error code: 400 - {'detail': "Invalid value: Model id │ │ 'meta/llama-3.1-70b-instruct-unknown' is not supported. Supported ids are: │ │ meta/llama-3.1-70b-instruct, snowflake/arctic-embed-l, meta/llama-3.2-1b-instruct, │ │ nvidia/nv-embedqa-mistral-7b-v2, meta/llama-3.2-90b-vision-instruct, meta/llama-3.2-3b-instruct, │ │ meta/llama-3.2-11b-vision-instruct, meta/llama-3.1-405b-instruct, meta/llama3-8b-instruct, │ │ meta/llama3-70b-instruct, nvidia/llama-3.2-nv-embedqa-1b-v2, meta/llama-3.1-8b-instruct, │ │ nvidia/nv-embedqa-e5-v5"} │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯ ``` new behavior for double register - ``` $ llama-stack-client models register test-model --provider-id nvidia --provider-model-id meta/llama-3.1-70b-instruct Successfully registered model test-model $ llama-stack-client models register test-model --provider-id nvidia --provider-model-id meta/llama-3.2-1b-instruct ╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ │ Failed to register model │ │ │ │ Error Type: BadRequestError │ │ Details: Error code: 400 - {'detail': "Invalid value: Model id 'test-model' is already │ │ registered. Please use a different id or unregister it first."} │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯ ``` ## Test Plan ``` uv run pytest -v tests/unit/providers/utils/test_model_registry.py ```
163 lines
5.7 KiB
Python
163 lines
5.7 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
#
|
|
# ModelRegistryHelper provides mixin functionality for registering and
|
|
# unregistering models. It maintains a mapping of model ID / aliases to
|
|
# provider model IDs.
|
|
#
|
|
# Test cases -
|
|
# - Looking up an alias that does not exist should return None.
|
|
# - Registering a model + provider ID should add the model to the registry. If
|
|
# provider ID is known or an alias for a provider ID.
|
|
# - Registering an existing model should return an error. Unless it's a
|
|
# dulicate entry.
|
|
# - Unregistering a model should remove it from the registry.
|
|
# - Unregistering a model that does not exist should return an error.
|
|
# - Supported model ID and their aliases are registered during initialization.
|
|
# Only aliases are added afterwards.
|
|
#
|
|
# Questions -
|
|
# - Should we be allowed to register models w/o provider model IDs? No.
|
|
# According to POST /v1/models, required params are
|
|
# - identifier
|
|
# - provider_resource_id
|
|
# - provider_id
|
|
# - type
|
|
# - metadata
|
|
# - model_type
|
|
#
|
|
# TODO: llama_model functionality
|
|
#
|
|
|
|
import pytest
|
|
|
|
from llama_stack.apis.models.models import Model
|
|
from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper, ProviderModelEntry
|
|
|
|
|
|
@pytest.fixture
|
|
def known_model() -> Model:
|
|
return Model(
|
|
provider_id="provider",
|
|
identifier="known-model",
|
|
provider_resource_id="known-provider-id",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def known_model2() -> Model:
|
|
return Model(
|
|
provider_id="provider",
|
|
identifier="known-model2",
|
|
provider_resource_id="known-provider-id2",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def known_provider_model(known_model: Model) -> ProviderModelEntry:
|
|
return ProviderModelEntry(
|
|
provider_model_id=known_model.provider_resource_id,
|
|
aliases=[known_model.model_id],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def known_provider_model2(known_model2: Model) -> ProviderModelEntry:
|
|
return ProviderModelEntry(
|
|
provider_model_id=known_model2.provider_resource_id,
|
|
# aliases=[],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def unknown_model() -> Model:
|
|
return Model(
|
|
provider_id="provider",
|
|
identifier="unknown-model",
|
|
provider_resource_id="unknown-provider-id",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def helper(known_provider_model: ProviderModelEntry, known_provider_model2: ProviderModelEntry) -> ModelRegistryHelper:
|
|
return ModelRegistryHelper([known_provider_model, known_provider_model2])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lookup_unknown_model(helper: ModelRegistryHelper, unknown_model: Model) -> None:
|
|
assert helper.get_provider_model_id(unknown_model.model_id) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_unknown_provider_model(helper: ModelRegistryHelper, unknown_model: Model) -> None:
|
|
with pytest.raises(ValueError):
|
|
await helper.register_model(unknown_model)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_model(helper: ModelRegistryHelper, known_model: Model) -> None:
|
|
model = Model(
|
|
provider_id=known_model.provider_id,
|
|
identifier="new-model",
|
|
provider_resource_id=known_model.provider_resource_id,
|
|
)
|
|
assert helper.get_provider_model_id(model.model_id) is None
|
|
await helper.register_model(model)
|
|
assert helper.get_provider_model_id(model.model_id) == model.provider_resource_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_model_from_alias(helper: ModelRegistryHelper, known_model: Model) -> None:
|
|
model = Model(
|
|
provider_id=known_model.provider_id,
|
|
identifier="new-model",
|
|
provider_resource_id=known_model.model_id, # use known model's id as an alias for the supported model id
|
|
)
|
|
assert helper.get_provider_model_id(model.model_id) is None
|
|
await helper.register_model(model)
|
|
assert helper.get_provider_model_id(model.model_id) == known_model.provider_resource_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_model_existing(helper: ModelRegistryHelper, known_model: Model) -> None:
|
|
await helper.register_model(known_model)
|
|
assert helper.get_provider_model_id(known_model.model_id) == known_model.provider_resource_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_model_existing_different(
|
|
helper: ModelRegistryHelper, known_model: Model, known_model2: Model
|
|
) -> None:
|
|
known_model.provider_resource_id = known_model2.provider_resource_id
|
|
with pytest.raises(ValueError):
|
|
await helper.register_model(known_model)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unregister_model(helper: ModelRegistryHelper, known_model: Model) -> None:
|
|
await helper.register_model(known_model) # duplicate entry
|
|
assert helper.get_provider_model_id(known_model.model_id) == known_model.provider_model_id
|
|
await helper.unregister_model(known_model.model_id)
|
|
assert helper.get_provider_model_id(known_model.model_id) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unregister_unknown_model(helper: ModelRegistryHelper, unknown_model: Model) -> None:
|
|
with pytest.raises(ValueError):
|
|
await helper.unregister_model(unknown_model.model_id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_model_during_init(helper: ModelRegistryHelper, known_model: Model) -> None:
|
|
assert helper.get_provider_model_id(known_model.provider_resource_id) == known_model.provider_model_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unregister_model_during_init(helper: ModelRegistryHelper, known_model: Model) -> None:
|
|
assert helper.get_provider_model_id(known_model.provider_resource_id) == known_model.provider_model_id
|
|
await helper.unregister_model(known_model.provider_resource_id)
|
|
assert helper.get_provider_model_id(known_model.provider_resource_id) is None
|