Add Optional SSL Support to controller.py (#2448)

fastchat/serve/controller.py

@@ -8,6 +8,7 @@
 from enum import Enum, auto
 import json
 import logging
+import os
 import time
 from typing import List, Union
 import threading
@@ -318,6 +319,13 @@ def create_controller():
         choices=["lottery", "shortest_queue"],
         default="shortest_queue",
     )
+    parser.add_argument(
+        "--ssl",
+        action="store_true",
+        required=False,
+        default=False,
+        help="Enable SSL. Requires OS Environment variables 'SSL_KEYFILE' and 'SSL_CERTFILE'.",
+    )
     args = parser.parse_args()
     logger.info(f"args: {args}")
 
@@ -327,4 +335,14 @@ def create_controller():
 
 if __name__ == "__main__":
     args, controller = create_controller()
-    uvicorn.run(app, host=args.host, port=args.port, log_level="info")
+    if args.ssl:
+        uvicorn.run(
+            app, 
+            host=args.host, 
+            port=args.port, 
+            log_level="info", 
+            ssl_keyfile=os.environ["SSL_KEYFILE"], 
+            ssl_certfile=os.environ["SSL_CERTFILE"]
+        )
+    else:
+        uvicorn.run(app, host=args.host, port=args.port, log_level="info")