Update main.py

main.py

@@ -5,6 +5,7 @@ import time
 import requests
 import asyncio
 import threading
+import platform
 from fastapi import FastAPI, Request, HTTPException, Depends
 from fastapi.responses import JSONResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
@@ -12,6 +13,11 @@ import uvicorn
 from pydantic import BaseModel
 from typing import List, Optional
 
+try:
+    import psutil
+except ImportError:
+    psutil = None
+
 try:
     from vietnamese_ai.serving.batch_server import MayChuBatch  # type: ignore
     from vietnamese_ai.llm.vietnamese_llm import VietnameseLLM  # type: ignore
@@ -54,12 +60,33 @@ batch_server = MayChuBatch(
     so_worker=4
 )
 
+def get_metrics():
+    """Thu thập chỉ số phần cứng thực tế của server"""
+    metrics = {
+        "total_memory_mb": 0,
+        "used_memory_mb": 0,
+        "cpu_percent": 0,
+        "cpu_count": os.cpu_count() or 1,
+        "platform": platform.system(),
+        "max_workers": 2,
+    }
+    if psutil:
+        mem = psutil.virtual_memory()
+        metrics["total_memory_mb"] = round(mem.total / 1024 / 1024)
+        metrics["used_memory_mb"] = round(mem.used / 1024 / 1024)
+        metrics["cpu_percent"] = psutil.cpu_percent(interval=0.1)
+        # Ước tính workers: mỗi worker cần ~2GB RAM
+        metrics["max_workers"] = max(1, int(mem.available / 1024 / 1024 / 2048))
+    return metrics
+
 def start_heartbeat(port):
     while True:
         try:
             requests.post(f"{CENTER_URL}/api/v1/admin/datacenter/register", json={
                 "role": "Gateway",
-                "port": port
+                "port": port,
+                "public_url": os.environ.get("SPACE_HOST", ""),
+                "metrics": get_metrics()
             }, timeout=5)
         except Exception:
             pass