phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-10-08 04:54:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Rename Distribution -> DistributionSpec, simplify RemoteProviders

Browse source
This commit is contained in:
Ashwin Bharambe 2024-08-06 10:45:06 -07:00
parent 0a67f3d3e6
commit 7cc0445517
9 changed files with 181 additions and 147 deletions
Download patch file Download diff file Expand all files Collapse all files

80
llama_toolchain/cli/distribution/configure.py
View file

@ -8,7 +8,7 @@ import argparse
import json
import shlex
from pathlib import Path
from typing import Any, Dict
import yaml
from termcolor import cprint
@ -32,83 +32,81 @@ class DistributionConfigure(Subcommand):
self.parser.set_defaults(func=self._run_distribution_configure_cmd)
def _add_arguments(self):
from llama_toolchain.distribution.registry import available_distributions
self.parser.add_argument(
"--name",
type=str,
help="Name of the distribution to configure",
default="local-source",
choices=[d.name for d in available_distributions()],
required=True,
)
def _run_distribution_configure_cmd(self, args: argparse.Namespace) -> None:
from llama_toolchain.distribution.registry import resolve_distribution
from llama_toolchain.distribution.registry import resolve_distribution_spec
dist = resolve_distribution(args.name)
if dist is None:
self.parser.error(f"Could not find distribution {args.name}")
config_file = DISTRIBS_BASE_DIR / args.name / "config.yaml"
if not config_file.exists():
self.parser.error(
f"Could not find {config_file}. Please run `llama distribution install` first"
)
return
env_file = DISTRIBS_BASE_DIR / dist.name / "conda.env"
# read this file to get the conda env name
assert env_file.exists(), f"Could not find conda env file {env_file}"
with open(env_file, "r") as f:
conda_env = f.read().strip()
# we need to find the spec from the name
with open(config_file, "r") as f:
config = yaml.safe_load(f)
configure_llama_distribution(dist, conda_env)
dist = resolve_distribution_spec(config["spec"])
if dist is None:
raise ValueError(f"Could not find any registered spec `{config['spec']}`")
configure_llama_distribution(dist, config)
def configure_llama_distribution(dist: "Distribution", conda_env: str):
def configure_llama_distribution(dist: "Distribution", config: Dict[str, Any]):
from llama_toolchain.common.exec import run_command
from llama_toolchain.common.prompt_for_config import prompt_for_config
from llama_toolchain.common.serialize import EnumEncoder
from llama_toolchain.distribution.datatypes import RemoteProviderSpec
from llama_toolchain.distribution.dynamic import instantiate_class_type
python_exe = run_command(shlex.split("which python"))
# simple check
conda_env = config["conda_env"]
if conda_env not in python_exe:
raise ValueError(
f"Please re-run configure by activating the `{conda_env}` conda environment"
)
existing_config = None
config_path = Path(DISTRIBS_BASE_DIR) / dist.name / "config.yaml"
if config_path.exists():
existing_config = config
if "providers" in existing_config:
cprint(
f"Configuration already exists for {dist.name}. Will overwrite...",
f"Configuration already exists for {config['name']}. Will overwrite...",
"yellow",
attrs=["bold"],
)
with open(config_path, "r") as fp:
existing_config = yaml.safe_load(fp)
provider_configs = {}
for api, provider_spec in dist.provider_specs.items():
if isinstance(provider_spec, RemoteProviderSpec):
provider_configs[api.value] = provider_spec.dict()
else:
cprint(f"Configuring API surface: {api.value}", "white", attrs=["bold"])
config_type = instantiate_class_type(provider_spec.config_class)
config = prompt_for_config(
config_type,
(
config_type(**existing_config["providers"][api.value])
if existing_config and api.value in existing_config["providers"]
else None
),
)
provider_configs[api.value] = {
"provider_id": provider_spec.provider_id,
**config.dict(),
}
cprint(f"Configuring API surface: {api.value}", "white", attrs=["bold"])
config_type = instantiate_class_type(provider_spec.config_class)
config = prompt_for_config(
config_type,
(
config_type(**existing_config["providers"][api.value])
if existing_config
and "providers" in existing_config
and api.value in existing_config["providers"]
else None
),
)
provider_configs[api.value] = {
"provider_id": provider_spec.provider_id,
**config.dict(),
}
dist_config = {
"providers": provider_configs,
"conda_env": conda_env,
**existing_config,
}
config_path = DISTRIBS_BASE_DIR / existing_config["name"] / "config.yaml"
with open(config_path, "w") as fp:
dist_config = json.loads(json.dumps(dist_config, cls=EnumEncoder))
fp.write(yaml.dump(dist_config, sort_keys=False))

4
llama_toolchain/cli/distribution/create.py
View file

@ -34,9 +34,9 @@ class DistributionCreate(Subcommand):
# wants to pick and then ask for their configuration.
def _run_distribution_create_cmd(self, args: argparse.Namespace) -> None:
from llama_toolchain.distribution.registry import resolve_distribution
from llama_toolchain.distribution.registry import resolve_distribution_spec
dist = resolve_distribution(args.name)
dist = resolve_distribution_spec(args.name)
if dist is not None:
self.parser.error(f"Distribution with name {args.name} already exists")
return

52
llama_toolchain/cli/distribution/install.py
View file

@ -10,6 +10,7 @@ import shlex
import textwrap
import pkg_resources
import yaml
from termcolor import cprint
@ -32,26 +33,31 @@ class DistributionInstall(Subcommand):
self.parser.set_defaults(func=self._run_distribution_install_cmd)
def _add_arguments(self):
from llama_toolchain.distribution.registry import available_distributions
from llama_toolchain.distribution.registry import available_distribution_specs
self.parser.add_argument(
"--spec",
type=str,
help="Distribution spec to install (try ollama-inline)",
required=True,
choices=[d.spec_id for d in available_distribution_specs()],
)
self.parser.add_argument(
"--name",
type=str,
help="Name of the distribution to install -- (try local-ollama)",
help="What should the installation be called locally?",
required=True,
choices=[d.name for d in available_distributions()],
)
self.parser.add_argument(
"--conda-env",
type=str,
help="Specify the name of the conda environment you would like to create or update",
required=True,
help="conda env in which this distribution will run (default = distribution name)",
)
def _run_distribution_install_cmd(self, args: argparse.Namespace) -> None:
from llama_toolchain.common.exec import run_with_pty
from llama_toolchain.distribution.distribution import distribution_dependencies
from llama_toolchain.distribution.registry import resolve_distribution
from llama_toolchain.distribution.registry import resolve_distribution_spec
os.makedirs(DISTRIBS_BASE_DIR, exist_ok=True)
script = pkg_resources.resource_filename(
@ -59,25 +65,36 @@ class DistributionInstall(Subcommand):
"distribution/install_distribution.sh",
)
dist = resolve_distribution(args.name)
dist = resolve_distribution_spec(args.spec)
if dist is None:
self.parser.error(f"Could not find distribution {args.name}")
self.parser.error(f"Could not find distribution {args.spec}")
return
os.makedirs(DISTRIBS_BASE_DIR / dist.name, exist_ok=True)
distrib_dir = DISTRIBS_BASE_DIR / args.name
os.makedirs(distrib_dir, exist_ok=True)
deps = distribution_dependencies(dist)
return_code = run_with_pty([script, args.conda_env, " ".join(deps)])
if not args.conda_env:
print(f"Using {args.name} as the Conda environment for this distribution")
conda_env = args.conda_env or args.name
return_code = run_with_pty([script, conda_env, " ".join(deps)])
assert return_code == 0, cprint(
f"Failed to install distribution {dist.name}", color="red"
f"Failed to install distribution {dist.spec_id}", color="red"
)
with open(DISTRIBS_BASE_DIR / dist.name / "conda.env", "w") as f:
f.write(f"{args.conda_env}\n")
config_file = distrib_dir / "config.yaml"
with open(config_file, "w") as f:
c = {
"conda_env": conda_env,
"spec": dist.spec_id,
"name": args.name,
}
f.write(yaml.dump(c))
cprint(
f"Distribution `{dist.name}` has been installed successfully!",
f"Distribution `{args.name}` (with spec {dist.spec_id}) has been installed successfully!",
color="green",
)
print(
@ -85,8 +102,7 @@ class DistributionInstall(Subcommand):
f"""
Update your conda environment and configure this distribution by running:
conda deactivate && conda activate {args.conda_env}
llama distribution configure --name {dist.name}
conda deactivate && conda activate {conda_env}
llama distribution configure --name {args.name}
"""
)
)
))

12
llama_toolchain/cli/distribution/list.py
View file

@ -28,23 +28,23 @@ class DistributionList(Subcommand):
def _run_distribution_list_cmd(self, args: argparse.Namespace) -> None:
from llama_toolchain.cli.table import print_table
from llama_toolchain.distribution.registry import available_distributions
from llama_toolchain.distribution.registry import available_distribution_specs
# eventually, this should query a registry at llama.meta.com/llamastack/distributions
headers = [
"Name",
"Spec ID",
"ProviderSpecs",
"Description",
]
rows = []
for dist in available_distributions():
providers = {k.value: v.provider_id for k, v in dist.provider_specs.items()}
for spec in available_distribution_specs():
providers = {k.value: v.provider_id for k, v in spec.provider_specs.items()}
rows.append(
[
dist.name,
spec.spec_id,
json.dumps(providers, indent=2),
dist.description,
spec.description,
]
)
print_table(

27
llama_toolchain/cli/distribution/start.py
View file

@ -6,7 +6,6 @@
import argparse
import shlex
from pathlib import Path
import yaml
@ -49,22 +48,23 @@ class DistributionStart(Subcommand):
def _run_distribution_start_cmd(self, args: argparse.Namespace) -> None:
from llama_toolchain.common.exec import run_command
from llama_toolchain.distribution.registry import resolve_distribution
from llama_toolchain.distribution.registry import resolve_distribution_spec
from llama_toolchain.distribution.server import main as distribution_server_init
dist = resolve_distribution(args.name)
if dist is None:
self.parser.error(f"Distribution with name {args.name} not found")
config_file = DISTRIBS_BASE_DIR / args.name / "config.yaml"
if not config_file.exists():
self.parser.error(
f"Could not find {config_file}. Please run `llama distribution install` first"
)
return
config_yaml = Path(DISTRIBS_BASE_DIR) / dist.name / "config.yaml"
if not config_yaml.exists():
raise ValueError(
f"Configuration {config_yaml} does not exist. Please run `llama distribution install` or `llama distribution configure` first"
)
# we need to find the spec from the name
with open(config_file, "r") as f:
config = yaml.safe_load(f)
with open(config_yaml, "r") as fp:
config = yaml.safe_load(fp)
dist = resolve_distribution_spec(config["spec"])
if dist is None:
raise ValueError(f"Could not find any registered spec `{config['spec']}`")
conda_env = config["conda_env"]
@ -76,8 +76,7 @@ class DistributionStart(Subcommand):
)
distribution_server_init(
dist.name,
config_yaml,
config_file,
args.port,
disable_ipv6=args.disable_ipv6,
)