mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-07-02 20:40:36 +00:00
Allow specifying resources in StackRunConfig (#425)
# What does this PR do? This PR brings back the facility to not force registration of resources onto the user. This is not just annoying but actually not feasible sometimes. For example, you may have a Stack which boots up with private providers for inference for models A and B. There is no way for the user to actually know which model is being served by these providers now (to be able to register it.) How will this avoid the users needing to do registration? In a follow-up diff, I will make sure I update the sample run.yaml files so they list the models served by the distributions explicitly. So when users do `llama stack build --template <...>` and run it, their distributions come up with the right set of models they expect. For self-hosted distributions, it also allows us to have a place to explicit list the models that need to be served to make the "complete" stack (including safety, e.g.) ## Test Plan Started ollama locally with two lightweight models: Llama3.2-3B-Instruct and Llama-Guard-3-1B. Updated all the tests including agents. Here's the tests I ran so far: ```bash pytest -s -v -m "fireworks and llama_3b" test_text_inference.py::TestInference \ --env FIREWORKS_API_KEY=... pytest -s -v -m "ollama and llama_3b" test_text_inference.py::TestInference pytest -s -v -m ollama test_safety.py pytest -s -v -m faiss test_memory.py pytest -s -v -m ollama test_agents.py \ --inference-model=Llama3.2-3B-Instruct --safety-model=Llama-Guard-3-1B ``` Found a few bugs here and there pre-existing that these test runs fixed.
This commit is contained in:
parent
8035fa1869
commit
d9d271a684
15 changed files with 221 additions and 124 deletions
|
@ -9,6 +9,7 @@ import tempfile
|
|||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from llama_stack.apis.models import Model
|
||||
from llama_stack.distribution.datatypes import Api, Provider
|
||||
|
||||
from llama_stack.providers.inline.agents.meta_reference import (
|
||||
|
@ -17,8 +18,18 @@ from llama_stack.providers.inline.agents.meta_reference import (
|
|||
|
||||
from llama_stack.providers.tests.resolver import resolve_impls_for_test_v2
|
||||
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
|
||||
|
||||
from ..conftest import ProviderFixture, remote_stack_fixture
|
||||
from ..safety.fixtures import get_shield_to_register
|
||||
|
||||
|
||||
def pick_inference_model(inference_model):
|
||||
# This is not entirely satisfactory. The fixture `inference_model` can correspond to
|
||||
# multiple models when you need to run a safety model in addition to normal agent
|
||||
# inference model. We filter off the safety model by looking for "Llama-Guard"
|
||||
if isinstance(inference_model, list):
|
||||
inference_model = next(m for m in inference_model if "Llama-Guard" not in m)
|
||||
assert inference_model is not None
|
||||
return inference_model
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
|
@ -49,7 +60,7 @@ AGENTS_FIXTURES = ["meta_reference", "remote"]
|
|||
|
||||
|
||||
@pytest_asyncio.fixture(scope="session")
|
||||
async def agents_stack(request):
|
||||
async def agents_stack(request, inference_model, safety_model):
|
||||
fixture_dict = request.param
|
||||
|
||||
providers = {}
|
||||
|
@ -60,9 +71,28 @@ async def agents_stack(request):
|
|||
if fixture.provider_data:
|
||||
provider_data.update(fixture.provider_data)
|
||||
|
||||
inf_provider_id = providers["inference"][0].provider_id
|
||||
safety_provider_id = providers["safety"][0].provider_id
|
||||
|
||||
shield = get_shield_to_register(
|
||||
providers["safety"][0].provider_type, safety_provider_id, safety_model
|
||||
)
|
||||
|
||||
inference_models = (
|
||||
inference_model if isinstance(inference_model, list) else [inference_model]
|
||||
)
|
||||
impls = await resolve_impls_for_test_v2(
|
||||
[Api.agents, Api.inference, Api.safety, Api.memory],
|
||||
providers,
|
||||
provider_data,
|
||||
models=[
|
||||
Model(
|
||||
identifier=model,
|
||||
provider_id=inf_provider_id,
|
||||
provider_resource_id=model,
|
||||
)
|
||||
for model in inference_models
|
||||
],
|
||||
shields=[shield],
|
||||
)
|
||||
return impls[Api.agents], impls[Api.memory]
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue