phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-29 15:23:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Ensure models are downloaded before serving in Ollama inference

Browse source
This commit is contained in:
Frieda (Jingying) Huang 2024-10-06 12:09:22 -04:00
parent 0edf24b227
commit 969a11fb8a
2 changed files with 8 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

10
llama_stack/providers/adapters/inference/ollama/__init__.py
View file

@ -15,12 +15,6 @@ async def get_adapter_impl(config: RemoteProviderConfig, _deps):
from .ollama import OllamaInferenceAdapter
impl = OllamaInferenceAdapter(config.url)
routing_key = _deps.get("routing_key")
if not routing_key:
raise ValueError(
"Routing key is required for the Ollama adapter but was not found."
)
await impl.initialize(routing_key)
impl._deps = _deps
await impl.initialize()
return impl

7
llama_stack/providers/adapters/inference/ollama/ollama.py
View file

@ -45,10 +45,15 @@ class OllamaInferenceAdapter(Inference, RoutableProviderForModels):
def client(self) -> AsyncClient:
return AsyncClient(host=self.url)
async def initialize(self, routing_key: str) -> None:
async def initialize(self) -> None:
print("Initializing Ollama, checking connectivity to server...")
try:
await self.client.ps()
routing_key = self._deps.get("routing_key")
if not routing_key:
raise ValueError(
"Routing key is required for the Ollama adapter but was not found."
)
ollama_model = self.map_to_provider_model(routing_key)
print(f"Connected to Ollama server. Pre-downloading {ollama_model}...")
await self.predownload_models(ollama_model)