phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-08 19:10:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Push registration methods onto the backing providers

Browse source
This commit is contained in:
Ashwin Bharambe 2024-10-05 22:17:06 -07:00 • committed by Ashwin Bharambe
parent 5a7b01d292
commit 4215cc9331
14 changed files with 269 additions and 220 deletions
Download patch file Download diff file Expand all files Collapse all files

17
llama_stack/distribution/datatypes.py
View file

@ -14,6 +14,9 @@ from llama_stack.providers.datatypes import * # noqa: F403
from llama_stack.apis.models import * # noqa: F403
from llama_stack.apis.shields import * # noqa: F403
from llama_stack.apis.memory_banks import * # noqa: F403
from llama_stack.apis.inference import Inference
from llama_stack.apis.memory import Memory
from llama_stack.apis.safety import Safety
LLAMA_STACK_BUILD_CONFIG_VERSION = "2"
@ -23,6 +26,19 @@ LLAMA_STACK_RUN_CONFIG_VERSION = "2"
RoutingKey = Union[str, List[str]]
RoutableObject = Union[
ModelDef,
ShieldDef,
MemoryBankDef,
]
RoutedProtocol = Union[
Inference,
Safety,
Memory,
]
class GenericProviderConfig(BaseModel):
provider_type: str
config: Dict[str, Any]
@ -56,6 +72,7 @@ class RoutingTableProviderSpec(ProviderSpec):
docker_image: Optional[str] = None
router_api: Api
registry: List[RoutableObject]
module: str
pip_packages: List[str] = Field(default_factory=list)

194
llama_stack/distribution/resolver.py
View file

@ -28,46 +28,48 @@ async def resolve_impls_with_routing(run_config: StackRunConfig) -> Dict[Api, An
"""
all_api_providers = get_provider_registry()
auto_routed_apis = builtin_automatically_routed_apis()
routing_table_apis = set(
x.routing_table_api for x in builtin_automatically_routed_apis()
)
router_apis = set(x.router_api for x in builtin_automatically_routed_apis())
providers_with_specs = {}
for api_str, instances in run_config.providers.items():
for api_str, providers in run_config.providers.items():
api = Api(api_str)
if api in [a.routing_table_api for a in auto_routed_apis]:
if api in routing_table_apis:
raise ValueError(
f"Provider for `{api_str}` is automatically provided and cannot be overridden"
)
providers_with_specs[api] = {}
for config in instances:
if config.provider_type not in all_api_providers[api]:
specs = {}
for provider in providers:
if provider.provider_type not in all_api_providers[api]:
raise ValueError(
f"Provider `{config.provider_type}` is not available for API `{api}`"
f"Provider `{provider.provider_type}` is not available for API `{api}`"
)
spec = ProviderWithSpec(
spec=all_api_providers[api][config.provider_type],
**config,
spec=all_api_providers[api][provider.provider_type],
**(provider.dict()),
)
providers_with_specs[api][spec.provider_id] = spec
specs[provider.provider_id] = spec
key = api_str if api not in router_apis else f"inner-{api_str}"
providers_with_specs[key] = specs
apis_to_serve = run_config.apis_to_serve or set(
list(providers_with_specs.keys())
+ [a.routing_table_api.value for a in auto_routed_apis]
list(providers_with_specs.keys()) + list(routing_table_apis)
)
for info in builtin_automatically_routed_apis():
if info.router_api.value not in apis_to_serve:
continue
if info.routing_table_api.value not in run_config:
raise ValueError(
f"Registry for `{info.routing_table_api.value}` is not provided?"
)
available_providers = providers_with_specs[info.router_api]
available_providers = providers_with_specs[f"inner-{info.router_api.value}"]
inner_deps = []
registry = run_config[info.routing_table_api.value]
registry = getattr(run_config, info.routing_table_api.value)
for entry in registry:
if entry.provider_id not in available_providers:
raise ValueError(
@ -77,74 +79,70 @@ async def resolve_impls_with_routing(run_config: StackRunConfig) -> Dict[Api, An
provider = available_providers[entry.provider_id]
inner_deps.extend(provider.spec.api_dependencies)
providers_with_specs[info.routing_table_api] = {
"__builtin__": [
ProviderWithSpec(
provider_id="__builtin__",
provider_type="__builtin__",
config=registry,
spec=RoutingTableProviderSpec(
api=info.routing_table_api,
router_api=info.router_api,
module="llama_stack.distribution.routers",
api_dependencies=inner_deps,
),
)
]
providers_with_specs[info.routing_table_api.value] = {
"__builtin__": ProviderWithSpec(
provider_id="__builtin__",
provider_type="__routing_table__",
config={},
spec=RoutingTableProviderSpec(
api=info.routing_table_api,
router_api=info.router_api,
registry=registry,
module="llama_stack.distribution.routers",
api_dependencies=inner_deps,
),
)
}
providers_with_specs[info.router_api] = {
"__builtin__": [
ProviderWithSpec(
provider_id="__builtin__",
provider_type="__builtin__",
config={},
spec=AutoRoutedProviderSpec(
api=info.router_api,
module="llama_stack.distribution.routers",
routing_table_api=source_api,
api_dependencies=[source_api],
),
)
]
providers_with_specs[info.router_api.value] = {
"__builtin__": ProviderWithSpec(
provider_id="__builtin__",
provider_type="__autorouted__",
config={},
spec=AutoRoutedProviderSpec(
api=info.router_api,
module="llama_stack.distribution.routers",
routing_table_api=info.routing_table_api,
api_dependencies=[info.routing_table_api],
),
)
}
sorted_providers = topological_sort(providers_with_specs)
sorted_providers = topological_sort(
{k: v.values() for k, v in providers_with_specs.items()}
)
sorted_providers.append(
ProviderWithSpec(
provider_id="__builtin__",
provider_type="__builtin__",
config={},
spec=InlineProviderSpec(
api=Api.inspect,
(
"inspect",
ProviderWithSpec(
provider_id="__builtin__",
provider_type="__builtin__",
config_class="llama_stack.distribution.inspect.DistributionInspectConfig",
module="llama_stack.distribution.inspect",
config={},
spec=InlineProviderSpec(
api=Api.inspect,
provider_type="__builtin__",
config_class="llama_stack.distribution.inspect.DistributionInspectConfig",
module="llama_stack.distribution.inspect",
),
),
)
)
print(f"Resolved {len(sorted_providers)} providers in topological order")
for provider in sorted_providers:
print(
f" {provider.spec.api}: ({provider.provider_id}) {provider.spec.provider_type}"
)
for api_str, provider in sorted_providers:
print(f" {api_str}: ({provider.provider_id}) {provider.spec.provider_type}")
print("")
impls = {}
impls_by_provider_id = {}
for provider in sorted_providers:
api = provider.spec.api
if api not in impls_by_provider_id:
impls_by_provider_id[api] = {}
deps = {api: impls[api] for api in provider.spec.api_dependencies}
inner_impls_by_provider_id = {f"inner-{x}": {} for x in router_apis}
for api_str, provider in sorted_providers:
deps = {a: impls[a] for a in provider.spec.api_dependencies}
inner_impls = {}
if isinstance(provider.spec, RoutingTableProviderSpec):
for entry in provider.config:
inner_impls[entry.provider_id] = impls_by_provider_id[
provider.spec.router_api
for entry in provider.spec.registry:
inner_impls[entry.provider_id] = inner_impls_by_provider_id[
f"inner-{provider.spec.router_api.value}"
][entry.provider_id]
impl = await instantiate_provider(
@ -152,37 +150,46 @@ async def resolve_impls_with_routing(run_config: StackRunConfig) -> Dict[Api, An
deps,
inner_impls,
)
impls[api] = impl
impls_by_provider_id[api][provider.provider_id] = impl
if "inner-" in api_str:
inner_impls_by_provider_id[api_str][provider.provider_id] = impl
else:
api = Api(api_str)
impls[api] = impl
return impls
def topological_sort(
providers_with_specs: Dict[Api, List[ProviderWithSpec]],
providers_with_specs: Dict[str, List[ProviderWithSpec]],
) -> List[ProviderWithSpec]:
def dfs(kv, visited: Set[Api], stack: List[Api]):
api, providers = kv
visited.add(api)
def dfs(kv, visited: Set[str], stack: List[str]):
api_str, providers = kv
visited.add(api_str)
deps = [dep for x in providers for dep in x.api_dependencies]
for api in deps:
if api not in visited:
dfs((api, providers_with_specs[api]), visited, stack)
deps = []
for provider in providers:
for dep in provider.spec.api_dependencies:
deps.append(dep.value)
if isinstance(provider, AutoRoutedProviderSpec):
deps.append(f"inner-{provider.api}")
stack.append(api)
for dep in deps:
if dep not in visited:
dfs((dep, providers_with_specs[dep]), visited, stack)
stack.append(api_str)
visited = set()
stack = []
for api, providers in providers_with_specs.items():
if api not in visited:
dfs((api, providers), visited, stack)
for api_str, providers in providers_with_specs.items():
if api_str not in visited:
dfs((api_str, providers), visited, stack)
flattened = []
for api in stack:
flattened.extend(providers_with_specs[api])
for api_str in stack:
for provider in providers_with_specs[api_str]:
flattened.append((api_str, provider))
return flattened
@ -202,9 +209,8 @@ async def instantiate_provider(
else:
method = "get_client_impl"
assert isinstance(provider_config, GenericProviderConfig)
config_type = instantiate_class_type(provider_spec.config_class)
config = config_type(**provider_config.config)
config = config_type(**provider.config)
args = [config, deps]
elif isinstance(provider_spec, AutoRoutedProviderSpec):
method = "get_auto_router_impl"
@ -214,17 +220,13 @@ async def instantiate_provider(
elif isinstance(provider_spec, RoutingTableProviderSpec):
method = "get_routing_table_impl"
assert isinstance(provider_config, list)
registry = provider_config
config = None
args = [provider_spec.api, registry, inner_impls, deps]
args = [provider_spec.api, provider_spec.registry, inner_impls, deps]
else:
method = "get_provider_impl"
assert isinstance(provider_config, GenericProviderConfig)
config_type = instantiate_class_type(provider_spec.config_class)
config = config_type(**provider_config.config)
config = config_type(**provider.config)
args = [config, deps]
fn = getattr(module, method)

2
llama_stack/distribution/routers/__init__.py
View file

@ -10,8 +10,6 @@ from llama_stack.distribution.datatypes import * # noqa: F403
from .routing_tables import (
MemoryBanksRoutingTable,
ModelsRoutingTable,
RoutableObject,
RoutedProtocol,
ShieldsRoutingTable,
)

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

@ -4,33 +4,17 @@
# 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, List, Optional, Union
from typing import Any, List, Optional
from llama_models.llama3.api.datatypes import * # noqa: F403
from llama_stack.apis.models import * # noqa: F403
from llama_stack.apis.shields import * # noqa: F403
from llama_stack.apis.memory_banks import * # noqa: F403
from llama_stack.apis.inference import Inference
from llama_stack.apis.memory import Memory
from llama_stack.apis.safety import Safety
from llama_stack.distribution.datatypes import * # noqa: F403
RoutableObject = Union[
ModelDef,
ShieldDef,
MemoryBankDef,
]
RoutedProtocol = Union[
Inference,
Safety,
Memory,
]
class CommonRoutingTableImpl(RoutingTable):
def __init__(
self,
@ -46,19 +30,14 @@ class CommonRoutingTableImpl(RoutingTable):
self.impls_by_provider_id = impls_by_provider_id
self.registry = registry
async def initialize(self) -> None:
keys_by_provider = {}
self.routing_key_to_object = {}
for obj in self.registry:
keys = keys_by_provider.setdefault(obj.provider_id, [])
keys.append(obj.routing_key)
self.routing_key_to_object[obj.identifier] = obj
for provider_id, keys in keys_by_provider.items():
p = self.impls_by_provider_id[provider_id]
spec = p.__provider_spec__
if is_passthrough(spec):
continue
await p.validate_routing_keys(keys)
async def initialize(self) -> None:
for obj in self.registry:
p = self.impls_by_provider_id[obj.provider_id]
await self.register_object(obj, p)
async def shutdown(self) -> None:
pass
@ -75,8 +54,24 @@ class CommonRoutingTableImpl(RoutingTable):
return obj
return None
def register_object(self, obj: RoutableObject) -> None:
if obj.identifier in self.routing_key_to_object:
raise ValueError(f"Object `{obj.identifier}` already registered")
if obj.provider_id not in self.impls_by_provider_id:
raise ValueError(f"Provider `{obj.provider_id}` not found")
p = self.impls_by_provider_id[obj.provider_id]
await p.register_object(obj)
self.routing_key_to_object[obj.identifier] = obj
self.registry.append(obj)
class ModelsRoutingTable(CommonRoutingTableImpl, Models):
async def register_object(self, obj: ModelDef, p: Inference) -> None:
await p.register_model(obj)
async def list_models(self) -> List[ModelDef]:
return self.registry
@ -84,10 +79,13 @@ class ModelsRoutingTable(CommonRoutingTableImpl, Models):
return self.get_object_by_identifier(identifier)
async def register_model(self, model: ModelDef) -> None:
raise NotImplementedError()
await self.register_object(model)
class ShieldsRoutingTable(CommonRoutingTableImpl, Shields):
async def register_object(self, obj: ShieldDef, p: Safety) -> None:
await p.register_shield(obj)
async def list_shields(self) -> List[ShieldDef]:
return self.registry
@ -95,10 +93,13 @@ class ShieldsRoutingTable(CommonRoutingTableImpl, Shields):
return self.get_object_by_identifier(shield_type)
async def register_shield(self, shield: ShieldDef) -> None:
raise NotImplementedError()
await self.register_object(shield)
class MemoryBanksRoutingTable(CommonRoutingTableImpl, MemoryBanks):
async def register_object(self, obj: MemoryBankDef, p: Memory) -> None:
await p.register_memory_bank(obj)
async def list_memory_banks(self) -> List[MemoryBankDef]:
return self.registry
@ -106,4 +107,4 @@ class MemoryBanksRoutingTable(CommonRoutingTableImpl, MemoryBanks):
return self.get_object_by_identifier(identifier)
async def register_memory_bank(self, bank: MemoryBankDef) -> None:
raise NotImplementedError()
await self.register_object(bank)