Add docstring for auto accelerator (#1956)

Signed-off-by: yiliu30 <[email protected]>

Author: yiliu30

.azure-pipelines/scripts/codeScan/pydocstyle/scan_path.txt

@@ -15,15 +15,15 @@
 /neural-compressor/neural_compressor/strategy
 /neural-compressor/neural_compressor/training.py
 /neural-compressor/neural_compressor/utils
-/neural_compressor/torch/algorithms/layer_wise
-/neural_compressor/torch/algorithms/mixed_precision
-/neural_compressor/torch/algorithms/mx_quant
-/neural_compressor/torch/algorithms/pt2e_quant
+/neural-compressor/neural_compressor/common
+/neural-compressor/neural_compressor/tensorflow
+/neural-compressor/neural_compressor/torch/algorithms/layer_wise
+/neural-compressor/neural_compressor/torch/algorithms/mixed_precision
+/neural-compressor/neural_compressor/torch/algorithms/mx_quant
+/neural-compressor/neural_compressor/torch/algorithms/pt2e_quant
 /neural-compressor/neural_compressor/torch/algorithms/smooth_quant
 /neural-compressor/neural_compressor/torch/algorithms/static_quant
-/neural_compressor/torch/algorithms/weight_only
-/neural_compressor/torch/export
-/neural_compressor/torch/quantization
-/neural_compressor/torch/utils
-/neural_compressor/common
-/neural_compressor/tensorflow
+/neural-compressor/neural_compressor/torch/algorithms/weight_only
+/neural-compressor/neural_compressor/torch/export
+/neural-compressor/neural_compressor/torch/quantization
+/neural-compressor/neural_compressor/torch/utils

neural_compressor/common/benchmark.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+"""Benchmark API for Intel Neural Compressor."""
 
 import argparse
 import os
@@ -242,7 +243,9 @@ def get_numa_node(core_list, reversed_numa_info):
 
 
 def set_cores_for_instance(args, numa_info):
-    """All use cases are listed below:
+    """Set cores for each instance based on the input args.
+
+    All use cases are listed below:
         Params: a=num_instance; b=num_cores_per_instance; c=cores;
             - no a, b, c: a=1, c=numa:0
             - no a, b: a=1, c=c
@@ -357,6 +360,7 @@ def generate_prefix(args, core_list):
     Args:
         args (argparse): arguments for setting different configurations
         core_list: ["node_index", "cpu_index", num_cpu]
+
     Returns:
         command_prefix (str): command_prefix with specific core list for Linux or Windows.
     """

neural_compressor/common/utils/utility.py

@@ -328,6 +328,8 @@ def wrapper(*args, **kwargs):
 
 
 class ProcessorType(enum.Enum):
+    """The processor type."""
+
     Client = "Client"
     Server = "Server"
 

neural_compressor/torch/utils/auto_accelerator.py

@@ -19,6 +19,7 @@
 
 # NOTICE: The design adapted from:
 # https://github.com/microsoft/DeepSpeed/blob/master/accelerator/abstract_accelerator.py.
+"""Auto Accelerator Module."""
 
 
 # To keep it simply, only add the APIs we need.
@@ -40,6 +41,8 @@
 
 
 class AcceleratorRegistry:
+    """Accelerator Registry."""
+
     registered_accelerators = {}
 
     @classmethod
@@ -94,171 +97,253 @@ class CUDA_Accelerator:
         name: the accelerator name.
         priority: the priority of the accelerator. A larger number indicates a higher priority,
     """
-
     return accelerator_registry.register_accelerator_impl(name=name, priority=priority)
 
 
 class Auto_Accelerator(ABC):  # pragma: no cover
+    """Auto Accelerator Base class."""
+
     @classmethod
     @abstractmethod
     def is_available(cls) -> bool:
+        """Check if the accelerator is available."""
         pass
 
     @abstractmethod
     def name(self) -> str:
+        """Get the accelerator name."""
         pass
 
     @abstractmethod
     def device_name(self, device_indx) -> str:
+        """Get the device name."""
         pass
 
     @abstractmethod
     def set_device(self, device_index):
+        """Set the device."""
         pass
 
     @abstractmethod
     def current_device(self):
+        """Get the current device."""
         pass
 
     @abstractmethod
     def current_device_name(self):
+        """Get the current device name."""
         pass
 
     @abstractmethod
     def device(self, device_index=None):
+        """Get the device."""
         pass
 
     @abstractmethod
     def empty_cache(self):
+        """Empty the cache."""
         pass
 
     @abstractmethod
     def synchronize(self):
+        """Synchronize the accelerator."""
         pass
 
     def mark_step(self):
+        """Trigger graph to run."""
         pass
 
 
 @register_accelerator(name="cpu", priority=PRIORITY_CPU)
 class CPU_Accelerator(Auto_Accelerator):
+    """CPU Accelerator."""
+
     def __init__(self) -> None:
+        """Initialize CPU Accelerator."""
         self._name = "cpu"
 
     def name(self) -> str:
+        """Get the accelerator name."""
         return self._name
 
     @classmethod
     def is_available(cls) -> bool:
+        """Always return True."""
         return True
 
     def device_name(self, device_indx) -> str:
+        """Get the device name."""
         return "cpu"
 
     def set_device(self, device_index):
+        """Do nothing."""
         pass
 
     def current_device(self):
+        """Get the current device."""
         return "cpu"
 
     def current_device_name(self):
+        """Get the current device name."""
         return "cpu"
 
     def device(self, device_index=None):
+        """Do nothing."""
         pass
 
     def empty_cache(self):
+        """Do nothing."""
         pass
 
     def synchronize(self):
+        """Do nothing."""
         pass
 
 
 @register_accelerator(name="cuda", priority=PRIORITY_CUDA)
 class CUDA_Accelerator(Auto_Accelerator):  # pragma: no cover
+    """CUDA Accelerator."""
+
     def __init__(self) -> None:
+        """Initialize CUDA Accelerator."""
         self._name = "cuda"
 
     def name(self) -> str:
+        """Get the accelerator name."""
         return self._name
 
     @classmethod
     def is_available(cls) -> bool:
+        """Check if the 'cuda' device is available."""
         return torch.cuda.is_available()
 
     def device_name(self, device_indx) -> str:
+        """Returns the name of the 'cuda' device with the given index."""
         if device_indx is None:
             return "cuda"
         return f"cuda:{device_indx}"
 
     def synchronize(self):
+        """Synchronizes the 'cuda' device."""
         return torch.cuda.synchronize()
 
     def set_device(self, device_index):
+        """Sets the current 'cuda' device to the one with the given index."""
         return torch.cuda.set_device(device_index)
 
     def current_device(self):
+        """Returns the index of the current 'cuda' device."""
         return torch.cuda.current_device()
 
     def current_device_name(self):
+        """Returns the name of the current 'cuda' device."""
         return "cuda:{}".format(torch.cuda.current_device())
 
     def device(self, device_index=None):
+        """Returns a torch.device object for the 'cuda' device with the given index."""
         return torch.cuda.device(device_index)
 
     def empty_cache(self):
+        """Empties the cuda cache."""
         return torch.cuda.empty_cache()
 
 
 @register_accelerator(name="xpu", priority=PRIORITY_XPU)
 class XPU_Accelerator(Auto_Accelerator):  # pragma: no cover
+    """XPU Accelerator."""
+
     def __init__(self) -> None:
+        """Initialize XPU Accelerator."""
         self._name = "xpu"
 
     def name(self) -> str:
+        """Get the accelerator name."""
         return self._name
 
     @classmethod
     def is_available(cls) -> bool:
+        """Checks if the 'xpu' device is available.
+
+        Returns:
+            bool: True if the 'xpu' device is available, False otherwise.
+        """
         if hasattr(torch, "xpu") and torch.xpu.is_available():
             return True
         else:
             return False
 
     def device_name(self, device_indx) -> str:
+        """Returns the name of the 'xpu' device with the given index.
+
+        Args:
+            device_indx (int): The index of the 'xpu' device.
+
+        Returns:
+            str: The name of the 'xpu' device.
+        """
         if device_indx is None:
             return "xpu"
         return f"xpu:{device_indx}"
 
     def synchronize(self):
+        """Synchronizes the 'xpu' device."""
         return torch.xpu.synchronize()
 
     def set_device(self, device_index):
+        """Sets the current 'xpu' device to the one with the given index.
+
+        Args:
+            device_index (int): The index of the 'xpu' device.
+        """
         return torch.xpu.set_device(device_index)
 
     def current_device(self):
+        """Returns the index of the current 'xpu' device.
+
+        Returns:
+            int: The index of the current 'xpu' device.
+        """
         return torch.xpu.current_device()
 
     def current_device_name(self):
+        """Returns the name of the current 'xpu' device.
+
+        Returns:
+            str: The name of the current 'xpu' device.
+        """
         return "xpu:{}".format(torch.xpu.current_device())
 
     def device(self, device_index=None):
+        """Returns a torch.device object for the 'xpu' device with the given index.
+
+        Args:
+            device_index (int, optional): The index of the 'xpu' device. Defaults to None.
+
+        Returns:
+            torch.device: The torch.device object for the 'xpu' device.
+        """
         return torch.xpu.device(device_index)
 
     def empty_cache(self):
+        """Empties the xpu cache."""
         return torch.xpu.empty_cache()
 
 
 @register_accelerator(name="hpu", priority=PRIORITY_HPU)
 class HPU_Accelerator(Auto_Accelerator):  # pragma: no cover
+    """HPU Accelerator."""
+
     def __init__(self) -> None:
+        """Initialize HPU Accelerator."""
         self._name = "hpu"
 
     def name(self) -> str:
+        """Get the accelerator name."""
         return self._name
 
     @classmethod
     def is_available(cls) -> bool:
+        """Checks if the 'hpu' device is available."""
         from .environ import is_hpex_available
 
         if is_hpex_available():
@@ -267,43 +352,54 @@ def is_available(cls) -> bool:
             return False
 
     def device_name(self, device_indx) -> str:
+        """Returns the name of the 'hpu' device with the given index."""
         if device_indx is None:
             return "hpu"
         return f"hpu:{device_indx}"
 
     def synchronize(self):
+        """Synchronizes the 'hpu' device."""
         return torch.hpu.synchronize()
 
     def set_device(self, device_index):
+        """Sets the current 'hpu' device to the one with the given index."""
         try:
             torch.hpu.set_device(device_index)
         except Exception as e:
             logger.warning(e)
 
     def current_device(self):
+        """Returns the index of the current 'hpu' device."""
         return torch.hpu.current_device()
 
     def current_device_name(self):
+        """Returns the name of the current 'hpu' device."""
         return "hpu:{}".format(torch.hpu.current_device())
 
     def device(self, device_index=None):
+        """Returns a torch.device object for the 'hpu' device with the given index."""
         return torch.hpu.device(device_index)
 
     def empty_cache(self):
+        """Empties the hpu cache."""
         try:
             torch.hpu.empty_cache()
         except Exception as e:
             logger.warning(e)
 
     def mark_step(self):
+        """Trigger graph to run."""
         return htcore.mark_step()
 
 
 def auto_detect_accelerator(device_name="auto") -> Auto_Accelerator:
-    # Force use the cpu on node has both cpu and gpu: `FORCE_DEVICE=cpu` python main.py ...
-    # The `FORCE_DEVICE` is case insensitive.
-    # The environment variable `FORCE_DEVICE` has higher priority than the `device_name`.
-    # TODO: refine the docs and logic later
+    """Automatically detects and selects the appropriate accelerator.
+
+    Force use the cpu on node has both cpu and gpu: `FORCE_DEVICE=cpu` python main.py ...
+    The `FORCE_DEVICE` is case insensitive.
+    The environment variable `FORCE_DEVICE` has higher priority than the `device_name`.
+    TODO: refine the docs and logic later
+    """
     # 1. Get the device setting from environment variable `FORCE_DEVICE`.
     FORCE_DEVICE = os.environ.get("FORCE_DEVICE", None)
     if FORCE_DEVICE: