From ebb57a0c67f46845a0105444c275b414cd164454 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Fri, 27 Sep 2024 13:52:23 -0700
Subject: [PATCH] wip

---
 llama_stack/cli/stack/configure.py            | 32 +++++++---
 llama_stack/distribution/build_container.sh   |  2 +-
 .../distribution/configure_container.sh       | 34 ++++++++++
 llama_stack/distribution/server/server.py     |  5 +-
 llama_stack/distribution/start_container.sh   | 62 +++++++++++++++++++
 5 files changed, 123 insertions(+), 12 deletions(-)
 create mode 100644 llama_stack/distribution/configure_container.sh
 create mode 100644 llama_stack/distribution/start_container.sh

diff --git a/llama_stack/cli/stack/configure.py b/llama_stack/cli/stack/configure.py
index 7e2dc5994..b96e1e6a5 100644
--- a/llama_stack/cli/stack/configure.py
+++ b/llama_stack/cli/stack/configure.py
@@ -78,15 +78,31 @@ class StackConfigure(Subcommand):
                 self._configure_llama_distribution(build_config, args.output_dir)
                 return
         
-        # if we get here, we need to prompt user to try configure inside docker image
-        self.parser.error(
-            f"""
-            Could not find {build_config_file}. Did you download a docker image?
-            Try running `docker run -it --entrypoint "/bin/bash" <image_name>`
-            `llama stack configure llamastack-build.yaml --output-dir  ./`
-            to set a new run configuration file. 
-            """,
+        # if we get here, we need to try to find the docker image
+        cprint(
+            f"Could not find {build_config_file}. Trying docker image name instead...",
+            color="green",
         )
+        docker_image = args.config
+        builds_dir = BUILDS_BASE_DIR / ImageType.docker.value
+        if args.output_dir:
+            builds_dir = Path(output_dir)
+        os.makedirs(builds_dir, exist_ok=True)
+
+        script = pkg_resources.resource_filename(
+            "llama_stack", "distribution/configure_container.sh"
+        )
+        script_args = [script, docker_image, str(builds_dir)]
+
+        return_code = run_with_pty(script_args)
+
+        # we have regenerated the build config file with script, now check if it exists
+        if return_code != 0:
+            self.parser.error(
+                f"Failed to configure container {docker_image} with return code {return_code}. Please run `llama stack build first`. "
+            )
+            return
+
         return
 
     def _configure_llama_distribution(
diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh
index 6a39ecb70..a4797dadc 100755
--- a/llama_stack/distribution/build_container.sh
+++ b/llama_stack/distribution/build_container.sh
@@ -102,7 +102,7 @@ add_to_docker <<EOF
 # This would be good in production but for debugging flexibility lets not add it right now
 # We need a more solid production ready entrypoint.sh anyway
 #
-ENTRYPOINT ["python", "-m", "llama_stack.distribution.server.server", "./llamastack-run.yaml"]
+CMD ["python", "-m", "llama_stack.distribution.server.server", "./llamastack-run.yaml"]
 
 EOF
 
diff --git a/llama_stack/distribution/configure_container.sh b/llama_stack/distribution/configure_container.sh
new file mode 100644
index 000000000..56e45db7e
--- /dev/null
+++ b/llama_stack/distribution/configure_container.sh
@@ -0,0 +1,34 @@
+#!/bin/bash
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+DOCKER_BINARY=${DOCKER_BINARY:-docker}
+DOCKER_OPTS=${DOCKER_OPTS:-}
+
+set -euo pipefail
+
+error_handler() {
+  echo "Error occurred in script at line: ${1}" >&2
+  exit 1
+}
+
+trap 'error_handler ${LINENO}' ERR
+
+if [ $# -lt 2 ]; then
+  echo "Usage: $0 <container name> <build file path>"
+  exit 1
+fi
+
+docker_image="$1"
+host_build_dir="$2"
+container_build_dir="/app/builds"
+
+set -x
+$DOCKER_BINARY run $DOCKER_OPTS -it \
+  -v $host_build_dir:$container_build_dir \
+  $docker_image \
+  llama stack configure ./llamastack-build.yaml --output-dir $container_build_dir
diff --git a/llama_stack/distribution/server/server.py b/llama_stack/distribution/server/server.py
index d665226cf..7a3e6276c 100644
--- a/llama_stack/distribution/server/server.py
+++ b/llama_stack/distribution/server/server.py
@@ -410,9 +410,8 @@ async def resolve_impls_with_routing(run_config: StackRunConfig) -> Dict[Api, An
     return impls, specs
 
 
-def main(default_yaml: str, port: int = 5000, disable_ipv6: bool = False, config: Optional[str] = None):
-    config_file = config or default_yaml
-    with open(config_file, "r") as fp:
+def main(yaml_config: str, port: int = 5000, disable_ipv6: bool = False):
+    with open(yaml_config, "r") as fp:
         config = StackRunConfig(**yaml.safe_load(fp))
 
     app = FastAPI()
diff --git a/llama_stack/distribution/start_container.sh b/llama_stack/distribution/start_container.sh
new file mode 100644
index 000000000..ee581cac4
--- /dev/null
+++ b/llama_stack/distribution/start_container.sh
@@ -0,0 +1,62 @@
+#!/bin/bash
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+DOCKER_BINARY=${DOCKER_BINARY:-docker}
+DOCKER_OPTS=${DOCKER_OPTS:-}
+LLAMA_CHECKPOINT_DIR=${LLAMA_CHECKPOINT_DIR:-}
+
+set -euo pipefail
+
+RED='\033[0;31m'
+NC='\033[0m' # No Color
+
+error_handler() {
+  echo "Error occurred in script at line: ${1}" >&2
+  exit 1
+}
+
+trap 'error_handler ${LINENO}' ERR
+
+if [ $# -lt 3 ]; then
+  echo "Usage: $0 <build_name> <yaml_config> <port> <other_args...>"
+  exit 1
+fi
+
+build_name="$1"
+docker_image="llamastack-$build_name"
+shift
+
+yaml_config="$1"
+shift
+
+port="$1"
+shift
+
+set -x
+
+if [ -n "$LLAMA_CHECKPOINT_DIR" ]; then
+  $DOCKER_BINARY run $DOCKER_OPTS -it \
+    -p $port:$port \
+    -v "$yaml_config:/app/config.yaml" \
+    -v "$LLAMA_CHECKPOINT_DIR:/root/.llama" \
+    --gpus=all \
+    $docker_image \
+    python -m llama_stack.distribution.server.server \
+    --yaml_config /app/config.yaml \
+    --port $port "$@"
+fi
+
+if [ -z "$LLAMA_CHECKPOINT_DIR" ]; then
+  $DOCKER_BINARY run $DOCKER_OPTS -it \
+    -p $port:$port \
+    -v "$yaml_config:/app/config.yaml" \
+    $docker_image \
+    python -m llama_stack.distribution.server.server \
+    --yaml_config /app/config.yaml \
+    --port $port "$@"
+fi