From 86d7c378fed433caf77a8782603892d8160f9ecb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Han?= <seb@redhat.com>
Date: Tue, 25 Feb 2025 15:08:58 +0100
Subject: [PATCH] feat(server): Use system packages for execution
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Users prefer to rely on the main CLI rather than invoking the server
through a Python module. Users interact with a high-level CLI rather
than needing to know internal module structures.

Now, when running llama stack run <path-to-config>, the server will
attempt to use the system package or a virtual environment if one is
active.

This also eliminates the current process dependency chain when running
from a virtual environment:

-> llama stack run
  -> start_env.sh
    -> python -m server...

Signed-off-by: Sébastien Han <seb@redhat.com>
---
 llama_stack/cli/stack/run.py | 60 ++++++++++++++++++++++++++----------
 1 file changed, 43 insertions(+), 17 deletions(-)

diff --git a/llama_stack/cli/stack/run.py b/llama_stack/cli/stack/run.py
index 79f604d1aa..7bb5f50fe1 100644
--- a/llama_stack/cli/stack/run.py
+++ b/llama_stack/cli/stack/run.py
@@ -7,6 +7,7 @@
 import argparse
 import logging
 import os
+import sys
 from pathlib import Path
 
 from llama_stack.cli.subcommand import Subcommand
@@ -130,20 +131,45 @@ def _run_stack_run_cmd(self, args: argparse.Namespace) -> None:
         except AttributeError as e:
             self.parser.error(f"failed to parse config file '{config_file}':\n {e}")
 
-        run_args = formulate_run_args(args.image_type, args.image_name, config, template_name)
-
-        run_args.extend([str(config_file), str(args.port)])
-        if args.disable_ipv6:
-            run_args.append("--disable-ipv6")
-
-        for env_var in args.env:
-            if "=" not in env_var:
-                self.parser.error(f"Environment variable '{env_var}' must be in KEY=VALUE format")
-            key, value = env_var.split("=", 1)  # split on first = only
-            if not key:
-                self.parser.error(f"Environment variable '{env_var}' has empty key")
-            run_args.extend(["--env", f"{key}={value}"])
-
-        if args.tls_keyfile and args.tls_certfile:
-            run_args.extend(["--tls-keyfile", args.tls_keyfile, "--tls-certfile", args.tls_certfile])
-        run_with_pty(run_args)
+        # If the image type and name are not provided, we assume the user wants to run using system
+        # packages and directly run the server code.
+        if not args.image_type and not args.image_name:
+            logger.info("No image type or image name provided. Assuming system packages.")
+            # Import and run the server directly
+            try:
+                from llama_stack.distribution.server.server import main as server_main
+
+                # Remove the first argument which is the command name, the second argument which is
+                # the config file, the third argument which is the pass to the config file.
+                sys.argv = sys.argv[3:]
+
+                # Build final sys.argv to pass to the server and add the config file
+                sys.argv = ["server"] + sys.argv[1:]
+                sys.argv.extend(["--yaml-config", str(config_file)])
+
+                # Run the server
+                server_main()
+            except Exception as e:
+                self.parser.error(f"Failed to run server: {e}")
+            except ImportError as e:
+                self.parser.error(f"Failed to import server module: {e}")
+        else:
+            run_args = formulate_run_args(args.image_type, args.image_name, config, template_name)
+
+            run_args.extend([str(config_file), str(args.port)])
+            if args.disable_ipv6:
+                run_args.append("--disable-ipv6")
+
+            for env_var in args.env:
+                if "=" not in env_var:
+                    self.parser.error(f"Environment variable '{env_var}' must be in KEY=VALUE format")
+                    return
+                key, value = env_var.split("=", 1)  # split on first = only
+                if not key:
+                    self.parser.error(f"Environment variable '{env_var}' has empty key")
+                    return
+                run_args.extend(["--env", f"{key}={value}"])
+
+            if args.tls_keyfile and args.tls_certfile:
+                run_args.extend(["--tls-keyfile", args.tls_keyfile, "--tls-certfile", args.tls_certfile])
+            run_with_pty(run_args)