Separate script for setting modes. (#848)

Author: emeryberger

scalene/scalene_nvidia_gpu.py

@@ -45,9 +45,13 @@ def __del__(self) -> None:
                 file=sys.stderr,
             )
             print(
-                "Run once as Administrator or root (i.e., prefixed with `sudo`) to enable per-process GPU accounting.",
+                "If you have sudo privileges, you can run this command (Linux only) to enable per-process GPU accounting:",
                 file=sys.stderr,
             )
+            print(
+                "  python3 -m scalene.set_nvidia_gpu_modes",
+                file=sys.stderr
+            )
 
     def _set_accounting_mode(self) -> bool:
         """Returns true iff the accounting mode was set already for all GPUs or is now set."""

scalene/set_nvidia_gpu_modes.py

@@ -0,0 +1,46 @@
+import os
+import sys
+import subprocess
+
+def set_nvidia_gpu_modes():
+    import pynvml
+    try:
+        # Initialize NVML
+        pynvml.nvmlInit()
+        
+        # Get the number of GPUs
+        device_count = pynvml.nvmlDeviceGetCount()
+        
+        for i in range(device_count):
+            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+            
+            # Enable persistence mode
+            pynvml.nvmlDeviceSetPersistenceMode(handle, pynvml.NVML_FEATURE_ENABLED)
+            
+            # Enable accounting mode
+            pynvml.nvmlDeviceSetAccountingMode(handle, pynvml.NVML_FEATURE_ENABLED)
+        
+        print("Persistence and accounting mode set for all GPUs.")
+        return True
+    
+    except pynvml.NVMLError as e:
+        print(f"An NVML error occurred: {e}")
+        return False
+    
+    finally:
+        # Shutdown NVML
+        pynvml.nvmlShutdown()
+
+if __name__ == "__main__":
+    # Check if the script is running as root
+    if os.geteuid() != 0:
+        print("This script needs to be run as root. Attempting to rerun with sudo...")
+        try:
+            # Attempt to rerun the script with sudo
+            subprocess.check_call(['sudo', sys.executable] + sys.argv)
+        except subprocess.CalledProcessError as e:
+            print(f"Failed to run as root: {e}")
+            sys.exit(1)
+    else:
+        # Run the function if already root
+        set_nvidia_gpu_modes()