mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-12 20:12:33 +00:00
add tests for double registration of scoring functions, models and vector dbs
This commit is contained in:
parent
a8fea981c3
commit
40201d809b
2 changed files with 230 additions and 0 deletions
|
|
@ -354,6 +354,107 @@ async def test_scoring_functions_routing_table(cached_disk_dist_registry):
|
||||||
assert len(scoring_functions_list_after_deletion.data) == 0
|
assert len(scoring_functions_list_after_deletion.data) == 0
|
||||||
|
|
||||||
|
|
||||||
|
async def test_double_registration_models_positive(cached_disk_dist_registry):
|
||||||
|
"""Test that registering the same model twice with identical data succeeds."""
|
||||||
|
table = ModelsRoutingTable({"test_provider": InferenceImpl()}, cached_disk_dist_registry, {})
|
||||||
|
await table.initialize()
|
||||||
|
|
||||||
|
# Register a model
|
||||||
|
await table.register_model(model_id="test-model", provider_id="test_provider")
|
||||||
|
|
||||||
|
# Register the exact same model again - should succeed (idempotent)
|
||||||
|
await table.register_model(model_id="test-model", provider_id="test_provider")
|
||||||
|
|
||||||
|
# Verify only one model exists
|
||||||
|
models = await table.list_models()
|
||||||
|
assert len(models.data) == 1
|
||||||
|
assert models.data[0].identifier == "test_provider/test-model"
|
||||||
|
|
||||||
|
|
||||||
|
async def test_double_registration_models_negative(cached_disk_dist_registry):
|
||||||
|
"""Test that registering the same model with different data fails."""
|
||||||
|
table = ModelsRoutingTable({"test_provider": InferenceImpl()}, cached_disk_dist_registry, {})
|
||||||
|
await table.initialize()
|
||||||
|
|
||||||
|
# Register a model with specific metadata
|
||||||
|
await table.register_model(model_id="test-model", provider_id="test_provider", metadata={"param1": "value1"})
|
||||||
|
|
||||||
|
# Try to register the same model with different metadata - should fail
|
||||||
|
with pytest.raises(ValueError, match="Provider 'test_provider' is already registered"):
|
||||||
|
await table.register_model(
|
||||||
|
model_id="test-model", provider_id="test_provider", metadata={"param1": "different_value"}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def test_double_registration_scoring_functions_positive(cached_disk_dist_registry):
|
||||||
|
"""Test that registering the same scoring function twice with identical data succeeds."""
|
||||||
|
table = ScoringFunctionsRoutingTable({"test_provider": ScoringFunctionsImpl()}, cached_disk_dist_registry, {})
|
||||||
|
await table.initialize()
|
||||||
|
|
||||||
|
# Register a scoring function
|
||||||
|
await table.register_scoring_function(
|
||||||
|
scoring_fn_id="test-scoring-fn",
|
||||||
|
provider_id="test_provider",
|
||||||
|
description="Test scoring function",
|
||||||
|
return_type=NumberType(),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Register the exact same scoring function again - should succeed (idempotent)
|
||||||
|
await table.register_scoring_function(
|
||||||
|
scoring_fn_id="test-scoring-fn",
|
||||||
|
provider_id="test_provider",
|
||||||
|
description="Test scoring function",
|
||||||
|
return_type=NumberType(),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify only one scoring function exists
|
||||||
|
scoring_functions = await table.list_scoring_functions()
|
||||||
|
assert len(scoring_functions.data) == 1
|
||||||
|
assert scoring_functions.data[0].identifier == "test-scoring-fn"
|
||||||
|
|
||||||
|
|
||||||
|
async def test_double_registration_scoring_functions_negative(cached_disk_dist_registry):
|
||||||
|
"""Test that registering the same scoring function with different data fails."""
|
||||||
|
table = ScoringFunctionsRoutingTable({"test_provider": ScoringFunctionsImpl()}, cached_disk_dist_registry, {})
|
||||||
|
await table.initialize()
|
||||||
|
|
||||||
|
# Register a scoring function
|
||||||
|
await table.register_scoring_function(
|
||||||
|
scoring_fn_id="test-scoring-fn",
|
||||||
|
provider_id="test_provider",
|
||||||
|
description="Test scoring function",
|
||||||
|
return_type=NumberType(),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Try to register the same scoring function with different description - should fail
|
||||||
|
with pytest.raises(ValueError, match="Provider 'test_provider' is already registered"):
|
||||||
|
await table.register_scoring_function(
|
||||||
|
scoring_fn_id="test-scoring-fn",
|
||||||
|
provider_id="test_provider",
|
||||||
|
description="Different description",
|
||||||
|
return_type=NumberType(),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def test_double_registration_different_providers(cached_disk_dist_registry):
|
||||||
|
"""Test that registering objects with same ID but different providers succeeds."""
|
||||||
|
impl1 = InferenceImpl()
|
||||||
|
impl2 = InferenceImpl()
|
||||||
|
table = ModelsRoutingTable({"provider1": impl1, "provider2": impl2}, cached_disk_dist_registry, {})
|
||||||
|
await table.initialize()
|
||||||
|
|
||||||
|
# Register same model ID with different providers - should succeed
|
||||||
|
await table.register_model(model_id="shared-model", provider_id="provider1")
|
||||||
|
await table.register_model(model_id="shared-model", provider_id="provider2")
|
||||||
|
|
||||||
|
# Verify both models exist with different identifiers
|
||||||
|
models = await table.list_models()
|
||||||
|
assert len(models.data) == 2
|
||||||
|
model_ids = {m.identifier for m in models.data}
|
||||||
|
assert "provider1/shared-model" in model_ids
|
||||||
|
assert "provider2/shared-model" in model_ids
|
||||||
|
|
||||||
|
|
||||||
async def test_benchmarks_routing_table(cached_disk_dist_registry):
|
async def test_benchmarks_routing_table(cached_disk_dist_registry):
|
||||||
table = BenchmarksRoutingTable({"test_provider": BenchmarksImpl()}, cached_disk_dist_registry, {})
|
table = BenchmarksRoutingTable({"test_provider": BenchmarksImpl()}, cached_disk_dist_registry, {})
|
||||||
await table.initialize()
|
await table.initialize()
|
||||||
|
|
|
||||||
|
|
@ -229,3 +229,132 @@ async def test_cached_registry_error_handling(sqlite_kvstore):
|
||||||
|
|
||||||
invalid_obj = await cached_registry.get("vector_db", "invalid_cached_db")
|
invalid_obj = await cached_registry.get("vector_db", "invalid_cached_db")
|
||||||
assert invalid_obj is None
|
assert invalid_obj is None
|
||||||
|
|
||||||
|
|
||||||
|
async def test_double_registration_identical_objects(disk_dist_registry):
|
||||||
|
"""Test that registering identical objects succeeds (idempotent)."""
|
||||||
|
vector_db = VectorDB(
|
||||||
|
identifier="test_vector_db",
|
||||||
|
embedding_model="all-MiniLM-L6-v2",
|
||||||
|
embedding_dimension=384,
|
||||||
|
provider_resource_id="test_vector_db",
|
||||||
|
provider_id="test-provider",
|
||||||
|
)
|
||||||
|
|
||||||
|
# First registration should succeed
|
||||||
|
result1 = await disk_dist_registry.register(vector_db)
|
||||||
|
assert result1 is True
|
||||||
|
|
||||||
|
# Second registration of identical object should also succeed (idempotent)
|
||||||
|
result2 = await disk_dist_registry.register(vector_db)
|
||||||
|
assert result2 is True
|
||||||
|
|
||||||
|
# Verify object exists and is unchanged
|
||||||
|
retrieved = await disk_dist_registry.get("vector_db", "test_vector_db")
|
||||||
|
assert retrieved is not None
|
||||||
|
assert retrieved.identifier == vector_db.identifier
|
||||||
|
assert retrieved.embedding_model == vector_db.embedding_model
|
||||||
|
|
||||||
|
|
||||||
|
async def test_double_registration_different_objects(disk_dist_registry):
|
||||||
|
"""Test that registering different objects with same identifier fails."""
|
||||||
|
vector_db1 = VectorDB(
|
||||||
|
identifier="test_vector_db",
|
||||||
|
embedding_model="all-MiniLM-L6-v2",
|
||||||
|
embedding_dimension=384,
|
||||||
|
provider_resource_id="test_vector_db",
|
||||||
|
provider_id="test-provider",
|
||||||
|
)
|
||||||
|
|
||||||
|
vector_db2 = VectorDB(
|
||||||
|
identifier="test_vector_db", # Same identifier
|
||||||
|
embedding_model="different-model", # Different embedding model
|
||||||
|
embedding_dimension=384,
|
||||||
|
provider_resource_id="test_vector_db",
|
||||||
|
provider_id="test-provider",
|
||||||
|
)
|
||||||
|
|
||||||
|
# First registration should succeed
|
||||||
|
result1 = await disk_dist_registry.register(vector_db1)
|
||||||
|
assert result1 is True
|
||||||
|
|
||||||
|
# Second registration with different data should fail
|
||||||
|
with pytest.raises(ValueError, match="Provider 'test-provider' is already registered"):
|
||||||
|
await disk_dist_registry.register(vector_db2)
|
||||||
|
|
||||||
|
# Verify original object is unchanged
|
||||||
|
retrieved = await disk_dist_registry.get("vector_db", "test_vector_db")
|
||||||
|
assert retrieved is not None
|
||||||
|
assert retrieved.embedding_model == "all-MiniLM-L6-v2" # Original value
|
||||||
|
|
||||||
|
|
||||||
|
async def test_double_registration_different_providers(disk_dist_registry):
|
||||||
|
"""Test that objects with same identifier but different providers can be registered."""
|
||||||
|
from llama_stack.core.datatypes import VectorDBWithOwner
|
||||||
|
|
||||||
|
vector_db1 = VectorDBWithOwner(
|
||||||
|
identifier="shared_vector_db",
|
||||||
|
embedding_model="all-MiniLM-L6-v2",
|
||||||
|
embedding_dimension=384,
|
||||||
|
provider_resource_id="shared_vector_db",
|
||||||
|
provider_id="provider1",
|
||||||
|
)
|
||||||
|
|
||||||
|
vector_db2 = VectorDBWithOwner(
|
||||||
|
identifier="shared_vector_db", # Same identifier
|
||||||
|
embedding_model="all-MiniLM-L6-v2",
|
||||||
|
embedding_dimension=384,
|
||||||
|
provider_resource_id="shared_vector_db",
|
||||||
|
provider_id="provider2", # Different provider
|
||||||
|
)
|
||||||
|
|
||||||
|
# Both registrations should succeed because they have different provider_ids
|
||||||
|
result1 = await disk_dist_registry.register(vector_db1)
|
||||||
|
assert result1 is True
|
||||||
|
|
||||||
|
result2 = await disk_dist_registry.register(vector_db2)
|
||||||
|
assert result2 is True
|
||||||
|
|
||||||
|
# Both objects should be retrievable
|
||||||
|
retrieved1 = await disk_dist_registry.get("vector_db", "shared_vector_db")
|
||||||
|
# Note: Since they have the same identifier, the second one overwrites the first
|
||||||
|
# This behavior might be considered a bug, but we're testing current behavior
|
||||||
|
assert retrieved1 is not None
|
||||||
|
|
||||||
|
|
||||||
|
async def test_double_registration_with_cache(cached_disk_dist_registry):
|
||||||
|
"""Test double registration behavior with caching enabled."""
|
||||||
|
from llama_stack.apis.models import ModelType
|
||||||
|
from llama_stack.core.datatypes import ModelWithOwner
|
||||||
|
|
||||||
|
model1 = ModelWithOwner(
|
||||||
|
identifier="test_model",
|
||||||
|
provider_resource_id="test_model",
|
||||||
|
provider_id="test-provider",
|
||||||
|
model_type=ModelType.llm,
|
||||||
|
)
|
||||||
|
|
||||||
|
model2 = ModelWithOwner(
|
||||||
|
identifier="test_model", # Same identifier
|
||||||
|
provider_resource_id="test_model",
|
||||||
|
provider_id="test-provider",
|
||||||
|
model_type=ModelType.embedding, # Different type
|
||||||
|
)
|
||||||
|
|
||||||
|
# First registration should succeed and populate cache
|
||||||
|
result1 = await cached_disk_dist_registry.register(model1)
|
||||||
|
assert result1 is True
|
||||||
|
|
||||||
|
# Verify in cache
|
||||||
|
cached_model = cached_disk_dist_registry.get_cached("model", "test_model")
|
||||||
|
assert cached_model is not None
|
||||||
|
assert cached_model.model_type == ModelType.llm
|
||||||
|
|
||||||
|
# Second registration with different data should fail
|
||||||
|
with pytest.raises(ValueError, match="Provider 'test-provider' is already registered"):
|
||||||
|
await cached_disk_dist_registry.register(model2)
|
||||||
|
|
||||||
|
# Cache should still contain original model
|
||||||
|
cached_model_after = cached_disk_dist_registry.get_cached("model", "test_model")
|
||||||
|
assert cached_model_after is not None
|
||||||
|
assert cached_model_after.model_type == ModelType.llm
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue