Merge branch 'main' into add-watsonx-inference-adapter

llama_stack/apis/batch_inference/batch_inference.py

@@ -6,11 +6,8 @@
 
 from typing import List, Optional, Protocol, runtime_checkable
 
-from pydantic import BaseModel
-
+from llama_stack.apis.common.job_types import Job
 from llama_stack.apis.inference import (
-    ChatCompletionResponse,
-    CompletionResponse,
     InterleavedContent,
     LogProbConfig,
     Message,
@@ -20,41 +17,39 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
-from llama_stack.schema_utils import json_schema_type, webmethod
-
-
-@json_schema_type
-class BatchCompletionResponse(BaseModel):
-    batch: List[CompletionResponse]
-
-
-@json_schema_type
-class BatchChatCompletionResponse(BaseModel):
-    batch: List[ChatCompletionResponse]
+from llama_stack.schema_utils import webmethod
 
 
 @runtime_checkable
 class BatchInference(Protocol):
+    """Batch inference API for generating completions and chat completions.
+
+    This is an asynchronous API. If the request is successful, the response will be a job which can be polled for completion.
+
+    NOTE: This API is not yet implemented and is subject to change in concert with other asynchronous APIs
+    including (post-training, evals, etc).
+    """
+
     @webmethod(route="/batch-inference/completion", method="POST")
-    async def batch_completion(
+    async def completion(
         self,
         model: str,
         content_batch: List[InterleavedContent],
         sampling_params: Optional[SamplingParams] = None,
         response_format: Optional[ResponseFormat] = None,
         logprobs: Optional[LogProbConfig] = None,
-    ) -> BatchCompletionResponse: ...
+    ) -> Job: ...
 
     @webmethod(route="/batch-inference/chat-completion", method="POST")
-    async def batch_chat_completion(
+    async def chat_completion(
         self,
         model: str,
         messages_batch: List[List[Message]],
         sampling_params: Optional[SamplingParams] = None,
         # zero-shot tool definitions as input to the model
-        tools: Optional[List[ToolDefinition]] = list,
+        tools: Optional[List[ToolDefinition]] = None,
         tool_choice: Optional[ToolChoice] = ToolChoice.auto,
         tool_prompt_format: Optional[ToolPromptFormat] = None,
         response_format: Optional[ResponseFormat] = None,
         logprobs: Optional[LogProbConfig] = None,
-    ) -> BatchChatCompletionResponse: ...
+    ) -> Job: ...

llama_stack/apis/inference/inference.py

@@ -18,7 +18,7 @@ from typing import (
 )
 
 from pydantic import BaseModel, Field, field_validator
-from typing_extensions import Annotated
+from typing_extensions import Annotated, TypedDict
 
 from llama_stack.apis.common.content_types import ContentDelta, InterleavedContent, InterleavedContentItem
 from llama_stack.apis.models import Model
@@ -442,6 +442,352 @@ class EmbeddingsResponse(BaseModel):
     embeddings: List[List[float]]
 
 
+@json_schema_type
+class OpenAIChatCompletionContentPartTextParam(BaseModel):
+    type: Literal["text"] = "text"
+    text: str
+
+
+@json_schema_type
+class OpenAIImageURL(BaseModel):
+    url: str
+    detail: Optional[str] = None
+
+
+@json_schema_type
+class OpenAIChatCompletionContentPartImageParam(BaseModel):
+    type: Literal["image_url"] = "image_url"
+    image_url: OpenAIImageURL
+
+
+OpenAIChatCompletionContentPartParam = Annotated[
+    Union[
+        OpenAIChatCompletionContentPartTextParam,
+        OpenAIChatCompletionContentPartImageParam,
+    ],
+    Field(discriminator="type"),
+]
+register_schema(OpenAIChatCompletionContentPartParam, name="OpenAIChatCompletionContentPartParam")
+
+
+OpenAIChatCompletionMessageContent = Union[str, List[OpenAIChatCompletionContentPartParam]]
+
+
+@json_schema_type
+class OpenAIUserMessageParam(BaseModel):
+    """A message from the user in an OpenAI-compatible chat completion request.
+
+    :param role: Must be "user" to identify this as a user message
+    :param content: The content of the message, which can include text and other media
+    :param name: (Optional) The name of the user message participant.
+    """
+
+    role: Literal["user"] = "user"
+    content: OpenAIChatCompletionMessageContent
+    name: Optional[str] = None
+
+
+@json_schema_type
+class OpenAISystemMessageParam(BaseModel):
+    """A system message providing instructions or context to the model.
+
+    :param role: Must be "system" to identify this as a system message
+    :param content: The content of the "system prompt". If multiple system messages are provided, they are concatenated. The underlying Llama Stack code may also add other system messages (for example, for formatting tool definitions).
+    :param name: (Optional) The name of the system message participant.
+    """
+
+    role: Literal["system"] = "system"
+    content: OpenAIChatCompletionMessageContent
+    name: Optional[str] = None
+
+
+@json_schema_type
+class OpenAIChatCompletionToolCallFunction(BaseModel):
+    name: Optional[str] = None
+    arguments: Optional[str] = None
+
+
+@json_schema_type
+class OpenAIChatCompletionToolCall(BaseModel):
+    index: Optional[int] = None
+    id: Optional[str] = None
+    type: Literal["function"] = "function"
+    function: Optional[OpenAIChatCompletionToolCallFunction] = None
+
+
+@json_schema_type
+class OpenAIAssistantMessageParam(BaseModel):
+    """A message containing the model's (assistant) response in an OpenAI-compatible chat completion request.
+
+    :param role: Must be "assistant" to identify this as the model's response
+    :param content: The content of the model's response
+    :param name: (Optional) The name of the assistant message participant.
+    :param tool_calls: List of tool calls. Each tool call is an OpenAIChatCompletionToolCall object.
+    """
+
+    role: Literal["assistant"] = "assistant"
+    content: OpenAIChatCompletionMessageContent
+    name: Optional[str] = None
+    tool_calls: Optional[List[OpenAIChatCompletionToolCall]] = Field(default_factory=list)
+
+
+@json_schema_type
+class OpenAIToolMessageParam(BaseModel):
+    """A message representing the result of a tool invocation in an OpenAI-compatible chat completion request.
+
+    :param role: Must be "tool" to identify this as a tool response
+    :param tool_call_id: Unique identifier for the tool call this response is for
+    :param content: The response content from the tool
+    """
+
+    role: Literal["tool"] = "tool"
+    tool_call_id: str
+    content: OpenAIChatCompletionMessageContent
+
+
+@json_schema_type
+class OpenAIDeveloperMessageParam(BaseModel):
+    """A message from the developer in an OpenAI-compatible chat completion request.
+
+    :param role: Must be "developer" to identify this as a developer message
+    :param content: The content of the developer message
+    :param name: (Optional) The name of the developer message participant.
+    """
+
+    role: Literal["developer"] = "developer"
+    content: OpenAIChatCompletionMessageContent
+    name: Optional[str] = None
+
+
+OpenAIMessageParam = Annotated[
+    Union[
+        OpenAIUserMessageParam,
+        OpenAISystemMessageParam,
+        OpenAIAssistantMessageParam,
+        OpenAIToolMessageParam,
+        OpenAIDeveloperMessageParam,
+    ],
+    Field(discriminator="role"),
+]
+register_schema(OpenAIMessageParam, name="OpenAIMessageParam")
+
+
+@json_schema_type
+class OpenAIResponseFormatText(BaseModel):
+    type: Literal["text"] = "text"
+
+
+@json_schema_type
+class OpenAIJSONSchema(TypedDict, total=False):
+    name: str
+    description: Optional[str] = None
+    strict: Optional[bool] = None
+
+    # Pydantic BaseModel cannot be used with a schema param, since it already
+    # has one. And, we don't want to alias here because then have to handle
+    # that alias when converting to OpenAI params. So, to support schema,
+    # we use a TypedDict.
+    schema: Optional[Dict[str, Any]] = None
+
+
+@json_schema_type
+class OpenAIResponseFormatJSONSchema(BaseModel):
+    type: Literal["json_schema"] = "json_schema"
+    json_schema: OpenAIJSONSchema
+
+
+@json_schema_type
+class OpenAIResponseFormatJSONObject(BaseModel):
+    type: Literal["json_object"] = "json_object"
+
+
+OpenAIResponseFormatParam = Annotated[
+    Union[
+        OpenAIResponseFormatText,
+        OpenAIResponseFormatJSONSchema,
+        OpenAIResponseFormatJSONObject,
+    ],
+    Field(discriminator="type"),
+]
+register_schema(OpenAIResponseFormatParam, name="OpenAIResponseFormatParam")
+
+
+@json_schema_type
+class OpenAITopLogProb(BaseModel):
+    """The top log probability for a token from an OpenAI-compatible chat completion response.
+
+    :token: The token
+    :bytes: (Optional) The bytes for the token
+    :logprob: The log probability of the token
+    """
+
+    token: str
+    bytes: Optional[List[int]] = None
+    logprob: float
+
+
+@json_schema_type
+class OpenAITokenLogProb(BaseModel):
+    """The log probability for a token from an OpenAI-compatible chat completion response.
+
+    :token: The token
+    :bytes: (Optional) The bytes for the token
+    :logprob: The log probability of the token
+    :top_logprobs: The top log probabilities for the token
+    """
+
+    token: str
+    bytes: Optional[List[int]] = None
+    logprob: float
+    top_logprobs: List[OpenAITopLogProb]
+
+
+@json_schema_type
+class OpenAIChoiceLogprobs(BaseModel):
+    """The log probabilities for the tokens in the message from an OpenAI-compatible chat completion response.
+
+    :param content: (Optional) The log probabilities for the tokens in the message
+    :param refusal: (Optional) The log probabilities for the tokens in the message
+    """
+
+    content: Optional[List[OpenAITokenLogProb]] = None
+    refusal: Optional[List[OpenAITokenLogProb]] = None
+
+
+@json_schema_type
+class OpenAIChoiceDelta(BaseModel):
+    """A delta from an OpenAI-compatible chat completion streaming response.
+
+    :param content: (Optional) The content of the delta
+    :param refusal: (Optional) The refusal of the delta
+    :param role: (Optional) The role of the delta
+    :param tool_calls: (Optional) The tool calls of the delta
+    """
+
+    content: Optional[str] = None
+    refusal: Optional[str] = None
+    role: Optional[str] = None
+    tool_calls: Optional[List[OpenAIChatCompletionToolCall]] = None
+
+
+@json_schema_type
+class OpenAIChunkChoice(BaseModel):
+    """A chunk choice from an OpenAI-compatible chat completion streaming response.
+
+    :param delta: The delta from the chunk
+    :param finish_reason: The reason the model stopped generating
+    :param index: The index of the choice
+    :param logprobs: (Optional) The log probabilities for the tokens in the message
+    """
+
+    delta: OpenAIChoiceDelta
+    finish_reason: str
+    index: int
+    logprobs: Optional[OpenAIChoiceLogprobs] = None
+
+
+@json_schema_type
+class OpenAIChoice(BaseModel):
+    """A choice from an OpenAI-compatible chat completion response.
+
+    :param message: The message from the model
+    :param finish_reason: The reason the model stopped generating
+    :param index: The index of the choice
+    :param logprobs: (Optional) The log probabilities for the tokens in the message
+    """
+
+    message: OpenAIMessageParam
+    finish_reason: str
+    index: int
+    logprobs: Optional[OpenAIChoiceLogprobs] = None
+
+
+@json_schema_type
+class OpenAIChatCompletion(BaseModel):
+    """Response from an OpenAI-compatible chat completion request.
+
+    :param id: The ID of the chat completion
+    :param choices: List of choices
+    :param object: The object type, which will be "chat.completion"
+    :param created: The Unix timestamp in seconds when the chat completion was created
+    :param model: The model that was used to generate the chat completion
+    """
+
+    id: str
+    choices: List[OpenAIChoice]
+    object: Literal["chat.completion"] = "chat.completion"
+    created: int
+    model: str
+
+
+@json_schema_type
+class OpenAIChatCompletionChunk(BaseModel):
+    """Chunk from a streaming response to an OpenAI-compatible chat completion request.
+
+    :param id: The ID of the chat completion
+    :param choices: List of choices
+    :param object: The object type, which will be "chat.completion.chunk"
+    :param created: The Unix timestamp in seconds when the chat completion was created
+    :param model: The model that was used to generate the chat completion
+    """
+
+    id: str
+    choices: List[OpenAIChunkChoice]
+    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
+    created: int
+    model: str
+
+
+@json_schema_type
+class OpenAICompletionLogprobs(BaseModel):
+    """The log probabilities for the tokens in the message from an OpenAI-compatible completion response.
+
+    :text_offset: (Optional) The offset of the token in the text
+    :token_logprobs: (Optional) The log probabilities for the tokens
+    :tokens: (Optional) The tokens
+    :top_logprobs: (Optional) The top log probabilities for the tokens
+    """
+
+    text_offset: Optional[List[int]] = None
+    token_logprobs: Optional[List[float]] = None
+    tokens: Optional[List[str]] = None
+    top_logprobs: Optional[List[Dict[str, float]]] = None
+
+
+@json_schema_type
+class OpenAICompletionChoice(BaseModel):
+    """A choice from an OpenAI-compatible completion response.
+
+    :finish_reason: The reason the model stopped generating
+    :text: The text of the choice
+    :index: The index of the choice
+    :logprobs: (Optional) The log probabilities for the tokens in the choice
+    """
+
+    finish_reason: str
+    text: str
+    index: int
+    logprobs: Optional[OpenAIChoiceLogprobs] = None
+
+
+@json_schema_type
+class OpenAICompletion(BaseModel):
+    """Response from an OpenAI-compatible completion request.
+
+    :id: The ID of the completion
+    :choices: List of choices
+    :created: The Unix timestamp in seconds when the completion was created
+    :model: The model that was used to generate the completion
+    :object: The object type, which will be "text_completion"
+    """
+
+    id: str
+    choices: List[OpenAICompletionChoice]
+    created: int
+    model: str
+    object: Literal["text_completion"] = "text_completion"
+
+
 class ModelStore(Protocol):
     async def get_model(self, identifier: str) -> Model: ...
 
@@ -470,6 +816,16 @@ class EmbeddingTaskType(Enum):
     document = "document"
 
 
+@json_schema_type
+class BatchCompletionResponse(BaseModel):
+    batch: List[CompletionResponse]
+
+
+@json_schema_type
+class BatchChatCompletionResponse(BaseModel):
+    batch: List[ChatCompletionResponse]
+
+
 @runtime_checkable
 @trace_protocol
 class Inference(Protocol):
@@ -505,6 +861,17 @@ class Inference(Protocol):
         """
         ...
 
+    @webmethod(route="/inference/batch-completion", method="POST", experimental=True)
+    async def batch_completion(
+        self,
+        model_id: str,
+        content_batch: List[InterleavedContent],
+        sampling_params: Optional[SamplingParams] = None,
+        response_format: Optional[ResponseFormat] = None,
+        logprobs: Optional[LogProbConfig] = None,
+    ) -> BatchCompletionResponse:
+        raise NotImplementedError("Batch completion is not implemented")
+
     @webmethod(route="/inference/chat-completion", method="POST")
     async def chat_completion(
         self,
@@ -545,6 +912,19 @@ class Inference(Protocol):
         """
         ...
 
+    @webmethod(route="/inference/batch-chat-completion", method="POST", experimental=True)
+    async def batch_chat_completion(
+        self,
+        model_id: str,
+        messages_batch: List[List[Message]],
+        sampling_params: Optional[SamplingParams] = None,
+        tools: Optional[List[ToolDefinition]] = None,
+        tool_config: Optional[ToolConfig] = None,
+        response_format: Optional[ResponseFormat] = None,
+        logprobs: Optional[LogProbConfig] = None,
+    ) -> BatchChatCompletionResponse:
+        raise NotImplementedError("Batch chat completion is not implemented")
+
     @webmethod(route="/inference/embeddings", method="POST")
     async def embeddings(
         self,
@@ -564,3 +944,105 @@ class Inference(Protocol):
         :returns: An array of embeddings, one for each content. Each embedding is a list of floats. The dimensionality of the embedding is model-specific; you can check model metadata using /models/{model_id}
         """
         ...
+
+    @webmethod(route="/openai/v1/completions", method="POST")
+    async def openai_completion(
+        self,
+        # Standard OpenAI completion parameters
+        model: str,
+        prompt: Union[str, List[str], List[int], List[List[int]]],
+        best_of: Optional[int] = None,
+        echo: Optional[bool] = None,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        # vLLM-specific parameters
+        guided_choice: Optional[List[str]] = None,
+        prompt_logprobs: Optional[int] = None,
+    ) -> OpenAICompletion:
+        """Generate an OpenAI-compatible completion for the given prompt using the specified model.
+
+        :param model: The identifier of the model to use. The model must be registered with Llama Stack and available via the /models endpoint.
+        :param prompt: The prompt to generate a completion for
+        :param best_of: (Optional) The number of completions to generate
+        :param echo: (Optional) Whether to echo the prompt
+        :param frequency_penalty: (Optional) The penalty for repeated tokens
+        :param logit_bias: (Optional) The logit bias to use
+        :param logprobs: (Optional) The log probabilities to use
+        :param max_tokens: (Optional) The maximum number of tokens to generate
+        :param n: (Optional) The number of completions to generate
+        :param presence_penalty: (Optional) The penalty for repeated tokens
+        :param seed: (Optional) The seed to use
+        :param stop: (Optional) The stop tokens to use
+        :param stream: (Optional) Whether to stream the response
+        :param stream_options: (Optional) The stream options to use
+        :param temperature: (Optional) The temperature to use
+        :param top_p: (Optional) The top p to use
+        :param user: (Optional) The user to use
+        """
+        ...
+
+    @webmethod(route="/openai/v1/chat/completions", method="POST")
+    async def openai_chat_completion(
+        self,
+        model: str,
+        messages: List[OpenAIMessageParam],
+        frequency_penalty: Optional[float] = None,
+        function_call: Optional[Union[str, Dict[str, Any]]] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        parallel_tool_calls: Optional[bool] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[OpenAIResponseFormatParam] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+    ) -> Union[OpenAIChatCompletion, AsyncIterator[OpenAIChatCompletionChunk]]:
+        """Generate an OpenAI-compatible chat completion for the given messages using the specified model.
+
+        :param model: The identifier of the model to use. The model must be registered with Llama Stack and available via the /models endpoint.
+        :param messages: List of messages in the conversation
+        :param frequency_penalty: (Optional) The penalty for repeated tokens
+        :param function_call: (Optional) The function call to use
+        :param functions: (Optional) List of functions to use
+        :param logit_bias: (Optional) The logit bias to use
+        :param logprobs: (Optional) The log probabilities to use
+        :param max_completion_tokens: (Optional) The maximum number of tokens to generate
+        :param max_tokens: (Optional) The maximum number of tokens to generate
+        :param n: (Optional) The number of completions to generate
+        :param parallel_tool_calls: (Optional) Whether to parallelize tool calls
+        :param presence_penalty: (Optional) The penalty for repeated tokens
+        :param response_format: (Optional) The response format to use
+        :param seed: (Optional) The seed to use
+        :param stop: (Optional) The stop tokens to use
+        :param stream: (Optional) Whether to stream the response
+        :param stream_options: (Optional) The stream options to use
+        :param temperature: (Optional) The temperature to use
+        :param tool_choice: (Optional) The tool choice to use
+        :param tools: (Optional) The tools to use
+        :param top_logprobs: (Optional) The top log probabilities to use
+        :param top_p: (Optional) The top p to use
+        :param user: (Optional) The user to use
+        """
+        ...

llama_stack/apis/inspect/inspect.py

@@ -8,6 +8,7 @@ from typing import List, Protocol, runtime_checkable
 
 from pydantic import BaseModel
 
+from llama_stack.providers.datatypes import HealthStatus
 from llama_stack.schema_utils import json_schema_type, webmethod
 
 
@@ -20,8 +21,7 @@ class RouteInfo(BaseModel):
 
 @json_schema_type
 class HealthInfo(BaseModel):
-    status: str
-    # TODO: add a provider level status
+    status: HealthStatus
 
 
 @json_schema_type

llama_stack/apis/models/models.py

@@ -56,12 +56,35 @@ class ListModelsResponse(BaseModel):
     data: List[Model]
 
 
+@json_schema_type
+class OpenAIModel(BaseModel):
+    """A model from OpenAI.
+
+    :id: The ID of the model
+    :object: The object type, which will be "model"
+    :created: The Unix timestamp in seconds when the model was created
+    :owned_by: The owner of the model
+    """
+
+    id: str
+    object: Literal["model"] = "model"
+    created: int
+    owned_by: str
+
+
+class OpenAIListModelsResponse(BaseModel):
+    data: List[OpenAIModel]
+
+
 @runtime_checkable
 @trace_protocol
 class Models(Protocol):
     @webmethod(route="/models", method="GET")
     async def list_models(self) -> ListModelsResponse: ...
 
+    @webmethod(route="/openai/v1/models", method="GET")
+    async def openai_list_models(self) -> OpenAIListModelsResponse: ...
+
     @webmethod(route="/models/{model_id:path}", method="GET")
     async def get_model(
         self,

llama_stack/apis/post_training/post_training.py

@@ -60,11 +60,11 @@ class EfficiencyConfig(BaseModel):
 @json_schema_type
 class TrainingConfig(BaseModel):
     n_epochs: int
-    max_steps_per_epoch: int
-    gradient_accumulation_steps: int
-    max_validation_steps: int
-    data_config: DataConfig
-    optimizer_config: OptimizerConfig
+    max_steps_per_epoch: int = 1
+    gradient_accumulation_steps: int = 1
+    max_validation_steps: Optional[int] = 1
+    data_config: Optional[DataConfig] = None
+    optimizer_config: Optional[OptimizerConfig] = None
     efficiency_config: Optional[EfficiencyConfig] = None
     dtype: Optional[str] = "bf16"
 
@@ -177,9 +177,9 @@ class PostTraining(Protocol):
         training_config: TrainingConfig,
         hyperparam_search_config: Dict[str, Any],
         logger_config: Dict[str, Any],
-        model: str = Field(
-            default="Llama3.2-3B-Instruct",
-            description="Model descriptor from `llama model list`",
+        model: Optional[str] = Field(
+            default=None,
+            description="Model descriptor for training if not in provider config`",
         ),
         checkpoint_dir: Optional[str] = None,
         algorithm_config: Optional[AlgorithmConfig] = None,

llama_stack/apis/providers/providers.py

@@ -8,6 +8,7 @@ from typing import Any, Dict, List, Protocol, runtime_checkable
 
 from pydantic import BaseModel
 
+from llama_stack.providers.datatypes import HealthResponse
 from llama_stack.schema_utils import json_schema_type, webmethod
 
 
@@ -17,6 +18,7 @@ class ProviderInfo(BaseModel):
     provider_id: str
     provider_type: str
     config: Dict[str, Any]
+    health: HealthResponse
 
 
 class ListProvidersResponse(BaseModel):

llama_stack/cli/stack/build.py

@@ -57,7 +57,7 @@ class StackBuild(Subcommand):
             type=str,
             help=textwrap.dedent(
                 f"""[for image-type={"|".join(e.value for e in ImageType)}] Name of the conda or virtual environment to use for
-the build. If not specified, currently active Conda environment will be used if found.
+the build. If not specified, currently active environment will be used if found.
             """
             ),
             default=None,

llama_stack/cli/stack/run.py

@@ -45,7 +45,7 @@ class StackRun(Subcommand):
             "--image-name",
             type=str,
             default=os.environ.get("CONDA_DEFAULT_ENV"),
-            help="Name of the image to run. Defaults to the current conda environment",
+            help="Name of the image to run. Defaults to the current environment",
         )
         self.parser.add_argument(
             "--disable-ipv6",

llama_stack/distribution/inspect.py

@@ -17,6 +17,7 @@ from llama_stack.apis.inspect import (
 )
 from llama_stack.distribution.datatypes import StackRunConfig
 from llama_stack.distribution.server.endpoints import get_all_api_endpoints
+from llama_stack.providers.datatypes import HealthStatus
 
 
 class DistributionInspectConfig(BaseModel):
@@ -58,7 +59,7 @@ class DistributionInspectImpl(Inspect):
         return ListRoutesResponse(data=ret)
 
     async def health(self) -> HealthInfo:
-        return HealthInfo(status="OK")
+        return HealthInfo(status=HealthStatus.OK)
 
     async def version(self) -> VersionInfo:
         return VersionInfo(version=version("llama-stack"))

llama_stack/distribution/library_client.py

@@ -43,9 +43,9 @@ from llama_stack.distribution.server.endpoints import (
 from llama_stack.distribution.stack import (
     construct_stack,
     get_stack_run_config_from_template,
-    redact_sensitive_fields,
     replace_env_vars,
 )
+from llama_stack.distribution.utils.config import redact_sensitive_fields
 from llama_stack.distribution.utils.context import preserve_contexts_async_generator
 from llama_stack.distribution.utils.exec import in_notebook
 from llama_stack.providers.utils.telemetry.tracing import (

llama_stack/distribution/providers.py

@@ -4,14 +4,17 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import asyncio
+from typing import Any, Dict
 
 from pydantic import BaseModel
 
 from llama_stack.apis.providers import ListProvidersResponse, ProviderInfo, Providers
 from llama_stack.log import get_logger
+from llama_stack.providers.datatypes import HealthResponse, HealthStatus
 
 from .datatypes import StackRunConfig
-from .stack import redact_sensitive_fields
+from .utils.config import redact_sensitive_fields
 
 logger = get_logger(name=__name__, category="core")
 
@@ -41,19 +44,24 @@ class ProviderImpl(Providers):
     async def list_providers(self) -> ListProvidersResponse:
         run_config = self.config.run_config
         safe_config = StackRunConfig(**redact_sensitive_fields(run_config.model_dump()))
+        providers_health = await self.get_providers_health()
         ret = []
         for api, providers in safe_config.providers.items():
-            ret.extend(
-                [
+            for p in providers:
+                ret.append(
                     ProviderInfo(
                         api=api,
                         provider_id=p.provider_id,
                         provider_type=p.provider_type,
                         config=p.config,
+                        health=providers_health.get(api, {}).get(
+                            p.provider_id,
+                            HealthResponse(
+                                status=HealthStatus.NOT_IMPLEMENTED, message="Provider does not implement health check"
+                            ),
+                        ),
                     )
-                    for p in providers
-                ]
-            )
+                )
 
         return ListProvidersResponse(data=ret)
 
@@ -64,3 +72,57 @@ class ProviderImpl(Providers):
                 return p
 
         raise ValueError(f"Provider {provider_id} not found")
+
+    async def get_providers_health(self) -> Dict[str, Dict[str, HealthResponse]]:
+        """Get health status for all providers.
+
+        Returns:
+            Dict[str, Dict[str, HealthResponse]]: A dictionary mapping API names to provider health statuses.
+                Each API maps to a dictionary of provider IDs to their health responses.
+        """
+        providers_health: Dict[str, Dict[str, HealthResponse]] = {}
+        timeout = 1.0
+
+        async def check_provider_health(impl: Any) -> tuple[str, HealthResponse] | None:
+            # Skip special implementations (inspect/providers) that don't have provider specs
+            if not hasattr(impl, "__provider_spec__"):
+                return None
+            api_name = impl.__provider_spec__.api.name
+            if not hasattr(impl, "health"):
+                return (
+                    api_name,
+                    HealthResponse(
+                        status=HealthStatus.NOT_IMPLEMENTED, message="Provider does not implement health check"
+                    ),
+                )
+
+            try:
+                health = await asyncio.wait_for(impl.health(), timeout=timeout)
+                return api_name, health
+            except asyncio.TimeoutError:
+                return (
+                    api_name,
+                    HealthResponse(
+                        status=HealthStatus.ERROR, message=f"Health check timed out after {timeout} seconds"
+                    ),
+                )
+            except Exception as e:
+                return (
+                    api_name,
+                    HealthResponse(status=HealthStatus.ERROR, message=f"Health check failed: {str(e)}"),
+                )
+
+        # Create tasks for all providers
+        tasks = [check_provider_health(impl) for impl in self.deps.values()]
+
+        # Wait for all health checks to complete
+        results = await asyncio.gather(*tasks)
+
+        # Organize results by API and provider ID
+        for result in results:
+            if result is None:  # Skip special implementations
+                continue
+            api_name, health_response = result
+            providers_health[api_name] = health_response
+
+        return providers_health

llama_stack/distribution/resolver.py

@@ -41,7 +41,6 @@ from llama_stack.providers.datatypes import (
     Api,
     BenchmarksProtocolPrivate,
     DatasetsProtocolPrivate,
-    InlineProviderSpec,
     ModelsProtocolPrivate,
     ProviderSpec,
     RemoteProviderConfig,
@@ -230,50 +229,9 @@ def sort_providers_by_deps(
         {k: list(v.values()) for k, v in providers_with_specs.items()}
     )
 
-    # Append built-in "inspect" provider
-    apis = [x[1].spec.api for x in sorted_providers]
-    sorted_providers.append(
-        (
-            "inspect",
-            ProviderWithSpec(
-                provider_id="__builtin__",
-                provider_type="__builtin__",
-                config={"run_config": run_config.model_dump()},
-                spec=InlineProviderSpec(
-                    api=Api.inspect,
-                    provider_type="__builtin__",
-                    config_class="llama_stack.distribution.inspect.DistributionInspectConfig",
-                    module="llama_stack.distribution.inspect",
-                    api_dependencies=apis,
-                    deps__=[x.value for x in apis],
-                ),
-            ),
-        )
-    )
-
-    sorted_providers.append(
-        (
-            "providers",
-            ProviderWithSpec(
-                provider_id="__builtin__",
-                provider_type="__builtin__",
-                config={"run_config": run_config.model_dump()},
-                spec=InlineProviderSpec(
-                    api=Api.providers,
-                    provider_type="__builtin__",
-                    config_class="llama_stack.distribution.providers.ProviderImplConfig",
-                    module="llama_stack.distribution.providers",
-                    api_dependencies=apis,
-                    deps__=[x.value for x in apis],
-                ),
-            ),
-        )
-    )
-
     logger.debug(f"Resolved {len(sorted_providers)} providers")
     for api_str, provider in sorted_providers:
         logger.debug(f" {api_str} => {provider.provider_id}")
-        logger.debug("")
     return sorted_providers
 
 
@@ -400,6 +358,8 @@ def check_protocol_compliance(obj: Any, protocol: Any) -> None:
     mro = type(obj).__mro__
     for name, value in inspect.getmembers(protocol):
         if inspect.isfunction(value) and hasattr(value, "__webmethod__"):
+            if value.__webmethod__.experimental:
+                continue
             if not hasattr(obj, name):
                 missing_methods.append((name, "missing"))
             elif not callable(getattr(obj, name)):

llama_stack/distribution/routers/routers.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import asyncio
 import time
 from typing import Any, AsyncGenerator, AsyncIterator, Dict, List, Optional, Union
 
@@ -17,6 +18,8 @@ from llama_stack.apis.datasetio import DatasetIO
 from llama_stack.apis.datasets import DatasetPurpose, DataSource
 from llama_stack.apis.eval import BenchmarkConfig, Eval, EvaluateResponse, Job
 from llama_stack.apis.inference import (
+    BatchChatCompletionResponse,
+    BatchCompletionResponse,
     ChatCompletionResponse,
     ChatCompletionResponseEventType,
     ChatCompletionResponseStreamChunk,
@@ -35,6 +38,13 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
+from llama_stack.apis.inference.inference import (
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAICompletion,
+    OpenAIMessageParam,
+    OpenAIResponseFormatParam,
+)
 from llama_stack.apis.models import Model, ModelType
 from llama_stack.apis.safety import RunShieldResponse, Safety
 from llama_stack.apis.scoring import (
@@ -57,7 +67,7 @@ from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO
 from llama_stack.log import get_logger
 from llama_stack.models.llama.llama3.chat_format import ChatFormat
 from llama_stack.models.llama.llama3.tokenizer import Tokenizer
-from llama_stack.providers.datatypes import RoutingTable
+from llama_stack.providers.datatypes import HealthResponse, HealthStatus, RoutingTable
 from llama_stack.providers.utils.telemetry.tracing import get_current_span
 
 logger = get_logger(name=__name__, category="core")
@@ -333,6 +343,30 @@ class InferenceRouter(Inference):
             response.metrics = metrics if response.metrics is None else response.metrics + metrics
             return response
 
+    async def batch_chat_completion(
+        self,
+        model_id: str,
+        messages_batch: List[List[Message]],
+        tools: Optional[List[ToolDefinition]] = None,
+        tool_config: Optional[ToolConfig] = None,
+        sampling_params: Optional[SamplingParams] = None,
+        response_format: Optional[ResponseFormat] = None,
+        logprobs: Optional[LogProbConfig] = None,
+    ) -> BatchChatCompletionResponse:
+        logger.debug(
+            f"InferenceRouter.batch_chat_completion: {model_id=}, {len(messages_batch)=}, {sampling_params=}, {response_format=}, {logprobs=}",
+        )
+        provider = self.routing_table.get_provider_impl(model_id)
+        return await provider.batch_chat_completion(
+            model_id=model_id,
+            messages_batch=messages_batch,
+            tools=tools,
+            tool_config=tool_config,
+            sampling_params=sampling_params,
+            response_format=response_format,
+            logprobs=logprobs,
+        )
+
     async def completion(
         self,
         model_id: str,
@@ -397,6 +431,20 @@ class InferenceRouter(Inference):
             response.metrics = metrics if response.metrics is None else response.metrics + metrics
             return response
 
+    async def batch_completion(
+        self,
+        model_id: str,
+        content_batch: List[InterleavedContent],
+        sampling_params: Optional[SamplingParams] = None,
+        response_format: Optional[ResponseFormat] = None,
+        logprobs: Optional[LogProbConfig] = None,
+    ) -> BatchCompletionResponse:
+        logger.debug(
+            f"InferenceRouter.batch_completion: {model_id=}, {len(content_batch)=}, {sampling_params=}, {response_format=}, {logprobs=}",
+        )
+        provider = self.routing_table.get_provider_impl(model_id)
+        return await provider.batch_completion(model_id, content_batch, sampling_params, response_format, logprobs)
+
     async def embeddings(
         self,
         model_id: str,
@@ -419,6 +467,149 @@ class InferenceRouter(Inference):
             task_type=task_type,
         )
 
+    async def openai_completion(
+        self,
+        model: str,
+        prompt: Union[str, List[str], List[int], List[List[int]]],
+        best_of: Optional[int] = None,
+        echo: Optional[bool] = None,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        guided_choice: Optional[List[str]] = None,
+        prompt_logprobs: Optional[int] = None,
+    ) -> OpenAICompletion:
+        logger.debug(
+            f"InferenceRouter.openai_completion: {model=}, {stream=}, {prompt=}",
+        )
+        model_obj = await self.routing_table.get_model(model)
+        if model_obj is None:
+            raise ValueError(f"Model '{model}' not found")
+        if model_obj.model_type == ModelType.embedding:
+            raise ValueError(f"Model '{model}' is an embedding model and does not support completions")
+
+        params = dict(
+            model=model_obj.identifier,
+            prompt=prompt,
+            best_of=best_of,
+            echo=echo,
+            frequency_penalty=frequency_penalty,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_tokens=max_tokens,
+            n=n,
+            presence_penalty=presence_penalty,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            top_p=top_p,
+            user=user,
+            guided_choice=guided_choice,
+            prompt_logprobs=prompt_logprobs,
+        )
+
+        provider = self.routing_table.get_provider_impl(model_obj.identifier)
+        return await provider.openai_completion(**params)
+
+    async def openai_chat_completion(
+        self,
+        model: str,
+        messages: List[OpenAIMessageParam],
+        frequency_penalty: Optional[float] = None,
+        function_call: Optional[Union[str, Dict[str, Any]]] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        parallel_tool_calls: Optional[bool] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[OpenAIResponseFormatParam] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+    ) -> Union[OpenAIChatCompletion, AsyncIterator[OpenAIChatCompletionChunk]]:
+        logger.debug(
+            f"InferenceRouter.openai_chat_completion: {model=}, {stream=}, {messages=}",
+        )
+        model_obj = await self.routing_table.get_model(model)
+        if model_obj is None:
+            raise ValueError(f"Model '{model}' not found")
+        if model_obj.model_type == ModelType.embedding:
+            raise ValueError(f"Model '{model}' is an embedding model and does not support chat completions")
+
+        params = dict(
+            model=model_obj.identifier,
+            messages=messages,
+            frequency_penalty=frequency_penalty,
+            function_call=function_call,
+            functions=functions,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_completion_tokens=max_completion_tokens,
+            max_tokens=max_tokens,
+            n=n,
+            parallel_tool_calls=parallel_tool_calls,
+            presence_penalty=presence_penalty,
+            response_format=response_format,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            tool_choice=tool_choice,
+            tools=tools,
+            top_logprobs=top_logprobs,
+            top_p=top_p,
+            user=user,
+        )
+
+        provider = self.routing_table.get_provider_impl(model_obj.identifier)
+        return await provider.openai_chat_completion(**params)
+
+    async def health(self) -> Dict[str, HealthResponse]:
+        health_statuses = {}
+        timeout = 0.5
+        for provider_id, impl in self.routing_table.impls_by_provider_id.items():
+            try:
+                # check if the provider has a health method
+                if not hasattr(impl, "health"):
+                    continue
+                health = await asyncio.wait_for(impl.health(), timeout=timeout)
+                health_statuses[provider_id] = health
+            except asyncio.TimeoutError:
+                health_statuses[provider_id] = HealthResponse(
+                    status=HealthStatus.ERROR,
+                    message=f"Health check timed out after {timeout} seconds",
+                )
+            except NotImplementedError:
+                health_statuses[provider_id] = HealthResponse(status=HealthStatus.NOT_IMPLEMENTED)
+            except Exception as e:
+                health_statuses[provider_id] = HealthResponse(
+                    status=HealthStatus.ERROR, message=f"Health check failed: {str(e)}"
+                )
+        return health_statuses
+
 
 class SafetyRouter(Safety):
     def __init__(

llama_stack/distribution/routers/routing_tables.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 import logging
+import time
 import uuid
 from typing import Any, Dict, List, Optional
 
@@ -23,7 +24,7 @@ from llama_stack.apis.datasets import (
     RowsDataSource,
     URIDataSource,
 )
-from llama_stack.apis.models import ListModelsResponse, Model, Models, ModelType
+from llama_stack.apis.models import ListModelsResponse, Model, Models, ModelType, OpenAIListModelsResponse, OpenAIModel
 from llama_stack.apis.resource import ResourceType
 from llama_stack.apis.scoring_functions import (
     ListScoringFunctionsResponse,
@@ -254,6 +255,19 @@ class ModelsRoutingTable(CommonRoutingTableImpl, Models):
     async def list_models(self) -> ListModelsResponse:
         return ListModelsResponse(data=await self.get_all_with_type("model"))
 
+    async def openai_list_models(self) -> OpenAIListModelsResponse:
+        models = await self.get_all_with_type("model")
+        openai_models = [
+            OpenAIModel(
+                id=model.identifier,
+                object="model",
+                created=int(time.time()),
+                owned_by="llama_stack",
+            )
+            for model in models
+        ]
+        return OpenAIListModelsResponse(data=openai_models)
+
     async def get_model(self, model_id: str) -> Model:
         model = await self.get_object_by_identifier("model", model_id)
         if model is None:

llama_stack/distribution/server/server.py

@@ -38,10 +38,10 @@ from llama_stack.distribution.server.endpoints import (
 )
 from llama_stack.distribution.stack import (
     construct_stack,
-    redact_sensitive_fields,
     replace_env_vars,
     validate_env_pair,
 )
+from llama_stack.distribution.utils.config import redact_sensitive_fields
 from llama_stack.distribution.utils.context import preserve_contexts_async_generator
 from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import Api
@@ -229,15 +229,30 @@ class TracingMiddleware:
     def __init__(self, app, impls):
         self.app = app
         self.impls = impls
+        # FastAPI built-in paths that should bypass custom routing
+        self.fastapi_paths = ("/docs", "/redoc", "/openapi.json", "/favicon.ico", "/static")
 
     async def __call__(self, scope, receive, send):
         if scope.get("type") == "lifespan":
             return await self.app(scope, receive, send)
 
         path = scope.get("path", "")
+
+        # Check if the path is a FastAPI built-in path
+        if path.startswith(self.fastapi_paths):
+            # Pass through to FastAPI's built-in handlers
+            logger.debug(f"Bypassing custom routing for FastAPI built-in path: {path}")
+            return await self.app(scope, receive, send)
+
         if not hasattr(self, "endpoint_impls"):
             self.endpoint_impls = initialize_endpoint_impls(self.impls)
-        _, _, trace_path = find_matching_endpoint(scope.get("method", "GET"), path, self.endpoint_impls)
+
+        try:
+            _, _, trace_path = find_matching_endpoint(scope.get("method", "GET"), path, self.endpoint_impls)
+        except ValueError:
+            # If no matching endpoint is found, pass through to FastAPI
+            logger.debug(f"No matching endpoint found for path: {path}, falling back to FastAPI")
+            return await self.app(scope, receive, send)
 
         trace_context = await start_trace(trace_path, {"__location__": "server", "raw_path": path})
 
@@ -388,7 +403,12 @@ def main(args: Optional[argparse.Namespace] = None):
     safe_config = redact_sensitive_fields(config.model_dump())
     logger.info(yaml.dump(safe_config, indent=2))
 
-    app = FastAPI(lifespan=lifespan)
+    app = FastAPI(
+        lifespan=lifespan,
+        docs_url="/docs",
+        redoc_url="/redoc",
+        openapi_url="/openapi.json",
+    )
     if not os.environ.get("LLAMA_STACK_DISABLE_VERSION_CHECK"):
         app.add_middleware(ClientVersionMiddleware)
 

llama_stack/distribution/stack.py

@@ -35,6 +35,8 @@ from llama_stack.apis.vector_dbs import VectorDBs
 from llama_stack.apis.vector_io import VectorIO
 from llama_stack.distribution.datatypes import Provider, StackRunConfig
 from llama_stack.distribution.distribution import get_provider_registry
+from llama_stack.distribution.inspect import DistributionInspectConfig, DistributionInspectImpl
+from llama_stack.distribution.providers import ProviderImpl, ProviderImplConfig
 from llama_stack.distribution.resolver import ProviderRegistry, resolve_impls
 from llama_stack.distribution.store.registry import create_dist_registry
 from llama_stack.distribution.utils.dynamic import instantiate_class_type
@@ -96,7 +98,10 @@ async def register_resources(run_config: StackRunConfig, impls: Dict[Api, Any]):
 
         method = getattr(impls[api], register_method)
         for obj in objects:
-            await method(**obj.model_dump())
+            # we want to maintain the type information in arguments to method.
+            # instead of method(**obj.model_dump()), which may convert a typed attr to a dict,
+            # we use model_dump() to find all the attrs and then getattr to get the still typed value.
+            await method(**{k: getattr(obj, k) for k in obj.model_dump().keys()})
 
         method = getattr(impls[api], list_method)
         response = await method()
@@ -116,26 +121,6 @@ class EnvVarError(Exception):
         super().__init__(f"Environment variable '{var_name}' not set or empty{f' at {path}' if path else ''}")
 
 
-def redact_sensitive_fields(data: Dict[str, Any]) -> Dict[str, Any]:
-    """Redact sensitive information from config before printing."""
-    sensitive_patterns = ["api_key", "api_token", "password", "secret"]
-
-    def _redact_dict(d: Dict[str, Any]) -> Dict[str, Any]:
-        result = {}
-        for k, v in d.items():
-            if isinstance(v, dict):
-                result[k] = _redact_dict(v)
-            elif isinstance(v, list):
-                result[k] = [_redact_dict(i) if isinstance(i, dict) else i for i in v]
-            elif any(pattern in k.lower() for pattern in sensitive_patterns):
-                result[k] = "********"
-            else:
-                result[k] = v
-        return result
-
-    return _redact_dict(data)
-
-
 def replace_env_vars(config: Any, path: str = "") -> Any:
     if isinstance(config, dict):
         result = {}
@@ -212,6 +197,26 @@ def validate_env_pair(env_pair: str) -> tuple[str, str]:
         ) from e
 
 
+def add_internal_implementations(impls: Dict[Api, Any], run_config: StackRunConfig) -> None:
+    """Add internal implementations (inspect and providers) to the implementations dictionary.
+
+    Args:
+        impls: Dictionary of API implementations
+        run_config: Stack run configuration
+    """
+    inspect_impl = DistributionInspectImpl(
+        DistributionInspectConfig(run_config=run_config),
+        deps=impls,
+    )
+    impls[Api.inspect] = inspect_impl
+
+    providers_impl = ProviderImpl(
+        ProviderImplConfig(run_config=run_config),
+        deps=impls,
+    )
+    impls[Api.providers] = providers_impl
+
+
 # Produces a stack of providers for the given run config. Not all APIs may be
 # asked for in the run config.
 async def construct_stack(
@@ -219,6 +224,10 @@ async def construct_stack(
 ) -> Dict[Api, Any]:
     dist_registry, _ = await create_dist_registry(run_config.metadata_store, run_config.image_name)
     impls = await resolve_impls(run_config, provider_registry or get_provider_registry(run_config), dist_registry)
+
+    # Add internal implementations after all other providers are resolved
+    add_internal_implementations(impls, run_config)
+
     await register_resources(run_config, impls)
     return impls
 

llama_stack/distribution/start_stack.sh

@@ -18,6 +18,7 @@ VIRTUAL_ENV=${VIRTUAL_ENV:-}
 set -euo pipefail
 
 RED='\033[0;31m'
+GREEN='\033[0;32m'
 NC='\033[0m' # No Color
 
 error_handler() {
@@ -73,7 +74,7 @@ done
 PYTHON_BINARY="python"
 case "$env_type" in
   "venv")
-    if [ -n "$VIRTUAL_ENV" && "$VIRTUAL_ENV" == "$env_path_or_name" ]; then
+    if [ -n "$VIRTUAL_ENV" ] && [ "$VIRTUAL_ENV" == "$env_path_or_name" ]; then
         echo -e "${GREEN}Virtual environment already activated${NC}" >&2
     else
         # Activate virtual environment

llama_stack/distribution/ui/page/playground/rag.py

@@ -9,6 +9,7 @@ import uuid
 import streamlit as st
 from llama_stack_client import Agent, AgentEventLogger, RAGDocument
 
+from llama_stack.apis.common.content_types import ToolCallDelta
 from llama_stack.distribution.ui.modules.api import llama_stack_api
 from llama_stack.distribution.ui.modules.utils import data_url_from_file
 
@@ -16,9 +17,16 @@ from llama_stack.distribution.ui.modules.utils import data_url_from_file
 def rag_chat_page():
     st.title("🦙 RAG")
 
+    def reset_agent_and_chat():
+        st.session_state.clear()
+        st.cache_resource.clear()
+
+    def should_disable_input():
+        return "displayed_messages" in st.session_state and len(st.session_state.displayed_messages) > 0
+
     with st.sidebar:
         # File/Directory Upload Section
-        st.subheader("Upload Documents")
+        st.subheader("Upload Documents", divider=True)
         uploaded_files = st.file_uploader(
             "Upload file(s) or directory",
             accept_multiple_files=True,
@@ -29,11 +37,11 @@ def rag_chat_page():
             st.success(f"Successfully uploaded {len(uploaded_files)} files")
             # Add memory bank name input field
             vector_db_name = st.text_input(
-                "Vector Database Name",
+                "Document Collection Name",
                 value="rag_vector_db",
-                help="Enter a unique identifier for this vector database",
+                help="Enter a unique identifier for this document collection",
             )
-            if st.button("Create Vector Database"):
+            if st.button("Create Document Collection"):
                 documents = [
                     RAGDocument(
                         document_id=uploaded_file.name,
@@ -64,26 +72,45 @@ def rag_chat_page():
                 )
                 st.success("Vector database created successfully!")
 
-        st.subheader("Configure Agent")
+        st.subheader("RAG Parameters", divider=True)
+
+        rag_mode = st.radio(
+            "RAG mode",
+            ["Direct", "Agent-based"],
+            captions=[
+                "RAG is performed by directly retrieving the information and augmenting the user query",
+                "RAG is performed by an agent activating a dedicated knowledge search tool.",
+            ],
+            on_change=reset_agent_and_chat,
+            disabled=should_disable_input(),
+        )
+
         # select memory banks
         vector_dbs = llama_stack_api.client.vector_dbs.list()
         vector_dbs = [vector_db.identifier for vector_db in vector_dbs]
         selected_vector_dbs = st.multiselect(
-            "Select Vector Databases",
-            vector_dbs,
+            label="Select Document Collections to use in RAG queries",
+            options=vector_dbs,
+            on_change=reset_agent_and_chat,
+            disabled=should_disable_input(),
         )
 
+        st.subheader("Inference Parameters", divider=True)
         available_models = llama_stack_api.client.models.list()
         available_models = [model.identifier for model in available_models if model.model_type == "llm"]
         selected_model = st.selectbox(
-            "Choose a model",
-            available_models,
+            label="Choose a model",
+            options=available_models,
             index=0,
+            on_change=reset_agent_and_chat,
+            disabled=should_disable_input(),
         )
         system_prompt = st.text_area(
             "System Prompt",
             value="You are a helpful assistant. ",
             help="Initial instructions given to the AI to set its behavior and context",
+            on_change=reset_agent_and_chat,
+            disabled=should_disable_input(),
         )
         temperature = st.slider(
             "Temperature",
@@ -92,6 +119,8 @@ def rag_chat_page():
             value=0.0,
             step=0.1,
             help="Controls the randomness of the response. Higher values make the output more creative and unexpected, lower values make it more conservative and predictable",
+            on_change=reset_agent_and_chat,
+            disabled=should_disable_input(),
         )
 
         top_p = st.slider(
@@ -100,19 +129,23 @@ def rag_chat_page():
             max_value=1.0,
             value=0.95,
             step=0.1,
+            on_change=reset_agent_and_chat,
+            disabled=should_disable_input(),
         )
 
         # Add clear chat button to sidebar
         if st.button("Clear Chat", use_container_width=True):
-            st.session_state.clear()
-            st.cache_resource.clear()
+            reset_agent_and_chat()
+            st.rerun()
 
     # Chat Interface
     if "messages" not in st.session_state:
         st.session_state.messages = []
+    if "displayed_messages" not in st.session_state:
+        st.session_state.displayed_messages = []
 
     # Display chat history
-    for message in st.session_state.messages:
+    for message in st.session_state.displayed_messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
 
@@ -144,22 +177,18 @@ def rag_chat_page():
             ],
         )
 
-    agent = create_agent()
+    if rag_mode == "Agent-based":
+        agent = create_agent()
+        if "agent_session_id" not in st.session_state:
+            st.session_state["agent_session_id"] = agent.create_session(session_name=f"rag_demo_{uuid.uuid4()}")
 
-    if "agent_session_id" not in st.session_state:
-        st.session_state["agent_session_id"] = agent.create_session(session_name=f"rag_demo_{uuid.uuid4()}")
+        session_id = st.session_state["agent_session_id"]
 
-    session_id = st.session_state["agent_session_id"]
-
-    # Chat input
-    if prompt := st.chat_input("Ask a question about your documents"):
+    def agent_process_prompt(prompt):
         # Add user message to chat history
         st.session_state.messages.append({"role": "user", "content": prompt})
 
-        # Display user message
-        with st.chat_message("user"):
-            st.markdown(prompt)
-
+        # Send the prompt to the agent
         response = agent.create_turn(
             messages=[
                 {
@@ -187,6 +216,79 @@ def rag_chat_page():
             message_placeholder.markdown(full_response)
 
             st.session_state.messages.append({"role": "assistant", "content": full_response})
+            st.session_state.displayed_messages.append({"role": "assistant", "content": full_response})
+
+    def direct_process_prompt(prompt):
+        # Add the system prompt in the beginning of the conversation
+        if len(st.session_state.messages) == 0:
+            st.session_state.messages.append({"role": "system", "content": system_prompt})
+
+        # Query the vector DB
+        rag_response = llama_stack_api.client.tool_runtime.rag_tool.query(
+            content=prompt, vector_db_ids=list(selected_vector_dbs)
+        )
+        prompt_context = rag_response.content
+
+        with st.chat_message("assistant"):
+            retrieval_message_placeholder = st.empty()
+            message_placeholder = st.empty()
+            full_response = ""
+            retrieval_response = ""
+
+            # Display the retrieved content
+            retrieval_response += str(prompt_context)
+            retrieval_message_placeholder.info(retrieval_response)
+
+            # Construct the extended prompt
+            extended_prompt = f"Please answer the following query using the context below.\n\nCONTEXT:\n{prompt_context}\n\nQUERY:\n{prompt}"
+
+            # Run inference directly
+            st.session_state.messages.append({"role": "user", "content": extended_prompt})
+            response = llama_stack_api.client.inference.chat_completion(
+                messages=st.session_state.messages,
+                model_id=selected_model,
+                sampling_params={
+                    "strategy": strategy,
+                },
+                stream=True,
+            )
+
+            # Display assistant response
+            for chunk in response:
+                response_delta = chunk.event.delta
+                if isinstance(response_delta, ToolCallDelta):
+                    retrieval_response += response_delta.tool_call.replace("====", "").strip()
+                    retrieval_message_placeholder.info(retrieval_response)
+                else:
+                    full_response += chunk.event.delta.text
+                    message_placeholder.markdown(full_response + "▌")
+            message_placeholder.markdown(full_response)
+
+        response_dict = {"role": "assistant", "content": full_response, "stop_reason": "end_of_message"}
+        st.session_state.messages.append(response_dict)
+        st.session_state.displayed_messages.append(response_dict)
+
+    # Chat input
+    if prompt := st.chat_input("Ask a question about your documents"):
+        # Add user message to chat history
+        st.session_state.displayed_messages.append({"role": "user", "content": prompt})
+
+        # Display user message
+        with st.chat_message("user"):
+            st.markdown(prompt)
+
+        # store the prompt to process it after page refresh
+        st.session_state.prompt = prompt
+
+        # force page refresh to disable the settings widgets
+        st.rerun()
+
+    if "prompt" in st.session_state and st.session_state.prompt is not None:
+        if rag_mode == "Agent-based":
+            agent_process_prompt(st.session_state.prompt)
+        else:  # rag_mode == "Direct"
+            direct_process_prompt(st.session_state.prompt)
+        st.session_state.prompt = None
 
 
 rag_chat_page()

llama_stack/distribution/utils/config.py

@@ -0,0 +1,30 @@
+# 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.
+
+from typing import Any, Dict
+
+
+def redact_sensitive_fields(data: Dict[str, Any]) -> Dict[str, Any]:
+    """Redact sensitive information from config before printing."""
+    sensitive_patterns = ["api_key", "api_token", "password", "secret"]
+
+    def _redact_value(v: Any) -> Any:
+        if isinstance(v, dict):
+            return _redact_dict(v)
+        elif isinstance(v, list):
+            return [_redact_value(i) for i in v]
+        return v
+
+    def _redact_dict(d: Dict[str, Any]) -> Dict[str, Any]:
+        result = {}
+        for k, v in d.items():
+            if any(pattern in k.lower() for pattern in sensitive_patterns):
+                result[k] = "********"
+            else:
+                result[k] = _redact_value(v)
+        return result
+
+    return _redact_dict(data)

llama_stack/models/llama/llama3/chat_format.py

@@ -226,7 +226,6 @@ class ChatFormat:
                     arguments_json=json.dumps(tool_arguments),
                 )
             )
-            content = ""
 
         return RawMessage(
             role="assistant",

llama_stack/models/llama/llama3/generation.py

@@ -140,7 +140,12 @@ class Llama3:
 
         return Llama3(model, tokenizer, model_args)
 
-    def __init__(self, model: Transformer | CrossAttentionTransformer, tokenizer: Tokenizer, args: ModelArgs):
+    def __init__(
+        self,
+        model: Transformer | CrossAttentionTransformer,
+        tokenizer: Tokenizer,
+        args: ModelArgs,
+    ):
         self.args = args
         self.model = model
         self.tokenizer = tokenizer
@@ -149,7 +154,7 @@ class Llama3:
     @torch.inference_mode()
     def generate(
         self,
-        model_inputs: List[LLMInput],
+        llm_inputs: List[LLMInput],
         temperature: float = 0.6,
         top_p: float = 0.9,
         max_gen_len: Optional[int] = None,
@@ -164,15 +169,15 @@ class Llama3:
 
         print_model_input = print_model_input or os.environ.get("LLAMA_MODELS_DEBUG", "0") == "1"
         if print_model_input:
-            for inp in model_inputs:
+            for inp in llm_inputs:
                 tokens_to_print = [self.formatter.vision_token if t == 128256 else t for t in inp.tokens]
                 cprint(
                     "Input to model:\n" + self.tokenizer.decode(tokens_to_print) + "\n",
                     "red",
                 )
-        prompt_tokens = [inp.tokens for inp in model_inputs]
+        prompt_tokens = [inp.tokens for inp in llm_inputs]
 
-        bsz = len(model_inputs)
+        bsz = len(llm_inputs)
         assert bsz <= params.max_batch_size, (bsz, params.max_batch_size)
 
         min_prompt_len = min(len(t) for t in prompt_tokens)
@@ -193,8 +198,8 @@ class Llama3:
 
         is_vision = not isinstance(self.model, Transformer)
         if is_vision:
-            images = [inp.vision.images if inp.vision is not None else [] for inp in model_inputs]
-            mask = [inp.vision.mask if inp.vision is not None else [] for inp in model_inputs]
+            images = [inp.vision.images if inp.vision is not None else [] for inp in llm_inputs]
+            mask = [inp.vision.mask if inp.vision is not None else [] for inp in llm_inputs]
 
             xattn_caches, cross_attention_masks, full_text_row_masked_out_mask = self.model.compute_vision_tokens_masks(
                 batch_images=images,
@@ -229,7 +234,7 @@ class Llama3:
         for cur_pos in range(min_prompt_len, total_len):
             if is_vision:
                 position_ids = torch.arange(prev_pos, cur_pos, dtype=torch.long)
-                text_only_inference = all(inp.vision is None for inp in model_inputs)
+                text_only_inference = all(inp.vision is None for inp in llm_inputs)
                 logits = self.model.forward(
                     position_ids,
                     tokens,
@@ -285,7 +290,7 @@ class Llama3:
                         source="output",
                         logprobs=(token_logprobs[idx, cur_pos : cur_pos + 1].tolist() if logprobs else None),
                         batch_idx=idx,
-                        finished=eos_reached[idx],
+                        finished=eos_reached[idx].item(),
                         ignore_token=cur_pos < len(prompt_tokens[idx]),
                     )
                 )

llama_stack/models/llama/llama3/prompt_templates/system_prompts.py

@@ -229,6 +229,11 @@ class PythonListCustomToolGenerator(PromptTemplateGeneratorBase):  # noqa: N801
         You are an expert in composing functions. You are given a question and a set of possible functions.
         Based on the question, you may or may not need to make one function/tool call to achieve the purpose.
 
+        If you decide to invoke any of the function(s), you MUST put it in the format of [func_name1(params_name1=params_value1, params_name2=params_value2...), func_name2(params)]
+        If you decide to invoke a function, you SHOULD NOT include any other text in the response. besides the function call in the above format.
+        For a boolean parameter, be sure to use `True` or `False` (capitalized) for the value.
+
+
         {{ function_description }}
         """.strip("\n")
     )
@@ -243,10 +248,6 @@ class PythonListCustomToolGenerator(PromptTemplateGeneratorBase):  # noqa: N801
     def _gen_function_description(self, custom_tools: List[ToolDefinition]) -> PromptTemplate:
         template_str = textwrap.dedent(
             """
-            If you decide to invoke any of the function(s), you MUST put it in the format of [func_name1(params_name1=params_value1, params_name2=params_value2...), func_name2(params)]
-            For a boolean parameter, be sure to use `True` or `False` (capitalized) for the value.
-            You SHOULD NOT include any other text in the response.
-
             Here is a list of functions in JSON format that you can invoke.
 
             [

llama_stack/models/llama/llama3/tool_utils.py

@@ -4,13 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# top-level folder for each specific model found within the models/ directory at
-# the top-level of this source tree.
-import ast
 import json
 import re
 from typing import Optional, Tuple
@@ -35,80 +28,141 @@ def is_json(s):
     return True
 
 
-def is_valid_python_list(input_string):
-    """Check if the input string is a valid Python list of function calls"""
-    try:
-        # Try to parse the string
-        tree = ast.parse(input_string)
-
-        # Check if it's a single expression
-        if len(tree.body) != 1 or not isinstance(tree.body[0], ast.Expr):
-            return False
-
-        # Check if the expression is a list
-        expr = tree.body[0].value
-        if not isinstance(expr, ast.List):
-            return False
-
-        # Check if the list is empty
-        if len(expr.elts) == 0:
-            return False
-
-        # Check if all elements in the list are function calls
-        for element in expr.elts:
-            if not isinstance(element, ast.Call):
-                return False
-
-            # Check if the function call has a valid name
-            if not isinstance(element.func, ast.Name):
-                return False
-
-            # Check if all arguments are keyword arguments
-            if element.args or not all(isinstance(arg, ast.keyword) for arg in element.keywords):
-                return False
-
-        return True
-
-    except SyntaxError:
-        # If parsing fails, it's not a valid Python expression
-        return False
-
-
-def parse_python_list_for_function_calls(input_string):
+def parse_llama_tool_call_format(input_string):
     """
-    Parse a Python list of function calls and
-    return a list of tuples containing the function name and arguments
-    """
-    # Parse the string into an AST
-    tree = ast.parse(input_string)
+    Parse tool calls in the format:
+    [func_name1(params_name1=params_value1, params_name2=params_value2...), func_name2(params)]
 
-    # Ensure the input is a list
-    if not isinstance(tree.body[0], ast.Expr) or not isinstance(tree.body[0].value, ast.List):
-        raise ValueError("Input must be a list of function calls")
+    Returns a list of (function_name, arguments_dict) tuples or None if parsing fails.
+    """
+    # Strip outer brackets and whitespace
+    input_string = input_string.strip()
+    if not (input_string.startswith("[") and input_string.endswith("]")):
+        return None
+
+    content = input_string[1:-1].strip()
+    if not content:
+        return None
 
     result = []
 
-    # Iterate through each function call in the list
-    for node in tree.body[0].value.elts:
-        if isinstance(node, ast.Call):
-            function_name = node.func.id
-            function_args = {}
+    # State variables for parsing
+    pos = 0
+    length = len(content)
 
-            # Extract keyword arguments
-            for keyword in node.keywords:
-                try:
-                    function_args[keyword.arg] = ast.literal_eval(keyword.value)
-                except ValueError as e:
-                    logger.error(
-                        f"Error parsing tool call argument '{keyword.arg}': {e}, full input string: '{input_string}'"
-                    )
-                    raise ValueError(
-                        f"Error parsing tool call argument '{keyword.arg}', full input string: '{input_string}'"
-                    ) from e
+    while pos < length:
+        # Find function name
+        name_end = content.find("(", pos)
+        if name_end == -1:
+            break
 
-            result.append((function_name, function_args))
+        func_name = content[pos:name_end].strip()
 
-    return result
+        # Find closing parenthesis for this function call
+        paren_level = 1
+        args_start = name_end + 1
+        args_end = args_start
+
+        while args_end < length and paren_level > 0:
+            if content[args_end] == "(":
+                paren_level += 1
+            elif content[args_end] == ")":
+                paren_level -= 1
+            args_end += 1
+
+        if paren_level != 0:
+            # Unmatched parentheses
+            return None
+
+        # Parse arguments
+        args_str = content[args_start : args_end - 1].strip()
+        args_dict = {}
+
+        if args_str:
+            # Split by commas, but respect nested structures
+            parts = []
+            part_start = 0
+            in_quotes = False
+            quote_char = None
+            nested_level = 0
+
+            for i, char in enumerate(args_str):
+                if char in ('"', "'") and (i == 0 or args_str[i - 1] != "\\"):
+                    if not in_quotes:
+                        in_quotes = True
+                        quote_char = char
+                    elif char == quote_char:
+                        in_quotes = False
+                        quote_char = None
+                elif not in_quotes:
+                    if char in ("{", "["):
+                        nested_level += 1
+                    elif char in ("}", "]"):
+                        nested_level -= 1
+                    elif char == "," and nested_level == 0:
+                        parts.append(args_str[part_start:i].strip())
+                        part_start = i + 1
+
+            parts.append(args_str[part_start:].strip())
+
+            # Process each key=value pair
+            for part in parts:
+                if "=" in part:
+                    key, value = part.split("=", 1)
+                    key = key.strip()
+                    value = value.strip()
+
+                    # Try to convert value to appropriate Python type
+                    if (value.startswith('"') and value.endswith('"')) or (
+                        value.startswith("'") and value.endswith("'")
+                    ):
+                        # String
+                        value = value[1:-1]
+                    elif value.lower() == "true":
+                        value = True
+                    elif value.lower() == "false":
+                        value = False
+                    elif value.lower() == "none":
+                        value = None
+                    elif value.startswith("{") and value.endswith("}"):
+                        # This is a nested dictionary
+                        try:
+                            # Try to parse as JSON
+                            value = json.loads(value.replace("'", '"'))
+                        except json.JSONDecodeError:
+                            # Keep as string if parsing fails
+                            pass
+                    elif value.startswith("[") and value.endswith("]"):
+                        # This is a nested list
+                        try:
+                            # Try to parse as JSON
+                            value = json.loads(value.replace("'", '"'))
+                        except json.JSONDecodeError:
+                            # Keep as string if parsing fails
+                            pass
+                    else:
+                        # Try to convert to number
+                        try:
+                            if "." in value:
+                                value = float(value)
+                            else:
+                                value = int(value)
+                        except ValueError:
+                            # Keep as string if not a valid number
+                            pass
+
+                    args_dict[key] = value
+
+        result.append((func_name, args_dict))
+
+        # Move to the next function call
+        pos = args_end
+
+        # Skip the comma between function calls if present
+        if pos < length and content[pos] == ",":
+            pos += 1
+
+    return result if result else None
 
 
 class ToolUtils:
@@ -150,17 +204,19 @@ class ToolUtils:
                 return None
         elif is_json(message_body):
             response = json.loads(message_body)
-            if ("type" in response and response["type"] == "function") or ("name" in response):
+            if ("type" in response and response["type"] == "function") or (
+                "name" in response and "parameters" in response
+            ):
                 function_name = response["name"]
                 args = response["parameters"]
                 return function_name, args
             else:
                 return None
-        elif is_valid_python_list(message_body):
-            res = parse_python_list_for_function_calls(message_body)
+        elif function_calls := parse_llama_tool_call_format(message_body):
             # FIXME: Enable multiple tool calls
-            return res[0]
+            return function_calls[0]
         else:
+            logger.debug(f"Did not parse tool call from message body: {message_body}")
             return None
 
     @staticmethod

llama_stack/models/llama/llama4/chat_format.py

@@ -301,7 +301,6 @@ class ChatFormat:
                     arguments=tool_arguments,
                 )
             )
-            content = ""
 
         return RawMessage(
             role="assistant",

llama_stack/models/llama/llama4/generation.py

@@ -233,7 +233,7 @@ class Llama4:
                         source="output",
                         logprobs=(token_logprobs[idx, cur_pos : cur_pos + 1].tolist() if logprobs else None),
                         batch_idx=idx,
-                        finished=eos_reached[idx],
+                        finished=eos_reached[idx].item(),
                         ignore_token=cur_pos < len(prompt_tokens[idx]),
                     )
                 )

llama_stack/models/llama/llama4/tokenizer.py

@@ -56,8 +56,8 @@ LLAMA4_TEXT_POST_TRAIN_SPECIAL_TOKENS = [
     "<|text_post_train_reserved_special_token_3|>",
     "<|text_post_train_reserved_special_token_4|>",
     "<|text_post_train_reserved_special_token_5|>",
-    "<|text_post_train_reserved_special_token_6|>",
-    "<|text_post_train_reserved_special_token_7|>",
+    "<|python_start|>",
+    "<|python_end|>",
     "<|finetune_right_pad|>",
 ] + get_reserved_special_tokens(
     "text_post_train", 61, 8

llama_stack/providers/datatypes.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+from enum import Enum
 from typing import Any, List, Optional, Protocol
 from urllib.parse import urlparse
 
@@ -201,3 +202,12 @@ def remote_provider_spec(
         adapter=adapter,
         api_dependencies=api_dependencies or [],
     )
+
+
+class HealthStatus(str, Enum):
+    OK = "OK"
+    ERROR = "Error"
+    NOT_IMPLEMENTED = "Not Implemented"
+
+
+HealthResponse = dict[str, Any]

llama_stack/providers/inline/inference/meta_reference/config.py

@@ -52,14 +52,17 @@ class MetaReferenceInferenceConfig(BaseModel):
         checkpoint_dir: str = "${env.CHECKPOINT_DIR:null}",
         quantization_type: str = "${env.QUANTIZATION_TYPE:bf16}",
         model_parallel_size: str = "${env.MODEL_PARALLEL_SIZE:0}",
+        max_batch_size: str = "${env.MAX_BATCH_SIZE:1}",
+        max_seq_len: str = "${env.MAX_SEQ_LEN:4096}",
         **kwargs,
     ) -> Dict[str, Any]:
         return {
             "model": model,
-            "max_seq_len": 4096,
             "checkpoint_dir": checkpoint_dir,
             "quantization": {
                 "type": quantization_type,
             },
             "model_parallel_size": model_parallel_size,
+            "max_batch_size": max_batch_size,
+            "max_seq_len": max_seq_len,
         }

llama_stack/providers/inline/inference/meta_reference/generators.py

@@ -22,7 +22,7 @@ from llama_stack.models.llama.llama3.generation import Llama3
 from llama_stack.models.llama.llama3.tokenizer import Tokenizer as Llama3Tokenizer
 from llama_stack.models.llama.llama4.generation import Llama4
 from llama_stack.models.llama.llama4.tokenizer import Tokenizer as Llama4Tokenizer
-from llama_stack.models.llama.sku_types import Model
+from llama_stack.models.llama.sku_types import Model, ModelFamily
 from llama_stack.providers.utils.inference.prompt_adapter import (
     ChatCompletionRequestWithRawContent,
     CompletionRequestWithRawContent,
@@ -113,8 +113,7 @@ def _infer_tool_prompt_format(request: ChatCompletionRequestWithRawContent):
         return get_default_tool_prompt_format(request.model)
 
 
-# TODO: combine Llama3 and Llama4 generators since they are almost identical now
-class Llama4Generator:
+class LlamaGenerator:
     def __init__(
         self,
         config: MetaReferenceInferenceConfig,
@@ -144,7 +143,8 @@ class Llama4Generator:
         else:
             quantization_mode = None
 
-        self.inner_generator = Llama4.build(
+        cls = Llama4 if llama_model.model_family == ModelFamily.llama4 else Llama3
+        self.inner_generator = cls.build(
             ckpt_dir=ckpt_dir,
             max_seq_len=config.max_seq_len,
             max_batch_size=config.max_batch_size,
@@ -158,142 +158,55 @@ class Llama4Generator:
 
     def completion(
         self,
-        request: CompletionRequestWithRawContent,
+        request_batch: List[CompletionRequestWithRawContent],
     ) -> Generator:
-        sampling_params = request.sampling_params or SamplingParams()
+        first_request = request_batch[0]
+        sampling_params = first_request.sampling_params or SamplingParams()
         max_gen_len = sampling_params.max_tokens
         if max_gen_len is None or max_gen_len == 0 or max_gen_len >= self.args.max_seq_len:
             max_gen_len = self.args.max_seq_len - 1
 
         temperature, top_p = _infer_sampling_params(sampling_params)
         for result in self.inner_generator.generate(
-            llm_inputs=[self.formatter.encode_content(request.content)],
+            llm_inputs=[self.formatter.encode_content(request.content) for request in request_batch],
             max_gen_len=max_gen_len,
             temperature=temperature,
             top_p=top_p,
-            logprobs=bool(request.logprobs),
+            logprobs=bool(first_request.logprobs),
             echo=False,
             logits_processor=get_logits_processor(
                 self.tokenizer,
                 self.args.vocab_size,
-                request.response_format,
+                first_request.response_format,
             ),
         ):
-            yield result[0]
+            yield result
 
     def chat_completion(
         self,
-        request: ChatCompletionRequestWithRawContent,
+        request_batch: List[ChatCompletionRequestWithRawContent],
     ) -> Generator:
-        sampling_params = request.sampling_params or SamplingParams()
+        first_request = request_batch[0]
+        sampling_params = first_request.sampling_params or SamplingParams()
         max_gen_len = sampling_params.max_tokens
         if max_gen_len is None or max_gen_len == 0 or max_gen_len >= self.args.max_seq_len:
             max_gen_len = self.args.max_seq_len - 1
 
         temperature, top_p = _infer_sampling_params(sampling_params)
         for result in self.inner_generator.generate(
-            llm_inputs=[self.formatter.encode_dialog_prompt(request.messages, _infer_tool_prompt_format(request))],
+            llm_inputs=[
+                self.formatter.encode_dialog_prompt(request.messages, _infer_tool_prompt_format(request))
+                for request in request_batch
+            ],
             max_gen_len=max_gen_len,
             temperature=temperature,
             top_p=top_p,
-            logprobs=bool(request.logprobs),
+            logprobs=bool(first_request.logprobs),
             echo=False,
             logits_processor=get_logits_processor(
                 self.tokenizer,
                 self.args.vocab_size,
-                request.response_format,
+                first_request.response_format,
             ),
         ):
-            yield result[0]
-
-
-class Llama3Generator:
-    def __init__(
-        self,
-        config: MetaReferenceInferenceConfig,
-        model_id: str,
-        llama_model: Model,
-    ):
-        if config.checkpoint_dir and config.checkpoint_dir != "null":
-            ckpt_dir = config.checkpoint_dir
-        else:
-            resolved_model = resolve_model(model_id)
-            if resolved_model is None:
-                # if the model is not a native llama model, get the default checkpoint_dir based on model id
-                ckpt_dir = model_checkpoint_dir(model_id)
-            else:
-                # if the model is a native llama model, get the default checkpoint_dir based on model core_model_id value
-                ckpt_dir = model_checkpoint_dir(resolved_model.descriptor())
-
-        if config.quantization:
-            if config.quantization.type == "fp8_mixed":
-                quantization_mode = QuantizationMode.fp8_mixed
-            elif config.quantization.type == "int4_mixed":
-                quantization_mode = QuantizationMode.int4_mixed
-            elif config.quantization.type == "bf16":
-                quantization_mode = None
-            else:
-                raise ValueError(f"Unsupported quantization mode {config.quantization}")
-        else:
-            quantization_mode = None
-
-        self.inner_generator = Llama3.build(
-            ckpt_dir=ckpt_dir,
-            max_seq_len=config.max_seq_len,
-            max_batch_size=config.max_batch_size,
-            world_size=config.model_parallel_size or llama_model.pth_file_count,
-            quantization_mode=quantization_mode,
-        )
-        self.tokenizer = self.inner_generator.tokenizer
-        self.args = self.inner_generator.args
-        self.formatter = self.inner_generator.formatter
-
-    def completion(
-        self,
-        request: CompletionRequestWithRawContent,
-    ) -> Generator:
-        sampling_params = request.sampling_params or SamplingParams()
-        max_gen_len = sampling_params.max_tokens
-        if max_gen_len is None or max_gen_len == 0 or max_gen_len >= self.args.max_seq_len:
-            max_gen_len = self.args.max_seq_len - 1
-
-        temperature, top_p = _infer_sampling_params(sampling_params)
-        for result in self.inner_generator.generate(
-            model_inputs=[self.formatter.encode_content(request.content)],
-            max_gen_len=max_gen_len,
-            temperature=temperature,
-            top_p=top_p,
-            logprobs=bool(request.logprobs),
-            echo=False,
-            logits_processor=get_logits_processor(
-                self.tokenizer,
-                self.args.vocab_size,
-                request.response_format,
-            ),
-        ):
-            yield result[0]
-
-    def chat_completion(
-        self,
-        request: ChatCompletionRequestWithRawContent,
-    ) -> Generator:
-        sampling_params = request.sampling_params or SamplingParams()
-        max_gen_len = sampling_params.max_tokens
-        if max_gen_len is None or max_gen_len == 0 or max_gen_len >= self.args.max_seq_len:
-            max_gen_len = self.args.max_seq_len - 1
-
-        temperature, top_p = _infer_sampling_params(sampling_params)
-        for result in self.inner_generator.generate(
-            model_inputs=[self.formatter.encode_dialog_prompt(request.messages, _infer_tool_prompt_format(request))],
-            max_gen_len=max_gen_len,
-            temperature=temperature,
-            top_p=top_p,
-            logprobs=bool(request.logprobs),
-            echo=False,
-            logits_processor=get_logits_processor(
-                self.tokenizer,
-                self.args.vocab_size,
-                request.response_format,
-            ),
-        ):
-            yield result[0]
+            yield result

llama_stack/providers/inline/inference/meta_reference/inference.py

@@ -5,10 +5,10 @@
 # the root directory of this source tree.
 
 import asyncio
-import logging
 import os
 from typing import AsyncGenerator, List, Optional, Union
 
+from pydantic import BaseModel
 from termcolor import cprint
 
 from llama_stack.apis.common.content_types import (
@@ -17,6 +17,8 @@ from llama_stack.apis.common.content_types import (
     ToolCallParseStatus,
 )
 from llama_stack.apis.inference import (
+    BatchChatCompletionResponse,
+    BatchCompletionResponse,
     ChatCompletionRequest,
     ChatCompletionResponse,
     ChatCompletionResponseEvent,
@@ -38,8 +40,10 @@ from llama_stack.apis.inference import (
     ToolConfig,
     ToolDefinition,
     ToolPromptFormat,
+    UserMessage,
 )
 from llama_stack.apis.models import Model, ModelType
+from llama_stack.log import get_logger
 from llama_stack.models.llama.llama3.chat_format import ChatFormat as Llama3ChatFormat
 from llama_stack.models.llama.llama3.tokenizer import Tokenizer as Llama3Tokenizer
 from llama_stack.models.llama.llama4.chat_format import ChatFormat as Llama4ChatFormat
@@ -54,6 +58,10 @@ from llama_stack.providers.utils.inference.model_registry import (
     ModelRegistryHelper,
     build_hf_repo_model_entry,
 )
+from llama_stack.providers.utils.inference.openai_compat import (
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
+)
 from llama_stack.providers.utils.inference.prompt_adapter import (
     augment_content_with_response_format_prompt,
     chat_completion_request_to_messages,
@@ -61,24 +69,22 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
 )
 
 from .config import MetaReferenceInferenceConfig
-from .generators import Llama3Generator, Llama4Generator
+from .generators import LlamaGenerator
 from .model_parallel import LlamaModelParallelGenerator
 
-log = logging.getLogger(__name__)
+log = get_logger(__name__, category="inference")
 # there's a single model parallel process running serving the model. for now,
 # we don't support multiple concurrent requests to this process.
 SEMAPHORE = asyncio.Semaphore(1)
 
 
-def llama3_builder_fn(config: MetaReferenceInferenceConfig, model_id: str, llama_model: Model) -> Llama3Generator:
-    return Llama3Generator(config, model_id, llama_model)
-
-
-def llama4_builder_fn(config: MetaReferenceInferenceConfig, model_id: str, llama_model: Model) -> Llama4Generator:
-    return Llama4Generator(config, model_id, llama_model)
+def llama_builder_fn(config: MetaReferenceInferenceConfig, model_id: str, llama_model: Model) -> LlamaGenerator:
+    return LlamaGenerator(config, model_id, llama_model)
 
 
 class MetaReferenceInferenceImpl(
+    OpenAICompletionToLlamaStackMixin,
+    OpenAIChatCompletionToLlamaStackMixin,
     SentenceTransformerEmbeddingMixin,
     Inference,
     ModelsProtocolPrivate,
@@ -133,24 +139,12 @@ class MetaReferenceInferenceImpl(
     async def load_model(self, model_id, llama_model) -> None:
         log.info(f"Loading model `{model_id}`")
 
-        if llama_model.model_family in {
-            ModelFamily.llama3,
-            ModelFamily.llama3_1,
-            ModelFamily.llama3_2,
-            ModelFamily.llama3_3,
-        }:
-            builder_fn = llama3_builder_fn
-        elif llama_model.model_family == ModelFamily.llama4:
-            builder_fn = llama4_builder_fn
-        else:
-            raise ValueError(f"Unsupported model family: {llama_model.model_family}")
-
         builder_params = [self.config, model_id, llama_model]
 
         if self.config.create_distributed_process_group:
             self.generator = LlamaModelParallelGenerator(
                 model_parallel_size=self.config.model_parallel_size or llama_model.pth_file_count,
-                builder_fn=builder_fn,
+                builder_fn=llama_builder_fn,
                 builder_params=builder_params,
                 formatter=(
                     Llama4ChatFormat(Llama4Tokenizer.get_instance())
@@ -160,11 +154,24 @@ class MetaReferenceInferenceImpl(
             )
             self.generator.start()
         else:
-            self.generator = builder_fn(*builder_params)
+            self.generator = llama_builder_fn(*builder_params)
 
         self.model_id = model_id
         self.llama_model = llama_model
 
+        log.info("Warming up...")
+        await self.completion(
+            model_id=model_id,
+            content="Hello, world!",
+            sampling_params=SamplingParams(max_tokens=10),
+        )
+        await self.chat_completion(
+            model_id=model_id,
+            messages=[UserMessage(content="Hi how are you?")],
+            sampling_params=SamplingParams(max_tokens=20),
+        )
+        log.info("Warmed up!")
+
     def check_model(self, request) -> None:
         if self.model_id is None or self.llama_model is None:
             raise RuntimeError(
@@ -202,7 +209,43 @@ class MetaReferenceInferenceImpl(
         if request.stream:
             return self._stream_completion(request)
         else:
-            return await self._nonstream_completion(request)
+            results = await self._nonstream_completion([request])
+            return results[0]
+
+    async def batch_completion(
+        self,
+        model_id: str,
+        content_batch: List[InterleavedContent],
+        sampling_params: Optional[SamplingParams] = None,
+        response_format: Optional[ResponseFormat] = None,
+        stream: Optional[bool] = False,
+        logprobs: Optional[LogProbConfig] = None,
+    ) -> BatchCompletionResponse:
+        if sampling_params is None:
+            sampling_params = SamplingParams()
+        if logprobs:
+            assert logprobs.top_k == 1, f"Unexpected top_k={logprobs.top_k}"
+
+        content_batch = [
+            augment_content_with_response_format_prompt(response_format, content) for content in content_batch
+        ]
+
+        request_batch = []
+        for content in content_batch:
+            request = CompletionRequest(
+                model=model_id,
+                content=content,
+                sampling_params=sampling_params,
+                response_format=response_format,
+                stream=stream,
+                logprobs=logprobs,
+            )
+            self.check_model(request)
+            request = await convert_request_to_raw(request)
+            request_batch.append(request)
+
+        results = await self._nonstream_completion(request_batch)
+        return BatchCompletionResponse(batch=results)
 
     async def _stream_completion(self, request: CompletionRequest) -> AsyncGenerator:
         tokenizer = self.generator.formatter.tokenizer
@@ -247,37 +290,54 @@ class MetaReferenceInferenceImpl(
             for x in impl():
                 yield x
 
-    async def _nonstream_completion(self, request: CompletionRequest) -> CompletionResponse:
+    async def _nonstream_completion(self, request_batch: List[CompletionRequest]) -> List[CompletionResponse]:
         tokenizer = self.generator.formatter.tokenizer
 
+        first_request = request_batch[0]
+
+        class ItemState(BaseModel):
+            tokens: List[int] = []
+            logprobs: List[TokenLogProbs] = []
+            stop_reason: StopReason | None = None
+            finished: bool = False
+
         def impl():
-            tokens = []
-            logprobs = []
-            stop_reason = None
+            states = [ItemState() for _ in request_batch]
 
-            for token_result in self.generator.completion(request):
-                tokens.append(token_result.token)
-                if token_result.token == tokenizer.eot_id:
-                    stop_reason = StopReason.end_of_turn
-                elif token_result.token == tokenizer.eom_id:
-                    stop_reason = StopReason.end_of_message
+            results = []
+            for token_results in self.generator.completion(request_batch):
+                for result in token_results:
+                    idx = result.batch_idx
+                    state = states[idx]
+                    if state.finished or result.ignore_token:
+                        continue
 
-                if request.logprobs:
-                    assert len(token_result.logprobs) == 1
+                    state.finished = result.finished
+                    if first_request.logprobs:
+                        state.logprobs.append(TokenLogProbs(logprobs_by_token={result.text: result.logprobs[0]}))
 
-                    logprobs.append(TokenLogProbs(logprobs_by_token={token_result.text: token_result.logprobs[0]}))
+                    state.tokens.append(result.token)
+                    if result.token == tokenizer.eot_id:
+                        state.stop_reason = StopReason.end_of_turn
+                    elif result.token == tokenizer.eom_id:
+                        state.stop_reason = StopReason.end_of_message
 
-            if stop_reason is None:
-                stop_reason = StopReason.out_of_tokens
+            for state in states:
+                if state.stop_reason is None:
+                    state.stop_reason = StopReason.out_of_tokens
 
-            if tokens[-1] in self.generator.formatter.tokenizer.stop_tokens:
-                tokens = tokens[:-1]
-            content = self.generator.formatter.tokenizer.decode(tokens)
-            return CompletionResponse(
-                content=content,
-                stop_reason=stop_reason,
-                logprobs=logprobs if request.logprobs else None,
-            )
+                if state.tokens[-1] in self.generator.formatter.tokenizer.stop_tokens:
+                    state.tokens = state.tokens[:-1]
+                content = self.generator.formatter.tokenizer.decode(state.tokens)
+                results.append(
+                    CompletionResponse(
+                        content=content,
+                        stop_reason=state.stop_reason,
+                        logprobs=state.logprobs if first_request.logprobs else None,
+                    )
+                )
+
+            return results
 
         if self.config.create_distributed_process_group:
             async with SEMAPHORE:
@@ -312,7 +372,7 @@ class MetaReferenceInferenceImpl(
             response_format=response_format,
             stream=stream,
             logprobs=logprobs,
-            tool_config=tool_config,
+            tool_config=tool_config or ToolConfig(),
         )
         self.check_model(request)
 
@@ -328,44 +388,110 @@ class MetaReferenceInferenceImpl(
         if request.stream:
             return self._stream_chat_completion(request)
         else:
-            return await self._nonstream_chat_completion(request)
+            results = await self._nonstream_chat_completion([request])
+            return results[0]
 
-    async def _nonstream_chat_completion(self, request: ChatCompletionRequest) -> ChatCompletionResponse:
+    async def batch_chat_completion(
+        self,
+        model_id: str,
+        messages_batch: List[List[Message]],
+        sampling_params: Optional[SamplingParams] = None,
+        response_format: Optional[ResponseFormat] = None,
+        tools: Optional[List[ToolDefinition]] = None,
+        stream: Optional[bool] = False,
+        logprobs: Optional[LogProbConfig] = None,
+        tool_config: Optional[ToolConfig] = None,
+    ) -> BatchChatCompletionResponse:
+        if sampling_params is None:
+            sampling_params = SamplingParams()
+        if logprobs:
+            assert logprobs.top_k == 1, f"Unexpected top_k={logprobs.top_k}"
+
+        # wrapper request to make it easier to pass around (internal only, not exposed to API)
+        request_batch = []
+        for messages in messages_batch:
+            request = ChatCompletionRequest(
+                model=model_id,
+                messages=messages,
+                sampling_params=sampling_params,
+                tools=tools or [],
+                response_format=response_format,
+                logprobs=logprobs,
+                tool_config=tool_config or ToolConfig(),
+            )
+            self.check_model(request)
+
+            # augment and rewrite messages depending on the model
+            request.messages = chat_completion_request_to_messages(request, self.llama_model.core_model_id.value)
+            # download media and convert to raw content so we can send it to the model
+            request = await convert_request_to_raw(request)
+            request_batch.append(request)
+
+        if self.config.create_distributed_process_group:
+            if SEMAPHORE.locked():
+                raise RuntimeError("Only one concurrent request is supported")
+
+        results = await self._nonstream_chat_completion(request_batch)
+        return BatchChatCompletionResponse(batch=results)
+
+    async def _nonstream_chat_completion(
+        self, request_batch: List[ChatCompletionRequest]
+    ) -> List[ChatCompletionResponse]:
         tokenizer = self.generator.formatter.tokenizer
 
+        first_request = request_batch[0]
+
+        class ItemState(BaseModel):
+            tokens: List[int] = []
+            logprobs: List[TokenLogProbs] = []
+            stop_reason: StopReason | None = None
+            finished: bool = False
+
         def impl():
-            tokens = []
-            logprobs = []
-            stop_reason = None
+            states = [ItemState() for _ in request_batch]
 
-            for token_result in self.generator.chat_completion(request):
-                if os.environ.get("LLAMA_MODELS_DEBUG", "0") == "1":
-                    cprint(token_result.text, "cyan", end="")
+            for token_results in self.generator.chat_completion(request_batch):
+                first = token_results[0]
+                if not first.finished and not first.ignore_token:
+                    if os.environ.get("LLAMA_MODELS_DEBUG", "0") in ("1", "2"):
+                        cprint(first.text, "cyan", end="")
+                    if os.environ.get("LLAMA_MODELS_DEBUG", "0") == "2":
+                        cprint(f"<{first.token}>", "magenta", end="")
 
-                tokens.append(token_result.token)
+                for result in token_results:
+                    idx = result.batch_idx
+                    state = states[idx]
+                    if state.finished or result.ignore_token:
+                        continue
 
-                if token_result.token == tokenizer.eot_id:
-                    stop_reason = StopReason.end_of_turn
-                elif token_result.token == tokenizer.eom_id:
-                    stop_reason = StopReason.end_of_message
+                    state.finished = result.finished
+                    if first_request.logprobs:
+                        state.logprobs.append(TokenLogProbs(logprobs_by_token={result.text: result.logprobs[0]}))
 
-                if request.logprobs:
-                    assert len(token_result.logprobs) == 1
+                    state.tokens.append(result.token)
+                    if result.token == tokenizer.eot_id:
+                        state.stop_reason = StopReason.end_of_turn
+                    elif result.token == tokenizer.eom_id:
+                        state.stop_reason = StopReason.end_of_message
 
-                    logprobs.append(TokenLogProbs(logprobs_by_token={token_result.text: token_result.logprobs[0]}))
+            results = []
+            for state in states:
+                if state.stop_reason is None:
+                    state.stop_reason = StopReason.out_of_tokens
 
-            if stop_reason is None:
-                stop_reason = StopReason.out_of_tokens
+                raw_message = self.generator.formatter.decode_assistant_message(state.tokens, state.stop_reason)
+                results.append(
+                    ChatCompletionResponse(
+                        completion_message=CompletionMessage(
+                            content=raw_message.content,
+                            stop_reason=raw_message.stop_reason,
+                            tool_calls=raw_message.tool_calls,
+                        ),
+                        logprobs=state.logprobs if first_request.logprobs else None,
+                    )
+                )
 
-            raw_message = self.generator.formatter.decode_assistant_message(tokens, stop_reason)
-            return ChatCompletionResponse(
-                completion_message=CompletionMessage(
-                    content=raw_message.content,
-                    stop_reason=raw_message.stop_reason,
-                    tool_calls=raw_message.tool_calls,
-                ),
-                logprobs=logprobs if request.logprobs else None,
-            )
+            return results
 
         if self.config.create_distributed_process_group:
             async with SEMAPHORE:
@@ -392,6 +518,22 @@ class MetaReferenceInferenceImpl(
             for token_result in self.generator.chat_completion(request):
                 if os.environ.get("LLAMA_MODELS_DEBUG", "0") == "1":
                     cprint(token_result.text, "cyan", end="")
+                if os.environ.get("LLAMA_MODELS_DEBUG", "0") == "2":
+                    cprint(f"<{token_result.token}>", "magenta", end="")
+
+                if token_result.token == tokenizer.eot_id:
+                    stop_reason = StopReason.end_of_turn
+                    text = ""
+                elif token_result.token == tokenizer.eom_id:
+                    stop_reason = StopReason.end_of_message
+                    text = ""
+                else:
+                    text = token_result.text
+
+                if request.logprobs:
+                    assert len(token_result.logprobs) == 1
+
+                    logprobs.append(TokenLogProbs(logprobs_by_token={token_result.text: token_result.logprobs[0]}))
 
                 tokens.append(token_result.token)
 

llama_stack/providers/inline/inference/meta_reference/model_parallel.py

@@ -6,7 +6,7 @@
 
 from copy import deepcopy
 from functools import partial
-from typing import Any, Callable, Generator
+from typing import Any, Callable, Generator, List
 
 from llama_stack.models.llama.llama3.chat_format import ChatFormat as Llama3ChatFormat
 from llama_stack.models.llama.llama4.chat_format import ChatFormat as Llama4ChatFormat
@@ -23,13 +23,13 @@ class ModelRunner:
         self.llama = llama
 
     # the `task` object is the same that is sent to `ModelParallelProcessGroup.run_inference()`
-    def __call__(self, req: Any):
-        if isinstance(req, ChatCompletionRequestWithRawContent):
-            return self.llama.chat_completion(req)
-        elif isinstance(req, CompletionRequestWithRawContent):
-            return self.llama.completion(req)
+    def __call__(self, task: Any):
+        if task[0] == "chat_completion":
+            return self.llama.chat_completion(task[1])
+        elif task[0] == "completion":
+            return self.llama.completion(task[1])
         else:
-            raise ValueError(f"Unexpected task type {type(req)}")
+            raise ValueError(f"Unexpected task type {task[0]}")
 
 
 def init_model_cb(
@@ -82,16 +82,16 @@ class LlamaModelParallelGenerator:
 
     def completion(
         self,
-        request: CompletionRequestWithRawContent,
+        request_batch: List[CompletionRequestWithRawContent],
     ) -> Generator:
-        req_obj = deepcopy(request)
-        gen = self.group.run_inference(req_obj)
+        req_obj = deepcopy(request_batch)
+        gen = self.group.run_inference(("completion", req_obj))
         yield from gen
 
     def chat_completion(
         self,
-        request: ChatCompletionRequestWithRawContent,
+        request_batch: List[ChatCompletionRequestWithRawContent],
     ) -> Generator:
-        req_obj = deepcopy(request)
-        gen = self.group.run_inference(req_obj)
+        req_obj = deepcopy(request_batch)
+        gen = self.group.run_inference(("chat_completion", req_obj))
         yield from gen

llama_stack/providers/inline/inference/meta_reference/parallel_utils.py

@@ -19,7 +19,7 @@ import tempfile
 import time
 import uuid
 from enum import Enum
-from typing import Callable, Generator, Literal, Optional, Union
+from typing import Callable, Generator, List, Literal, Optional, Tuple, Union
 
 import torch
 import zmq
@@ -69,12 +69,12 @@ class CancelSentinel(BaseModel):
 
 class TaskRequest(BaseModel):
     type: Literal[ProcessingMessageName.task_request] = ProcessingMessageName.task_request
-    task: Union[CompletionRequestWithRawContent, ChatCompletionRequestWithRawContent]
+    task: Tuple[str, List[CompletionRequestWithRawContent] | List[ChatCompletionRequestWithRawContent]]
 
 
 class TaskResponse(BaseModel):
     type: Literal[ProcessingMessageName.task_response] = ProcessingMessageName.task_response
-    result: GenerationResult
+    result: List[GenerationResult]
 
 
 class ExceptionResponse(BaseModel):
@@ -331,7 +331,7 @@ class ModelParallelProcessGroup:
 
     def run_inference(
         self,
-        req: Union[CompletionRequestWithRawContent, ChatCompletionRequestWithRawContent],
+        req: Tuple[str, List[CompletionRequestWithRawContent] | List[ChatCompletionRequestWithRawContent]],
     ) -> Generator:
         assert not self.running, "inference already running"
 

llama_stack/providers/inline/inference/sentence_transformers/sentence_transformers.py

@@ -10,6 +10,7 @@ from typing import AsyncGenerator, List, Optional, Union
 from llama_stack.apis.inference import (
     CompletionResponse,
     Inference,
+    InterleavedContent,
     LogProbConfig,
     Message,
     ResponseFormat,
@@ -23,6 +24,10 @@ from llama_stack.providers.datatypes import Model, ModelsProtocolPrivate
 from llama_stack.providers.utils.inference.embedding_mixin import (
     SentenceTransformerEmbeddingMixin,
 )
+from llama_stack.providers.utils.inference.openai_compat import (
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
+)
 
 from .config import SentenceTransformersInferenceConfig
 
@@ -30,6 +35,8 @@ log = logging.getLogger(__name__)
 
 
 class SentenceTransformersInferenceImpl(
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
     SentenceTransformerEmbeddingMixin,
     Inference,
     ModelsProtocolPrivate,
@@ -74,3 +81,25 @@ class SentenceTransformersInferenceImpl(
         tool_config: Optional[ToolConfig] = None,
     ) -> AsyncGenerator:
         raise ValueError("Sentence transformers don't support chat completion")
+
+    async def batch_completion(
+        self,
+        model_id: str,
+        content_batch: List[InterleavedContent],
+        sampling_params: Optional[SamplingParams] = None,
+        response_format: Optional[ResponseFormat] = None,
+        logprobs: Optional[LogProbConfig] = None,
+    ):
+        raise NotImplementedError("Batch completion is not supported for Sentence Transformers")
+
+    async def batch_chat_completion(
+        self,
+        model_id: str,
+        messages_batch: List[List[Message]],
+        sampling_params: Optional[SamplingParams] = None,
+        tools: Optional[List[ToolDefinition]] = None,
+        tool_config: Optional[ToolConfig] = None,
+        response_format: Optional[ResponseFormat] = None,
+        logprobs: Optional[LogProbConfig] = None,
+    ):
+        raise NotImplementedError("Batch chat completion is not supported for Sentence Transformers")

llama_stack/providers/inline/inference/vllm/vllm.py

@@ -66,8 +66,10 @@ from llama_stack.providers.utils.inference.model_registry import (
     ModelsProtocolPrivate,
 )
 from llama_stack.providers.utils.inference.openai_compat import (
+    OpenAIChatCompletionToLlamaStackMixin,
     OpenAICompatCompletionChoice,
     OpenAICompatCompletionResponse,
+    OpenAICompletionToLlamaStackMixin,
     get_stop_reason,
     process_chat_completion_stream_response,
 )
@@ -172,7 +174,12 @@ def _convert_sampling_params(
     return vllm_sampling_params
 
 
-class VLLMInferenceImpl(Inference, ModelsProtocolPrivate):
+class VLLMInferenceImpl(
+    Inference,
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
+    ModelsProtocolPrivate,
+):
     """
     vLLM-based inference model adapter for Llama Stack with support for multiple models.
 

llama_stack/providers/inline/post_training/torchtune/post_training.py

@@ -3,13 +3,14 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-from datetime import datetime, timezone
+from enum import Enum
 from typing import Any, Dict, Optional
 
 from llama_stack.apis.datasetio import DatasetIO
 from llama_stack.apis.datasets import Datasets
 from llama_stack.apis.post_training import (
     AlgorithmConfig,
+    Checkpoint,
     DPOAlignmentConfig,
     JobStatus,
     ListPostTrainingJobsResponse,
@@ -25,9 +26,19 @@ from llama_stack.providers.inline.post_training.torchtune.config import (
 from llama_stack.providers.inline.post_training.torchtune.recipes.lora_finetuning_single_device import (
     LoraFinetuningSingleDevice,
 )
+from llama_stack.providers.utils.scheduler import JobArtifact, Scheduler
+from llama_stack.providers.utils.scheduler import JobStatus as SchedulerJobStatus
 from llama_stack.schema_utils import webmethod
 
 
+class TrainingArtifactType(Enum):
+    CHECKPOINT = "checkpoint"
+    RESOURCES_STATS = "resources_stats"
+
+
+_JOB_TYPE_SUPERVISED_FINE_TUNE = "supervised-fine-tune"
+
+
 class TorchtunePostTrainingImpl:
     def __init__(
         self,
@@ -38,13 +49,27 @@ class TorchtunePostTrainingImpl:
         self.config = config
         self.datasetio_api = datasetio_api
         self.datasets_api = datasets
+        self._scheduler = Scheduler()
 
-        # TODO: assume sync job, will need jobs API for async scheduling
-        self.jobs = {}
-        self.checkpoints_dict = {}
+    async def shutdown(self) -> None:
+        await self._scheduler.shutdown()
 
-    async def shutdown(self):
-        pass
+    @staticmethod
+    def _checkpoint_to_artifact(checkpoint: Checkpoint) -> JobArtifact:
+        return JobArtifact(
+            type=TrainingArtifactType.CHECKPOINT.value,
+            name=checkpoint.identifier,
+            uri=checkpoint.path,
+            metadata=dict(checkpoint),
+        )
+
+    @staticmethod
+    def _resources_stats_to_artifact(resources_stats: Dict[str, Any]) -> JobArtifact:
+        return JobArtifact(
+            type=TrainingArtifactType.RESOURCES_STATS.value,
+            name=TrainingArtifactType.RESOURCES_STATS.value,
+            metadata=resources_stats,
+        )
 
     async def supervised_fine_tune(
         self,
@@ -56,20 +81,11 @@ class TorchtunePostTrainingImpl:
         checkpoint_dir: Optional[str],
         algorithm_config: Optional[AlgorithmConfig],
     ) -> PostTrainingJob:
-        if job_uuid in self.jobs:
-            raise ValueError(f"Job {job_uuid} already exists")
-
-        post_training_job = PostTrainingJob(job_uuid=job_uuid)
-
-        job_status_response = PostTrainingJobStatusResponse(
-            job_uuid=job_uuid,
-            status=JobStatus.scheduled,
-            scheduled_at=datetime.now(timezone.utc),
-        )
-        self.jobs[job_uuid] = job_status_response
-
         if isinstance(algorithm_config, LoraFinetuningConfig):
-            try:
+
+            async def handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb):
+                on_log_message_cb("Starting Lora finetuning")
+
                 recipe = LoraFinetuningSingleDevice(
                     self.config,
                     job_uuid,
@@ -82,26 +98,22 @@ class TorchtunePostTrainingImpl:
                     self.datasetio_api,
                     self.datasets_api,
                 )
-
-                job_status_response.status = JobStatus.in_progress
-                job_status_response.started_at = datetime.now(timezone.utc)
-
                 await recipe.setup()
+
                 resources_allocated, checkpoints = await recipe.train()
 
-                self.checkpoints_dict[job_uuid] = checkpoints
-                job_status_response.resources_allocated = resources_allocated
-                job_status_response.checkpoints = checkpoints
-                job_status_response.status = JobStatus.completed
-                job_status_response.completed_at = datetime.now(timezone.utc)
+                on_artifact_collected_cb(self._resources_stats_to_artifact(resources_allocated))
+                for checkpoint in checkpoints:
+                    artifact = self._checkpoint_to_artifact(checkpoint)
+                    on_artifact_collected_cb(artifact)
 
-            except Exception:
-                job_status_response.status = JobStatus.failed
-                raise
+                on_status_change_cb(SchedulerJobStatus.completed)
+                on_log_message_cb("Lora finetuning completed")
         else:
             raise NotImplementedError()
 
-        return post_training_job
+        job_uuid = self._scheduler.schedule(_JOB_TYPE_SUPERVISED_FINE_TUNE, job_uuid, handler)
+        return PostTrainingJob(job_uuid=job_uuid)
 
     async def preference_optimize(
         self,
@@ -114,19 +126,55 @@ class TorchtunePostTrainingImpl:
     ) -> PostTrainingJob: ...
 
     async def get_training_jobs(self) -> ListPostTrainingJobsResponse:
-        return ListPostTrainingJobsResponse(data=[PostTrainingJob(job_uuid=uuid_) for uuid_ in self.jobs])
+        return ListPostTrainingJobsResponse(
+            data=[PostTrainingJob(job_uuid=job.id) for job in self._scheduler.get_jobs()]
+        )
+
+    @staticmethod
+    def _get_artifacts_metadata_by_type(job, artifact_type):
+        return [artifact.metadata for artifact in job.artifacts if artifact.type == artifact_type]
+
+    @classmethod
+    def _get_checkpoints(cls, job):
+        return cls._get_artifacts_metadata_by_type(job, TrainingArtifactType.CHECKPOINT.value)
+
+    @classmethod
+    def _get_resources_allocated(cls, job):
+        data = cls._get_artifacts_metadata_by_type(job, TrainingArtifactType.RESOURCES_STATS.value)
+        return data[0] if data else None
 
     @webmethod(route="/post-training/job/status")
     async def get_training_job_status(self, job_uuid: str) -> Optional[PostTrainingJobStatusResponse]:
-        return self.jobs.get(job_uuid, None)
+        job = self._scheduler.get_job(job_uuid)
+
+        match job.status:
+            # TODO: Add support for other statuses to API
+            case SchedulerJobStatus.new | SchedulerJobStatus.scheduled:
+                status = JobStatus.scheduled
+            case SchedulerJobStatus.running:
+                status = JobStatus.in_progress
+            case SchedulerJobStatus.completed:
+                status = JobStatus.completed
+            case SchedulerJobStatus.failed:
+                status = JobStatus.failed
+            case _:
+                raise NotImplementedError()
+
+        return PostTrainingJobStatusResponse(
+            job_uuid=job_uuid,
+            status=status,
+            scheduled_at=job.scheduled_at,
+            started_at=job.started_at,
+            completed_at=job.completed_at,
+            checkpoints=self._get_checkpoints(job),
+            resources_allocated=self._get_resources_allocated(job),
+        )
 
     @webmethod(route="/post-training/job/cancel")
     async def cancel_training_job(self, job_uuid: str) -> None:
-        raise NotImplementedError("Job cancel is not implemented yet")
+        self._scheduler.cancel(job_uuid)
 
     @webmethod(route="/post-training/job/artifacts")
     async def get_training_job_artifacts(self, job_uuid: str) -> Optional[PostTrainingJobArtifactsResponse]:
-        if job_uuid in self.checkpoints_dict:
-            checkpoints = self.checkpoints_dict.get(job_uuid, [])
-            return PostTrainingJobArtifactsResponse(job_uuid=job_uuid, checkpoints=checkpoints)
-        return None
+        job = self._scheduler.get_job(job_uuid)
+        return PostTrainingJobArtifactsResponse(job_uuid=job_uuid, checkpoints=self._get_checkpoints(job))

llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py

@@ -38,6 +38,8 @@ from llama_stack.apis.datasetio import DatasetIO
 from llama_stack.apis.datasets import Datasets
 from llama_stack.apis.post_training import (
     Checkpoint,
+    DataConfig,
+    EfficiencyConfig,
     LoraFinetuningConfig,
     OptimizerConfig,
     QATFinetuningConfig,
@@ -89,6 +91,10 @@ class LoraFinetuningSingleDevice:
         datasetio_api: DatasetIO,
         datasets_api: Datasets,
     ) -> None:
+        assert isinstance(training_config.data_config, DataConfig), "DataConfig must be initialized"
+
+        assert isinstance(training_config.efficiency_config, EfficiencyConfig), "EfficiencyConfig must be initialized"
+
         self.job_uuid = job_uuid
         self.training_config = training_config
         if not isinstance(algorithm_config, LoraFinetuningConfig):
@@ -188,6 +194,7 @@ class LoraFinetuningSingleDevice:
         self._tokenizer = await self._setup_tokenizer()
         log.info("Tokenizer is initialized.")
 
+        assert isinstance(self.training_config.optimizer_config, OptimizerConfig), "OptimizerConfig must be initialized"
         self._optimizer = await self._setup_optimizer(optimizer_config=self.training_config.optimizer_config)
         log.info("Optimizer is initialized.")
 
@@ -195,6 +202,8 @@ class LoraFinetuningSingleDevice:
         self._model.set_num_output_chunks(self._loss_fn.num_output_chunks)
         log.info("Loss is initialized.")
 
+        assert isinstance(self.training_config.data_config, DataConfig), "DataConfig must be initialized"
+
         self._training_sampler, self._training_dataloader = await self._setup_data(
             dataset_id=self.training_config.data_config.dataset_id,
             tokenizer=self._tokenizer,
@@ -452,6 +461,7 @@ class LoraFinetuningSingleDevice:
         """
         The core training loop.
         """
+        assert isinstance(self.training_config.data_config, DataConfig), "DataConfig must be initialized"
         # Initialize tokens count and running loss (for grad accumulation)
         t0 = time.perf_counter()
         running_loss: float = 0.0

llama_stack/providers/inline/safety/llama_guard/llama_guard.py

@@ -10,7 +10,6 @@ from typing import Any, Dict, List, Optional
 
 from llama_stack.apis.common.content_types import ImageContentItem, TextContentItem
 from llama_stack.apis.inference import (
-    ChatCompletionResponseEventType,
     Inference,
     Message,
     UserMessage,
@@ -239,16 +238,12 @@ class LlamaGuardShield:
             shield_input_message = self.build_text_shield_input(messages)
 
         # TODO: llama-stack inference protocol has issues with non-streaming inference code
-        content = ""
-        async for chunk in await self.inference_api.chat_completion(
+        response = await self.inference_api.chat_completion(
             model_id=self.model,
             messages=[shield_input_message],
-            stream=True,
-        ):
-            event = chunk.event
-            if event.event_type == ChatCompletionResponseEventType.progress and event.delta.type == "text":
-                content += event.delta.text
-
+            stream=False,
+        )
+        content = response.completion_message.content
         content = content.strip()
         return self.get_shield_response(content)
 

llama_stack/providers/registry/inference.py

@@ -24,7 +24,7 @@ META_REFERENCE_DEPS = [
     "zmq",
     "lm-format-enforcer",
     "sentence-transformers",
-    "torchao==0.5.0",
+    "torchao==0.8.0",
     "fbgemm-gpu-genai==1.1.2",
 ]
 

llama_stack/providers/remote/inference/bedrock/bedrock.py

@@ -36,8 +36,10 @@ from llama_stack.providers.utils.inference.model_registry import (
     ModelRegistryHelper,
 )
 from llama_stack.providers.utils.inference.openai_compat import (
+    OpenAIChatCompletionToLlamaStackMixin,
     OpenAICompatCompletionChoice,
     OpenAICompatCompletionResponse,
+    OpenAICompletionToLlamaStackMixin,
     get_sampling_strategy_options,
     process_chat_completion_response,
     process_chat_completion_stream_response,
@@ -51,7 +53,12 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
 from .models import MODEL_ENTRIES
 
 
-class BedrockInferenceAdapter(ModelRegistryHelper, Inference):
+class BedrockInferenceAdapter(
+    ModelRegistryHelper,
+    Inference,
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
+):
     def __init__(self, config: BedrockConfig) -> None:
         ModelRegistryHelper.__init__(self, MODEL_ENTRIES)
         self._config = config

llama_stack/providers/remote/inference/cerebras/cerebras.py

@@ -34,6 +34,8 @@ from llama_stack.providers.utils.inference.model_registry import (
     ModelRegistryHelper,
 )
 from llama_stack.providers.utils.inference.openai_compat import (
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
     get_sampling_options,
     process_chat_completion_response,
     process_chat_completion_stream_response,
@@ -49,7 +51,12 @@ from .config import CerebrasImplConfig
 from .models import MODEL_ENTRIES
 
 
-class CerebrasInferenceAdapter(ModelRegistryHelper, Inference):
+class CerebrasInferenceAdapter(
+    ModelRegistryHelper,
+    Inference,
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
+):
     def __init__(self, config: CerebrasImplConfig) -> None:
         ModelRegistryHelper.__init__(
             self,

llama_stack/providers/remote/inference/databricks/databricks.py

@@ -34,6 +34,8 @@ from llama_stack.providers.utils.inference.model_registry import (
     build_hf_repo_model_entry,
 )
 from llama_stack.providers.utils.inference.openai_compat import (
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
     get_sampling_options,
     process_chat_completion_response,
     process_chat_completion_stream_response,
@@ -56,7 +58,12 @@ model_entries = [
 ]
 
 
-class DatabricksInferenceAdapter(ModelRegistryHelper, Inference):
+class DatabricksInferenceAdapter(
+    ModelRegistryHelper,
+    Inference,
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
+):
     def __init__(self, config: DatabricksImplConfig) -> None:
         ModelRegistryHelper.__init__(self, model_entries=model_entries)
         self.config = config

llama_stack/providers/remote/inference/fireworks/fireworks.py

@@ -4,9 +4,10 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import AsyncGenerator, List, Optional, Union
+from typing import Any, AsyncGenerator, AsyncIterator, Dict, List, Optional, Union
 
 from fireworks.client import Fireworks
+from openai import AsyncOpenAI
 
 from llama_stack.apis.common.content_types import (
     InterleavedContent,
@@ -31,14 +32,23 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
+from llama_stack.apis.inference.inference import (
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAICompletion,
+    OpenAIMessageParam,
+    OpenAIResponseFormatParam,
+)
 from llama_stack.distribution.request_headers import NeedsRequestProviderData
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.model_registry import (
     ModelRegistryHelper,
 )
 from llama_stack.providers.utils.inference.openai_compat import (
+    OpenAIChatCompletionToLlamaStackMixin,
     convert_message_to_openai_dict,
     get_sampling_options,
+    prepare_openai_completion_params,
     process_chat_completion_response,
     process_chat_completion_stream_response,
     process_completion_response,
@@ -81,10 +91,16 @@ class FireworksInferenceAdapter(ModelRegistryHelper, Inference, NeedsRequestProv
                 )
             return provider_data.fireworks_api_key
 
+    def _get_base_url(self) -> str:
+        return "https://api.fireworks.ai/inference/v1"
+
     def _get_client(self) -> Fireworks:
         fireworks_api_key = self._get_api_key()
         return Fireworks(api_key=fireworks_api_key)
 
+    def _get_openai_client(self) -> AsyncOpenAI:
+        return AsyncOpenAI(base_url=self._get_base_url(), api_key=self._get_api_key())
+
     async def completion(
         self,
         model_id: str,
@@ -268,3 +284,114 @@ class FireworksInferenceAdapter(ModelRegistryHelper, Inference, NeedsRequestProv
 
         embeddings = [data.embedding for data in response.data]
         return EmbeddingsResponse(embeddings=embeddings)
+
+    async def openai_completion(
+        self,
+        model: str,
+        prompt: Union[str, List[str], List[int], List[List[int]]],
+        best_of: Optional[int] = None,
+        echo: Optional[bool] = None,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        guided_choice: Optional[List[str]] = None,
+        prompt_logprobs: Optional[int] = None,
+    ) -> OpenAICompletion:
+        model_obj = await self.model_store.get_model(model)
+
+        # Fireworks always prepends with BOS
+        if isinstance(prompt, str) and prompt.startswith("<|begin_of_text|>"):
+            prompt = prompt[len("<|begin_of_text|>") :]
+
+        params = await prepare_openai_completion_params(
+            model=model_obj.provider_resource_id,
+            prompt=prompt,
+            best_of=best_of,
+            echo=echo,
+            frequency_penalty=frequency_penalty,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_tokens=max_tokens,
+            n=n,
+            presence_penalty=presence_penalty,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            top_p=top_p,
+            user=user,
+        )
+
+        return await self._get_openai_client().completions.create(**params)
+
+    async def openai_chat_completion(
+        self,
+        model: str,
+        messages: List[OpenAIMessageParam],
+        frequency_penalty: Optional[float] = None,
+        function_call: Optional[Union[str, Dict[str, Any]]] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        parallel_tool_calls: Optional[bool] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[OpenAIResponseFormatParam] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+    ) -> Union[OpenAIChatCompletion, AsyncIterator[OpenAIChatCompletionChunk]]:
+        model_obj = await self.model_store.get_model(model)
+        params = await prepare_openai_completion_params(
+            messages=messages,
+            frequency_penalty=frequency_penalty,
+            function_call=function_call,
+            functions=functions,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_completion_tokens=max_completion_tokens,
+            max_tokens=max_tokens,
+            n=n,
+            parallel_tool_calls=parallel_tool_calls,
+            presence_penalty=presence_penalty,
+            response_format=response_format,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            tool_choice=tool_choice,
+            tools=tools,
+            top_logprobs=top_logprobs,
+            top_p=top_p,
+            user=user,
+        )
+
+        # Divert Llama Models through Llama Stack inference APIs because
+        # Fireworks chat completions OpenAI-compatible API does not support
+        # tool calls properly.
+        llama_model = self.get_llama_model(model_obj.provider_resource_id)
+        if llama_model:
+            return await OpenAIChatCompletionToLlamaStackMixin.openai_chat_completion(self, model=model, **params)
+
+        return await self._get_openai_client().chat.completions.create(model=model_obj.provider_resource_id, **params)

llama_stack/providers/remote/inference/groq/groq.py

@@ -4,8 +4,24 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+from typing import Any, AsyncIterator, Dict, List, Optional, Union
+
+from openai import AsyncOpenAI
+
+from llama_stack.apis.inference.inference import (
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAIChoiceDelta,
+    OpenAIChunkChoice,
+    OpenAIMessageParam,
+    OpenAIResponseFormatParam,
+    OpenAISystemMessageParam,
+)
 from llama_stack.providers.remote.inference.groq.config import GroqConfig
 from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
+from llama_stack.providers.utils.inference.openai_compat import (
+    prepare_openai_completion_params,
+)
 
 from .models import MODEL_ENTRIES
 
@@ -21,9 +37,129 @@ class GroqInferenceAdapter(LiteLLMOpenAIMixin):
             provider_data_api_key_field="groq_api_key",
         )
         self.config = config
+        self._openai_client = None
 
     async def initialize(self):
         await super().initialize()
 
     async def shutdown(self):
         await super().shutdown()
+        if self._openai_client:
+            await self._openai_client.close()
+            self._openai_client = None
+
+    def _get_openai_client(self) -> AsyncOpenAI:
+        if not self._openai_client:
+            self._openai_client = AsyncOpenAI(
+                base_url=f"{self.config.url}/openai/v1",
+                api_key=self.config.api_key,
+            )
+        return self._openai_client
+
+    async def openai_chat_completion(
+        self,
+        model: str,
+        messages: List[OpenAIMessageParam],
+        frequency_penalty: Optional[float] = None,
+        function_call: Optional[Union[str, Dict[str, Any]]] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        parallel_tool_calls: Optional[bool] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[OpenAIResponseFormatParam] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+    ) -> Union[OpenAIChatCompletion, AsyncIterator[OpenAIChatCompletionChunk]]:
+        model_obj = await self.model_store.get_model(model)
+
+        # Groq does not support json_schema response format, so we need to convert it to json_object
+        if response_format and response_format.type == "json_schema":
+            response_format.type = "json_object"
+            schema = response_format.json_schema.get("schema", {})
+            response_format.json_schema = None
+            json_instructions = f"\nYour response should be a JSON object that matches the following schema: {schema}"
+            if messages and messages[0].role == "system":
+                messages[0].content = messages[0].content + json_instructions
+            else:
+                messages.insert(0, OpenAISystemMessageParam(content=json_instructions))
+
+        # Groq returns a 400 error if tools are provided but none are called
+        # So, set tool_choice to "required" to attempt to force a call
+        if tools and (not tool_choice or tool_choice == "auto"):
+            tool_choice = "required"
+
+        params = await prepare_openai_completion_params(
+            model=model_obj.provider_resource_id.replace("groq/", ""),
+            messages=messages,
+            frequency_penalty=frequency_penalty,
+            function_call=function_call,
+            functions=functions,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_completion_tokens=max_completion_tokens,
+            max_tokens=max_tokens,
+            n=n,
+            parallel_tool_calls=parallel_tool_calls,
+            presence_penalty=presence_penalty,
+            response_format=response_format,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            tool_choice=tool_choice,
+            tools=tools,
+            top_logprobs=top_logprobs,
+            top_p=top_p,
+            user=user,
+        )
+
+        # Groq does not support streaming requests that set response_format
+        fake_stream = False
+        if stream and response_format:
+            params["stream"] = False
+            fake_stream = True
+
+        response = await self._get_openai_client().chat.completions.create(**params)
+
+        if fake_stream:
+            chunk_choices = []
+            for choice in response.choices:
+                delta = OpenAIChoiceDelta(
+                    content=choice.message.content,
+                    role=choice.message.role,
+                    tool_calls=choice.message.tool_calls,
+                )
+                chunk_choice = OpenAIChunkChoice(
+                    delta=delta,
+                    finish_reason=choice.finish_reason,
+                    index=choice.index,
+                    logprobs=None,
+                )
+                chunk_choices.append(chunk_choice)
+            chunk = OpenAIChatCompletionChunk(
+                id=response.id,
+                choices=chunk_choices,
+                object="chat.completion.chunk",
+                created=response.created,
+                model=response.model,
+            )
+
+            async def _fake_stream_generator():
+                yield chunk
+
+            return _fake_stream_generator()
+        else:
+            return response

llama_stack/providers/remote/inference/groq/models.py

@@ -39,8 +39,16 @@ MODEL_ENTRIES = [
         "groq/llama-4-scout-17b-16e-instruct",
         CoreModelId.llama4_scout_17b_16e_instruct.value,
     ),
+    build_hf_repo_model_entry(
+        "groq/meta-llama/llama-4-scout-17b-16e-instruct",
+        CoreModelId.llama4_scout_17b_16e_instruct.value,
+    ),
     build_hf_repo_model_entry(
         "groq/llama-4-maverick-17b-128e-instruct",
         CoreModelId.llama4_maverick_17b_128e_instruct.value,
     ),
+    build_hf_repo_model_entry(
+        "groq/meta-llama/llama-4-maverick-17b-128e-instruct",
+        CoreModelId.llama4_maverick_17b_128e_instruct.value,
+    ),
 ]

llama_stack/providers/remote/inference/nvidia/nvidia.py

@@ -7,7 +7,7 @@
 import logging
 import warnings
 from functools import lru_cache
-from typing import AsyncIterator, List, Optional, Union
+from typing import Any, AsyncIterator, Dict, List, Optional, Union
 
 from openai import APIConnectionError, AsyncOpenAI, BadRequestError
 
@@ -35,6 +35,13 @@ from llama_stack.apis.inference import (
     ToolConfig,
     ToolDefinition,
 )
+from llama_stack.apis.inference.inference import (
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAICompletion,
+    OpenAIMessageParam,
+    OpenAIResponseFormatParam,
+)
 from llama_stack.models.llama.datatypes import ToolPromptFormat
 from llama_stack.providers.utils.inference.model_registry import (
     ModelRegistryHelper,
@@ -42,6 +49,7 @@ from llama_stack.providers.utils.inference.model_registry import (
 from llama_stack.providers.utils.inference.openai_compat import (
     convert_openai_chat_completion_choice,
     convert_openai_chat_completion_stream,
+    prepare_openai_completion_params,
 )
 from llama_stack.providers.utils.inference.prompt_adapter import content_has_media
 
@@ -263,3 +271,111 @@ class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
         else:
             # we pass n=1 to get only one completion
             return convert_openai_chat_completion_choice(response.choices[0])
+
+    async def openai_completion(
+        self,
+        model: str,
+        prompt: Union[str, List[str], List[int], List[List[int]]],
+        best_of: Optional[int] = None,
+        echo: Optional[bool] = None,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        guided_choice: Optional[List[str]] = None,
+        prompt_logprobs: Optional[int] = None,
+    ) -> OpenAICompletion:
+        provider_model_id = self.get_provider_model_id(model)
+
+        params = await prepare_openai_completion_params(
+            model=provider_model_id,
+            prompt=prompt,
+            best_of=best_of,
+            echo=echo,
+            frequency_penalty=frequency_penalty,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_tokens=max_tokens,
+            n=n,
+            presence_penalty=presence_penalty,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            top_p=top_p,
+            user=user,
+        )
+
+        try:
+            return await self._get_client(provider_model_id).completions.create(**params)
+        except APIConnectionError as e:
+            raise ConnectionError(f"Failed to connect to NVIDIA NIM at {self._config.url}: {e}") from e
+
+    async def openai_chat_completion(
+        self,
+        model: str,
+        messages: List[OpenAIMessageParam],
+        frequency_penalty: Optional[float] = None,
+        function_call: Optional[Union[str, Dict[str, Any]]] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        parallel_tool_calls: Optional[bool] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[OpenAIResponseFormatParam] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+    ) -> Union[OpenAIChatCompletion, AsyncIterator[OpenAIChatCompletionChunk]]:
+        provider_model_id = self.get_provider_model_id(model)
+
+        params = await prepare_openai_completion_params(
+            model=provider_model_id,
+            messages=messages,
+            frequency_penalty=frequency_penalty,
+            function_call=function_call,
+            functions=functions,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_completion_tokens=max_completion_tokens,
+            max_tokens=max_tokens,
+            n=n,
+            parallel_tool_calls=parallel_tool_calls,
+            presence_penalty=presence_penalty,
+            response_format=response_format,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            tool_choice=tool_choice,
+            tools=tools,
+            top_logprobs=top_logprobs,
+            top_p=top_p,
+            user=user,
+        )
+
+        try:
+            return await self._get_client(provider_model_id).chat.completions.create(**params)
+        except APIConnectionError as e:
+            raise ConnectionError(f"Failed to connect to NVIDIA NIM at {self._config.url}: {e}") from e

llama_stack/providers/remote/inference/ollama/ollama.py

@@ -5,10 +5,11 @@
 # the root directory of this source tree.
 
 
-from typing import Any, AsyncGenerator, List, Optional, Union
+from typing import Any, AsyncGenerator, AsyncIterator, Dict, List, Optional, Union
 
 import httpx
 from ollama import AsyncClient
+from openai import AsyncOpenAI
 
 from llama_stack.apis.common.content_types import (
     ImageContentItem,
@@ -38,9 +39,20 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
+from llama_stack.apis.inference.inference import (
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAICompletion,
+    OpenAIMessageParam,
+    OpenAIResponseFormatParam,
+)
 from llama_stack.apis.models import Model, ModelType
 from llama_stack.log import get_logger
-from llama_stack.providers.datatypes import ModelsProtocolPrivate
+from llama_stack.providers.datatypes import (
+    HealthResponse,
+    HealthStatus,
+    ModelsProtocolPrivate,
+)
 from llama_stack.providers.utils.inference.model_registry import (
     ModelRegistryHelper,
 )
@@ -67,7 +79,10 @@ from .models import model_entries
 logger = get_logger(name=__name__, category="inference")
 
 
-class OllamaInferenceAdapter(Inference, ModelsProtocolPrivate):
+class OllamaInferenceAdapter(
+    Inference,
+    ModelsProtocolPrivate,
+):
     def __init__(self, url: str) -> None:
         self.register_helper = ModelRegistryHelper(model_entries)
         self.url = url
@@ -76,10 +91,25 @@ class OllamaInferenceAdapter(Inference, ModelsProtocolPrivate):
     def client(self) -> AsyncClient:
         return AsyncClient(host=self.url)
 
+    @property
+    def openai_client(self) -> AsyncOpenAI:
+        return AsyncOpenAI(base_url=f"{self.url}/v1", api_key="ollama")
+
     async def initialize(self) -> None:
         logger.info(f"checking connectivity to Ollama at `{self.url}`...")
+        await self.health()
+
+    async def health(self) -> HealthResponse:
+        """
+        Performs a health check by verifying connectivity to the Ollama server.
+        This method is used by initialize() and the Provider API to verify that the service is running
+        correctly.
+        Returns:
+            HealthResponse: A dictionary containing the health status.
+        """
         try:
             await self.client.ps()
+            return HealthResponse(status=HealthStatus.OK)
         except httpx.ConnectError as e:
             raise RuntimeError(
                 "Ollama Server is not running, start it using `ollama serve` in a separate terminal"
@@ -313,12 +343,149 @@ class OllamaInferenceAdapter(Inference, ModelsProtocolPrivate):
         response = await self.client.list()
         available_models = [m["model"] for m in response["models"]]
         if model.provider_resource_id not in available_models:
+            available_models_latest = [m["model"].split(":latest")[0] for m in response["models"]]
+            if model.provider_resource_id in available_models_latest:
+                logger.warning(
+                    f"Imprecise provider resource id was used but 'latest' is available in Ollama - using '{model.provider_resource_id}:latest'"
+                )
+                return model
             raise ValueError(
                 f"Model '{model.provider_resource_id}' is not available in Ollama. Available models: {', '.join(available_models)}"
             )
 
         return model
 
+    async def openai_completion(
+        self,
+        model: str,
+        prompt: Union[str, List[str], List[int], List[List[int]]],
+        best_of: Optional[int] = None,
+        echo: Optional[bool] = None,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        guided_choice: Optional[List[str]] = None,
+        prompt_logprobs: Optional[int] = None,
+    ) -> OpenAICompletion:
+        if not isinstance(prompt, str):
+            raise ValueError("Ollama does not support non-string prompts for completion")
+
+        model_obj = await self._get_model(model)
+        params = {
+            k: v
+            for k, v in {
+                "model": model_obj.provider_resource_id,
+                "prompt": prompt,
+                "best_of": best_of,
+                "echo": echo,
+                "frequency_penalty": frequency_penalty,
+                "logit_bias": logit_bias,
+                "logprobs": logprobs,
+                "max_tokens": max_tokens,
+                "n": n,
+                "presence_penalty": presence_penalty,
+                "seed": seed,
+                "stop": stop,
+                "stream": stream,
+                "stream_options": stream_options,
+                "temperature": temperature,
+                "top_p": top_p,
+                "user": user,
+            }.items()
+            if v is not None
+        }
+        return await self.openai_client.completions.create(**params)  # type: ignore
+
+    async def openai_chat_completion(
+        self,
+        model: str,
+        messages: List[OpenAIMessageParam],
+        frequency_penalty: Optional[float] = None,
+        function_call: Optional[Union[str, Dict[str, Any]]] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        parallel_tool_calls: Optional[bool] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[OpenAIResponseFormatParam] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+    ) -> Union[OpenAIChatCompletion, AsyncIterator[OpenAIChatCompletionChunk]]:
+        model_obj = await self._get_model(model)
+        params = {
+            k: v
+            for k, v in {
+                "model": model_obj.provider_resource_id,
+                "messages": messages,
+                "frequency_penalty": frequency_penalty,
+                "function_call": function_call,
+                "functions": functions,
+                "logit_bias": logit_bias,
+                "logprobs": logprobs,
+                "max_completion_tokens": max_completion_tokens,
+                "max_tokens": max_tokens,
+                "n": n,
+                "parallel_tool_calls": parallel_tool_calls,
+                "presence_penalty": presence_penalty,
+                "response_format": response_format,
+                "seed": seed,
+                "stop": stop,
+                "stream": stream,
+                "stream_options": stream_options,
+                "temperature": temperature,
+                "tool_choice": tool_choice,
+                "tools": tools,
+                "top_logprobs": top_logprobs,
+                "top_p": top_p,
+                "user": user,
+            }.items()
+            if v is not None
+        }
+        return await self.openai_client.chat.completions.create(**params)  # type: ignore
+
+    async def batch_completion(
+        self,
+        model_id: str,
+        content_batch: List[InterleavedContent],
+        sampling_params: Optional[SamplingParams] = None,
+        response_format: Optional[ResponseFormat] = None,
+        logprobs: Optional[LogProbConfig] = None,
+    ):
+        raise NotImplementedError("Batch completion is not supported for Ollama")
+
+    async def batch_chat_completion(
+        self,
+        model_id: str,
+        messages_batch: List[List[Message]],
+        sampling_params: Optional[SamplingParams] = None,
+        tools: Optional[List[ToolDefinition]] = None,
+        tool_config: Optional[ToolConfig] = None,
+        response_format: Optional[ResponseFormat] = None,
+        logprobs: Optional[LogProbConfig] = None,
+    ):
+        raise NotImplementedError("Batch chat completion is not supported for Ollama")
+
 
 async def convert_message_to_openai_dict_for_ollama(message: Message) -> List[dict]:
     async def _convert_content(content) -> dict:

llama_stack/providers/remote/inference/passthrough/passthrough.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Any, AsyncGenerator, Dict, List, Optional
+from typing import Any, AsyncGenerator, AsyncIterator, Dict, List, Optional, Union
 
 from llama_stack_client import AsyncLlamaStackClient
 
@@ -26,9 +26,17 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
+from llama_stack.apis.inference.inference import (
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAICompletion,
+    OpenAIMessageParam,
+    OpenAIResponseFormatParam,
+)
 from llama_stack.apis.models import Model
 from llama_stack.distribution.library_client import convert_pydantic_to_json_value, convert_to_pydantic
 from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
+from llama_stack.providers.utils.inference.openai_compat import prepare_openai_completion_params
 
 from .config import PassthroughImplConfig
 
@@ -201,6 +209,112 @@ class PassthroughInferenceAdapter(Inference):
             task_type=task_type,
         )
 
+    async def openai_completion(
+        self,
+        model: str,
+        prompt: Union[str, List[str], List[int], List[List[int]]],
+        best_of: Optional[int] = None,
+        echo: Optional[bool] = None,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        guided_choice: Optional[List[str]] = None,
+        prompt_logprobs: Optional[int] = None,
+    ) -> OpenAICompletion:
+        client = self._get_client()
+        model_obj = await self.model_store.get_model(model)
+
+        params = await prepare_openai_completion_params(
+            model=model_obj.provider_resource_id,
+            prompt=prompt,
+            best_of=best_of,
+            echo=echo,
+            frequency_penalty=frequency_penalty,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_tokens=max_tokens,
+            n=n,
+            presence_penalty=presence_penalty,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            top_p=top_p,
+            user=user,
+            guided_choice=guided_choice,
+            prompt_logprobs=prompt_logprobs,
+        )
+
+        return await client.inference.openai_completion(**params)
+
+    async def openai_chat_completion(
+        self,
+        model: str,
+        messages: List[OpenAIMessageParam],
+        frequency_penalty: Optional[float] = None,
+        function_call: Optional[Union[str, Dict[str, Any]]] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        parallel_tool_calls: Optional[bool] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[OpenAIResponseFormatParam] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+    ) -> Union[OpenAIChatCompletion, AsyncIterator[OpenAIChatCompletionChunk]]:
+        client = self._get_client()
+        model_obj = await self.model_store.get_model(model)
+
+        params = await prepare_openai_completion_params(
+            model=model_obj.provider_resource_id,
+            messages=messages,
+            frequency_penalty=frequency_penalty,
+            function_call=function_call,
+            functions=functions,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_completion_tokens=max_completion_tokens,
+            max_tokens=max_tokens,
+            n=n,
+            parallel_tool_calls=parallel_tool_calls,
+            presence_penalty=presence_penalty,
+            response_format=response_format,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            tool_choice=tool_choice,
+            tools=tools,
+            top_logprobs=top_logprobs,
+            top_p=top_p,
+            user=user,
+        )
+
+        return await client.inference.openai_chat_completion(**params)
+
     def cast_value_to_json_dict(self, request_params: Dict[str, Any]) -> Dict[str, Any]:
         json_params = {}
         for key, value in request_params.items():

llama_stack/providers/remote/inference/runpod/runpod.py

@@ -12,6 +12,8 @@ from llama_stack.apis.inference import *  # noqa: F403
 # from llama_stack.providers.datatypes import ModelsProtocolPrivate
 from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
 from llama_stack.providers.utils.inference.openai_compat import (
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
     get_sampling_options,
     process_chat_completion_response,
     process_chat_completion_stream_response,
@@ -38,7 +40,12 @@ RUNPOD_SUPPORTED_MODELS = {
 }
 
 
-class RunpodInferenceAdapter(ModelRegistryHelper, Inference):
+class RunpodInferenceAdapter(
+    ModelRegistryHelper,
+    Inference,
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
+):
     def __init__(self, config: RunpodImplConfig) -> None:
         ModelRegistryHelper.__init__(self, stack_to_provider_models_map=RUNPOD_SUPPORTED_MODELS)
         self.config = config

llama_stack/providers/remote/inference/sambanova/sambanova.py

@@ -42,6 +42,8 @@ from llama_stack.apis.inference import (
 )
 from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
 from llama_stack.providers.utils.inference.openai_compat import (
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
     process_chat_completion_stream_response,
 )
 from llama_stack.providers.utils.inference.prompt_adapter import (
@@ -52,7 +54,12 @@ from .config import SambaNovaImplConfig
 from .models import MODEL_ENTRIES
 
 
-class SambaNovaInferenceAdapter(ModelRegistryHelper, Inference):
+class SambaNovaInferenceAdapter(
+    ModelRegistryHelper,
+    Inference,
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
+):
     def __init__(self, config: SambaNovaImplConfig) -> None:
         ModelRegistryHelper.__init__(self, model_entries=MODEL_ENTRIES)
         self.config = config

llama_stack/providers/remote/inference/tgi/tgi.py

@@ -40,8 +40,10 @@ from llama_stack.providers.utils.inference.model_registry import (
     build_hf_repo_model_entry,
 )
 from llama_stack.providers.utils.inference.openai_compat import (
+    OpenAIChatCompletionToLlamaStackMixin,
     OpenAICompatCompletionChoice,
     OpenAICompatCompletionResponse,
+    OpenAICompletionToLlamaStackMixin,
     get_sampling_options,
     process_chat_completion_response,
     process_chat_completion_stream_response,
@@ -69,7 +71,12 @@ def build_hf_repo_model_entries():
     ]
 
 
-class _HfAdapter(Inference, ModelsProtocolPrivate):
+class _HfAdapter(
+    Inference,
+    OpenAIChatCompletionToLlamaStackMixin,
+    OpenAICompletionToLlamaStackMixin,
+    ModelsProtocolPrivate,
+):
     client: AsyncInferenceClient
     max_tokens: int
     model_id: str

llama_stack/providers/remote/inference/together/together.py

@@ -4,8 +4,9 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import AsyncGenerator, List, Optional, Union
+from typing import Any, AsyncGenerator, AsyncIterator, Dict, List, Optional, Union
 
+from openai import AsyncOpenAI
 from together import AsyncTogether
 
 from llama_stack.apis.common.content_types import (
@@ -30,12 +31,20 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
+from llama_stack.apis.inference.inference import (
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAICompletion,
+    OpenAIMessageParam,
+    OpenAIResponseFormatParam,
+)
 from llama_stack.distribution.request_headers import NeedsRequestProviderData
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
 from llama_stack.providers.utils.inference.openai_compat import (
     convert_message_to_openai_dict,
     get_sampling_options,
+    prepare_openai_completion_params,
     process_chat_completion_response,
     process_chat_completion_stream_response,
     process_completion_response,
@@ -60,6 +69,7 @@ class TogetherInferenceAdapter(ModelRegistryHelper, Inference, NeedsRequestProvi
         ModelRegistryHelper.__init__(self, MODEL_ENTRIES)
         self.config = config
         self._client = None
+        self._openai_client = None
 
     async def initialize(self) -> None:
         pass
@@ -110,6 +120,15 @@ class TogetherInferenceAdapter(ModelRegistryHelper, Inference, NeedsRequestProvi
             self._client = AsyncTogether(api_key=together_api_key)
         return self._client
 
+    def _get_openai_client(self) -> AsyncOpenAI:
+        if not self._openai_client:
+            together_client = self._get_client().client
+            self._openai_client = AsyncOpenAI(
+                base_url=together_client.base_url,
+                api_key=together_client.api_key,
+            )
+        return self._openai_client
+
     async def _nonstream_completion(self, request: CompletionRequest) -> ChatCompletionResponse:
         params = await self._get_params(request)
         client = self._get_client()
@@ -243,3 +262,123 @@ class TogetherInferenceAdapter(ModelRegistryHelper, Inference, NeedsRequestProvi
         )
         embeddings = [item.embedding for item in r.data]
         return EmbeddingsResponse(embeddings=embeddings)
+
+    async def openai_completion(
+        self,
+        model: str,
+        prompt: Union[str, List[str], List[int], List[List[int]]],
+        best_of: Optional[int] = None,
+        echo: Optional[bool] = None,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        guided_choice: Optional[List[str]] = None,
+        prompt_logprobs: Optional[int] = None,
+    ) -> OpenAICompletion:
+        model_obj = await self.model_store.get_model(model)
+        params = await prepare_openai_completion_params(
+            model=model_obj.provider_resource_id,
+            prompt=prompt,
+            best_of=best_of,
+            echo=echo,
+            frequency_penalty=frequency_penalty,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_tokens=max_tokens,
+            n=n,
+            presence_penalty=presence_penalty,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            top_p=top_p,
+            user=user,
+        )
+        return await self._get_openai_client().completions.create(**params)  # type: ignore
+
+    async def openai_chat_completion(
+        self,
+        model: str,
+        messages: List[OpenAIMessageParam],
+        frequency_penalty: Optional[float] = None,
+        function_call: Optional[Union[str, Dict[str, Any]]] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        parallel_tool_calls: Optional[bool] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[OpenAIResponseFormatParam] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+    ) -> Union[OpenAIChatCompletion, AsyncIterator[OpenAIChatCompletionChunk]]:
+        model_obj = await self.model_store.get_model(model)
+        params = await prepare_openai_completion_params(
+            model=model_obj.provider_resource_id,
+            messages=messages,
+            frequency_penalty=frequency_penalty,
+            function_call=function_call,
+            functions=functions,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_completion_tokens=max_completion_tokens,
+            max_tokens=max_tokens,
+            n=n,
+            parallel_tool_calls=parallel_tool_calls,
+            presence_penalty=presence_penalty,
+            response_format=response_format,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            tool_choice=tool_choice,
+            tools=tools,
+            top_logprobs=top_logprobs,
+            top_p=top_p,
+            user=user,
+        )
+        if params.get("stream", True):
+            return self._stream_openai_chat_completion(params)
+        return await self._get_openai_client().chat.completions.create(**params)  # type: ignore
+
+    async def _stream_openai_chat_completion(self, params: dict) -> AsyncGenerator:
+        # together.ai sometimes adds usage data to the stream, even if include_usage is False
+        # This causes an unexpected final chunk with empty choices array to be sent
+        # to clients that may not handle it gracefully.
+        include_usage = False
+        if params.get("stream_options", None):
+            include_usage = params["stream_options"].get("include_usage", False)
+        stream = await self._get_openai_client().chat.completions.create(**params)
+
+        seen_finish_reason = False
+        async for chunk in stream:
+            # Final usage chunk with no choices that the user didn't request, so discard
+            if not include_usage and seen_finish_reason and len(chunk.choices) == 0:
+                break
+            yield chunk
+            for choice in chunk.choices:
+                if choice.finish_reason:
+                    seen_finish_reason = True
+                    break

llama_stack/providers/remote/inference/vllm/vllm.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 import json
 import logging
-from typing import Any, AsyncGenerator, List, Optional, Union
+from typing import Any, AsyncGenerator, AsyncIterator, Dict, List, Optional, Union
 
 import httpx
 from openai import AsyncOpenAI
@@ -45,6 +45,12 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
+from llama_stack.apis.inference.inference import (
+    OpenAIChatCompletion,
+    OpenAICompletion,
+    OpenAIMessageParam,
+    OpenAIResponseFormatParam,
+)
 from llama_stack.apis.models import Model, ModelType
 from llama_stack.models.llama.datatypes import BuiltinTool, StopReason, ToolCall
 from llama_stack.models.llama.sku_list import all_registered_models
@@ -58,6 +64,7 @@ from llama_stack.providers.utils.inference.openai_compat import (
     convert_message_to_openai_dict,
     convert_tool_call,
     get_sampling_options,
+    prepare_openai_completion_params,
     process_chat_completion_stream_response,
     process_completion_response,
     process_completion_stream_response,
@@ -418,3 +425,131 @@ class VLLMInferenceAdapter(Inference, ModelsProtocolPrivate):
 
         embeddings = [data.embedding for data in response.data]
         return EmbeddingsResponse(embeddings=embeddings)
+
+    async def openai_completion(
+        self,
+        model: str,
+        prompt: Union[str, List[str], List[int], List[List[int]]],
+        best_of: Optional[int] = None,
+        echo: Optional[bool] = None,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        guided_choice: Optional[List[str]] = None,
+        prompt_logprobs: Optional[int] = None,
+    ) -> OpenAICompletion:
+        model_obj = await self._get_model(model)
+
+        extra_body: Dict[str, Any] = {}
+        if prompt_logprobs is not None and prompt_logprobs >= 0:
+            extra_body["prompt_logprobs"] = prompt_logprobs
+        if guided_choice:
+            extra_body["guided_choice"] = guided_choice
+
+        params = await prepare_openai_completion_params(
+            model=model_obj.provider_resource_id,
+            prompt=prompt,
+            best_of=best_of,
+            echo=echo,
+            frequency_penalty=frequency_penalty,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_tokens=max_tokens,
+            n=n,
+            presence_penalty=presence_penalty,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            top_p=top_p,
+            user=user,
+            extra_body=extra_body,
+        )
+        return await self.client.completions.create(**params)  # type: ignore
+
+    async def openai_chat_completion(
+        self,
+        model: str,
+        messages: List[OpenAIMessageParam],
+        frequency_penalty: Optional[float] = None,
+        function_call: Optional[Union[str, Dict[str, Any]]] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        parallel_tool_calls: Optional[bool] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[OpenAIResponseFormatParam] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+    ) -> Union[OpenAIChatCompletion, AsyncIterator[OpenAIChatCompletionChunk]]:
+        model_obj = await self._get_model(model)
+        params = await prepare_openai_completion_params(
+            model=model_obj.provider_resource_id,
+            messages=messages,
+            frequency_penalty=frequency_penalty,
+            function_call=function_call,
+            functions=functions,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_completion_tokens=max_completion_tokens,
+            max_tokens=max_tokens,
+            n=n,
+            parallel_tool_calls=parallel_tool_calls,
+            presence_penalty=presence_penalty,
+            response_format=response_format,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            tool_choice=tool_choice,
+            tools=tools,
+            top_logprobs=top_logprobs,
+            top_p=top_p,
+            user=user,
+        )
+        return await self.client.chat.completions.create(**params)  # type: ignore
+
+    async def batch_completion(
+        self,
+        model_id: str,
+        content_batch: List[InterleavedContent],
+        sampling_params: Optional[SamplingParams] = None,
+        response_format: Optional[ResponseFormat] = None,
+        logprobs: Optional[LogProbConfig] = None,
+    ):
+        raise NotImplementedError("Batch completion is not supported for Ollama")
+
+    async def batch_chat_completion(
+        self,
+        model_id: str,
+        messages_batch: List[List[Message]],
+        sampling_params: Optional[SamplingParams] = None,
+        tools: Optional[List[ToolDefinition]] = None,
+        tool_config: Optional[ToolConfig] = None,
+        response_format: Optional[ResponseFormat] = None,
+        logprobs: Optional[LogProbConfig] = None,
+    ):
+        raise NotImplementedError("Batch chat completion is not supported for Ollama")

llama_stack/providers/remote/post_training/nvidia/post_training.py

@@ -206,10 +206,6 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
         model: str,
         checkpoint_dir: Optional[str],
         algorithm_config: Optional[AlgorithmConfig] = None,
-        extra_json: Optional[Dict[str, Any]] = None,
-        params: Optional[Dict[str, Any]] = None,
-        headers: Optional[Dict[str, Any]] = None,
-        **kwargs,
     ) -> NvidiaPostTrainingJob:
         """
         Fine-tunes a model on a dataset.

llama_stack/providers/remote/safety/nvidia/nvidia.py

@@ -104,6 +104,15 @@ class NeMoGuardrails:
         self.threshold = threshold
         self.guardrails_service_url = config.guardrails_service_url
 
+    async def _guardrails_post(self, path: str, data: Any | None):
+        """Helper for making POST requests to the guardrails service."""
+        headers = {
+            "Accept": "application/json",
+        }
+        response = requests.post(url=f"{self.guardrails_service_url}{path}", headers=headers, json=data)
+        response.raise_for_status()
+        return response.json()
+
     async def run(self, messages: List[Message]) -> RunShieldResponse:
         """
         Queries the /v1/guardrails/checks endpoint of the NeMo guardrails deployed API.
@@ -118,9 +127,6 @@ class NeMoGuardrails:
         Raises:
             requests.HTTPError: If the POST request fails.
         """
-        headers = {
-            "Accept": "application/json",
-        }
         request_data = {
             "model": self.model,
             "messages": convert_pydantic_to_json_value(messages),
@@ -134,15 +140,11 @@ class NeMoGuardrails:
                 "config_id": self.config_id,
             },
         }
-        response = requests.post(
-            url=f"{self.guardrails_service_url}/v1/guardrail/checks", headers=headers, json=request_data
-        )
-        response.raise_for_status()
-        if "Content-Type" in response.headers and response.headers["Content-Type"].startswith("application/json"):
-            response_json = response.json()
-        if response_json["status"] == "blocked":
+        response = await self._guardrails_post(path="/v1/guardrail/checks", data=request_data)
+
+        if response["status"] == "blocked":
             user_message = "Sorry I cannot do this."
-            metadata = response_json["rails_status"]
+            metadata = response["rails_status"]
 
             return RunShieldResponse(
                 violation=SafetyViolation(
@@ -151,4 +153,5 @@ class NeMoGuardrails:
                     metadata=metadata,
                 )
             )
+
         return RunShieldResponse(violation=None)

llama_stack/providers/utils/inference/litellm_openai_mixin.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import AsyncGenerator, AsyncIterator, List, Optional, Union
+from typing import Any, AsyncGenerator, AsyncIterator, Dict, List, Optional, Union
 
 import litellm
 
@@ -30,6 +30,13 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
+from llama_stack.apis.inference.inference import (
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAICompletion,
+    OpenAIMessageParam,
+    OpenAIResponseFormatParam,
+)
 from llama_stack.apis.models.models import Model
 from llama_stack.distribution.request_headers import NeedsRequestProviderData
 from llama_stack.log import get_logger
@@ -40,6 +47,7 @@ from llama_stack.providers.utils.inference.openai_compat import (
     convert_openai_chat_completion_stream,
     convert_tooldef_to_openai_tool,
     get_sampling_options,
+    prepare_openai_completion_params,
 )
 from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,
@@ -245,3 +253,125 @@ class LiteLLMOpenAIMixin(
 
         embeddings = [data["embedding"] for data in response["data"]]
         return EmbeddingsResponse(embeddings=embeddings)
+
+    async def openai_completion(
+        self,
+        model: str,
+        prompt: Union[str, List[str], List[int], List[List[int]]],
+        best_of: Optional[int] = None,
+        echo: Optional[bool] = None,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        guided_choice: Optional[List[str]] = None,
+        prompt_logprobs: Optional[int] = None,
+    ) -> OpenAICompletion:
+        model_obj = await self.model_store.get_model(model)
+        params = await prepare_openai_completion_params(
+            model=model_obj.provider_resource_id,
+            prompt=prompt,
+            best_of=best_of,
+            echo=echo,
+            frequency_penalty=frequency_penalty,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_tokens=max_tokens,
+            n=n,
+            presence_penalty=presence_penalty,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            top_p=top_p,
+            user=user,
+            guided_choice=guided_choice,
+            prompt_logprobs=prompt_logprobs,
+        )
+        return await litellm.atext_completion(**params)
+
+    async def openai_chat_completion(
+        self,
+        model: str,
+        messages: List[OpenAIMessageParam],
+        frequency_penalty: Optional[float] = None,
+        function_call: Optional[Union[str, Dict[str, Any]]] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        parallel_tool_calls: Optional[bool] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[OpenAIResponseFormatParam] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+    ) -> Union[OpenAIChatCompletion, AsyncIterator[OpenAIChatCompletionChunk]]:
+        model_obj = await self.model_store.get_model(model)
+        params = await prepare_openai_completion_params(
+            model=model_obj.provider_resource_id,
+            messages=messages,
+            frequency_penalty=frequency_penalty,
+            function_call=function_call,
+            functions=functions,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_completion_tokens=max_completion_tokens,
+            max_tokens=max_tokens,
+            n=n,
+            parallel_tool_calls=parallel_tool_calls,
+            presence_penalty=presence_penalty,
+            response_format=response_format,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=stream_options,
+            temperature=temperature,
+            tool_choice=tool_choice,
+            tools=tools,
+            top_logprobs=top_logprobs,
+            top_p=top_p,
+            user=user,
+        )
+        return await litellm.acompletion(**params)
+
+    async def batch_completion(
+        self,
+        model_id: str,
+        content_batch: List[InterleavedContent],
+        sampling_params: Optional[SamplingParams] = None,
+        response_format: Optional[ResponseFormat] = None,
+        logprobs: Optional[LogProbConfig] = None,
+    ):
+        raise NotImplementedError("Batch completion is not supported for OpenAI Compat")
+
+    async def batch_chat_completion(
+        self,
+        model_id: str,
+        messages_batch: List[List[Message]],
+        sampling_params: Optional[SamplingParams] = None,
+        tools: Optional[List[ToolDefinition]] = None,
+        tool_config: Optional[ToolConfig] = None,
+        response_format: Optional[ResponseFormat] = None,
+        logprobs: Optional[LogProbConfig] = None,
+    ):
+        raise NotImplementedError("Batch chat completion is not supported for OpenAI Compat")

llama_stack/providers/utils/inference/openai_compat.py

@@ -5,8 +5,10 @@
 # the root directory of this source tree.
 import json
 import logging
+import time
+import uuid
 import warnings
-from typing import AsyncGenerator, Dict, Iterable, List, Optional, Union
+from typing import Any, AsyncGenerator, AsyncIterator, Awaitable, Dict, Iterable, List, Optional, Union
 
 from openai import AsyncStream
 from openai.types.chat import (
@@ -48,6 +50,18 @@ from openai.types.chat.chat_completion import (
 from openai.types.chat.chat_completion import (
     ChoiceLogprobs as OpenAIChoiceLogprobs,  # same as chat_completion_chunk ChoiceLogprobs
 )
+from openai.types.chat.chat_completion_chunk import (
+    Choice as OpenAIChatCompletionChunkChoice,
+)
+from openai.types.chat.chat_completion_chunk import (
+    ChoiceDelta as OpenAIChoiceDelta,
+)
+from openai.types.chat.chat_completion_chunk import (
+    ChoiceDeltaToolCall as OpenAIChoiceDeltaToolCall,
+)
+from openai.types.chat.chat_completion_chunk import (
+    ChoiceDeltaToolCallFunction as OpenAIChoiceDeltaToolCallFunction,
+)
 from openai.types.chat.chat_completion_content_part_image_param import (
     ImageURL as OpenAIImageURL,
 )
@@ -57,6 +71,7 @@ from openai.types.chat.chat_completion_message_tool_call_param import (
 from pydantic import BaseModel
 
 from llama_stack.apis.common.content_types import (
+    URL,
     ImageContentItem,
     InterleavedContent,
     TextContentItem,
@@ -83,11 +98,24 @@ from llama_stack.apis.inference import (
     TopPSamplingStrategy,
     UserMessage,
 )
+from llama_stack.apis.inference.inference import (
+    JsonSchemaResponseFormat,
+    OpenAIChatCompletion,
+    OpenAICompletion,
+    OpenAICompletionChoice,
+    OpenAIMessageParam,
+    OpenAIResponseFormatParam,
+    ToolConfig,
+)
+from llama_stack.apis.inference.inference import (
+    OpenAIChoice as OpenAIChatCompletionChoice,
+)
 from llama_stack.models.llama.datatypes import (
     BuiltinTool,
     StopReason,
     ToolCall,
     ToolDefinition,
+    ToolParamDefinition,
 )
 from llama_stack.providers.utils.inference.prompt_adapter import (
     convert_image_content_to_url,
@@ -748,6 +776,17 @@ def convert_tooldef_to_openai_tool(tool: ToolDefinition) -> dict:
     return out
 
 
+def _convert_stop_reason_to_openai_finish_reason(stop_reason: StopReason) -> str:
+    """
+    Convert a StopReason to an OpenAI chat completion finish_reason.
+    """
+    return {
+        StopReason.end_of_turn: "stop",
+        StopReason.end_of_message: "tool_calls",
+        StopReason.out_of_tokens: "length",
+    }.get(stop_reason, "stop")
+
+
 def _convert_openai_finish_reason(finish_reason: str) -> StopReason:
     """
     Convert an OpenAI chat completion finish_reason to a StopReason.
@@ -773,6 +812,56 @@ def _convert_openai_finish_reason(finish_reason: str) -> StopReason:
     }.get(finish_reason, StopReason.end_of_turn)
 
 
+def _convert_openai_request_tool_config(tool_choice: Optional[Union[str, Dict[str, Any]]] = None) -> ToolConfig:
+    tool_config = ToolConfig()
+    if tool_choice:
+        tool_config.tool_choice = tool_choice
+    return tool_config
+
+
+def _convert_openai_request_tools(tools: Optional[List[Dict[str, Any]]] = None) -> List[ToolDefinition]:
+    lls_tools = []
+    if not tools:
+        return lls_tools
+
+    for tool in tools:
+        tool_fn = tool.get("function", {})
+        tool_name = tool_fn.get("name", None)
+        tool_desc = tool_fn.get("description", None)
+
+        tool_params = tool_fn.get("parameters", None)
+        lls_tool_params = {}
+        if tool_params is not None:
+            tool_param_properties = tool_params.get("properties", {})
+            for tool_param_key, tool_param_value in tool_param_properties.items():
+                tool_param_def = ToolParamDefinition(
+                    param_type=tool_param_value.get("type", None),
+                    description=tool_param_value.get("description", None),
+                )
+                lls_tool_params[tool_param_key] = tool_param_def
+
+        lls_tool = ToolDefinition(
+            tool_name=tool_name,
+            description=tool_desc,
+            parameters=lls_tool_params,
+        )
+        lls_tools.append(lls_tool)
+    return lls_tools
+
+
+def _convert_openai_request_response_format(response_format: OpenAIResponseFormatParam = None):
+    if not response_format:
+        return None
+    # response_format can be a dict or a pydantic model
+    response_format = dict(response_format)
+    if response_format.get("type", "") == "json_schema":
+        return JsonSchemaResponseFormat(
+            type="json_schema",
+            json_schema=response_format.get("json_schema", {}).get("schema", ""),
+        )
+    return None
+
+
 def _convert_openai_tool_calls(
     tool_calls: List[OpenAIChatCompletionMessageToolCall],
 ) -> List[ToolCall]:
@@ -843,6 +932,65 @@ def _convert_openai_logprobs(
     ]
 
 
+def _convert_openai_sampling_params(
+    max_tokens: Optional[int] = None,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+) -> SamplingParams:
+    sampling_params = SamplingParams()
+
+    if max_tokens:
+        sampling_params.max_tokens = max_tokens
+
+    # Map an explicit temperature of 0 to greedy sampling
+    if temperature == 0:
+        strategy = GreedySamplingStrategy()
+    else:
+        # OpenAI defaults to 1.0 for temperature and top_p if unset
+        if temperature is None:
+            temperature = 1.0
+        if top_p is None:
+            top_p = 1.0
+        strategy = TopPSamplingStrategy(temperature=temperature, top_p=top_p)
+
+    sampling_params.strategy = strategy
+    return sampling_params
+
+
+def _convert_openai_request_messages(messages: List[OpenAIMessageParam]):
+    # Llama Stack messages and OpenAI messages are similar, but not identical.
+    lls_messages = []
+    for message in messages:
+        lls_message = dict(message)
+
+        #  Llama Stack expects `call_id` but OpenAI uses `tool_call_id`
+        tool_call_id = lls_message.pop("tool_call_id", None)
+        if tool_call_id:
+            lls_message["call_id"] = tool_call_id
+
+        content = lls_message.get("content", None)
+        if isinstance(content, list):
+            lls_content = []
+            for item in content:
+                # items can either by pydantic models or dicts here...
+                item = dict(item)
+                if item.get("type", "") == "image_url":
+                    lls_item = ImageContentItem(
+                        type="image",
+                        image=URL(uri=item.get("image_url", {}).get("url", "")),
+                    )
+                elif item.get("type", "") == "text":
+                    lls_item = TextContentItem(
+                        type="text",
+                        text=item.get("text", ""),
+                    )
+                lls_content.append(lls_item)
+            lls_message["content"] = lls_content
+        lls_messages.append(lls_message)
+
+    return lls_messages
+
+
 def convert_openai_chat_completion_choice(
     choice: OpenAIChoice,
 ) -> ChatCompletionResponse:
@@ -1049,3 +1197,218 @@ async def convert_openai_chat_completion_stream(
             stop_reason=stop_reason,
         )
     )
+
+
+async def prepare_openai_completion_params(**params):
+    async def _prepare_value(value: Any) -> Any:
+        new_value = value
+        if isinstance(value, list):
+            new_value = [await _prepare_value(v) for v in value]
+        elif isinstance(value, dict):
+            new_value = {k: await _prepare_value(v) for k, v in value.items()}
+        elif isinstance(value, BaseModel):
+            new_value = value.model_dump(exclude_none=True)
+        return new_value
+
+    completion_params = {}
+    for k, v in params.items():
+        if v is not None:
+            completion_params[k] = await _prepare_value(v)
+    return completion_params
+
+
+class OpenAICompletionToLlamaStackMixin:
+    async def openai_completion(
+        self,
+        model: str,
+        prompt: Union[str, List[str], List[int], List[List[int]]],
+        best_of: Optional[int] = None,
+        echo: Optional[bool] = None,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        guided_choice: Optional[List[str]] = None,
+        prompt_logprobs: Optional[int] = None,
+    ) -> OpenAICompletion:
+        if stream:
+            raise ValueError(f"{self.__class__.__name__} doesn't support streaming openai completions")
+
+        # This is a pretty hacky way to do emulate completions -
+        # basically just de-batches them...
+        prompts = [prompt] if not isinstance(prompt, list) else prompt
+
+        sampling_params = _convert_openai_sampling_params(
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+        )
+
+        choices = []
+        # "n" is the number of completions to generate per prompt
+        n = n or 1
+        for _i in range(0, n):
+            # and we may have multiple prompts, if batching was used
+
+            for prompt in prompts:
+                result = self.completion(
+                    model_id=model,
+                    content=prompt,
+                    sampling_params=sampling_params,
+                )
+
+                index = len(choices)
+                text = result.content
+                finish_reason = _convert_stop_reason_to_openai_finish_reason(result.stop_reason)
+
+                choice = OpenAICompletionChoice(
+                    index=index,
+                    text=text,
+                    finish_reason=finish_reason,
+                )
+                choices.append(choice)
+
+        return OpenAICompletion(
+            id=f"cmpl-{uuid.uuid4()}",
+            choices=choices,
+            created=int(time.time()),
+            model=model,
+            object="text_completion",
+        )
+
+
+class OpenAIChatCompletionToLlamaStackMixin:
+    async def openai_chat_completion(
+        self,
+        model: str,
+        messages: List[OpenAIChatCompletionMessage],
+        frequency_penalty: Optional[float] = None,
+        function_call: Optional[Union[str, Dict[str, Any]]] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        logprobs: Optional[bool] = None,
+        max_completion_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        parallel_tool_calls: Optional[bool] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[OpenAIResponseFormatParam] = None,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stream: Optional[bool] = None,
+        stream_options: Optional[Dict[str, Any]] = None,
+        temperature: Optional[float] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+    ) -> Union[OpenAIChatCompletion, AsyncIterator[OpenAIChatCompletionChunk]]:
+        messages = _convert_openai_request_messages(messages)
+        response_format = _convert_openai_request_response_format(response_format)
+        sampling_params = _convert_openai_sampling_params(
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+        )
+        tool_config = _convert_openai_request_tool_config(tool_choice)
+        tools = _convert_openai_request_tools(tools)
+
+        outstanding_responses = []
+        # "n" is the number of completions to generate per prompt
+        n = n or 1
+        for _i in range(0, n):
+            response = self.chat_completion(
+                model_id=model,
+                messages=messages,
+                sampling_params=sampling_params,
+                response_format=response_format,
+                stream=stream,
+                tool_config=tool_config,
+                tools=tools,
+            )
+            outstanding_responses.append(response)
+
+        if stream:
+            return OpenAIChatCompletionToLlamaStackMixin._process_stream_response(self, model, outstanding_responses)
+
+        return await OpenAIChatCompletionToLlamaStackMixin._process_non_stream_response(
+            self, model, outstanding_responses
+        )
+
+    async def _process_stream_response(
+        self, model: str, outstanding_responses: List[Awaitable[AsyncIterator[ChatCompletionResponseStreamChunk]]]
+    ):
+        id = f"chatcmpl-{uuid.uuid4()}"
+        for outstanding_response in outstanding_responses:
+            response = await outstanding_response
+            i = 0
+            async for chunk in response:
+                event = chunk.event
+                finish_reason = _convert_stop_reason_to_openai_finish_reason(event.stop_reason)
+
+                if isinstance(event.delta, TextDelta):
+                    text_delta = event.delta.text
+                    delta = OpenAIChoiceDelta(content=text_delta)
+                    yield OpenAIChatCompletionChunk(
+                        id=id,
+                        choices=[OpenAIChatCompletionChunkChoice(index=i, finish_reason=finish_reason, delta=delta)],
+                        created=int(time.time()),
+                        model=model,
+                        object="chat.completion.chunk",
+                    )
+                elif isinstance(event.delta, ToolCallDelta):
+                    if event.delta.parse_status == ToolCallParseStatus.succeeded:
+                        tool_call = event.delta.tool_call
+                        openai_tool_call = OpenAIChoiceDeltaToolCall(
+                            index=0,
+                            id=tool_call.call_id,
+                            function=OpenAIChoiceDeltaToolCallFunction(
+                                name=tool_call.tool_name, arguments=tool_call.arguments_json
+                            ),
+                        )
+                        delta = OpenAIChoiceDelta(tool_calls=[openai_tool_call])
+                        yield OpenAIChatCompletionChunk(
+                            id=id,
+                            choices=[
+                                OpenAIChatCompletionChunkChoice(index=i, finish_reason=finish_reason, delta=delta)
+                            ],
+                            created=int(time.time()),
+                            model=model,
+                            object="chat.completion.chunk",
+                        )
+                i = i + 1
+
+    async def _process_non_stream_response(
+        self, model: str, outstanding_responses: List[Awaitable[ChatCompletionResponse]]
+    ) -> OpenAIChatCompletion:
+        choices = []
+        for outstanding_response in outstanding_responses:
+            response = await outstanding_response
+            completion_message = response.completion_message
+            message = await convert_message_to_openai_dict_new(completion_message)
+            finish_reason = _convert_stop_reason_to_openai_finish_reason(completion_message.stop_reason)
+
+            choice = OpenAIChatCompletionChoice(
+                index=len(choices),
+                message=message,
+                finish_reason=finish_reason,
+            )
+            choices.append(choice)
+
+        return OpenAIChatCompletion(
+            id=f"chatcmpl-{uuid.uuid4()}",
+            choices=choices,
+            created=int(time.time()),
+            model=model,
+            object="chat.completion",
+        )

llama_stack/providers/utils/scheduler.py

@@ -0,0 +1,265 @@
+# 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.
+
+import abc
+import asyncio
+import functools
+import threading
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any, Callable, Coroutine, Dict, Iterable, Tuple, TypeAlias
+
+from pydantic import BaseModel
+
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="scheduler")
+
+
+# TODO: revisit the list of possible statuses when defining a more coherent
+# Jobs API for all API flows; e.g. do we need new vs scheduled?
+class JobStatus(Enum):
+    new = "new"
+    scheduled = "scheduled"
+    running = "running"
+    failed = "failed"
+    completed = "completed"
+
+
+JobID: TypeAlias = str
+JobType: TypeAlias = str
+
+
+class JobArtifact(BaseModel):
+    type: JobType
+    name: str
+    # TODO: uri should be a reference to /files API; revisit when /files is implemented
+    uri: str | None = None
+    metadata: Dict[str, Any]
+
+
+JobHandler = Callable[
+    [Callable[[str], None], Callable[[JobStatus], None], Callable[[JobArtifact], None]], Coroutine[Any, Any, None]
+]
+
+
+LogMessage: TypeAlias = Tuple[datetime, str]
+
+
+_COMPLETED_STATUSES = {JobStatus.completed, JobStatus.failed}
+
+
+class Job:
+    def __init__(self, job_type: JobType, job_id: JobID, handler: JobHandler):
+        super().__init__()
+        self.id = job_id
+        self._type = job_type
+        self._handler = handler
+        self._artifacts: list[JobArtifact] = []
+        self._logs: list[LogMessage] = []
+        self._state_transitions: list[Tuple[datetime, JobStatus]] = [(datetime.now(timezone.utc), JobStatus.new)]
+
+    @property
+    def handler(self) -> JobHandler:
+        return self._handler
+
+    @property
+    def status(self) -> JobStatus:
+        return self._state_transitions[-1][1]
+
+    @status.setter
+    def status(self, status: JobStatus):
+        if status in _COMPLETED_STATUSES and self.status in _COMPLETED_STATUSES:
+            raise ValueError(f"Job is already in a completed state ({self.status})")
+        if self.status == status:
+            return
+        self._state_transitions.append((datetime.now(timezone.utc), status))
+
+    @property
+    def artifacts(self) -> list[JobArtifact]:
+        return self._artifacts
+
+    def register_artifact(self, artifact: JobArtifact) -> None:
+        self._artifacts.append(artifact)
+
+    def _find_state_transition_date(self, status: Iterable[JobStatus]) -> datetime | None:
+        for date, s in reversed(self._state_transitions):
+            if s in status:
+                return date
+        return None
+
+    @property
+    def scheduled_at(self) -> datetime | None:
+        return self._find_state_transition_date([JobStatus.scheduled])
+
+    @property
+    def started_at(self) -> datetime | None:
+        return self._find_state_transition_date([JobStatus.running])
+
+    @property
+    def completed_at(self) -> datetime | None:
+        return self._find_state_transition_date(_COMPLETED_STATUSES)
+
+    @property
+    def logs(self) -> list[LogMessage]:
+        return self._logs[:]
+
+    def append_log(self, message: LogMessage) -> None:
+        self._logs.append(message)
+
+    # TODO: implement
+    def cancel(self) -> None:
+        raise NotImplementedError
+
+
+class _SchedulerBackend(abc.ABC):
+    @abc.abstractmethod
+    def on_log_message_cb(self, job: Job, message: LogMessage) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def on_status_change_cb(self, job: Job, status: JobStatus) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def on_artifact_collected_cb(self, job: Job, artifact: JobArtifact) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def shutdown(self) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def schedule(
+        self,
+        job: Job,
+        on_log_message_cb: Callable[[str], None],
+        on_status_change_cb: Callable[[JobStatus], None],
+        on_artifact_collected_cb: Callable[[JobArtifact], None],
+    ) -> None:
+        raise NotImplementedError
+
+
+class _NaiveSchedulerBackend(_SchedulerBackend):
+    def __init__(self, timeout: int = 5):
+        self._timeout = timeout
+        self._loop = asyncio.new_event_loop()
+        # There may be performance implications of using threads due to Python
+        # GIL; may need to measure if it's a real problem though
+        self._thread = threading.Thread(target=self._run_loop, daemon=True)
+        self._thread.start()
+
+    def _run_loop(self) -> None:
+        asyncio.set_event_loop(self._loop)
+        self._loop.run_forever()
+
+        # When stopping the loop, give tasks a chance to finish
+        # TODO: should we explicitly inform jobs of pending stoppage?
+        for task in asyncio.all_tasks(self._loop):
+            self._loop.run_until_complete(task)
+        self._loop.close()
+
+    async def shutdown(self) -> None:
+        self._loop.call_soon_threadsafe(self._loop.stop)
+        self._thread.join()
+
+    # TODO: decouple scheduling and running the job
+    def schedule(
+        self,
+        job: Job,
+        on_log_message_cb: Callable[[str], None],
+        on_status_change_cb: Callable[[JobStatus], None],
+        on_artifact_collected_cb: Callable[[JobArtifact], None],
+    ) -> None:
+        async def do():
+            try:
+                job.status = JobStatus.running
+                await job.handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb)
+            except Exception as e:
+                on_log_message_cb(str(e))
+                job.status = JobStatus.failed
+                logger.exception(f"Job {job.id} failed.")
+
+        asyncio.run_coroutine_threadsafe(do(), self._loop)
+
+    def on_log_message_cb(self, job: Job, message: LogMessage) -> None:
+        pass
+
+    def on_status_change_cb(self, job: Job, status: JobStatus) -> None:
+        pass
+
+    def on_artifact_collected_cb(self, job: Job, artifact: JobArtifact) -> None:
+        pass
+
+
+_BACKENDS = {
+    "naive": _NaiveSchedulerBackend,
+}
+
+
+def _get_backend_impl(backend: str) -> _SchedulerBackend:
+    try:
+        return _BACKENDS[backend]()
+    except KeyError as e:
+        raise ValueError(f"Unknown backend {backend}") from e
+
+
+class Scheduler:
+    def __init__(self, backend: str = "naive"):
+        # TODO: if server crashes, job states are lost; we need to persist jobs on disc
+        self._jobs: dict[JobID, Job] = {}
+        self._backend = _get_backend_impl(backend)
+
+    def _on_log_message_cb(self, job: Job, message: str) -> None:
+        msg = (datetime.now(timezone.utc), message)
+        # At least for the time being, until there's a better way to expose
+        # logs to users, log messages on console
+        logger.info(f"Job {job.id}: {message}")
+        job.append_log(msg)
+        self._backend.on_log_message_cb(job, msg)
+
+    def _on_status_change_cb(self, job: Job, status: JobStatus) -> None:
+        job.status = status
+        self._backend.on_status_change_cb(job, status)
+
+    def _on_artifact_collected_cb(self, job: Job, artifact: JobArtifact) -> None:
+        job.register_artifact(artifact)
+        self._backend.on_artifact_collected_cb(job, artifact)
+
+    def schedule(self, type_: JobType, job_id: JobID, handler: JobHandler) -> JobID:
+        job = Job(type_, job_id, handler)
+        if job.id in self._jobs:
+            raise ValueError(f"Job {job.id} already exists")
+
+        self._jobs[job.id] = job
+        job.status = JobStatus.scheduled
+        self._backend.schedule(
+            job,
+            functools.partial(self._on_log_message_cb, job),
+            functools.partial(self._on_status_change_cb, job),
+            functools.partial(self._on_artifact_collected_cb, job),
+        )
+
+        return job.id
+
+    def cancel(self, job_id: JobID) -> None:
+        self.get_job(job_id).cancel()
+
+    def get_job(self, job_id: JobID) -> Job:
+        try:
+            return self._jobs[job_id]
+        except KeyError as e:
+            raise ValueError(f"Job {job_id} not found") from e
+
+    def get_jobs(self, type_: JobType | None = None) -> list[Job]:
+        jobs = list(self._jobs.values())
+        if type_:
+            jobs = [job for job in jobs if job._type == type_]
+        return jobs
+
+    async def shutdown(self):
+        # TODO: also cancel jobs once implemented
+        await self._backend.shutdown()

llama_stack/schema_utils.py

@@ -20,6 +20,7 @@ class WebMethod:
     raw_bytes_request_body: Optional[bool] = False
     # A descriptive name of the corresponding span created by tracing
     descriptive_name: Optional[str] = None
+    experimental: Optional[bool] = False
 
 
 T = TypeVar("T", bound=Callable[..., Any])
@@ -33,6 +34,7 @@ def webmethod(
     response_examples: Optional[List[Any]] = None,
     raw_bytes_request_body: Optional[bool] = False,
     descriptive_name: Optional[str] = None,
+    experimental: Optional[bool] = False,
 ) -> Callable[[T], T]:
     """
     Decorator that supplies additional metadata to an endpoint operation function.
@@ -41,6 +43,7 @@ def webmethod(
     :param public: True if the operation can be invoked without prior authentication.
     :param request_examples: Sample requests that the operation might take. Pass a list of objects, not JSON.
     :param response_examples: Sample responses that the operation might produce. Pass a list of objects, not JSON.
+    :param experimental: True if the operation is experimental and subject to change.
     """
 
     def wrap(func: T) -> T:
@@ -52,6 +55,7 @@ def webmethod(
             response_examples=response_examples,
             raw_bytes_request_body=raw_bytes_request_body,
             descriptive_name=descriptive_name,
+            experimental=experimental,
         )
         return func
 

llama_stack/templates/dependencies.json

@@ -381,7 +381,7 @@
     "sentence-transformers",
     "sentencepiece",
     "torch",
-    "torchao==0.5.0",
+    "torchao==0.8.0",
     "torchvision",
     "tqdm",
     "transformers",

llama_stack/templates/dev/run.yaml

@@ -386,6 +386,16 @@ models:
   provider_id: groq
   provider_model_id: groq/llama-4-scout-17b-16e-instruct
   model_type: llm
+- metadata: {}
+  model_id: groq/meta-llama/llama-4-scout-17b-16e-instruct
+  provider_id: groq
+  provider_model_id: groq/meta-llama/llama-4-scout-17b-16e-instruct
+  model_type: llm
+- metadata: {}
+  model_id: meta-llama/Llama-4-Scout-17B-16E-Instruct
+  provider_id: groq
+  provider_model_id: groq/meta-llama/llama-4-scout-17b-16e-instruct
+  model_type: llm
 - metadata: {}
   model_id: groq/llama-4-maverick-17b-128e-instruct
   provider_id: groq
@@ -396,6 +406,16 @@ models:
   provider_id: groq
   provider_model_id: groq/llama-4-maverick-17b-128e-instruct
   model_type: llm
+- metadata: {}
+  model_id: groq/meta-llama/llama-4-maverick-17b-128e-instruct
+  provider_id: groq
+  provider_model_id: groq/meta-llama/llama-4-maverick-17b-128e-instruct
+  model_type: llm
+- metadata: {}
+  model_id: meta-llama/Llama-4-Maverick-17B-128E-Instruct
+  provider_id: groq
+  provider_model_id: groq/meta-llama/llama-4-maverick-17b-128e-instruct
+  model_type: llm
 - metadata:
     embedding_dimension: 384
   model_id: all-MiniLM-L6-v2

llama_stack/templates/groq/run.yaml

@@ -158,6 +158,16 @@ models:
   provider_id: groq
   provider_model_id: groq/llama-4-scout-17b-16e-instruct
   model_type: llm
+- metadata: {}
+  model_id: groq/meta-llama/llama-4-scout-17b-16e-instruct
+  provider_id: groq
+  provider_model_id: groq/meta-llama/llama-4-scout-17b-16e-instruct
+  model_type: llm
+- metadata: {}
+  model_id: meta-llama/Llama-4-Scout-17B-16E-Instruct
+  provider_id: groq
+  provider_model_id: groq/meta-llama/llama-4-scout-17b-16e-instruct
+  model_type: llm
 - metadata: {}
   model_id: groq/llama-4-maverick-17b-128e-instruct
   provider_id: groq
@@ -168,6 +178,16 @@ models:
   provider_id: groq
   provider_model_id: groq/llama-4-maverick-17b-128e-instruct
   model_type: llm
+- metadata: {}
+  model_id: groq/meta-llama/llama-4-maverick-17b-128e-instruct
+  provider_id: groq
+  provider_model_id: groq/meta-llama/llama-4-maverick-17b-128e-instruct
+  model_type: llm
+- metadata: {}
+  model_id: meta-llama/Llama-4-Maverick-17B-128E-Instruct
+  provider_id: groq
+  provider_model_id: groq/meta-llama/llama-4-maverick-17b-128e-instruct
+  model_type: llm
 - metadata:
     embedding_dimension: 384
   model_id: all-MiniLM-L6-v2

llama_stack/templates/meta-reference-gpu/run-with-safety.yaml

@@ -16,11 +16,12 @@ providers:
     provider_type: inline::meta-reference
     config:
       model: ${env.INFERENCE_MODEL}
-      max_seq_len: 4096
       checkpoint_dir: ${env.INFERENCE_CHECKPOINT_DIR:null}
       quantization:
         type: ${env.QUANTIZATION_TYPE:bf16}
       model_parallel_size: ${env.MODEL_PARALLEL_SIZE:0}
+      max_batch_size: ${env.MAX_BATCH_SIZE:1}
+      max_seq_len: ${env.MAX_SEQ_LEN:4096}
   - provider_id: sentence-transformers
     provider_type: inline::sentence-transformers
     config: {}
@@ -28,11 +29,12 @@ providers:
     provider_type: inline::meta-reference
     config:
       model: ${env.SAFETY_MODEL}
-      max_seq_len: 4096
       checkpoint_dir: ${env.SAFETY_CHECKPOINT_DIR:null}
       quantization:
         type: ${env.QUANTIZATION_TYPE:bf16}
       model_parallel_size: ${env.MODEL_PARALLEL_SIZE:0}
+      max_batch_size: ${env.MAX_BATCH_SIZE:1}
+      max_seq_len: ${env.MAX_SEQ_LEN:4096}
   vector_io:
   - provider_id: faiss
     provider_type: inline::faiss

llama_stack/templates/meta-reference-gpu/run.yaml

@@ -16,11 +16,12 @@ providers:
     provider_type: inline::meta-reference
     config:
       model: ${env.INFERENCE_MODEL}
-      max_seq_len: 4096
       checkpoint_dir: ${env.INFERENCE_CHECKPOINT_DIR:null}
       quantization:
         type: ${env.QUANTIZATION_TYPE:bf16}
       model_parallel_size: ${env.MODEL_PARALLEL_SIZE:0}
+      max_batch_size: ${env.MAX_BATCH_SIZE:1}
+      max_seq_len: ${env.MAX_SEQ_LEN:4096}
   - provider_id: sentence-transformers
     provider_type: inline::sentence-transformers
     config: {}

llama_stack/templates/remote-vllm/doc_template.md

@@ -13,7 +13,7 @@ The `llamastack/distribution-{{ name }}` distribution consists of the following
 
 {{ providers_table }}
 
-You can use this distribution if you have GPUs and want to run an independent vLLM server container for running inference.
+You can use this distribution if you want to run an independent vLLM server for inference.
 
 {% if run_config_env_vars %}
 ### Environment Variables
@@ -28,7 +28,10 @@ The following environment variables can be configured:
 
 ## Setting up vLLM server
 
-Both AMD and NVIDIA GPUs can serve as accelerators for the vLLM server, which acts as both the LLM inference provider and the safety provider.
+In the following sections, we'll use either AMD and NVIDIA GPUs to serve as hardware accelerators for the vLLM
+server, which acts as both the LLM inference provider and the safety provider. Note that vLLM also
+[supports many other hardware accelerators](https://docs.vllm.ai/en/latest/getting_started/installation.html) and
+that we only use GPUs here for demonstration purposes.
 
 ### Setting up vLLM server on AMD GPU
 

llama_stack/templates/verification/run.yaml

@@ -474,6 +474,16 @@ models:
   provider_id: groq-openai-compat
   provider_model_id: groq/llama-4-scout-17b-16e-instruct
   model_type: llm
+- metadata: {}
+  model_id: groq/meta-llama/llama-4-scout-17b-16e-instruct
+  provider_id: groq-openai-compat
+  provider_model_id: groq/meta-llama/llama-4-scout-17b-16e-instruct
+  model_type: llm
+- metadata: {}
+  model_id: meta-llama/Llama-4-Scout-17B-16E-Instruct
+  provider_id: groq-openai-compat
+  provider_model_id: groq/meta-llama/llama-4-scout-17b-16e-instruct
+  model_type: llm
 - metadata: {}
   model_id: groq/llama-4-maverick-17b-128e-instruct
   provider_id: groq-openai-compat
@@ -484,6 +494,16 @@ models:
   provider_id: groq-openai-compat
   provider_model_id: groq/llama-4-maverick-17b-128e-instruct
   model_type: llm
+- metadata: {}
+  model_id: groq/meta-llama/llama-4-maverick-17b-128e-instruct
+  provider_id: groq-openai-compat
+  provider_model_id: groq/meta-llama/llama-4-maverick-17b-128e-instruct
+  model_type: llm
+- metadata: {}
+  model_id: meta-llama/Llama-4-Maverick-17B-128E-Instruct
+  provider_id: groq-openai-compat
+  provider_model_id: groq/meta-llama/llama-4-maverick-17b-128e-instruct
+  model_type: llm
 - metadata: {}
   model_id: Meta-Llama-3.1-8B-Instruct
   provider_id: sambanova-openai-compat