update torchserve and streamlit (#15)

Author: cceyda

README.md

@@ -1,4 +1,7 @@
 # Torchserve Dashboard
+[![Total Downloads](https://pepy.tech/badge/torchserve-dashboard)](https://pepy.tech/project/torchserve-dashboard)
+![](https://img.shields.io/pypi/dm/torchserve-dashboard)
+
 
 Torchserve Dashboard using Streamlit
 
@@ -8,7 +11,7 @@ Related blog [post](https://cceyda.github.io/blog/torchserve/streamlit/dashboard
 
 # Usage
 Additional Requirement: 
-[torchserve](https://github.com/pytorch/serve/tree/v0.3.1#install-torchserve-and-torch-model-archiver) (recommended:v0.3.1)
+[torchserve](https://github.com/pytorch/serve/tree/v0.5.0#install-torchserve-and-torch-model-archiver) (recommended:v0.5.0)
 
 Simply run:
 
@@ -22,7 +25,7 @@ torchserve-dashboard --server.port 8105 -- --config_path ./torchserve.properties
 torchserve-dashboard -- --config_path ./torchserve.properties --model_store ./model_store
 ```
 
-:exclamation: Keep in mind that If you change any of the `--config_path`,`--model_store`,`--metrics_location`,`--log_location` options while there is a torchserver already running before starting torch-dashboard they won't come into effect until you stop&start torchserve.
+:exclamation: Keep in mind that If you change any of the `--config_path`,`--model_store`,`--metrics_location`,`--log_location` options while there is a torchserver already running before starting torch-dashboard they won't come into effect until you stop&start torchserve. These options are used instead of their respective environment variables `TS_CONFIG_FILE, METRICS_LOCATION, LOG_LOCATION`.
 
 OR 
 ```bash
@@ -56,6 +59,8 @@ If the server doesn't start for some reason check if your ports are already in u
 
 [31-may-2021] Update to v0.4 (Add workflow API) Refactor out streamlit from api.py.  
 
+[30-nov-2021] Update to v0.5, adding support for [encrypted model serving](https://github.com/pytorch/serve/blob/v0.5.0/docs/management_api.md#encrypted-model-serving) (not tested). Update streamlit to v1+
+
 # FAQs
 - **Does torchserver keep running in the background?**
 
@@ -65,7 +70,7 @@ If the server doesn't start for some reason check if your ports are already in u
 
     These environment variables are passed to the torchserve command:
 
-    `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","AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_DEFAULT_REGION"]`
 
 - **How to change the logging format of torchserve?**
 

requirements/install.txt

@@ -1,3 +1,2 @@
-streamlit==0.82.0
-click<8.0,>=7.0
-httpx>=0.16.0
+httpx >= 0.16.0
+streamlit == 1.2.0

setup.py

@@ -70,7 +70,7 @@ def _prepare_extras():
     install_requires=setup_tools.load_requirements(
         file_name='install.txt'),  # Optional
     extras_require=_prepare_extras(),
-    python_requires='>=3.6',
+    python_requires='>=3.7',
 
     # test_suite='setup.get_test_suite',
     # tests_require=["coverage"],

torchserve_dashboard/api.py

@@ -9,7 +9,8 @@
 
 ENVIRON_WHITELIST = [
     "LD_LIBRARY_PATH", "LC_CTYPE", "LC_ALL", "PATH", "JAVA_HOME", "PYTHONPATH",
-    "TS_CONFIG_FILE", "LOG_LOCATION", "METRICS_LOCATION"
+    "TS_CONFIG_FILE", "LOG_LOCATION", "METRICS_LOCATION",
+    "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_DEFAULT_REGION"
 ]
 
 log = logging.getLogger(__name__)
@@ -18,14 +19,17 @@
 class LocalTS:
     def __init__(self,
                  model_store: str,
-                 config_path: str,
+                 config_path: Optional[str] = None,
                  log_location: Optional[str] = None,
-                 metrics_location: Optional[str] = None) -> None:
+                 metrics_location: Optional[str] = None,
+                 log_config: Optional[str] = None) -> None:
         new_env = {}
         env = os.environ
         for x in ENVIRON_WHITELIST:
             if x in env:
                 new_env[x] = env[x]
+        if config_path:
+            new_env["TS_CONFIG_FILE"] = config_path
         if log_location:
             new_env["LOG_LOCATION"] = log_location
             if not os.path.isdir(log_location):
@@ -39,6 +43,7 @@ def __init__(self,
         self.config_path = config_path
         self.log_location = log_location
         self.metrics_location = metrics_location
+        self.log_config = log_config
         self.env = new_env
 
     def check_version(self) -> Tuple[str, Union[str, Exception]]:
@@ -62,6 +67,8 @@ def start_torchserve(self) -> str:
             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}"
+        if self.log_config:
+            torchserve_cmd += f" --log-config {self.log_config}"
         p = subprocess.Popen(
             torchserve_cmd.split(" "),
             env=self.env,
@@ -132,6 +139,7 @@ def register_model(
         max_batch_delay: Optional[int] = None,
         initial_workers: Optional[int] = None,
         response_timeout: Optional[int] = None,
+        is_encrypted: Optional[bool] = None,
     ) -> Dict[str, str]:
 
         req_url = self.address + "/models?url=" + mar_path + "&synchronous=false"
@@ -149,6 +157,8 @@ def register_model(
             req_url += "&initial_workers=" + str(initial_workers)
         if response_timeout:
             req_url += "&response_timeout=" + str(response_timeout)
+        if is_encrypted:
+            req_url += "&s3_sse_kms=true"
 
         res = self.client.post(req_url)
         return res.json()

torchserve_dashboard/dash.py

@@ -34,18 +34,24 @@
     default="./logs/metrics/",
     help="Passed as environment variable METRICS_LOCATION to Torchserve",
 )
+parser.add_argument(
+    "--log_config",
+    default=None,
+    help="Overrides the default log4j.properties",
+)
 try:
     args = parser.parse_args()
 except SystemExit as e:
     os._exit(e.code)
 
-
+# Clean this up later
 def check_args(args):
     M_API = "http://127.0.0.1:8081"
     model_store = args.model_store
     config_path = args.config_path
     log_location = args.log_location
     metrics_location = args.metrics_location
+    log_config = args.log_config
     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")
@@ -57,6 +63,8 @@ def check_args(args):
                     model_store = c.split("=")[-1].strip()
             if c.startswith("management_address"):
                 M_API = c.split("=")[-1].strip()
+    else:
+        st.write("Config file can't be found!")
 
     if log_location:
         log_location = str(Path(log_location).resolve())
@@ -73,12 +81,12 @@ def check_args(args):
         st.write(f"Created model store directory {model_store}")
         os.makedirs(model_store, exist_ok=True)
 
-    return M_API, config, model_store, config_path, log_location, metrics_location
+    return M_API, config, model_store, config_path, log_location, metrics_location, log_config
 
 
 st.title("Torchserve Management Dashboard")
 default_key = "None"
-api_address, config, model_store, config_path, log_location, metrics_location = check_args(args)
+api_address, config, model_store, config_path, log_location, metrics_location, log_config = check_args(args)
 
 
 def rerun():
@@ -101,7 +109,7 @@ def get_model_store():
 
 
 api = ManagementAPI(api_address, error_callback)
-ts = LocalTS(model_store, config_path, log_location, metrics_location)
+ts = LocalTS(model_store, config_path, log_location, metrics_location, log_config)
 ts_version,ts_error=ts.check_version() # doing it this way rather than ts.__version__ on purpose
 if ts_error:
     st.error(ts_error)
@@ -114,7 +122,7 @@ def get_model_store():
 # As a design choice I'm leaving config_path,log_location,metrics_location non-editable from the UI as a semi-security measure (maybe?:/)
 ##########Sidebar##########
 st.sidebar.markdown("## Help")
-with st.sidebar.beta_expander(label="Show Paths:", expanded=False):
+with st.sidebar.expander(label="Show Paths:", expanded=False):
     st.markdown(f"### Model Store Path: \n {model_store}")
     st.markdown(f"### Config Path: \n {config_path}")
     st.markdown(f"### Log Location: \n {log_location}")
@@ -145,13 +153,13 @@ def get_model_store():
 
 st.markdown(f"**Last Message**: {last_res()[0]}")
 
-with st.beta_expander(label="Show torchserve config", expanded=False):
+with st.expander(label="Show torchserve config", expanded=False):
     st.write(config)
     st.markdown("[configuration docs](https://pytorch.org/serve/configuration.html)")
 
 if torchserve_status:
 
-    with st.beta_expander(label="Register a model", expanded=False):
+    with st.expander(label="Register a model", expanded=False):
 
         st.markdown(
             "# Register a model [(docs)](https://pytorch.org/serve/management_api.html#register-a-model)"
@@ -160,12 +168,12 @@ def get_model_store():
         mar_path = placeholder.selectbox(
             "Choose mar file *", [default_key] + stored_models, index=0
         )
-        # mar_path = os.path.join(model_store,mar_path)
-        p = st.checkbox("manually enter path")
+        p = st.checkbox("manually enter location")
         if p:
             mar_path = placeholder.text_input("Input mar file path*")
+        
         model_name = st.text_input(label="Model name (overrides predefined)")
-        col1, col2 = st.beta_columns(2)
+        col1, col2 = st.columns(2)
         batch_size = col1.number_input(label="batch_size", value=0, min_value=0, step=1)
         max_batch_delay = col2.number_input(
             label="max_batch_delay", value=0, min_value=0, step=100
@@ -178,7 +186,7 @@ def get_model_store():
         )
         handler = col1.text_input(label="handler")
         runtime = col2.text_input(label="runtime")
-
+        is_encrypted = st.checkbox("SSE-KMS Encrypted", help="Refer to https://github.com/pytorch/serve/blob/v0.5.0/docs/management_api.md#encrypted-model-serving")
         proceed = st.button("Register")
         if proceed:
             if mar_path != default_key:
@@ -192,13 +200,14 @@ def get_model_store():
                     max_batch_delay=max_batch_delay,
                     initial_workers=initial_workers,
                     response_timeout=response_timeout,
+                    is_encrypted=is_encrypted,
                 )
                 last_res()[0] = res
                 rerun()
             else:
                 st.warning(":octagonal_sign: Fill the required fileds!")
 
-    with st.beta_expander(label="Remove a model", expanded=False):
+    with st.expander(label="Remove a model", expanded=False):
 
         st.header("Remove a model")
         model_name = st.selectbox(
@@ -221,7 +230,7 @@ def get_model_store():
                 else:
                     st.warning(":octagonal_sign: Pick a model & version!")
 
-    with st.beta_expander(label="Get model details", expanded=False):
+    with st.expander(label="Get model details", expanded=False):
 
         st.header("Get model details")
         model_name = st.selectbox(
@@ -243,7 +252,7 @@ def get_model_store():
                     res = api.get_model(model_name, version)
                     st.write(res)
 
-    with st.beta_expander(label="Scale workers", expanded=False):
+    with st.expander(label="Scale workers", expanded=False):
         st.markdown(
             "# Scale workers [(docs)](https://pytorch.org/serve/management_api.html#scale-workers)"
         )
@@ -257,7 +266,7 @@ def get_model_store():
             versions = [m["modelVersion"] for m in versions]
             version = st.selectbox("Choose version", ["All"] + versions, index=0)
 
-            col1, col2, col3 = st.beta_columns(3)
+            col1, col2, col3 = st.columns(3)
             min_worker = col1.number_input(
                 label="min_worker(optional)", value=-1, min_value=-1, step=1
             )
@@ -286,7 +295,7 @@ def get_model_store():
                 last_res()[0] = res
                 rerun()
     if support_workflow:
-        with st.beta_expander(label="Register Workflow", expanded=False):
+        with st.expander(label="Register Workflow", expanded=False):
             st.markdown(
                 "# Register a workflow [(docs)](https://pytorch.org/serve/workflow_management_api.html#register-a-workflow)"
             )
@@ -296,7 +305,7 @@ def get_model_store():
                 res = api.register_workflow(url, workflow_name)
                 st.write(res)
 
-        with st.beta_expander(label="Show Workflow Details", expanded=False):
+        with st.expander(label="Show Workflow Details", expanded=False):
             st.markdown(
                 "# Describe a workflow [(docs)](https://pytorch.org/serve/workflow_management_api.html#describe-workflow)"
             )
@@ -310,7 +319,7 @@ def get_model_store():
                     res = api.get_workflow(workflow_name)
                     st.write(res)
 
-        with st.beta_expander(label="Unregister Workflow", expanded=False):
+        with st.expander(label="Unregister Workflow", expanded=False):
             st.markdown(
                 "# Unregister a Workflow [(docs)](https://pytorch.org/serve/workflow_management_api.html#unregister-a-workflow)"
             )
@@ -324,7 +333,7 @@ def get_model_store():
                     res = api.unregister_workflow(workflow_name)
                     st.write(res)
 
-        with st.beta_expander(label="List Workflows", expanded=False):
+        with st.expander(label="List Workflows", expanded=False):
             st.markdown(
                 "# List Workflows"
             )

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist=py36, pypy, flake8
+envlist=py37, pypy, flake8
 
 [testenv]
 commands=py.test --cov torchserve_dashboard {posargs}
@@ -8,7 +8,7 @@ deps=
     pytest-cov
 
 [testenv:flake8]
-basepython = python3.6
+basepython = python3.7
 deps =
     flake8
 commands =