local imports for faster cli

This commit is contained in:
Hardik Shah 2024-08-02 16:34:29 -07:00
parent af4710c959
commit 67229f23a4
9 changed files with 44 additions and 47 deletions

View file

@ -18,15 +18,8 @@ from pydantic import BaseModel
from termcolor import cprint
from llama_toolchain.cli.subcommand import Subcommand
from llama_toolchain.distribution.datatypes import Distribution, PassthroughApiAdapter
from llama_toolchain.distribution.registry import (
available_distributions,
resolve_distribution,
)
from llama_toolchain.utils import DISTRIBS_BASE_DIR, EnumEncoder
from .utils import run_command
class DistributionConfigure(Subcommand):
"""Llama cli for configuring llama toolchain configs"""
@ -43,6 +36,7 @@ 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,
@ -52,6 +46,8 @@ class DistributionConfigure(Subcommand):
)
def _run_distribution_configure_cmd(self, args: argparse.Namespace) -> None:
from llama_toolchain.distribution.registry import resolve_distribution
dist = resolve_distribution(args.name)
if dist is None:
self.parser.error(f"Could not find distribution {args.name}")
@ -66,7 +62,10 @@ class DistributionConfigure(Subcommand):
configure_llama_distribution(dist, conda_env)
def configure_llama_distribution(dist: Distribution, conda_env: str):
def configure_llama_distribution(dist: "Distribution", conda_env: str):
from llama_toolchain.distribution.datatypes import PassthroughApiAdapter
from .utils import run_command
python_exe = run_command(shlex.split("which python"))
# simple check
if conda_env not in python_exe:

View file

@ -7,7 +7,6 @@
import argparse
from llama_toolchain.cli.subcommand import Subcommand
from llama_toolchain.distribution.registry import resolve_distribution
class DistributionCreate(Subcommand):
@ -35,6 +34,8 @@ 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
dist = resolve_distribution(args.name)
if dist is not None:
self.parser.error(f"Distribution with name {args.name} already exists")

View file

@ -11,17 +11,8 @@ import shlex
import pkg_resources
from llama_toolchain.cli.subcommand import Subcommand
from llama_toolchain.distribution.distribution import distribution_dependencies
from llama_toolchain.distribution.registry import (
available_distributions,
resolve_distribution,
)
from llama_toolchain.utils import DISTRIBS_BASE_DIR
from .utils import run_command, run_with_pty
DISTRIBS = available_distributions()
class DistributionInstall(Subcommand):
"""Llama cli for configuring llama toolchain configs"""
@ -38,12 +29,13 @@ 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
self.parser.add_argument(
"--name",
type=str,
help="Name of the distribution to install -- (try local-ollama)",
required=True,
choices=[d.name for d in DISTRIBS],
choices=[d.name for d in available_distributions()],
)
self.parser.add_argument(
"--conda-env",
@ -53,6 +45,10 @@ class DistributionInstall(Subcommand):
)
def _run_distribution_install_cmd(self, args: argparse.Namespace) -> None:
from llama_toolchain.distribution.distribution import distribution_dependencies
from llama_toolchain.distribution.registry import resolve_distribution
from .utils import run_command, run_with_pty
os.makedirs(DISTRIBS_BASE_DIR, exist_ok=True)
script = pkg_resources.resource_filename(
"llama_toolchain",

View file

@ -7,10 +7,6 @@
import argparse
from llama_toolchain.cli.subcommand import Subcommand
from llama_toolchain.cli.table import print_table
from llama_toolchain.distribution.distribution import distribution_dependencies
from llama_toolchain.distribution.registry import available_distributions
class DistributionList(Subcommand):
@ -30,6 +26,10 @@ class DistributionList(Subcommand):
pass
def _run_distribution_list_cmd(self, args: argparse.Namespace) -> None:
from llama_toolchain.cli.table import print_table
from llama_toolchain.distribution.distribution import distribution_dependencies
from llama_toolchain.distribution.registry import available_distributions
# eventually, this should query a registry at llama.meta.com/llamastack/distributions
headers = [
"Name",

View file

@ -11,12 +11,8 @@ from pathlib import Path
import yaml
from llama_toolchain.cli.subcommand import Subcommand
from llama_toolchain.distribution.registry import resolve_distribution
from llama_toolchain.distribution.server import main as distribution_server_init
from llama_toolchain.utils import DISTRIBS_BASE_DIR
from .utils import run_command
class DistributionStart(Subcommand):
@ -52,6 +48,10 @@ class DistributionStart(Subcommand):
)
def _run_distribution_start_cmd(self, args: argparse.Namespace) -> None:
from llama_toolchain.distribution.registry import resolve_distribution
from llama_toolchain.distribution.server import main as distribution_server_init
from .utils import run_command
dist = resolve_distribution(args.name)
if dist is None:
self.parser.error(f"Distribution with name {args.name} not found")