llama-stack-mirror/llama_stack/distribution/utils/exec.py
ehhuang 0d7a90b8bc
chore: merge --config and --template in server.py (#2716)
# What does this PR do?
Part of #2696 

## Test Plan
Run `llama stack run starter`

Error:
```

myenv ❯ llama stack run starters
WARNING  2025-07-10 12:12:43,052 llama_stack.cli.stack.run:82 server: Conda detected. Using conda environment myenv for the run.
usage: llama stack run [-h] [--port PORT] [--image-name IMAGE_NAME] [--env KEY=VALUE]
                       [--image-type {conda,venv}] [--enable-ui]
                       [config | template]
llama stack run: error: Could not resolve config or template 'starters'.

Tried the following locations:
  1. As file path: /Users/erichuang/projects/llama-stack-git/starters
  2. As template: /Users/erichuang/projects/llama-stack-git/llama_stack/templates/starters/run.yaml
  3. As built distribution: (/Users/erichuang/.llama/distributions/llamastack-starters/starters-run.yaml, /Users/erichuang/.llama/distributions/starters/starters-run.yaml)

Available templates: dell, test-env, vllm-gpu, test-template, cerebras, openai-api-verification, sambanova, passthrough, direct-config, together, openai, fireworks, meta-reference-gpu, __pycache__, dev, ollama, watsonx, remote-vllm, llama_api, groq, dummy, oracle, nvidia, ci-tests, postgres-demo, test-stack, bedrock, starter, hf-serverless, hf-endpoint, tgi, open-benchmark, verification

Did you mean one of these templates?
  - starter
  - together
  - postgres-demo
```
2025-07-21 13:19:27 -07:00

143 lines
4.4 KiB
Python

# 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.
import logging
import os
import signal
import subprocess
import sys
from termcolor import cprint
log = logging.getLogger(__name__)
import importlib
import json
from pathlib import Path
from llama_stack.distribution.utils.image_types import LlamaStackImageType
def formulate_run_args(image_type: str, image_name: str) -> list[str]:
env_name = ""
if image_type == LlamaStackImageType.CONDA.value:
current_conda_env = os.environ.get("CONDA_DEFAULT_ENV")
env_name = image_name or current_conda_env
if not env_name:
cprint(
"No current conda environment detected, please specify a conda environment name with --image-name",
color="red",
file=sys.stderr,
)
return
def get_conda_prefix(env_name):
# Conda "base" environment does not end with "base" in the
# prefix, so should be handled separately.
if env_name == "base":
return os.environ.get("CONDA_PREFIX")
# Get conda environments info
conda_env_info = json.loads(subprocess.check_output(["conda", "info", "--envs", "--json"]).decode())
envs = conda_env_info["envs"]
for envpath in envs:
if os.path.basename(envpath) == env_name:
return envpath
return None
cprint(f"Using conda environment: {env_name}", color="green", file=sys.stderr)
conda_prefix = get_conda_prefix(env_name)
if not conda_prefix:
cprint(
f"Conda environment {env_name} does not exist.",
color="red",
file=sys.stderr,
)
return
build_file = Path(conda_prefix) / "llamastack-build.yaml"
if not build_file.exists():
cprint(
f"Build file {build_file} does not exist.\n\nPlease run `llama stack build` or specify the correct conda environment name with --image-name",
color="red",
file=sys.stderr,
)
return
else:
# else must be venv since that is the only valid option left.
current_venv = os.environ.get("VIRTUAL_ENV")
env_name = image_name or current_venv
if not env_name:
cprint(
"No current virtual environment detected, please specify a virtual environment name with --image-name",
color="red",
file=sys.stderr,
)
return
cprint(f"Using virtual environment: {env_name}", file=sys.stderr)
script = importlib.resources.files("llama_stack") / "distribution/start_stack.sh"
run_args = [
script,
image_type,
env_name,
]
return run_args
def in_notebook():
try:
from IPython import get_ipython
if "IPKernelApp" not in get_ipython().config: # pragma: no cover
return False
except ImportError:
return False
except AttributeError:
return False
return True
def run_command(command: list[str]) -> int:
"""
Run a command with interrupt handling and output capture.
Uses subprocess.run with direct stream piping for better performance.
Args:
command (list): The command to run.
Returns:
int: The return code of the command.
"""
original_sigint = signal.getsignal(signal.SIGINT)
ctrl_c_pressed = False
def sigint_handler(signum, frame):
nonlocal ctrl_c_pressed
ctrl_c_pressed = True
log.info("\nCtrl-C detected. Aborting...")
try:
# Set up the signal handler
signal.signal(signal.SIGINT, sigint_handler)
# Run the command with stdout/stderr piped directly to system streams
result = subprocess.run(
command,
text=True,
check=False,
)
return result.returncode
except subprocess.SubprocessError as e:
log.error(f"Subprocess error: {e}")
return 1
except Exception as e:
log.exception(f"Unexpected error: {e}")
return 1
finally:
# Restore the original signal handler
signal.signal(signal.SIGINT, original_sigint)