refactor(cli): move out to its own packages (#619)

Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>

openllm-python/src/openllm_cli/extension/__init__.py

@@ -0,0 +1,16 @@
+"""OpenLLM CLI Extension.
+
+The following directory contains all possible extensions for OpenLLM CLI
+For adding new extension, just simply name that ext to `<name_ext>.py` and define
+a ``click.command()`` with the following format:
+
+```python
+import click
+
+@click.command(<name_ext>)
+...
+def cli(...): # <- this is important here, it should always name CLI in order for the extension resolver to know how to import this extensions.
+```
+
+NOTE: Make sure to keep this file blank such that it won't mess with the import order.
+"""

openllm-python/src/openllm_cli/extension/build_base_container.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+import typing as t
+
+import click
+import orjson
+
+import openllm
+
+from openllm_cli import termui
+from openllm_cli._factory import container_registry_option
+from openllm_cli._factory import machine_option
+
+
+if t.TYPE_CHECKING:
+  from openllm_core._typing_compat import LiteralContainerRegistry
+  from openllm_core._typing_compat import LiteralContainerVersionStrategy
+
+
+@click.command(
+  'build_base_container',
+  context_settings=termui.CONTEXT_SETTINGS,
+  help="""Base image builder for BentoLLM.
+
+                By default, the base image will include custom kernels (PagedAttention via vllm, FlashAttention-v2, etc.) built with CUDA 11.8, Python 3.9 on Ubuntu22.04.
+                Optionally, this can also be pushed directly to remote registry. Currently support ``docker.io``, ``ghcr.io`` and ``quay.io``.
+
+                \b
+                If '--machine' is passed, then it will run the process quietly, and output a JSON to the current running terminal.
+                This command is only useful for debugging and for building custom base image for extending BentoML with custom base images and custom kernels.
+
+                Note that we already release images on our CI to ECR and GHCR, so you don't need to build it yourself.
+                """,
+)
+@container_registry_option
+@click.option(
+  '--version-strategy',
+  type=click.Choice(['release', 'latest', 'nightly']),
+  default='nightly',
+  help='Version strategy to use for tagging the image.',
+)
+@click.option('--push/--no-push', help='Whether to push to remote repository', is_flag=True, default=False)
+@machine_option
+def cli(
+  container_registry: tuple[LiteralContainerRegistry, ...] | None,
+  version_strategy: LiteralContainerVersionStrategy,
+  push: bool,
+  machine: bool,
+) -> dict[str, str]:
+  mapping = openllm.bundle.build_container(container_registry, version_strategy, push, machine)
+  if machine:
+    termui.echo(orjson.dumps(mapping, option=orjson.OPT_INDENT_2).decode(), fg='white')
+  return mapping

openllm-python/src/openllm_cli/extension/dive_bentos.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+import shutil
+import subprocess
+import typing as t
+
+import click
+import psutil
+
+from simple_di import Provide
+from simple_di import inject
+
+import bentoml
+
+from bentoml._internal.configuration.containers import BentoMLContainer
+from openllm_cli import termui
+from openllm_cli._factory import bento_complete_envvar
+from openllm_cli._factory import machine_option
+
+
+if t.TYPE_CHECKING:
+  from bentoml._internal.bento import BentoStore
+
+
+@click.command('dive_bentos', context_settings=termui.CONTEXT_SETTINGS)
+@click.argument('bento', type=str, shell_complete=bento_complete_envvar)
+@machine_option
+@click.pass_context
+@inject
+def cli(
+  ctx: click.Context, bento: str, machine: bool, _bento_store: BentoStore = Provide[BentoMLContainer.bento_store]
+) -> str | None:
+  """Dive into a BentoLLM. This is synonymous to cd $(b get <bento>:<tag> -o path)."""
+  try:
+    bentomodel = _bento_store.get(bento)
+  except bentoml.exceptions.NotFound:
+    ctx.fail(f'Bento {bento} not found. Make sure to call `openllm build` first.')
+  if 'bundler' not in bentomodel.info.labels or bentomodel.info.labels['bundler'] != 'openllm.bundle':
+    ctx.fail(
+      f"Bento is either too old or not built with OpenLLM. Make sure to use ``openllm build {bentomodel.info.labels['start_name']}`` for correctness."
+    )
+  if machine:
+    return bentomodel.path
+  # copy and paste this into a new shell
+  if psutil.WINDOWS:
+    subprocess.check_call([shutil.which('dir') or 'dir'], cwd=bentomodel.path)
+  else:
+    subprocess.check_call([shutil.which('ls') or 'ls', '-Rrthla'], cwd=bentomodel.path)
+  ctx.exit(0)

openllm-python/src/openllm_cli/extension/get_containerfile.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+import typing as t
+
+import click
+
+from simple_di import Provide
+from simple_di import inject
+
+import bentoml
+
+from bentoml._internal.bento.bento import BentoInfo
+from bentoml._internal.bento.build_config import DockerOptions
+from bentoml._internal.configuration.containers import BentoMLContainer
+from bentoml._internal.container.generate import generate_containerfile
+from openllm_cli import termui
+from openllm_cli._factory import bento_complete_envvar
+from openllm_core.utils import converter
+
+
+if t.TYPE_CHECKING:
+  from bentoml._internal.bento import BentoStore
+
+
+@click.command(
+  'get_containerfile', context_settings=termui.CONTEXT_SETTINGS, help='Return Containerfile of any given Bento.'
+)
+@click.argument('bento', type=str, shell_complete=bento_complete_envvar)
+@click.pass_context
+@inject
+def cli(ctx: click.Context, bento: str, _bento_store: BentoStore = Provide[BentoMLContainer.bento_store]) -> str:
+  try:
+    bentomodel = _bento_store.get(bento)
+  except bentoml.exceptions.NotFound:
+    ctx.fail(f'Bento {bento} not found. Make sure to call `openllm build` first.')
+  # The logic below are similar to bentoml._internal.container.construct_containerfile
+  with open(bentomodel.path_of('bento.yaml'), 'r') as f:
+    options = BentoInfo.from_yaml_file(f)
+    # NOTE: dockerfile_template is already included in the
+    # Dockerfile inside bento, and it is not relevant to
+    # construct_containerfile. Hence it is safe to set it to None here.
+    # See https://github.com/bentoml/BentoML/issues/3399.
+    docker_attrs = converter.unstructure(options.docker)
+    # NOTE: if users specify a dockerfile_template, we will
+    # save it to /env/docker/Dockerfile.template. This is necessary
+    # for the reconstruction of the Dockerfile.
+    if 'dockerfile_template' in docker_attrs and docker_attrs['dockerfile_template'] is not None:
+      docker_attrs['dockerfile_template'] = 'env/docker/Dockerfile.template'
+    doc = generate_containerfile(
+      docker=DockerOptions(**docker_attrs),
+      build_ctx=bentomodel.path,
+      conda=options.conda,
+      bento_fs=bentomodel._fs,
+      enable_buildkit=True,
+      add_header=True,
+    )
+    termui.echo(doc, fg='white')
+  return bentomodel.path

openllm-python/src/openllm_cli/extension/get_prompt.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+import logging
+import traceback
+import typing as t
+
+import click
+import inflection
+import orjson
+
+from bentoml_cli.utils import opt_callback
+
+import openllm
+import openllm_core
+
+from openllm_cli import termui
+from openllm_cli._factory import model_complete_envvar
+from openllm_core.prompts import process_prompt
+
+
+logger = logging.getLogger(__name__)
+
+
+@click.command('get_prompt', context_settings=termui.CONTEXT_SETTINGS)
+@click.argument(
+  'model_name',
+  type=click.Choice([inflection.dasherize(name) for name in openllm.CONFIG_MAPPING.keys()]),
+  shell_complete=model_complete_envvar,
+)
+@click.argument('prompt', type=click.STRING)
+@click.option('--format', type=click.STRING, default=None)
+@click.option(
+  '--opt',
+  help="Define additional prompt variables. (format: ``--opt system_prompt='You are a useful assistant'``)",
+  required=False,
+  multiple=True,
+  callback=opt_callback,
+  metavar='ARG=VALUE[,ARG=VALUE]',
+)
+@click.pass_context
+def cli(
+  ctx: click.Context, /, model_name: str, prompt: str, format: str | None, _memoized: dict[str, t.Any], **_: t.Any
+) -> str | None:
+  """Get the default prompt used by OpenLLM."""
+  module = getattr(openllm_core.config, f'configuration_{model_name}')
+  _memoized = {k: v[0] for k, v in _memoized.items() if v}
+  try:
+    template = getattr(module, 'DEFAULT_PROMPT_TEMPLATE', None)
+    prompt_mapping = getattr(module, 'PROMPT_MAPPING', None)
+    if template is None:
+      raise click.BadArgumentUsage(f'model {model_name} does not have a default prompt template') from None
+    if callable(template):
+      if format is None:
+        if not hasattr(module, 'PROMPT_MAPPING') or module.PROMPT_MAPPING is None:
+          raise RuntimeError('Failed to find prompt mapping while DEFAULT_PROMPT_TEMPLATE is a function.')
+        raise click.BadOptionUsage(
+          'format',
+          f"{model_name} prompt requires passing '--format' (available format: {list(module.PROMPT_MAPPING)})",
+        )
+      if prompt_mapping is None:
+        raise click.BadArgumentUsage(
+          f'Failed to fine prompt mapping while the default prompt for {model_name} is a callable.'
+        ) from None
+      if format not in prompt_mapping:
+        raise click.BadOptionUsage(
+          'format', f'Given format {format} is not valid for {model_name} (available format: {list(prompt_mapping)})'
+        )
+      _prompt_template = template(format)
+    else:
+      _prompt_template = template
+    try:
+      # backward-compatible. TO BE REMOVED once every model has default system message and prompt template.
+      fully_formatted = process_prompt(prompt, _prompt_template, True, **_memoized)
+    except RuntimeError as err:
+      logger.debug('Exception caught while formatting prompt: %s', err)
+      fully_formatted = openllm.AutoConfig.for_model(model_name).sanitize_parameters(
+        prompt, prompt_template=_prompt_template
+      )[0]
+    termui.echo(orjson.dumps({'prompt': fully_formatted}, option=orjson.OPT_INDENT_2).decode(), fg='white')
+  except Exception as err:
+    traceback.print_exc()
+    raise click.ClickException(f'Failed to determine a default prompt template for {model_name}.') from err
+  ctx.exit(0)

openllm-python/src/openllm_cli/extension/list_bentos.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import click
+import inflection
+import orjson
+
+import bentoml
+import openllm
+
+from bentoml._internal.utils import human_readable_size
+from openllm_cli import termui
+
+
+@click.command('list_bentos', context_settings=termui.CONTEXT_SETTINGS)
+@click.pass_context
+def cli(ctx: click.Context) -> None:
+  """List available bentos built by OpenLLM."""
+  mapping = {
+    k: [
+      {
+        'tag': str(b.tag),
+        'size': human_readable_size(openllm.utils.calc_dir_size(b.path)),
+        'models': [
+          {'tag': str(m.tag), 'size': human_readable_size(openllm.utils.calc_dir_size(m.path))}
+          for m in (bentoml.models.get(_.tag) for _ in b.info.models)
+        ],
+      }
+      for b in tuple(i for i in bentoml.list() if all(k in i.info.labels for k in {'start_name', 'bundler'}))
+      if b.info.labels['start_name'] == k
+    ]
+    for k in tuple(inflection.dasherize(key) for key in openllm.CONFIG_MAPPING.keys())
+  }
+  mapping = {k: v for k, v in mapping.items() if v}
+  termui.echo(orjson.dumps(mapping, option=orjson.OPT_INDENT_2).decode(), fg='white')
+  ctx.exit(0)

openllm-python/src/openllm_cli/extension/list_models.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+import typing as t
+
+import click
+import inflection
+import orjson
+
+import bentoml
+import openllm
+
+from bentoml._internal.utils import human_readable_size
+from openllm_cli import termui
+from openllm_cli._factory import model_complete_envvar
+from openllm_cli._factory import model_name_argument
+
+
+if t.TYPE_CHECKING:
+  from openllm_core._typing_compat import DictStrAny
+
+
+@click.command('list_models', context_settings=termui.CONTEXT_SETTINGS)
+@model_name_argument(required=False, shell_complete=model_complete_envvar)
+def cli(model_name: str | None) -> DictStrAny:
+  """This is equivalent to openllm models --show-available less the nice table."""
+  models = tuple(inflection.dasherize(key) for key in openllm.CONFIG_MAPPING.keys())
+  ids_in_local_store = {
+    k: [
+      i
+      for i in bentoml.models.list()
+      if 'framework' in i.info.labels
+      and i.info.labels['framework'] == 'openllm'
+      and 'model_name' in i.info.labels
+      and i.info.labels['model_name'] == k
+    ]
+    for k in models
+  }
+  if model_name is not None:
+    ids_in_local_store = {
+      k: [
+        i
+        for i in v
+        if 'model_name' in i.info.labels and i.info.labels['model_name'] == inflection.dasherize(model_name)
+      ]
+      for k, v in ids_in_local_store.items()
+    }
+  ids_in_local_store = {k: v for k, v in ids_in_local_store.items() if v}
+  local_models = {
+    k: [{'tag': str(i.tag), 'size': human_readable_size(openllm.utils.calc_dir_size(i.path))} for i in val]
+    for k, val in ids_in_local_store.items()
+  }
+  termui.echo(orjson.dumps(local_models, option=orjson.OPT_INDENT_2).decode(), fg='white')
+  return local_models

openllm-python/src/openllm_cli/extension/playground.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+import importlib.machinery
+import logging
+import os
+import pkgutil
+import subprocess
+import sys
+import tempfile
+import typing as t
+
+import click
+import jupytext
+import nbformat
+import yaml
+
+from openllm import playground
+from openllm_cli import termui
+from openllm_core.utils import is_jupyter_available
+from openllm_core.utils import is_jupytext_available
+from openllm_core.utils import is_notebook_available
+
+
+if t.TYPE_CHECKING:
+  from openllm_core._typing_compat import DictStrAny
+
+logger = logging.getLogger(__name__)
+
+
+def load_notebook_metadata() -> DictStrAny:
+  with open(os.path.join(os.path.dirname(playground.__file__), '_meta.yml'), 'r') as f:
+    content = yaml.safe_load(f)
+  if not all('description' in k for k in content.values()):
+    raise ValueError("Invalid metadata file. All entries must have a 'description' key.")
+  return content
+
+
+@click.command('playground', context_settings=termui.CONTEXT_SETTINGS)
+@click.argument('output-dir', default=None, required=False)
+@click.option(
+  '--port',
+  envvar='JUPYTER_PORT',
+  show_envvar=True,
+  show_default=True,
+  default=8888,
+  help='Default port for Jupyter server',
+)
+@click.pass_context
+def cli(ctx: click.Context, output_dir: str | None, port: int) -> None:
+  """OpenLLM Playground.
+
+  A collections of notebooks to explore the capabilities of OpenLLM.
+  This includes notebooks for fine-tuning, inference, and more.
+
+  All of the script available in the playground can also be run directly as a Python script:
+  For example:
+
+  \b
+  ```bash
+  python -m openllm.playground.falcon_tuned --help
+  ```
+
+  \b
+  > [!NOTE]
+  > This command requires Jupyter to be installed. Install it with 'pip install "openllm[playground]"'
+  """
+  if not is_jupyter_available() or not is_jupytext_available() or not is_notebook_available():
+    raise RuntimeError(
+      "Playground requires 'jupyter', 'jupytext', and 'notebook'. Install it with 'pip install \"openllm[playground]\"'"
+    )
+  metadata = load_notebook_metadata()
+  _temp_dir = False
+  if output_dir is None:
+    _temp_dir = True
+    output_dir = tempfile.mkdtemp(prefix='openllm-playground-')
+  else:
+    os.makedirs(os.path.abspath(os.path.expandvars(os.path.expanduser(output_dir))), exist_ok=True)
+
+  termui.echo('The playground notebooks will be saved to: ' + os.path.abspath(output_dir), fg='blue')
+  for module in pkgutil.iter_modules(playground.__path__):
+    if module.ispkg or os.path.exists(os.path.join(output_dir, module.name + '.ipynb')):
+      logger.debug(
+        'Skipping: %s (%s)', module.name, 'File already exists' if not module.ispkg else f'{module.name} is a module'
+      )
+      continue
+    if not isinstance(module.module_finder, importlib.machinery.FileFinder):
+      continue
+    termui.echo('Generating notebook for: ' + module.name, fg='magenta')
+    markdown_cell = nbformat.v4.new_markdown_cell(metadata[module.name]['description'])
+    f = jupytext.read(os.path.join(module.module_finder.path, module.name + '.py'))
+    f.cells.insert(0, markdown_cell)
+    jupytext.write(f, os.path.join(output_dir, module.name + '.ipynb'), fmt='notebook')
+  try:
+    subprocess.check_output(
+      [
+        sys.executable,
+        '-m',
+        'jupyter',
+        'notebook',
+        '--notebook-dir',
+        output_dir,
+        '--port',
+        str(port),
+        '--no-browser',
+        '--debug',
+      ]
+    )
+  except subprocess.CalledProcessError as e:
+    termui.echo(e.output, fg='red')
+    raise click.ClickException(f'Failed to start a jupyter server:\n{e}') from None
+  except KeyboardInterrupt:
+    termui.echo('\nShutting down Jupyter server...', fg='yellow')
+    if _temp_dir:
+      termui.echo('Note: You can access the generated notebooks in: ' + output_dir, fg='blue')
+  ctx.exit(0)