From 94f83382ebcdc77a034d9f55f02111e0f504d371 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Han?= Date: Fri, 18 Apr 2025 17:18:28 +0200 Subject: [PATCH] feat: allow building distro with external providers (#1967) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # What does this PR do? We can now build a distribution that includes external providers. Closes: https://github.com/meta-llama/llama-stack/issues/1948 ## Test Plan Build a distro with an external provider following the doc instructions. [//]: # (## Documentation) Added. Rendered: ![Screenshot 2025-04-18 at 11 26 39](https://github.com/user-attachments/assets/afcf3d50-8d30-48c3-8d24-06a4b3662881) Signed-off-by: Sébastien Han --- .github/workflows/test-external-providers.yml | 23 ++++++-- docs/source/distributions/building_distro.md | 55 +++++++++++++++++++ llama_stack/cli/stack/_build.py | 31 +++++++---- llama_stack/distribution/build.py | 25 +++++---- llama_stack/distribution/build_container.sh | 2 +- llama_stack/distribution/datatypes.py | 9 +++ llama_stack/distribution/distribution.py | 10 ++-- scripts/distro_codegen.py | 2 +- .../custom-distro.yaml | 9 +++ .../custom_ollama.yaml | 2 +- .../llama-stack-provider-ollama/run.yaml | 38 +------------ 11 files changed, 137 insertions(+), 69 deletions(-) create mode 100644 tests/external-provider/llama-stack-provider-ollama/custom-distro.yaml diff --git a/.github/workflows/test-external-providers.yml b/.github/workflows/test-external-providers.yml index 2ead8f845..f7801c8d3 100644 --- a/.github/workflows/test-external-providers.yml +++ b/.github/workflows/test-external-providers.yml @@ -9,6 +9,11 @@ on: jobs: test-external-providers: runs-on: ubuntu-latest + strategy: + matrix: + image-type: [venv] + # We don't do container yet, it's tricky to install a package from the host into the + # container and point 'uv pip install' to the correct path... steps: - name: Checkout repository uses: actions/checkout@v4 @@ -35,17 +40,25 @@ jobs: uv sync --extra dev --extra test uv pip install -e . - - name: Install Ollama custom provider + - name: Apply image type to config file + run: | + yq -i '.image_type = "${{ matrix.image-type }}"' tests/external-provider/llama-stack-provider-ollama/custom-distro.yaml + cat tests/external-provider/llama-stack-provider-ollama/custom-distro.yaml + + - name: Setup directory for Ollama custom provider run: | mkdir -p tests/external-provider/llama-stack-provider-ollama/src/ cp -a llama_stack/providers/remote/inference/ollama/ tests/external-provider/llama-stack-provider-ollama/src/llama_stack_provider_ollama - uv pip install tests/external-provider/llama-stack-provider-ollama - name: Create provider configuration run: | mkdir -p /tmp/providers.d/remote/inference cp tests/external-provider/llama-stack-provider-ollama/custom_ollama.yaml /tmp/providers.d/remote/inference/custom_ollama.yaml + - name: Build distro from config file + run: | + USE_COPY_NOT_MOUNT=true LLAMA_STACK_DIR=. uv run llama stack build --config tests/external-provider/llama-stack-provider-ollama/custom-distro.yaml + - name: Wait for Ollama to start run: | echo "Waiting for Ollama..." @@ -62,11 +75,13 @@ jobs: exit 1 - name: Start Llama Stack server in background + if: ${{ matrix.image-type }} == 'venv' env: INFERENCE_MODEL: "meta-llama/Llama-3.2-3B-Instruct" run: | - source .venv/bin/activate - nohup uv run llama stack run tests/external-provider/llama-stack-provider-ollama/run.yaml --image-type venv > server.log 2>&1 & + source ci-test/bin/activate + uv run pip list + nohup uv run --active llama stack run tests/external-provider/llama-stack-provider-ollama/run.yaml --image-type ${{ matrix.image-type }} > server.log 2>&1 & - name: Wait for Llama Stack server to be ready run: | diff --git a/docs/source/distributions/building_distro.md b/docs/source/distributions/building_distro.md index ad5d3bff4..4c342b14b 100644 --- a/docs/source/distributions/building_distro.md +++ b/docs/source/distributions/building_distro.md @@ -176,7 +176,11 @@ distribution_spec: safety: inline::llama-guard agents: inline::meta-reference telemetry: inline::meta-reference +image_name: ollama image_type: conda + +# If some providers are external, you can specify the path to the implementation +external_providers_dir: /etc/llama-stack/providers.d ``` ``` @@ -184,6 +188,57 @@ llama stack build --config llama_stack/templates/ollama/build.yaml ``` ::: +:::{tab-item} Building with External Providers + +Llama Stack supports external providers that live outside of the main codebase. This allows you to create and maintain your own providers independently or use community-provided providers. + +To build a distribution with external providers, you need to: + +1. Configure the `external_providers_dir` in your build configuration file: + +```yaml +# Example my-external-stack.yaml with external providers +version: '2' +distribution_spec: + description: Custom distro for CI tests + providers: + inference: + - remote::custom_ollama +# Add more providers as needed +image_type: container +image_name: ci-test +# Path to external provider implementations +external_providers_dir: /etc/llama-stack/providers.d +``` + +Here's an example for a custom Ollama provider: + +```yaml +adapter: + adapter_type: custom_ollama + pip_packages: + - ollama + - aiohttp + - llama-stack-provider-ollama # This is the provider package + config_class: llama_stack_ollama_provider.config.OllamaImplConfig + module: llama_stack_ollama_provider +api_dependencies: [] +optional_api_dependencies: [] +``` + +The `pip_packages` section lists the Python packages required by the provider, as well as the +provider package itself. The package must be available on PyPI or can be provided from a local +directory or a git repository (git must be installed on the build environment). + +2. Build your distribution using the config file: + +``` +llama stack build --config my-external-stack.yaml +``` + +For more information on external providers, including directory structure, provider types, and implementation requirements, see the [External Providers documentation](../providers/external.md). +::: + :::{tab-item} Building Container ```{admonition} Podman Alternative diff --git a/llama_stack/cli/stack/_build.py b/llama_stack/cli/stack/_build.py index 760ba2e5a..26c09af4e 100644 --- a/llama_stack/cli/stack/_build.py +++ b/llama_stack/cli/stack/_build.py @@ -210,16 +210,9 @@ def run_stack_build_command(args: argparse.Namespace) -> None: ) sys.exit(1) - if build_config.image_type == LlamaStackImageType.CONTAINER.value and not args.image_name: - cprint( - "Please specify --image-name when building a container from a config file", - color="red", - ) - sys.exit(1) - if args.print_deps_only: print(f"# Dependencies for {args.template or args.config or image_name}") - normal_deps, special_deps = get_provider_dependencies(build_config.distribution_spec.providers) + normal_deps, special_deps = get_provider_dependencies(build_config) normal_deps += SERVER_DEPENDENCIES print(f"uv pip install {' '.join(normal_deps)}") for special_dep in special_deps: @@ -274,9 +267,10 @@ def _generate_run_config( image_name=image_name, apis=apis, providers={}, + external_providers_dir=build_config.external_providers_dir if build_config.external_providers_dir else None, ) # build providers dict - provider_registry = get_provider_registry() + provider_registry = get_provider_registry(build_config) for api in apis: run_config.providers[api] = [] provider_types = build_config.distribution_spec.providers[api] @@ -290,8 +284,22 @@ def _generate_run_config( if p.deprecation_error: raise InvalidProviderError(p.deprecation_error) - config_type = instantiate_class_type(provider_registry[Api(api)][provider_type].config_class) - if hasattr(config_type, "sample_run_config"): + try: + config_type = instantiate_class_type(provider_registry[Api(api)][provider_type].config_class) + except ModuleNotFoundError: + # HACK ALERT: + # This code executes after building is done, the import cannot work since the + # package is either available in the venv or container - not available on the host. + # TODO: use a "is_external" flag in ProviderSpec to check if the provider is + # external + cprint( + f"Failed to import provider {provider_type} for API {api} - assuming it's external, skipping", + color="yellow", + ) + # Set config_type to None to avoid UnboundLocalError + config_type = None + + if config_type is not None and hasattr(config_type, "sample_run_config"): config = config_type.sample_run_config(__distro_dir__=f"~/.llama/distributions/{image_name}") else: config = {} @@ -323,6 +331,7 @@ def _run_stack_build_command_from_build_config( template_name: Optional[str] = None, config_path: Optional[str] = None, ) -> str: + image_name = image_name or build_config.image_name if build_config.image_type == LlamaStackImageType.CONTAINER.value: if template_name: image_name = f"distribution-{template_name}" diff --git a/llama_stack/distribution/build.py b/llama_stack/distribution/build.py index a8ee372da..5b61ae081 100644 --- a/llama_stack/distribution/build.py +++ b/llama_stack/distribution/build.py @@ -7,16 +7,16 @@ import importlib.resources import logging from pathlib import Path -from typing import Dict, List from pydantic import BaseModel from termcolor import cprint -from llama_stack.distribution.datatypes import BuildConfig, Provider +from llama_stack.distribution.datatypes import BuildConfig from llama_stack.distribution.distribution import get_provider_registry from llama_stack.distribution.utils.exec import run_command from llama_stack.distribution.utils.image_types import LlamaStackImageType from llama_stack.providers.datatypes import Api +from llama_stack.templates.template import DistributionTemplate log = logging.getLogger(__name__) @@ -37,19 +37,24 @@ class ApiInput(BaseModel): def get_provider_dependencies( - config_providers: Dict[str, List[Provider]], + config: BuildConfig | DistributionTemplate, ) -> tuple[list[str], list[str]]: """Get normal and special dependencies from provider configuration.""" - all_providers = get_provider_registry() + # Extract providers based on config type + if isinstance(config, DistributionTemplate): + providers = config.providers + elif isinstance(config, BuildConfig): + providers = config.distribution_spec.providers deps = [] + registry = get_provider_registry(config) - for api_str, provider_or_providers in config_providers.items(): - providers_for_api = all_providers[Api(api_str)] + for api_str, provider_or_providers in providers.items(): + providers_for_api = registry[Api(api_str)] providers = provider_or_providers if isinstance(provider_or_providers, list) else [provider_or_providers] for provider in providers: - # Providers from BuildConfig and RunConfig are subtly different – not great + # Providers from BuildConfig and RunConfig are subtly different – not great provider_type = provider if isinstance(provider, str) else provider.provider_type if provider_type not in providers_for_api: @@ -71,8 +76,8 @@ def get_provider_dependencies( return list(set(normal_deps)), list(set(special_deps)) -def print_pip_install_help(providers: Dict[str, List[Provider]]): - normal_deps, special_deps = get_provider_dependencies(providers) +def print_pip_install_help(config: BuildConfig): + normal_deps, special_deps = get_provider_dependencies(config) cprint( f"Please install needed dependencies using the following commands:\n\nuv pip install {' '.join(normal_deps)}", @@ -91,7 +96,7 @@ def build_image( ): container_base = build_config.distribution_spec.container_image or "python:3.10-slim" - normal_deps, special_deps = get_provider_dependencies(build_config.distribution_spec.providers) + normal_deps, special_deps = get_provider_dependencies(build_config) normal_deps += SERVER_DEPENDENCIES if build_config.image_type == LlamaStackImageType.CONTAINER.value: diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh index 97259ed0a..fb4780432 100755 --- a/llama_stack/distribution/build_container.sh +++ b/llama_stack/distribution/build_container.sh @@ -90,7 +90,7 @@ WORKDIR /app RUN apt-get update && apt-get install -y \ iputils-ping net-tools iproute2 dnsutils telnet \ - curl wget telnet \ + curl wget telnet git\ procps psmisc lsof \ traceroute \ bubblewrap \ diff --git a/llama_stack/distribution/datatypes.py b/llama_stack/distribution/datatypes.py index b24b0ec50..38353c1ff 100644 --- a/llama_stack/distribution/datatypes.py +++ b/llama_stack/distribution/datatypes.py @@ -326,3 +326,12 @@ class BuildConfig(BaseModel): default="conda", description="Type of package to build (conda | container | venv)", ) + image_name: Optional[str] = Field( + default=None, + description="Name of the distribution to build", + ) + external_providers_dir: Optional[str] = Field( + default=None, + description="Path to directory containing external provider implementations. The providers packages will be resolved from this directory. " + "pip_packages MUST contain the provider package name.", + ) diff --git a/llama_stack/distribution/distribution.py b/llama_stack/distribution/distribution.py index d4447139c..f948ddf1c 100644 --- a/llama_stack/distribution/distribution.py +++ b/llama_stack/distribution/distribution.py @@ -12,7 +12,6 @@ from typing import Any, Dict, List import yaml from pydantic import BaseModel -from llama_stack.distribution.datatypes import StackRunConfig from llama_stack.log import get_logger from llama_stack.providers.datatypes import ( AdapterSpec, @@ -97,7 +96,9 @@ def _load_inline_provider_spec(spec_data: Dict[str, Any], api: Api, provider_nam return spec -def get_provider_registry(config: StackRunConfig | None = None) -> Dict[Api, Dict[str, ProviderSpec]]: +def get_provider_registry( + config=None, +) -> Dict[Api, Dict[str, ProviderSpec]]: """Get the provider registry, optionally including external providers. This function loads both built-in providers and external providers from YAML files. @@ -122,7 +123,7 @@ def get_provider_registry(config: StackRunConfig | None = None) -> Dict[Api, Dic llama-guard.yaml Args: - config: Optional StackRunConfig containing the external providers directory path + config: Optional object containing the external providers directory path Returns: A dictionary mapping APIs to their available providers @@ -142,7 +143,8 @@ def get_provider_registry(config: StackRunConfig | None = None) -> Dict[Api, Dic except ImportError as e: logger.warning(f"Failed to import module {name}: {e}") - if config and config.external_providers_dir: + # Check if config has the external_providers_dir attribute + if config and hasattr(config, "external_providers_dir") and config.external_providers_dir: external_providers_dir = os.path.abspath(config.external_providers_dir) if not os.path.exists(external_providers_dir): raise FileNotFoundError(f"External providers directory not found: {external_providers_dir}") diff --git a/scripts/distro_codegen.py b/scripts/distro_codegen.py index 98faa53a3..a65e2c80d 100755 --- a/scripts/distro_codegen.py +++ b/scripts/distro_codegen.py @@ -98,7 +98,7 @@ def collect_template_dependencies(template_dir: Path) -> tuple[str | None, list[ if template_func := getattr(module, "get_distribution_template", None): template = template_func() - normal_deps, special_deps = get_provider_dependencies(template.providers) + normal_deps, special_deps = get_provider_dependencies(template) # Combine all dependencies in order: normal deps, special deps, server deps all_deps = sorted(set(normal_deps + SERVER_DEPENDENCIES)) + sorted(set(special_deps)) diff --git a/tests/external-provider/llama-stack-provider-ollama/custom-distro.yaml b/tests/external-provider/llama-stack-provider-ollama/custom-distro.yaml new file mode 100644 index 000000000..eb3b85e52 --- /dev/null +++ b/tests/external-provider/llama-stack-provider-ollama/custom-distro.yaml @@ -0,0 +1,9 @@ +version: '2' +distribution_spec: + description: Custom distro for CI tests + providers: + inference: + - remote::custom_ollama +image_type: container +image_name: ci-test +external_providers_dir: /tmp/providers.d diff --git a/tests/external-provider/llama-stack-provider-ollama/custom_ollama.yaml b/tests/external-provider/llama-stack-provider-ollama/custom_ollama.yaml index f0960b4d8..2ae1e2cf3 100644 --- a/tests/external-provider/llama-stack-provider-ollama/custom_ollama.yaml +++ b/tests/external-provider/llama-stack-provider-ollama/custom_ollama.yaml @@ -1,6 +1,6 @@ adapter: adapter_type: custom_ollama - pip_packages: ["ollama", "aiohttp"] + pip_packages: ["ollama", "aiohttp", "tests/external-provider/llama-stack-provider-ollama"] config_class: llama_stack_provider_ollama.config.OllamaImplConfig module: llama_stack_provider_ollama api_dependencies: [] diff --git a/tests/external-provider/llama-stack-provider-ollama/run.yaml b/tests/external-provider/llama-stack-provider-ollama/run.yaml index 7a3636c4d..a070a6dbb 100644 --- a/tests/external-provider/llama-stack-provider-ollama/run.yaml +++ b/tests/external-provider/llama-stack-provider-ollama/run.yaml @@ -1,14 +1,10 @@ version: '2' image_name: ollama apis: -- agents -- datasetio -- eval - inference -- safety -- scoring - telemetry - tool_runtime +- datasetio - vector_io providers: inference: @@ -24,19 +20,6 @@ providers: type: sqlite namespace: null db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/ollama}/faiss_store.db - safety: - - provider_id: llama-guard - provider_type: inline::llama-guard - config: - excluded_categories: [] - agents: - - provider_id: meta-reference - provider_type: inline::meta-reference - config: - persistence_store: - type: sqlite - namespace: null - db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/ollama}/agents_store.db telemetry: - provider_id: meta-reference provider_type: inline::meta-reference @@ -44,14 +27,6 @@ providers: service_name: ${env.OTEL_SERVICE_NAME:llama-stack} sinks: ${env.TELEMETRY_SINKS:console,sqlite} sqlite_db_path: ${env.SQLITE_DB_PATH:~/.llama/distributions/ollama/trace_store.db} - eval: - - provider_id: meta-reference - provider_type: inline::meta-reference - config: - kvstore: - type: sqlite - namespace: null - db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/ollama}/meta_reference_eval.db datasetio: - provider_id: huggingface provider_type: remote::huggingface @@ -67,17 +42,6 @@ providers: type: sqlite namespace: null db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/ollama}/localfs_datasetio.db - scoring: - - provider_id: basic - provider_type: inline::basic - config: {} - - provider_id: llm-as-judge - provider_type: inline::llm-as-judge - config: {} - - provider_id: braintrust - provider_type: inline::braintrust - config: - openai_api_key: ${env.OPENAI_API_KEY:} tool_runtime: - provider_id: brave-search provider_type: remote::brave-search