Update app.py

app.py

@@ -1,37 +1,78 @@
-# ui/main.py
+
 import gradio as gr
 import requests
 import logging
-from typing import Dict, Any
+from typing import Optional, Dict, Any
 from pydantic import BaseModel
 
-# Configure logging
 logging.basicConfig(filename="ui.log", level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Pydantic models (mirroring backend)
 class ResourceLimits(BaseModel):
     cpu_limit: Optional[int] = None
     memory_limit: Optional[int] = None
     gpu_power_limit: Optional[int] = None
 
-# API client class
 class APIClient:
     def __init__(self, base_url: str, api_key: str = None):
-        self.base_url = base_url
+        self.base_url = base_url.rstrip("/")
         self.api_key = api_key
-    # Methods for each endpoint (get_metrics, set_limits, etc.)
 
-# UI components
-def build_dashboard():
-    # Gradio components for metrics and health
-    pass
+    def _headers(self):
+        return {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
+
+    def get_metrics(self) -> Dict[str, Any]:
+        try:
+            resp = requests.get(f"{self.base_url}/metrics", headers=self._headers(), timeout=10)
+            resp.raise_for_status()
+            return resp.json()
+        except Exception as e:
+            logger.error(f"Error fetching metrics: {e}")
+            return {}
+
+    def set_limits(self, limits: ResourceLimits) -> bool:
+        try:
+            resp = requests.post(
+                f"{self.base_url}/limits",
+                json=limits.dict(exclude_none=True),
+                headers=self._headers(),
+                timeout=10,
+            )
+            resp.raise_for_status()
+            return True
+        except Exception as e:
+            logger.error(f"Error setting limits: {e}")
+            return False
+
+def build_dashboard(api_client: APIClient):
+    def fetch_metrics():
+        return api_client.get_metrics()
+
+    def update_limits(cpu, memory, gpu):
+        limits = ResourceLimits(cpu_limit=cpu, memory_limit=memory, gpu_power_limit=gpu)
+        success = api_client.set_limits(limits)
+        return "Limits updated!" if success else "Failed to update limits."
+
+    with gr.Row():
+        with gr.Column():
+            metrics_output = gr.JSON(label="Current Metrics")
+            refresh_btn = gr.Button("Refresh Metrics")
+        with gr.Column():
+            cpu = gr.Number(label="CPU Limit")
+            memory = gr.Number(label="Memory Limit")
+            gpu = gr.Number(label="GPU Power Limit")
+            set_btn = gr.Button("Set Limits")
+            status = gr.Textbox(label="Status", interactive=False)
+
+    refresh_btn.click(fetch_metrics, outputs=metrics_output)
+    set_btn.click(update_limits, inputs=[cpu, memory, gpu], outputs=status)
 
 def main():
+    api_client = APIClient(base_url="http://localhost:8000")
     with gr.Blocks(theme=gr.themes.Default()) as app:
-        # Define tabs and components
-        pass
+        gr.Markdown("# Resource Dashboard")
+        build_dashboard(api_client)
     app.launch(server_name="0.0.0.0", server_port=7860, share=False)
 
 if __name__ == "__main__":
-    main()
+    main()