Add platform for docker

docs/source/distributions/building_distro.md

@@ -20,6 +20,7 @@ llama stack build -h
 We will start build our distribution (in the form of a Conda environment, or Docker image). In this step, we will specify:
 - `name`: the name for our distribution (e.g. `my-stack`)
 - `image_type`: our build image type (`conda | docker`)
+- `platform`: the platform for docker image (e.g. `linux/arm64`)
 - `distribution_spec`: our distribution specs for specifying API providers
   - `description`: a short description of the configurations for the distribution
   - `providers`: specifies the underlying implementation for serving each API endpoint
@@ -36,6 +37,7 @@ llama stack build
 
 > Enter a name for your Llama Stack (e.g. my-local-stack): my-stack
 > Enter the image type you want your Llama Stack to be built as (docker or conda): conda
+> Enter the target platform you want your Llama Stack to be built for:
 
 Llama Stack is composed of several APIs working together. Let's select
 the provider types (implementations) you want to use for these APIs.

llama_stack/cli/stack/build.py

@@ -77,6 +77,13 @@ def _add_arguments(self):
             default="conda",
         )
 
+        self.parser.add_argument(
+            "--platform",
+            type=str,
+            default=None,
+            help="Platform to use for the build. Required when using docker as image type, defaults to host if no platform is specified",
+        )
+
     def _run_stack_build_command(self, args: argparse.Namespace) -> None:
         import textwrap
 
@@ -96,6 +103,8 @@ def _run_stack_build_command(self, args: argparse.Namespace) -> None:
             available_templates = available_templates_specs()
             for build_config in available_templates:
                 if build_config.name == args.template:
+                    if args.platform:
+                        build_config.platform = args.platform
                     if args.image_type:
                         build_config.image_type = args.image_type
                     else:
@@ -130,6 +139,10 @@ def _run_stack_build_command(self, args: argparse.Namespace) -> None:
                 default="conda",
             )
 
+            platform = prompt(
+                "> Enter the target platform you want your Llama Stack to be built for: "
+            )
+
             cprint(
                 textwrap.dedent(
                     """
@@ -174,6 +187,8 @@ def _run_stack_build_command(self, args: argparse.Namespace) -> None:
             build_config = BuildConfig(
                 name=name, image_type=image_type, distribution_spec=distribution_spec
             )
+            if platform.strip():
+                build_config.platform = platform
             self._run_stack_build_command_from_build_config(build_config)
             return
 

llama_stack/distribution/build.py

@@ -121,6 +121,8 @@ def build_image(build_config: BuildConfig, build_file_path: Path):
             str(BUILDS_BASE_DIR / ImageType.docker.value),
             " ".join(normal_deps),
         ]
+        if build_config.platform is not None:
+            args.append(build_config.platform)
     elif build_config.image_type == ImageType.conda.value:
         script = pkg_resources.resource_filename(
             "llama_stack", "distribution/build_conda_env.sh"

llama_stack/distribution/build_container.sh

@@ -11,9 +11,9 @@ LLAMA_STACK_DIR=${LLAMA_STACK_DIR:-}
 TEST_PYPI_VERSION=${TEST_PYPI_VERSION:-}
 BUILD_PLATFORM=${BUILD_PLATFORM:-}
 
-if [ "$#" -lt 4 ]; then
-  echo "Usage: $0 <build_name> <docker_base> <pip_dependencies> [<special_pip_deps>]" >&2
-  echo "Example: $0 my-fastapi-app python:3.9-slim 'fastapi uvicorn' " >&2
+if [ "$#" -lt 5 ]; then
+  echo "Usage: $0 <build_name> <docker_base> <platform> <pip_dependencies> [<special_pip_deps>]" >&2
+  echo "Example: $0 my-fastapi-app python:3.9-slim linux/arm64 'fastapi uvicorn' " >&2
   exit 1
 fi
 
@@ -24,9 +24,10 @@ set -euo pipefail
 build_name="$1"
 image_name="distribution-$build_name"
 docker_base=$2
-build_file_path=$3
-host_build_dir=$4
-pip_dependencies=$5
+platform=$3
+build_file_path=$4
+host_build_dir=$5
+pip_dependencies=$6
 
 # Define color codes
 RED='\033[0;31m'

llama_stack/distribution/datatypes.py

@@ -167,3 +167,7 @@ class BuildConfig(BaseModel):
         default="conda",
         description="Type of package to build (conda | docker | venv)",
     )
+    platform: Optional[str] = Field(
+        default=None,
+        description="The platform for docker image",
+    )