add build and dt types for platfrom

meta-llama · Dec 13, 2024 · 24f0f91 · 24f0f91
1 parent fe48b9f
commit 24f0f91
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/llama_stack/cli/stack/build.py b/llama_stack/cli/stack/build.py
@@ -77,6 +77,12 @@ def _add_arguments(self):
             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 @@ 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 +138,15 @@ 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: ",
+                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 @@ def _run_stack_build_command(self, args: argparse.Namespace) -> None:
             )
 
             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
@@ -25,6 +25,7 @@
 
 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 | container)",
     )
+    platform: str = Field(
+        default=LLAMA_STACK_DEFAULT_PLATFORM,
+        description="The platform for docker image, defaults to linux/arm64",
+    )