chore: add nightly and check default for model list

Signed-off-by: Aaron Pham <contact@aarnphm.xyz>

Author: aarnphm

src/openllm/cloud.py

@@ -117,6 +117,15 @@ def _get_deploy_cmd(
   return cmd, env
 
 
+def get_current_context() -> str | None:
+  cmd = ['bentoml', 'cloud', 'current-context']
+  try:
+    result = subprocess.check_output(cmd, stderr=subprocess.DEVNULL)
+    return typing.cast(str, json.loads(result)['name'])
+  except subprocess.CalledProcessError:
+    return None
+
+
 def ensure_cloud_context() -> None:
   import questionary
 
@@ -167,6 +176,10 @@ def get_cloud_machine_spec(context: typing.Optional[str] = None) -> list[Deploym
   cmd = ['bentoml', 'deployment', 'list-instance-types', '-o', 'json']
   if context:
     cmd += ['--context', context]
+
+  if context is None:
+    context = get_current_context()
+
   try:
     result = subprocess.check_output(cmd, stderr=subprocess.DEVNULL)
     instance_types = json.loads(result)
@@ -185,7 +198,10 @@ def get_cloud_machine_spec(context: typing.Optional[str] = None) -> list[Deploym
       for it in instance_types
     ]
   except (subprocess.CalledProcessError, json.JSONDecodeError):
-    output('Failed to get cloud instance types', style='red')
+    output(
+      f'Failed to get cloud instance types{"" if context is None else f" for context {context}"}',
+      style='red',
+    )
     return []
 
 

src/openllm/common.py

@@ -73,7 +73,10 @@ def output(
 
 class Config(pydantic.BaseModel):
   repos: dict[str, str] = pydantic.Field(
-    default_factory=lambda: {'default': 'https://github.com/bentoml/openllm-models@main'}
+    default_factory=lambda: {
+      'default': 'https://github.com/bentoml/openllm-models@main',
+      'nightly': 'https://github.com/bentoml/openllm-models@nightly',
+    }
   )
   default_repo: str = 'default'
 

src/openllm/model.py

@@ -6,7 +6,7 @@
 from openllm.accelerator_spec import can_run
 from openllm.common import DeploymentTarget
 from openllm.analytic import OpenLLMTyper
-from openllm.common import VERBOSE_LEVEL, BentoInfo, output as output_
+from openllm.common import VERBOSE_LEVEL, BentoInfo, output as output_, load_config
 from openllm.repo import ensure_repo_updated, list_repo
 
 app = OpenLLMTyper(help='manage models')
@@ -31,6 +31,9 @@ def list_model(
   if verbose:
     VERBOSE_LEVEL.set(20)
 
+  if repo is None:
+    repo = load_config().default_repo
+
   bentos = list_bento(tag=tag, repo_name=repo)
   bentos.sort(key=lambda x: x.name)
 
@@ -79,6 +82,11 @@ def ensure_bento(
   target: typing.Optional[DeploymentTarget] = None,
   repo_name: typing.Optional[str] = None,
 ) -> BentoInfo:
+  if repo_name is None:
+    from openllm.common import load_config
+
+    repo_name = load_config().default_repo
+
   bentos = list_bento(model, repo_name=repo_name)
   if len(bentos) == 0:
     output_(f'No model found for {model}', style='red')