From f1b9578f8d80d395ecc955f77cefdcf19a2542e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vladimir=20Ivi=C4=87?= <vladimir.ivic@me.com>
Date: Fri, 15 Nov 2024 23:16:42 -0800
Subject: [PATCH] Extend shorthand support for the `llama stack run` command
 (#465)

**Summary:**
Extend the shorthand run command so it can run successfully when config
exists under DISTRIBS_BASE_DIR (i.e. ~/.llama/distributions).

For example, imagine you created a new stack using the `llama stack
build` command where you named it "my-awesome-llama-stack".

```
$ llama stack build

> Enter a name for your Llama Stack (e.g. my-local-stack): my-awesome-llama-stack
```

To run the stack you created you will have to use long config path:
```
llama stack run ~/.llama/distributions/llamastack-my-awesome-llama-stack/my-awesome-llama-stack-run.yaml
```

With this change, you can start it using the stack name instead of full
path:
```
llama stack run my-awesome-llama-stack
```

**Test Plan:**
Verify command fails when stack doesn't exist
```
python3 -m llama_stack.cli.llama stack run my-test-stack
```

Output [FAILURE]
```
usage: llama stack run [-h] [--port PORT] [--disable-ipv6] config
llama stack run: error: File /Users/vladimirivic/.llama/distributions/llamastack-my-test-stack/my-test-stack-run.yaml does not exist. Please run `llama stack build` to generate (and optionally edit) a run.yaml file
```

Create a new stack using `llama stack build`.
Name it `my-test-stack`.

Verify command runs successfully
```
python3 -m llama_stack.cli.llama stack run my-test-stack
```

Output [SUCCESS]
```
Listening on ['::', '0.0.0.0']:5000
INFO:     Started server process [80146]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://['::', '0.0.0.0']:5000 (Press CTRL+C to quit)
```
---
 llama_stack/cli/stack/run.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/llama_stack/cli/stack/run.py b/llama_stack/cli/stack/run.py
index 842703d4c..5fce8c92c 100644
--- a/llama_stack/cli/stack/run.py
+++ b/llama_stack/cli/stack/run.py
@@ -48,7 +48,10 @@ class StackRun(Subcommand):
 
         from llama_stack.distribution.build import ImageType
         from llama_stack.distribution.configure import parse_and_maybe_upgrade_config
-        from llama_stack.distribution.utils.config_dirs import BUILDS_BASE_DIR
+        from llama_stack.distribution.utils.config_dirs import (
+            BUILDS_BASE_DIR,
+            DISTRIBS_BASE_DIR,
+        )
         from llama_stack.distribution.utils.exec import run_with_pty
 
         if not args.config:
@@ -68,6 +71,14 @@ class StackRun(Subcommand):
                 BUILDS_BASE_DIR / ImageType.docker.value / f"{args.config}-run.yaml"
             )
 
+        if not config_file.exists() and not args.config.endswith(".yaml"):
+            # check if it's a build config saved to ~/.llama dir
+            config_file = Path(
+                DISTRIBS_BASE_DIR
+                / f"llamastack-{args.config}"
+                / f"{args.config}-run.yaml"
+            )
+
         if not config_file.exists():
             self.parser.error(
                 f"File {str(config_file)} does not exist. Please run `llama stack build` to generate (and optionally edit) a run.yaml file"