175 lines
7.1 KiB
Python
175 lines
7.1 KiB
Python
try:
|
|
import gevent.monkey
|
|
|
|
gevent.monkey.patch_all()
|
|
except ImportError:
|
|
pass
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import simplejson as json
|
|
from flask import Flask, jsonify, render_template, request
|
|
|
|
from llm_server.cluster.cluster_config import cluster_config
|
|
from llm_server.cluster.model_choices import get_model_choices
|
|
from llm_server.config.config import mode_ui_names
|
|
from llm_server.config.load import load_config
|
|
from llm_server.database.conn import database
|
|
from llm_server.database.create import create_db
|
|
from llm_server.pre_fork import server_startup
|
|
from llm_server.routes.openai import openai_bp
|
|
from llm_server.routes.server_error import handle_server_error
|
|
from llm_server.routes.v1 import bp
|
|
from llm_server.sock import init_socketio
|
|
|
|
# TODO: what happens when all backends are offline? What about the "online" key in the stats page?
|
|
# TODO: redis SCAN vs KEYS??
|
|
# TODO: implement blind RRD controlled via header and only used when there is a queue on the primary backend(s)
|
|
# TODO: is frequency penalty the same as ooba repetition penalty???
|
|
# TODO: make sure openai_moderation_enabled works on websockets, completions, and chat completions
|
|
# TODO: if a backend is at its limit of concurrent requests, choose a different one
|
|
|
|
# Lower priority
|
|
# TODO: fix moderation freezing after a while
|
|
# TODO: support logit_bias on OpenAI and Ooba endpoints.
|
|
# TODO: add a way to cancel VLLM gens. Maybe use websockets?
|
|
# TODO: validate openai_silent_trim works as expected and only when enabled
|
|
# TODO: rewrite config storage. Store in redis so we can reload it.
|
|
# TODO: set VLLM to stream ALL data using socket.io. If the socket disconnects, cancel generation.
|
|
# TODO: estiamted wait time needs to account for full concurrent_gens but the queue is less than concurrent_gens
|
|
# TODO: the estiamted wait time lags behind the stats
|
|
# TODO: simulate OpenAI error messages regardless of endpoint
|
|
# TODO: send extra headers when ratelimited?
|
|
# TODO: make sure log_prompt() is used everywhere, including errors and invalid requests
|
|
# TODO: unify logging thread in a function and use async/await instead
|
|
# TODO: move the netdata stats to a seperate part of the stats and have it set to the currently selected backend
|
|
# TODO: have VLLM reply with stats (TPS, generated token count, processing time)
|
|
# TODO: add config reloading via stored redis variables
|
|
|
|
# Done, but need to verify
|
|
# TODO: add more excluding to SYSTEM__ tokens
|
|
# TODO: return 200 when returning formatted sillytavern error
|
|
|
|
try:
|
|
import vllm
|
|
except ModuleNotFoundError as e:
|
|
print('Could not import vllm-gptq:', e)
|
|
print('Please see README.md for install instructions.')
|
|
sys.exit(1)
|
|
|
|
import config
|
|
from llm_server import opts
|
|
from llm_server.helpers import auto_set_base_client_api
|
|
from llm_server.llm.vllm.info import vllm_info
|
|
from llm_server.custom_redis import flask_cache
|
|
from llm_server.llm import redis
|
|
from llm_server.routes.v1.generate_stats import generate_stats
|
|
|
|
app = Flask(__name__)
|
|
app.register_blueprint(bp, url_prefix='/api/')
|
|
app.register_blueprint(openai_bp, url_prefix='/api/openai/v1/')
|
|
init_socketio(app)
|
|
flask_cache.init_app(app)
|
|
flask_cache.clear()
|
|
|
|
script_path = os.path.dirname(os.path.realpath(__file__))
|
|
config_path_environ = os.getenv("CONFIG_PATH")
|
|
if config_path_environ:
|
|
config_path = config_path_environ
|
|
else:
|
|
config_path = Path(script_path, 'config', 'config.yml')
|
|
|
|
success, config, msg = load_config(config_path)
|
|
if not success:
|
|
print('Failed to load config:', msg)
|
|
sys.exit(1)
|
|
|
|
database.init_db(config['mysql']['host'], config['mysql']['username'], config['mysql']['password'], config['mysql']['database'])
|
|
create_db()
|
|
|
|
|
|
@app.route('/')
|
|
@app.route('/api')
|
|
@app.route('/api/openai')
|
|
@flask_cache.cached(timeout=10)
|
|
def home():
|
|
base_client_api = redis.get('base_client_api', dtype=str)
|
|
stats = generate_stats()
|
|
model_choices, default_model = get_model_choices()
|
|
default_model_info = model_choices[default_model]
|
|
|
|
if default_model_info['queued'] == 0 and default_model_info['queued'] >= default_model_info['concurrent_gens']:
|
|
# There will be a wait if the queue is empty but prompts are processing, but we don't
|
|
# know how long.
|
|
default_estimated_wait_sec = f"less than {int(default_model_info['estimated_wait'])} seconds"
|
|
else:
|
|
default_estimated_wait_sec = f"{int(default_model_info['estimated_wait'])} seconds"
|
|
|
|
if len(config['analytics_tracking_code']):
|
|
analytics_tracking_code = f"<script>\n{config['analytics_tracking_code']}\n</script>"
|
|
else:
|
|
analytics_tracking_code = ''
|
|
|
|
if config['info_html']:
|
|
info_html = config['info_html']
|
|
else:
|
|
info_html = ''
|
|
|
|
mode_info = ''
|
|
for k, v in cluster_config.all().items():
|
|
if v['mode'] == 'vllm':
|
|
mode_info = vllm_info
|
|
break
|
|
|
|
return render_template('home.html',
|
|
llm_middleware_name=opts.llm_middleware_name,
|
|
analytics_tracking_code=analytics_tracking_code,
|
|
info_html=info_html,
|
|
default_model=default_model_info['model'],
|
|
default_active_gen_workers=default_model_info['processing'],
|
|
default_proompters_in_queue=default_model_info['queued'],
|
|
current_model=opts.manual_model_name if opts.manual_model_name else None, # else running_model,
|
|
client_api=f'https://{base_client_api}',
|
|
ws_client_api=f'wss://{base_client_api}/v1/stream' if opts.enable_streaming else 'disabled',
|
|
default_estimated_wait=default_estimated_wait_sec,
|
|
mode_name=mode_ui_names[opts.frontend_api_mode][0],
|
|
api_input_textbox=mode_ui_names[opts.frontend_api_mode][1],
|
|
streaming_input_textbox=mode_ui_names[opts.frontend_api_mode][2],
|
|
default_context_size=default_model_info['context_size'],
|
|
stats_json=json.dumps(stats, indent=4, ensure_ascii=False),
|
|
extra_info=mode_info,
|
|
openai_client_api=f'https://{base_client_api}/openai/v1' if opts.enable_openi_compatible_backend else 'disabled',
|
|
expose_openai_system_prompt=opts.expose_openai_system_prompt,
|
|
enable_streaming=opts.enable_streaming,
|
|
model_choices=model_choices,
|
|
proompters_5_min=stats['stats']['proompters']['5_min'],
|
|
proompters_24_hrs=stats['stats']['proompters']['24_hrs'],
|
|
)
|
|
|
|
|
|
@app.route('/<first>')
|
|
@app.route('/<first>/<path:rest>')
|
|
def fallback(first=None, rest=None):
|
|
return jsonify({
|
|
'code': 404,
|
|
'msg': 'not found'
|
|
}), 404
|
|
|
|
|
|
@app.errorhandler(500)
|
|
def server_error(e):
|
|
return handle_server_error(e)
|
|
|
|
|
|
@app.before_request
|
|
def before_app_request():
|
|
auto_set_base_client_api(request)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
server_startup(None)
|
|
print('FLASK MODE - Startup complete!')
|
|
app.run(host='0.0.0.0', threaded=False, processes=15)
|