From e8224c542f9910a820b154af037f7b7162c54b5c Mon Sep 17 00:00:00 2001 From: agent Date: Thu, 17 Oct 2024 20:43:31 +0800 Subject: [PATCH] chore: add OPENLLM_TEST_REPO --- src/openllm/common.py | 1 + src/openllm/model.py | 19 +++-- src/openllm/repo.py | 161 +++++++++++++++++++++++++----------------- 3 files changed, 106 insertions(+), 75 deletions(-) diff --git a/src/openllm/common.py b/src/openllm/common.py index 89da93583..0fc022a86 100644 --- a/src/openllm/common.py +++ b/src/openllm/common.py @@ -23,6 +23,7 @@ SUCCESS_STYLE = 'green' OPENLLM_HOME = pathlib.Path(os.getenv('OPENLLM_HOME', pathlib.Path.home() / '.openllm')) + REPO_DIR = OPENLLM_HOME / 'repos' TEMP_DIR = OPENLLM_HOME / 'temp' VENV_DIR = OPENLLM_HOME / 'venv' diff --git a/src/openllm/model.py b/src/openllm/model.py index 8e78edbe7..858322966 100644 --- a/src/openllm/model.py +++ b/src/openllm/model.py @@ -7,8 +7,8 @@ from openllm.accelerator_spec import DeploymentTarget, can_run from openllm.analytic import OpenLLMTyper -from openllm.common import VERBOSE_LEVEL, BentoInfo, load_config, output -from openllm.repo import ensure_repo_updated, parse_repo_url +from openllm.common import VERBOSE_LEVEL, BentoInfo, output +from openllm.repo import ensure_repo_updated, list_repo app = OpenLLMTyper(help='manage models') @@ -95,11 +95,12 @@ def list_bento( if repo_name is None and tag and '/' in tag: repo_name, tag = tag.split('/', 1) + repo_list = list_repo(repo_name) if repo_name is not None: - config = load_config() - if repo_name not in config.repos: + repo_map = {repo.name: repo for repo in repo_list} + if repo_name not in repo_map: output(f'Repo `{repo_name}` not found, did you mean one of these?') - for repo_name in config.repos: + for repo_name in repo_map: output(f' {repo_name}') raise typer.Exit(1) @@ -112,12 +113,8 @@ def list_bento( glob_pattern = f'bentoml/bentos/{tag}/*' model_list = [] - config = load_config() - for _repo_name, repo_url in config.repos.items(): - if repo_name is not None and _repo_name != repo_name: - continue - repo = parse_repo_url(repo_url, _repo_name) - + repo_list = list_repo(repo_name) + for repo in repo_list: paths = sorted( repo.path.glob(glob_pattern), key=lambda x: (x.parent.name, _extract_first_number(x.name), len(x.name), x.name), diff --git a/src/openllm/repo.py b/src/openllm/repo.py index e6e43f390..e3b7e16b9 100644 --- a/src/openllm/repo.py +++ b/src/openllm/repo.py @@ -3,6 +3,9 @@ import shutil import typing +from pathlib import Path +import os + import pyaml import questionary import typer @@ -11,22 +14,25 @@ from openllm.common import INTERACTIVE, REPO_DIR, VERBOSE_LEVEL, RepoInfo, load_config, output, save_config UPDATE_INTERVAL = datetime.timedelta(days=3) +TEST_REPO = os.getenv('OPENLLM_TEST_REPO', None) # for testing + app = OpenLLMTyper(help='manage repos') @app.command(name='list', help='list available repo') -def list_repo(verbose: bool = False): +def cmd_list(verbose: bool = False): if verbose: VERBOSE_LEVEL.set(20) - config = load_config() pyaml.pprint( - [parse_repo_url(repo, name) for name, repo in config.repos.items()], sort_dicts=False, sort_keys=False + list_repo(), sort_dicts=False, sort_keys=False ) -@app.command(help='remove given repo') -def remove(name: str): +@app.command(name='remove', help='remove given repo') +def cmd_remove(name: str): + if TEST_REPO: + return config = load_config() if name not in config.repos: output(f'Repo {name} does not exist', style='red') @@ -37,6 +43,92 @@ def remove(name: str): output(f'Repo {name} removed', style='green') +@app.command(name='update', help='update default repo') +def cmd_update(): + if TEST_REPO: + return + repos_in_use = set() + for repo in list_repo(): + repos_in_use.add((repo.server, repo.owner, repo.repo, repo.branch)) + if repo.path.exists(): + shutil.rmtree(repo.path, ignore_errors=True) + repo.path.parent.mkdir(parents=True, exist_ok=True) + try: + _clone_repo(repo) + output('') + output(f'Repo `{repo.name}` updated', style='green') + except Exception as e: + shutil.rmtree(repo.path, ignore_errors=True) + output(f'Failed to clone repo {repo.name}', style='red') + output(e) + for c in REPO_DIR.glob('*/*/*/*'): + repo_spec = tuple(c.parts[-4:]) + if repo_spec not in repos_in_use: + shutil.rmtree(c, ignore_errors=True) + output(f'Removed unused repo cache {c}') + with open(REPO_DIR / 'last_update', 'w') as f: + f.write(datetime.datetime.now().isoformat()) + for repo in list_repo(): + _complete_alias(repo.name) + + +@app.command(name='add', help='add new repo') +def cmd_add(name: str, repo: str): + if TEST_REPO: + return + name = name.lower() + if not name.isidentifier(): + output(f'Invalid repo name: {name}, should only contain letters, numbers and underscores', style='red') + return + + try: + parse_repo_url(repo) + except ValueError: + output(f'Invalid repo url: {repo}', style='red') + return + + config = load_config() + if name in config.repos: + override = questionary.confirm(f'Repo {name} already exists({config.repos[name]}), override?').ask() + if not override: + return + + config.repos[name] = repo + save_config(config) + output(f'Repo {name} added', style='green') + + +@app.command(name='default', help='get default repo path') +def default(): + if TEST_REPO: + return + output((info := parse_repo_url(load_config().repos['default'], 'default')).path) + return info.path + + +def list_repo(repo_name: typing.Optional[str]=None) -> typing.List[RepoInfo]: + if TEST_REPO: + return [ + RepoInfo( + name='default', + url='', + server='test', + owner='test', + repo='test', + branch='main', + path=Path(TEST_REPO), + ) + ] + config = load_config() + repos = [] + for _repo_name, repo_url in config.repos.items(): + if repo_name is not None and _repo_name != repo_name: + continue + repo = parse_repo_url(repo_url, _repo_name) + repos.append(repo) + return repos + + def _complete_alias(repo_name: str): from openllm.model import list_bento @@ -63,35 +155,6 @@ def _clone_repo(repo: RepoInfo): dulwich.porcelain.clone(repo.url, str(repo.path), checkout=True, depth=1, branch=repo.branch) -@app.command(help='update default repo') -def update(): - config = load_config() - repos_in_use = set() - for repo_name, repo in config.repos.items(): - repo = parse_repo_url(repo, repo_name) - repos_in_use.add((repo.server, repo.owner, repo.repo, repo.branch)) - if repo.path.exists(): - shutil.rmtree(repo.path, ignore_errors=True) - repo.path.parent.mkdir(parents=True, exist_ok=True) - try: - _clone_repo(repo) - output('') - output(f'Repo `{repo.name}` updated', style='green') - except Exception as e: - shutil.rmtree(repo.path, ignore_errors=True) - output(f'Failed to clone repo {repo.name}', style='red') - output(e) - for c in REPO_DIR.glob('*/*/*/*'): - repo_spec = tuple(c.parts[-4:]) - if repo_spec not in repos_in_use: - shutil.rmtree(c, ignore_errors=True) - output(f'Removed unused repo cache {c}') - with open(REPO_DIR / 'last_update', 'w') as f: - f.write(datetime.datetime.now().isoformat()) - for repo_name in config.repos: - _complete_alias(repo_name) - - def ensure_repo_updated(): last_update_file = REPO_DIR / 'last_update' if not last_update_file.exists(): @@ -182,35 +245,5 @@ def parse_repo_url(repo_url: str, repo_name: typing.Optional[str] = None) -> Rep ) -@app.command(help='add new repo') -def add(name: str, repo: str): - name = name.lower() - if not name.isidentifier(): - output(f'Invalid repo name: {name}, should only contain letters, numbers and underscores', style='red') - return - - try: - parse_repo_url(repo) - except ValueError: - output(f'Invalid repo url: {repo}', style='red') - return - - config = load_config() - if name in config.repos: - override = questionary.confirm(f'Repo {name} already exists({config.repos[name]}), override?').ask() - if not override: - return - - config.repos[name] = repo - save_config(config) - output(f'Repo {name} added', style='green') - - -@app.command(help='get default repo path') -def default(): - output((info := parse_repo_url(load_config().repos['default'], 'default')).path) - return info.path - - if __name__ == '__main__': app()