phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-06-28 02:53:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

model registration in ollama and vllm check against the available models in the provider (#446)

Browse source 
tests:
pytest -v -s -m "ollama"
llama_stack/providers/tests/inference/test_text_inference.py

pytest -v -s -m vllm_remote
llama_stack/providers/tests/inference/test_text_inference.py --env
VLLM_URL="http://localhost:9798/v1"

---------
This commit is contained in:
Dinesh Yeduguru 2024-11-13 13:04:06 -08:00 • committed by GitHub
parent 7f6ac2fbd7
commit 787e2034b7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 73 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

24
llama_stack/providers/remote/inference/vllm/vllm.py
View file

@ -45,12 +45,9 @@ def build_model_aliases():
]
class VLLMInferenceAdapter(Inference, ModelRegistryHelper, ModelsProtocolPrivate):
class VLLMInferenceAdapter(Inference, ModelsProtocolPrivate):
def __init__(self, config: VLLMInferenceAdapterConfig) -> None:
ModelRegistryHelper.__init__(
self,
model_aliases=build_model_aliases(),
)
self.register_helper = ModelRegistryHelper(build_model_aliases())
self.config = config
self.formatter = ChatFormat(Tokenizer.get_instance())
self.client = None
@ -131,6 +128,17 @@ class VLLMInferenceAdapter(Inference, ModelRegistryHelper, ModelsProtocolPrivate
):
yield chunk
async def register_model(self, model: Model) -> Model:
model = await self.register_helper.register_model(model)
res = self.client.models.list()
available_models = [m.id for m in res]
if model.provider_resource_id not in available_models:
raise ValueError(
f"Model {model.provider_resource_id} is not being served by vLLM. "
f"Available models: {', '.join(available_models)}"
)
return model
async def _get_params(
self, request: Union[ChatCompletionRequest, CompletionRequest]
) -> dict:
@ -149,7 +157,9 @@ class VLLMInferenceAdapter(Inference, ModelRegistryHelper, ModelsProtocolPrivate
]
else:
input_dict["prompt"] = chat_completion_request_to_prompt(
request, self.get_llama_model(request.model), self.formatter
request,
self.register_helper.get_llama_model(request.model),
self.formatter,
)
else:
assert (
@ -157,7 +167,7 @@ class VLLMInferenceAdapter(Inference, ModelRegistryHelper, ModelsProtocolPrivate
), "Together does not support media for Completion requests"
input_dict["prompt"] = completion_request_to_prompt(
request,
self.get_llama_model(request.model),
self.register_helper.get_llama_model(request.model),
self.formatter,
)