Add serve command line options to list all supported model-ids (#221)

* add modelid serve

* add test

* fix

* fix

* fix

* test error

* test error

* fix

* fix test bentchmark

* fix review

* fix

Author: yutianchen666

README.md

@@ -71,7 +71,14 @@ Deploy a model on Ray and expose an endpoint for serving. This command uses GPT2
 ```bash
 llm_on_ray-serve --config_file llm_on_ray/inference/models/gpt2.yaml
 ```
-
+You can also use model_ids to serve directly through:
+```bash
+llm_on_ray-serve --models gpt2
+```
+List all support model_ids with config file path:
+```bash
+llm_on_ray-serve --list_model_ids
+```
 The default served method is to provide an OpenAI-compatible API server ([OpenAI API Reference](https://platform.openai.com/docs/api-reference/chat)), you can access and test it in many ways:
 ```bash
 # using curl

benchmarks/run_benchmark.sh

@@ -229,4 +229,4 @@ then
     fi
     output_tokens_length=32
     get_best_latency $iter "${input_tokens_length[*]}" $output_tokens_length $benchmark_dir
-fi
+fi

llm_on_ray/inference/serve.py

@@ -20,7 +20,11 @@
 from llm_on_ray.inference.api_server_simple import serve_run
 from llm_on_ray.inference.api_server_openai import openai_serve_run
 from llm_on_ray.inference.predictor_deployment import PredictorDeployment
-from llm_on_ray.inference.inference_config import ModelDescription, InferenceConfig, all_models
+from llm_on_ray.inference.inference_config import (
+    ModelDescription,
+    InferenceConfig,
+    all_models,
+)
 
 
 def get_deployed_models(args):
@@ -90,6 +94,11 @@ def main(argv=None):
         type=str,
         help=f"Only used when config_file is None, valid values can be any items in {list(all_models.keys())}.",
     )
+    parser.add_argument(
+        "--list_model_ids",
+        action="store_true",
+        help="List all supported model IDs with config file path",
+    )
     parser.add_argument(
         "--simple",
         action="store_true",
@@ -130,6 +139,12 @@ def main(argv=None):
 
     args = parser.parse_args(argv)
 
+    all_models_name = list(all_models.keys())
+    if args.list_model_ids:
+        for model in all_models_name:
+            print(f"{model}: \tllm_on_ray/inference/models/{model}.yaml")
+        sys.exit(0)
+
     ray.init(address="auto")
     deployments, model_list = get_deployed_models(args)
     if args.simple:

tests/inference/test_serve.py

@@ -20,20 +20,22 @@
 
 # Parametrize the test function with different combinations of parameters
 @pytest.mark.parametrize(
-    "config_file, models, port, simple, keep_serve_termimal",
+    "config_file, models, port, simple, keep_serve_termimal, list_model_ids",
     [
         (
             config_file,
             models,
             port,
             simple,
             keep_serve_termimal,
+            list_model_ids,
         )
         for config_file in ["../.github/workflows/config/gpt2-ci.yaml"]
         for models in ["gpt2"]
         for port in [8000]
         for simple in [False]
         for keep_serve_termimal in [False]
+        for list_model_ids in [False, True]
     ],
 )
 def test_script(
@@ -42,25 +44,41 @@ def test_script(
     port,
     simple,
     keep_serve_termimal,
+    list_model_ids,
 ):
-    cmd_serve = ["python", "../llm_on_ray/inference/serve.py"]
-    if config_file is not None:
-        cmd_serve.append("--config_file")
-        cmd_serve.append(str(config_file))
-    if models is not None:
-        cmd_serve.append("--models")
-        cmd_serve.append(str(models))
-    if port is not None:
-        cmd_serve.append("--port")
-        cmd_serve.append(str(port))
-    if simple:
-        cmd_serve.append("--simple")
-    if keep_serve_termimal:
-        cmd_serve.append("--keep_serve_termimal")
+    cmd_serve = ["llm_on_ray-serve"]
+    if list_model_ids:
+        cmd_serve.append("--list_model_ids")
+    else:
+        if config_file is not None:
+            cmd_serve.append("--config_file")
+            cmd_serve.append(str(config_file))
+        elif models is not None:
+            cmd_serve.append("--models")
+            cmd_serve.append(str(models))
+        if port is not None:
+            cmd_serve.append("--port")
+            cmd_serve.append(str(port))
+        if simple:
+            cmd_serve.append("--simple")
+        if keep_serve_termimal:
+            cmd_serve.append("--keep_serve_termimal")
 
+    print(cmd_serve)
     result_serve = subprocess.run(cmd_serve, capture_output=True, text=True)
+    if list_model_ids:
+        output = result_serve.stdout.strip()
+        lines = output.split("\n")
+        assert len(lines) > 0, "No model IDs found in the output"
 
-    assert "Error" not in result_serve.stderr
-    assert result_serve.returncode == 0
-    print("Output of stderr:")
-    print(result_serve.stderr)
+        # Check if the model IDs are listed
+        for line in lines:
+            parts = line.split()
+            assert len(parts) == 2, f"Invalid line format: {line}"
+            model_id, config_path = parts
+
+            assert config_path.endswith(".yaml"), f"Invalid config path format: {config_path}"
+
+    assert result_serve.returncode == 0, print(
+        "\n" + "Output of stderr: " + "\n", result_serve.stderr
+    )