From f9fffd42e0b8de61f775a12db0a5dfa9831b27a1 Mon Sep 17 00:00:00 2001 From: kingbri Date: Sat, 16 Nov 2024 23:28:44 -0500 Subject: [PATCH] OAI: Fix inline model loading errors when disabled The admin key check was running even if inline loading was disabled. Fix this bug, but also preserve the existing permission system when inline loading is enabled. Signed-off-by: kingbri --- endpoints/OAI/utils/completion.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/endpoints/OAI/utils/completion.py b/endpoints/OAI/utils/completion.py index e9395255..cfcdebad 100644 --- a/endpoints/OAI/utils/completion.py +++ b/endpoints/OAI/utils/completion.py @@ -119,8 +119,19 @@ async def load_inline_model(model_name: str, request: Request): ): return - # Inline model loading isn't enabled or the user isn't an admin - if not get_key_permission(request) == "admin": + # Return if inline loading is disabled + # Also warn if an admin key is used + if not config.model.inline_model_loading: + if get_key_permission(request) == "admin": + logger.warning( + f"Unable to switch model to {model_name} because " + '"inline_model_loading" is not True in config.yml.' + ) + + return + + # Error if an invalid key is passed + if get_key_permission(request) != "admin": error_message = handle_request_error( f"Unable to switch model to {model_name} because " + "an admin key isn't provided", @@ -129,14 +140,8 @@ async def load_inline_model(model_name: str, request: Request): raise HTTPException(401, error_message) - if not config.model.inline_model_loading: - logger.warning( - f"Unable to switch model to {model_name} because " - '"inline_model_loading" is not True in config.yml.' - ) - - return - + # Start inline loading + # Past here, user is assumed to be admin model_path = pathlib.Path(config.model.model_dir) model_path = model_path / model_name