Skip to content

Commit 8675a7b

Browse files
bojiangbojiang
committed
fix: openllm hello nits
1 parent 690d682 commit 8675a7b

1 file changed

Lines changed: 116 additions & 80 deletions

File tree

‎src/openllm/__main__.py‎

Lines changed: 116 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1+
import importlib.metadata
12
import os
3+
import platform
24
import random
35
import sys
4-
import platform
5-
import importlib.metadata
66
from collections import defaultdict
77
from typing import Annotated, Optional
88

@@ -22,35 +22,42 @@
2222
from openllm.repo import app as repo_app
2323

2424
app = OpenLLMTyper(
25-
help='`openllm hello` to get started. '
26-
'OpenLLM is a CLI tool to manage and deploy open source LLMs and'
27-
' get an OpenAI API compatible chat server in seconds.'
25+
help="`openllm hello` to get started. "
26+
"OpenLLM is a CLI tool to manage and deploy open source LLMs and"
27+
" get an OpenAI API compatible chat server in seconds."
2828
)
2929

30-
app.add_typer(repo_app, name='repo')
31-
app.add_typer(model_app, name='model')
32-
app.add_typer(clean_app, name='clean')
30+
app.add_typer(repo_app, name="repo")
31+
app.add_typer(model_app, name="model")
32+
app.add_typer(clean_app, name="clean")
3333

3434

3535
def _select_bento_name(models: list[BentoInfo], target: DeploymentTarget):
3636
from tabulate import tabulate
3737

3838
options = []
39-
model_infos = [(model.repo.name, model.name, can_run(model, target)) for model in models]
39+
model_infos = [
40+
(model.repo.name, model.name, can_run(model, target)) for model in models
41+
]
4042
model_name_groups = defaultdict(lambda: 0.0)
4143
for repo, name, score in model_infos:
4244
model_name_groups[(repo, name)] += score
43-
table_data = [(name, repo, CHECKED if score > 0 else '') for (repo, name), score in model_name_groups.items()]
45+
table_data = [
46+
(name, repo, CHECKED if score > 0 else "")
47+
for (repo, name), score in model_name_groups.items()
48+
]
4449
if not table_data:
45-
output('No model found', style='red')
50+
output("No model found", style="red")
4651
raise typer.Exit(1)
47-
table = tabulate(table_data, headers=['model', 'repo', 'locally runnable']).split('\n')
48-
headers = f'{table[0]}\n {table[1]}'
52+
table = tabulate(table_data, headers=["model", "repo", "locally runnable"]).split(
53+
"\n"
54+
)
55+
headers = f"{table[0]}\n {table[1]}"
4956

5057
options.append(questionary.Separator(headers))
5158
for table_data, table_line in zip(table_data, table[2:]):
5259
options.append(questionary.Choice(table_line, value=table_data[:2]))
53-
selected = questionary.select('Select a model', options).ask()
60+
selected = questionary.select("Select a model", options).ask()
5461
if selected is None:
5562
raise typer.Exit(1)
5663
return selected
@@ -60,24 +67,26 @@ def _select_bento_version(models, target, bento_name, repo):
6067
from tabulate import tabulate
6168

6269
model_infos = [
63-
[model, can_run(model, target)] for model in models if model.name == bento_name and model.repo.name == repo
70+
[model, can_run(model, target)]
71+
for model in models
72+
if model.name == bento_name and model.repo.name == repo
6473
]
6574

6675
table_data = [
67-
[model.tag, CHECKED if score > 0 else '']
76+
[model.tag, CHECKED if score > 0 else ""]
6877
for model, score in model_infos
6978
if model.name == bento_name and model.repo.name == repo
7079
]
7180
if not table_data:
72-
output(f'No model found for {bento_name} in {repo}', style='red')
81+
output(f"No model found for {bento_name} in {repo}", style="red")
7382
raise typer.Exit(1)
74-
table = tabulate(table_data, headers=['version', 'locally runnable']).split('\n')
83+
table = tabulate(table_data, headers=["version", "locally runnable"]).split("\n")
7584

7685
options = []
77-
options.append(questionary.Separator(f'{table[0]}\n {table[1]}'))
86+
options.append(questionary.Separator(f"{table[0]}\n {table[1]}"))
7887
for table_data, table_line in zip(model_infos, table[2:]):
7988
options.append(questionary.Choice(table_line, value=table_data))
80-
selected = questionary.select('Select a version', options).ask()
89+
selected = questionary.select("Select a version", options).ask()
8190
if selected is None:
8291
raise typer.Exit(1)
8392
return selected
@@ -89,121 +98,143 @@ def _select_target(bento, targets):
8998
options = []
9099
targets.sort(key=lambda x: can_run(bento, x), reverse=True)
91100
if not targets:
92-
output('No available instance type, check your bentocloud account', style='red')
101+
output("No available instance type, check your bentocloud account", style="red")
93102
raise typer.Exit(1)
94103

95104
table = tabulate(
96105
[
97106
[
98107
target.name,
99108
target.accelerators_repr,
100-
f'${target.price}',
101-
CHECKED if can_run(bento, target) else 'insufficient res.',
109+
f"${target.price}",
110+
CHECKED if can_run(bento, target) else "insufficient res.",
102111
]
103112
for target in targets
104113
],
105-
headers=['instance type', 'accelerator', 'price/hr', 'deployable'],
106-
).split('\n')
107-
options.append(questionary.Separator(f'{table[0]}\n {table[1]}'))
114+
headers=["instance type", "accelerator", "price/hr", "deployable"],
115+
).split("\n")
116+
options.append(questionary.Separator(f"{table[0]}\n {table[1]}"))
108117

109118
for target, line in zip(targets, table[2:]):
110-
options.append(questionary.Choice(f'{line}', value=target))
111-
selected = questionary.select('Select an instance type', options).ask()
119+
options.append(questionary.Choice(f"{line}", value=target))
120+
selected = questionary.select("Select an instance type", options).ask()
112121
if selected is None:
113122
raise typer.Exit(1)
114123
return selected
115124

116125

117-
def _select_action(bento, score):
126+
def _select_action(bento: BentoInfo, score):
118127
if score > 0:
119128
options = [
120-
questionary.Separator('Available actions'),
121-
questionary.Choice('0. Run the model in terminal', value='run', shortcut_key='0'),
122-
questionary.Separator(f' $ openllm run {bento}'),
123-
questionary.Separator(' '),
124-
questionary.Choice('1. Serve the model locally and get a chat server', value='serve', shortcut_key='1'),
125-
questionary.Separator(f' $ openllm serve {bento}'),
126-
questionary.Separator(' '),
129+
questionary.Separator("Available actions"),
127130
questionary.Choice(
128-
'2. Deploy the model to bentocloud and get a scalable chat server', value='deploy', shortcut_key='2'
131+
"0. Run the model in terminal", value="run", shortcut_key="0"
129132
),
130-
questionary.Separator(f' $ openllm deploy {bento}'),
133+
questionary.Separator(f" $ openllm run {bento}"),
134+
questionary.Separator(" "),
135+
questionary.Choice(
136+
"1. Serve the model locally and get a chat server",
137+
value="serve",
138+
shortcut_key="1",
139+
),
140+
questionary.Separator(f" $ openllm serve {bento}"),
141+
questionary.Separator(" "),
142+
questionary.Choice(
143+
"2. Deploy the model to bentocloud and get a scalable chat server",
144+
value="deploy",
145+
shortcut_key="2",
146+
),
147+
questionary.Separator(f" $ openllm deploy {bento}"),
131148
]
132149
else:
133150
options = [
134-
questionary.Separator('Available actions'),
151+
questionary.Separator("Available actions"),
135152
questionary.Choice(
136-
'0. Run the model in terminal', value='run', disabled='insufficient res.', shortcut_key='0'
153+
"0. Run the model in terminal",
154+
value="run",
155+
disabled="insufficient res.",
156+
shortcut_key="0",
137157
),
138-
questionary.Separator(f' $ openllm run {bento}'),
139-
questionary.Separator(' '),
158+
questionary.Separator(f" $ openllm run {bento}"),
159+
questionary.Separator(" "),
140160
questionary.Choice(
141-
'1. Serve the model locally and get a chat server',
142-
value='serve',
143-
disabled='insufficient res.',
144-
shortcut_key='1',
161+
"1. Serve the model locally and get a chat server",
162+
value="serve",
163+
disabled="insufficient res.",
164+
shortcut_key="1",
145165
),
146-
questionary.Separator(f' $ openllm serve {bento}'),
147-
questionary.Separator(' '),
166+
questionary.Separator(f" $ openllm serve {bento}"),
167+
questionary.Separator(" "),
148168
questionary.Choice(
149-
'2. Deploy the model to bentocloud and get a scalable chat server', value='deploy', shortcut_key='2'
169+
"2. Deploy the model to bentocloud and get a scalable chat server",
170+
value="deploy",
171+
shortcut_key="2",
150172
),
151-
questionary.Separator(f' $ openllm deploy {bento}'),
173+
questionary.Separator(f" $ openllm deploy {bento}"),
152174
]
153-
action = questionary.select('Select an action', options).ask()
175+
action = questionary.select("Select an action", options).ask()
154176
if action is None:
155177
raise typer.Exit(1)
156-
if action == 'run':
178+
if action == "run":
157179
try:
158-
local_run(bento)
180+
port = random.randint(30000, 40000)
181+
local_run(bento, port=port)
159182
finally:
160-
output('\nUse this command to run the action again:', style='green')
161-
output(f' $ openllm run {bento}', style='orange')
162-
elif action == 'serve':
183+
output("\nUse this command to run the action again:", style="green")
184+
output(f" $ openllm run {bento}", style="orange")
185+
elif action == "serve":
163186
try:
164187
local_serve(bento)
165188
finally:
166-
output('\nUse this command to run the action again:', style='green')
167-
output(f' $ openllm serve {bento}', style='orange')
168-
elif action == 'deploy':
189+
output("\nUse this command to run the action again:", style="green")
190+
output(f" $ openllm serve {bento}", style="orange")
191+
elif action == "deploy":
169192
ensure_cloud_context()
170193
targets = get_cloud_machine_spec()
171194
target = _select_target(bento, targets)
172195
try:
173196
cloud_deploy(bento, target)
174197
finally:
175-
output('\nUse this command to run the action again:', style='green')
176-
output(f' $ openllm deploy {bento} --instance-type {target.name}', style='orange')
198+
output("\nUse this command to run the action again:", style="green")
199+
output(
200+
f" $ openllm deploy {bento} --instance-type {target.name}",
201+
style="orange",
202+
)
177203

178204

179-
@app.command(help='get started interactively')
205+
@app.command(help="get started interactively")
180206
def hello():
181207
INTERACTIVE.set(True)
182-
VERBOSE_LEVEL.set(20)
208+
# VERBOSE_LEVEL.set(20)
183209

184210
target = get_local_machine_spec()
185-
output(f' Detected Platform: {target.platform}', style='green')
211+
output(f" Detected Platform: {target.platform}", style="green")
186212
if target.accelerators:
187-
output(' Detected Accelerators: ', style='green')
213+
output(" Detected Accelerators: ", style="green")
188214
for a in target.accelerators:
189-
output(f' - {a.model} {a.memory_size}GB', style='green')
215+
output(f" - {a.model} {a.memory_size}GB", style="green")
190216
else:
191-
output(' Detected Accelerators: None', style='yellow')
217+
output(" Detected Accelerators: None", style="yellow")
192218

193219
models = list_bento()
194220
if not models:
195-
output('No model found, you probably need to update the model repo:', style='red')
196-
output(' $ openllm repo update', style='orange')
221+
output(
222+
"No model found, you probably need to update the model repo:", style="red"
223+
)
224+
output(" $ openllm repo update", style="orange")
197225
raise typer.Exit(1)
198226

199227
bento_name, repo = _select_bento_name(models, target)
200228
bento, score = _select_bento_version(models, target, bento_name, repo)
201229
_select_action(bento, score)
202230

203231

204-
@app.command(help='start an OpenAI API compatible chat server and chat in browser')
232+
@app.command(help="start an OpenAI API compatible chat server and chat in browser")
205233
def serve(
206-
model: Annotated[str, typer.Argument()] = '', repo: Optional[str] = None, port: int = 3000, verbose: bool = False
234+
model: Annotated[str, typer.Argument()] = "",
235+
repo: Optional[str] = None,
236+
port: int = 3000,
237+
verbose: bool = False,
207238
):
208239
if verbose:
209240
VERBOSE_LEVEL.set(20)
@@ -212,9 +243,9 @@ def serve(
212243
local_serve(bento, port=port)
213244

214245

215-
@app.command(help='run the model and chat in terminal')
246+
@app.command(help="run the model and chat in terminal")
216247
def run(
217-
model: Annotated[str, typer.Argument()] = '',
248+
model: Annotated[str, typer.Argument()] = "",
218249
repo: Optional[str] = None,
219250
port: Optional[int] = None,
220251
timeout: int = 600,
@@ -229,9 +260,11 @@ def run(
229260
local_run(bento, port=port, timeout=timeout)
230261

231262

232-
@app.command(help='deploy an production-ready OpenAI API compatible chat server to bentocloud ($100 free credit)')
263+
@app.command(
264+
help="deploy an production-ready OpenAI API compatible chat server to bentocloud ($100 free credit)"
265+
)
233266
def deploy(
234-
model: Annotated[str, typer.Argument()] = '',
267+
model: Annotated[str, typer.Argument()] = "",
235268
instance_type: Optional[str] = None,
236269
repo: Optional[str] = None,
237270
verbose: bool = False,
@@ -246,20 +279,23 @@ def deploy(
246279
targets = filter(lambda x: can_run(bento, x) > 0, targets)
247280
targets = sorted(targets, key=lambda x: can_run(bento, x), reverse=True)
248281
if not targets:
249-
output('No available instance type, check your bentocloud account', style='red')
282+
output("No available instance type, check your bentocloud account", style="red")
250283
raise typer.Exit(1)
251284
target = targets[0]
252-
output(f'Recommended instance type: {target.name}', style='green')
285+
output(f"Recommended instance type: {target.name}", style="green")
253286
cloud_deploy(bento, target)
254287

255288

256289
@app.callback(invoke_without_command=True)
257290
def typer_callback(
258291
verbose: int = 0,
259292
do_not_track: bool = typer.Option(
260-
False, '--do-not-track', help='Whether to disable usage tracking', envvar=DO_NOT_TRACK
293+
False,
294+
"--do-not-track",
295+
help="Whether to disable usage tracking",
296+
envvar=DO_NOT_TRACK,
261297
),
262-
version: bool = typer.Option(False, '--version', '-v', help='Show version'),
298+
version: bool = typer.Option(False, "--version", "-v", help="Show version"),
263299
):
264300
if verbose:
265301
VERBOSE_LEVEL.set(verbose)
@@ -272,5 +308,5 @@ def typer_callback(
272308
os.environ[DO_NOT_TRACK] = str(True)
273309

274310

275-
if __name__ == '__main__':
311+
if __name__ == "__main__":
276312
app()

0 commit comments

Comments
 (0)