Fix serve config parsing

Co-authored-by: Praveen <[email protected]>
ray-project · Oct 4, 2023 · 8ae595b · 8ae595b
1 parent a54ef33
commit 8ae595b
Show file tree

Hide file tree

Showing 7 changed files with 10 additions and 9 deletions.
diff --git a/aviary/backend/server/models.py b/aviary/backend/server/models.py
@@ -508,7 +508,7 @@ class AppArgs(BaseModel):
 
 
 class RouterArgs(BaseModel):
-    models: Dict[str, Union[str, LLMApp]]
+    models: Union[str, LLMApp, List[Union[LLMApp, str]]]
 
 
 class PlacementConfig(BaseModel):

diff --git a/aviary/backend/server/run.py b/aviary/backend/server/run.py
@@ -7,7 +7,7 @@
 from aviary.backend.llm.vllm.vllm_engine import VLLMEngine
 from aviary.backend.llm.vllm.vllm_models import VLLMApp
 from aviary.backend.server.app import RouterDeployment
-from aviary.backend.server.models import LLMApp, ScalingConfig
+from aviary.backend.server.models import LLMApp, RouterArgs, ScalingConfig
 from aviary.backend.server.plugins.deployment_base_client import DeploymentBaseClient
 from aviary.backend.server.plugins.execution_hooks import (
     ExecutionHooks,
@@ -111,7 +111,8 @@ def router_deployment(
 
 
 def router_application(args):
-    llm_apps = parse_args(args, llm_app_cls=VLLMApp)
+    router_args = RouterArgs.parse_obj(args)
+    llm_apps = parse_args(router_args.models, llm_app_cls=VLLMApp)
     return router_deployment(llm_apps, enable_duplicate_models=False)
 
 
@@ -125,7 +126,7 @@ def run(
 
     """
     ray._private.usage.usage_lib.record_library_usage("aviary")
-    router_app = router_application(vllm_base_args)
+    router_app = router_application({"models": vllm_base_args})
 
     host = "0.0.0.0"
 

diff --git a/serve/amazon--LightGPT.yaml b/serve/amazon--LightGPT.yaml
@@ -3,5 +3,5 @@ applications:
   route_prefix: /amazon--LightGPT
   import_path: aviary.backend:router_application
   args:
-    args:
+    models:
       - "./models/continuous_batching/amazon--LightGPT.yaml"
diff --git a/serve/codellama--CodeLlama-34b-Instruct-hf.yaml b/serve/codellama--CodeLlama-34b-Instruct-hf.yaml
@@ -3,5 +3,5 @@ applications:
   route_prefix: /codellama--CodeLlama-34b-Instruct-hf
   import_path: aviary.backend:router_application
   args:
-    args:
+    models:
       - "./models/continuous_batching/codellama--CodeLlama-34b-Instruct-hf.yaml"
diff --git a/serve/meta-llama--Llama-2-13b-chat-hf.yaml b/serve/meta-llama--Llama-2-13b-chat-hf.yaml
@@ -3,5 +3,5 @@ applications:
   route_prefix: /meta-llama--Llama-2-13b-chat-hf
   import_path: aviary.backend:router_application
   args:
-    args:
+    models:
       - "./models/continuous_batching/meta-llama--Llama-2-13b-chat-hf.yaml"
diff --git a/serve/meta-llama--Llama-2-70b-chat-hf.yaml b/serve/meta-llama--Llama-2-70b-chat-hf.yaml
@@ -3,5 +3,5 @@ applications:
   route_prefix: /meta-llama--Llama-2-70b-chat-hf
   import_path: aviary.backend:router_application
   args:
-    args:
+    models:
       - "./models/continuous_batching/meta-llama--Llama-2-70b-chat-hf.yaml"
diff --git a/serve/meta-llama--Llama-2-7b-chat-hf.yaml b/serve/meta-llama--Llama-2-7b-chat-hf.yaml
@@ -3,5 +3,5 @@ applications:
   route_prefix: /meta-llama--Llama-2-7b-chat-hf
   import_path: aviary.backend:router_application
   args:
-    args:
+    models:
       - "./models/continuous_batching/meta-llama--Llama-2-7b-chat-hf.yaml"