From 4747c844a2386157b0c79d9409a83873e36a1c61 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Tue, 14 Jan 2025 13:59:20 -0800
Subject: [PATCH] tmp dockerfile_only flag

---
 llama_stack/cli/stack/build.py              | 33 +++++++++++++++++----
 llama_stack/distribution/build.py           |  5 +++-
 llama_stack/distribution/build_container.sh |  3 ++
 3 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/llama_stack/cli/stack/build.py b/llama_stack/cli/stack/build.py
index 85e6cb962..5afc96be1 100644
--- a/llama_stack/cli/stack/build.py
+++ b/llama_stack/cli/stack/build.py
@@ -81,6 +81,14 @@ class StackBuild(Subcommand):
             default="conda",
         )
 
+        self.parser.add_argument(
+            "--dockerfile-only",
+            type=bool,
+            default=False,
+            action=argparse.BooleanOptionalAction,
+            help="Whether to generate a Dockerfile only (only valid when image-type is docker)",
+        )
+
     def _run_stack_build_command(self, args: argparse.Namespace) -> None:
         import textwrap
 
@@ -92,6 +100,12 @@ class StackBuild(Subcommand):
 
         from llama_stack.distribution.distribution import get_provider_registry
 
+        if args.dockerfile_only and args.image_type != "docker":
+            self.parser.error(
+                "dockerfile-only flag is only valid when image-type is docker"
+            )
+            return
+
         if args.list_templates:
             self._run_template_list_cmd(args)
             return
@@ -107,7 +121,9 @@ class StackBuild(Subcommand):
                             f"Please specify a image-type (docker | conda | venv) for {args.template}"
                         )
                     self._run_stack_build_command_from_build_config(
-                        build_config, template_name=args.template
+                        build_config,
+                        template_name=args.template,
+                        dockerfile_only=args.dockerfile_only,
                     )
                     return
 
@@ -178,7 +194,9 @@ class StackBuild(Subcommand):
             build_config = BuildConfig(
                 name=name, image_type=image_type, distribution_spec=distribution_spec
             )
-            self._run_stack_build_command_from_build_config(build_config)
+            self._run_stack_build_command_from_build_config(
+                build_config, dockerfile_only=args.dockerfile_only
+            )
             return
 
         with open(args.config, "r") as f:
@@ -187,7 +205,9 @@ class StackBuild(Subcommand):
             except Exception as e:
                 self.parser.error(f"Could not parse config file {args.config}: {e}")
                 return
-            self._run_stack_build_command_from_build_config(build_config)
+            self._run_stack_build_command_from_build_config(
+                build_config, dockerfile_only=args.dockerfile_only
+            )
 
     def _generate_run_config(self, build_config: BuildConfig, build_dir: Path) -> None:
         """
@@ -261,7 +281,10 @@ class StackBuild(Subcommand):
         )
 
     def _run_stack_build_command_from_build_config(
-        self, build_config: BuildConfig, template_name: Optional[str] = None
+        self,
+        build_config: BuildConfig,
+        template_name: Optional[str] = None,
+        dockerfile_only: bool = False,
     ) -> None:
         import json
         import os
@@ -281,7 +304,7 @@ class StackBuild(Subcommand):
             to_write = json.loads(build_config.model_dump_json())
             f.write(yaml.dump(to_write, sort_keys=False))
 
-        return_code = build_image(build_config, build_file_path)
+        return_code = build_image(build_config, build_file_path, dockerfile_only)
         if return_code != 0:
             return
 
diff --git a/llama_stack/distribution/build.py b/llama_stack/distribution/build.py
index 5a7dfba11..4f63be383 100644
--- a/llama_stack/distribution/build.py
+++ b/llama_stack/distribution/build.py
@@ -102,7 +102,9 @@ def print_pip_install_help(providers: Dict[str, List[Provider]]):
     print()
 
 
-def build_image(build_config: BuildConfig, build_file_path: Path):
+def build_image(
+    build_config: BuildConfig, build_file_path: Path, dockerfile_only: bool = False
+):
     docker_image = build_config.distribution_spec.docker_image or "python:3.10-slim"
 
     normal_deps, special_deps = get_provider_dependencies(
@@ -121,6 +123,7 @@ def build_image(build_config: BuildConfig, build_file_path: Path):
             str(build_file_path),
             str(BUILDS_BASE_DIR / ImageType.docker.value),
             " ".join(normal_deps),
+            dockerfile_only,
         ]
     elif build_config.image_type == ImageType.conda.value:
         script = (
diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh
index 06cb19c32..bde549af3 100755
--- a/llama_stack/distribution/build_container.sh
+++ b/llama_stack/distribution/build_container.sh
@@ -28,6 +28,7 @@ docker_base=$2
 build_file_path=$3
 host_build_dir=$4
 pip_dependencies=$5
+dockerfile_only=$7
 
 # Define color codes
 RED='\033[0;31m'
@@ -146,6 +147,8 @@ ENTRYPOINT ["python", "-m", "llama_stack.distribution.server.server", "--templat
 EOF
 
 printf "Dockerfile created successfully in $TEMP_DIR/Dockerfile\n\n"
+cp "$TEMP_DIR/Dockerfile" "./Dockerfile"
+printf "Dockerfile saved to ./Dockerfile\n\n"
 cat $TEMP_DIR/Dockerfile
 printf "\n"