Merge pull request #7 from FlorianMF/typing

WIP: Add type annotations

Author: cceyda

torchserve_dashboard/api.py

@@ -1,59 +1,73 @@
 import os
 import subprocess
+from typing import Any, Dict, List, Optional, Tuple, Union, Callable
 
 import httpx
+from httpx import Response
 
 import logging
 
-ENVIRON_WHITELIST = ["LD_LIBRARY_PATH", "LC_CTYPE", "LC_ALL", "PATH", "JAVA_HOME", "PYTHONPATH", "TS_CONFIG_FILE", "LOG_LOCATION", "METRICS_LOCATION"]
+ENVIRON_WHITELIST = [
+    "LD_LIBRARY_PATH", "LC_CTYPE", "LC_ALL", "PATH", "JAVA_HOME", "PYTHONPATH",
+    "TS_CONFIG_FILE", "LOG_LOCATION", "METRICS_LOCATION"
+]
 
 log = logging.getLogger(__name__)
 
 
 class LocalTS:
-    def __init__(self, model_store, config_path, log_location=None, metrics_location=None):
+    def __init__(self,
+                 model_store: str,
+                 config_path: str,
+                 log_location: Optional[str] = None,
+                 metrics_location: Optional[str] = None) -> None:
         new_env = {}
         env = os.environ
         for x in ENVIRON_WHITELIST:
             if x in env:
                 new_env[x] = env[x]
         if log_location:
             new_env["LOG_LOCATION"] = log_location
+            if not os.path.isdir(log_location):
+                os.makedirs(log_location, exist_ok=True)
         if metrics_location:
             new_env["METRICS_LOCATION"] = metrics_location
-        if not os.path.isdir(metrics_location):
-            os.makedirs(metrics_location, exist_ok=True)
-        if not os.path.isdir(log_location):
-            os.makedirs(log_location, exist_ok=True)
+            if not os.path.isdir(metrics_location):
+                os.makedirs(metrics_location, exist_ok=True)
 
         self.model_store = model_store
         self.config_path = config_path
         self.log_location = log_location
         self.metrics_location = metrics_location
         self.env = new_env
-    
-    def check_version(self):
+
+    def check_version(self) -> Tuple[str, Union[str, Exception]]:
         try:
-            p=subprocess.run(["torchserve","--version"], check=True,
-                            stdout=subprocess.PIPE,stderr=subprocess.PIPE,
-                            universal_newlines=True)
-            return p.stdout ,p.stderr
-        except (subprocess.CalledProcessError,OSError) as e:
-            return "",e
-         
-    def start_torchserve(self):
+            p = subprocess.run(["torchserve", "--version"],
+                               check=True,
+                               stdout=subprocess.PIPE,
+                               stderr=subprocess.PIPE,
+                               universal_newlines=True)
+            return p.stdout, p.stderr
+        except (subprocess.CalledProcessError, OSError) as e:
+            return "", e
+
+    def start_torchserve(self) -> str:
 
         if not os.path.exists(self.model_store):
             return "Can't find model store path"
         elif not os.path.exists(self.config_path):
             return "Can't find configuration path"
-        dashboard_log_path = os.path.join(self.log_location, "torchserve_dashboard.log")
+        dashboard_log_path = os.path.join(
+            self.log_location, "torchserve_dashboard.log"
+        ) if self.log_location is not None else None
         torchserve_cmd = f"torchserve --start --ncs --model-store {self.model_store} --ts-config {self.config_path}"
         p = subprocess.Popen(
             torchserve_cmd.split(" "),
             env=self.env,
             stdout=subprocess.DEVNULL,
-            stderr=open(dashboard_log_path, "a+"),
+            stderr=open(dashboard_log_path, "a+")
+            if dashboard_log_path else subprocess.DEVNULL,
             start_new_session=True,
             close_fds=True  # IDK stackoverflow told me to do it
         )
@@ -63,33 +77,41 @@ def start_torchserve(self):
         else:
             return f"Torchserve is already started. Check {dashboard_log_path} for errors"
 
-    def stop_torchserve(self):
+    def stop_torchserve(self) -> Union[str, Exception]:
         try:
-            p=subprocess.run(["torchserve","--stop"], check=True,
-                            stdout=subprocess.PIPE,stderr=subprocess.PIPE,
-                            universal_newlines=True)
+            p = subprocess.run(["torchserve", "--stop"],
+                               check=True,
+                               stdout=subprocess.PIPE,
+                               stderr=subprocess.PIPE,
+                               universal_newlines=True)
             return p.stdout
-        except (subprocess.CalledProcessError,OSError) as e:
+        except (subprocess.CalledProcessError, OSError) as e:
             return e
 
-class ManagementAPI:
 
-    def __init__(self, address, error_callback):
+class ManagementAPI:
+    def __init__(self, address: str, error_callback: Callable = None) -> None:
         self.address = address
-        self.client = httpx.Client(timeout=1000, event_hooks={"response": [error_callback]})
-
-    def default_error_callback(response):
+        if not error_callback:
+            error_callback=self.default_error_callback
+        self.client = httpx.Client(timeout=1000,
+                                   event_hooks={"response": [error_callback]})
+    @staticmethod
+    def default_error_callback(response: Response) -> None:
         if response.status_code != 200:
             log.info(f"Warn - status code: {response.status_code},{response}")
 
-    def get_loaded_models(self):
+    def get_loaded_models(self) -> Optional[Dict[str, Any]]:
         try:
             res = self.client.get(self.address + "/models")
             return res.json()
         except httpx.HTTPError:
             return None
 
-    def get_model(self, model_name, version=None, list_all=False):
+    def get_model(self,
+                  model_name: str,
+                  version: Optional[str] = None,
+                  list_all: bool = False) -> List[Dict[str, Any]]:
         req_url = self.address + "/models/" + model_name
         if version:
             req_url += "/" + version
@@ -102,15 +124,15 @@ def get_model(self, model_name, version=None, list_all=False):
     # Doesn't have version
     def register_model(
         self,
-        mar_path,
-        model_name=None,
-        handler=None,
-        runtime=None,
-        batch_size=None,
-        max_batch_delay=None,
-        initial_workers=None,
-        response_timeout=None,
-    ):
+        mar_path: str,
+        model_name: Optional[str] = None,
+        handler: Optional[str] = None,
+        runtime: Optional[str] = None,
+        batch_size: Optional[int] = None,
+        max_batch_delay: Optional[int] = None,
+        initial_workers: Optional[int] = None,
+        response_timeout: Optional[int] = None,
+    ) -> Dict[str, str]:
 
         req_url = self.address + "/models?url=" + mar_path + "&synchronous=false"
         if model_name:
@@ -131,22 +153,32 @@ def register_model(
         res = self.client.post(req_url)
         return res.json()
 
-    def delete_model(self, model_name, version):
+    def delete_model(self,
+                     model_name: str,
+                     version: Optional[str] = None) -> Dict[str, str]:
         req_url = self.address + "/models/" + model_name
         if version:
             req_url += "/" + version
         res = self.client.delete(req_url)
         return res.json()
 
-    def change_model_default(self, model_name, version):
+    def change_model_default(self,
+                             model_name: str,
+                             version: Optional[str] = None):
         req_url = self.address + "/models/" + model_name
         if version:
             req_url += "/" + version
         req_url += "/set-default"
         res = self.client.put(req_url)
         return res.json()
 
-    def change_model_workers(self, model_name, version=None, min_worker=None, max_worker=None, number_gpu=None):
+    def change_model_workers(
+            self,
+            model_name: str,
+            version: Optional[str] = None,
+            min_worker: Optional[int] = None,
+            max_worker: Optional[int] = None,
+            number_gpu: Optional[int] = None) -> Dict[str, str]:
         req_url = self.address + "/models/" + model_name
         if version:
             req_url += "/" + version
@@ -160,28 +192,30 @@ def change_model_workers(self, model_name, version=None, min_worker=None, max_wo
         res = self.client.put(req_url)
         return res.json()
 
-    def register_workflow(
-        self,
-        url,
-        workflow_name=None
-    ):
+    def register_workflow(self,
+                          url: str,
+                          workflow_name: Optional[str] = None) -> Dict[str, str]:
         req_url = self.address + "/workflows/" + url
         if workflow_name:
             req_url += "&workflow_name=" + workflow_name
         res = self.client.post(req_url)
         return res.json()
 
-    def get_workflow(self, workflow_name):
+    def get_workflow(self, workflow_name: str) -> Dict[str, str]:
         req_url = self.address + "/workflows/" + workflow_name
         res = self.client.get(req_url)
         return res.json()
 
-    def unregister_workflow(self, workflow_name):
+    def unregister_workflow(self, workflow_name: str) -> Dict[str, str]:
         req_url = self.address + "/workflows/" + workflow_name
         res = self.client.delete(req_url)
         return res.json()
 
-    def list_workflows(self, limit=None, next_page_token=None):
+    def list_workflows(
+        self,
+        limit: Optional[int] = None,
+        next_page_token: Optional[int] = None
+    ) -> Optional[Dict[str, Any]]:
         req_url = self.address + "/workflows/"
         if limit:
             req_url += "&limit=" + str(limit)

torchserve_dashboard/cli.py

@@ -1,17 +1,16 @@
+from typing import Any
 import click
 import streamlit.cli
 from streamlit.cli import configurator_options
 import os
 
-@click.command(context_settings=dict(
-    ignore_unknown_options=True,
-    allow_extra_args = True
-))
+
+@click.command(context_settings=dict(ignore_unknown_options=True,
+                                     allow_extra_args=True))
 @configurator_options
 @click.argument("args", nargs=-1)
 @click.pass_context
-def main(ctx,args,**kwargs):
+def main(ctx: click.Context, args: Any, **kwargs: Any):
     dirname = os.path.dirname(__file__)
     filename = os.path.join(dirname, 'dash.py')
     ctx.forward(streamlit.cli.main_run, target=filename, args=args, *kwargs)
-

torchserve_dashboard/dash.py

@@ -2,8 +2,9 @@
 import os
 
 import streamlit as st
+from httpx import Response
 
-from api import ManagementAPI, LocalTS
+from torchserve_dashboard.api import ManagementAPI, LocalTS
 from pathlib import Path
 
 st.set_page_config(
@@ -48,7 +49,6 @@ def check_args(args):
     if not os.path.exists(config_path):
         st.write(f"Can't find config file at {config_path}. Using default config instead")
         config_path = os.path.join(os.path.dirname(__file__), "default.torchserve.properties")
-
     if os.path.exists(config_path):
         config = open(config_path, "r").readlines()
         for c in config:
@@ -85,7 +85,7 @@ def rerun():
     st.experimental_rerun()
 
 
-def error_callback(response):
+def error_callback(response:Response):
     if response.status_code != 200:
         st.write("There was an error!")
         st.write(response)