Refactor

nvidia_pstate/__init__.py

@@ -1 +1 @@
-from nvidia_pstate.pstate import set_pstate, set_pstate_low, set_pstate_high
+from .pstate import set_pstate, set_pstate_low, set_pstate_high

nvidia_pstate/pstate.py

@@ -1,79 +1,27 @@
 import atexit
-import ctypes
-import os
 import sys
 
-from .nvapi import NvAPI_Initialize, NvAPI_Unload, NvAPI_EnumPhysicalGPUs, NvAPI_GPU_SetForcePstate
+# Try to use NvAPI
+try:
+  from .nvapi import NvAPI_Initialize, NvAPI_Unload
 
+  # Initialize NvAPI
+  if NvAPI_Initialize() != 0:
+    raise Exception("Failed to initialize NvAPI")
 
-# Initialize NvAPI
-if NvAPI_Initialize() != 0:
-  print("Failed to initialize NvAPI", file=sys.stderr)
+  # Unload NvAPI at exit
+  atexit.register(NvAPI_Unload)
 
-def NvAPI_Unload_atexit():
-  # Unload NvAPI
-  if NvAPI_Unload() != 0:
-    print("Failed to unload NvAPI")
-    return
+  # Import functions
+  from .pstate_nvapi import set_pstate, set_pstate_high, set_pstate_low
+except Exception as e:
+  # Print exception
+  print(e, file=sys.stderr)
 
-# Unload NvAPI at exit
-atexit.register(NvAPI_Unload_atexit)
+  # Define noop function
+  noop = lambda *args, **kwargs: None
 
-
-def set_pstate(ids, pstate, *, silent = False):
-  # Default ids: all or limited by CUDA_VISIBLE_DEVICES
-  if ids is None:
-    ids = []
-
-  # Default pstate: let driver decide
-  if pstate is None:
-    pstate = 16
-
-  # Array to hold GPU handles
-  gpu_array = (ctypes.c_void_p * 64)()
-
-  # Integer to hold GPU count
-  gpu_count = ctypes.pointer(ctypes.c_int32())
-
-  # Enumerate GPUs
-  if NvAPI_EnumPhysicalGPUs(gpu_array, gpu_count) != 0:
-    print("Failed to enumerate GPUs", file=sys.stderr)
-    return
-
-  # GPU count as int
-  gpu_count = gpu_count.contents.value
-
-  # Function to actually set performance state
-  def _set_performance_state(gpu_id, pstate):
-    if not (0 <= gpu_id < gpu_count):
-      if not silent:
-        print(f"Invalid GPU ID: {gpu_id}", file=sys.stderr)
-      return
-
-    if NvAPI_GPU_SetForcePstate(gpu_array[gpu_id], pstate, 2) == 0:
-      if not silent:
-        print(f"Performance state has been set successfully for gpu #{gpu_id}", file=sys.stderr)
-    else:
-      if not silent:
-        print(f"Failed to set performance state for gpu #{gpu_id}", file=sys.stderr)
-
-  # If GPUs is not specified, try to use CUDA_VISIBLE_DEVICES
-  if len(ids) == 0:
-    visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
-    if visible_devices:
-      for device in visible_devices.split(","):
-        ids.append(int(device))
-
-  # Set performance state for specified GPU or all GPUs
-  if len(ids) == 0:
-    for i in range(gpu_count):
-      _set_performance_state(i, pstate)
-  else:
-    for i in ids:
-      _set_performance_state(i, pstate)
-
-def set_pstate_low():
-  set_pstate([], int(os.getenv("NVIDIA_PSTATE_LOW", "8")), silent=True)
-
-def set_pstate_high():
-  set_pstate([], int(os.getenv("NVIDIA_PSTATE_HIGH", "16")), silent=True)
+  # Use noop function instead of real functions
+  set_pstate = noop
+  set_pstate_high = noop
+  set_pstate_low = noop

nvidia_pstate/pstate_nvapi.py

@@ -0,0 +1,64 @@
+import ctypes
+import os
+import sys
+
+from .nvapi import NvAPI_EnumPhysicalGPUs, NvAPI_GPU_SetForcePstate
+
+
+def set_pstate(ids, pstate, *, silent = False):
+  # Default ids: all or limited by CUDA_VISIBLE_DEVICES
+  if ids is None:
+    ids = []
+
+  # Default pstate: let driver decide
+  if pstate is None:
+    pstate = 16
+
+  # Array to hold GPU handles
+  gpu_array = (ctypes.c_void_p * 64)()
+
+  # Integer to hold GPU count
+  gpu_count = ctypes.pointer(ctypes.c_int32())
+
+  # Enumerate GPUs
+  if NvAPI_EnumPhysicalGPUs(gpu_array, gpu_count) != 0:
+    print("Failed to enumerate GPUs", file=sys.stderr)
+    return
+
+  # GPU count as int
+  gpu_count = gpu_count.contents.value
+
+  # Function to actually set performance state
+  def _set_performance_state(gpu_id, pstate):
+    if not (0 <= gpu_id < gpu_count):
+      if not silent:
+        print(f"Invalid GPU ID: {gpu_id}", file=sys.stderr)
+      return
+
+    if NvAPI_GPU_SetForcePstate(gpu_array[gpu_id], pstate, 2) == 0:
+      if not silent:
+        print(f"Performance state has been set successfully for gpu #{gpu_id}", file=sys.stderr)
+    else:
+      if not silent:
+        print(f"Failed to set performance state for gpu #{gpu_id}", file=sys.stderr)
+
+  # If GPUs is not specified, try to use CUDA_VISIBLE_DEVICES
+  if len(ids) == 0:
+    visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
+    if visible_devices:
+      for device in visible_devices.split(","):
+        ids.append(int(device))
+
+  # Set performance state for specified GPU or all GPUs
+  if len(ids) == 0:
+    for i in range(gpu_count):
+      _set_performance_state(i, pstate)
+  else:
+    for i in ids:
+      _set_performance_state(i, pstate)
+
+def set_pstate_high():
+  set_pstate([], int(os.getenv("NVIDIA_PSTATE_HIGH", "16")), silent=True)
+
+def set_pstate_low():
+  set_pstate([], int(os.getenv("NVIDIA_PSTATE_LOW", "8")), silent=True)

nvidia_pstate/scripts/nvidia_pstate.py

@@ -1,34 +1,41 @@
 import argparse
+import importlib.metadata
 
-from nvidia_pstate import set_pstate
+import nvidia_pstate
 
 
-def list_command(args):
+# Function to list available performance states
+def query_pstate(args):
   # TODO
   print("Does not work for now")
 
-# Function to set performance state
-def set_command(args):
-  set_pstate(args.id, args.performance_state)
+# Function to specify performance state
+def set_pstate(args):
+  nvidia_pstate.set_pstate(args.id, args.perf_state)
 
-def main():
-  # Create the main argument parser with a description
-  parser = argparse.ArgumentParser(description="CLI utility for managing performance states of NVIDIA GPUs.")
 
-  # Create subparsers for different commands
-  subparsers = parser.add_subparsers(title="Commands", dest="command", required=True)
+def main():
+  # Determine program version
+  try:
+    __version__ = importlib.metadata.version("nvidia-pstate")
+  except:
+    __version__ = "unknown"
 
-  # "list" command parser
-  list_parser = subparsers.add_parser("list-pstates", help="List available performance states of GPU")
-  list_parser.add_argument("-i", "--id", type=int, nargs="+", help="Target a specific GPU")
-  list_parser.set_defaults(func=list_command)
+  # Create the main argument parser with a description
+  parser = argparse.ArgumentParser(description="CLI utility for managing performance states of NVIDIA GPUs")
 
-  # "set" command parser
-  set_parser = subparsers.add_parser("set-pstate", help="Set performance state for GPU")
-  set_parser.add_argument("-i", "--id", type=int, nargs="+", help="Target a specific GPU")
-  set_parser.add_argument("-ps", "--performance-state", type=int, required=True, help="Specifies performance state.")
-  set_parser.set_defaults(func=set_command)
+  parser.add_argument("-i", "--id", type=int, nargs="+", help="target a specific GPU")
+  parser.add_argument("-ps", "--perf-state", type=int, help="specify performance state")
+  parser.add_argument("-q", "--query", action="store_true", help="list available performance states")
+  parser.add_argument("-v", "--version", action="version", version=f"%(prog)s {__version__}")
 
-  # Parse command-line arguments and execute the appropriate function
+  # Parse command-line arguments
   args = parser.parse_args()
-  args.func(args)
+
+  # Execute the appropriate function
+  if args.perf_state is not None:
+    set_pstate(args)
+  elif args.query is not None:
+    query_pstate(args)
+  else:
+    parser.print_help()

pyproject.toml

@@ -9,7 +9,7 @@ requires = [
 [project]
 description = "A library and CLI utilities for managing performance states of NVIDIA GPUs."
 name = "nvidia_pstate"
-version = "1.0.2"
+version = "1.0.3"
 readme = "README.md"
 
 authors = [