phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-06-27 18:50:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

migrate model to Resource and new registration signature (#410)

Browse source 
* resource oriented object design for models

* add back llama_model field

* working tests

* register singature fix

* address feedback

---------

Co-authored-by: Dinesh Yeduguru <dineshyv@fb.com>
This commit is contained in:
Dinesh Yeduguru 2024-11-08 16:12:57 -08:00 • committed by GitHub
parent bd0622ef10
commit ec644d3418
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 99 additions and 90 deletions
Download patch file Download diff file Expand all files Collapse all files

33
llama_stack/distribution/routers/routing_tables.py
View file

@ -84,8 +84,6 @@ class CommonRoutingTableImpl(RoutingTable):
api = get_impl_api(p)
if api == Api.inference:
p.model_store = self
models = await p.list_models()
await add_objects(models, pid, ModelDefWithProvider)
elif api == Api.safety:
p.shield_store = self
@ -198,14 +196,39 @@ class CommonRoutingTableImpl(RoutingTable):
class ModelsRoutingTable(CommonRoutingTableImpl, Models):
async def list_models(self) -> List[ModelDefWithProvider]:
async def list_models(self) -> List[Model]:
return await self.get_all_with_type("model")
async def get_model(self, identifier: str) -> Optional[ModelDefWithProvider]:
async def get_model(self, identifier: str) -> Optional[Model]:
return await self.get_object_by_identifier(identifier)
async def register_model(self, model: ModelDefWithProvider) -> None:
async def register_model(
self,
model_id: str,
provider_model_id: Optional[str] = None,
provider_id: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
) -> Model:
if provider_model_id is None:
provider_model_id = model_id
if provider_id is None:
# If provider_id not specified, use the only provider if it supports this model
if len(self.impls_by_provider_id) == 1:
provider_id = list(self.impls_by_provider_id.keys())[0]
else:
raise ValueError(
"No provider specified and multiple providers available. Please specify a provider_id. Available providers: {self.impls_by_provider_id.keys()}"
)
if metadata is None:
metadata = {}
model = Model(
identifier=model_id,
provider_resource_id=provider_model_id,
provider_id=provider_id,
metadata=metadata,
)
await self.register_object(model)
return model
class ShieldsRoutingTable(CommonRoutingTableImpl, Shields):