From 064602bc9701d4aac6f42bf282d0dd9837f798c7 Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 30 Sep 2025 15:20:05 -0700 Subject: [PATCH 1/3] Update runpod.py --- .../remote/inference/runpod/runpod.py | 236 ++++++++++++++---- 1 file changed, 188 insertions(+), 48 deletions(-) diff --git a/llama_stack/providers/remote/inference/runpod/runpod.py b/llama_stack/providers/remote/inference/runpod/runpod.py index 82252b04da..3294d8eab0 100644 --- a/llama_stack/providers/remote/inference/runpod/runpod.py +++ b/llama_stack/providers/remote/inference/runpod/runpod.py @@ -3,52 +3,29 @@ # # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator from openai import OpenAI - -from llama_stack.apis.inference import * # noqa: F403 +from llama_stack.apis.inference import * from llama_stack.apis.inference import OpenAIEmbeddingsResponse - -# from llama_stack.providers.datatypes import ModelsProtocolPrivate -from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper, build_hf_repo_model_entry +from llama_stack.apis.models import Model, ModelType +from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper from llama_stack.providers.utils.inference.openai_compat import ( OpenAIChatCompletionToLlamaStackMixin, OpenAICompletionToLlamaStackMixin, get_sampling_options, process_chat_completion_response, process_chat_completion_stream_response, + process_completion_response, + process_completion_stream_response, ) from llama_stack.providers.utils.inference.prompt_adapter import ( - chat_completion_request_to_prompt, + completion_request_to_prompt, + interleaved_content_as_str, ) - from .config import RunpodImplConfig -# https://docs.runpod.io/serverless/vllm/overview#compatible-models -# https://github.com/runpod-workers/worker-vllm/blob/main/README.md#compatible-model-architectures -RUNPOD_SUPPORTED_MODELS = { - "Llama3.1-8B": "meta-llama/Llama-3.1-8B", - "Llama3.1-70B": "meta-llama/Llama-3.1-70B", - "Llama3.1-405B:bf16-mp8": "meta-llama/Llama-3.1-405B", - "Llama3.1-405B": "meta-llama/Llama-3.1-405B-FP8", - "Llama3.1-405B:bf16-mp16": "meta-llama/Llama-3.1-405B", - "Llama3.1-8B-Instruct": "meta-llama/Llama-3.1-8B-Instruct", - "Llama3.1-70B-Instruct": "meta-llama/Llama-3.1-70B-Instruct", - "Llama3.1-405B-Instruct:bf16-mp8": "meta-llama/Llama-3.1-405B-Instruct", - "Llama3.1-405B-Instruct": "meta-llama/Llama-3.1-405B-Instruct-FP8", - "Llama3.1-405B-Instruct:bf16-mp16": "meta-llama/Llama-3.1-405B-Instruct", - "Llama3.2-1B": "meta-llama/Llama-3.2-1B", - "Llama3.2-3B": "meta-llama/Llama-3.2-3B", -} - -SAFETY_MODELS_ENTRIES = [] - -# Create MODEL_ENTRIES from RUNPOD_SUPPORTED_MODELS for compatibility with starter template -MODEL_ENTRIES = [ - build_hf_repo_model_entry(provider_model_id, model_descriptor) - for provider_model_id, model_descriptor in RUNPOD_SUPPORTED_MODELS.items() -] + SAFETY_MODELS_ENTRIES +MODEL_ENTRIES = [] class RunpodInferenceAdapter( @@ -57,30 +34,83 @@ class RunpodInferenceAdapter( OpenAIChatCompletionToLlamaStackMixin, OpenAICompletionToLlamaStackMixin, ): + """ + Adapter for RunPod's OpenAI-compatible API endpoints. + Supports VLLM for serverless endpoint self-hosted or public endpoints. + Can work with any runpod endpoints that support OpenAI-compatible API + """ + def __init__(self, config: RunpodImplConfig) -> None: - ModelRegistryHelper.__init__(self, stack_to_provider_models_map=RUNPOD_SUPPORTED_MODELS) + ModelRegistryHelper.__init__(self, MODEL_ENTRIES) self.config = config async def initialize(self) -> None: - return + pass async def shutdown(self) -> None: pass + async def register_model(self, model: Model) -> Model: + """ + Register any model with the runpod provider_id. + + Pass-through registration - accepts any model string that the RunPod endpoint serves. + No static model validation since RunPod endpoints can serve arbitrary vLLM models. + + YAML Configuration Example: + models: + - metadata: {} + model_id: runpod/qwen/qwen3-8b + model_type: llm + provider_id: runpod + provider_model_id: qwen/qwen3-8b + - metadata: {} + model_id: runpod/deepcogito/cogito-v2-preview-llama-70B + model_type: llm + provider_id: runpod + provider_model_id: deepcogito/cogito-v2-preview-llama-70B + + The provider strips 'runpod/' prefix before API calls: + "runpod/qwen/qwen3-8b" -> "qwen/qwen3-8b" + """ + if model.provider_id == "runpod": + logger.info( + f"Registering model: {model.identifier} -> {model.provider_resource_id}" + ) + return model + return await super().register_model(model) + async def completion( self, - model: str, + model_id: str, content: InterleavedContent, sampling_params: SamplingParams | None = None, response_format: ResponseFormat | None = None, stream: bool | None = False, logprobs: LogProbConfig | None = None, ) -> AsyncGenerator: - raise NotImplementedError() + if sampling_params is None: + sampling_params = SamplingParams() + + request = CompletionRequest( + model=model_id, + content=content, + sampling_params=sampling_params, + response_format=response_format, + stream=stream, + logprobs=logprobs, + ) + + client = OpenAI(base_url=self.config.url, api_key=self.config.api_token) + + if stream: + return self._stream_completion(request, client) + else: + return await self._nonstream_completion(request, client) async def chat_completion( self, - model: str, + model_id: str, messages: list[Message], sampling_params: SamplingParams | None = None, response_format: ResponseFormat | None = None, @@ -91,10 +121,12 @@ async def chat_completion( logprobs: LogProbConfig | None = None, tool_config: ToolConfig | None = None, ) -> AsyncGenerator: + """Process chat completion requests using RunPod's OpenAI-compatible API.""" if sampling_params is None: sampling_params = SamplingParams() + request = ChatCompletionRequest( - model=model, + model=model_id, messages=messages, sampling_params=sampling_params, tools=tools or [], @@ -104,6 +136,7 @@ async def chat_completion( ) client = OpenAI(base_url=self.config.url, api_key=self.config.api_token) + if stream: return self._stream_chat_completion(request, client) else: @@ -112,15 +145,17 @@ async def chat_completion( async def _nonstream_chat_completion( self, request: ChatCompletionRequest, client: OpenAI ) -> ChatCompletionResponse: - params = self._get_params(request) - r = client.completions.create(**params) + params = await self._get_chat_params(request) + r = client.chat.completions.create(**params) return process_chat_completion_response(r, request) - async def _stream_chat_completion(self, request: ChatCompletionRequest, client: OpenAI) -> AsyncGenerator: - params = self._get_params(request) + async def _stream_chat_completion( + self, request: ChatCompletionRequest, client: OpenAI + ) -> AsyncGenerator: + params = await self._get_chat_params(request) async def _to_async_generator(): - s = client.completions.create(**params) + s = client.chat.completions.create(**params) for chunk in s: yield chunk @@ -128,14 +163,102 @@ async def _to_async_generator(): async for chunk in process_chat_completion_stream_response(stream, request): yield chunk - def _get_params(self, request: ChatCompletionRequest) -> dict: - return { - "model": self.map_to_provider_model(request.model), - "prompt": chat_completion_request_to_prompt(request), + async def _get_chat_params(self, request: ChatCompletionRequest) -> dict: + """Convert Llama Stack request to RunPod API parameters.""" + messages = [ + {"role": msg.role, "content": msg.content} for msg in request.messages + ] + + # Resolve model_id to provider_resource_id + model_obj = await self.model_store.get_model(request.model) + model = model_obj.provider_resource_id or request.model + + if model.startswith("runpod/"): + model = model.replace("runpod/", "", 1) + + params = { + "model": model, + "messages": messages, "stream": request.stream, **get_sampling_options(request.sampling_params), } + if request.stream: + params["stream_options"] = {"include_usage": True} + + return params + + async def _nonstream_completion( + self, request: CompletionRequest, client: OpenAI + ) -> CompletionResponse: + params = await self._get_completion_params(request) + r = client.completions.create(**params) + return process_completion_response(r) + + async def _stream_completion( + self, request: CompletionRequest, client: OpenAI + ) -> AsyncGenerator: + params = await self._get_completion_params(request) + + async def _to_async_generator(): + s = client.completions.create(**params) + for chunk in s: + yield chunk + + stream = _to_async_generator() + async for chunk in process_completion_stream_response(stream): + yield chunk + + async def _get_completion_params(self, request: CompletionRequest) -> dict: + # Resolve model_id to provider_resource_id + model_obj = await self.model_store.get_model(request.model) + model = model_obj.provider_resource_id or request.model + + if model.startswith("runpod/"): + model = model.replace("runpod/", "", 1) + + params = { + "model": model, + "prompt": completion_request_to_prompt(request), + "stream": request.stream, + **get_sampling_options(request.sampling_params), + } + + if request.stream: + params["stream_options"] = {"include_usage": True} + + return params + + 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: + # Resolve model_id to provider_resource_id + model_obj = await self.model_store.get_model(model_id) + model = model_obj.provider_resource_id or model_id + + if model.startswith("runpod/"): + model = model.replace("runpod/", "", 1) + + client = OpenAI(base_url=self.config.url, api_key=self.config.api_token) + + kwargs = {} + if output_dimension: + kwargs["dimensions"] = output_dimension + + response = client.embeddings.create( + model=model, + input=[interleaved_content_as_str(content) for content in contents], + **kwargs, + ) + + embeddings = [data.embedding for data in response.data] + return EmbeddingsResponse(embeddings=embeddings) + async def openai_embeddings( self, model: str, @@ -144,4 +267,21 @@ async def openai_embeddings( dimensions: int | None = None, user: str | None = None, ) -> OpenAIEmbeddingsResponse: - raise NotImplementedError() + # Resolve model_id to provider_resource_id + model_obj = await self.model_store.get_model(model) + model_stripped = model_obj.provider_resource_id or model + + if model_stripped.startswith("runpod/"): + model_stripped = model_stripped.replace("runpod/", "", 1) + + client = OpenAI(base_url=self.config.url, api_key=self.config.api_token) + + response = client.embeddings.create( + model=model_stripped, + input=input, + encoding_format=encoding_format, + dimensions=dimensions, + user=user, + ) + + return response \ No newline at end of file From 3236f8222397f9115b334c9b774b3e780dde3a48 Mon Sep 17 00:00:00 2001 From: Justin Date: Wed, 1 Oct 2025 14:33:15 -0700 Subject: [PATCH 2/3] Stable update --- .../remote/inference/runpod/runpod.py | 228 ++++++++++-------- 1 file changed, 129 insertions(+), 99 deletions(-) diff --git a/llama_stack/providers/remote/inference/runpod/runpod.py b/llama_stack/providers/remote/inference/runpod/runpod.py index 3294d8eab0..be7238628e 100644 --- a/llama_stack/providers/remote/inference/runpod/runpod.py +++ b/llama_stack/providers/remote/inference/runpod/runpod.py @@ -5,14 +5,21 @@ # the root directory of this source tree. from collections.abc import AsyncGenerator -from openai import OpenAI +import asyncio +from typing import Any + +from openai import AsyncOpenAI + from llama_stack.apis.inference import * -from llama_stack.apis.inference import OpenAIEmbeddingsResponse +from llama_stack.apis.inference import ( + OpenAIMessageParam, + OpenAIResponseFormatParam, +) +from llama_stack.apis.common.content_types import InterleavedContentItem from llama_stack.apis.models import Model, ModelType from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper from llama_stack.providers.utils.inference.openai_compat import ( - OpenAIChatCompletionToLlamaStackMixin, - OpenAICompletionToLlamaStackMixin, + convert_message_to_openai_dict, get_sampling_options, process_chat_completion_response, process_chat_completion_stream_response, @@ -23,16 +30,16 @@ completion_request_to_prompt, interleaved_content_as_str, ) +from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin from .config import RunpodImplConfig MODEL_ENTRIES = [] class RunpodInferenceAdapter( + OpenAIMixin, ModelRegistryHelper, Inference, - OpenAIChatCompletionToLlamaStackMixin, - OpenAICompletionToLlamaStackMixin, ): """ Adapter for RunPod's OpenAI-compatible API endpoints. @@ -41,44 +48,96 @@ class RunpodInferenceAdapter( """ def __init__(self, config: RunpodImplConfig) -> None: + OpenAIMixin.__init__(self) ModelRegistryHelper.__init__(self, MODEL_ENTRIES) self.config = config + def get_api_key(self) -> str: + """Get API key for OpenAI client.""" + return self.config.api_token + + def get_base_url(self) -> str: + """Get base URL for OpenAI client.""" + return self.config.url + async def initialize(self) -> None: pass async def shutdown(self) -> None: pass - async def register_model(self, model: Model) -> Model: - """ - Register any model with the runpod provider_id. + def get_extra_client_params(self) -> dict[str, Any]: + """Override to add RunPod-specific client parameters if needed.""" + return {} - Pass-through registration - accepts any model string that the RunPod endpoint serves. - No static model validation since RunPod endpoints can serve arbitrary vLLM models. + 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, + ): + """Override to add RunPod-specific stream_options requirement.""" + if stream and not stream_options: + stream_options = {"include_usage": True} - YAML Configuration Example: - models: - - metadata: {} - model_id: runpod/qwen/qwen3-8b - model_type: llm - provider_id: runpod - provider_model_id: qwen/qwen3-8b - - metadata: {} - model_id: runpod/deepcogito/cogito-v2-preview-llama-70B - model_type: llm - provider_id: runpod - provider_model_id: deepcogito/cogito-v2-preview-llama-70B + return await super().openai_chat_completion( + model=model, + 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, + ) - The provider strips 'runpod/' prefix before API calls: - "runpod/qwen/qwen3-8b" -> "qwen/qwen3-8b" + async def register_model(self, model: Model) -> Model: """ - if model.provider_id == "runpod": - logger.info( - f"Registering model: {model.identifier} -> {model.provider_resource_id}" - ) - return model - return await super().register_model(model) + Pass-through registration - accepts any model that the RunPod endpoint serves. + In the .yaml file the model: can be defined as example + models: + - metadata: {} + model_id: qwen3-32b-awq + model_type: llm + provider_id: runpod + provider_model_id: Qwen/Qwen3-32B-AWQ + """ + return model async def completion( self, @@ -88,12 +147,16 @@ async def completion( response_format: ResponseFormat | None = None, stream: bool | None = False, logprobs: LogProbConfig | None = None, - ) -> AsyncGenerator: + ) -> CompletionResponse | AsyncGenerator[CompletionResponseStreamChunk, None]: if sampling_params is None: sampling_params = SamplingParams() + # Resolve model_id to provider_resource_id + model = await self.model_store.get_model(model_id) + provider_model_id = model.provider_resource_id or model_id + request = CompletionRequest( - model=model_id, + model=provider_model_id, content=content, sampling_params=sampling_params, response_format=response_format, @@ -101,12 +164,10 @@ async def completion( logprobs=logprobs, ) - client = OpenAI(base_url=self.config.url, api_key=self.config.api_token) - if stream: - return self._stream_completion(request, client) + return self._stream_completion(request, self.client) else: - return await self._nonstream_completion(request, client) + return await self._nonstream_completion(request, self.client) async def chat_completion( self, @@ -120,13 +181,17 @@ async def chat_completion( stream: bool | None = False, logprobs: LogProbConfig | None = None, tool_config: ToolConfig | None = None, - ) -> AsyncGenerator: + ) -> ChatCompletionResponse | AsyncGenerator[ChatCompletionResponseStreamChunk, None]: """Process chat completion requests using RunPod's OpenAI-compatible API.""" if sampling_params is None: sampling_params = SamplingParams() + # Resolve model_id to provider_resource_id + model = await self.model_store.get_model(model_id) + provider_model_id = model.provider_resource_id or model_id + request = ChatCompletionRequest( - model=model_id, + model=provider_model_id, messages=messages, sampling_params=sampling_params, tools=tools or [], @@ -135,49 +200,34 @@ async def chat_completion( tool_config=tool_config, ) - client = OpenAI(base_url=self.config.url, api_key=self.config.api_token) - if stream: - return self._stream_chat_completion(request, client) + return self._stream_chat_completion(request, self.client) else: - return await self._nonstream_chat_completion(request, client) + return await self._nonstream_chat_completion(request, self.client) async def _nonstream_chat_completion( - self, request: ChatCompletionRequest, client: OpenAI + self, request: ChatCompletionRequest, client: AsyncOpenAI ) -> ChatCompletionResponse: params = await self._get_chat_params(request) - r = client.chat.completions.create(**params) + # Make actual RunPod API call + r = await client.chat.completions.create(**params) return process_chat_completion_response(r, request) async def _stream_chat_completion( - self, request: ChatCompletionRequest, client: OpenAI - ) -> AsyncGenerator: + self, request: ChatCompletionRequest, client: AsyncOpenAI + ) -> AsyncGenerator[ChatCompletionResponseStreamChunk, None]: params = await self._get_chat_params(request) - - async def _to_async_generator(): - s = client.chat.completions.create(**params) - for chunk in s: - yield chunk - - stream = _to_async_generator() + # Make actual RunPod API call for streaming + stream = await client.chat.completions.create(**params) async for chunk in process_chat_completion_stream_response(stream, request): yield chunk async def _get_chat_params(self, request: ChatCompletionRequest) -> dict: """Convert Llama Stack request to RunPod API parameters.""" - messages = [ - {"role": msg.role, "content": msg.content} for msg in request.messages - ] - - # Resolve model_id to provider_resource_id - model_obj = await self.model_store.get_model(request.model) - model = model_obj.provider_resource_id or request.model - - if model.startswith("runpod/"): - model = model.replace("runpod/", "", 1) + messages = [await convert_message_to_openai_dict(m, download=False) for m in request.messages] params = { - "model": model, + "model": request.model, "messages": messages, "stream": request.stream, **get_sampling_options(request.sampling_params), @@ -189,37 +239,27 @@ async def _get_chat_params(self, request: ChatCompletionRequest) -> dict: return params async def _nonstream_completion( - self, request: CompletionRequest, client: OpenAI + self, request: CompletionRequest, client: AsyncOpenAI ) -> CompletionResponse: params = await self._get_completion_params(request) - r = client.completions.create(**params) + # Make actual RunPod API call + r = await client.completions.create(**params) return process_completion_response(r) async def _stream_completion( - self, request: CompletionRequest, client: OpenAI + self, request: CompletionRequest, client: AsyncOpenAI ) -> AsyncGenerator: params = await self._get_completion_params(request) - - async def _to_async_generator(): - s = client.completions.create(**params) - for chunk in s: - yield chunk - - stream = _to_async_generator() + # Make actual RunPod API call for streaming + stream = await client.completions.create(**params) async for chunk in process_completion_stream_response(stream): yield chunk async def _get_completion_params(self, request: CompletionRequest) -> dict: - # Resolve model_id to provider_resource_id - model_obj = await self.model_store.get_model(request.model) - model = model_obj.provider_resource_id or request.model - - if model.startswith("runpod/"): - model = model.replace("runpod/", "", 1) - + """Convert Llama Stack request to RunPod API parameters.""" params = { - "model": model, - "prompt": completion_request_to_prompt(request), + "model": request.model, + "prompt": await completion_request_to_prompt(request), "stream": request.stream, **get_sampling_options(request.sampling_params), } @@ -241,16 +281,11 @@ async def embeddings( model_obj = await self.model_store.get_model(model_id) model = model_obj.provider_resource_id or model_id - if model.startswith("runpod/"): - model = model.replace("runpod/", "", 1) - - client = OpenAI(base_url=self.config.url, api_key=self.config.api_token) - kwargs = {} if output_dimension: kwargs["dimensions"] = output_dimension - response = client.embeddings.create( + response = await self.client.embeddings.create( model=model, input=[interleaved_content_as_str(content) for content in contents], **kwargs, @@ -269,19 +304,14 @@ async def openai_embeddings( ) -> OpenAIEmbeddingsResponse: # Resolve model_id to provider_resource_id model_obj = await self.model_store.get_model(model) - model_stripped = model_obj.provider_resource_id or model - - if model_stripped.startswith("runpod/"): - model_stripped = model_stripped.replace("runpod/", "", 1) - - client = OpenAI(base_url=self.config.url, api_key=self.config.api_token) + provider_model_id = model_obj.provider_resource_id or model - response = client.embeddings.create( - model=model_stripped, + response = await self.client.embeddings.create( + model=provider_model_id, input=input, encoding_format=encoding_format, dimensions=dimensions, user=user, ) - return response \ No newline at end of file + return response From 11b75a6b149b29d18db61ce1b0aa16901adb606d Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 6 Oct 2025 12:32:23 -0400 Subject: [PATCH 3/3] Remove unnecessary code --- .../remote/inference/runpod/runpod.py | 191 ++---------------- 1 file changed, 13 insertions(+), 178 deletions(-) diff --git a/llama_stack/providers/remote/inference/runpod/runpod.py b/llama_stack/providers/remote/inference/runpod/runpod.py index be7238628e..1a2af37b28 100644 --- a/llama_stack/providers/remote/inference/runpod/runpod.py +++ b/llama_stack/providers/remote/inference/runpod/runpod.py @@ -4,33 +4,18 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -from collections.abc import AsyncGenerator -import asyncio from typing import Any -from openai import AsyncOpenAI - -from llama_stack.apis.inference import * from llama_stack.apis.inference import ( + Inference, + OpenAIEmbeddingsResponse, OpenAIMessageParam, OpenAIResponseFormatParam, ) -from llama_stack.apis.common.content_types import InterleavedContentItem -from llama_stack.apis.models import Model, ModelType +from llama_stack.apis.models import Model from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper -from llama_stack.providers.utils.inference.openai_compat import ( - convert_message_to_openai_dict, - get_sampling_options, - process_chat_completion_response, - process_chat_completion_stream_response, - process_completion_response, - process_completion_stream_response, -) -from llama_stack.providers.utils.inference.prompt_adapter import ( - completion_request_to_prompt, - interleaved_content_as_str, -) from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin + from .config import RunpodImplConfig MODEL_ENTRIES = [] @@ -66,10 +51,6 @@ async def initialize(self) -> None: async def shutdown(self) -> None: pass - def get_extra_client_params(self) -> dict[str, Any]: - """Override to add RunPod-specific client parameters if needed.""" - return {} - async def openai_chat_completion( self, model: str, @@ -128,7 +109,7 @@ async def openai_chat_completion( async def register_model(self, model: Model) -> Model: """ - Pass-through registration - accepts any model that the RunPod endpoint serves. + Register a model and verify it's available on the RunPod endpoint. In the .yaml file the model: can be defined as example models: - metadata: {} @@ -137,162 +118,16 @@ async def register_model(self, model: Model) -> Model: provider_id: runpod provider_model_id: Qwen/Qwen3-32B-AWQ """ - return model - - 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 | AsyncGenerator[CompletionResponseStreamChunk, None]: - if sampling_params is None: - sampling_params = SamplingParams() - - # Resolve model_id to provider_resource_id - model = await self.model_store.get_model(model_id) - provider_model_id = model.provider_resource_id or model_id - - request = CompletionRequest( - model=provider_model_id, - content=content, - sampling_params=sampling_params, - response_format=response_format, - stream=stream, - logprobs=logprobs, - ) - - if stream: - return self._stream_completion(request, self.client) - else: - return await self._nonstream_completion(request, self.client) - - async def chat_completion( - self, - model_id: str, - messages: list[Message], - sampling_params: SamplingParams | None = None, - response_format: ResponseFormat | None = None, - 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 | AsyncGenerator[ChatCompletionResponseStreamChunk, None]: - """Process chat completion requests using RunPod's OpenAI-compatible API.""" - if sampling_params is None: - sampling_params = SamplingParams() - - # Resolve model_id to provider_resource_id - model = await self.model_store.get_model(model_id) - provider_model_id = model.provider_resource_id or model_id - - request = ChatCompletionRequest( - model=provider_model_id, - messages=messages, - sampling_params=sampling_params, - tools=tools or [], - stream=stream, - logprobs=logprobs, - tool_config=tool_config, - ) - - if stream: - return self._stream_chat_completion(request, self.client) - else: - return await self._nonstream_chat_completion(request, self.client) - - async def _nonstream_chat_completion( - self, request: ChatCompletionRequest, client: AsyncOpenAI - ) -> ChatCompletionResponse: - params = await self._get_chat_params(request) - # Make actual RunPod API call - r = await client.chat.completions.create(**params) - return process_chat_completion_response(r, request) - - async def _stream_chat_completion( - self, request: ChatCompletionRequest, client: AsyncOpenAI - ) -> AsyncGenerator[ChatCompletionResponseStreamChunk, None]: - params = await self._get_chat_params(request) - # Make actual RunPod API call for streaming - stream = await client.chat.completions.create(**params) - async for chunk in process_chat_completion_stream_response(stream, request): - yield chunk - - async def _get_chat_params(self, request: ChatCompletionRequest) -> dict: - """Convert Llama Stack request to RunPod API parameters.""" - messages = [await convert_message_to_openai_dict(m, download=False) for m in request.messages] - - params = { - "model": request.model, - "messages": messages, - "stream": request.stream, - **get_sampling_options(request.sampling_params), - } - - if request.stream: - params["stream_options"] = {"include_usage": True} - - return params - - async def _nonstream_completion( - self, request: CompletionRequest, client: AsyncOpenAI - ) -> CompletionResponse: - params = await self._get_completion_params(request) - # Make actual RunPod API call - r = await client.completions.create(**params) - return process_completion_response(r) + provider_model_id = model.provider_resource_id or model.identifier + is_available = await self.check_model_availability(provider_model_id) - async def _stream_completion( - self, request: CompletionRequest, client: AsyncOpenAI - ) -> AsyncGenerator: - params = await self._get_completion_params(request) - # Make actual RunPod API call for streaming - stream = await client.completions.create(**params) - async for chunk in process_completion_stream_response(stream): - yield chunk + if not is_available: + raise ValueError( + f"Model {provider_model_id} is not available on RunPod endpoint. " + f"Check your RunPod endpoint configuration." + ) - async def _get_completion_params(self, request: CompletionRequest) -> dict: - """Convert Llama Stack request to RunPod API parameters.""" - params = { - "model": request.model, - "prompt": await completion_request_to_prompt(request), - "stream": request.stream, - **get_sampling_options(request.sampling_params), - } - - if request.stream: - params["stream_options"] = {"include_usage": True} - - return params - - 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: - # Resolve model_id to provider_resource_id - model_obj = await self.model_store.get_model(model_id) - model = model_obj.provider_resource_id or model_id - - kwargs = {} - if output_dimension: - kwargs["dimensions"] = output_dimension - - response = await self.client.embeddings.create( - model=model, - input=[interleaved_content_as_str(content) for content in contents], - **kwargs, - ) - - embeddings = [data.embedding for data in response.data] - return EmbeddingsResponse(embeddings=embeddings) + return model async def openai_embeddings( self,