feat: add postgres deps to starter distro (#2360)

Once we have this, we can use the starter distro for the Kubernetes
cluster demos.
This commit is contained in:
Ashwin Bharambe 2025-06-03 11:04:23 -07:00 committed by GitHub
parent e743257d1d
commit b380cb463f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 68 additions and 16 deletions

View file

@ -101,15 +101,7 @@ def get_distribution_template() -> DistributionTemplate:
provider_id="vllm-inference",
)
)
postgres_config = {
"type": "postgres",
"host": "${env.POSTGRES_HOST:localhost}",
"port": "${env.POSTGRES_PORT:5432}",
"db": "${env.POSTGRES_DB:llamastack}",
"user": "${env.POSTGRES_USER:llamastack}",
"password": "${env.POSTGRES_PASSWORD:llamastack}",
}
postgres_config = PostgresSqlStoreConfig.sample_run_config()
return DistributionTemplate(
name=name,
distro_type="self_hosted",
@ -147,8 +139,8 @@ def get_distribution_template() -> DistributionTemplate:
default_models=default_models,
default_tool_groups=default_tool_groups,
default_shields=[ShieldInput(shield_id="meta-llama/Llama-Guard-3-8B")],
metadata_store=PostgresKVStoreConfig.model_validate(postgres_config),
inference_store=PostgresSqlStoreConfig.model_validate(postgres_config),
metadata_store=PostgresKVStoreConfig.sample_run_config(),
inference_store=postgres_config,
),
},
run_config_env_vars={

View file

@ -79,7 +79,7 @@ metadata_store:
db: ${env.POSTGRES_DB:llamastack}
user: ${env.POSTGRES_USER:llamastack}
password: ${env.POSTGRES_PASSWORD:llamastack}
table_name: llamastack_kvstore
table_name: ${env.POSTGRES_TABLE_NAME:llamastack_kvstore}
inference_store:
type: postgres
host: ${env.POSTGRES_HOST:localhost}

View file

@ -9,6 +9,7 @@ distribution_spec:
- remote::gemini
- remote::groq
- remote::sambanova
- remote::remote-vllm
- inline::sentence-transformers
vector_io:
- inline::sqlite-vec
@ -37,4 +38,5 @@ distribution_spec:
image_type: conda
additional_pip_packages:
- aiosqlite
- asyncpg
- sqlalchemy[asyncio]

View file

@ -39,6 +39,13 @@ providers:
config:
url: https://api.sambanova.ai/v1
api_key: ${env.SAMBANOVA_API_KEY:}
- provider_id: remote-vllm
provider_type: remote::remote-vllm
config:
url: ${env.VLLM_URL:http://localhost:8000/v1}
max_tokens: ${env.VLLM_MAX_TOKENS:4096}
api_token: ${env.VLLM_API_TOKEN:fake}
tls_verify: ${env.VLLM_TLS_VERIFY:true}
- provider_id: sentence-transformers
provider_type: inline::sentence-transformers
config: {}

View file

@ -42,11 +42,13 @@ from llama_stack.providers.remote.inference.sambanova.config import SambaNovaImp
from llama_stack.providers.remote.inference.sambanova.models import (
MODEL_ENTRIES as SAMBANOVA_MODEL_ENTRIES,
)
from llama_stack.providers.remote.inference.vllm import VLLMInferenceAdapterConfig
from llama_stack.providers.remote.vector_io.chroma.config import ChromaVectorIOConfig
from llama_stack.providers.remote.vector_io.pgvector.config import (
PGVectorVectorIOConfig,
)
from llama_stack.providers.utils.inference.model_registry import ProviderModelEntry
from llama_stack.providers.utils.sqlstore.sqlstore import PostgresSqlStoreConfig
from llama_stack.templates.template import (
DistributionTemplate,
RunConfigSettings,
@ -87,6 +89,13 @@ def get_inference_providers() -> tuple[list[Provider], dict[str, list[ProviderMo
SAMBANOVA_MODEL_ENTRIES,
SambaNovaImplConfig.sample_run_config(api_key="${env.SAMBANOVA_API_KEY:}"),
),
(
"remote-vllm",
[],
VLLMInferenceAdapterConfig.sample_run_config(
url="${env.VLLM_URL:http://localhost:8000/v1}",
),
),
]
inference_providers = []
available_models = {}
@ -169,6 +178,8 @@ def get_distribution_template() -> DistributionTemplate:
)
default_models = get_model_registry(available_models)
postgres_store = PostgresSqlStoreConfig.sample_run_config()
return DistributionTemplate(
name=name,
distro_type="self_hosted",
@ -177,6 +188,7 @@ def get_distribution_template() -> DistributionTemplate:
template_path=None,
providers=providers,
available_models_by_provider=available_models,
additional_pip_packages=postgres_store.pip_packages,
run_configs={
"run.yaml": RunConfigSettings(
provider_overrides={
@ -201,5 +213,25 @@ def get_distribution_template() -> DistributionTemplate:
"",
"OpenAI API Key",
),
"GROQ_API_KEY": (
"",
"Groq API Key",
),
"ANTHROPIC_API_KEY": (
"",
"Anthropic API Key",
),
"GEMINI_API_KEY": (
"",
"Gemini API Key",
),
"SAMBANOVA_API_KEY": (
"",
"SambaNova API Key",
),
"VLLM_URL": (
"http://localhost:8000/v1",
"VLLM URL",
),
},
)

View file

@ -154,6 +154,11 @@ class DistributionTemplate(BaseModel):
available_models_by_provider: dict[str, list[ProviderModelEntry]] | None = None
# we may want to specify additional pip packages without necessarily indicating a
# specific "default" inference store (which is what typically used to dictate additional
# pip packages)
additional_pip_packages: list[str] | None = None
def build_config(self) -> BuildConfig:
additional_pip_packages: list[str] = []
for run_config in self.run_configs.values():
@ -161,6 +166,9 @@ class DistributionTemplate(BaseModel):
if run_config_.inference_store:
additional_pip_packages.extend(run_config_.inference_store.pip_packages)
if self.additional_pip_packages:
additional_pip_packages.extend(self.additional_pip_packages)
return BuildConfig(
distribution_spec=DistributionSpec(
description=self.description,