Added: Server Health Check Endpoint

client.py

@@ -19,6 +19,9 @@ parser.add_argument(
 parser.add_argument(
     "--duration", type=int, default=10, help="Duration of generated music in seconds"
 )
+parser.add_argument(
+    "--check_health", action='store_true', help="Check server health"
+)
 
 args = parser.parse_args()
 
@@ -36,5 +39,27 @@ def generate_music(server_url, prompts, duration, output_file):
     else:
         print(f"Failed to generate music: {response.status_code}, {response.text}")
 
+def check_server_health(server_url):
+    url = f"{server_url}/health"
+    response = requests.get(url)
+    
+    if response.status_code == 200:
+        health_status = response.json()
+        print("Server Health Check:")
+        print(f"Server Running: {health_status['server_running']}")
+        print(f"Model Loaded: {health_status['model_loaded']}")
+        print(f"CPU Usage: {health_status['cpu_usage_percent']}%")
+        print(f"RAM Usage: {health_status['ram_usage_percent']}%")
+        if 'gpu_memory_allocated' in health_status:
+            gpu_memory_allocated_gb = health_status['gpu_memory_allocated'] / (1024 ** 3)
+            gpu_memory_reserved_gb = health_status['gpu_memory_reserved'] / (1024 ** 3)
+            print(f"GPU Memory Allocated: {gpu_memory_allocated_gb:.2f} GB")
+            print(f"GPU Memory Reserved: {gpu_memory_reserved_gb:.2f} GB")
+    else:
+        print(f"Failed to check server health: {response.status_code}, {response.text}")
+
 if __name__ == "__main__":
-    generate_music(args.server_url, args.prompts, args.duration, args.output_file)
+    if args.check_health:
+        check_server_health(args.server_url)
+    else:
+        generate_music(args.server_url, args.prompts, args.duration, args.output_file)

requirements.txt

@@ -6,3 +6,4 @@ Requests==2.32.3
 scipy==1.13.1
 torch==2.1.0
 uvicorn==0.30.1
+psutil==6.0.0

server.py

@@ -4,12 +4,14 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from typing import List, Optional
 import torch
+from torch.cuda import memory_allocated, memory_reserved
 from audiocraft.models import musicgen
 import numpy as np
 import io
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, JSONResponse
 from scipy.io.wavfile import write as wav_write
 import uvicorn
+import psutil
 
 warnings.simplefilter('ignore')
 
@@ -33,7 +35,12 @@ else:
     args.model_name = f"facebook/{args.model}"
 
 # Load the model with the provided arguments
-musicgen_model = musicgen.MusicGen.get_pretrained(args.model_name, device=args.device)
+try:
+    musicgen_model = musicgen.MusicGen.get_pretrained(args.model_name, device=args.device)
+    model_loaded = True
+except Exception as e:
+    musicgen_model = None
+    model_loaded = False
 
 class MusicRequest(BaseModel):
     prompts: List[str]
@@ -41,6 +48,9 @@ class MusicRequest(BaseModel):
 
 @app.post("/generate_music")
 def generate_music(request: MusicRequest):
+    if not model_loaded:
+        raise HTTPException(status_code=500, detail="Model is not loaded.")
+    
     try:
         musicgen_model.set_generation_params(duration=request.duration)
         result = musicgen_model.generate(request.prompts, progress=False)
@@ -57,5 +67,24 @@ def generate_music(request: MusicRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.get("/health")
+def health_check():
+    cpu_usage = psutil.cpu_percent(interval=1)
+    ram_usage = psutil.virtual_memory().percent
+    stats = {
+        "server_running": True,
+        "model_loaded": model_loaded,
+        "cpu_usage_percent": cpu_usage,
+        "ram_usage_percent": ram_usage
+    }
+    if args.device == "cuda" and torch.cuda.is_available():
+        gpu_memory_allocated = memory_allocated()
+        gpu_memory_reserved = memory_reserved()
+        stats.update({
+            "gpu_memory_allocated": gpu_memory_allocated,
+            "gpu_memory_reserved": gpu_memory_reserved
+        })
+    return JSONResponse(content=stats)
+
 if __name__ == "__main__":
     uvicorn.run(app, host=args.host, port=args.port)