Do not include power-related results unless requested (#54)

Querying power status (current draw and limit -- may be cached?) is
quite expensive. Thus we can skip querying power status unless
it is explicitly requested. This leads to 50% less CPU usage in my case.

The default behavior of `new_query()` still includes power information.

TODO: Add unit test case for the behavior

gpustat/__main__.py

@@ -4,19 +4,27 @@
 
 import sys
 import time
+from functools import partial
 
 from blessings import Terminal
 
 from gpustat import __version__
 from .core import GPUStatCollection
 
 
-def print_gpustat(json=False, debug=False, **kwargs):
+def print_gpustat(json=False, debug=False,
+                  gpustat_factory=GPUStatCollection.new_query,
+                  **kwargs):
     '''
     Display the GPU query results into standard output.
     '''
+
+    if not kwargs.get('show_power', None):
+        gpustat_factory = partial(gpustat_factory, include_power=False)
+
     try:
-        gpu_stats = GPUStatCollection.new_query()
+        gpu_stats = gpustat_factory()
+        assert isinstance(gpu_stats, GPUStatCollection)
     except Exception as e:
         sys.stderr.write('Error on querying NVIDIA devices.'
                          ' Use --debug flag for details\n')

gpustat/core.py

@@ -200,7 +200,7 @@ def _repr(v, none_value='??'):
             "%(CTemp)s{entry[temperature.gpu]:>3}'C%(C0)s, " \
             "%(CUtil)s{entry[utilization.gpu]:>3} %%%(C0)s"
 
-        if show_power:
+        if show_power and 'power.draw' in self.entry:
             reps += ",  %(CPowU)s{entry[power.draw]:>3}%(C0)s "
             if show_power is True or 'limit' in show_power:
                 reps += "/ %(CPowL)s{entry[enforced.power.limit]:>3}%(C0)s "
@@ -267,7 +267,7 @@ def __init__(self, gpu_list, driver_version=None):
         self.driver_version = driver_version
 
     @staticmethod
-    def new_query():
+    def new_query(include_power=True):
         """Query the information of all the GPUs on local machine"""
 
         if N is None:
@@ -321,15 +321,16 @@ def get_process_info(nv_process):
             except N.NVMLError:
                 utilization = None  # Not supported
 
-            try:
-                power = N.nvmlDeviceGetPowerUsage(handle)
-            except N.NVMLError:
-                power = None
+            if include_power:
+                try:
+                    power = N.nvmlDeviceGetPowerUsage(handle)
+                except N.NVMLError:
+                    power = None
 
-            try:
-                power_limit = N.nvmlDeviceGetEnforcedPowerLimit(handle)
-            except N.NVMLError:
-                power_limit = None
+                try:
+                    power_limit = N.nvmlDeviceGetEnforcedPowerLimit(handle)
+                except N.NVMLError:
+                    power_limit = None
 
             try:
                 nv_comp_processes = \
@@ -366,14 +367,21 @@ def get_process_info(nv_process):
                 'name': name,
                 'temperature.gpu': temperature,
                 'utilization.gpu': utilization.gpu if utilization else None,
-                'power.draw': power // 1000 if power is not None else None,
-                'enforced.power.limit': power_limit // 1000
-                if power_limit is not None else None,
                 # Convert bytes into MBytes
                 'memory.used': memory.used // MB if memory else None,
                 'memory.total': memory.total // MB if memory else None,
                 'processes': processes,
             }
+
+            # optional entries
+            if include_power:
+                gpu_info.update({
+                    'power.draw': power // 1000
+                        if power is not None else None,  # noqa
+                    'enforced.power.limit': power_limit // 1000
+                        if power_limit is not None else None,  # noqa
+                })
+
             return gpu_info
 
         # 1. get the list of gpu and status
@@ -490,9 +498,9 @@ def date_handler(obj):
         fp.flush()
 
 
-def new_query():
+def new_query(*args, **kwargs):
     '''
     Obtain a new GPUStatCollection instance by querying nvidia-smi
     to get the list of GPUs and running process information.
     '''
-    return GPUStatCollection.new_query()
+    return GPUStatCollection.new_query(*args, **kwargs)