From 12858d6d617eaebba8477b403b56847e8a31e14d Mon Sep 17 00:00:00 2001
From: Vladislav <vladislav.bronzov@gmail.com>
Date: Fri, 13 Dec 2024 18:15:17 +0100
Subject: [PATCH] add build and dt types for platfrom

---
 llama_stack/cli/stack/build.py        | 22 +++++++++++++++++++++-
 llama_stack/distribution/datatypes.py |  5 +++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/llama_stack/cli/stack/build.py b/llama_stack/cli/stack/build.py
index 0cb873b57..d723d30db 100644
--- a/llama_stack/cli/stack/build.py
+++ b/llama_stack/cli/stack/build.py
@@ -77,6 +77,12 @@ class StackBuild(Subcommand):
             default="conda",
         )
 
+        self.parser.add_argument(
+            "--platform",
+            type=str,
+            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 +102,8 @@ class StackBuild(Subcommand):
             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 +138,15 @@ class StackBuild(Subcommand):
                 default="conda",
             )
 
+            platform = prompt(
+                "> Enter the target platform you want your Llama Stack to be built for: ",
+                validator=Validator.from_callable(
+                    lambda x: len(x) > 0,
+                    error_message="Platform cannot be empty, please enter a platform",
+                ),
+                default="linux/arm64",
+            )
+
             cprint(
                 textwrap.dedent(
                     """
@@ -172,7 +189,10 @@ class StackBuild(Subcommand):
             )
 
             build_config = BuildConfig(
-                name=name, image_type=image_type, distribution_spec=distribution_spec
+                name=name,
+                image_type=image_type,
+                distribution_spec=distribution_spec,
+                platform=platform,
             )
             self._run_stack_build_command_from_build_config(build_config)
             return
diff --git a/llama_stack/distribution/datatypes.py b/llama_stack/distribution/datatypes.py
index 1159372d4..e77aa14bf 100644
--- a/llama_stack/distribution/datatypes.py
+++ b/llama_stack/distribution/datatypes.py
@@ -25,6 +25,7 @@ from llama_stack.providers.utils.kvstore.config import KVStoreConfig
 
 LLAMA_STACK_BUILD_CONFIG_VERSION = "2"
 LLAMA_STACK_RUN_CONFIG_VERSION = "2"
+LLAMA_STACK_DEFAULT_PLATFORM = "linux/arm64"
 
 
 RoutingKey = Union[str, List[str]]
@@ -167,3 +168,7 @@ class BuildConfig(BaseModel):
         default="conda",
         description="Type of package to build (conda | docker | venv)",
     )
+    platform: str = Field(
+        default=LLAMA_STACK_DEFAULT_PLATFORM,
+        description="The platform for docker image, defaults to linux/arm64",
+    )