update: server (#121)

* update: server

* update: server

Author: Vincentqyw

imcui/api/config/api.yaml

@@ -1,31 +1,15 @@
-# This file was generated using the `serve build` command on Ray v2.38.0.
-
-proxy_location: EveryNode
-http_options:
-  host: 0.0.0.0
-  port: 8001
-
-grpc_options:
-  port: 9000
-  grpc_servicer_functions: []
-
-logging_config:
-  encoding: TEXT
-  log_level: INFO
-  logs_dir: null
-  enable_access_log: true
-
-applications:
-- name: app1
-  route_prefix: /
-  import_path: api.server:service
-  runtime_env: {}
-  deployments:
-  - name: ImageMatchingService
-    num_replicas: 4
-    ray_actor_options:
-      num_cpus: 2.0
-      num_gpus: 1.0
+service:
+  num_replicas: 4
+  ray_actor_options:
+    num_cpus: 2.0
+    num_gpus: 1.0
+  host: &default_host
+    "0.0.0.0"
+  http_options:
+    host: *default_host
+    port: 8001
+  route_prefix: "/"
+  dashboard_port: 8265
 
 api:
   feature:

imcui/api/server.py

@@ -6,38 +6,61 @@
 import numpy as np
 import ray
 import torch
-import yaml
 from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import JSONResponse
 from PIL import Image
 from ray import serve
+import argparse
 
 from . import ImagesInput, to_base64_nparray
 from .core import ImageMatchingAPI
 from ..hloc import DEVICE
+from ..hloc.utils.io import read_yaml
 from ..ui import get_version
 
 warnings.simplefilter("ignore")
 app = FastAPI()
 if ray.is_initialized():
     ray.shutdown()
+
+
+# read some configs
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--config",
+    type=Path,
+    required=False,
+    default=Path(__file__).parent / "config/api.yaml",
+)
+args = parser.parse_args()
+config_path = args.config
+config = read_yaml(config_path)
+num_gpus = 1 if torch.cuda.is_available() else 0
+ray_actor_options = config["service"].get("ray_actor_options", {})
+ray_actor_options.update({"num_gpus": num_gpus})
+dashboard_port = config["service"].get("dashboard_port", 8265)
+http_options = config["service"].get(
+    "http_options",
+    {
+        "host": "0.0.0.0",
+        "port": 8001,
+    },
+)
+num_replicas = config["service"].get("num_replicas", 4)
 ray.init(
-    dashboard_port=8265,
+    dashboard_port=dashboard_port,
     ignore_reinit_error=True,
 )
-serve.start(
-    http_options={"host": "0.0.0.0", "port": 8001},
-)
-
-num_gpus = 1 if torch.cuda.is_available() else 0
+serve.start(http_options=http_options)
 
 
 @serve.deployment(
-    num_replicas=4, ray_actor_options={"num_cpus": 2, "num_gpus": num_gpus}
+    num_replicas=num_replicas,
+    ray_actor_options=ray_actor_options,
 )
 @serve.ingress(app)
 class ImageMatchingService:
-    def __init__(self, conf: dict, device: str):
+    def __init__(self, conf: dict, device: str, **kwargs):
         self.conf = conf
         self.api = ImageMatchingAPI(conf=conf, device=device)
 
@@ -137,7 +160,7 @@ def load_image(self, file_path: Union[str, UploadFile]) -> np.ndarray:
             image_array = np.array(img)
         return image_array
 
-    def postprocess(self, output: dict, skip_keys: list, binarize: bool = True) -> dict:
+    def postprocess(self, output: dict, skip_keys: list, **kwargs) -> dict:
         pred = {}
         for key, value in output.items():
             if key in skip_keys:
@@ -152,19 +175,12 @@ def run(self, host: str = "0.0.0.0", port: int = 8001):
         uvicorn.run(app, host=host, port=port)
 
 
-def read_config(config_path: Path) -> dict:
-    with open(config_path, "r") as f:
-        conf = yaml.safe_load(f)
-    return conf
-
-
-# api server
-conf = read_config(Path(__file__).parent / "config/api.yaml")
-service = ImageMatchingService.bind(conf=conf["api"], device=DEVICE)
-handle = serve.run(service, route_prefix="/")
+if __name__ == "__main__":
+    # api server
+    service = ImageMatchingService.bind(conf=config["api"], device=DEVICE)
+    handle = serve.run(service, route_prefix="/", blocking=False)
 
 # serve run api.server_ray:service
-
 # build to generate config file
 # serve build api.server_ray:service -o api/config/ray.yaml
 # serve run api/config/ray.yaml

imcui/hloc/utils/io.py

@@ -3,6 +3,7 @@
 import numpy as np
 import cv2
 import h5py
+import yaml
 
 from .parsers import names_to_pair, names_to_pair_old
 
@@ -75,3 +76,9 @@ def get_matches(path: Path, name0: str, name1: str) -> Tuple[np.ndarray]:
         matches = np.flip(matches, -1)
     scores = scores[idx]
     return matches, scores
+
+
+def read_yaml(config_path: Path) -> dict:
+    with open(config_path, "r") as f:
+        conf = yaml.safe_load(f)
+    return conf

imcui/ui/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.0.1"
+__version__ = "1.3.0"
 
 
 def get_version():