Merge pull request #17 from ryanontheinside/feat/maybites-fix

feat: mps compatability

metrics.py

@@ -1,38 +1,62 @@
-"""Memory measurement utilities for ProfilerX"""
+"""Memory measurement utilities for ProfilerX — supports CUDA, MPS, and CPU-only."""
 import psutil
 import os
 
 _process = psutil.Process(os.getpid())
 
-def _has_cuda():
-    try:
-        import torch
-        return torch.cuda.is_available()
-    except Exception:
-        return False
+# Detect GPU backend once at import time
+_GPU_BACKEND = None  # 'cuda', 'mps', or None
+try:
+    import torch
+    if torch.cuda.is_available():
+        _GPU_BACKEND = 'cuda'
+    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+        _GPU_BACKEND = 'mps'
+except Exception:
+    pass
+
 
 def snapshot_memory():
     """Return (vram_bytes, ram_bytes) current allocation."""
     ram = _process.memory_info().rss
-    if _has_cuda():
-        import torch
-        vram = torch.cuda.memory_allocated()
-    else:
-        vram = 0
+    vram = 0
+    try:
+        if _GPU_BACKEND == 'cuda':
+            import torch
+            vram = torch.cuda.memory_allocated()
+        elif _GPU_BACKEND == 'mps':
+            import torch
+            vram = torch.mps.current_allocated_memory()
+    except Exception:
+        pass
     return vram, ram
 
+
 def reset_peak():
-    """Reset CUDA peak memory stats."""
-    if _has_cuda():
-        import torch
-        torch.cuda.reset_peak_memory_stats()
+    """Reset peak memory stats for the active GPU backend."""
+    try:
+        if _GPU_BACKEND == 'cuda':
+            import torch
+            torch.cuda.reset_peak_memory_stats()
+        elif _GPU_BACKEND == 'mps':
+            # MPS has no peak reset API; this is a no-op.
+            pass
+    except Exception:
+        pass
+
 
 def get_peak():
-    """Return (vram_peak_bytes, ram_bytes). RAM peak is just current RSS (no true peak without sampling)."""
+    """Return (vram_peak_bytes, ram_bytes). RAM peak is just current RSS."""
     ram = _process.memory_info().rss
-    if _has_cuda():
-        import torch
-        vram = torch.cuda.max_memory_allocated()
-    else:
-        vram = 0
+    vram = 0
+    try:
+        if _GPU_BACKEND == 'cuda':
+            import torch
+            vram = torch.cuda.max_memory_allocated()
+        elif _GPU_BACKEND == 'mps':
+            # MPS has no peak tracking; return current allocation as best estimate.
+            import torch
+            vram = torch.mps.current_allocated_memory()
+    except Exception:
+        pass
     return vram, ram