Addonに適用

server/main.py

@@ -7,6 +7,7 @@ GPU-accelerated face detection using ONNX Runtime.
 
 import os
 import sys
+import platform
 import threading
 import uuid
 import queue
@@ -27,6 +28,9 @@ from server.detector import YOLOFaceDetector, get_detector
 
 app = FastAPI(title="Face Mask Inference Server")
 
+# GPU status cache
+_gpu_status_cache = None
+
 # Task storage
 class TaskStatus:
     PENDING = "pending"
@@ -146,9 +150,146 @@ def process_video_task(task_id: str, req: GenerateRequest):
         if task_id in cancel_events:
             del cancel_events[task_id]
 
+def check_gpu_available() -> dict:
+    """
+    Check if GPU is available for inference.
+
+    Returns a dict with GPU information:
+    {
+        "available": bool,
+        "device_name": str or None,
+        "device_count": int,
+        "rocm_version": str or None
+    }
+    """
+    global _gpu_status_cache
+
+    # Return cached result if available
+    if _gpu_status_cache is not None:
+        return _gpu_status_cache
+
+    result = {
+        "available": False,
+        "device_name": None,
+        "device_count": 0,
+        "rocm_version": None
+    }
+
+    try:
+        import torch
+
+        result["available"] = torch.cuda.is_available()
+
+        if result["available"]:
+            result["device_count"] = torch.cuda.device_count()
+            if result["device_count"] > 0:
+                result["device_name"] = torch.cuda.get_device_name(0)
+
+            if hasattr(torch.version, 'hip'):
+                result["rocm_version"] = torch.version.hip
+
+    except Exception as e:
+        print(f"[FaceMask] Warning: GPU detection failed: {e}")
+        result["available"] = False
+
+    # Cache the result
+    _gpu_status_cache = result
+
+    return result
+
+
+def log_startup_diagnostics():
+    """Log diagnostic information about the environment and GPU."""
+    print("=" * 70)
+    print("[FaceMask Server] Startup Diagnostics")
+    print("=" * 70)
+
+    # Python Environment
+    print("\n[Python Environment]")
+    print(f"  Python Version: {sys.version.split()[0]}")
+    print(f"  Python Executable: {sys.executable}")
+    print(f"  Platform: {platform.platform()}")
+    print(f"  Working Directory: {os.getcwd()}")
+
+    # Check if in venv
+    in_venv = sys.prefix != sys.base_prefix
+    print(f"  Virtual Environment: {'Yes' if in_venv else 'No'}")
+    if in_venv:
+        print(f"    venv path: {sys.prefix}")
+
+    # ROCm Environment Variables
+    print("\n[ROCm Environment Variables]")
+    rocm_vars = [
+        'ROCM_PATH',
+        'HSA_OVERRIDE_GFX_VERSION',
+        'PYTORCH_ROCM_ARCH',
+        'ROCBLAS_TENSILE_LIBPATH',
+        'LD_LIBRARY_PATH'
+    ]
+
+    for var in rocm_vars:
+        value = os.environ.get(var)
+        if value:
+            # Truncate very long values
+            if len(value) > 200:
+                display_value = value[:200] + "... (truncated)"
+            else:
+                display_value = value
+            print(f"  {var}: {display_value}")
+        else:
+            print(f"  {var}: (not set)")
+
+    # GPU Detection
+    print("\n[GPU Detection]")
+    try:
+        import torch
+
+        cuda_available = torch.cuda.is_available()
+        print(f"  torch.cuda.is_available(): {cuda_available}")
+
+        if cuda_available:
+            device_count = torch.cuda.device_count()
+            print(f"  GPU Device Count: {device_count}")
+
+            if device_count > 0:
+                device_name = torch.cuda.get_device_name(0)
+                print(f"  GPU Device 0: {device_name}")
+
+                # ROCm version
+                if hasattr(torch.version, 'hip'):
+                    print(f"  ROCm Version (HIP): {torch.version.hip}")
+
+                # CUDA version (might be emulated by ROCm)
+                if torch.version.cuda:
+                    print(f"  CUDA Version: {torch.version.cuda}")
+        else:
+            print("  WARNING: GPU not detected!")
+            print("  Server will use CPU for inference (slower)")
+            print("  Troubleshooting:")
+            print("    - Check ROCm environment variables above")
+            print("    - Run: python -c 'import torch; print(torch.cuda.is_available())'")
+
+    except ImportError as e:
+        print(f"  ERROR: Cannot import torch: {e}")
+        print("  PyTorch must be installed for inference")
+    except Exception as e:
+        print(f"  ERROR during GPU detection: {e}")
+
+    print("=" * 70)
+    print()
+
+
 @app.get("/status")
 def get_status():
-    return {"status": "running", "gpu_available": True} # TODO: check GPU
+    gpu_info = check_gpu_available()
+
+    return {
+        "status": "running",
+        "gpu_available": gpu_info["available"],
+        "gpu_device": gpu_info["device_name"],
+        "gpu_count": gpu_info["device_count"],
+        "rocm_version": gpu_info["rocm_version"]
+    }
 
 @app.post("/generate", response_model=Task)
 def generate_mask_endpoint(req: GenerateRequest, background_tasks: BackgroundTasks):
@@ -177,4 +318,5 @@ def cancel_task(task_id: str):
     return {"message": "Cancellation requested"}
 
 if __name__ == "__main__":
+    log_startup_diagnostics()
     uvicorn.run(app, host="127.0.0.1", port=8181)