Merge branch 'main' into opengauss-add

llama_stack/providers/datatypes.py

@@ -43,14 +43,30 @@ class ModelsProtocolPrivate(Protocol):
        -> Provider uses provider-model-id for inference
     """
 
+    # this should be called `on_model_register` or something like that.
+    # the provider should _not_ be able to change the object in this
+    # callback
     async def register_model(self, model: Model) -> Model: ...
 
     async def unregister_model(self, model_id: str) -> None: ...
 
+    # the Stack router will query each provider for their list of models
+    # if a `refresh_interval_seconds` is provided, this method will be called
+    # periodically to refresh the list of models
+    #
+    # NOTE: each model returned will be registered with the model registry. this means
+    # a callback to the `register_model()` method will be made. this is duplicative and
+    # may be removed in the future.
+    async def list_models(self) -> list[Model] | None: ...
+
+    async def should_refresh_models(self) -> bool: ...
+
 
 class ShieldsProtocolPrivate(Protocol):
     async def register_shield(self, shield: Shield) -> None: ...
 
+    async def unregister_shield(self, identifier: str) -> None: ...
+
 
 class VectorDBsProtocolPrivate(Protocol):
     async def register_vector_db(self, vector_db: VectorDB) -> None: ...
@@ -104,6 +120,19 @@ class ProviderSpec(BaseModel):
         description="If this provider is deprecated and does NOT work, specify the error message here",
     )
 
+    module: str | None = Field(
+        default=None,
+        description="""
+ Fully-qualified name of the module to import. The module is expected to have:
+
+  - `get_adapter_impl(config, deps)`: returns the adapter implementation
+
+  Example: `module: ramalama_stack`
+ """,
+    )
+
+    is_external: bool = Field(default=False, description="Notes whether this provider is an external provider.")
+
     # used internally by the resolver; this is a hack for now
     deps__: list[str] = Field(default_factory=list)
 
@@ -113,7 +142,7 @@ class ProviderSpec(BaseModel):
 
 
 class RoutingTable(Protocol):
-    def get_provider_impl(self, routing_key: str) -> Any: ...
+    async def get_provider_impl(self, routing_key: str) -> Any: ...
 
 
 # TODO: this can now be inlined into RemoteProviderSpec
@@ -124,7 +153,7 @@ class AdapterSpec(BaseModel):
         description="Unique identifier for this adapter",
     )
     module: str = Field(
-        ...,
+        default_factory=str,
         description="""
 Fully-qualified name of the module to import. The module is expected to have:
 
@@ -162,14 +191,7 @@ The container image to use for this implementation. If one is provided, pip_pack
 If a provider depends on other providers, the dependencies MUST NOT specify a container image.
 """,
     )
-    module: str = Field(
-        ...,
-        description="""
-Fully-qualified name of the module to import. The module is expected to have:
-
- - `get_provider_impl(config, deps)`: returns the local implementation
-""",
-    )
+    # module field is inherited from ProviderSpec
     provider_data_validator: str | None = Field(
         default=None,
     )
@@ -212,9 +234,7 @@ API responses, specify the adapter here.
     def container_image(self) -> str | None:
         return None
 
-    @property
-    def module(self) -> str:
-        return self.adapter.module
+    # module field is inherited from ProviderSpec
 
     @property
     def pip_packages(self) -> list[str]:
@@ -226,14 +246,19 @@ API responses, specify the adapter here.
 
 
 def remote_provider_spec(
-    api: Api, adapter: AdapterSpec, api_dependencies: list[Api] | None = None
+    api: Api,
+    adapter: AdapterSpec,
+    api_dependencies: list[Api] | None = None,
+    optional_api_dependencies: list[Api] | None = None,
 ) -> RemoteProviderSpec:
     return RemoteProviderSpec(
         api=api,
         provider_type=f"remote::{adapter.adapter_type}",
         config_class=adapter.config_class,
+        module=adapter.module,
         adapter=adapter,
         api_dependencies=api_dependencies or [],
+        optional_api_dependencies=optional_api_dependencies or [],
     )
 
 

llama_stack/providers/inline/agents/meta_reference/__init__.py

@@ -6,7 +6,7 @@
 
 from typing import Any
 
-from llama_stack.distribution.datatypes import AccessRule, Api
+from llama_stack.core.datatypes import AccessRule, Api
 
 from .config import MetaReferenceAgentsImplConfig
 

llama_stack/providers/inline/agents/meta_reference/agent_instance.py

@@ -10,6 +10,7 @@ import re
 import secrets
 import string
 import uuid
+import warnings
 from collections.abc import AsyncGenerator
 from datetime import UTC, datetime
 
@@ -43,6 +44,7 @@ from llama_stack.apis.common.content_types import (
     ToolCallDelta,
     ToolCallParseStatus,
 )
+from llama_stack.apis.common.errors import SessionNotFoundError
 from llama_stack.apis.inference import (
     ChatCompletionResponseEventType,
     CompletionMessage,
@@ -60,7 +62,7 @@ from llama_stack.apis.inference import (
 from llama_stack.apis.safety import Safety
 from llama_stack.apis.tools import ToolGroups, ToolInvocationResult, ToolRuntime
 from llama_stack.apis.vector_io import VectorIO
-from llama_stack.distribution.datatypes import AccessRule
+from llama_stack.core.datatypes import AccessRule
 from llama_stack.log import get_logger
 from llama_stack.models.llama.datatypes import (
     BuiltinTool,
@@ -213,7 +215,7 @@ class ChatAgent(ShieldRunnerMixin):
         is_resume = isinstance(request, AgentTurnResumeRequest)
         session_info = await self.storage.get_session_info(request.session_id)
         if session_info is None:
-            raise ValueError(f"Session {request.session_id} not found")
+            raise SessionNotFoundError(request.session_id)
 
         turns = await self.storage.get_session_turns(request.session_id)
         if is_resume and len(turns) == 0:
@@ -911,8 +913,16 @@ async def load_data_from_url(url: str) -> str:
 
 
 async def get_raw_document_text(document: Document) -> str:
-    if not document.mime_type.startswith("text/"):
+    # Handle deprecated text/yaml mime type with warning
+    if document.mime_type == "text/yaml":
+        warnings.warn(
+            "The 'text/yaml' MIME type is deprecated. Please use 'application/yaml' instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+    elif not (document.mime_type.startswith("text/") or document.mime_type == "application/yaml"):
         raise ValueError(f"Unexpected document mime type: {document.mime_type}")
+
     if isinstance(document.content, URL):
         return await load_data_from_url(document.content.uri)
     elif isinstance(document.content, str):

llama_stack/providers/inline/agents/meta_reference/agents.py

@@ -41,7 +41,7 @@ from llama_stack.apis.inference import (
 from llama_stack.apis.safety import Safety
 from llama_stack.apis.tools import ToolGroups, ToolRuntime
 from llama_stack.apis.vector_io import VectorIO
-from llama_stack.distribution.datatypes import AccessRule
+from llama_stack.core.datatypes import AccessRule
 from llama_stack.providers.utils.kvstore import InmemoryKVStoreImpl, kvstore_impl
 from llama_stack.providers.utils.pagination import paginate_records
 from llama_stack.providers.utils.responses.responses_store import ResponsesStore
@@ -230,8 +230,6 @@ class MetaReferenceAgentsImpl(Agents):
         agent = await self._get_agent_impl(agent_id)
 
         session_info = await agent.storage.get_session_info(session_id)
-        if session_info is None:
-            raise ValueError(f"Session {session_id} not found")
         turns = await agent.storage.get_session_turns(session_id)
         if turn_ids:
             turns = [turn for turn in turns if turn.turn_id in turn_ids]
@@ -244,9 +242,6 @@ class MetaReferenceAgentsImpl(Agents):
 
     async def delete_agents_session(self, agent_id: str, session_id: str) -> None:
         agent = await self._get_agent_impl(agent_id)
-        session_info = await agent.storage.get_session_info(session_id)
-        if session_info is None:
-            raise ValueError(f"Session {session_id} not found")
 
         # Delete turns first, then the session
         await agent.storage.delete_session_turns(session_id)

llama_stack/providers/inline/agents/meta_reference/persistence.py

@@ -10,10 +10,11 @@ import uuid
 from datetime import UTC, datetime
 
 from llama_stack.apis.agents import AgentConfig, Session, ToolExecutionStep, Turn
-from llama_stack.distribution.access_control.access_control import AccessDeniedError, is_action_allowed
-from llama_stack.distribution.access_control.datatypes import AccessRule
-from llama_stack.distribution.datatypes import User
-from llama_stack.distribution.request_headers import get_authenticated_user
+from llama_stack.apis.common.errors import SessionNotFoundError
+from llama_stack.core.access_control.access_control import AccessDeniedError, is_action_allowed
+from llama_stack.core.access_control.datatypes import AccessRule
+from llama_stack.core.datatypes import User
+from llama_stack.core.request_headers import get_authenticated_user
 from llama_stack.providers.utils.kvstore import KVStore
 
 log = logging.getLogger(__name__)
@@ -61,12 +62,12 @@ class AgentPersistence:
         )
         return session_id
 
-    async def get_session_info(self, session_id: str) -> AgentSessionInfo | None:
+    async def get_session_info(self, session_id: str) -> AgentSessionInfo:
         value = await self.kvstore.get(
             key=f"session:{self.agent_id}:{session_id}",
         )
         if not value:
-            return None
+            raise SessionNotFoundError(session_id)
 
         session_info = AgentSessionInfo(**json.loads(value))
 
@@ -95,7 +96,7 @@ class AgentPersistence:
     async def add_vector_db_to_session(self, session_id: str, vector_db_id: str):
         session_info = await self.get_session_if_accessible(session_id)
         if session_info is None:
-            raise ValueError(f"Session {session_id} not found or access denied")
+            raise SessionNotFoundError(session_id)
 
         session_info.vector_db_id = vector_db_id
         await self.kvstore.set(
@@ -105,7 +106,7 @@ class AgentPersistence:
 
     async def add_turn_to_session(self, session_id: str, turn: Turn):
         if not await self.get_session_if_accessible(session_id):
-            raise ValueError(f"Session {session_id} not found or access denied")
+            raise SessionNotFoundError(session_id)
 
         await self.kvstore.set(
             key=f"session:{self.agent_id}:{session_id}:{turn.turn_id}",
@@ -114,7 +115,7 @@ class AgentPersistence:
 
     async def get_session_turns(self, session_id: str) -> list[Turn]:
         if not await self.get_session_if_accessible(session_id):
-            raise ValueError(f"Session {session_id} not found or access denied")
+            raise SessionNotFoundError(session_id)
 
         values = await self.kvstore.values_in_range(
             start_key=f"session:{self.agent_id}:{session_id}:",
@@ -128,11 +129,16 @@ class AgentPersistence:
             except Exception as e:
                 log.error(f"Error parsing turn: {e}")
                 continue
+
+        # The kvstore does not guarantee order, so we sort by started_at
+        # to ensure consistent ordering of turns.
+        turns.sort(key=lambda t: t.started_at)
+
         return turns
 
     async def get_session_turn(self, session_id: str, turn_id: str) -> Turn | None:
         if not await self.get_session_if_accessible(session_id):
-            raise ValueError(f"Session {session_id} not found or access denied")
+            raise SessionNotFoundError(session_id)
 
         value = await self.kvstore.get(
             key=f"session:{self.agent_id}:{session_id}:{turn_id}",
@@ -143,7 +149,7 @@ class AgentPersistence:
 
     async def set_in_progress_tool_call_step(self, session_id: str, turn_id: str, step: ToolExecutionStep):
         if not await self.get_session_if_accessible(session_id):
-            raise ValueError(f"Session {session_id} not found or access denied")
+            raise SessionNotFoundError(session_id)
 
         await self.kvstore.set(
             key=f"in_progress_tool_call_step:{self.agent_id}:{session_id}:{turn_id}",
@@ -161,7 +167,7 @@ class AgentPersistence:
 
     async def set_num_infer_iters_in_turn(self, session_id: str, turn_id: str, num_infer_iters: int):
         if not await self.get_session_if_accessible(session_id):
-            raise ValueError(f"Session {session_id} not found or access denied")
+            raise SessionNotFoundError(session_id)
 
         await self.kvstore.set(
             key=f"num_infer_iters_in_turn:{self.agent_id}:{session_id}:{turn_id}",
@@ -213,6 +219,6 @@ class AgentPersistence:
         """
         session_info = await self.get_session_info(session_id)
         if session_info is None:
-            raise ValueError(f"Session {session_id} not found")
+            raise SessionNotFoundError(session_id)
 
         await self.kvstore.delete(key=f"session:{self.agent_id}:{session_id}")

llama_stack/providers/inline/datasetio/localfs/datasetio.py

@@ -5,8 +5,6 @@
 # the root directory of this source tree.
 from typing import Any
 
-import pandas
-
 from llama_stack.apis.common.responses import PaginatedResponse
 from llama_stack.apis.datasetio import DatasetIO
 from llama_stack.apis.datasets import Dataset
@@ -44,6 +42,8 @@ class PandasDataframeDataset:
         if self.dataset_def.source.type == "uri":
             self.df = await get_dataframe_from_uri(self.dataset_def.source.uri)
         elif self.dataset_def.source.type == "rows":
+            import pandas
+
             self.df = pandas.DataFrame(self.dataset_def.source.rows)
         else:
             raise ValueError(f"Unsupported dataset source type: {self.dataset_def.source.type}")
@@ -103,6 +103,8 @@ class LocalFSDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
         return paginate_records(records, start_index, limit)
 
     async def append_rows(self, dataset_id: str, rows: list[dict[str, Any]]) -> None:
+        import pandas
+
         dataset_def = self.dataset_infos[dataset_id]
         dataset_impl = PandasDataframeDataset(dataset_def)
         await dataset_impl.load()

llama_stack/providers/inline/eval/meta_reference/__init__.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 from typing import Any
 
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import Api
 
 from .config import MetaReferenceEvalConfig
 

llama_stack/providers/inline/files/localfs/__init__.py

@@ -6,7 +6,7 @@
 
 from typing import Any
 
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import AccessRule, Api
 
 from .config import LocalfsFilesImplConfig
 from .files import LocalfsFilesImpl
@@ -14,7 +14,7 @@ from .files import LocalfsFilesImpl
 __all__ = ["LocalfsFilesImpl", "LocalfsFilesImplConfig"]
 
 
-async def get_provider_impl(config: LocalfsFilesImplConfig, deps: dict[Api, Any]):
-    impl = LocalfsFilesImpl(config)
+async def get_provider_impl(config: LocalfsFilesImplConfig, deps: dict[Api, Any], policy: list[AccessRule]):
+    impl = LocalfsFilesImpl(config, policy)
     await impl.initialize()
     return impl

llama_stack/providers/inline/files/localfs/files.py

@@ -19,16 +19,19 @@ from llama_stack.apis.files import (
     OpenAIFileObject,
     OpenAIFilePurpose,
 )
+from llama_stack.core.datatypes import AccessRule
 from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
-from llama_stack.providers.utils.sqlstore.sqlstore import SqlStore, sqlstore_impl
+from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl
 
 from .config import LocalfsFilesImplConfig
 
 
 class LocalfsFilesImpl(Files):
-    def __init__(self, config: LocalfsFilesImplConfig) -> None:
+    def __init__(self, config: LocalfsFilesImplConfig, policy: list[AccessRule]) -> None:
         self.config = config
-        self.sql_store: SqlStore | None = None
+        self.policy = policy
+        self.sql_store: AuthorizedSqlStore | None = None
 
     async def initialize(self) -> None:
         """Initialize the files provider by setting up storage directory and metadata database."""
@@ -37,7 +40,7 @@ class LocalfsFilesImpl(Files):
         storage_path.mkdir(parents=True, exist_ok=True)
 
         # Initialize SQL store for metadata
-        self.sql_store = sqlstore_impl(self.config.metadata_store)
+        self.sql_store = AuthorizedSqlStore(sqlstore_impl(self.config.metadata_store))
         await self.sql_store.create_table(
             "openai_files",
             {
@@ -51,6 +54,9 @@ class LocalfsFilesImpl(Files):
             },
         )
 
+    async def shutdown(self) -> None:
+        pass
+
     def _generate_file_id(self) -> str:
         """Generate a unique file ID for OpenAI API."""
         return f"file-{uuid.uuid4().hex}"
@@ -123,6 +129,7 @@ class LocalfsFilesImpl(Files):
 
         paginated_result = await self.sql_store.fetch_all(
             table="openai_files",
+            policy=self.policy,
             where=where_conditions if where_conditions else None,
             order_by=[("created_at", order.value)],
             cursor=("id", after) if after else None,
@@ -153,7 +160,7 @@ class LocalfsFilesImpl(Files):
         if not self.sql_store:
             raise RuntimeError("Files provider not initialized")
 
-        row = await self.sql_store.fetch_one("openai_files", where={"id": file_id})
+        row = await self.sql_store.fetch_one("openai_files", policy=self.policy, where={"id": file_id})
         if not row:
             raise ValueError(f"File with id {file_id} not found")
 
@@ -171,7 +178,7 @@ class LocalfsFilesImpl(Files):
         if not self.sql_store:
             raise RuntimeError("Files provider not initialized")
 
-        row = await self.sql_store.fetch_one("openai_files", where={"id": file_id})
+        row = await self.sql_store.fetch_one("openai_files", policy=self.policy, where={"id": file_id})
         if not row:
             raise ValueError(f"File with id {file_id} not found")
 
@@ -194,7 +201,7 @@ class LocalfsFilesImpl(Files):
             raise RuntimeError("Files provider not initialized")
 
         # Get file metadata
-        row = await self.sql_store.fetch_one("openai_files", where={"id": file_id})
+        row = await self.sql_store.fetch_one("openai_files", policy=self.policy, where={"id": file_id})
         if not row:
             raise ValueError(f"File with id {file_id} not found")
 

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

@@ -6,7 +6,7 @@
 
 from pathlib import Path
 
-from llama_stack.distribution.utils.model_utils import model_local_dir
+from llama_stack.core.utils.model_utils import model_local_dir
 
 
 def model_checkpoint_dir(model_id) -> str:

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

@@ -102,6 +102,12 @@ class MetaReferenceInferenceImpl(
         if self.config.create_distributed_process_group:
             self.generator.stop()
 
+    async def should_refresh_models(self) -> bool:
+        return False
+
+    async def list_models(self) -> list[Model] | None:
+        return None
+
     async def unregister_model(self, model_id: str) -> None:
         pass
 

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

@@ -98,7 +98,7 @@ class ProcessingMessageWrapper(BaseModel):
 
 
 def mp_rank_0() -> bool:
-    return get_model_parallel_rank() == 0
+    return bool(get_model_parallel_rank() == 0)
 
 
 def encode_msg(msg: ProcessingMessage) -> bytes:
@@ -125,7 +125,7 @@ def retrieve_requests(reply_socket_url: str):
         reply_socket.send_multipart([client_id, encode_msg(obj)])
 
     while True:
-        tasks = [None]
+        tasks: list[ProcessingMessage | None] = [None]
         if mp_rank_0():
             client_id, maybe_task_json = maybe_get_work(reply_socket)
             if maybe_task_json is not None:
@@ -152,7 +152,7 @@ def retrieve_requests(reply_socket_url: str):
                     break
 
                 for obj in out:
-                    updates = [None]
+                    updates: list[ProcessingMessage | None] = [None]
                     if mp_rank_0():
                         _, update_json = maybe_get_work(reply_socket)
                         update = maybe_parse_message(update_json)

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

@@ -20,6 +20,7 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
+from llama_stack.apis.models import ModelType
 from llama_stack.providers.datatypes import Model, ModelsProtocolPrivate
 from llama_stack.providers.utils.inference.embedding_mixin import (
     SentenceTransformerEmbeddingMixin,
@@ -41,6 +42,8 @@ class SentenceTransformersInferenceImpl(
     InferenceProvider,
     ModelsProtocolPrivate,
 ):
+    __provider_id__: str
+
     def __init__(self, config: SentenceTransformersInferenceConfig) -> None:
         self.config = config
 
@@ -50,6 +53,22 @@ class SentenceTransformersInferenceImpl(
     async def shutdown(self) -> None:
         pass
 
+    async def should_refresh_models(self) -> bool:
+        return False
+
+    async def list_models(self) -> list[Model] | None:
+        return [
+            Model(
+                identifier="all-MiniLM-L6-v2",
+                provider_resource_id="all-MiniLM-L6-v2",
+                provider_id=self.__provider_id__,
+                metadata={
+                    "embedding_dimension": 384,
+                },
+                model_type=ModelType.embedding,
+            ),
+        ]
+
     async def register_model(self, model: Model) -> Model:
         return model
 

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

@@ -1,17 +0,0 @@
-# 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
-
-from .config import VLLMConfig
-
-
-async def get_provider_impl(config: VLLMConfig, _deps: dict[str, Any]):
-    from .vllm import VLLMInferenceImpl
-
-    impl = VLLMInferenceImpl(config)
-    await impl.initialize()
-    return impl

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

@@ -1,53 +0,0 @@
-# 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
-
-from pydantic import BaseModel, Field
-
-from llama_stack.schema_utils import json_schema_type
-
-
-@json_schema_type
-class VLLMConfig(BaseModel):
-    """Configuration for the vLLM inference provider.
-
-    Note that the model name is no longer part of this static configuration.
-    You can bind an instance of this provider to a specific model with the
-    ``models.register()`` API call."""
-
-    tensor_parallel_size: int = Field(
-        default=1,
-        description="Number of tensor parallel replicas (number of GPUs to use).",
-    )
-    max_tokens: int = Field(
-        default=4096,
-        description="Maximum number of tokens to generate.",
-    )
-    max_model_len: int = Field(default=4096, description="Maximum context length to use during serving.")
-    max_num_seqs: int = Field(default=4, description="Maximum parallel batch size for generation.")
-    enforce_eager: bool = Field(
-        default=False,
-        description="Whether to use eager mode for inference (otherwise cuda graphs are used).",
-    )
-    gpu_memory_utilization: float = Field(
-        default=0.3,
-        description=(
-            "How much GPU memory will be allocated when this provider has finished "
-            "loading, including memory that was already allocated before loading."
-        ),
-    )
-
-    @classmethod
-    def sample_run_config(cls, **kwargs: Any) -> dict[str, Any]:
-        return {
-            "tensor_parallel_size": "${env.TENSOR_PARALLEL_SIZE:=1}",
-            "max_tokens": "${env.MAX_TOKENS:=4096}",
-            "max_model_len": "${env.MAX_MODEL_LEN:=4096}",
-            "max_num_seqs": "${env.MAX_NUM_SEQS:=4}",
-            "enforce_eager": "${env.ENFORCE_EAGER:=False}",
-            "gpu_memory_utilization": "${env.GPU_MEMORY_UTILIZATION:=0.3}",
-        }

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

@@ -1,170 +0,0 @@
-# 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 vllm
-
-from llama_stack.apis.inference import (
-    ChatCompletionRequest,
-    GrammarResponseFormat,
-    JsonSchemaResponseFormat,
-    Message,
-    ToolChoice,
-    ToolDefinition,
-    UserMessage,
-)
-from llama_stack.models.llama.datatypes import BuiltinTool
-from llama_stack.providers.utils.inference.openai_compat import (
-    convert_message_to_openai_dict,
-    get_sampling_options,
-)
-
-###############################################################################
-# This file contains OpenAI compatibility code that is currently only used
-# by the inline vLLM connector. Some or all of this code may be moved to a
-# central location at a later date.
-
-
-def _merge_context_into_content(message: Message) -> Message:  # type: ignore
-    """
-    Merge the ``context`` field of a Llama Stack ``Message`` object into
-    the content field for compabilitiy with OpenAI-style APIs.
-
-    Generates a content string that emulates the current behavior
-    of ``llama_models.llama3.api.chat_format.encode_message()``.
-
-    :param message: Message that may include ``context`` field
-
-    :returns: A version of ``message`` with any context merged into the
-     ``content`` field.
-    """
-    if not isinstance(message, UserMessage):  # Separate type check for linter
-        return message
-    if message.context is None:
-        return message
-    return UserMessage(
-        role=message.role,
-        # Emumate llama_models.llama3.api.chat_format.encode_message()
-        content=message.content + "\n\n" + message.context,
-        context=None,
-    )
-
-
-def _llama_stack_tools_to_openai_tools(
-    tools: list[ToolDefinition] | None = None,
-) -> list[vllm.entrypoints.openai.protocol.ChatCompletionToolsParam]:
-    """
-    Convert the list of available tools from Llama Stack's format to vLLM's
-    version of OpenAI's format.
-    """
-    if tools is None:
-        return []
-
-    result = []
-    for t in tools:
-        if isinstance(t.tool_name, BuiltinTool):
-            raise NotImplementedError("Built-in tools not yet implemented")
-        if t.parameters is None:
-            parameters = None
-        else:  # if t.parameters is not None
-            # Convert the "required" flags to a list of required params
-            required_params = [k for k, v in t.parameters.items() if v.required]
-            parameters = {
-                "type": "object",  # Mystery value that shows up in OpenAI docs
-                "properties": {
-                    k: {"type": v.param_type, "description": v.description} for k, v in t.parameters.items()
-                },
-                "required": required_params,
-            }
-
-        function_def = vllm.entrypoints.openai.protocol.FunctionDefinition(
-            name=t.tool_name, description=t.description, parameters=parameters
-        )
-
-        # Every tool definition is double-boxed in a ChatCompletionToolsParam
-        result.append(vllm.entrypoints.openai.protocol.ChatCompletionToolsParam(function=function_def))
-    return result
-
-
-async def llama_stack_chat_completion_to_openai_chat_completion_dict(
-    request: ChatCompletionRequest,
-) -> dict:
-    """
-    Convert a chat completion request in Llama Stack format into an
-    equivalent set of arguments to pass to an OpenAI-compatible
-    chat completions API.
-
-    :param request: Bundled request parameters in Llama Stack format.
-
-    :returns: Dictionary of key-value pairs to use as an initializer
-     for a dataclass or to be converted directly to JSON and sent
-     over the wire.
-    """
-
-    converted_messages = [
-        # This mystery async call makes the parent function also be async
-        await convert_message_to_openai_dict(_merge_context_into_content(m), download=True)
-        for m in request.messages
-    ]
-    converted_tools = _llama_stack_tools_to_openai_tools(request.tools)
-
-    # Llama will try to use built-in tools with no tool catalog, so don't enable
-    # tool choice unless at least one tool is enabled.
-    converted_tool_choice = "none"
-    if (
-        request.tool_config is not None
-        and request.tool_config.tool_choice == ToolChoice.auto
-        and request.tools is not None
-        and len(request.tools) > 0
-    ):
-        converted_tool_choice = "auto"
-
-    # TODO: Figure out what to do with the tool_prompt_format argument.
-    #  Other connectors appear to drop it quietly.
-
-    # Use Llama Stack shared code to translate sampling parameters.
-    sampling_options = get_sampling_options(request.sampling_params)
-
-    # get_sampling_options() translates repetition penalties to an option that
-    # OpenAI's APIs don't know about.
-    # vLLM's OpenAI-compatible API also handles repetition penalties wrong.
-    # For now, translate repetition penalties into a format that vLLM's broken
-    # API will handle correctly. Two wrongs make a right...
-    if "repeat_penalty" in sampling_options:
-        del sampling_options["repeat_penalty"]
-    if request.sampling_params.repetition_penalty is not None and request.sampling_params.repetition_penalty != 1.0:
-        sampling_options["repetition_penalty"] = request.sampling_params.repetition_penalty
-
-    # Convert a single response format into four different parameters, per
-    # the OpenAI spec
-    guided_decoding_options = dict()
-    if request.response_format is None:
-        # Use defaults
-        pass
-    elif isinstance(request.response_format, JsonSchemaResponseFormat):
-        guided_decoding_options["guided_json"] = request.response_format.json_schema
-    elif isinstance(request.response_format, GrammarResponseFormat):
-        guided_decoding_options["guided_grammar"] = request.response_format.bnf
-    else:
-        raise TypeError(f"ResponseFormat object is of unexpected subtype '{type(request.response_format)}'")
-
-    logprob_options = dict()
-    if request.logprobs is not None:
-        logprob_options["logprobs"] = request.logprobs.top_k
-
-    # Marshall together all the arguments for a ChatCompletionRequest
-    request_options = {
-        "model": request.model,
-        "messages": converted_messages,
-        "tools": converted_tools,
-        "tool_choice": converted_tool_choice,
-        "stream": request.stream,
-        **sampling_options,
-        **guided_decoding_options,
-        **logprob_options,
-    }
-
-    return request_options

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

@@ -1,811 +0,0 @@
-# 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 json
-import re
-import uuid
-from collections.abc import AsyncGenerator, AsyncIterator
-
-# These vLLM modules contain names that overlap with Llama Stack names, so we import
-# fully-qualified names
-import vllm.entrypoints.openai.protocol
-import vllm.sampling_params
-from vllm.engine.arg_utils import AsyncEngineArgs
-from vllm.engine.async_llm_engine import AsyncLLMEngine
-from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
-from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels
-
-from llama_stack.apis.common.content_types import (
-    InterleavedContent,
-    InterleavedContentItem,
-    TextDelta,
-    ToolCallDelta,
-)
-from llama_stack.apis.inference import (
-    ChatCompletionRequest,
-    ChatCompletionResponse,
-    ChatCompletionResponseEvent,
-    ChatCompletionResponseEventType,
-    ChatCompletionResponseStreamChunk,
-    CompletionMessage,
-    CompletionResponse,
-    CompletionResponseStreamChunk,
-    EmbeddingsResponse,
-    EmbeddingTaskType,
-    GrammarResponseFormat,
-    Inference,
-    JsonSchemaResponseFormat,
-    LogProbConfig,
-    Message,
-    OpenAIEmbeddingsResponse,
-    ResponseFormat,
-    SamplingParams,
-    TextTruncation,
-    TokenLogProbs,
-    ToolChoice,
-    ToolConfig,
-    TopKSamplingStrategy,
-    TopPSamplingStrategy,
-)
-from llama_stack.apis.models import Model
-from llama_stack.log import get_logger
-from llama_stack.models.llama import sku_list
-from llama_stack.models.llama.datatypes import (
-    StopReason,
-    ToolCall,
-    ToolDefinition,
-    ToolPromptFormat,
-)
-from llama_stack.models.llama.llama3.chat_format import ChatFormat
-from llama_stack.models.llama.llama3.tokenizer import Tokenizer
-from llama_stack.providers.remote.inference.vllm.vllm import build_hf_repo_model_entries
-from llama_stack.providers.utils.inference.model_registry import (
-    ModelRegistryHelper,
-    ModelsProtocolPrivate,
-)
-from llama_stack.providers.utils.inference.openai_compat import (
-    OpenAIChatCompletionToLlamaStackMixin,
-    OpenAICompatCompletionChoice,
-    OpenAICompatCompletionResponse,
-    OpenAICompletionToLlamaStackMixin,
-    get_stop_reason,
-    process_chat_completion_stream_response,
-)
-from llama_stack.providers.utils.inference.prompt_adapter import (
-    chat_completion_request_to_prompt,
-)
-
-from .config import VLLMConfig
-from .openai_utils import llama_stack_chat_completion_to_openai_chat_completion_dict
-
-# Map from Hugging Face model architecture name to appropriate tool parser.
-# See vllm.entrypoints.openai.tool_parsers.ToolParserManager.tool_parsers for the full list of
-# available parsers.
-# TODO: Expand this list
-CONFIG_TYPE_TO_TOOL_PARSER = {
-    "GraniteConfig": "granite",
-    "MllamaConfig": "llama3_json",
-    "LlamaConfig": "llama3_json",
-}
-DEFAULT_TOOL_PARSER = "pythonic"
-
-
-logger = get_logger(__name__, category="inference")
-
-
-def _random_uuid_str() -> str:
-    return str(uuid.uuid4().hex)
-
-
-def _response_format_to_guided_decoding_params(
-    response_format: ResponseFormat | None,  # type: ignore
-) -> vllm.sampling_params.GuidedDecodingParams:
-    """
-    Translate constrained decoding parameters from Llama Stack's format to vLLM's format.
-
-    :param response_format: Llama Stack version of constrained decoding info. Can be ``None``,
-     indicating no constraints.
-    :returns: The equivalent dataclass object for the low-level inference layer of vLLM.
-    """
-    if response_format is None:
-        # As of vLLM 0.6.3, the default constructor for GuidedDecodingParams() returns an invalid
-        # value that crashes the executor on some code paths. Use ``None`` instead.
-        return None
-
-    # Llama Stack currently implements fewer types of constrained decoding than vLLM does.
-    # Translate the types that exist and detect if Llama Stack adds new ones.
-    if isinstance(response_format, JsonSchemaResponseFormat):
-        return vllm.sampling_params.GuidedDecodingParams(json=response_format.json_schema)
-    elif isinstance(response_format, GrammarResponseFormat):
-        # BNF grammar.
-        # Llama Stack uses the parse tree of the grammar, while vLLM uses the string
-        # representation of the grammar.
-        raise TypeError(
-            "Constrained decoding with BNF grammars is not currently implemented, because the "
-            "reference implementation does not implement it."
-        )
-    else:
-        raise TypeError(f"ResponseFormat object is of unexpected subtype '{type(response_format)}'")
-
-
-def _convert_sampling_params(
-    sampling_params: SamplingParams | None,
-    response_format: ResponseFormat | None,  # type: ignore
-    log_prob_config: LogProbConfig | None,
-) -> vllm.SamplingParams:
-    """Convert sampling and constrained decoding configuration from Llama Stack's format to vLLM's
-    format."""
-    # In the absence of provided config values, use Llama Stack defaults as encoded in the Llama
-    # Stack dataclasses. These defaults are different from vLLM's defaults.
-    if sampling_params is None:
-        sampling_params = SamplingParams()
-    if log_prob_config is None:
-        log_prob_config = LogProbConfig()
-
-    if isinstance(sampling_params.strategy, TopKSamplingStrategy):
-        if sampling_params.strategy.top_k == 0:
-            # vLLM treats "k" differently for top-k sampling
-            vllm_top_k = -1
-        else:
-            vllm_top_k = sampling_params.strategy.top_k
-    else:
-        vllm_top_k = -1
-
-    if isinstance(sampling_params.strategy, TopPSamplingStrategy):
-        vllm_top_p = sampling_params.strategy.top_p
-        # Llama Stack only allows temperature with top-P.
-        vllm_temperature = sampling_params.strategy.temperature
-    else:
-        vllm_top_p = 1.0
-        vllm_temperature = 0.0
-
-    # vLLM allows top-p and top-k at the same time.
-    vllm_sampling_params = vllm.SamplingParams.from_optional(
-        max_tokens=(None if sampling_params.max_tokens == 0 else sampling_params.max_tokens),
-        temperature=vllm_temperature,
-        top_p=vllm_top_p,
-        top_k=vllm_top_k,
-        repetition_penalty=sampling_params.repetition_penalty,
-        guided_decoding=_response_format_to_guided_decoding_params(response_format),
-        logprobs=log_prob_config.top_k,
-    )
-    return vllm_sampling_params
-
-
-class VLLMInferenceImpl(
-    Inference,
-    OpenAIChatCompletionToLlamaStackMixin,
-    OpenAICompletionToLlamaStackMixin,
-    ModelsProtocolPrivate,
-):
-    """
-    vLLM-based inference model adapter for Llama Stack with support for multiple models.
-
-    Requires the configuration parameters documented in the :class:`VllmConfig2` class.
-    """
-
-    config: VLLMConfig
-    register_helper: ModelRegistryHelper
-    model_ids: set[str]
-    resolved_model_id: str | None
-    engine: AsyncLLMEngine | None
-    chat: OpenAIServingChat | None
-    is_meta_llama_model: bool
-
-    def __init__(self, config: VLLMConfig):
-        self.config = config
-        logger.info(f"Config is: {self.config}")
-
-        self.register_helper = ModelRegistryHelper(build_hf_repo_model_entries())
-        self.formatter = ChatFormat(Tokenizer.get_instance())
-
-        # The following are initialized when paths are bound to this provider
-        self.resolved_model_id = None
-        self.model_ids = set()
-        self.engine = None
-        self.chat = None
-        self.is_meta_llama_model = False
-
-    ###########################################################################
-    # METHODS INHERITED FROM IMPLICIT BASE CLASS.
-    # TODO: Make this class inherit from the new base class ProviderBase once that class exists.
-
-    async def initialize(self) -> None:
-        """
-        Callback that is invoked through many levels of indirection during provider class
-        instantiation, sometime after when __init__() is called and before any model registration
-        methods or methods connected to a REST API are called.
-
-        It's not clear what assumptions the class can make about the platform's initialization
-        state here that can't be made during __init__(), and vLLM can't be started until we know
-        what model it's supposed to be serving, so nothing happens here currently.
-        """
-        pass
-
-    async def shutdown(self) -> None:
-        logger.info(f"Shutting down inline vLLM inference provider {self}.")
-        if self.engine is not None:
-            self.engine.shutdown_background_loop()
-            self.engine = None
-            self.chat = None
-            self.model_ids = set()
-            self.resolved_model_id = None
-
-    ###########################################################################
-    # METHODS INHERITED FROM ModelsProtocolPrivate INTERFACE
-
-    # Note that the return type of the superclass method is WRONG
-    async def register_model(self, model: Model) -> Model:
-        """
-        Callback that is called when the server associates an inference endpoint with an
-        inference provider.
-
-        :param model: Object that encapsulates parameters necessary for identifying a specific
-         LLM.
-
-        :returns: The input ``Model`` object. It may or may not be permissible to change fields
-         before returning this object.
-        """
-        logger.debug(f"In register_model({model})")
-
-        # First attempt to interpret the model coordinates as a Llama model name
-        resolved_llama_model = sku_list.resolve_model(model.provider_model_id)
-        if resolved_llama_model is not None:
-            # Load from Hugging Face repo into default local cache dir
-            model_id_for_vllm = resolved_llama_model.huggingface_repo
-
-            # Detect a genuine Meta Llama model to trigger Meta-specific preprocessing.
-            # Don't set self.is_meta_llama_model until we actually load the model.
-            is_meta_llama_model = True
-        else:  # if resolved_llama_model is None
-            # Not a Llama model name. Pass the model id through to vLLM's loader
-            model_id_for_vllm = model.provider_model_id
-            is_meta_llama_model = False
-
-        if self.resolved_model_id is not None:
-            if model_id_for_vllm != self.resolved_model_id:
-                raise ValueError(
-                    f"Attempted to serve two LLMs (ids '{self.resolved_model_id}') and "
-                    f"'{model_id_for_vllm}') from one copy of provider '{self}'. Use multiple "
-                    f"copies of the provider instead."
-                )
-            else:
-                # Model already loaded
-                logger.info(
-                    f"Requested id {model} resolves to {model_id_for_vllm}, which is already loaded. Continuing."
-                )
-                self.model_ids.add(model.model_id)
-                return model
-
-        logger.info(f"Requested id {model} resolves to {model_id_for_vllm}. Loading {model_id_for_vllm}.")
-        if is_meta_llama_model:
-            logger.info(f"Model {model_id_for_vllm} is a Meta Llama model.")
-        self.is_meta_llama_model = is_meta_llama_model
-
-        # If we get here, this is the first time registering a model.
-        # Preload so that the first inference request won't time out.
-        engine_args = AsyncEngineArgs(
-            model=model_id_for_vllm,
-            tokenizer=model_id_for_vllm,
-            tensor_parallel_size=self.config.tensor_parallel_size,
-            enforce_eager=self.config.enforce_eager,
-            gpu_memory_utilization=self.config.gpu_memory_utilization,
-            max_num_seqs=self.config.max_num_seqs,
-            max_model_len=self.config.max_model_len,
-        )
-        self.engine = AsyncLLMEngine.from_engine_args(engine_args)
-
-        # vLLM currently requires the user to specify the tool parser manually. To choose a tool
-        # parser, we need to determine what model architecture is being used. For now, we infer
-        # that information from what config class the model uses.
-        low_level_model_config = self.engine.engine.get_model_config()
-        hf_config = low_level_model_config.hf_config
-        hf_config_class_name = hf_config.__class__.__name__
-        if hf_config_class_name in CONFIG_TYPE_TO_TOOL_PARSER:
-            tool_parser = CONFIG_TYPE_TO_TOOL_PARSER[hf_config_class_name]
-        else:
-            # No info -- choose a default so we can at least attempt tool
-            # use.
-            tool_parser = DEFAULT_TOOL_PARSER
-        logger.debug(f"{hf_config_class_name=}")
-        logger.debug(f"{tool_parser=}")
-
-        # Wrap the lower-level engine in an OpenAI-compatible chat API
-        model_config = await self.engine.get_model_config()
-        self.chat = OpenAIServingChat(
-            engine_client=self.engine,
-            model_config=model_config,
-            models=OpenAIServingModels(
-                engine_client=self.engine,
-                model_config=model_config,
-                base_model_paths=[
-                    # The layer below us will only see resolved model IDs
-                    BaseModelPath(model_id_for_vllm, model_id_for_vllm)
-                ],
-            ),
-            response_role="assistant",
-            request_logger=None,  # Use default logging
-            chat_template=None,  # Use default template from model checkpoint
-            enable_auto_tools=True,
-            tool_parser=tool_parser,
-            chat_template_content_format="auto",
-        )
-        self.resolved_model_id = model_id_for_vllm
-        self.model_ids.add(model.model_id)
-
-        logger.info(f"Finished preloading model: {model_id_for_vllm}")
-
-        return model
-
-    async def unregister_model(self, model_id: str) -> None:
-        """
-        Callback that is called when the server removes an inference endpoint from an inference
-        provider.
-
-        :param model_id: The same external ID that the higher layers of the stack previously passed
-        to :func:`register_model()`
-        """
-        if model_id not in self.model_ids:
-            raise ValueError(
-                f"Attempted to unregister model ID '{model_id}', but that ID is not registered to this provider."
-            )
-        self.model_ids.remove(model_id)
-
-        if len(self.model_ids) == 0:
-            # Last model was just unregistered. Shut down the connection to vLLM and free up
-            # resources.
-            # Note that this operation may cause in-flight chat completion requests on the
-            # now-unregistered model to return errors.
-            self.resolved_model_id = None
-            self.chat = None
-            self.engine.shutdown_background_loop()
-            self.engine = None
-
-    ###########################################################################
-    # METHODS INHERITED FROM Inference INTERFACE
-
-    async def completion(
-        self,
-        model_id: str,
-        content: InterleavedContent,
-        sampling_params: SamplingParams | None = None,
-        response_format: ResponseFormat | None = None,
-        stream: bool | None = False,
-        logprobs: LogProbConfig | None = None,
-    ) -> CompletionResponse | AsyncIterator[CompletionResponseStreamChunk]:
-        if model_id not in self.model_ids:
-            raise ValueError(
-                f"This adapter is not registered to model id '{model_id}'. Registered IDs are: {self.model_ids}"
-            )
-        if not isinstance(content, str):
-            raise NotImplementedError("Multimodal input not currently supported")
-        if sampling_params is None:
-            sampling_params = SamplingParams()
-
-        converted_sampling_params = _convert_sampling_params(sampling_params, response_format, logprobs)
-
-        logger.debug(f"{converted_sampling_params=}")
-
-        if stream:
-            return self._streaming_completion(content, converted_sampling_params)
-        else:
-            streaming_result = None
-            async for _ in self._streaming_completion(content, converted_sampling_params):
-                pass
-            return CompletionResponse(
-                content=streaming_result.delta,
-                stop_reason=streaming_result.stop_reason,
-                logprobs=streaming_result.logprobs,
-            )
-
-    async def embeddings(
-        self,
-        model_id: str,
-        contents: list[str] | list[InterleavedContentItem],
-        text_truncation: TextTruncation | None = TextTruncation.none,
-        output_dimension: int | None = None,
-        task_type: EmbeddingTaskType | None = None,
-    ) -> EmbeddingsResponse:
-        raise NotImplementedError()
-
-    async def openai_embeddings(
-        self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
-    ) -> OpenAIEmbeddingsResponse:
-        raise NotImplementedError()
-
-    async def chat_completion(
-        self,
-        model_id: str,
-        messages: list[Message],  # type: ignore
-        sampling_params: SamplingParams | None = None,
-        response_format: ResponseFormat | None = None,  # type: ignore
-        tools: list[ToolDefinition] | None = None,
-        tool_choice: ToolChoice | None = ToolChoice.auto,
-        tool_prompt_format: ToolPromptFormat | None = None,
-        stream: bool | None = False,
-        logprobs: LogProbConfig | None = None,
-        tool_config: ToolConfig | None = None,
-    ) -> ChatCompletionResponse | ChatCompletionResponseStreamChunk:
-        sampling_params = sampling_params or SamplingParams()
-        if model_id not in self.model_ids:
-            raise ValueError(
-                f"This adapter is not registered to model id '{model_id}'. Registered IDs are: {self.model_ids}"
-            )
-
-        # Convert to Llama Stack internal format for consistency
-        request = ChatCompletionRequest(
-            model=self.resolved_model_id,
-            messages=messages,
-            sampling_params=sampling_params,
-            response_format=response_format,
-            tools=tools,
-            tool_choice=tool_choice,
-            tool_prompt_format=tool_prompt_format,
-            stream=stream,
-            logprobs=logprobs,
-        )
-
-        if self.is_meta_llama_model:
-            # Bypass vLLM chat templating layer for Meta Llama models, because the
-            # templating layer in Llama Stack currently produces better results.
-            logger.debug(
-                f"Routing {self.resolved_model_id} chat completion through "
-                f"Llama Stack's templating layer instead of vLLM's."
-            )
-            return await self._chat_completion_for_meta_llama(request)
-
-        logger.debug(f"{self.resolved_model_id} is not a Meta Llama model")
-
-        # Arguments to the vLLM call must be packaged as a ChatCompletionRequest dataclass.
-        # Note that this dataclass has the same name as a similar dataclass in Llama Stack.
-        request_options = await llama_stack_chat_completion_to_openai_chat_completion_dict(request)
-        chat_completion_request = vllm.entrypoints.openai.protocol.ChatCompletionRequest(**request_options)
-
-        logger.debug(f"Converted request: {chat_completion_request}")
-
-        vllm_result = await self.chat.create_chat_completion(chat_completion_request)
-        logger.debug(f"Result from vLLM: {vllm_result}")
-        if isinstance(vllm_result, vllm.entrypoints.openai.protocol.ErrorResponse):
-            raise ValueError(f"Error from vLLM layer: {vllm_result}")
-
-        # Return type depends on "stream" argument
-        if stream:
-            if not isinstance(vllm_result, AsyncGenerator):
-                raise TypeError(f"Unexpected result type {type(vllm_result)} for streaming inference call")
-            # vLLM client returns a stream of strings, which need to be parsed.
-            # Stream comes in the form of an async generator.
-            return self._convert_streaming_results(vllm_result)
-        else:
-            if not isinstance(vllm_result, vllm.entrypoints.openai.protocol.ChatCompletionResponse):
-                raise TypeError(f"Unexpected result type {type(vllm_result)} for non-streaming inference call")
-            return self._convert_non_streaming_results(vllm_result)
-
-    ###########################################################################
-    # INTERNAL METHODS
-
-    async def _streaming_completion(
-        self, content: str, sampling_params: vllm.SamplingParams
-    ) -> AsyncIterator[CompletionResponseStreamChunk]:
-        """Internal implementation of :func:`completion()` API for the streaming case. Assumes
-        that arguments have been validated upstream.
-
-        :param content: Must be a string
-        :param sampling_params: Paramters from  public API's ``response_format``
-         and ``sampling_params`` arguments, converted to VLLM format
-        """
-        # We run agains the vLLM generate() call directly instead of using the OpenAI-compatible
-        # layer, because doing so simplifies the code here.
-
-        # The vLLM engine requires a unique identifier for each call to generate()
-        request_id = _random_uuid_str()
-
-        # The vLLM generate() API is streaming-only and returns an async generator.
-        # The generator returns objects of type vllm.RequestOutput.
-        results_generator = self.engine.generate(content, sampling_params, request_id)
-
-        # Need to know the model's EOS token ID for the conversion code below.
-        # AsyncLLMEngine is a wrapper around LLMEngine, and the tokenizer is only available if
-        # we drill down to the LLMEngine inside the AsyncLLMEngine.
-        # Similarly, the tokenizer in an LLMEngine is a wrapper around a BaseTokenizerGroup,
-        # and we need to drill down to the Hugging Face tokenizer inside the BaseTokenizerGroup.
-        llm_engine = self.engine.engine
-        tokenizer_group = llm_engine.tokenizer
-        eos_token_id = tokenizer_group.tokenizer.eos_token_id
-
-        request_output: vllm.RequestOutput = None
-        async for request_output in results_generator:
-            # Check for weird inference failures
-            if request_output.outputs is None or len(request_output.outputs) == 0:
-                # This case also should never happen
-                raise ValueError("Inference produced empty result")
-
-            # If we get here, then request_output contains the final output of the generate() call.
-            # The result may include multiple alternate outputs, but Llama Stack APIs only allow
-            # us to return one.
-            output: vllm.CompletionOutput = request_output.outputs[0]
-            completion_string = output.text
-
-            # Convert logprobs from vLLM's format to Llama Stack's format
-            logprobs = [
-                TokenLogProbs(logprobs_by_token={v.decoded_token: v.logprob for _, v in logprob_dict.items()})
-                for logprob_dict in output.logprobs
-            ]
-
-            # The final output chunk should be labeled with the reason that the overall generate()
-            # call completed.
-            logger.debug(f"{output.stop_reason=}; {type(output.stop_reason)=}")
-            if output.stop_reason is None:
-                stop_reason = None  # Still going
-            elif output.stop_reason == "stop":
-                stop_reason = StopReason.end_of_turn
-            elif output.stop_reason == "length":
-                stop_reason = StopReason.out_of_tokens
-            elif isinstance(output.stop_reason, int):
-                # If the model config specifies multiple end-of-sequence tokens, then vLLM
-                # will return the token ID of the EOS token in the stop_reason field.
-                stop_reason = StopReason.end_of_turn
-            else:
-                raise ValueError(f"Unrecognized stop reason '{output.stop_reason}'")
-
-            # vLLM's protocol outputs the stop token, then sets end of message on the next step for
-            # some reason.
-            if request_output.outputs[-1].token_ids[-1] == eos_token_id:
-                stop_reason = StopReason.end_of_message
-
-            yield CompletionResponseStreamChunk(delta=completion_string, stop_reason=stop_reason, logprobs=logprobs)
-
-        # Llama Stack requires that the last chunk have a stop reason, but vLLM doesn't always
-        # provide one if it runs out of tokens.
-        if stop_reason is None:
-            yield CompletionResponseStreamChunk(
-                delta=completion_string,
-                stop_reason=StopReason.out_of_tokens,
-                logprobs=logprobs,
-            )
-
-    def _convert_non_streaming_results(
-        self, vllm_result: vllm.entrypoints.openai.protocol.ChatCompletionResponse
-    ) -> ChatCompletionResponse:
-        """
-        Subroutine to convert the non-streaming output of vLLM's OpenAI-compatible API into an
-        equivalent Llama Stack object.
-
-        The result from vLLM's non-streaming API is a dataclass with the same name as the Llama
-        Stack ChatCompletionResponse dataclass, but with more and different field names. We ignore
-        the fields that aren't currently present in the Llama Stack dataclass.
-        """
-
-        # There may be multiple responses, but we can only pass through the first one.
-        if len(vllm_result.choices) == 0:
-            raise ValueError("Don't know how to convert response object without any responses")
-        vllm_message = vllm_result.choices[0].message
-        vllm_finish_reason = vllm_result.choices[0].finish_reason
-
-        converted_message = CompletionMessage(
-            role=vllm_message.role,
-            # Llama Stack API won't accept None for content field.
-            content=("" if vllm_message.content is None else vllm_message.content),
-            stop_reason=get_stop_reason(vllm_finish_reason),
-            tool_calls=[
-                ToolCall(
-                    call_id=t.id,
-                    tool_name=t.function.name,
-                    # vLLM function args come back as a string. Llama Stack expects JSON.
-                    arguments=json.loads(t.function.arguments),
-                    arguments_json=t.function.arguments,
-                )
-                for t in vllm_message.tool_calls
-            ],
-        )
-
-        # TODO: Convert logprobs
-
-        logger.debug(f"Converted message: {converted_message}")
-
-        return ChatCompletionResponse(
-            completion_message=converted_message,
-        )
-
-    async def _chat_completion_for_meta_llama(
-        self, request: ChatCompletionRequest
-    ) -> ChatCompletionResponse | AsyncIterator[ChatCompletionResponseStreamChunk]:
-        """
-        Subroutine that routes chat completions for Meta Llama models through Llama Stack's
-        chat template instead of using vLLM's version of that template. The Llama Stack version
-        of the chat template currently produces more reliable outputs.
-
-        Once vLLM's support for Meta Llama models has matured more, we should consider routing
-        Meta Llama requests through the vLLM chat completions API instead of using this method.
-        """
-        formatter = ChatFormat(Tokenizer.get_instance())
-
-        # Note that this function call modifies `request` in place.
-        prompt = await chat_completion_request_to_prompt(request, self.resolved_model_id)
-
-        model_id = list(self.model_ids)[0]  # Any model ID will do here
-        completion_response_or_iterator = await self.completion(
-            model_id=model_id,
-            content=prompt,
-            sampling_params=request.sampling_params,
-            response_format=request.response_format,
-            stream=request.stream,
-            logprobs=request.logprobs,
-        )
-
-        if request.stream:
-            if not isinstance(completion_response_or_iterator, AsyncIterator):
-                raise TypeError(
-                    f"Received unexpected result type {type(completion_response_or_iterator)}for streaming request."
-                )
-            return self._chat_completion_for_meta_llama_streaming(completion_response_or_iterator, request)
-
-        # elsif not request.stream:
-        if not isinstance(completion_response_or_iterator, CompletionResponse):
-            raise TypeError(
-                f"Received unexpected result type {type(completion_response_or_iterator)}for non-streaming request."
-            )
-        completion_response: CompletionResponse = completion_response_or_iterator
-        raw_message = formatter.decode_assistant_message_from_content(
-            completion_response.content, completion_response.stop_reason
-        )
-        return ChatCompletionResponse(
-            completion_message=CompletionMessage(
-                content=raw_message.content,
-                stop_reason=raw_message.stop_reason,
-                tool_calls=raw_message.tool_calls,
-            ),
-            logprobs=completion_response.logprobs,
-        )
-
-    async def _chat_completion_for_meta_llama_streaming(
-        self, results_iterator: AsyncIterator, request: ChatCompletionRequest
-    ) -> AsyncIterator:
-        """
-        Code from :func:`_chat_completion_for_meta_llama()` that needs to be a separate
-        method to keep asyncio happy.
-        """
-
-        # Convert to OpenAI format, then use shared code to convert to Llama Stack format.
-        async def _generate_and_convert_to_openai_compat():
-            chunk: CompletionResponseStreamChunk  # Make Pylance happy
-            last_text_len = 0
-            async for chunk in results_iterator:
-                if chunk.stop_reason == StopReason.end_of_turn:
-                    finish_reason = "stop"
-                elif chunk.stop_reason == StopReason.end_of_message:
-                    finish_reason = "eos"
-                elif chunk.stop_reason == StopReason.out_of_tokens:
-                    finish_reason = "length"
-                else:
-                    finish_reason = None
-
-                # Convert delta back to an actual delta
-                text_delta = chunk.delta[last_text_len:]
-                last_text_len = len(chunk.delta)
-
-                logger.debug(f"{text_delta=}; {finish_reason=}")
-
-                yield OpenAICompatCompletionResponse(
-                    choices=[OpenAICompatCompletionChoice(finish_reason=finish_reason, text=text_delta)]
-                )
-
-        stream = _generate_and_convert_to_openai_compat()
-        async for chunk in process_chat_completion_stream_response(stream, request):
-            logger.debug(f"Returning chunk: {chunk}")
-            yield chunk
-
-    async def _convert_streaming_results(self, vllm_result: AsyncIterator) -> AsyncIterator:
-        """
-        Subroutine that wraps the streaming outputs of vLLM's OpenAI-compatible
-        API into a second async iterator that returns Llama Stack objects.
-
-        :param vllm_result: Stream of strings that need to be parsed
-        """
-        # Tool calls come in pieces, but Llama Stack expects them in bigger chunks. We build up
-        # those chunks and output them at the end.
-        # This data structure holds the current set of partial tool calls.
-        index_to_tool_call: dict[int, dict] = dict()
-
-        # The Llama Stack event stream must always start with a start event. Use an empty one to
-        # simplify logic below
-        yield ChatCompletionResponseStreamChunk(
-            event=ChatCompletionResponseEvent(
-                event_type=ChatCompletionResponseEventType.start,
-                delta=TextDelta(text=""),
-                stop_reason=None,
-            )
-        )
-
-        converted_stop_reason = None
-        async for chunk_str in vllm_result:
-            # Due to OpenAI compatibility, each event in the stream will start with "data: " and
-            # end with "\n\n".
-            _prefix = "data: "
-            _suffix = "\n\n"
-            if not chunk_str.startswith(_prefix) or not chunk_str.endswith(_suffix):
-                raise ValueError(f"Can't parse result string from vLLM: '{re.escape(chunk_str)}'")
-
-            # In between the "data: " and newlines is an event record
-            data_str = chunk_str[len(_prefix) : -len(_suffix)]
-
-            # The end of the stream is indicated with "[DONE]"
-            if data_str == "[DONE]":
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.complete,
-                        delta=TextDelta(text=""),
-                        stop_reason=converted_stop_reason,
-                    )
-                )
-                return
-
-            # Anything that is not "[DONE]" should be a JSON record
-            parsed_chunk = json.loads(data_str)
-
-            logger.debug(f"Parsed JSON event to:\n{json.dumps(parsed_chunk, indent=2)}")
-
-            # The result may contain multiple completions, but Llama Stack APIs only support
-            # returning one.
-            first_choice = parsed_chunk["choices"][0]
-            converted_stop_reason = get_stop_reason(first_choice["finish_reason"])
-            delta_record = first_choice["delta"]
-
-            if "content" in delta_record:
-                # Text delta
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=TextDelta(text=delta_record["content"]),
-                        stop_reason=converted_stop_reason,
-                    )
-                )
-            elif "tool_calls" in delta_record:
-                # Tool call(s). Llama Stack APIs do not have a clear way to return partial tool
-                # calls, so buffer until we get a "tool calls" stop reason
-                for tc in delta_record["tool_calls"]:
-                    index = tc["index"]
-                    if index not in index_to_tool_call:
-                        # First time this tool call is showing up
-                        index_to_tool_call[index] = dict()
-                    tool_call = index_to_tool_call[index]
-                    if "id" in tc:
-                        tool_call["call_id"] = tc["id"]
-                    if "function" in tc:
-                        if "name" in tc["function"]:
-                            tool_call["tool_name"] = tc["function"]["name"]
-                        if "arguments" in tc["function"]:
-                            # Arguments comes in as pieces of a string
-                            if "arguments_str" not in tool_call:
-                                tool_call["arguments_str"] = ""
-                            tool_call["arguments_str"] += tc["function"]["arguments"]
-            else:
-                raise ValueError(f"Don't know how to parse event delta: {delta_record}")
-
-            if first_choice["finish_reason"] == "tool_calls":
-                # Special OpenAI code for "tool calls complete".
-                # Output the buffered tool calls. Llama Stack requires a separate event per tool
-                # call.
-                for tool_call_record in index_to_tool_call.values():
-                    # Arguments come in as a string. Parse the completed string.
-                    tool_call_record["arguments"] = json.loads(tool_call_record["arguments_str"])
-                    del tool_call_record["arguments_str"]
-
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=ToolCallDelta(tool_call=tool_call_record, parse_status="succeeded"),
-                            stop_reason=converted_stop_reason,
-                        )
-                    )
-
-        # If we get here, we've lost the connection with the vLLM event stream before it ended
-        # normally.
-        raise ValueError("vLLM event stream ended without [DONE] message.")

llama_stack/providers/inline/post_training/huggingface/__init__.py

@@ -6,7 +6,7 @@
 
 from typing import Any
 
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import Api
 
 from .config import HuggingFacePostTrainingConfig
 

llama_stack/providers/inline/post_training/huggingface/config.py

@@ -67,6 +67,17 @@ class HuggingFacePostTrainingConfig(BaseModel):
     # Can improve data transfer speed to GPU but uses more memory
     dataloader_pin_memory: bool = True
 
+    # DPO-specific parameters
+    dpo_beta: float = 0.1
+    use_reference_model: bool = True
+    dpo_loss_type: Literal["sigmoid", "hinge", "ipo", "kto_pair"] = "sigmoid"
+    dpo_output_dir: str
+
     @classmethod
     def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
-        return {"checkpoint_format": "huggingface", "distributed_backend": None, "device": "cpu"}
+        return {
+            "checkpoint_format": "huggingface",
+            "distributed_backend": None,
+            "device": "cpu",
+            "dpo_output_dir": __distro_dir__ + "/dpo_output",
+        }

llama_stack/providers/inline/post_training/huggingface/post_training.py

@@ -22,12 +22,8 @@ from llama_stack.apis.post_training import (
 from llama_stack.providers.inline.post_training.huggingface.config import (
     HuggingFacePostTrainingConfig,
 )
-from llama_stack.providers.inline.post_training.huggingface.recipes.finetune_single_device import (
-    HFFinetuningSingleDevice,
-)
 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):
@@ -36,6 +32,7 @@ class TrainingArtifactType(Enum):
 
 
 _JOB_TYPE_SUPERVISED_FINE_TUNE = "supervised-fine-tune"
+_JOB_TYPE_DPO_TRAINING = "dpo-training"
 
 
 class HuggingFacePostTrainingImpl:
@@ -81,6 +78,10 @@ class HuggingFacePostTrainingImpl:
         algorithm_config: AlgorithmConfig | None = None,
     ) -> PostTrainingJob:
         async def handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb):
+            from llama_stack.providers.inline.post_training.huggingface.recipes.finetune_single_device import (
+                HFFinetuningSingleDevice,
+            )
+
             on_log_message_cb("Starting HF finetuning")
 
             recipe = HFFinetuningSingleDevice(
@@ -119,12 +120,41 @@ class HuggingFacePostTrainingImpl:
         hyperparam_search_config: dict[str, Any],
         logger_config: dict[str, Any],
     ) -> PostTrainingJob:
-        raise NotImplementedError("DPO alignment is not implemented yet")
+        async def handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb):
+            from llama_stack.providers.inline.post_training.huggingface.recipes.finetune_single_device_dpo import (
+                HFDPOAlignmentSingleDevice,
+            )
 
-    async def get_training_jobs(self) -> ListPostTrainingJobsResponse:
-        return ListPostTrainingJobsResponse(
-            data=[PostTrainingJob(job_uuid=job.id) for job in self._scheduler.get_jobs()]
-        )
+            on_log_message_cb("Starting HF DPO alignment")
+
+            recipe = HFDPOAlignmentSingleDevice(
+                job_uuid=job_uuid,
+                datasetio_api=self.datasetio_api,
+                datasets_api=self.datasets_api,
+            )
+
+            resources_allocated, checkpoints = await recipe.train(
+                model=finetuned_model,
+                output_dir=f"{self.config.dpo_output_dir}/{job_uuid}",
+                job_uuid=job_uuid,
+                dpo_config=algorithm_config,
+                config=training_config,
+                provider_config=self.config,
+            )
+
+            on_artifact_collected_cb(self._resources_stats_to_artifact(resources_allocated))
+            if checkpoints:
+                for checkpoint in checkpoints:
+                    artifact = self._checkpoint_to_artifact(checkpoint)
+                    on_artifact_collected_cb(artifact)
+            else:
+                on_log_message_cb("Warning: No checkpoints were saved during DPO training")
+
+            on_status_change_cb(SchedulerJobStatus.completed)
+            on_log_message_cb("HF DPO alignment completed")
+
+        job_uuid = self._scheduler.schedule(_JOB_TYPE_DPO_TRAINING, job_uuid, handler)
+        return PostTrainingJob(job_uuid=job_uuid)
 
     @staticmethod
     def _get_artifacts_metadata_by_type(job, artifact_type):
@@ -139,7 +169,6 @@ class HuggingFacePostTrainingImpl:
         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) -> PostTrainingJobStatusResponse | None:
         job = self._scheduler.get_job(job_uuid)
 
@@ -166,11 +195,14 @@ class HuggingFacePostTrainingImpl:
             resources_allocated=self._get_resources_allocated(job),
         )
 
-    @webmethod(route="/post-training/job/cancel")
     async def cancel_training_job(self, job_uuid: str) -> None:
         self._scheduler.cancel(job_uuid)
 
-    @webmethod(route="/post-training/job/artifacts")
     async def get_training_job_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse | None:
         job = self._scheduler.get_job(job_uuid)
         return PostTrainingJobArtifactsResponse(job_uuid=job_uuid, checkpoints=self._get_checkpoints(job))
+
+    async def get_training_jobs(self) -> ListPostTrainingJobsResponse:
+        return ListPostTrainingJobsResponse(
+            data=[PostTrainingJob(job_uuid=job.id) for job in self._scheduler.get_jobs()]
+        )

llama_stack/providers/inline/post_training/huggingface/recipes/finetune_single_device.py

@@ -8,30 +8,13 @@ import gc
 import json
 import logging
 import multiprocessing
-import os
-import signal
-import sys
-from datetime import UTC, datetime
 from pathlib import Path
 from typing import Any
 
-import psutil
-
-from llama_stack.providers.inline.post_training.common.utils import evacuate_model_from_device
-
-# Set tokenizer parallelism environment variable
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-# Force PyTorch to use OpenBLAS instead of MKL
-os.environ["MKL_THREADING_LAYER"] = "GNU"
-os.environ["MKL_SERVICE_FORCE_INTEL"] = "0"
-os.environ["MKL_NUM_THREADS"] = "1"
-
 import torch
 from datasets import Dataset
 from peft import LoraConfig
 from transformers import (
-    AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
 )
@@ -45,93 +28,25 @@ from llama_stack.apis.post_training import (
     LoraFinetuningConfig,
     TrainingConfig,
 )
+from llama_stack.providers.inline.post_training.common.utils import evacuate_model_from_device
 
 from ..config import HuggingFacePostTrainingConfig
+from ..utils import (
+    calculate_training_steps,
+    create_checkpoints,
+    get_memory_stats,
+    get_save_strategy,
+    load_model,
+    load_rows_from_dataset,
+    setup_environment,
+    setup_signal_handlers,
+    setup_torch_device,
+    split_dataset,
+)
 
 logger = logging.getLogger(__name__)
 
 
-def get_gb(to_convert: int) -> str:
-    """Converts memory stats to GB and formats to 2 decimal places.
-    Args:
-        to_convert: Memory value in bytes
-    Returns:
-        str: Memory value in GB formatted to 2 decimal places
-    """
-    return f"{(to_convert / (1024**3)):.2f}"
-
-
-def get_memory_stats(device: torch.device) -> dict[str, Any]:
-    """Get memory statistics for the given device."""
-    stats = {
-        "system_memory": {
-            "total": get_gb(psutil.virtual_memory().total),
-            "available": get_gb(psutil.virtual_memory().available),
-            "used": get_gb(psutil.virtual_memory().used),
-            "percent": psutil.virtual_memory().percent,
-        }
-    }
-
-    if device.type == "cuda":
-        stats["device_memory"] = {
-            "allocated": get_gb(torch.cuda.memory_allocated(device)),
-            "reserved": get_gb(torch.cuda.memory_reserved(device)),
-            "max_allocated": get_gb(torch.cuda.max_memory_allocated(device)),
-        }
-    elif device.type == "mps":
-        # MPS doesn't provide direct memory stats, but we can track system memory
-        stats["device_memory"] = {
-            "note": "MPS memory stats not directly available",
-            "system_memory_used": get_gb(psutil.virtual_memory().used),
-        }
-    elif device.type == "cpu":
-        # For CPU, we track process memory usage
-        process = psutil.Process()
-        stats["device_memory"] = {
-            "process_rss": get_gb(process.memory_info().rss),
-            "process_vms": get_gb(process.memory_info().vms),
-            "process_percent": process.memory_percent(),
-        }
-
-    return stats
-
-
-def setup_torch_device(device_str: str) -> torch.device:
-    """Initialize and validate a PyTorch device.
-    This function handles device initialization and validation for different device types:
-    - CUDA: Validates CUDA availability and handles device selection
-    - MPS: Validates MPS availability for Apple Silicon
-    - CPU: Basic validation
-    - HPU: Raises error as it's not supported
-    Args:
-        device_str: String specifying the device ('cuda', 'cpu', 'mps')
-    Returns:
-        torch.device: The initialized and validated device
-    Raises:
-        RuntimeError: If device initialization fails or device is not supported
-    """
-    try:
-        device = torch.device(device_str)
-    except RuntimeError as e:
-        raise RuntimeError(f"Error getting Torch Device {str(e)}") from e
-
-    # Validate device capabilities
-    if device.type == "cuda":
-        if not torch.cuda.is_available():
-            raise RuntimeError(
-                f"{device.type}: Torch has no CUDA/ROCm support or could not detect a compatible device."
-            )
-        if device.index is None:
-            device = torch.device(device.type, torch.cuda.current_device())
-    elif device.type == "mps":
-        if not torch.backends.mps.is_available():
-            raise RuntimeError(f"{device.type}: Torch has no MPS support or could not detect a compatible device.")
-    elif device.type == "hpu":
-        raise RuntimeError(f"{device.type}: training does not support Intel Gaudi.")
-
-    return device
-
-
 class HFFinetuningSingleDevice:
     def __init__(
         self,
@@ -262,19 +177,6 @@ class HFFinetuningSingleDevice:
             remove_columns=ds.column_names,
         )
 
-    async def _setup_data(self, dataset_id: str) -> list[dict[str, Any]]:
-        """Load dataset from llama stack dataset provider"""
-        try:
-            all_rows = await self.datasetio_api.iterrows(
-                dataset_id=dataset_id,
-                limit=-1,
-            )
-            if not isinstance(all_rows.data, list):
-                raise RuntimeError("Expected dataset data to be a list")
-            return all_rows.data
-        except Exception as e:
-            raise RuntimeError(f"Failed to load dataset: {str(e)}") from e
-
     def _run_training_sync(
         self,
         model: str,
@@ -327,7 +229,7 @@ class HFFinetuningSingleDevice:
 
         # Load dataset
         logger.info(f"Loading dataset: {config.data_config.dataset_id}")
-        rows = await self._setup_data(config.data_config.dataset_id)
+        rows = await load_rows_from_dataset(self.datasetio_api, config.data_config.dataset_id)
         if not self.validate_dataset_format(rows):
             raise ValueError("Dataset is missing required fields: input_query, expected_answer, chat_completion_input")
         logger.info(f"Loaded {len(rows)} rows from dataset")
@@ -369,47 +271,10 @@ class HFFinetuningSingleDevice:
             raise ValueError(f"Failed to create dataset: {str(e)}") from e
 
         # Split dataset
-        logger.info("Splitting dataset into train and validation sets")
-        train_val_split = ds.train_test_split(test_size=0.1, seed=42)
-        train_dataset = train_val_split["train"]
-        eval_dataset = train_val_split["test"]
-        logger.info(f"Split dataset into {len(train_dataset)} training and {len(eval_dataset)} validation examples")
+        train_dataset, eval_dataset = split_dataset(ds)
 
         return train_dataset, eval_dataset, tokenizer
 
-    def load_model(
-        self,
-        model: str,
-        device: torch.device,
-        provider_config: HuggingFacePostTrainingConfig,
-    ) -> AutoModelForCausalLM:
-        """Load and initialize the model for training.
-        Args:
-            model: The model identifier to load
-            device: The device to load the model onto
-            provider_config: Provider-specific configuration
-        Returns:
-            The loaded and initialized model
-        Raises:
-            RuntimeError: If model loading fails
-        """
-        logger.info("Loading the base model")
-        try:
-            model_config = AutoConfig.from_pretrained(model, **provider_config.model_specific_config)
-            model_obj = AutoModelForCausalLM.from_pretrained(
-                model,
-                torch_dtype="auto" if device.type != "cpu" else "float32",
-                quantization_config=None,
-                config=model_config,
-                **provider_config.model_specific_config,
-            )
-            # Always move model to specified device
-            model_obj = model_obj.to(device)
-            logger.info(f"Model loaded and moved to device: {model_obj.device}")
-            return model_obj
-        except Exception as e:
-            raise RuntimeError(f"Failed to load model: {str(e)}") from e
-
     def setup_training_args(
         self,
         config: TrainingConfig,
@@ -439,27 +304,12 @@ class HFFinetuningSingleDevice:
             raise ValueError("DataConfig is required for training")
         data_config = config.data_config
 
-        # Calculate steps
-        total_steps = steps_per_epoch * config.n_epochs
-        max_steps = min(config.max_steps_per_epoch, total_steps)
-        logging_steps = max(1, steps_per_epoch // 50)  # Log 50 times per epoch
-
-        logger.info("Training configuration:")
-        logger.info(f"- Steps per epoch: {steps_per_epoch}")
-        logger.info(f"- Total steps: {total_steps}")
-        logger.info(f"- Max steps: {max_steps}")
-        logger.info(f"- Logging steps: {logging_steps}")
-
-        # Configure save strategy
-        save_strategy = "no"
-        eval_strategy = "no"
-        if output_dir_path:
-            save_strategy = "epoch"
-            eval_strategy = "epoch"
-            logger.info(f"Will save checkpoints to {output_dir_path}")
+        # Calculate steps and get save strategy
+        step_info = calculate_training_steps(steps_per_epoch, config)
+        save_strategy, eval_strategy = get_save_strategy(output_dir_path)
 
         return SFTConfig(
-            max_steps=max_steps,
+            max_steps=step_info["max_steps"],
             output_dir=str(output_dir_path) if output_dir_path is not None else None,
             num_train_epochs=config.n_epochs,
             per_device_train_batch_size=data_config.batch_size,
@@ -469,7 +319,7 @@ class HFFinetuningSingleDevice:
             use_cpu=True if device.type == "cpu" and not torch.backends.mps.is_available() else False,
             save_strategy=save_strategy,
             report_to="none",
-            max_seq_length=provider_config.max_seq_length,
+            max_length=provider_config.max_seq_length,
             gradient_accumulation_steps=config.gradient_accumulation_steps,
             gradient_checkpointing=provider_config.gradient_checkpointing,
             learning_rate=lr,
@@ -483,7 +333,7 @@ class HFFinetuningSingleDevice:
             load_best_model_at_end=True if output_dir_path else False,
             metric_for_best_model="eval_loss",
             greater_is_better=False,
-            logging_steps=logging_steps,
+            logging_steps=step_info["logging_steps"],
         )
 
     def save_model(
@@ -523,13 +373,11 @@ class HFFinetuningSingleDevice:
     ) -> None:
         """Run the training process with signal handling."""
 
-        def signal_handler(signum, frame):
-            """Handle termination signals gracefully."""
-            logger.info(f"Received signal {signum}, initiating graceful shutdown")
-            sys.exit(0)
+        # Setup environment variables
+        setup_environment()
 
-        signal.signal(signal.SIGTERM, signal_handler)
-        signal.signal(signal.SIGINT, signal_handler)
+        # Setup signal handlers
+        setup_signal_handlers()
 
         # Convert config dicts back to objects
         logger.info("Initializing configuration objects")
@@ -558,7 +406,7 @@ class HFFinetuningSingleDevice:
         )
 
         # Load model
-        model_obj = self.load_model(model, device, provider_config_obj)
+        model_obj = load_model(model, device, provider_config_obj)
 
         # Initialize trainer
         logger.info("Initializing SFTTrainer")
@@ -633,7 +481,7 @@ class HFFinetuningSingleDevice:
         # Train in a separate process
         logger.info("Starting training in separate process")
         try:
-            # Set multiprocessing start method to 'spawn' for CUDA/MPS compatibility
+            # Setup multiprocessing for device
             if device.type in ["cuda", "mps"]:
                 multiprocessing.set_start_method("spawn", force=True)
 
@@ -663,37 +511,7 @@ class HFFinetuningSingleDevice:
 
             checkpoints = []
             if output_dir_path:
-                # Get all checkpoint directories and sort them numerically
-                checkpoint_dirs = sorted(
-                    [d for d in output_dir_path.glob("checkpoint-*") if d.is_dir()],
-                    key=lambda x: int(x.name.split("-")[1]),
-                )
-
-                # Add all checkpoint directories
-                for epoch_number, checkpoint_dir in enumerate(checkpoint_dirs, start=1):
-                    # Get the creation time of the directory
-                    created_time = datetime.fromtimestamp(os.path.getctime(checkpoint_dir), tz=UTC)
-
-                    checkpoint = Checkpoint(
-                        identifier=checkpoint_dir.name,
-                        created_at=created_time,
-                        epoch=epoch_number,
-                        post_training_job_id=job_uuid,
-                        path=str(checkpoint_dir),
-                    )
-                    checkpoints.append(checkpoint)
-
-                # Add the merged model as a checkpoint
-                merged_model_path = output_dir_path / "merged_model"
-                if merged_model_path.exists():
-                    checkpoint = Checkpoint(
-                        identifier=f"{model}-sft-{config.n_epochs}",
-                        created_at=datetime.now(UTC),
-                        epoch=config.n_epochs,
-                        post_training_job_id=job_uuid,
-                        path=str(merged_model_path),
-                    )
-                    checkpoints.append(checkpoint)
+                checkpoints = create_checkpoints(output_dir_path, job_uuid, model, config, "merged_model")
 
             return memory_stats, checkpoints if checkpoints else None
         finally:

llama_stack/providers/inline/post_training/huggingface/recipes/finetune_single_device_dpo.py

@@ -0,0 +1,485 @@
+# 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 gc
+import logging
+import multiprocessing
+from pathlib import Path
+from typing import Any
+
+import torch
+from datasets import Dataset
+from transformers import (
+    AutoTokenizer,
+)
+from trl import DPOConfig, DPOTrainer
+
+from llama_stack.apis.datasetio import DatasetIO
+from llama_stack.apis.datasets import Datasets
+from llama_stack.apis.post_training import (
+    Checkpoint,
+    DPOAlignmentConfig,
+    TrainingConfig,
+)
+from llama_stack.providers.inline.post_training.common.utils import evacuate_model_from_device
+
+from ..config import HuggingFacePostTrainingConfig
+from ..utils import (
+    calculate_training_steps,
+    create_checkpoints,
+    get_memory_stats,
+    get_save_strategy,
+    load_model,
+    load_rows_from_dataset,
+    setup_environment,
+    setup_signal_handlers,
+    setup_torch_device,
+    split_dataset,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class HFDPOAlignmentSingleDevice:
+    def __init__(
+        self,
+        job_uuid: str,
+        datasetio_api: DatasetIO,
+        datasets_api: Datasets,
+    ):
+        self.datasetio_api = datasetio_api
+        self.datasets_api = datasets_api
+        self.job_uuid = job_uuid
+
+    def validate_dataset_format(self, rows: list[dict]) -> None:
+        """Validate that the dataset has the required fields for DPO training."""
+        required_fields = ["prompt", "chosen", "rejected"]
+
+        if not rows:
+            logger.warning("Dataset is empty")
+            raise ValueError("Dataset is empty")
+
+        for i, row in enumerate(rows):
+            if not isinstance(row, dict):
+                logger.warning(f"Row {i} is not a dictionary")
+                raise ValueError(f"Row {i} is not a dictionary")
+
+            for field in required_fields:
+                if field not in row:
+                    logger.warning(f"Row {i} missing required DPO field: {field}")
+                    raise ValueError(f"Row {i} missing required DPO field: {field}")
+
+                # Handle both string and list formats
+                if field == "prompt":
+                    # Prompt should be a string
+                    if not isinstance(row[field], str):
+                        logger.warning(f"Row {i} field '{field}' is not a string")
+                        raise ValueError(f"Row {i} field '{field}' is not a string")
+                    if not row[field].strip():
+                        logger.warning(f"Row {i} field '{field}' is empty")
+                        raise ValueError(f"Row {i} field '{field}' is empty")
+                else:
+                    # chosen/rejected can be either strings or lists of messages
+                    if isinstance(row[field], str):
+                        if not row[field].strip():
+                            logger.warning(f"Row {i} field '{field}' is empty")
+                            raise ValueError(f"Row {i} field '{field}' is empty")
+                    elif isinstance(row[field], list):
+                        if not row[field]:
+                            logger.warning(f"Row {i} field '{field}' is empty list")
+                            raise ValueError(f"Row {i} field '{field}' is empty list")
+                    else:
+                        logger.warning(f"Row {i} field '{field}' is neither string nor list")
+                        raise ValueError(f"Row {i} field '{field}' is neither string nor list")
+
+        logger.info(f"DPO dataset validation passed: {len(rows)} preference examples")
+
+    def _process_dpo_format(self, row: dict) -> tuple[str | None, str | None, str | None]:
+        """Process a row in DPO format, handling both string and conversation list formats."""
+        if all(field in row for field in ["prompt", "chosen", "rejected"]):
+            prompt = row["prompt"]
+
+            # Handle chosen field - convert list to string if needed
+            if isinstance(row["chosen"], list):
+                # For conversation format, concatenate messages
+                chosen = "\n".join(
+                    [msg.get("content", "") if isinstance(msg, dict) else str(msg) for msg in row["chosen"]]
+                )
+            else:
+                chosen = row["chosen"]
+
+            # Handle rejected field - convert list to string if needed
+            if isinstance(row["rejected"], list):
+                # For conversation format, concatenate messages
+                rejected = "\n".join(
+                    [msg.get("content", "") if isinstance(msg, dict) else str(msg) for msg in row["rejected"]]
+                )
+            else:
+                rejected = row["rejected"]
+
+            return prompt, chosen, rejected
+        return None, None, None
+
+    def _format_text_for_dpo(self, prompt: str, response: str, provider_config: HuggingFacePostTrainingConfig) -> str:
+        """Format prompt and response text based on model requirements."""
+        if hasattr(provider_config, "chat_template") and provider_config.chat_template:
+            # Use the chat template, supporting both {prompt}/{response} and {input}/{output}
+            template = provider_config.chat_template
+            # Try prompt/response first (DPO style)
+            if "{prompt}" in template and "{response}" in template:
+                return template.format(prompt=prompt, response=response)
+            # Fall back to input/output (SFT style)
+            elif "{input}" in template and "{output}" in template:
+                return template.format(input=prompt, output=response)
+            else:
+                # If template doesn't have expected placeholders, use default
+                return f"{prompt}\n{response}"
+        return f"{prompt}\n{response}"
+
+    def _create_dataset(
+        self, rows: list[dict], config: TrainingConfig, provider_config: HuggingFacePostTrainingConfig
+    ) -> Dataset:
+        """Create and preprocess the dataset for DPO."""
+        dpo_examples = []
+        for row in rows:
+            prompt, chosen, rejected = self._process_dpo_format(row)
+
+            if prompt and chosen and rejected:
+                # Format the texts
+                chosen_formatted = self._format_text_for_dpo(prompt, chosen, provider_config)
+                rejected_formatted = self._format_text_for_dpo(prompt, rejected, provider_config)
+
+                dpo_examples.append(
+                    {
+                        "prompt": prompt,
+                        "chosen": chosen_formatted,
+                        "rejected": rejected_formatted,
+                    }
+                )
+
+        if not dpo_examples:
+            raise ValueError("No valid preference examples found in dataset")
+
+        logger.info(f"Created DPO dataset with {len(dpo_examples)} preference pairs")
+        return Dataset.from_list(dpo_examples)
+
+    def _preprocess_dataset(
+        self, ds: Dataset, tokenizer: AutoTokenizer, provider_config: HuggingFacePostTrainingConfig
+    ) -> Dataset:
+        """Preprocess the dataset with tokenizer for DPO."""
+        # DPOTrainer expects raw text, so we don't tokenize here
+        # Just return the dataset as is
+        return ds
+
+    def _run_training_sync(
+        self,
+        model: str,
+        provider_config: dict[str, Any],
+        dpo_config: dict[str, Any],
+        config: dict[str, Any],
+        output_dir_path: Path | None,
+    ) -> None:
+        """Synchronous wrapper for running DPO training process."""
+        import asyncio
+
+        logger.info("Starting DPO training process with async wrapper")
+        asyncio.run(
+            self._run_training(
+                model=model,
+                provider_config=provider_config,
+                dpo_config=dpo_config,
+                config=config,
+                output_dir_path=output_dir_path,
+            )
+        )
+
+    async def load_dataset(
+        self,
+        model: str,
+        config: TrainingConfig,
+        provider_config: HuggingFacePostTrainingConfig,
+    ) -> tuple[Dataset, Dataset, AutoTokenizer]:
+        """Load and prepare the dataset for DPO training."""
+        # Validate data config
+        if not config.data_config:
+            raise ValueError("DataConfig is required for DPO training")
+
+        # Load dataset
+        logger.info(f"Loading dataset: {config.data_config.dataset_id}")
+        rows = await load_rows_from_dataset(self.datasetio_api, config.data_config.dataset_id)
+        self.validate_dataset_format(rows)
+        logger.info(f"Loaded {len(rows)} rows from dataset")
+
+        # Initialize tokenizer
+        logger.info(f"Initializing tokenizer for model: {model}")
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model, **provider_config.model_specific_config)
+
+            # Set pad token to eos token if not present
+            if not tokenizer.pad_token:
+                tokenizer.pad_token = tokenizer.eos_token
+
+            # Set padding side to left for DPO
+            tokenizer.padding_side = "left"
+
+            # Set truncation side to right to keep the beginning of the sequence
+            tokenizer.truncation_side = "right"
+
+            # Set model max length to match provider config
+            tokenizer.model_max_length = provider_config.max_seq_length
+
+            logger.info("Tokenizer initialized successfully for DPO")
+        except Exception as e:
+            raise RuntimeError(f"Failed to initialize tokenizer: {str(e)}") from e
+
+        # Create and preprocess dataset
+        logger.info("Creating and preprocessing dataset for DPO")
+        try:
+            ds = self._create_dataset(rows, config, provider_config)
+            ds = self._preprocess_dataset(ds, tokenizer, provider_config)
+            logger.info(f"Dataset created with {len(ds)} examples")
+        except Exception as e:
+            raise ValueError(f"Failed to create dataset: {str(e)}") from e
+
+        # Split dataset
+        train_dataset, eval_dataset = split_dataset(ds)
+
+        return train_dataset, eval_dataset, tokenizer
+
+    def setup_training_args(
+        self,
+        config: TrainingConfig,
+        provider_config: HuggingFacePostTrainingConfig,
+        dpo_config: DPOAlignmentConfig,
+        device: torch.device,
+        output_dir_path: Path | None,
+        steps_per_epoch: int,
+    ) -> DPOConfig:
+        """Setup DPO training arguments."""
+        logger.info("Configuring DPO training arguments")
+        lr = 5e-7  # Lower learning rate for DPO
+        if config.optimizer_config:
+            lr = config.optimizer_config.lr
+            logger.info(f"Using custom learning rate: {lr}")
+
+        # Validate data config
+        if not config.data_config:
+            raise ValueError("DataConfig is required for training")
+        data_config = config.data_config
+
+        # Calculate steps and get save strategy
+        step_info = calculate_training_steps(steps_per_epoch, config)
+        save_strategy, eval_strategy = get_save_strategy(output_dir_path)
+
+        logger.info("DPO training configuration:")
+        logger.info(f"- DPO beta: {dpo_config.beta}")
+        logger.info(f"- DPO loss type: {provider_config.dpo_loss_type}")
+
+        # Calculate max prompt length as half of max sequence length
+        max_prompt_length = provider_config.max_seq_length // 2
+
+        return DPOConfig(
+            max_steps=step_info["max_steps"],
+            output_dir=str(output_dir_path) if output_dir_path is not None else None,
+            num_train_epochs=config.n_epochs,
+            per_device_train_batch_size=data_config.batch_size,
+            fp16=device.type == "cuda",
+            bf16=False,  # Causes CPU issues.
+            eval_strategy=eval_strategy,
+            use_cpu=True if device.type == "cpu" and not torch.backends.mps.is_available() else False,
+            save_strategy=save_strategy,
+            report_to="none",
+            max_length=provider_config.max_seq_length,
+            max_prompt_length=max_prompt_length,
+            gradient_accumulation_steps=config.gradient_accumulation_steps,
+            gradient_checkpointing=provider_config.gradient_checkpointing,
+            learning_rate=lr,
+            warmup_ratio=provider_config.warmup_ratio,
+            weight_decay=provider_config.weight_decay,
+            remove_unused_columns=False,
+            dataloader_pin_memory=provider_config.dataloader_pin_memory,
+            dataloader_num_workers=provider_config.dataloader_num_workers,
+            load_best_model_at_end=True if output_dir_path else False,
+            metric_for_best_model="eval_loss",
+            greater_is_better=False,
+            logging_steps=step_info["logging_steps"],
+            save_total_limit=provider_config.save_total_limit,
+            # DPO specific parameters
+            beta=dpo_config.beta,
+            loss_type=provider_config.dpo_loss_type,
+        )
+
+    def save_model(
+        self,
+        trainer: DPOTrainer,
+        output_dir_path: Path,
+    ) -> None:
+        """Save the trained DPO model."""
+        logger.info("Saving final DPO model")
+
+        save_path = output_dir_path / "dpo_model"
+        logger.info(f"Saving model to {save_path}")
+
+        # Save model and tokenizer
+        trainer.save_model(str(save_path))
+
+    async def _run_training(
+        self,
+        model: str,
+        provider_config: dict[str, Any],
+        dpo_config: dict[str, Any],
+        config: dict[str, Any],
+        output_dir_path: Path | None,
+    ) -> None:
+        """Run the DPO training process with signal handling."""
+
+        # Setup environment variables
+        setup_environment()
+
+        # Setup signal handlers
+        setup_signal_handlers()
+
+        # Convert config dicts back to objects
+        logger.info("Initializing configuration objects")
+        provider_config_obj = HuggingFacePostTrainingConfig(**provider_config)
+        config_obj = TrainingConfig(**config)
+        dpo_config_obj = DPOAlignmentConfig(**dpo_config)
+
+        # Initialize and validate device
+        device = setup_torch_device(provider_config_obj.device)
+        logger.info(f"Using device '{device}'")
+
+        # Load dataset and tokenizer
+        train_dataset, eval_dataset, tokenizer = await self.load_dataset(model, config_obj, provider_config_obj)
+
+        # Calculate steps per epoch
+        if not config_obj.data_config:
+            raise ValueError("DataConfig is required for training")
+        steps_per_epoch = len(train_dataset) // config_obj.data_config.batch_size
+
+        # Setup training arguments
+        training_args = self.setup_training_args(
+            config_obj,
+            provider_config_obj,
+            dpo_config_obj,
+            device,
+            output_dir_path,
+            steps_per_epoch,
+        )
+
+        # Load model and reference model
+        model_obj = load_model(model, device, provider_config_obj)
+        ref_model = None
+        if provider_config_obj.use_reference_model:
+            logger.info("Loading separate reference model for DPO")
+            ref_model = load_model(model, device, provider_config_obj)
+        else:
+            logger.info("Using shared reference model for DPO")
+
+        # Initialize DPO trainer
+        logger.info("Initializing DPOTrainer")
+        trainer = DPOTrainer(
+            model=model_obj,
+            ref_model=ref_model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            processing_class=tokenizer,
+        )
+
+        try:
+            # Train
+            logger.info("Starting DPO training")
+            trainer.train()
+            logger.info("DPO training completed successfully")
+
+            # Save final model if output directory is provided
+            if output_dir_path:
+                logger.info(f"Saving model to output directory: {output_dir_path}")
+                self.save_model(trainer, output_dir_path)
+                logger.info("Model save completed")
+
+        finally:
+            # Clean up resources
+            logger.info("Cleaning up resources")
+            if hasattr(trainer, "model"):
+                evacuate_model_from_device(trainer.model, device.type)
+            if ref_model:
+                evacuate_model_from_device(ref_model, device.type)
+            del trainer
+            del ref_model
+            gc.collect()
+            logger.info("Cleanup completed")
+            logger.info("DPO training process finishing successfully")
+
+    async def train(
+        self,
+        model: str,
+        output_dir: str | None,
+        job_uuid: str,
+        dpo_config: DPOAlignmentConfig,
+        config: TrainingConfig,
+        provider_config: HuggingFacePostTrainingConfig,
+    ) -> tuple[dict[str, Any], list[Checkpoint] | None]:
+        """Train a model using HuggingFace's DPOTrainer"""
+        # Initialize and validate device
+        device = setup_torch_device(provider_config.device)
+        logger.info(f"Using device '{device}'")
+
+        output_dir_path = None
+        if output_dir:
+            output_dir_path = Path(output_dir)
+
+        # Track memory stats
+        memory_stats = {
+            "initial": get_memory_stats(device),
+            "after_training": None,
+            "final": None,
+        }
+
+        # Validate data config
+        if not config.data_config:
+            raise ValueError("DataConfig is required for training")
+
+        # Train in a separate process
+        logger.info("Starting DPO training in separate process")
+        try:
+            # Setup multiprocessing for device
+            if device.type in ["cuda", "mps"]:
+                multiprocessing.set_start_method("spawn", force=True)
+
+            process = multiprocessing.Process(
+                target=self._run_training_sync,
+                kwargs={
+                    "model": model,
+                    "provider_config": provider_config.model_dump(),
+                    "dpo_config": dpo_config.model_dump(),
+                    "config": config.model_dump(),
+                    "output_dir_path": output_dir_path,
+                },
+            )
+            process.start()
+
+            # Monitor the process
+            while process.is_alive():
+                process.join(timeout=1)  # Check every second
+                if not process.is_alive():
+                    break
+
+            # Get the return code
+            if process.exitcode != 0:
+                raise RuntimeError(f"DPO training failed with exit code {process.exitcode}")
+
+            memory_stats["after_training"] = get_memory_stats(device)
+
+            checkpoints = []
+            if output_dir_path:
+                checkpoints = create_checkpoints(output_dir_path, job_uuid, model, config, "dpo_model")
+
+            return memory_stats, checkpoints if checkpoints else None
+        finally:
+            memory_stats["final"] = get_memory_stats(device)
+            gc.collect()

llama_stack/providers/inline/post_training/huggingface/utils.py

@@ -0,0 +1,269 @@
+# 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 logging
+import os
+import signal
+import sys
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+import psutil
+import torch
+from datasets import Dataset
+from transformers import AutoConfig, AutoModelForCausalLM
+
+from llama_stack.apis.datasetio import DatasetIO
+from llama_stack.apis.post_training import Checkpoint, TrainingConfig
+
+from .config import HuggingFacePostTrainingConfig
+
+logger = logging.getLogger(__name__)
+
+
+def setup_environment():
+    """Setup common environment variables for training."""
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    os.environ["MKL_THREADING_LAYER"] = "GNU"
+    os.environ["MKL_SERVICE_FORCE_INTEL"] = "0"
+    os.environ["MKL_NUM_THREADS"] = "1"
+
+
+def bytes_to_gb(to_convert: int) -> str:
+    """Converts memory stats to GB and formats to 2 decimal places.
+    Args:
+        to_convert: Memory value in bytes
+    Returns:
+        str: Memory value in GB formatted to 2 decimal places
+    """
+    return f"{(to_convert / (1024**3)):.2f}"
+
+
+def get_memory_stats(device: torch.device) -> dict[str, Any]:
+    """Get memory statistics for the given device."""
+    stats = {
+        "system_memory": {
+            "total": bytes_to_gb(psutil.virtual_memory().total),
+            "available": bytes_to_gb(psutil.virtual_memory().available),
+            "used": bytes_to_gb(psutil.virtual_memory().used),
+            "percent": psutil.virtual_memory().percent,
+        }
+    }
+
+    if device.type == "cuda":
+        stats["device_memory"] = {
+            "allocated": bytes_to_gb(torch.cuda.memory_allocated(device)),
+            "reserved": bytes_to_gb(torch.cuda.memory_reserved(device)),
+            "max_allocated": bytes_to_gb(torch.cuda.max_memory_allocated(device)),
+        }
+    elif device.type == "mps":
+        # MPS doesn't provide direct memory stats, but we can track system memory
+        stats["device_memory"] = {
+            "note": "MPS memory stats not directly available",
+            "system_memory_used": bytes_to_gb(psutil.virtual_memory().used),
+        }
+    elif device.type == "cpu":
+        # For CPU, we track process memory usage
+        process = psutil.Process()
+        stats["device_memory"] = {
+            "process_rss": bytes_to_gb(process.memory_info().rss),
+            "process_vms": bytes_to_gb(process.memory_info().vms),
+            "process_percent": process.memory_percent(),
+        }
+
+    return stats
+
+
+def setup_torch_device(device_str: str) -> torch.device:
+    """Initialize and validate a PyTorch device.
+    This function handles device initialization and validation for different device types:
+    - CUDA: Validates CUDA availability and handles device selection
+    - MPS: Validates MPS availability for Apple Silicon
+    - CPU: Basic validation
+    - HPU: Raises error as it's not supported
+    Args:
+        device_str: String specifying the device ('cuda', 'cpu', 'mps')
+    Returns:
+        torch.device: The initialized and validated device
+    Raises:
+        RuntimeError: If device initialization fails or device is not supported
+    """
+    try:
+        device = torch.device(device_str)
+    except RuntimeError as e:
+        raise RuntimeError(f"Error getting Torch Device {str(e)}") from e
+
+    # Validate device capabilities
+    if device.type == "cuda":
+        if not torch.cuda.is_available():
+            raise RuntimeError(
+                f"{device.type}: Torch has no CUDA/ROCm support or could not detect a compatible device."
+            )
+        if device.index is None:
+            device = torch.device(device.type, torch.cuda.current_device())
+    elif device.type == "mps":
+        if not torch.backends.mps.is_available():
+            raise RuntimeError(f"{device.type}: Torch has no MPS support or could not detect a compatible device.")
+    elif device.type == "hpu":
+        raise RuntimeError(f"{device.type}: training does not support Intel Gaudi.")
+
+    return device
+
+
+async def load_rows_from_dataset(datasetio_api: DatasetIO, dataset_id: str) -> list[dict[str, Any]]:
+    """Load dataset from llama stack dataset provider"""
+    try:
+        all_rows = await datasetio_api.iterrows(
+            dataset_id=dataset_id,
+            limit=-1,
+        )
+        if not isinstance(all_rows.data, list):
+            raise RuntimeError("Expected dataset data to be a list")
+        return all_rows.data
+    except Exception as e:
+        raise RuntimeError(f"Failed to load dataset: {str(e)}") from e
+
+
+def load_model(
+    model: str,
+    device: torch.device,
+    provider_config: HuggingFacePostTrainingConfig,
+) -> AutoModelForCausalLM:
+    """Load and initialize the model for training.
+    Args:
+        model: The model identifier to load
+        device: The device to load the model onto
+        provider_config: Provider-specific configuration
+    Returns:
+        The loaded and initialized model
+    Raises:
+        RuntimeError: If model loading fails
+    """
+    logger.info("Loading the base model")
+    try:
+        model_config = AutoConfig.from_pretrained(model, **provider_config.model_specific_config)
+        model_obj = AutoModelForCausalLM.from_pretrained(
+            model,
+            torch_dtype="auto" if device.type != "cpu" else "float32",
+            quantization_config=None,
+            config=model_config,
+            **provider_config.model_specific_config,
+        )
+        # Always move model to specified device
+        model_obj = model_obj.to(device)
+        logger.info(f"Model loaded and moved to device: {model_obj.device}")
+        return model_obj
+    except Exception as e:
+        raise RuntimeError(f"Failed to load model: {str(e)}") from e
+
+
+def split_dataset(ds: Dataset) -> tuple[Dataset, Dataset]:
+    """Split dataset into train and validation sets.
+    Args:
+        ds: Dataset to split
+    Returns:
+        tuple: (train_dataset, eval_dataset)
+    """
+    logger.info("Splitting dataset into train and validation sets")
+    train_val_split = ds.train_test_split(test_size=0.1, seed=42)
+    train_dataset = train_val_split["train"]
+    eval_dataset = train_val_split["test"]
+    logger.info(f"Split dataset into {len(train_dataset)} training and {len(eval_dataset)} validation examples")
+    return train_dataset, eval_dataset
+
+
+def setup_signal_handlers():
+    """Setup signal handlers for graceful shutdown."""
+
+    def signal_handler(signum, frame):
+        logger.info(f"Received signal {signum}, initiating graceful shutdown")
+        sys.exit(0)
+
+    signal.signal(signal.SIGTERM, signal_handler)
+    signal.signal(signal.SIGINT, signal_handler)
+
+
+def calculate_training_steps(steps_per_epoch: int, config: TrainingConfig) -> dict[str, int]:
+    """Calculate training steps and logging configuration.
+    Args:
+        steps_per_epoch: Number of training steps per epoch
+        config: Training configuration
+    Returns:
+        dict: Dictionary with calculated step values
+    """
+    total_steps = steps_per_epoch * config.n_epochs
+    max_steps = min(config.max_steps_per_epoch, total_steps)
+    logging_steps = max(1, steps_per_epoch // 50)  # Log 50 times per epoch
+
+    logger.info("Training configuration:")
+    logger.info(f"- Steps per epoch: {steps_per_epoch}")
+    logger.info(f"- Total steps: {total_steps}")
+    logger.info(f"- Max steps: {max_steps}")
+    logger.info(f"- Logging steps: {logging_steps}")
+
+    return {"total_steps": total_steps, "max_steps": max_steps, "logging_steps": logging_steps}
+
+
+def get_save_strategy(output_dir_path: Path | None) -> tuple[str, str]:
+    """Get save and evaluation strategy based on output directory.
+    Args:
+        output_dir_path: Optional path to save the model
+    Returns:
+        tuple: (save_strategy, eval_strategy)
+    """
+    if output_dir_path:
+        logger.info(f"Will save checkpoints to {output_dir_path}")
+        return "epoch", "epoch"
+    return "no", "no"
+
+
+def create_checkpoints(
+    output_dir_path: Path, job_uuid: str, model: str, config: TrainingConfig, final_model_name: str
+) -> list[Checkpoint]:
+    """Create checkpoint objects from training output.
+    Args:
+        output_dir_path: Path to the training output directory
+        job_uuid: Unique identifier for the training job
+        model: Model identifier
+        config: Training configuration
+        final_model_name: Name of the final model directory ("merged_model" for SFT, "dpo_model" for DPO)
+    Returns:
+        List of Checkpoint objects
+    """
+    checkpoints = []
+
+    # Add checkpoint directories
+    checkpoint_dirs = sorted(
+        [d for d in output_dir_path.glob("checkpoint-*") if d.is_dir()],
+        key=lambda x: int(x.name.split("-")[1]),
+    )
+
+    for epoch_number, checkpoint_dir in enumerate(checkpoint_dirs, start=1):
+        created_time = datetime.fromtimestamp(os.path.getctime(checkpoint_dir), tz=UTC)
+        checkpoint = Checkpoint(
+            identifier=checkpoint_dir.name,
+            created_at=created_time,
+            epoch=epoch_number,
+            post_training_job_id=job_uuid,
+            path=str(checkpoint_dir),
+        )
+        checkpoints.append(checkpoint)
+
+    # Add final model
+    final_model_path = output_dir_path / final_model_name
+    if final_model_path.exists():
+        training_type = "sft" if final_model_name == "merged_model" else "dpo"
+        checkpoint = Checkpoint(
+            identifier=f"{model}-{training_type}-{config.n_epochs}",
+            created_at=datetime.now(UTC),
+            epoch=config.n_epochs,
+            post_training_job_id=job_uuid,
+            path=str(final_model_path),
+        )
+        checkpoints.append(checkpoint)
+
+    return checkpoints

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

@@ -6,7 +6,7 @@
 
 from typing import Any
 
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import Api
 
 from .config import TorchtunePostTrainingConfig
 

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

@@ -23,12 +23,8 @@ from llama_stack.apis.post_training import (
 from llama_stack.providers.inline.post_training.torchtune.config import (
     TorchtunePostTrainingConfig,
 )
-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):
@@ -84,6 +80,10 @@ class TorchtunePostTrainingImpl:
         if isinstance(algorithm_config, LoraFinetuningConfig):
 
             async def handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb):
+                from llama_stack.providers.inline.post_training.torchtune.recipes.lora_finetuning_single_device import (
+                    LoraFinetuningSingleDevice,
+                )
+
                 on_log_message_cb("Starting Lora finetuning")
 
                 recipe = LoraFinetuningSingleDevice(
@@ -144,7 +144,6 @@ class TorchtunePostTrainingImpl:
         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) -> PostTrainingJobStatusResponse | None:
         job = self._scheduler.get_job(job_uuid)
 
@@ -171,11 +170,9 @@ class TorchtunePostTrainingImpl:
             resources_allocated=self._get_resources_allocated(job),
         )
 
-    @webmethod(route="/post-training/job/cancel")
     async def cancel_training_job(self, job_uuid: str) -> None:
         self._scheduler.cancel(job_uuid)
 
-    @webmethod(route="/post-training/job/artifacts")
     async def get_training_job_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse | 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

@@ -43,8 +43,8 @@ from llama_stack.apis.post_training import (
     QATFinetuningConfig,
     TrainingConfig,
 )
-from llama_stack.distribution.utils.config_dirs import DEFAULT_CHECKPOINT_DIR
-from llama_stack.distribution.utils.model_utils import model_local_dir
+from llama_stack.core.utils.config_dirs import DEFAULT_CHECKPOINT_DIR
+from llama_stack.core.utils.model_utils import model_local_dir
 from llama_stack.models.llama.sku_list import resolve_model
 from llama_stack.providers.inline.post_training.common.utils import evacuate_model_from_device
 from llama_stack.providers.inline.post_training.torchtune.common import utils

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

@@ -4,7 +4,9 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import logging
 import re
+import uuid
 from string import Template
 from typing import Any
 
@@ -20,8 +22,9 @@ from llama_stack.apis.safety import (
     SafetyViolation,
     ViolationLevel,
 )
+from llama_stack.apis.safety.safety import ModerationObject, ModerationObjectResults, OpenAICategories
 from llama_stack.apis.shields import Shield
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import Api
 from llama_stack.models.llama.datatypes import Role
 from llama_stack.models.llama.sku_types import CoreModelId
 from llama_stack.providers.datatypes import ShieldsProtocolPrivate
@@ -67,6 +70,31 @@ SAFETY_CATEGORIES_TO_CODE_MAP = {
     CAT_ELECTIONS: "S13",
     CAT_CODE_INTERPRETER_ABUSE: "S14",
 }
+SAFETY_CODE_TO_CATEGORIES_MAP = {v: k for k, v in SAFETY_CATEGORIES_TO_CODE_MAP.items()}
+
+OPENAI_TO_LLAMA_CATEGORIES_MAP = {
+    OpenAICategories.VIOLENCE: [CAT_VIOLENT_CRIMES],
+    OpenAICategories.VIOLENCE_GRAPHIC: [CAT_VIOLENT_CRIMES],
+    OpenAICategories.HARRASMENT: [CAT_CHILD_EXPLOITATION],
+    OpenAICategories.HARRASMENT_THREATENING: [CAT_VIOLENT_CRIMES, CAT_CHILD_EXPLOITATION],
+    OpenAICategories.HATE: [CAT_HATE],
+    OpenAICategories.HATE_THREATENING: [CAT_HATE, CAT_VIOLENT_CRIMES],
+    OpenAICategories.ILLICIT: [CAT_NON_VIOLENT_CRIMES],
+    OpenAICategories.ILLICIT_VIOLENT: [CAT_VIOLENT_CRIMES, CAT_INDISCRIMINATE_WEAPONS],
+    OpenAICategories.SEXUAL: [CAT_SEX_CRIMES, CAT_SEXUAL_CONTENT],
+    OpenAICategories.SEXUAL_MINORS: [CAT_CHILD_EXPLOITATION],
+    OpenAICategories.SELF_HARM: [CAT_SELF_HARM],
+    OpenAICategories.SELF_HARM_INTENT: [CAT_SELF_HARM],
+    OpenAICategories.SELF_HARM_INSTRUCTIONS: [CAT_SELF_HARM, CAT_SPECIALIZED_ADVICE],
+    # These are custom categories that are not in the OpenAI moderation categories
+    "custom/defamation": [CAT_DEFAMATION],
+    "custom/specialized_advice": [CAT_SPECIALIZED_ADVICE],
+    "custom/privacy_violation": [CAT_PRIVACY],
+    "custom/intellectual_property": [CAT_INTELLECTUAL_PROPERTY],
+    "custom/weapons": [CAT_INDISCRIMINATE_WEAPONS],
+    "custom/elections": [CAT_ELECTIONS],
+    "custom/code_interpreter_abuse": [CAT_CODE_INTERPRETER_ABUSE],
+}
 
 
 DEFAULT_LG_V3_SAFETY_CATEGORIES = [
@@ -146,8 +174,13 @@ class LlamaGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
         pass
 
     async def register_shield(self, shield: Shield) -> None:
-        # Allow any model to be registered as a shield
-        # The model will be validated during runtime when making inference calls
+        model_id = shield.provider_resource_id
+        if not model_id:
+            raise ValueError("Llama Guard shield must have a model id")
+
+    async def unregister_shield(self, identifier: str) -> None:
+        # LlamaGuard doesn't need to do anything special for unregistration
+        # The routing table handles the removal from the registry
         pass
 
     async def run_shield(
@@ -189,6 +222,34 @@ class LlamaGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
 
         return await impl.run(messages)
 
+    async def run_moderation(self, input: str | list[str], model: str) -> ModerationObject:
+        if isinstance(input, list):
+            messages = input.copy()
+        else:
+            messages = [input]
+
+        # convert to user messages format with role
+        messages = [UserMessage(content=m) for m in messages]
+
+        # Determine safety categories based on the model type
+        # For known Llama Guard models, use specific categories
+        if model in LLAMA_GUARD_MODEL_IDS:
+            # Use the mapped model for categories but the original model_id for inference
+            mapped_model = LLAMA_GUARD_MODEL_IDS[model]
+            safety_categories = MODEL_TO_SAFETY_CATEGORIES_MAP.get(mapped_model, DEFAULT_LG_V3_SAFETY_CATEGORIES)
+        else:
+            # For unknown models, use default Llama Guard 3 8B categories
+            safety_categories = DEFAULT_LG_V3_SAFETY_CATEGORIES + [CAT_CODE_INTERPRETER_ABUSE]
+
+        impl = LlamaGuardShield(
+            model=model,
+            inference_api=self.inference_api,
+            excluded_categories=self.config.excluded_categories,
+            safety_categories=safety_categories,
+        )
+
+        return await impl.run_moderation(messages)
+
 
 class LlamaGuardShield:
     def __init__(
@@ -335,3 +396,117 @@ class LlamaGuardShield:
             )
 
         raise ValueError(f"Unexpected response: {response}")
+
+    async def run_moderation(self, messages: list[Message]) -> ModerationObject:
+        if not messages:
+            return self.create_moderation_object(self.model)
+
+        # TODO: Add Image based support for OpenAI Moderations
+        shield_input_message = self.build_text_shield_input(messages)
+
+        response = await self.inference_api.openai_chat_completion(
+            model=self.model,
+            messages=[shield_input_message],
+            stream=False,
+        )
+        content = response.choices[0].message.content
+        content = content.strip()
+        return self.get_moderation_object(content)
+
+    def create_moderation_object(self, model: str, unsafe_code: str | None = None) -> ModerationObject:
+        """Create a ModerationObject for either safe or unsafe content.
+
+        Args:
+            model: The model name
+            unsafe_code: Optional comma-separated list of safety codes. If None, creates safe object.
+
+        Returns:
+            ModerationObject with appropriate configuration
+        """
+        # Set default values for safe case
+        categories = dict.fromkeys(OPENAI_TO_LLAMA_CATEGORIES_MAP.keys(), False)
+        category_scores = dict.fromkeys(OPENAI_TO_LLAMA_CATEGORIES_MAP.keys(), 1.0)
+        category_applied_input_types = {key: [] for key in OPENAI_TO_LLAMA_CATEGORIES_MAP.keys()}
+        flagged = False
+        user_message = None
+        metadata = {}
+
+        # Handle unsafe case
+        if unsafe_code:
+            unsafe_code_list = [code.strip() for code in unsafe_code.split(",")]
+            invalid_codes = [code for code in unsafe_code_list if code not in SAFETY_CODE_TO_CATEGORIES_MAP]
+            if invalid_codes:
+                logging.warning(f"Invalid safety codes returned: {invalid_codes}")
+                # just returning safe object, as we don't know what the invalid codes can map to
+                return ModerationObject(
+                    id=f"modr-{uuid.uuid4()}",
+                    model=model,
+                    results=[
+                        ModerationObjectResults(
+                            flagged=flagged,
+                            categories=categories,
+                            category_applied_input_types=category_applied_input_types,
+                            category_scores=category_scores,
+                            user_message=user_message,
+                            metadata=metadata,
+                        )
+                    ],
+                )
+
+            # Get OpenAI categories for the unsafe codes
+            openai_categories = []
+            for code in unsafe_code_list:
+                llama_guard_category = SAFETY_CODE_TO_CATEGORIES_MAP[code]
+                openai_categories.extend(
+                    k for k, v_l in OPENAI_TO_LLAMA_CATEGORIES_MAP.items() if llama_guard_category in v_l
+                )
+
+            # Update categories for unsafe content
+            categories = {k: k in openai_categories for k in OPENAI_TO_LLAMA_CATEGORIES_MAP}
+            category_scores = {k: 1.0 if k in openai_categories else 0.0 for k in OPENAI_TO_LLAMA_CATEGORIES_MAP}
+            category_applied_input_types = {
+                k: ["text"] if k in openai_categories else [] for k in OPENAI_TO_LLAMA_CATEGORIES_MAP
+            }
+            flagged = True
+            user_message = CANNED_RESPONSE_TEXT
+            metadata = {"violation_type": unsafe_code_list}
+
+        return ModerationObject(
+            id=f"modr-{uuid.uuid4()}",
+            model=model,
+            results=[
+                ModerationObjectResults(
+                    flagged=flagged,
+                    categories=categories,
+                    category_applied_input_types=category_applied_input_types,
+                    category_scores=category_scores,
+                    user_message=user_message,
+                    metadata=metadata,
+                )
+            ],
+        )
+
+    def is_content_safe(self, response: str, unsafe_code: str | None = None) -> bool:
+        """Check if content is safe based on response and unsafe code."""
+        if response.strip() == SAFE_RESPONSE:
+            return True
+
+        if unsafe_code:
+            unsafe_code_list = unsafe_code.split(",")
+            if set(unsafe_code_list).issubset(set(self.excluded_categories)):
+                return True
+
+        return False
+
+    def get_moderation_object(self, response: str) -> ModerationObject:
+        response = response.strip()
+        if self.is_content_safe(response):
+            return self.create_moderation_object(self.model)
+        unsafe_code = self.check_unsafe_response(response)
+        if not unsafe_code:
+            raise ValueError(f"Unexpected response: {response}")
+
+        if self.is_content_safe(response, unsafe_code):
+            return self.create_moderation_object(self.model)
+        else:
+            return self.create_moderation_object(self.model, unsafe_code)

llama_stack/providers/inline/safety/prompt_guard/prompt_guard.py

@@ -18,7 +18,7 @@ from llama_stack.apis.safety import (
     ViolationLevel,
 )
 from llama_stack.apis.shields import Shield
-from llama_stack.distribution.utils.model_utils import model_local_dir
+from llama_stack.core.utils.model_utils import model_local_dir
 from llama_stack.providers.datatypes import ShieldsProtocolPrivate
 from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,
@@ -46,6 +46,9 @@ class PromptGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
         if shield.provider_resource_id != PROMPT_GUARD_MODEL:
             raise ValueError(f"Only {PROMPT_GUARD_MODEL} is supported for Prompt Guard. ")
 
+    async def unregister_shield(self, identifier: str) -> None:
+        pass
+
     async def run_shield(
         self,
         shield_id: str,

llama_stack/providers/inline/scoring/basic/__init__.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 from typing import Any
 
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import Api
 
 from .config import BasicScoringConfig
 

llama_stack/providers/inline/scoring/basic/scoring.py

@@ -14,7 +14,7 @@ from llama_stack.apis.scoring import (
     ScoringResult,
 )
 from llama_stack.apis.scoring_functions import ScoringFn, ScoringFnParams
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import Api
 from llama_stack.providers.datatypes import ScoringFunctionsProtocolPrivate
 from llama_stack.providers.utils.common.data_schema_validator import (
     get_valid_schemas,

llama_stack/providers/inline/scoring/braintrust/__init__.py

@@ -7,7 +7,7 @@ from typing import Any
 
 from pydantic import BaseModel
 
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import Api
 
 from .config import BraintrustScoringConfig
 

llama_stack/providers/inline/scoring/braintrust/braintrust.py

@@ -29,8 +29,8 @@ from llama_stack.apis.scoring import (
     ScoringResultRow,
 )
 from llama_stack.apis.scoring_functions import ScoringFn, ScoringFnParams
-from llama_stack.distribution.datatypes import Api
-from llama_stack.distribution.request_headers import NeedsRequestProviderData
+from llama_stack.core.datatypes import Api
+from llama_stack.core.request_headers import NeedsRequestProviderData
 from llama_stack.providers.datatypes import ScoringFunctionsProtocolPrivate
 from llama_stack.providers.utils.common.data_schema_validator import (
     get_valid_schemas,

llama_stack/providers/inline/scoring/llm_as_judge/__init__.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 from typing import Any
 
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import Api
 
 from .config import LlmAsJudgeScoringConfig
 

llama_stack/providers/inline/scoring/llm_as_judge/scoring.py

@@ -15,7 +15,7 @@ from llama_stack.apis.scoring import (
     ScoringResult,
 )
 from llama_stack.apis.scoring_functions import ScoringFn, ScoringFnParams
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import Api
 from llama_stack.providers.datatypes import ScoringFunctionsProtocolPrivate
 from llama_stack.providers.utils.common.data_schema_validator import (
     get_valid_schemas,

llama_stack/providers/inline/telemetry/meta_reference/__init__.py

@@ -6,7 +6,7 @@
 
 from typing import Any
 
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import Api
 
 from .config import TelemetryConfig, TelemetrySink
 

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

@@ -9,7 +9,7 @@ from typing import Any
 
 from pydantic import BaseModel, Field, field_validator
 
-from llama_stack.distribution.utils.config_dirs import RUNTIME_BASE_DIR
+from llama_stack.core.utils.config_dirs import RUNTIME_BASE_DIR
 
 
 class TelemetrySink(StrEnum):

llama_stack/providers/inline/telemetry/meta_reference/console_span_processor.py

@@ -11,19 +11,9 @@ from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanProcessor
 from opentelemetry.trace.status import StatusCode
 
-# Colors for console output
-COLORS = {
-    "reset": "\033[0m",
-    "bold": "\033[1m",
-    "dim": "\033[2m",
-    "red": "\033[31m",
-    "green": "\033[32m",
-    "yellow": "\033[33m",
-    "blue": "\033[34m",
-    "magenta": "\033[35m",
-    "cyan": "\033[36m",
-    "white": "\033[37m",
-}
+from llama_stack.log import get_logger
+
+logger = get_logger(name="console_span_processor", category="telemetry")
 
 
 class ConsoleSpanProcessor(SpanProcessor):
@@ -35,34 +25,18 @@ class ConsoleSpanProcessor(SpanProcessor):
             return
 
         timestamp = datetime.fromtimestamp(span.start_time / 1e9, tz=UTC).strftime("%H:%M:%S.%f")[:-3]
-
-        print(
-            f"{COLORS['dim']}{timestamp}{COLORS['reset']} "
-            f"{COLORS['magenta']}[START]{COLORS['reset']} "
-            f"{COLORS['dim']}{span.name}{COLORS['reset']}"
-        )
+        logger.info(f"[dim]{timestamp}[/dim] [bold magenta][START][/bold magenta] [dim]{span.name}[/dim]")
 
     def on_end(self, span: ReadableSpan) -> None:
-        if span.attributes and span.attributes.get("__autotraced__"):
-            return
-
         timestamp = datetime.fromtimestamp(span.end_time / 1e9, tz=UTC).strftime("%H:%M:%S.%f")[:-3]
-
-        span_context = (
-            f"{COLORS['dim']}{timestamp}{COLORS['reset']} "
-            f"{COLORS['magenta']}[END]{COLORS['reset']} "
-            f"{COLORS['dim']}{span.name}{COLORS['reset']}"
-        )
-
+        span_context = f"[dim]{timestamp}[/dim] [bold magenta][END][/bold magenta] [dim]{span.name}[/dim]"
         if span.status.status_code == StatusCode.ERROR:
-            span_context += f"{COLORS['reset']} {COLORS['red']}[ERROR]{COLORS['reset']}"
+            span_context += " [bold red][ERROR][/bold red]"
         elif span.status.status_code != StatusCode.UNSET:
-            span_context += f"{COLORS['reset']} [{span.status.status_code}]"
-
+            span_context += f" [{span.status.status_code}]"
         duration_ms = (span.end_time - span.start_time) / 1e6
-        span_context += f"{COLORS['reset']} ({duration_ms:.2f}ms)"
-
-        print(span_context)
+        span_context += f" ({duration_ms:.2f}ms)"
+        logger.info(span_context)
 
         if self.print_attributes and span.attributes:
             for key, value in span.attributes.items():
@@ -71,31 +45,26 @@ class ConsoleSpanProcessor(SpanProcessor):
                 str_value = str(value)
                 if len(str_value) > 1000:
                     str_value = str_value[:997] + "..."
-                print(f"    {COLORS['dim']}{key}: {str_value}{COLORS['reset']}")
+                logger.info(f"    [dim]{key}[/dim]: {str_value}")
 
         for event in span.events:
             event_time = datetime.fromtimestamp(event.timestamp / 1e9, tz=UTC).strftime("%H:%M:%S.%f")[:-3]
-
             severity = event.attributes.get("severity", "info")
             message = event.attributes.get("message", event.name)
-            if isinstance(message, dict | list):
+            if isinstance(message, dict) or isinstance(message, list):
                 message = json.dumps(message, indent=2)
-
-            severity_colors = {
-                "error": f"{COLORS['bold']}{COLORS['red']}",
-                "warn": f"{COLORS['bold']}{COLORS['yellow']}",
-                "info": COLORS["white"],
-                "debug": COLORS["dim"],
-            }
-            msg_color = severity_colors.get(severity, COLORS["white"])
-
-            print(f" {event_time} {msg_color}[{severity.upper()}] {message}{COLORS['reset']}")
-
+            severity_color = {
+                "error": "red",
+                "warn": "yellow",
+                "info": "white",
+                "debug": "dim",
+            }.get(severity, "white")
+            logger.info(f" {event_time} [bold {severity_color}][{severity.upper()}][/bold {severity_color}] {message}")
             if event.attributes:
                 for key, value in event.attributes.items():
                     if key.startswith("__") or key in ["message", "severity"]:
                         continue
-                    print(f"   {COLORS['dim']}{key}: {value}{COLORS['reset']}")
+                    logger.info(f"[dim]{key}[/dim]: {value}")
 
     def shutdown(self) -> None:
         """Shutdown the processor."""

llama_stack/providers/inline/telemetry/meta_reference/telemetry.py

@@ -4,10 +4,13 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import logging
 import threading
 from typing import Any
 
 from opentelemetry import metrics, trace
+
+logger = logging.getLogger(__name__)
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk.metrics import MeterProvider
@@ -36,7 +39,7 @@ from llama_stack.apis.telemetry import (
     Trace,
     UnstructuredLogEvent,
 )
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import Api
 from llama_stack.providers.inline.telemetry.meta_reference.console_span_processor import (
     ConsoleSpanProcessor,
 )
@@ -110,7 +113,7 @@ class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
             if TelemetrySink.SQLITE in self.config.sinks:
                 trace.get_tracer_provider().add_span_processor(SQLiteSpanProcessor(self.config.sqlite_db_path))
             if TelemetrySink.CONSOLE in self.config.sinks:
-                trace.get_tracer_provider().add_span_processor(ConsoleSpanProcessor())
+                trace.get_tracer_provider().add_span_processor(ConsoleSpanProcessor(print_attributes=True))
 
         if TelemetrySink.OTEL_METRIC in self.config.sinks:
             self.meter = metrics.get_meter(__name__)
@@ -126,9 +129,11 @@ class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
         trace.get_tracer_provider().force_flush()
 
     async def log_event(self, event: Event, ttl_seconds: int = 604800) -> None:
+        logger.debug(f"DEBUG: log_event called with event type: {type(event).__name__}")
         if isinstance(event, UnstructuredLogEvent):
             self._log_unstructured(event, ttl_seconds)
         elif isinstance(event, MetricEvent):
+            logger.debug("DEBUG: Routing MetricEvent to _log_metric")
             self._log_metric(event)
         elif isinstance(event, StructuredLogEvent):
             self._log_structured(event, ttl_seconds)
@@ -188,6 +193,38 @@ class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
         return _GLOBAL_STORAGE["gauges"][name]
 
     def _log_metric(self, event: MetricEvent) -> None:
+        # Always log to console if console sink is enabled (debug)
+        if TelemetrySink.CONSOLE in self.config.sinks:
+            logger.debug(f"METRIC: {event.metric}={event.value} {event.unit} {event.attributes}")
+
+        # Add metric as an event to the current span
+        try:
+            with self._lock:
+                # Only try to add to span if we have a valid span_id
+                if event.span_id:
+                    try:
+                        span_id = int(event.span_id, 16)
+                        span = _GLOBAL_STORAGE["active_spans"].get(span_id)
+
+                        if span:
+                            timestamp_ns = int(event.timestamp.timestamp() * 1e9)
+                            span.add_event(
+                                name=f"metric.{event.metric}",
+                                attributes={
+                                    "value": event.value,
+                                    "unit": event.unit,
+                                    **(event.attributes or {}),
+                                },
+                                timestamp=timestamp_ns,
+                            )
+                    except (ValueError, KeyError):
+                        # Invalid span_id or span not found, but we already logged to console above
+                        pass
+        except Exception:
+            # Lock acquisition failed
+            logger.debug("Failed to acquire lock to add metric to span")
+
+        # Log to OpenTelemetry meter if available
         if self.meter is None:
             return
         if isinstance(event.value, int):

llama_stack/providers/inline/vector_io/chroma/__init__.py

@@ -16,6 +16,6 @@ async def get_provider_impl(config: ChromaVectorIOConfig, deps: dict[Api, Any]):
         ChromaVectorIOAdapter,
     )
 
-    impl = ChromaVectorIOAdapter(config, deps[Api.inference])
+    impl = ChromaVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
     await impl.initialize()
     return impl

llama_stack/providers/inline/vector_io/chroma/config.py

@@ -6,12 +6,25 @@
 
 from typing import Any
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+
+from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig
+from llama_stack.schema_utils import json_schema_type
 
 
+@json_schema_type
 class ChromaVectorIOConfig(BaseModel):
     db_path: str
+    kvstore: KVStoreConfig = Field(description="Config for KV store backend")
 
     @classmethod
-    def sample_run_config(cls, db_path: str = "${env.CHROMADB_PATH}", **kwargs: Any) -> dict[str, Any]:
-        return {"db_path": db_path}
+    def sample_run_config(
+        cls, __distro_dir__: str, db_path: str = "${env.CHROMADB_PATH}", **kwargs: Any
+    ) -> dict[str, Any]:
+        return {
+            "db_path": db_path,
+            "kvstore": SqliteKVStoreConfig.sample_run_config(
+                __distro_dir__=__distro_dir__,
+                db_name="chroma_inline_registry.db",
+            ),
+        }

llama_stack/providers/inline/vector_io/faiss/faiss.py

@@ -15,6 +15,7 @@ import faiss
 import numpy as np
 from numpy.typing import NDArray
 
+from llama_stack.apis.common.errors import VectorStoreNotFoundError
 from llama_stack.apis.files import Files
 from llama_stack.apis.inference import Inference, InterleavedContent
 from llama_stack.apis.vector_dbs import VectorDB
@@ -55,6 +56,11 @@ class FaissIndex(EmbeddingIndex):
         self.kvstore = kvstore
         self.bank_id = bank_id
 
+        # A list of chunk id's in the same order as they are in the index,
+        # must be updated when chunks are added or removed
+        self.chunk_id_lock = asyncio.Lock()
+        self.chunk_ids: list[Any] = []
+
     @classmethod
     async def create(cls, dimension: int, kvstore: KVStore | None = None, bank_id: str | None = None):
         instance = cls(dimension, kvstore, bank_id)
@@ -75,6 +81,7 @@ class FaissIndex(EmbeddingIndex):
             buffer = io.BytesIO(base64.b64decode(data["faiss_index"]))
             try:
                 self.index = faiss.deserialize_index(np.load(buffer, allow_pickle=False))
+                self.chunk_ids = [chunk.chunk_id for chunk in self.chunk_by_index.values()]
             except Exception as e:
                 logger.debug(e, exc_info=True)
                 raise ValueError(
@@ -114,11 +121,33 @@ class FaissIndex(EmbeddingIndex):
         for i, chunk in enumerate(chunks):
             self.chunk_by_index[indexlen + i] = chunk
 
-        self.index.add(np.array(embeddings).astype(np.float32))
+        async with self.chunk_id_lock:
+            self.index.add(np.array(embeddings).astype(np.float32))
+            self.chunk_ids.extend([chunk.chunk_id for chunk in chunks])
 
         # Save updated index
         await self._save_index()
 
+    async def delete_chunk(self, chunk_id: str) -> None:
+        if chunk_id not in self.chunk_ids:
+            return
+
+        async with self.chunk_id_lock:
+            index = self.chunk_ids.index(chunk_id)
+            self.index.remove_ids(np.array([index]))
+
+            new_chunk_by_index = {}
+            for idx, chunk in self.chunk_by_index.items():
+                # Shift all chunks after the removed chunk to the left
+                if idx > index:
+                    new_chunk_by_index[idx - 1] = chunk
+                else:
+                    new_chunk_by_index[idx] = chunk
+            self.chunk_by_index = new_chunk_by_index
+            self.chunk_ids.pop(index)
+
+        await self._save_index()
+
     async def query_vector(
         self,
         embedding: NDArray,
@@ -131,8 +160,11 @@ class FaissIndex(EmbeddingIndex):
         for d, i in zip(distances[0], indices[0], strict=False):
             if i < 0:
                 continue
+            score = 1.0 / float(d) if d != 0 else float("inf")
+            if score < score_threshold:
+                continue
             chunks.append(self.chunk_by_index[int(i)])
-            scores.append(1.0 / float(d) if d != 0 else float("inf"))
+            scores.append(score)
 
         return QueryChunksResponse(chunks=chunks, scores=scores)
 
@@ -181,8 +213,8 @@ class FaissVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPr
             )
             self.cache[vector_db.identifier] = index
 
-        # Load existing OpenAI vector stores using the mixin method
-        self.openai_vector_stores = await self._load_openai_vector_stores()
+        # Load existing OpenAI vector stores into the in-memory cache
+        await self.initialize_openai_vector_stores()
 
     async def shutdown(self) -> None:
         # Cleanup if needed
@@ -257,74 +289,12 @@ class FaissVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPr
     ) -> QueryChunksResponse:
         index = self.cache.get(vector_db_id)
         if index is None:
-            raise ValueError(f"Vector DB {vector_db_id} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         return await index.query_chunks(query, params)
 
-    # OpenAI Vector Store Mixin abstract method implementations
-    async def _save_openai_vector_store(self, store_id: str, store_info: dict[str, Any]) -> None:
-        """Save vector store metadata to kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}"
-        await self.kvstore.set(key=key, value=json.dumps(store_info))
-
-    async def _load_openai_vector_stores(self) -> dict[str, dict[str, Any]]:
-        """Load all vector store metadata from kvstore."""
-        assert self.kvstore is not None
-        start_key = OPENAI_VECTOR_STORES_PREFIX
-        end_key = f"{OPENAI_VECTOR_STORES_PREFIX}\xff"
-        stored_openai_stores = await self.kvstore.values_in_range(start_key, end_key)
-
-        stores = {}
-        for store_data in stored_openai_stores:
-            store_info = json.loads(store_data)
-            stores[store_info["id"]] = store_info
-        return stores
-
-    async def _update_openai_vector_store(self, store_id: str, store_info: dict[str, Any]) -> None:
-        """Update vector store metadata in kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}"
-        await self.kvstore.set(key=key, value=json.dumps(store_info))
-
-    async def _delete_openai_vector_store_from_storage(self, store_id: str) -> None:
-        """Delete vector store metadata from kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}"
-        await self.kvstore.delete(key)
-
-    async def _save_openai_vector_store_file(
-        self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
-    ) -> None:
-        """Save vector store file metadata to kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
-        await self.kvstore.set(key=key, value=json.dumps(file_info))
-        content_key = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}"
-        await self.kvstore.set(key=content_key, value=json.dumps(file_contents))
-
-    async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
-        """Load vector store file metadata from kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
-        stored_data = await self.kvstore.get(key)
-        return json.loads(stored_data) if stored_data else {}
-
-    async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
-        """Load vector store file contents from kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}"
-        stored_data = await self.kvstore.get(key)
-        return json.loads(stored_data) if stored_data else []
-
-    async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
-        """Update vector store file metadata in kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
-        await self.kvstore.set(key=key, value=json.dumps(file_info))
-
-    async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
-        """Delete vector store file metadata from kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
-        await self.kvstore.delete(key)
+    async def delete_chunks(self, store_id: str, chunk_ids: list[str]) -> None:
+        """Delete a chunk from a faiss index"""
+        faiss_index = self.cache[store_id].index
+        for chunk_id in chunk_ids:
+            await faiss_index.delete_chunk(chunk_id)

llama_stack/providers/inline/vector_io/qdrant/__init__.py

@@ -4,14 +4,18 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_stack.providers.datatypes import Api, ProviderSpec
+from typing import Any
+
+from llama_stack.providers.datatypes import Api
 
 from .config import QdrantVectorIOConfig
 
 
-async def get_adapter_impl(config: QdrantVectorIOConfig, deps: dict[Api, ProviderSpec]):
+async def get_provider_impl(config: QdrantVectorIOConfig, deps: dict[Api, Any]):
     from llama_stack.providers.remote.vector_io.qdrant.qdrant import QdrantVectorIOAdapter
 
-    impl = QdrantVectorIOAdapter(config, deps[Api.inference])
+    assert isinstance(config, QdrantVectorIOConfig), f"Unexpected config type: {type(config)}"
+    files_api = deps.get(Api.files)
+    impl = QdrantVectorIOAdapter(config, deps[Api.inference], files_api)
     await impl.initialize()
     return impl

llama_stack/providers/inline/vector_io/qdrant/config.py

@@ -9,15 +9,23 @@ from typing import Any
 
 from pydantic import BaseModel
 
+from llama_stack.providers.utils.kvstore.config import (
+    KVStoreConfig,
+    SqliteKVStoreConfig,
+)
 from llama_stack.schema_utils import json_schema_type
 
 
 @json_schema_type
 class QdrantVectorIOConfig(BaseModel):
     path: str
+    kvstore: KVStoreConfig
 
     @classmethod
     def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]:
         return {
             "path": "${env.QDRANT_PATH:=~/.llama/" + __distro_dir__ + "}/" + "qdrant.db",
+            "kvstore": SqliteKVStoreConfig.sample_run_config(
+                __distro_dir__=__distro_dir__, db_name="qdrant_registry.db"
+            ),
         }

llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py

@@ -5,8 +5,8 @@
 # the root directory of this source tree.
 
 import asyncio
-import json
 import logging
+import re
 import sqlite3
 import struct
 from typing import Any
@@ -15,6 +15,7 @@ import numpy as np
 import sqlite_vec
 from numpy.typing import NDArray
 
+from llama_stack.apis.common.errors import VectorStoreNotFoundError
 from llama_stack.apis.files import Files
 from llama_stack.apis.inference import Inference
 from llama_stack.apis.vector_dbs import VectorDB
@@ -117,6 +118,10 @@ def _rrf_rerank(
     return rrf_scores
 
 
+def _make_sql_identifier(name: str) -> str:
+    return re.sub(r"[^a-zA-Z0-9_]", "_", name)
+
+
 class SQLiteVecIndex(EmbeddingIndex):
     """
     An index implementation that stores embeddings in a SQLite virtual table using sqlite-vec.
@@ -130,9 +135,9 @@ class SQLiteVecIndex(EmbeddingIndex):
         self.dimension = dimension
         self.db_path = db_path
         self.bank_id = bank_id
-        self.metadata_table = f"chunks_{bank_id}".replace("-", "_")
-        self.vector_table = f"vec_chunks_{bank_id}".replace("-", "_")
-        self.fts_table = f"fts_chunks_{bank_id}".replace("-", "_")
+        self.metadata_table = _make_sql_identifier(f"chunks_{bank_id}")
+        self.vector_table = _make_sql_identifier(f"vec_chunks_{bank_id}")
+        self.fts_table = _make_sql_identifier(f"fts_chunks_{bank_id}")
         self.kvstore = kvstore
 
     @classmethod
@@ -148,14 +153,14 @@ class SQLiteVecIndex(EmbeddingIndex):
             try:
                 # Create the table to store chunk metadata.
                 cur.execute(f"""
-                    CREATE TABLE IF NOT EXISTS {self.metadata_table} (
+                    CREATE TABLE IF NOT EXISTS [{self.metadata_table}] (
                         id TEXT PRIMARY KEY,
                         chunk TEXT
                     );
                 """)
                 # Create the virtual table for embeddings.
                 cur.execute(f"""
-                    CREATE VIRTUAL TABLE IF NOT EXISTS {self.vector_table}
+                    CREATE VIRTUAL TABLE IF NOT EXISTS [{self.vector_table}]
                     USING vec0(embedding FLOAT[{self.dimension}], id TEXT);
                 """)
                 connection.commit()
@@ -163,7 +168,7 @@ class SQLiteVecIndex(EmbeddingIndex):
                 # based on query. Implementation of the change on client side will allow passing the search_mode option
                 # during initialization to make it easier to create the table that is required.
                 cur.execute(f"""
-                            CREATE VIRTUAL TABLE IF NOT EXISTS {self.fts_table}
+                            CREATE VIRTUAL TABLE IF NOT EXISTS [{self.fts_table}]
                             USING fts5(id, content);
                         """)
                 connection.commit()
@@ -178,9 +183,9 @@ class SQLiteVecIndex(EmbeddingIndex):
             connection = _create_sqlite_connection(self.db_path)
             cur = connection.cursor()
             try:
-                cur.execute(f"DROP TABLE IF EXISTS {self.metadata_table};")
-                cur.execute(f"DROP TABLE IF EXISTS {self.vector_table};")
-                cur.execute(f"DROP TABLE IF EXISTS {self.fts_table};")
+                cur.execute(f"DROP TABLE IF EXISTS [{self.metadata_table}];")
+                cur.execute(f"DROP TABLE IF EXISTS [{self.vector_table}];")
+                cur.execute(f"DROP TABLE IF EXISTS [{self.fts_table}];")
                 connection.commit()
             finally:
                 cur.close()
@@ -212,7 +217,7 @@ class SQLiteVecIndex(EmbeddingIndex):
                     metadata_data = [(chunk.chunk_id, chunk.model_dump_json()) for chunk in batch_chunks]
                     cur.executemany(
                         f"""
-                        INSERT INTO {self.metadata_table} (id, chunk)
+                        INSERT INTO [{self.metadata_table}] (id, chunk)
                         VALUES (?, ?)
                         ON CONFLICT(id) DO UPDATE SET chunk = excluded.chunk;
                         """,
@@ -230,7 +235,7 @@ class SQLiteVecIndex(EmbeddingIndex):
                         for chunk, emb in zip(batch_chunks, batch_embeddings, strict=True)
                     ]
                     cur.executemany(
-                        f"INSERT INTO {self.vector_table} (id, embedding) VALUES (?, ?);",
+                        f"INSERT INTO [{self.vector_table}] (id, embedding) VALUES (?, ?);",
                         embedding_data,
                     )
 
@@ -238,13 +243,13 @@ class SQLiteVecIndex(EmbeddingIndex):
                     fts_data = [(chunk.chunk_id, chunk.content) for chunk in batch_chunks]
                     # DELETE existing entries with same IDs (FTS5 doesn't support ON CONFLICT)
                     cur.executemany(
-                        f"DELETE FROM {self.fts_table} WHERE id = ?;",
+                        f"DELETE FROM [{self.fts_table}] WHERE id = ?;",
                         [(row[0],) for row in fts_data],
                     )
 
                     # INSERT new entries
                     cur.executemany(
-                        f"INSERT INTO {self.fts_table} (id, content) VALUES (?, ?);",
+                        f"INSERT INTO [{self.fts_table}] (id, content) VALUES (?, ?);",
                         fts_data,
                     )
 
@@ -280,8 +285,8 @@ class SQLiteVecIndex(EmbeddingIndex):
                 emb_blob = serialize_vector(emb_list)
                 query_sql = f"""
                     SELECT m.id, m.chunk, v.distance
-                    FROM {self.vector_table} AS v
-                    JOIN {self.metadata_table} AS m ON m.id = v.id
+                    FROM [{self.vector_table}] AS v
+                    JOIN [{self.metadata_table}] AS m ON m.id = v.id
                     WHERE v.embedding MATCH ? AND k = ?
                     ORDER BY v.distance;
                 """
@@ -322,9 +327,9 @@ class SQLiteVecIndex(EmbeddingIndex):
             cur = connection.cursor()
             try:
                 query_sql = f"""
-                    SELECT DISTINCT m.id, m.chunk, bm25({self.fts_table}) AS score
-                    FROM {self.fts_table} AS f
-                    JOIN {self.metadata_table} AS m ON m.id = f.id
+                    SELECT DISTINCT m.id, m.chunk, bm25([{self.fts_table}]) AS score
+                    FROM [{self.fts_table}] AS f
+                    JOIN [{self.metadata_table}] AS m ON m.id = f.id
                     WHERE f.content MATCH ?
                     ORDER BY score ASC
                     LIMIT ?;
@@ -421,6 +426,35 @@ class SQLiteVecIndex(EmbeddingIndex):
 
         return QueryChunksResponse(chunks=chunks, scores=scores)
 
+    async def delete_chunk(self, chunk_id: str) -> None:
+        """Remove a chunk from the SQLite vector store."""
+
+        def _delete_chunk():
+            connection = _create_sqlite_connection(self.db_path)
+            cur = connection.cursor()
+            try:
+                cur.execute("BEGIN TRANSACTION")
+
+                # Delete from metadata table
+                cur.execute(f"DELETE FROM {self.metadata_table} WHERE id = ?", (chunk_id,))
+
+                # Delete from vector table
+                cur.execute(f"DELETE FROM {self.vector_table} WHERE id = ?", (chunk_id,))
+
+                # Delete from FTS table
+                cur.execute(f"DELETE FROM {self.fts_table} WHERE id = ?", (chunk_id,))
+
+                connection.commit()
+            except Exception as e:
+                connection.rollback()
+                logger.error(f"Error deleting chunk {chunk_id}: {e}")
+                raise
+            finally:
+                cur.close()
+                connection.close()
+
+        await asyncio.to_thread(_delete_chunk)
+
 
 class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPrivate):
     """
@@ -452,8 +486,8 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoc
             )
             self.cache[vector_db.identifier] = VectorDBWithIndex(vector_db, index, self.inference_api)
 
-        # load any existing OpenAI vector stores
-        self.openai_vector_stores = await self._load_openai_vector_stores()
+        # Load existing OpenAI vector stores into the in-memory cache
+        await self.initialize_openai_vector_stores()
 
     async def shutdown(self) -> None:
         # nothing to do since we don't maintain a persistent connection
@@ -475,11 +509,11 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoc
             return self.cache[vector_db_id]
 
         if self.vector_db_store is None:
-            raise ValueError(f"Vector DB {vector_db_id} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         vector_db = self.vector_db_store.get_vector_db(vector_db_id)
         if not vector_db:
-            raise ValueError(f"Vector DB {vector_db_id} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         index = VectorDBWithIndex(
             vector_db=vector_db,
@@ -501,179 +535,10 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoc
         await self.cache[vector_db_id].index.delete()
         del self.cache[vector_db_id]
 
-    # OpenAI Vector Store Mixin abstract method implementations
-    async def _save_openai_vector_store(self, store_id: str, store_info: dict[str, Any]) -> None:
-        """Save vector store metadata to SQLite database."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}"
-        await self.kvstore.set(key=key, value=json.dumps(store_info))
-        self.openai_vector_stores[store_id] = store_info
-
-    async def _load_openai_vector_stores(self) -> dict[str, dict[str, Any]]:
-        """Load all vector store metadata from SQLite database."""
-        assert self.kvstore is not None
-        start_key = OPENAI_VECTOR_STORES_PREFIX
-        end_key = f"{OPENAI_VECTOR_STORES_PREFIX}\xff"
-        stored_openai_stores = await self.kvstore.values_in_range(start_key, end_key)
-        stores = {}
-        for store_data in stored_openai_stores:
-            store_info = json.loads(store_data)
-            stores[store_info["id"]] = store_info
-        return stores
-
-    async def _update_openai_vector_store(self, store_id: str, store_info: dict[str, Any]) -> None:
-        """Update vector store metadata in SQLite database."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}"
-        await self.kvstore.set(key=key, value=json.dumps(store_info))
-        self.openai_vector_stores[store_id] = store_info
-
-    async def _delete_openai_vector_store_from_storage(self, store_id: str) -> None:
-        """Delete vector store metadata from SQLite database."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}"
-        await self.kvstore.delete(key)
-        if store_id in self.openai_vector_stores:
-            del self.openai_vector_stores[store_id]
-
-    async def _save_openai_vector_store_file(
-        self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
-    ) -> None:
-        """Save vector store file metadata to SQLite database."""
-
-        def _create_or_store():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                # Create a table to persist OpenAI vector store files.
-                cur.execute("""
-                    CREATE TABLE IF NOT EXISTS openai_vector_store_files (
-                        store_id TEXT,
-                        file_id TEXT,
-                        metadata TEXT,
-                        PRIMARY KEY (store_id, file_id)
-                    );
-                """)
-                cur.execute("""
-                    CREATE TABLE IF NOT EXISTS openai_vector_store_files_contents (
-                        store_id TEXT,
-                        file_id TEXT,
-                        contents TEXT,
-                        PRIMARY KEY (store_id, file_id)
-                    );
-                """)
-                connection.commit()
-                cur.execute(
-                    "INSERT OR REPLACE INTO openai_vector_store_files (store_id, file_id, metadata) VALUES (?, ?, ?)",
-                    (store_id, file_id, json.dumps(file_info)),
-                )
-                cur.execute(
-                    "INSERT OR REPLACE INTO openai_vector_store_files_contents (store_id, file_id, contents) VALUES (?, ?, ?)",
-                    (store_id, file_id, json.dumps(file_contents)),
-                )
-                connection.commit()
-            except Exception as e:
-                logger.error(f"Error saving openai vector store file {store_id} {file_id}: {e}")
-                raise
-            finally:
-                cur.close()
-                connection.close()
-
-        try:
-            await asyncio.to_thread(_create_or_store)
-        except Exception as e:
-            logger.error(f"Error saving openai vector store file {store_id} {file_id}: {e}")
-            raise
-
-    async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
-        """Load vector store file metadata from SQLite database."""
-
-        def _load():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                cur.execute(
-                    "SELECT metadata FROM openai_vector_store_files WHERE store_id = ? AND file_id = ?",
-                    (store_id, file_id),
-                )
-                row = cur.fetchone()
-                if row is None:
-                    return None
-                (metadata,) = row
-                return metadata
-            finally:
-                cur.close()
-                connection.close()
-
-        stored_data = await asyncio.to_thread(_load)
-        return json.loads(stored_data) if stored_data else {}
-
-    async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
-        """Load vector store file contents from SQLite database."""
-
-        def _load():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                cur.execute(
-                    "SELECT contents FROM openai_vector_store_files_contents WHERE store_id = ? AND file_id = ?",
-                    (store_id, file_id),
-                )
-                row = cur.fetchone()
-                if row is None:
-                    return None
-                (contents,) = row
-                return contents
-            finally:
-                cur.close()
-                connection.close()
-
-        stored_contents = await asyncio.to_thread(_load)
-        return json.loads(stored_contents) if stored_contents else []
-
-    async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
-        """Update vector store file metadata in SQLite database."""
-
-        def _update():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                cur.execute(
-                    "UPDATE openai_vector_store_files SET metadata = ? WHERE store_id = ? AND file_id = ?",
-                    (json.dumps(file_info), store_id, file_id),
-                )
-                connection.commit()
-            finally:
-                cur.close()
-                connection.close()
-
-        await asyncio.to_thread(_update)
-
-    async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
-        """Delete vector store file metadata from SQLite database."""
-
-        def _delete():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                cur.execute(
-                    "DELETE FROM openai_vector_store_files WHERE store_id = ? AND file_id = ?", (store_id, file_id)
-                )
-                cur.execute(
-                    "DELETE FROM openai_vector_store_files_contents WHERE store_id = ? AND file_id = ?",
-                    (store_id, file_id),
-                )
-                connection.commit()
-            finally:
-                cur.close()
-                connection.close()
-
-        await asyncio.to_thread(_delete)
-
     async def insert_chunks(self, vector_db_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None:
         index = await self._get_and_cache_vector_db_index(vector_db_id)
         if not index:
-            raise ValueError(f"Vector DB {vector_db_id} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
         # The VectorDBWithIndex helper is expected to compute embeddings via the inference_api
         # and then call our index's add_chunks.
         await index.insert_chunks(chunks)
@@ -683,5 +548,15 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoc
     ) -> QueryChunksResponse:
         index = await self._get_and_cache_vector_db_index(vector_db_id)
         if not index:
-            raise ValueError(f"Vector DB {vector_db_id} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
         return await index.query_chunks(query, params)
+
+    async def delete_chunks(self, store_id: str, chunk_ids: list[str]) -> None:
+        """Delete a chunk from a sqlite_vec index."""
+        index = await self._get_and_cache_vector_db_index(store_id)
+        if not index:
+            raise VectorStoreNotFoundError(store_id)
+
+        for chunk_id in chunk_ids:
+            # Use the index's delete_chunk method
+            await index.index.delete_chunk(chunk_id)

llama_stack/providers/registry/inference.py

@@ -37,16 +37,6 @@ def available_providers() -> list[ProviderSpec]:
             config_class="llama_stack.providers.inline.inference.meta_reference.MetaReferenceInferenceConfig",
             description="Meta's reference implementation of inference with support for various model formats and optimization techniques.",
         ),
-        InlineProviderSpec(
-            api=Api.inference,
-            provider_type="inline::vllm",
-            pip_packages=[
-                "vllm",
-            ],
-            module="llama_stack.providers.inline.inference.vllm",
-            config_class="llama_stack.providers.inline.inference.vllm.VLLMConfig",
-            description="vLLM inference provider for high-performance model serving with PagedAttention and continuous batching.",
-        ),
         InlineProviderSpec(
             api=Api.inference,
             provider_type="inline::sentence-transformers",
@@ -234,17 +224,6 @@ def available_providers() -> list[ProviderSpec]:
                 description="Groq inference provider for ultra-fast inference using Groq's LPU technology.",
             ),
         ),
-        remote_provider_spec(
-            api=Api.inference,
-            adapter=AdapterSpec(
-                adapter_type="fireworks-openai-compat",
-                pip_packages=["litellm"],
-                module="llama_stack.providers.remote.inference.fireworks_openai_compat",
-                config_class="llama_stack.providers.remote.inference.fireworks_openai_compat.config.FireworksCompatConfig",
-                provider_data_validator="llama_stack.providers.remote.inference.fireworks_openai_compat.config.FireworksProviderDataValidator",
-                description="Fireworks AI OpenAI-compatible provider for using Fireworks models with OpenAI API format.",
-            ),
-        ),
         remote_provider_spec(
             api=Api.inference,
             adapter=AdapterSpec(
@@ -256,50 +235,6 @@ def available_providers() -> list[ProviderSpec]:
                 description="Llama OpenAI-compatible provider for using Llama models with OpenAI API format.",
             ),
         ),
-        remote_provider_spec(
-            api=Api.inference,
-            adapter=AdapterSpec(
-                adapter_type="together-openai-compat",
-                pip_packages=["litellm"],
-                module="llama_stack.providers.remote.inference.together_openai_compat",
-                config_class="llama_stack.providers.remote.inference.together_openai_compat.config.TogetherCompatConfig",
-                provider_data_validator="llama_stack.providers.remote.inference.together_openai_compat.config.TogetherProviderDataValidator",
-                description="Together AI OpenAI-compatible provider for using Together models with OpenAI API format.",
-            ),
-        ),
-        remote_provider_spec(
-            api=Api.inference,
-            adapter=AdapterSpec(
-                adapter_type="groq-openai-compat",
-                pip_packages=["litellm"],
-                module="llama_stack.providers.remote.inference.groq_openai_compat",
-                config_class="llama_stack.providers.remote.inference.groq_openai_compat.config.GroqCompatConfig",
-                provider_data_validator="llama_stack.providers.remote.inference.groq_openai_compat.config.GroqProviderDataValidator",
-                description="Groq OpenAI-compatible provider for using Groq models with OpenAI API format.",
-            ),
-        ),
-        remote_provider_spec(
-            api=Api.inference,
-            adapter=AdapterSpec(
-                adapter_type="sambanova-openai-compat",
-                pip_packages=["litellm"],
-                module="llama_stack.providers.remote.inference.sambanova_openai_compat",
-                config_class="llama_stack.providers.remote.inference.sambanova_openai_compat.config.SambaNovaCompatConfig",
-                provider_data_validator="llama_stack.providers.remote.inference.sambanova_openai_compat.config.SambaNovaProviderDataValidator",
-                description="SambaNova OpenAI-compatible provider for using SambaNova models with OpenAI API format.",
-            ),
-        ),
-        remote_provider_spec(
-            api=Api.inference,
-            adapter=AdapterSpec(
-                adapter_type="cerebras-openai-compat",
-                pip_packages=["litellm"],
-                module="llama_stack.providers.remote.inference.cerebras_openai_compat",
-                config_class="llama_stack.providers.remote.inference.cerebras_openai_compat.config.CerebrasCompatConfig",
-                provider_data_validator="llama_stack.providers.remote.inference.cerebras_openai_compat.config.CerebrasProviderDataValidator",
-                description="Cerebras OpenAI-compatible provider for using Cerebras models with OpenAI API format.",
-            ),
-        ),
         remote_provider_spec(
             api=Api.inference,
             adapter=AdapterSpec(

llama_stack/providers/registry/vector_io.py

@@ -395,7 +395,7 @@ That means you'll get fast and efficient vector retrieval.
 To use PGVector in your Llama Stack project, follow these steps:
 
 1. Install the necessary dependencies.
-2. Configure your Llama Stack project to use Faiss.
+2. Configure your Llama Stack project to use pgvector. (e.g. remote::pgvector).
 3. Start storing and querying vectors.
 
 ## Installation
@@ -410,6 +410,7 @@ See [PGVector's documentation](https://github.com/pgvector/pgvector) for more de
 """,
             ),
             api_dependencies=[Api.inference],
+            optional_api_dependencies=[Api.files],
         ),
         remote_provider_spec(
             Api.vector_io,
@@ -497,6 +498,7 @@ See [Weaviate's documentation](https://weaviate.io/developers/weaviate) for more
             module="llama_stack.providers.inline.vector_io.qdrant",
             config_class="llama_stack.providers.inline.vector_io.qdrant.QdrantVectorIOConfig",
             api_dependencies=[Api.inference],
+            optional_api_dependencies=[Api.files],
             description=r"""
 [Qdrant](https://qdrant.tech/documentation/) is an inline and remote vector database provider for Llama Stack. It
 allows you to store and query vectors directly in memory.
@@ -553,6 +555,7 @@ Please refer to the inline provider documentation.
 """,
             ),
             api_dependencies=[Api.inference],
+            optional_api_dependencies=[Api.files],
         ),
         remote_provider_spec(
             Api.vector_io,

llama_stack/providers/remote/datasetio/huggingface/huggingface.py

@@ -6,8 +6,6 @@
 from typing import Any
 from urllib.parse import parse_qs, urlparse
 
-import datasets as hf_datasets
-
 from llama_stack.apis.common.responses import PaginatedResponse
 from llama_stack.apis.datasetio import DatasetIO
 from llama_stack.apis.datasets import Dataset
@@ -73,6 +71,8 @@ class HuggingfaceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
         start_index: int | None = None,
         limit: int | None = None,
     ) -> PaginatedResponse:
+        import datasets as hf_datasets
+
         dataset_def = self.dataset_infos[dataset_id]
         path, params = parse_hf_params(dataset_def)
         loaded_dataset = hf_datasets.load_dataset(path, **params)
@@ -81,6 +81,8 @@ class HuggingfaceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
         return paginate_records(records, start_index, limit)
 
     async def append_rows(self, dataset_id: str, rows: list[dict[str, Any]]) -> None:
+        import datasets as hf_datasets
+
         dataset_def = self.dataset_infos[dataset_id]
         path, params = parse_hf_params(dataset_def)
         loaded_dataset = hf_datasets.load_dataset(path, **params)

llama_stack/providers/remote/datasetio/nvidia/README.md

@@ -20,7 +20,7 @@ This provider enables dataset management using NVIDIA's NeMo Customizer service.
 Build the NVIDIA environment:
 
 ```bash
-llama stack build --template nvidia --image-type conda
+llama stack build --distro nvidia --image-type venv
 ```
 
 ### Basic Usage using the LlamaStack Python Client
@@ -34,7 +34,7 @@ os.environ["NVIDIA_API_KEY"] = "your-api-key"
 os.environ["NVIDIA_CUSTOMIZER_URL"] = "http://nemo.test"
 os.environ["NVIDIA_DATASET_NAMESPACE"] = "default"
 os.environ["NVIDIA_PROJECT_ID"] = "test-project"
-from llama_stack.distribution.library_client import LlamaStackAsLibraryClient
+from llama_stack.core.library_client import LlamaStackAsLibraryClient
 
 client = LlamaStackAsLibraryClient("nvidia")
 client.initialize()

llama_stack/providers/remote/eval/nvidia/__init__.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 from typing import Any
 
-from llama_stack.distribution.datatypes import Api
+from llama_stack.core.datatypes import Api
 
 from .config import NVIDIAEvalConfig
 

llama_stack/providers/remote/inference/anthropic/anthropic.py

@@ -15,6 +15,7 @@ class AnthropicInferenceAdapter(LiteLLMOpenAIMixin):
         LiteLLMOpenAIMixin.__init__(
             self,
             MODEL_ENTRIES,
+            litellm_provider_name="anthropic",
             api_key_from_config=config.api_key,
             provider_data_api_key_field="anthropic_api_key",
         )

llama_stack/providers/remote/inference/anthropic/config.py

@@ -26,7 +26,7 @@ class AnthropicConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.ANTHROPIC_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.ANTHROPIC_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "api_key": api_key,
         }

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

@@ -10,9 +10,9 @@ from llama_stack.providers.utils.inference.model_registry import (
 )
 
 LLM_MODEL_IDS = [
-    "anthropic/claude-3-5-sonnet-latest",
-    "anthropic/claude-3-7-sonnet-latest",
-    "anthropic/claude-3-5-haiku-latest",
+    "claude-3-5-sonnet-latest",
+    "claude-3-7-sonnet-latest",
+    "claude-3-5-haiku-latest",
 ]
 
 SAFETY_MODELS_ENTRIES = []
@@ -21,17 +21,17 @@ MODEL_ENTRIES = (
     [ProviderModelEntry(provider_model_id=m) for m in LLM_MODEL_IDS]
     + [
         ProviderModelEntry(
-            provider_model_id="anthropic/voyage-3",
+            provider_model_id="voyage-3",
             model_type=ModelType.embedding,
             metadata={"embedding_dimension": 1024, "context_length": 32000},
         ),
         ProviderModelEntry(
-            provider_model_id="anthropic/voyage-3-lite",
+            provider_model_id="voyage-3-lite",
             model_type=ModelType.embedding,
             metadata={"embedding_dimension": 512, "context_length": 32000},
         ),
         ProviderModelEntry(
-            provider_model_id="anthropic/voyage-code-3",
+            provider_model_id="voyage-code-3",
             model_type=ModelType.embedding,
             metadata={"embedding_dimension": 1024, "context_length": 32000},
         ),

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

@@ -63,18 +63,20 @@ class BedrockInferenceAdapter(
     def __init__(self, config: BedrockConfig) -> None:
         ModelRegistryHelper.__init__(self, MODEL_ENTRIES)
         self._config = config
-
-        self._client = create_bedrock_client(config)
+        self._client = None
 
     @property
     def client(self) -> BaseClient:
+        if self._client is None:
+            self._client = create_bedrock_client(self._config)
         return self._client
 
     async def initialize(self) -> None:
         pass
 
     async def shutdown(self) -> None:
-        self.client.close()
+        if self._client is not None:
+            self._client.close()
 
     async def completion(
         self,

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

@@ -65,6 +65,7 @@ class CerebrasInferenceAdapter(
         )
         self.config = config
 
+        # TODO: make this use provider data, etc. like other providers
         self.client = AsyncCerebras(
             base_url=self.config.base_url,
             api_key=self.config.api_key.get_secret_value(),

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

@@ -26,7 +26,7 @@ class CerebrasImplConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.CEREBRAS_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.CEREBRAS_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "base_url": DEFAULT_BASE_URL,
             "api_key": api_key,

llama_stack/providers/remote/inference/cerebras_openai_compat/__init__.py

@@ -1,17 +0,0 @@
-# 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 llama_stack.apis.inference import InferenceProvider
-
-from .config import CerebrasCompatConfig
-
-
-async def get_adapter_impl(config: CerebrasCompatConfig, _deps) -> InferenceProvider:
-    # import dynamically so the import is used only when it is needed
-    from .cerebras import CerebrasCompatInferenceAdapter
-
-    adapter = CerebrasCompatInferenceAdapter(config)
-    return adapter

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

@@ -1,30 +0,0 @@
-# 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 llama_stack.providers.remote.inference.cerebras_openai_compat.config import CerebrasCompatConfig
-from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-
-from ..cerebras.models import MODEL_ENTRIES
-
-
-class CerebrasCompatInferenceAdapter(LiteLLMOpenAIMixin):
-    _config: CerebrasCompatConfig
-
-    def __init__(self, config: CerebrasCompatConfig):
-        LiteLLMOpenAIMixin.__init__(
-            self,
-            model_entries=MODEL_ENTRIES,
-            api_key_from_config=config.api_key,
-            provider_data_api_key_field="cerebras_api_key",
-            openai_compat_api_base=config.openai_compat_api_base,
-        )
-        self.config = config
-
-    async def initialize(self):
-        await super().initialize()
-
-    async def shutdown(self):
-        await super().shutdown()

llama_stack/providers/remote/inference/cerebras_openai_compat/config.py

@@ -1,38 +0,0 @@
-# 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
-
-from pydantic import BaseModel, Field
-
-from llama_stack.schema_utils import json_schema_type
-
-
-class CerebrasProviderDataValidator(BaseModel):
-    cerebras_api_key: str | None = Field(
-        default=None,
-        description="API key for Cerebras models",
-    )
-
-
-@json_schema_type
-class CerebrasCompatConfig(BaseModel):
-    api_key: str | None = Field(
-        default=None,
-        description="The Cerebras API key",
-    )
-
-    openai_compat_api_base: str = Field(
-        default="https://api.cerebras.ai/v1",
-        description="The URL for the Cerebras API server",
-    )
-
-    @classmethod
-    def sample_run_config(cls, api_key: str = "${env.CEREBRAS_API_KEY}", **kwargs) -> dict[str, Any]:
-        return {
-            "openai_compat_api_base": "https://api.cerebras.ai/v1",
-            "api_key": api_key,
-        }

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

@@ -25,8 +25,8 @@ class DatabricksImplConfig(BaseModel):
     @classmethod
     def sample_run_config(
         cls,
-        url: str = "${env.DATABRICKS_URL}",
-        api_token: str = "${env.DATABRICKS_API_TOKEN}",
+        url: str = "${env.DATABRICKS_URL:=}",
+        api_token: str = "${env.DATABRICKS_API_TOKEN:=}",
         **kwargs: Any,
     ) -> dict[str, Any]:
         return {

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

@@ -6,13 +6,14 @@
 
 from typing import Any
 
-from pydantic import BaseModel, Field, SecretStr
+from pydantic import Field, SecretStr
 
+from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
 from llama_stack.schema_utils import json_schema_type
 
 
 @json_schema_type
-class FireworksImplConfig(BaseModel):
+class FireworksImplConfig(RemoteInferenceProviderConfig):
     url: str = Field(
         default="https://api.fireworks.ai/inference/v1",
         description="The URL for the Fireworks server",
@@ -23,7 +24,7 @@ class FireworksImplConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.FIREWORKS_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.FIREWORKS_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "url": "https://api.fireworks.ai/inference/v1",
             "api_key": api_key,

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

@@ -39,7 +39,7 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
-from llama_stack.distribution.request_headers import NeedsRequestProviderData
+from llama_stack.core.request_headers import NeedsRequestProviderData
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.model_registry import (
     ModelRegistryHelper,
@@ -70,7 +70,7 @@ logger = get_logger(name=__name__, category="inference")
 
 class FireworksInferenceAdapter(ModelRegistryHelper, Inference, NeedsRequestProviderData):
     def __init__(self, config: FireworksImplConfig) -> None:
-        ModelRegistryHelper.__init__(self, MODEL_ENTRIES)
+        ModelRegistryHelper.__init__(self, MODEL_ENTRIES, config.allowed_models)
         self.config = config
 
     async def initialize(self) -> None:

llama_stack/providers/remote/inference/fireworks_openai_compat/__init__.py

@@ -1,17 +0,0 @@
-# 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 llama_stack.apis.inference import InferenceProvider
-
-from .config import FireworksCompatConfig
-
-
-async def get_adapter_impl(config: FireworksCompatConfig, _deps) -> InferenceProvider:
-    # import dynamically so the import is used only when it is needed
-    from .fireworks import FireworksCompatInferenceAdapter
-
-    adapter = FireworksCompatInferenceAdapter(config)
-    return adapter

llama_stack/providers/remote/inference/fireworks_openai_compat/config.py

@@ -1,38 +0,0 @@
-# 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
-
-from pydantic import BaseModel, Field
-
-from llama_stack.schema_utils import json_schema_type
-
-
-class FireworksProviderDataValidator(BaseModel):
-    fireworks_api_key: str | None = Field(
-        default=None,
-        description="API key for Fireworks models",
-    )
-
-
-@json_schema_type
-class FireworksCompatConfig(BaseModel):
-    api_key: str | None = Field(
-        default=None,
-        description="The Fireworks API key",
-    )
-
-    openai_compat_api_base: str = Field(
-        default="https://api.fireworks.ai/inference/v1",
-        description="The URL for the Fireworks API server",
-    )
-
-    @classmethod
-    def sample_run_config(cls, api_key: str = "${env.FIREWORKS_API_KEY}", **kwargs) -> dict[str, Any]:
-        return {
-            "openai_compat_api_base": "https://api.fireworks.ai/inference/v1",
-            "api_key": api_key,
-        }

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

@@ -1,30 +0,0 @@
-# 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 llama_stack.providers.remote.inference.fireworks_openai_compat.config import FireworksCompatConfig
-from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-
-from ..fireworks.models import MODEL_ENTRIES
-
-
-class FireworksCompatInferenceAdapter(LiteLLMOpenAIMixin):
-    _config: FireworksCompatConfig
-
-    def __init__(self, config: FireworksCompatConfig):
-        LiteLLMOpenAIMixin.__init__(
-            self,
-            model_entries=MODEL_ENTRIES,
-            api_key_from_config=config.api_key,
-            provider_data_api_key_field="fireworks_api_key",
-            openai_compat_api_base=config.openai_compat_api_base,
-        )
-        self.config = config
-
-    async def initialize(self):
-        await super().initialize()
-
-    async def shutdown(self):
-        await super().shutdown()

llama_stack/providers/remote/inference/gemini/config.py

@@ -26,7 +26,7 @@ class GeminiConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.GEMINI_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.GEMINI_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "api_key": api_key,
         }

llama_stack/providers/remote/inference/gemini/gemini.py

@@ -15,6 +15,7 @@ class GeminiInferenceAdapter(LiteLLMOpenAIMixin):
         LiteLLMOpenAIMixin.__init__(
             self,
             MODEL_ENTRIES,
+            litellm_provider_name="gemini",
             api_key_from_config=config.api_key,
             provider_data_api_key_field="gemini_api_key",
         )

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

@@ -10,11 +10,11 @@ from llama_stack.providers.utils.inference.model_registry import (
 )
 
 LLM_MODEL_IDS = [
-    "gemini/gemini-1.5-flash",
-    "gemini/gemini-1.5-pro",
-    "gemini/gemini-2.0-flash",
-    "gemini/gemini-2.5-flash",
-    "gemini/gemini-2.5-pro",
+    "gemini-1.5-flash",
+    "gemini-1.5-pro",
+    "gemini-2.0-flash",
+    "gemini-2.5-flash",
+    "gemini-2.5-pro",
 ]
 
 SAFETY_MODELS_ENTRIES = []
@@ -23,7 +23,7 @@ MODEL_ENTRIES = (
     [ProviderModelEntry(provider_model_id=m) for m in LLM_MODEL_IDS]
     + [
         ProviderModelEntry(
-            provider_model_id="gemini/text-embedding-004",
+            provider_model_id="text-embedding-004",
             model_type=ModelType.embedding,
             metadata={"embedding_dimension": 768, "context_length": 2048},
         ),

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

@@ -32,7 +32,7 @@ class GroqConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.GROQ_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.GROQ_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "url": "https://api.groq.com",
             "api_key": api_key,

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

@@ -34,28 +34,23 @@ class GroqInferenceAdapter(LiteLLMOpenAIMixin):
         LiteLLMOpenAIMixin.__init__(
             self,
             model_entries=MODEL_ENTRIES,
+            litellm_provider_name="groq",
             api_key_from_config=config.api_key,
             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
+        return AsyncOpenAI(
+            base_url=f"{self.config.url}/openai/v1",
+            api_key=self.get_api_key(),
+        )
 
     async def openai_chat_completion(
         self,
@@ -102,7 +97,7 @@ class GroqInferenceAdapter(LiteLLMOpenAIMixin):
             tool_choice = "required"
 
         params = await prepare_openai_completion_params(
-            model=model_obj.provider_resource_id.replace("groq/", ""),
+            model=model_obj.provider_resource_id,
             messages=messages,
             frequency_penalty=frequency_penalty,
             function_call=function_call,

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

@@ -14,19 +14,19 @@ SAFETY_MODELS_ENTRIES = []
 
 MODEL_ENTRIES = [
     build_hf_repo_model_entry(
-        "groq/llama3-8b-8192",
+        "llama3-8b-8192",
         CoreModelId.llama3_1_8b_instruct.value,
     ),
     build_model_entry(
-        "groq/llama-3.1-8b-instant",
+        "llama-3.1-8b-instant",
         CoreModelId.llama3_1_8b_instruct.value,
     ),
     build_hf_repo_model_entry(
-        "groq/llama3-70b-8192",
+        "llama3-70b-8192",
         CoreModelId.llama3_70b_instruct.value,
     ),
     build_hf_repo_model_entry(
-        "groq/llama-3.3-70b-versatile",
+        "llama-3.3-70b-versatile",
         CoreModelId.llama3_3_70b_instruct.value,
     ),
     # Groq only contains a preview version for llama-3.2-3b
@@ -34,23 +34,15 @@ MODEL_ENTRIES = [
     # to pass the test fixture
     # TODO(aidand): Replace this with a stable model once Groq supports it
     build_hf_repo_model_entry(
-        "groq/llama-3.2-3b-preview",
+        "llama-3.2-3b-preview",
         CoreModelId.llama3_2_3b_instruct.value,
     ),
     build_hf_repo_model_entry(
-        "groq/llama-4-scout-17b-16e-instruct",
+        "meta-llama/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",
+        "meta-llama/llama-4-maverick-17b-128e-instruct",
         CoreModelId.llama4_maverick_17b_128e_instruct.value,
     ),
 ] + SAFETY_MODELS_ENTRIES

llama_stack/providers/remote/inference/groq_openai_compat/__init__.py

@@ -1,17 +0,0 @@
-# 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 llama_stack.apis.inference import InferenceProvider
-
-from .config import GroqCompatConfig
-
-
-async def get_adapter_impl(config: GroqCompatConfig, _deps) -> InferenceProvider:
-    # import dynamically so the import is used only when it is needed
-    from .groq import GroqCompatInferenceAdapter
-
-    adapter = GroqCompatInferenceAdapter(config)
-    return adapter

llama_stack/providers/remote/inference/groq_openai_compat/config.py

@@ -1,38 +0,0 @@
-# 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
-
-from pydantic import BaseModel, Field
-
-from llama_stack.schema_utils import json_schema_type
-
-
-class GroqProviderDataValidator(BaseModel):
-    groq_api_key: str | None = Field(
-        default=None,
-        description="API key for Groq models",
-    )
-
-
-@json_schema_type
-class GroqCompatConfig(BaseModel):
-    api_key: str | None = Field(
-        default=None,
-        description="The Groq API key",
-    )
-
-    openai_compat_api_base: str = Field(
-        default="https://api.groq.com/openai/v1",
-        description="The URL for the Groq API server",
-    )
-
-    @classmethod
-    def sample_run_config(cls, api_key: str = "${env.GROQ_API_KEY}", **kwargs) -> dict[str, Any]:
-        return {
-            "openai_compat_api_base": "https://api.groq.com/openai/v1",
-            "api_key": api_key,
-        }

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

@@ -1,30 +0,0 @@
-# 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 llama_stack.providers.remote.inference.groq_openai_compat.config import GroqCompatConfig
-from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-
-from ..groq.models import MODEL_ENTRIES
-
-
-class GroqCompatInferenceAdapter(LiteLLMOpenAIMixin):
-    _config: GroqCompatConfig
-
-    def __init__(self, config: GroqCompatConfig):
-        LiteLLMOpenAIMixin.__init__(
-            self,
-            model_entries=MODEL_ENTRIES,
-            api_key_from_config=config.api_key,
-            provider_data_api_key_field="groq_api_key",
-            openai_compat_api_base=config.openai_compat_api_base,
-        )
-        self.config = config
-
-    async def initialize(self):
-        await super().initialize()
-
-    async def shutdown(self):
-        await super().shutdown()

llama_stack/providers/remote/inference/llama_openai_compat/llama.py

@@ -3,30 +3,53 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+import logging
 
-from llama_stack.providers.remote.inference.llama_openai_compat.config import (
-    LlamaCompatConfig,
-)
-from llama_stack.providers.utils.inference.litellm_openai_mixin import (
-    LiteLLMOpenAIMixin,
-)
+from llama_stack.providers.remote.inference.llama_openai_compat.config import LlamaCompatConfig
+from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
+from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 from .models import MODEL_ENTRIES
 
+logger = logging.getLogger(__name__)
+
+
+class LlamaCompatInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin):
+    """
+    Llama API Inference Adapter for Llama Stack.
+
+    Note: The inheritance order is important here. OpenAIMixin must come before
+    LiteLLMOpenAIMixin to ensure that OpenAIMixin.check_model_availability()
+    is used instead of ModelRegistryHelper.check_model_availability().
+
+    - OpenAIMixin.check_model_availability() queries the Llama API to check if a model exists
+    - ModelRegistryHelper.check_model_availability() (inherited by LiteLLMOpenAIMixin) just returns False and shows a warning
+    """
 
-class LlamaCompatInferenceAdapter(LiteLLMOpenAIMixin):
     _config: LlamaCompatConfig
 
     def __init__(self, config: LlamaCompatConfig):
         LiteLLMOpenAIMixin.__init__(
             self,
             model_entries=MODEL_ENTRIES,
+            litellm_provider_name="meta_llama",
             api_key_from_config=config.api_key,
             provider_data_api_key_field="llama_api_key",
             openai_compat_api_base=config.openai_compat_api_base,
         )
         self.config = config
 
+    # Delegate the client data handling get_api_key method to LiteLLMOpenAIMixin
+    get_api_key = LiteLLMOpenAIMixin.get_api_key
+
+    def get_base_url(self) -> str:
+        """
+        Get the base URL for OpenAI mixin.
+
+        :return: The Llama API base URL
+        """
+        return self.config.openai_compat_api_base
+
     async def initialize(self):
         await super().initialize()
 

llama_stack/providers/remote/inference/nvidia/NVIDIA.md

@@ -18,7 +18,7 @@ This provider enables running inference using NVIDIA NIM.
 Build the NVIDIA environment:
 
 ```bash
-llama stack build --template nvidia --image-type conda
+llama stack build --distro nvidia --image-type venv
 ```
 
 ### Basic Usage using the LlamaStack Python Client
@@ -33,7 +33,7 @@ os.environ["NVIDIA_API_KEY"] = (
 )
 os.environ["NVIDIA_BASE_URL"] = "http://nim.test"  # NIM URL
 
-from llama_stack.distribution.library_client import LlamaStackAsLibraryClient
+from llama_stack.core.library_client import LlamaStackAsLibraryClient
 
 client = LlamaStackAsLibraryClient("nvidia")
 client.initialize()
@@ -42,8 +42,8 @@ client.initialize()
 ### Create Completion
 
 ```python
-response = client.completion(
-    model_id="meta-llama/Llama-3.1-8b-Instruct",
+response = client.inference.completion(
+    model_id="meta-llama/Llama-3.1-8B-Instruct",
     content="Complete the sentence using one word: Roses are red, violets are :",
     stream=False,
     sampling_params={
@@ -56,8 +56,8 @@ print(f"Response: {response.content}")
 ### Create Chat Completion
 
 ```python
-response = client.chat_completion(
-    model_id="meta-llama/Llama-3.1-8b-Instruct",
+response = client.inference.chat_completion(
+    model_id="meta-llama/Llama-3.1-8B-Instruct",
     messages=[
         {
             "role": "system",
@@ -78,8 +78,10 @@ print(f"Response: {response.completion_message.content}")
 
 ### Create Embeddings
 ```python
-response = client.embeddings(
-    model_id="meta-llama/Llama-3.1-8b-Instruct", contents=["foo", "bar", "baz"]
+response = client.inference.embeddings(
+    model_id="nvidia/llama-3.2-nv-embedqa-1b-v2",
+    contents=["What is the capital of France?"],
+    task_type="query",
 )
 print(f"Embeddings: {response.embeddings}")
-```
+```

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

@@ -7,10 +7,8 @@
 import logging
 import warnings
 from collections.abc import AsyncIterator
-from functools import lru_cache
-from typing import Any
 
-from openai import APIConnectionError, AsyncOpenAI, BadRequestError
+from openai import APIConnectionError, BadRequestError
 
 from llama_stack.apis.common.content_types import (
     InterleavedContent,
@@ -29,31 +27,21 @@ from llama_stack.apis.inference import (
     Inference,
     LogProbConfig,
     Message,
-    OpenAIChatCompletion,
-    OpenAIChatCompletionChunk,
-    OpenAICompletion,
-    OpenAIEmbeddingsResponse,
-    OpenAIMessageParam,
-    OpenAIResponseFormatParam,
     ResponseFormat,
     SamplingParams,
     TextTruncation,
     ToolChoice,
     ToolConfig,
 )
-from llama_stack.apis.models import Model, ModelType
 from llama_stack.models.llama.datatypes import ToolDefinition, ToolPromptFormat
-from llama_stack.providers.utils.inference import (
-    ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR,
-)
 from llama_stack.providers.utils.inference.model_registry import (
     ModelRegistryHelper,
 )
 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.openai_mixin import OpenAIMixin
 from llama_stack.providers.utils.inference.prompt_adapter import content_has_media
 
 from . import NVIDIAConfig
@@ -69,7 +57,20 @@ from .utils import _is_nvidia_hosted
 logger = logging.getLogger(__name__)
 
 
-class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
+class NVIDIAInferenceAdapter(OpenAIMixin, Inference, ModelRegistryHelper):
+    """
+    NVIDIA Inference Adapter for Llama Stack.
+
+    Note: The inheritance order is important here. OpenAIMixin must come before
+    ModelRegistryHelper to ensure that OpenAIMixin.check_model_availability()
+    is used instead of ModelRegistryHelper.check_model_availability(). It also
+    must come before Inference to ensure that OpenAIMixin methods are available
+    in the Inference interface.
+
+    - OpenAIMixin.check_model_availability() queries the NVIDIA API to check if a model exists
+    - ModelRegistryHelper.check_model_availability() just returns False and shows a warning
+    """
+
     def __init__(self, config: NVIDIAConfig) -> None:
         # TODO(mf): filter by available models
         ModelRegistryHelper.__init__(self, model_entries=MODEL_ENTRIES)
@@ -93,49 +94,21 @@ class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
 
         self._config = config
 
-    @lru_cache  # noqa: B019
-    def _get_client(self, provider_model_id: str) -> AsyncOpenAI:
+    def get_api_key(self) -> str:
         """
-        For hosted models, https://integrate.api.nvidia.com/v1 is the primary base_url. However,
-        some models are hosted on different URLs. This function returns the appropriate client
-        for the given provider_model_id.
+        Get the API key for OpenAI mixin.
 
-        This relies on lru_cache and self._default_client to avoid creating a new client for each request
-        or for each model that is hosted on https://integrate.api.nvidia.com/v1.
-
-        :param provider_model_id: The provider model ID
-        :return: An OpenAI client
+        :return: The NVIDIA API key
         """
+        return self._config.api_key.get_secret_value() if self._config.api_key else "NO KEY"
 
-        @lru_cache  # noqa: B019
-        def _get_client_for_base_url(base_url: str) -> AsyncOpenAI:
-            """
-            Maintain a single OpenAI client per base_url.
-            """
-            return AsyncOpenAI(
-                base_url=base_url,
-                api_key=(self._config.api_key.get_secret_value() if self._config.api_key else "NO KEY"),
-                timeout=self._config.timeout,
-            )
+    def get_base_url(self) -> str:
+        """
+        Get the base URL for OpenAI mixin.
 
-        special_model_urls = {
-            "meta/llama-3.2-11b-vision-instruct": "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-11b-vision-instruct",
-            "meta/llama-3.2-90b-vision-instruct": "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct",
-        }
-
-        base_url = f"{self._config.url}/v1" if self._config.append_api_version else self._config.url
-
-        if _is_nvidia_hosted(self._config) and provider_model_id in special_model_urls:
-            base_url = special_model_urls[provider_model_id]
-        return _get_client_for_base_url(base_url)
-
-    async def _get_provider_model_id(self, model_id: str) -> str:
-        if not self.model_store:
-            raise RuntimeError("Model store is not set")
-        model = await self.model_store.get_model(model_id)
-        if model is None:
-            raise ValueError(f"Model {model_id} is unknown")
-        return model.provider_model_id
+        :return: The NVIDIA API base URL
+        """
+        return f"{self._config.url}/v1" if self._config.append_api_version else self._config.url
 
     async def completion(
         self,
@@ -169,7 +142,7 @@ class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
         )
 
         try:
-            response = await self._get_client(provider_model_id).completions.create(**request)
+            response = await self.client.completions.create(**request)
         except APIConnectionError as e:
             raise ConnectionError(f"Failed to connect to NVIDIA NIM at {self._config.url}: {e}") from e
 
@@ -222,7 +195,7 @@ class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
             extra_body["input_type"] = task_type_options[task_type]
 
         try:
-            response = await self._get_client(provider_model_id).embeddings.create(
+            response = await self.client.embeddings.create(
                 model=provider_model_id,
                 input=input,
                 extra_body=extra_body,
@@ -237,16 +210,6 @@ class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
         #
         return EmbeddingsResponse(embeddings=[embedding.embedding for embedding in response.data])
 
-    async def openai_embeddings(
-        self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
-    ) -> OpenAIEmbeddingsResponse:
-        raise NotImplementedError()
-
     async def chat_completion(
         self,
         model_id: str,
@@ -283,7 +246,7 @@ class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
         )
 
         try:
-            response = await self._get_client(provider_model_id).chat.completions.create(**request)
+            response = await self.client.chat.completions.create(**request)
         except APIConnectionError as e:
             raise ConnectionError(f"Failed to connect to NVIDIA NIM at {self._config.url}: {e}") from e
 
@@ -292,153 +255,3 @@ 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: str | list[str] | list[int] | list[list[int]],
-        best_of: int | None = None,
-        echo: bool | None = None,
-        frequency_penalty: float | None = None,
-        logit_bias: dict[str, float] | None = None,
-        logprobs: bool | None = None,
-        max_tokens: int | None = None,
-        n: int | None = None,
-        presence_penalty: float | None = None,
-        seed: int | None = None,
-        stop: str | list[str] | None = None,
-        stream: bool | None = None,
-        stream_options: dict[str, Any] | None = None,
-        temperature: float | None = None,
-        top_p: float | None = None,
-        user: str | None = None,
-        guided_choice: list[str] | None = None,
-        prompt_logprobs: int | None = None,
-        suffix: str | None = None,
-    ) -> OpenAICompletion:
-        provider_model_id = await 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: float | None = None,
-        function_call: str | dict[str, Any] | None = None,
-        functions: list[dict[str, Any]] | None = None,
-        logit_bias: dict[str, float] | None = None,
-        logprobs: bool | None = None,
-        max_completion_tokens: int | None = None,
-        max_tokens: int | None = None,
-        n: int | None = None,
-        parallel_tool_calls: bool | None = None,
-        presence_penalty: float | None = None,
-        response_format: OpenAIResponseFormatParam | None = None,
-        seed: int | None = None,
-        stop: str | list[str] | None = None,
-        stream: bool | None = None,
-        stream_options: dict[str, Any] | None = None,
-        temperature: float | None = None,
-        tool_choice: str | dict[str, Any] | None = None,
-        tools: list[dict[str, Any]] | None = None,
-        top_logprobs: int | None = None,
-        top_p: float | None = None,
-        user: str | None = None,
-    ) -> OpenAIChatCompletion | AsyncIterator[OpenAIChatCompletionChunk]:
-        provider_model_id = await 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
-
-    async def register_model(self, model: Model) -> Model:
-        """
-        Allow non-llama model registration.
-
-        Non-llama model registration: API Catalogue models, post-training models, etc.
-            client = LlamaStackAsLibraryClient("nvidia")
-            client.models.register(
-                    model_id="mistralai/mixtral-8x7b-instruct-v0.1",
-                    model_type=ModelType.llm,
-                    provider_id="nvidia",
-                    provider_model_id="mistralai/mixtral-8x7b-instruct-v0.1"
-            )
-
-            NOTE: Only supports models endpoints compatible with AsyncOpenAI base_url format.
-        """
-        if model.model_type == ModelType.embedding:
-            # embedding models are always registered by their provider model id and does not need to be mapped to a llama model
-            provider_resource_id = model.provider_resource_id
-        else:
-            provider_resource_id = self.get_provider_model_id(model.provider_resource_id)
-
-        if provider_resource_id:
-            model.provider_resource_id = provider_resource_id
-        else:
-            llama_model = model.metadata.get("llama_model")
-            existing_llama_model = self.get_llama_model(model.provider_resource_id)
-            if existing_llama_model:
-                if existing_llama_model != llama_model:
-                    raise ValueError(
-                        f"Provider model id '{model.provider_resource_id}' is already registered to a different llama model: '{existing_llama_model}'"
-                    )
-            else:
-                # not llama model
-                if llama_model in ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR:
-                    self.provider_id_to_llama_model_map[model.provider_resource_id] = (
-                        ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR[llama_model]
-                    )
-                else:
-                    self.alias_to_provider_id_map[model.provider_model_id] = model.provider_model_id
-        return model

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

@@ -6,13 +6,17 @@
 
 from typing import Any
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 DEFAULT_OLLAMA_URL = "http://localhost:11434"
 
 
 class OllamaImplConfig(BaseModel):
     url: str = DEFAULT_OLLAMA_URL
+    refresh_models: bool = Field(
+        default=False,
+        description="Whether to refresh models periodically",
+    )
 
     @classmethod
     def sample_run_config(cls, url: str = "${env.OLLAMA_URL:=http://localhost:11434}", **kwargs) -> dict[str, Any]:

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

@@ -12,6 +12,19 @@ from llama_stack.providers.utils.inference.model_registry import (
     build_model_entry,
 )
 
+SAFETY_MODELS_ENTRIES = [
+    # The Llama Guard models don't have their full fp16 versions
+    # so we are going to alias their default version to the canonical SKU
+    build_hf_repo_model_entry(
+        "llama-guard3:8b",
+        CoreModelId.llama_guard_3_8b.value,
+    ),
+    build_hf_repo_model_entry(
+        "llama-guard3:1b",
+        CoreModelId.llama_guard_3_1b.value,
+    ),
+]
+
 MODEL_ENTRIES = [
     build_hf_repo_model_entry(
         "llama3.1:8b-instruct-fp16",
@@ -73,16 +86,6 @@ MODEL_ENTRIES = [
         "llama3.3:70b",
         CoreModelId.llama3_3_70b_instruct.value,
     ),
-    # The Llama Guard models don't have their full fp16 versions
-    # so we are going to alias their default version to the canonical SKU
-    build_hf_repo_model_entry(
-        "llama-guard3:8b",
-        CoreModelId.llama_guard_3_8b.value,
-    ),
-    build_hf_repo_model_entry(
-        "llama-guard3:1b",
-        CoreModelId.llama_guard_3_1b.value,
-    ),
     ProviderModelEntry(
         provider_model_id="all-minilm:l6-v2",
         aliases=["all-minilm"],
@@ -100,4 +103,4 @@ MODEL_ENTRIES = [
             "context_length": 8192,
         },
     ),
-]
+] + SAFETY_MODELS_ENTRIES

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

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 
+import asyncio
 import base64
 import uuid
 from collections.abc import AsyncGenerator, AsyncIterator
@@ -91,23 +92,93 @@ class OllamaInferenceAdapter(
     InferenceProvider,
     ModelsProtocolPrivate,
 ):
+    # automatically set by the resolver when instantiating the provider
+    __provider_id__: str
+
     def __init__(self, config: OllamaImplConfig) -> None:
         self.register_helper = ModelRegistryHelper(MODEL_ENTRIES)
-        self.url = config.url
+        self.config = config
+        self._clients: dict[asyncio.AbstractEventLoop, AsyncClient] = {}
+        self._openai_client = None
 
     @property
     def client(self) -> AsyncClient:
-        return AsyncClient(host=self.url)
+        # ollama client attaches itself to the current event loop (sadly?)
+        loop = asyncio.get_running_loop()
+        if loop not in self._clients:
+            self._clients[loop] = AsyncClient(host=self.config.url)
+        return self._clients[loop]
 
     @property
     def openai_client(self) -> AsyncOpenAI:
-        return AsyncOpenAI(base_url=f"{self.url}/v1", api_key="ollama")
+        if self._openai_client is None:
+            url = self.config.url.rstrip("/")
+            self._openai_client = AsyncOpenAI(base_url=f"{url}/v1", api_key="ollama")
+        return self._openai_client
 
     async def initialize(self) -> None:
-        logger.debug(f"checking connectivity to Ollama at `{self.url}`...")
+        logger.info(f"checking connectivity to Ollama at `{self.config.url}`...")
         health_response = await self.health()
         if health_response["status"] == HealthStatus.ERROR:
-            raise RuntimeError("Ollama Server is not running, start it using `ollama serve` in a separate terminal")
+            logger.warning(
+                "Ollama Server is not running, make sure to start it using `ollama serve` in a separate terminal"
+            )
+
+    async def should_refresh_models(self) -> bool:
+        return self.config.refresh_models
+
+    async def list_models(self) -> list[Model] | None:
+        provider_id = self.__provider_id__
+        response = await self.client.list()
+
+        # always add the two embedding models which can be pulled on demand
+        models = [
+            Model(
+                identifier="all-minilm:l6-v2",
+                provider_resource_id="all-minilm:l6-v2",
+                provider_id=provider_id,
+                metadata={
+                    "embedding_dimension": 384,
+                    "context_length": 512,
+                },
+                model_type=ModelType.embedding,
+            ),
+            # add all-minilm alias
+            Model(
+                identifier="all-minilm",
+                provider_resource_id="all-minilm:l6-v2",
+                provider_id=provider_id,
+                metadata={
+                    "embedding_dimension": 384,
+                    "context_length": 512,
+                },
+                model_type=ModelType.embedding,
+            ),
+            Model(
+                identifier="nomic-embed-text",
+                provider_resource_id="nomic-embed-text",
+                provider_id=provider_id,
+                metadata={
+                    "embedding_dimension": 768,
+                    "context_length": 8192,
+                },
+                model_type=ModelType.embedding,
+            ),
+        ]
+        for m in response.models:
+            # kill embedding models since we don't know dimensions for them
+            if "bert" in m.details.family:
+                continue
+            models.append(
+                Model(
+                    identifier=m.model,
+                    provider_resource_id=m.model,
+                    provider_id=provider_id,
+                    metadata={},
+                    model_type=ModelType.llm,
+                )
+            )
+        return models
 
     async def health(self) -> HealthResponse:
         """
@@ -124,7 +195,7 @@ class OllamaInferenceAdapter(
             return HealthResponse(status=HealthStatus.ERROR, message=f"Health check failed: {str(e)}")
 
     async def shutdown(self) -> None:
-        pass
+        self._clients.clear()
 
     async def unregister_model(self, model_id: str) -> None:
         pass
@@ -350,12 +421,7 @@ class OllamaInferenceAdapter(
         except ValueError:
             pass  # Ignore statically unknown model, will check live listing
 
-        if model.provider_resource_id is None:
-            raise ValueError("Model provider_resource_id cannot be None")
-
         if model.model_type == ModelType.embedding:
-            logger.info(f"Pulling embedding model `{model.provider_resource_id}` if necessary...")
-            # TODO: you should pull here only if the model is not found in a list
             response = await self.client.list()
             if model.provider_resource_id not in [m.model for m in response.models]:
                 await self.client.pull(model.provider_resource_id)
@@ -365,9 +431,9 @@ class OllamaInferenceAdapter(
         #  - models not currently running are run by the ollama server as needed
         response = await self.client.list()
         available_models = [m.model for m in response.models]
-        provider_resource_id = self.register_helper.get_provider_model_id(model.provider_resource_id)
-        if provider_resource_id is None:
-            provider_resource_id = model.provider_resource_id
+
+        provider_resource_id = model.provider_resource_id
+        assert provider_resource_id is not None  # mypy
         if provider_resource_id not in available_models:
             available_models_latest = [m.model.split(":latest")[0] for m in response.models]
             if provider_resource_id in available_models_latest:
@@ -375,7 +441,9 @@ class OllamaInferenceAdapter(
                     f"Imprecise provider resource id was used but 'latest' is available in Ollama - using '{model.provider_resource_id}:latest'"
                 )
                 return model
-            raise UnsupportedModelError(model.provider_resource_id, available_models)
+            raise UnsupportedModelError(provider_resource_id, available_models)
+
+        # mutating this should be considered an anti-pattern
         model.provider_resource_id = provider_resource_id
 
         return model

llama_stack/providers/remote/inference/openai/config.py

@@ -24,9 +24,19 @@ class OpenAIConfig(BaseModel):
         default=None,
         description="API key for OpenAI models",
     )
+    base_url: str = Field(
+        default="https://api.openai.com/v1",
+        description="Base URL for OpenAI API",
+    )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.OPENAI_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(
+        cls,
+        api_key: str = "${env.OPENAI_API_KEY:=}",
+        base_url: str = "${env.OPENAI_BASE_URL:=https://api.openai.com/v1}",
+        **kwargs,
+    ) -> dict[str, Any]:
         return {
             "api_key": api_key,
+            "base_url": base_url,
         }

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

@@ -12,11 +12,6 @@ from llama_stack.providers.utils.inference.model_registry import (
 )
 
 LLM_MODEL_IDS = [
-    # the models w/ "openai/" prefix are the litellm specific model names.
-    # they should be deprecated in favor of the canonical openai model names.
-    "openai/gpt-4o",
-    "openai/gpt-4o-mini",
-    "openai/chatgpt-4o-latest",
     "gpt-3.5-turbo-0125",
     "gpt-3.5-turbo",
     "gpt-3.5-turbo-instruct",
@@ -43,8 +38,6 @@ class EmbeddingModelInfo:
 
 
 EMBEDDING_MODEL_IDS: dict[str, EmbeddingModelInfo] = {
-    "openai/text-embedding-3-small": EmbeddingModelInfo(1536, 8192),
-    "openai/text-embedding-3-large": EmbeddingModelInfo(3072, 8192),
     "text-embedding-3-small": EmbeddingModelInfo(1536, 8192),
     "text-embedding-3-large": EmbeddingModelInfo(3072, 8192),
 }

llama_stack/providers/remote/inference/openai/openai.py

@@ -5,23 +5,9 @@
 # the root directory of this source tree.
 
 import logging
-from collections.abc import AsyncIterator
-from typing import Any
 
-from openai import AsyncOpenAI
-
-from llama_stack.apis.inference import (
-    OpenAIChatCompletion,
-    OpenAIChatCompletionChunk,
-    OpenAICompletion,
-    OpenAIEmbeddingData,
-    OpenAIEmbeddingsResponse,
-    OpenAIEmbeddingUsage,
-    OpenAIMessageParam,
-    OpenAIResponseFormatParam,
-)
 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 llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 from .config import OpenAIConfig
 from .models import MODEL_ENTRIES
@@ -30,7 +16,7 @@ logger = logging.getLogger(__name__)
 
 
 #
-# This OpenAI adapter implements Inference methods using two clients -
+# This OpenAI adapter implements Inference methods using two mixins -
 #
 # | Inference Method           | Implementation Source    |
 # |----------------------------|--------------------------|
@@ -39,15 +25,27 @@ logger = logging.getLogger(__name__)
 # | embedding                  | LiteLLMOpenAIMixin       |
 # | batch_completion           | LiteLLMOpenAIMixin       |
 # | batch_chat_completion      | LiteLLMOpenAIMixin       |
-# | openai_completion          | AsyncOpenAI              |
-# | openai_chat_completion     | AsyncOpenAI              |
-# | openai_embeddings          | AsyncOpenAI              |
+# | openai_completion          | OpenAIMixin              |
+# | openai_chat_completion     | OpenAIMixin              |
+# | openai_embeddings          | OpenAIMixin              |
 #
-class OpenAIInferenceAdapter(LiteLLMOpenAIMixin):
+class OpenAIInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin):
+    """
+    OpenAI Inference Adapter for Llama Stack.
+
+    Note: The inheritance order is important here. OpenAIMixin must come before
+    LiteLLMOpenAIMixin to ensure that OpenAIMixin.check_model_availability()
+    is used instead of ModelRegistryHelper.check_model_availability().
+
+    - OpenAIMixin.check_model_availability() queries the OpenAI API to check if a model exists
+    - ModelRegistryHelper.check_model_availability() (inherited by LiteLLMOpenAIMixin) just returns False and shows a warning
+    """
+
     def __init__(self, config: OpenAIConfig) -> None:
         LiteLLMOpenAIMixin.__init__(
             self,
             MODEL_ENTRIES,
+            litellm_provider_name="openai",
             api_key_from_config=config.api_key,
             provider_data_api_key_field="openai_api_key",
         )
@@ -59,169 +57,20 @@ class OpenAIInferenceAdapter(LiteLLMOpenAIMixin):
         # if we do not set this, users will be exposed to the
         # litellm specific model names, an abstraction leak.
         self.is_openai_compat = True
-        self._openai_client = AsyncOpenAI(
-            api_key=self.config.api_key,
-        )
+
+    # Delegate the client data handling get_api_key method to LiteLLMOpenAIMixin
+    get_api_key = LiteLLMOpenAIMixin.get_api_key
+
+    def get_base_url(self) -> str:
+        """
+        Get the OpenAI API base URL.
+
+        Returns the OpenAI API base URL from the configuration.
+        """
+        return self.config.base_url
 
     async def initialize(self) -> None:
         await super().initialize()
 
     async def shutdown(self) -> None:
         await super().shutdown()
-
-    async def openai_completion(
-        self,
-        model: str,
-        prompt: str | list[str] | list[int] | list[list[int]],
-        best_of: int | None = None,
-        echo: bool | None = None,
-        frequency_penalty: float | None = None,
-        logit_bias: dict[str, float] | None = None,
-        logprobs: bool | None = None,
-        max_tokens: int | None = None,
-        n: int | None = None,
-        presence_penalty: float | None = None,
-        seed: int | None = None,
-        stop: str | list[str] | None = None,
-        stream: bool | None = None,
-        stream_options: dict[str, Any] | None = None,
-        temperature: float | None = None,
-        top_p: float | None = None,
-        user: str | None = None,
-        guided_choice: list[str] | None = None,
-        prompt_logprobs: int | None = None,
-        suffix: str | None = None,
-    ) -> OpenAICompletion:
-        if guided_choice is not None:
-            logging.warning("guided_choice is not supported by the OpenAI API. Ignoring.")
-        if prompt_logprobs is not None:
-            logging.warning("prompt_logprobs is not supported by the OpenAI API. Ignoring.")
-
-        model_id = (await self.model_store.get_model(model)).provider_resource_id
-        if model_id.startswith("openai/"):
-            model_id = model_id[len("openai/") :]
-        params = await prepare_openai_completion_params(
-            model=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,
-            suffix=suffix,
-        )
-        return await self._openai_client.completions.create(**params)
-
-    async def openai_chat_completion(
-        self,
-        model: str,
-        messages: list[OpenAIMessageParam],
-        frequency_penalty: float | None = None,
-        function_call: str | dict[str, Any] | None = None,
-        functions: list[dict[str, Any]] | None = None,
-        logit_bias: dict[str, float] | None = None,
-        logprobs: bool | None = None,
-        max_completion_tokens: int | None = None,
-        max_tokens: int | None = None,
-        n: int | None = None,
-        parallel_tool_calls: bool | None = None,
-        presence_penalty: float | None = None,
-        response_format: OpenAIResponseFormatParam | None = None,
-        seed: int | None = None,
-        stop: str | list[str] | None = None,
-        stream: bool | None = None,
-        stream_options: dict[str, Any] | None = None,
-        temperature: float | None = None,
-        tool_choice: str | dict[str, Any] | None = None,
-        tools: list[dict[str, Any]] | None = None,
-        top_logprobs: int | None = None,
-        top_p: float | None = None,
-        user: str | None = None,
-    ) -> OpenAIChatCompletion | AsyncIterator[OpenAIChatCompletionChunk]:
-        model_id = (await self.model_store.get_model(model)).provider_resource_id
-        if model_id.startswith("openai/"):
-            model_id = model_id[len("openai/") :]
-        params = await prepare_openai_completion_params(
-            model=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,
-        )
-        return await self._openai_client.chat.completions.create(**params)
-
-    async def openai_embeddings(
-        self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
-    ) -> OpenAIEmbeddingsResponse:
-        model_id = (await self.model_store.get_model(model)).provider_resource_id
-        if model_id.startswith("openai/"):
-            model_id = model_id[len("openai/") :]
-
-        # Prepare parameters for OpenAI embeddings API
-        params = {
-            "model": model_id,
-            "input": input,
-        }
-
-        if encoding_format is not None:
-            params["encoding_format"] = encoding_format
-        if dimensions is not None:
-            params["dimensions"] = dimensions
-        if user is not None:
-            params["user"] = user
-
-        # Call OpenAI embeddings API
-        response = await self._openai_client.embeddings.create(**params)
-
-        data = []
-        for i, embedding_data in enumerate(response.data):
-            data.append(
-                OpenAIEmbeddingData(
-                    embedding=embedding_data.embedding,
-                    index=i,
-                )
-            )
-
-        usage = OpenAIEmbeddingUsage(
-            prompt_tokens=response.usage.prompt_tokens,
-            total_tokens=response.usage.total_tokens,
-        )
-
-        return OpenAIEmbeddingsResponse(
-            data=data,
-            model=response.model,
-            usage=usage,
-        )

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

@@ -34,7 +34,7 @@ from llama_stack.apis.inference import (
     ToolPromptFormat,
 )
 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.core.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
 

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

@@ -30,7 +30,7 @@ class SambaNovaImplConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.SAMBANOVA_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.SAMBANOVA_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "url": "https://api.sambanova.ai/v1",
             "api_key": api_key,

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

@@ -9,49 +9,20 @@ from llama_stack.providers.utils.inference.model_registry import (
     build_hf_repo_model_entry,
 )
 
-SAFETY_MODELS_ENTRIES = [
-    build_hf_repo_model_entry(
-        "sambanova/Meta-Llama-Guard-3-8B",
-        CoreModelId.llama_guard_3_8b.value,
-    ),
-]
+SAFETY_MODELS_ENTRIES = []
 
 
 MODEL_ENTRIES = [
     build_hf_repo_model_entry(
-        "sambanova/Meta-Llama-3.1-8B-Instruct",
+        "Meta-Llama-3.1-8B-Instruct",
         CoreModelId.llama3_1_8b_instruct.value,
     ),
     build_hf_repo_model_entry(
-        "sambanova/Meta-Llama-3.1-405B-Instruct",
-        CoreModelId.llama3_1_405b_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "sambanova/Meta-Llama-3.2-1B-Instruct",
-        CoreModelId.llama3_2_1b_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "sambanova/Meta-Llama-3.2-3B-Instruct",
-        CoreModelId.llama3_2_3b_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "sambanova/Meta-Llama-3.3-70B-Instruct",
+        "Meta-Llama-3.3-70B-Instruct",
         CoreModelId.llama3_3_70b_instruct.value,
     ),
     build_hf_repo_model_entry(
-        "sambanova/Llama-3.2-11B-Vision-Instruct",
-        CoreModelId.llama3_2_11b_vision_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "sambanova/Llama-3.2-90B-Vision-Instruct",
-        CoreModelId.llama3_2_90b_vision_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "sambanova/Llama-4-Scout-17B-16E-Instruct",
-        CoreModelId.llama4_scout_17b_16e_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "sambanova/Llama-4-Maverick-17B-128E-Instruct",
+        "Llama-4-Maverick-17B-128E-Instruct",
         CoreModelId.llama4_maverick_17b_128e_instruct.value,
     ),
 ] + SAFETY_MODELS_ENTRIES

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

@@ -4,250 +4,23 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import json
-from collections.abc import Iterable
-
-from openai.types.chat import (
-    ChatCompletionAssistantMessageParam as OpenAIChatCompletionAssistantMessage,
-)
-from openai.types.chat import (
-    ChatCompletionContentPartImageParam as OpenAIChatCompletionContentPartImageParam,
-)
-from openai.types.chat import (
-    ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam,
-)
-from openai.types.chat import (
-    ChatCompletionContentPartTextParam as OpenAIChatCompletionContentPartTextParam,
-)
-from openai.types.chat import (
-    ChatCompletionMessageParam as OpenAIChatCompletionMessage,
-)
-from openai.types.chat import (
-    ChatCompletionMessageToolCallParam as OpenAIChatCompletionMessageToolCall,
-)
-from openai.types.chat import (
-    ChatCompletionSystemMessageParam as OpenAIChatCompletionSystemMessage,
-)
-from openai.types.chat import (
-    ChatCompletionToolMessageParam as OpenAIChatCompletionToolMessage,
-)
-from openai.types.chat import (
-    ChatCompletionUserMessageParam as OpenAIChatCompletionUserMessage,
-)
-from openai.types.chat.chat_completion_content_part_image_param import (
-    ImageURL as OpenAIImageURL,
-)
-from openai.types.chat.chat_completion_message_tool_call_param import (
-    Function as OpenAIFunction,
-)
-
-from llama_stack.apis.common.content_types import (
-    ImageContentItem,
-    InterleavedContent,
-    TextContentItem,
-)
-from llama_stack.apis.inference import (
-    ChatCompletionRequest,
-    CompletionMessage,
-    JsonSchemaResponseFormat,
-    Message,
-    SystemMessage,
-    ToolChoice,
-    ToolResponseMessage,
-    UserMessage,
-)
-from llama_stack.log import get_logger
-from llama_stack.models.llama.datatypes import BuiltinTool
 from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-from llama_stack.providers.utils.inference.openai_compat import (
-    convert_tooldef_to_openai_tool,
-    get_sampling_options,
-)
-from llama_stack.providers.utils.inference.prompt_adapter import convert_image_content_to_url
 
 from .config import SambaNovaImplConfig
 from .models import MODEL_ENTRIES
 
-logger = get_logger(name=__name__, category="inference")
-
-
-async def convert_message_to_openai_dict_with_b64_images(
-    message: Message | dict,
-) -> OpenAIChatCompletionMessage:
-    """
-    Convert a Message to an OpenAI API-compatible dictionary.
-    """
-    # users can supply a dict instead of a Message object, we'll
-    # convert it to a Message object and proceed with some type safety.
-    if isinstance(message, dict):
-        if "role" not in message:
-            raise ValueError("role is required in message")
-        if message["role"] == "user":
-            message = UserMessage(**message)
-        elif message["role"] == "assistant":
-            message = CompletionMessage(**message)
-        elif message["role"] == "tool":
-            message = ToolResponseMessage(**message)
-        elif message["role"] == "system":
-            message = SystemMessage(**message)
-        else:
-            raise ValueError(f"Unsupported message role: {message['role']}")
-
-    # Map Llama Stack spec to OpenAI spec -
-    #  str -> str
-    #  {"type": "text", "text": ...} -> {"type": "text", "text": ...}
-    #  {"type": "image", "image": {"url": {"uri": ...}}} -> {"type": "image_url", "image_url": {"url": ...}}
-    #  {"type": "image", "image": {"data": ...}} -> {"type": "image_url", "image_url": {"url": "data:image/?;base64,..."}}
-    #  List[...] -> List[...]
-    async def _convert_message_content(
-        content: InterleavedContent,
-    ) -> str | Iterable[OpenAIChatCompletionContentPartParam]:
-        async def impl(
-            content_: InterleavedContent,
-        ) -> str | OpenAIChatCompletionContentPartParam | list[OpenAIChatCompletionContentPartParam]:
-            # Llama Stack and OpenAI spec match for str and text input
-            if isinstance(content_, str):
-                return content_
-            elif isinstance(content_, TextContentItem):
-                return OpenAIChatCompletionContentPartTextParam(
-                    type="text",
-                    text=content_.text,
-                )
-            elif isinstance(content_, ImageContentItem):
-                return OpenAIChatCompletionContentPartImageParam(
-                    type="image_url",
-                    image_url=OpenAIImageURL(url=await convert_image_content_to_url(content_, download=True)),
-                )
-            elif isinstance(content_, list):
-                return [await impl(item) for item in content_]
-            else:
-                raise ValueError(f"Unsupported content type: {type(content_)}")
-
-        ret = await impl(content)
-
-        # OpenAI*Message expects a str or list
-        if isinstance(ret, str) or isinstance(ret, list):
-            return ret
-        else:
-            return [ret]
-
-    out: OpenAIChatCompletionMessage = None
-    if isinstance(message, UserMessage):
-        out = OpenAIChatCompletionUserMessage(
-            role="user",
-            content=await _convert_message_content(message.content),
-        )
-    elif isinstance(message, CompletionMessage):
-        out = OpenAIChatCompletionAssistantMessage(
-            role="assistant",
-            content=await _convert_message_content(message.content),
-            tool_calls=[
-                OpenAIChatCompletionMessageToolCall(
-                    id=tool.call_id,
-                    function=OpenAIFunction(
-                        name=tool.tool_name if not isinstance(tool.tool_name, BuiltinTool) else tool.tool_name.value,
-                        arguments=json.dumps(tool.arguments),
-                    ),
-                    type="function",
-                )
-                for tool in message.tool_calls
-            ]
-            or None,
-        )
-    elif isinstance(message, ToolResponseMessage):
-        out = OpenAIChatCompletionToolMessage(
-            role="tool",
-            tool_call_id=message.call_id,
-            content=await _convert_message_content(message.content),
-        )
-    elif isinstance(message, SystemMessage):
-        out = OpenAIChatCompletionSystemMessage(
-            role="system",
-            content=await _convert_message_content(message.content),
-        )
-    else:
-        raise ValueError(f"Unsupported message type: {type(message)}")
-
-    return out
-
 
 class SambaNovaInferenceAdapter(LiteLLMOpenAIMixin):
-    _config: SambaNovaImplConfig
-
     def __init__(self, config: SambaNovaImplConfig):
         self.config = config
+        self.environment_available_models = []
         LiteLLMOpenAIMixin.__init__(
             self,
             model_entries=MODEL_ENTRIES,
-            api_key_from_config=self.config.api_key,
+            litellm_provider_name="sambanova",
+            api_key_from_config=self.config.api_key.get_secret_value() if self.config.api_key else None,
             provider_data_api_key_field="sambanova_api_key",
+            openai_compat_api_base=self.config.url,
+            download_images=True,  # SambaNova requires base64 image encoding
+            json_schema_strict=False,  # SambaNova doesn't support strict=True yet
         )
-
-    def _get_api_key(self) -> str:
-        config_api_key = self.config.api_key if self.config.api_key else None
-        if config_api_key:
-            return config_api_key.get_secret_value()
-        else:
-            provider_data = self.get_request_provider_data()
-            if provider_data is None or not provider_data.sambanova_api_key:
-                raise ValueError(
-                    'Pass Sambanova API Key in the header X-LlamaStack-Provider-Data as { "sambanova_api_key": <your api key> }'
-                )
-            return provider_data.sambanova_api_key
-
-    async def _get_params(self, request: ChatCompletionRequest) -> dict:
-        input_dict = {}
-
-        input_dict["messages"] = [await convert_message_to_openai_dict_with_b64_images(m) for m in request.messages]
-        if fmt := request.response_format:
-            if not isinstance(fmt, JsonSchemaResponseFormat):
-                raise ValueError(
-                    f"Unsupported response format: {type(fmt)}. Only JsonSchemaResponseFormat is supported."
-                )
-
-            fmt = fmt.json_schema
-            name = fmt["title"]
-            del fmt["title"]
-            fmt["additionalProperties"] = False
-
-            # Apply additionalProperties: False recursively to all objects
-            fmt = self._add_additional_properties_recursive(fmt)
-
-            input_dict["response_format"] = {
-                "type": "json_schema",
-                "json_schema": {
-                    "name": name,
-                    "schema": fmt,
-                    "strict": False,
-                },
-            }
-        if request.tools:
-            input_dict["tools"] = [convert_tooldef_to_openai_tool(tool) for tool in request.tools]
-            if request.tool_config.tool_choice:
-                input_dict["tool_choice"] = (
-                    request.tool_config.tool_choice.value
-                    if isinstance(request.tool_config.tool_choice, ToolChoice)
-                    else request.tool_config.tool_choice
-                )
-
-        provider_data = self.get_request_provider_data()
-        key_field = self.provider_data_api_key_field
-        if provider_data and getattr(provider_data, key_field, None):
-            api_key = getattr(provider_data, key_field)
-        else:
-            api_key = self._get_api_key()
-
-        return {
-            "model": request.model,
-            "api_key": api_key,
-            "api_base": self.config.url,
-            **input_dict,
-            "stream": request.stream,
-            **get_sampling_options(request.sampling_params),
-        }
-
-    async def initialize(self):
-        await super().initialize()
-
-    async def shutdown(self):
-        await super().shutdown()

llama_stack/providers/remote/inference/sambanova_openai_compat/__init__.py

@@ -1,17 +0,0 @@
-# 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 llama_stack.apis.inference import InferenceProvider
-
-from .config import SambaNovaCompatConfig
-
-
-async def get_adapter_impl(config: SambaNovaCompatConfig, _deps) -> InferenceProvider:
-    # import dynamically so the import is used only when it is needed
-    from .sambanova import SambaNovaCompatInferenceAdapter
-
-    adapter = SambaNovaCompatInferenceAdapter(config)
-    return adapter

llama_stack/providers/remote/inference/sambanova_openai_compat/config.py

@@ -1,38 +0,0 @@
-# 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
-
-from pydantic import BaseModel, Field
-
-from llama_stack.schema_utils import json_schema_type
-
-
-class SambaNovaProviderDataValidator(BaseModel):
-    sambanova_api_key: str | None = Field(
-        default=None,
-        description="API key for SambaNova models",
-    )
-
-
-@json_schema_type
-class SambaNovaCompatConfig(BaseModel):
-    api_key: str | None = Field(
-        default=None,
-        description="The SambaNova API key",
-    )
-
-    openai_compat_api_base: str = Field(
-        default="https://api.sambanova.ai/v1",
-        description="The URL for the SambaNova API server",
-    )
-
-    @classmethod
-    def sample_run_config(cls, api_key: str = "${env.SAMBANOVA_API_KEY}", **kwargs) -> dict[str, Any]:
-        return {
-            "openai_compat_api_base": "https://api.sambanova.ai/v1",
-            "api_key": api_key,
-        }

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

@@ -1,30 +0,0 @@
-# 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 llama_stack.providers.remote.inference.sambanova_openai_compat.config import SambaNovaCompatConfig
-from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-
-from ..sambanova.models import MODEL_ENTRIES
-
-
-class SambaNovaCompatInferenceAdapter(LiteLLMOpenAIMixin):
-    _config: SambaNovaCompatConfig
-
-    def __init__(self, config: SambaNovaCompatConfig):
-        LiteLLMOpenAIMixin.__init__(
-            self,
-            model_entries=MODEL_ENTRIES,
-            api_key_from_config=config.api_key,
-            provider_data_api_key_field="sambanova_api_key",
-            openai_compat_api_base=config.openai_compat_api_base,
-        )
-        self.config = config
-
-    async def initialize(self):
-        await super().initialize()
-
-    async def shutdown(self):
-        await super().shutdown()

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

@@ -19,7 +19,7 @@ class TGIImplConfig(BaseModel):
     @classmethod
     def sample_run_config(
         cls,
-        url: str = "${env.TGI_URL}",
+        url: str = "${env.TGI_URL:=}",
         **kwargs,
     ):
         return {

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

@@ -305,6 +305,8 @@ class _HfAdapter(
 
 class TGIAdapter(_HfAdapter):
     async def initialize(self, config: TGIImplConfig) -> None:
+        if not config.url:
+            raise ValueError("You must provide a URL in run.yaml (or via the TGI_URL environment variable) to use TGI.")
         log.info(f"Initializing TGI client with url={config.url}")
         self.client = AsyncInferenceClient(
             model=config.url,

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

@@ -6,13 +6,14 @@
 
 from typing import Any
 
-from pydantic import BaseModel, Field, SecretStr
+from pydantic import Field, SecretStr
 
+from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
 from llama_stack.schema_utils import json_schema_type
 
 
 @json_schema_type
-class TogetherImplConfig(BaseModel):
+class TogetherImplConfig(RemoteInferenceProviderConfig):
     url: str = Field(
         default="https://api.together.xyz/v1",
         description="The URL for the Together AI server",
@@ -26,5 +27,5 @@ class TogetherImplConfig(BaseModel):
     def sample_run_config(cls, **kwargs) -> dict[str, Any]:
         return {
             "url": "https://api.together.xyz/v1",
-            "api_key": "${env.TOGETHER_API_KEY}",
+            "api_key": "${env.TOGETHER_API_KEY:=}",
         }

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

@@ -69,15 +69,9 @@ MODEL_ENTRIES = [
     build_hf_repo_model_entry(
         "meta-llama/Llama-4-Scout-17B-16E-Instruct",
         CoreModelId.llama4_scout_17b_16e_instruct.value,
-        additional_aliases=[
-            "together/meta-llama/Llama-4-Scout-17B-16E-Instruct",
-        ],
     ),
     build_hf_repo_model_entry(
         "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
         CoreModelId.llama4_maverick_17b_128e_instruct.value,
-        additional_aliases=[
-            "together/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
-        ],
     ),
 ] + SAFETY_MODELS_ENTRIES

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

@@ -38,7 +38,7 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
-from llama_stack.distribution.request_headers import NeedsRequestProviderData
+from llama_stack.core.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 (
@@ -66,21 +66,14 @@ logger = get_logger(name=__name__, category="inference")
 
 class TogetherInferenceAdapter(ModelRegistryHelper, Inference, NeedsRequestProviderData):
     def __init__(self, config: TogetherImplConfig) -> None:
-        ModelRegistryHelper.__init__(self, MODEL_ENTRIES)
+        ModelRegistryHelper.__init__(self, MODEL_ENTRIES, config.allowed_models)
         self.config = config
-        self._client = None
-        self._openai_client = None
 
     async def initialize(self) -> None:
         pass
 
     async def shutdown(self) -> None:
-        if self._client:
-            # Together client has no close method, so just set to None
-            self._client = None
-        if self._openai_client:
-            await self._openai_client.close()
-            self._openai_client = None
+        pass
 
     async def completion(
         self,
@@ -108,29 +101,25 @@ class TogetherInferenceAdapter(ModelRegistryHelper, Inference, NeedsRequestProvi
             return await self._nonstream_completion(request)
 
     def _get_client(self) -> AsyncTogether:
-        if not self._client:
-            together_api_key = None
-            config_api_key = self.config.api_key.get_secret_value() if self.config.api_key else None
-            if config_api_key:
-                together_api_key = config_api_key
-            else:
-                provider_data = self.get_request_provider_data()
-                if provider_data is None or not provider_data.together_api_key:
-                    raise ValueError(
-                        'Pass Together API Key in the header X-LlamaStack-Provider-Data as { "together_api_key": <your api key>}'
-                    )
-                together_api_key = provider_data.together_api_key
-            self._client = AsyncTogether(api_key=together_api_key)
-        return self._client
+        together_api_key = None
+        config_api_key = self.config.api_key.get_secret_value() if self.config.api_key else None
+        if config_api_key:
+            together_api_key = config_api_key
+        else:
+            provider_data = self.get_request_provider_data()
+            if provider_data is None or not provider_data.together_api_key:
+                raise ValueError(
+                    'Pass Together API Key in the header X-LlamaStack-Provider-Data as { "together_api_key": <your api key>}'
+                )
+            together_api_key = provider_data.together_api_key
+        return AsyncTogether(api_key=together_api_key)
 
     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
+        together_client = self._get_client().client
+        return AsyncOpenAI(
+            base_url=together_client.base_url,
+            api_key=together_client.api_key,
+        )
 
     async def _nonstream_completion(self, request: CompletionRequest) -> ChatCompletionResponse:
         params = await self._get_params(request)

llama_stack/providers/remote/inference/together_openai_compat/__init__.py

@@ -1,17 +0,0 @@
-# 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 llama_stack.apis.inference import InferenceProvider
-
-from .config import TogetherCompatConfig
-
-
-async def get_adapter_impl(config: TogetherCompatConfig, _deps) -> InferenceProvider:
-    # import dynamically so the import is used only when it is needed
-    from .together import TogetherCompatInferenceAdapter
-
-    adapter = TogetherCompatInferenceAdapter(config)
-    return adapter

llama_stack/providers/remote/inference/together_openai_compat/config.py

@@ -1,38 +0,0 @@
-# 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
-
-from pydantic import BaseModel, Field
-
-from llama_stack.schema_utils import json_schema_type
-
-
-class TogetherProviderDataValidator(BaseModel):
-    together_api_key: str | None = Field(
-        default=None,
-        description="API key for Together models",
-    )
-
-
-@json_schema_type
-class TogetherCompatConfig(BaseModel):
-    api_key: str | None = Field(
-        default=None,
-        description="The Together API key",
-    )
-
-    openai_compat_api_base: str = Field(
-        default="https://api.together.xyz/v1",
-        description="The URL for the Together API server",
-    )
-
-    @classmethod
-    def sample_run_config(cls, api_key: str = "${env.TOGETHER_API_KEY}", **kwargs) -> dict[str, Any]:
-        return {
-            "openai_compat_api_base": "https://api.together.xyz/v1",
-            "api_key": api_key,
-        }

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

@@ -1,30 +0,0 @@
-# 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 llama_stack.providers.remote.inference.together_openai_compat.config import TogetherCompatConfig
-from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-
-from ..together.models import MODEL_ENTRIES
-
-
-class TogetherCompatInferenceAdapter(LiteLLMOpenAIMixin):
-    _config: TogetherCompatConfig
-
-    def __init__(self, config: TogetherCompatConfig):
-        LiteLLMOpenAIMixin.__init__(
-            self,
-            model_entries=MODEL_ENTRIES,
-            api_key_from_config=config.api_key,
-            provider_data_api_key_field="together_api_key",
-            openai_compat_api_base=config.openai_compat_api_base,
-        )
-        self.config = config
-
-    async def initialize(self):
-        await super().initialize()
-
-    async def shutdown(self):
-        await super().shutdown()

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

@@ -29,6 +29,10 @@ class VLLMInferenceAdapterConfig(BaseModel):
         default=True,
         description="Whether to verify TLS certificates. Can be a boolean or a path to a CA certificate file.",
     )
+    refresh_models: bool = Field(
+        default=False,
+        description="Whether to refresh models periodically",
+    )
 
     @field_validator("tls_verify")
     @classmethod
@@ -46,7 +50,7 @@ class VLLMInferenceAdapterConfig(BaseModel):
     @classmethod
     def sample_run_config(
         cls,
-        url: str = "${env.VLLM_URL}",
+        url: str = "${env.VLLM_URL:=}",
         **kwargs,
     ):
         return {

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

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 import json
-import logging
 from collections.abc import AsyncGenerator, AsyncIterator
 from typing import Any
 
@@ -38,6 +37,7 @@ from llama_stack.apis.inference import (
     JsonSchemaResponseFormat,
     LogProbConfig,
     Message,
+    ModelStore,
     OpenAIChatCompletion,
     OpenAICompletion,
     OpenAIEmbeddingData,
@@ -54,6 +54,7 @@ from llama_stack.apis.inference import (
     ToolPromptFormat,
 )
 from llama_stack.apis.models import Model, ModelType
+from llama_stack.log import get_logger
 from llama_stack.models.llama.datatypes import BuiltinTool, StopReason, ToolCall
 from llama_stack.models.llama.sku_list import all_registered_models
 from llama_stack.providers.datatypes import (
@@ -84,7 +85,7 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
 
 from .config import VLLMInferenceAdapterConfig
 
-log = logging.getLogger(__name__)
+log = get_logger(name=__name__, category="inference")
 
 
 def build_hf_repo_model_entries():
@@ -288,13 +289,40 @@ async def _process_vllm_chat_completion_stream_response(
 
 
 class VLLMInferenceAdapter(Inference, ModelsProtocolPrivate):
+    # automatically set by the resolver when instantiating the provider
+    __provider_id__: str
+    model_store: ModelStore | None = None
+
     def __init__(self, config: VLLMInferenceAdapterConfig) -> None:
         self.register_helper = ModelRegistryHelper(build_hf_repo_model_entries())
         self.config = config
         self.client = None
 
     async def initialize(self) -> None:
-        pass
+        if not self.config.url:
+            raise ValueError(
+                "You must provide a URL in run.yaml (or via the VLLM_URL environment variable) to use vLLM."
+            )
+
+    async def should_refresh_models(self) -> bool:
+        return self.config.refresh_models
+
+    async def list_models(self) -> list[Model] | None:
+        self._lazy_initialize_client()
+        assert self.client is not None  # mypy
+        models = []
+        async for m in self.client.models.list():
+            model_type = ModelType.llm  # unclear how to determine embedding vs. llm models
+            models.append(
+                Model(
+                    identifier=m.id,
+                    provider_resource_id=m.id,
+                    provider_id=self.__provider_id__,
+                    metadata={},
+                    model_type=model_type,
+                )
+            )
+        return models
 
     async def shutdown(self) -> None:
         pass

llama_stack/providers/remote/post_training/nvidia/README.md

@@ -22,7 +22,7 @@ This provider enables fine-tuning of LLMs using NVIDIA's NeMo Customizer service
 Build the NVIDIA environment:
 
 ```bash
-llama stack build --template nvidia --image-type conda
+llama stack build --distro nvidia --image-type venv
 ```
 
 ### Basic Usage using the LlamaStack Python Client
@@ -40,7 +40,7 @@ os.environ["NVIDIA_DATASET_NAMESPACE"] = "default"
 os.environ["NVIDIA_PROJECT_ID"] = "test-project"
 os.environ["NVIDIA_OUTPUT_MODEL_DIR"] = "test-example-model@v1"
 
-from llama_stack.distribution.library_client import LlamaStackAsLibraryClient
+from llama_stack.core.library_client import LlamaStackAsLibraryClient
 
 client = LlamaStackAsLibraryClient("nvidia")
 client.initialize()