chore: add OPENLLM_TEST_REPO

Author: agent

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'

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),

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()