phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2026-01-03 00:52:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

ruff fix

Browse source
This commit is contained in:
Honglin Cao 2025-03-12 22:23:25 -04:00
parent 0cef9adda5
commit 8943755156
3 changed files with 48 additions and 86 deletions
Download patch file Download diff file Expand all files Collapse all files

4
llama_stack/providers/remote/inference/centml/__init__.py
View file

@ -23,9 +23,7 @@ async def get_adapter_impl(config: CentMLImplConfig, _deps):
from .centml import CentMLInferenceAdapter from .centml import CentMLInferenceAdapter
# Ensure the provided config is indeed a CentMLImplConfig # Ensure the provided config is indeed a CentMLImplConfig
assert isinstance(config, CentMLImplConfig), ( assert isinstance(config, CentMLImplConfig), f"Unexpected config type: {type(config)}"
f"Unexpected config type: {type(config)}"
)
# Instantiate and initialize the adapter # Instantiate and initialize the adapter
adapter = CentMLInferenceAdapter(config) adapter = CentMLInferenceAdapter(config)

114
llama_stack/providers/remote/inference/centml/centml.py
View file

@ -6,12 +6,8 @@
from typing import AsyncGenerator, List, Optional, Union from typing import AsyncGenerator, List, Optional, Union
from openai import OpenAI
from pydantic import parse_obj_as
from llama_models.datatypes import CoreModelId from llama_models.datatypes import CoreModelId
from llama_models.llama3.api.chat_format import ChatFormat from openai import OpenAI
from llama_models.llama3.api.tokenizer import Tokenizer
from llama_stack.apis.common.content_types import InterleavedContent from llama_stack.apis.common.content_types import InterleavedContent
from llama_stack.apis.inference import ( from llama_stack.apis.inference import (
@ -33,8 +29,8 @@ from llama_stack.apis.inference import (
) )
from llama_stack.distribution.request_headers import NeedsRequestProviderData from llama_stack.distribution.request_headers import NeedsRequestProviderData
from llama_stack.providers.utils.inference.model_registry import ( from llama_stack.providers.utils.inference.model_registry import (
build_model_entry,
ModelRegistryHelper, ModelRegistryHelper,
build_model_entry,
) )
from llama_stack.providers.utils.inference.openai_compat import ( from llama_stack.providers.utils.inference.openai_compat import (
convert_message_to_openai_dict, convert_message_to_openai_dict,
@ -48,7 +44,6 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
chat_completion_request_to_prompt, chat_completion_request_to_prompt,
completion_request_to_prompt, completion_request_to_prompt,
content_has_media, content_has_media,
interleaved_content_as_str,
request_has_media, request_has_media,
) )
@ -67,8 +62,7 @@ MODEL_ALIASES = [
] ]
class CentMLInferenceAdapter(ModelRegistryHelper, Inference, class CentMLInferenceAdapter(ModelRegistryHelper, Inference, NeedsRequestProviderData):
NeedsRequestProviderData):
""" """
Adapter to use CentML's serverless inference endpoints, Adapter to use CentML's serverless inference endpoints,
which adhere to the OpenAI chat/completions API spec, which adhere to the OpenAI chat/completions API spec,
@ -116,7 +110,7 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
self, self,
model_id: str, model_id: str,
content: InterleavedContent, content: InterleavedContent,
sampling_params: Optional[SamplingParams] = SamplingParams(), sampling_params: Optional[SamplingParams] = None, # Avoid function call in default argument.
response_format: Optional[ResponseFormat] = None, response_format: Optional[ResponseFormat] = None,
stream: Optional[bool] = False, stream: Optional[bool] = False,
logprobs: Optional[LogProbConfig] = None, logprobs: Optional[LogProbConfig] = None,
@ -124,6 +118,9 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
""" """
For "completion" style requests (non-chat). For "completion" style requests (non-chat).
""" """
# Instantiate sampling_params if not provided.
if sampling_params is None:
sampling_params = SamplingParams()
model = await self.model_store.get_model(model_id) model = await self.model_store.get_model(model_id)
request = CompletionRequest( request = CompletionRequest(
model=model.provider_resource_id, model=model.provider_resource_id,
@ -138,8 +135,7 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
else: else:
return await self._nonstream_completion(request) return await self._nonstream_completion(request)
async def _nonstream_completion( async def _nonstream_completion(self, request: CompletionRequest) -> CompletionResponse:
self, request: CompletionRequest) -> CompletionResponse:
""" """
Process non-streaming completion requests. Process non-streaming completion requests.
@ -157,12 +153,10 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
choice = response.choices[0] choice = response.choices[0]
message = choice.message message = choice.message
# If message.content is returned as a list of tokens, join them into a string. # If message.content is returned as a list of tokens, join them into a string.
content = message.content if not isinstance( content = message.content if not isinstance(message.content, list) else "".join(message.content)
message.content, list) else "".join(message.content)
return CompletionResponse( return CompletionResponse(
content=content, content=content,
stop_reason= stop_reason="end_of_message", # ***** HACK: Hard-coded stop_reason because the chat API doesn't return one.
"end_of_message", # ***** HACK: Hard-coded stop_reason because the chat API doesn't return one.
logprobs=None, logprobs=None,
) )
else: else:
@ -180,8 +174,7 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
result.content = "".join(result.content) result.content = "".join(result.content)
return result return result
async def _stream_completion(self, async def _stream_completion(self, request: CompletionRequest) -> AsyncGenerator:
request: CompletionRequest) -> AsyncGenerator:
params = await self._get_params(request) params = await self._get_params(request)
async def _to_async_generator(): async def _to_async_generator():
@ -196,8 +189,7 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
stream = _to_async_generator() stream = _to_async_generator()
if request.response_format is not None: if request.response_format is not None:
async for chunk in process_chat_completion_stream_response( async for chunk in process_chat_completion_stream_response(stream, request):
stream, request):
yield chunk yield chunk
else: else:
async for chunk in process_completion_stream_response(stream): async for chunk in process_completion_stream_response(stream):
@ -211,7 +203,7 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
self, self,
model_id: str, model_id: str,
messages: List[Message], messages: List[Message],
sampling_params: Optional[SamplingParams] = SamplingParams(), sampling_params: Optional[SamplingParams] = None, # Avoid function call in default argument.
tools: Optional[List[ToolDefinition]] = None, tools: Optional[List[ToolDefinition]] = None,
tool_choice: Optional[ToolChoice] = ToolChoice.auto, tool_choice: Optional[ToolChoice] = ToolChoice.auto,
tool_prompt_format: Optional[ToolPromptFormat] = None, tool_prompt_format: Optional[ToolPromptFormat] = None,
@ -223,6 +215,9 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
""" """
For "chat completion" style requests. For "chat completion" style requests.
""" """
# Instantiate sampling_params if not provided.
if sampling_params is None:
sampling_params = SamplingParams()
model = await self.model_store.get_model(model_id) model = await self.model_store.get_model(model_id)
request = ChatCompletionRequest( request = ChatCompletionRequest(
model=model.provider_resource_id, model=model.provider_resource_id,
@ -240,8 +235,7 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
else: else:
return await self._nonstream_chat_completion(request) return await self._nonstream_chat_completion(request)
async def _nonstream_chat_completion( async def _nonstream_chat_completion(self, request: ChatCompletionRequest) -> ChatCompletionResponse:
self, request: ChatCompletionRequest) -> ChatCompletionResponse:
params = await self._get_params(request) params = await self._get_params(request)
# Use the chat completions endpoint if "messages" key is present. # Use the chat completions endpoint if "messages" key is present.
if "messages" in params: if "messages" in params:
@ -258,16 +252,13 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
result.completion_message["content"] = "".join(content) result.completion_message["content"] = "".join(content)
else: else:
if isinstance(result.completion_message.content, list): if isinstance(result.completion_message.content, list):
updated_msg = result.completion_message.copy(update={ updated_msg = result.completion_message.copy(
"content": update={"content": "".join(result.completion_message.content)}
"".join(result.completion_message.content) )
}) result = result.copy(update={"completion_message": updated_msg})
result = result.copy(
update={"completion_message": updated_msg})
return result return result
async def _stream_chat_completion( async def _stream_chat_completion(self, request: ChatCompletionRequest) -> AsyncGenerator:
self, request: ChatCompletionRequest) -> AsyncGenerator:
params = await self._get_params(request) params = await self._get_params(request)
async def _to_async_generator(): async def _to_async_generator():
@ -280,17 +271,14 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
yield chunk yield chunk
stream = _to_async_generator() stream = _to_async_generator()
async for chunk in process_chat_completion_stream_response( async for chunk in process_chat_completion_stream_response(stream, request):
stream, request):
yield chunk yield chunk
# #
# HELPER METHODS # HELPER METHODS
# #
async def _get_params( async def _get_params(self, request: Union[ChatCompletionRequest, CompletionRequest]) -> dict:
self, request: Union[ChatCompletionRequest,
CompletionRequest]) -> dict:
""" """
Build a unified set of parameters for both chat and non-chat requests. Build a unified set of parameters for both chat and non-chat requests.
When a structured output is specified (response_format is not None), we force When a structured output is specified (response_format is not None), we force
@ -301,38 +289,24 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
llama_model = self.get_llama_model(request.model) llama_model = self.get_llama_model(request.model)
if request.response_format is not None: if request.response_format is not None:
if isinstance(request, ChatCompletionRequest): if isinstance(request, ChatCompletionRequest):
input_dict["messages"] = [ input_dict["messages"] = [await convert_message_to_openai_dict(m) for m in request.messages]
await convert_message_to_openai_dict(m)
for m in request.messages
]
else: else:
# ***** HACK: For CompletionRequests with structured output, # ***** HACK: For CompletionRequests with structured output,
# we simulate a chat conversation by wrapping the prompt as a single user message. # we simulate a chat conversation by wrapping the prompt as a single user message.
prompt_str = await completion_request_to_prompt(request) prompt_str = await completion_request_to_prompt(request)
input_dict["messages"] = [{ input_dict["messages"] = [{"role": "user", "content": prompt_str}]
"role": "user",
"content": prompt_str
}]
else: else:
if isinstance(request, ChatCompletionRequest): if isinstance(request, ChatCompletionRequest):
if media_present or not llama_model: if media_present or not llama_model:
input_dict["messages"] = [ input_dict["messages"] = [await convert_message_to_openai_dict(m) for m in request.messages]
await convert_message_to_openai_dict(m)
for m in request.messages
]
else: else:
input_dict[ input_dict["prompt"] = await chat_completion_request_to_prompt(request, llama_model)
"prompt"] = await chat_completion_request_to_prompt(
request, llama_model)
else: else:
input_dict["prompt"] = await completion_request_to_prompt( input_dict["prompt"] = await completion_request_to_prompt(request)
request)
params = { params = {
"model": "model": request.model,
request.model,
**input_dict, **input_dict,
"stream": "stream": request.stream,
request.stream,
**self._build_options(request.sampling_params, request.logprobs, request.response_format), **self._build_options(request.sampling_params, request.logprobs, request.response_format),
} }
return params return params
@ -352,14 +326,10 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
if fmt.type == ResponseFormatType.json_schema.value: if fmt.type == ResponseFormatType.json_schema.value:
options["response_format"] = { options["response_format"] = {
"type": "json_schema", "type": "json_schema",
"json_schema": { "json_schema": {"name": "schema", "schema": fmt.json_schema},
"name": "schema",
"schema": fmt.json_schema
},
} }
elif fmt.type == ResponseFormatType.grammar.value: elif fmt.type == ResponseFormatType.grammar.value:
raise NotImplementedError( raise NotImplementedError("Grammar response format not supported yet")
"Grammar response format not supported yet")
else: else:
raise ValueError(f"Unknown response format {fmt.type}") raise ValueError(f"Unknown response format {fmt.type}")
if logprobs and logprobs.top_k: if logprobs and logprobs.top_k:
@ -378,13 +348,11 @@ class CentMLInferenceAdapter(ModelRegistryHelper, Inference,
output_dimension: Optional[int], output_dimension: Optional[int],
contents: List[InterleavedContent], contents: List[InterleavedContent],
) -> EmbeddingsResponse: ) -> EmbeddingsResponse:
# this will come in future updates # ***** HACK/ASSERT: CentML does not support media for embeddings.
model = await self.model_store.get_model(model_id) # We assert here to catch any cases where media is inadvertently included.
assert all(not content_has_media(c) for c in contents), ( # model = await self.model_store.get_model(model_id)
"CentML does not support media for embeddings") assert all(not content_has_media(c) for c in contents), "CentML does not support media for embeddings"
resp = self._get_client().embeddings.create( # resp = self._get_client().embeddings.create(
model=model.provider_resource_id, # model=model.provider_resource_id,
input=[interleaved_content_as_str(c) for c in contents], # input=[interleaved_content_as_str(c) for c in contents],
) # )
embeddings = [item.embedding for item in resp.data]
return EmbeddingsResponse(embeddings=embeddings)

16
llama_stack/templates/centml/centml.py
View file

@ -7,19 +7,19 @@
from pathlib import Path from pathlib import Path
from llama_models.sku_list import all_registered_models from llama_models.sku_list import all_registered_models
from llama_stack.apis.models.models import ModelType from llama_stack.apis.models.models import ModelType
from llama_stack.distribution.datatypes import ModelInput, Provider, ShieldInput from llama_stack.distribution.datatypes import ModelInput, Provider, ShieldInput
from llama_stack.providers.inline.inference.sentence_transformers import ( from llama_stack.providers.inline.inference.sentence_transformers import (
SentenceTransformersInferenceConfig, SentenceTransformersInferenceConfig,
) )
from llama_stack.providers.inline.memory.faiss.config import FaissImplConfig from llama_stack.providers.inline.memory.faiss.config import FaissImplConfig
from llama_stack.providers.remote.inference.centml.config import (
CentMLImplConfig,
)
# If your CentML adapter has a MODEL_ALIASES constant with known model mappings: # If your CentML adapter has a MODEL_ALIASES constant with known model mappings:
from llama_stack.providers.remote.inference.centml.centml import MODEL_ALIASES from llama_stack.providers.remote.inference.centml.centml import MODEL_ALIASES
from llama_stack.providers.remote.inference.centml.config import (
CentMLImplConfig,
)
from llama_stack.templates.template import ( from llama_stack.templates.template import (
DistributionTemplate, DistributionTemplate,
RunConfigSettings, RunConfigSettings,
@ -68,9 +68,7 @@ def get_distribution_template() -> DistributionTemplate:
) )
# Map Llama Models to provider IDs if needed # Map Llama Models to provider IDs if needed
core_model_to_hf_repo = { core_model_to_hf_repo = {m.descriptor(): m.huggingface_repo for m in all_registered_models()}
m.descriptor(): m.huggingface_repo for m in all_registered_models()
}
default_models = [ default_models = [
ModelInput( ModelInput(
model_id=core_model_to_hf_repo[m.llama_model], model_id=core_model_to_hf_repo[m.llama_model],
@ -103,9 +101,7 @@ def get_distribution_template() -> DistributionTemplate:
"memory": [memory_provider], "memory": [memory_provider],
}, },
default_models=default_models + [embedding_model], default_models=default_models + [embedding_model],
default_shields=[ default_shields=[ShieldInput(shield_id="meta-llama/Llama-Guard-3-8B")],
ShieldInput(shield_id="meta-llama/Llama-Guard-3-8B")
],
), ),
}, },
run_config_env_vars={ run_config_env_vars={