Added a page to see and control (pause/resume) the model inference endpoints

pages/900_System_Status.py

@@ -0,0 +1,47 @@
+import streamlit as st
+
+from src.st_helpers import st_setup
+from src.common import *
+
+
+if st_setup('LLM Arch'):
+    st.write("# System Status")
+
+    st.write("## HF Inference Endpoint Statuses")
+    st.write("The following endpoints need to be running to run all the demonstrations.")
+
+    endpoints = st.secrets['endpoints'].split(',')
+
+    for i, e in enumerate(endpoints):
+        status = hf_endpoint_status('alfraser', e)
+        message = f'{e} ({status})'
+
+        if i != 0:
+            st.divider()
+
+        status_col, button_col = st.columns([2, 1])
+
+        if status == HF_RUNNING:
+            with status_col:
+                st.success(message)
+            with button_col:
+                if st.button("Pause", key=f'action_{i}'):
+                    pause_hf_endpoint('alfraser', e)
+                    st.rerun()
+        elif status == HF_SCALEDTOZERO:
+            with status_col:
+                st.error(message)
+        elif status == HF_PAUSED:
+            with status_col:
+                st.warning(message)
+            with button_col:
+                if st.button("Resume", key=f'action_{i}'):
+                    resume_hf_endpoint('alfraser', e)
+                    st.rerun()
+        else:
+            with status_col:
+                st.info(message)
+            with button_col:
+                if st.button("Pause", key=f'action_{i}'):
+                    pause_hf_endpoint('alfraser', e)
+                    st.rerun()

src/common.py

@@ -1,7 +1,8 @@
 import os
+import requests
 import streamlit as st
 
-from typing import List
+from typing import List, Tuple
 
 
 data_dir = os.path.join(os.path.dirname(__file__), '..', 'data')
@@ -19,6 +20,41 @@ def hf_api_token() -> str:
     return token
 
 
+HF_RUNNING = 'running'
+HF_SCALEDTOZERO = 'scaleToZero'
+HF_PAUSED = 'paused'
+
+
+def hf_endpoint_status(username: str, endpoint_name: str, api_token: str = None) -> Tuple[bool, bool]:
+    """
+    Utility to check the status of a hugging face inference endpoint
+    :return: tuple of booleans first is if it is available, second is if it fully shutdown.  False, True indicates it is not fully available or fully shutdown so is in process of starting up or shutting down
+    """
+    url = f'https://api.endpoints.huggingface.cloud/v2/endpoint/{username}/{endpoint_name}'
+    if api_token is None:
+        api_token = hf_api_token()
+    response = requests.get(url, headers={"Authorization" : f"Bearer {api_token}"})
+    if response.status_code != 200:
+        raise ValueError(f"Likely config error. Received status code {response.status_code} - {response.text}")
+    else:
+        state = response.json()['status']['state']
+        return state
+
+
+def pause_hf_endpoint(username: str, endpoint_name: str, api_token: str = None) -> None:
+    url = f'https://api.endpoints.huggingface.cloud/v2/endpoint/{username}/{endpoint_name}/pause'
+    if api_token is None:
+        api_token = hf_api_token()
+    requests.post(url, headers={"Authorization": f"Bearer {api_token}"})
+
+
+def resume_hf_endpoint(username: str, endpoint_name: str, api_token: str = None) -> None:
+    url = f'https://api.endpoints.huggingface.cloud/v2/endpoint/{username}/{endpoint_name}/resume'
+    if api_token is None:
+        api_token = hf_api_token()
+    requests.post(url, headers={"Authorization": f"Bearer {api_token}"})
+
+
 def join_items_comma_and(items: List[str]) -> str:
     """
     Utility to convert a list of items to lowercase strings, comma separated and ending with and
@@ -39,4 +75,3 @@ def escape_dollars(text: str) -> str:
     if text is None:
         return text
     return text.replace("$", "\\$")
-