Merge pull request #4 from Kawaeee/streamlit-logging

Update Streamlit logging and server usage

requirements.txt

@@ -5,4 +5,5 @@ tqdm==4.60.0
 pillow==8.2.0
 streamlit==0.80.0
 torch==1.8.1
-torchvision==0.9.1
+torchvision==0.9.1
+psutil==5.8.0

streamlit_app.py

@@ -4,6 +4,7 @@
 import time
 import os
 import requests
+import psutil
 
 from PIL import Image
 import numpy as np
@@ -94,7 +95,7 @@
 )
 
 
-@st.cache(allow_output_mutation=True, max_entries=5, ttl=3600)
+@st.cache(allow_output_mutation=True, max_entries=3, ttl=1800)
 def initialize_model(device=processing_device):
     """Retrieves the butt_bread trained model and maps it to the CPU by default, can also specify GPU here."""
     model = models.resnet152(pretrained=False).to(device)
@@ -147,14 +148,30 @@ def download_model():
     return True
 
 
+def health_check():
+    """ "Check CPU/Memory/Disk usage of deployed machine"""
+    cpu_percent = psutil.cpu_percent(0.15)
+    total_memory = psutil.virtual_memory().total / float(1 << 30)
+    used_memory = psutil.virtual_memory().used / float(1 << 30)
+    total_disk = psutil.disk_usage("/").total / float(1 << 30)
+    used_disk = psutil.disk_usage("/").used / float(1 << 30)
+
+    cpu_usage = "CPU Usage: {:.2f}%".format(cpu_percent)
+    memory_usage = "Memory usage: {:,.2f}G/{:,.2f}G".format(used_memory, total_memory)
+    disk_usage = "Disk usage: {:,.2f}G/{:,.2f}G".format(used_disk, total_disk)
+
+    return " | ".join([cpu_usage, memory_usage, disk_usage])
+
+
 if __name__ == "__main__":
     img_file = None
     img = None
     prediction = None
 
     download_model()
     model = initialize_model()
-
+
+    st_logger.info("[DEBUG] %s",health_check(),exc_info=0)
     st_logger.info("[INFO] Initialize %s model successfully", "buttbread_resnet152_3.h5", exc_info=0)
 
     st.title("Corgi butt or loaf of bread? 🐕🍞")
@@ -189,7 +206,7 @@ def download_model():
                 img.filename = os.path.basename(img_file)
 
             prediction = predict(img, model)
-
+            st_logger.info("[DEBUG] %s",health_check(),exc_info=0)
             st_logger.info("[INFO] Predict %s image successfully", img.filename, exc_info=0)
 
         except Exception as e:
@@ -205,3 +222,6 @@ def download_model():
         st.image(resized_image)
         st.write("Prediction:")
         st.json(prediction)
+
+    # Reset model after used
+    model = None