From b380cb463fdca94b682c9d9be34f03b9b475c5b4 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Tue, 3 Jun 2025 11:04:23 -0700 Subject: [PATCH] feat: add postgres deps to starter distro (#2360) Once we have this, we can use the starter distro for the Kubernetes cluster demos. --- llama_stack/providers/utils/kvstore/config.py | 8 ++--- .../providers/utils/sqlstore/sqlstore.py | 11 +++++++ .../templates/postgres-demo/postgres_demo.py | 14 ++------ llama_stack/templates/postgres-demo/run.yaml | 2 +- llama_stack/templates/starter/build.yaml | 2 ++ llama_stack/templates/starter/run.yaml | 7 ++++ llama_stack/templates/starter/starter.py | 32 +++++++++++++++++++ llama_stack/templates/template.py | 8 +++++ 8 files changed, 68 insertions(+), 16 deletions(-) diff --git a/llama_stack/providers/utils/kvstore/config.py b/llama_stack/providers/utils/kvstore/config.py index bbb0c5c0a..18a6eb8dd 100644 --- a/llama_stack/providers/utils/kvstore/config.py +++ b/llama_stack/providers/utils/kvstore/config.py @@ -72,15 +72,15 @@ class PostgresKVStoreConfig(CommonConfig): table_name: str = "llamastack_kvstore" @classmethod - def sample_run_config(cls, table_name: str = "llamastack_kvstore"): + def sample_run_config(cls, table_name: str = "llamastack_kvstore", **kwargs): return { "type": "postgres", "namespace": None, "host": "${env.POSTGRES_HOST:localhost}", "port": "${env.POSTGRES_PORT:5432}", - "db": "${env.POSTGRES_DB}", - "user": "${env.POSTGRES_USER}", - "password": "${env.POSTGRES_PASSWORD}", + "db": "${env.POSTGRES_DB:llamastack}", + "user": "${env.POSTGRES_USER:llamastack}", + "password": "${env.POSTGRES_PASSWORD:llamastack}", "table_name": "${env.POSTGRES_TABLE_NAME:" + table_name + "}", } diff --git a/llama_stack/providers/utils/sqlstore/sqlstore.py b/llama_stack/providers/utils/sqlstore/sqlstore.py index bd78137c0..edc7672a3 100644 --- a/llama_stack/providers/utils/sqlstore/sqlstore.py +++ b/llama_stack/providers/utils/sqlstore/sqlstore.py @@ -74,6 +74,17 @@ class PostgresSqlStoreConfig(SqlAlchemySqlStoreConfig): def pip_packages(self) -> list[str]: return super().pip_packages + ["asyncpg"] + @classmethod + def sample_run_config(cls, **kwargs): + return cls( + 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}", + ) + SqlStoreConfig = Annotated[ SqliteSqlStoreConfig | PostgresSqlStoreConfig, diff --git a/llama_stack/templates/postgres-demo/postgres_demo.py b/llama_stack/templates/postgres-demo/postgres_demo.py index d2e352320..2edd6d0e7 100644 --- a/llama_stack/templates/postgres-demo/postgres_demo.py +++ b/llama_stack/templates/postgres-demo/postgres_demo.py @@ -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={ diff --git a/llama_stack/templates/postgres-demo/run.yaml b/llama_stack/templates/postgres-demo/run.yaml index 889b8eaa7..9dc366434 100644 --- a/llama_stack/templates/postgres-demo/run.yaml +++ b/llama_stack/templates/postgres-demo/run.yaml @@ -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} diff --git a/llama_stack/templates/starter/build.yaml b/llama_stack/templates/starter/build.yaml index ec97c7d3e..26b48e9c9 100644 --- a/llama_stack/templates/starter/build.yaml +++ b/llama_stack/templates/starter/build.yaml @@ -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] diff --git a/llama_stack/templates/starter/run.yaml b/llama_stack/templates/starter/run.yaml index 04cfedcab..ec4a81085 100644 --- a/llama_stack/templates/starter/run.yaml +++ b/llama_stack/templates/starter/run.yaml @@ -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: {} diff --git a/llama_stack/templates/starter/starter.py b/llama_stack/templates/starter/starter.py index 0932bfdfe..efec5fe69 100644 --- a/llama_stack/templates/starter/starter.py +++ b/llama_stack/templates/starter/starter.py @@ -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", + ), }, ) diff --git a/llama_stack/templates/template.py b/llama_stack/templates/template.py index 4013f08f9..f00be28ad 100644 --- a/llama_stack/templates/template.py +++ b/llama_stack/templates/template.py @@ -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,