Control ds_report output

bin/ds_report

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-from deepspeed.env_report import main
+from deepspeed.env_report import cli_main
 
 if __name__ == '__main__':
-    main()
+    cli_main()

deepspeed/env_report.py

@@ -1,6 +1,7 @@
 import torch
 import deepspeed
 import subprocess
+import argparse
 from .ops.op_builder import ALL_OPS
 from .git_version_info import installed_ops, torch_info
 from .ops import __compatible_ops__ as compatible_ops
@@ -20,7 +21,7 @@
 warning = f"{YELLOW}[WARNING]{END}"
 
 
-def op_report():
+def op_report(verbose=True):
     max_dots = 23
     max_dots2 = 11
     h = ["op name", "installed", "compatible"]
@@ -43,7 +44,7 @@ def op_report():
     no = f"{YELLOW}[NO]{END}"
     for op_name, builder in ALL_OPS.items():
         dots = "." * (max_dots - len(op_name))
-        is_compatible = OKAY if builder.is_compatible() else no
+        is_compatible = OKAY if builder.is_compatible(verbose) else no
         is_installed = installed if installed_ops[op_name] else no
         dots2 = '.' * ((len(h[1]) + (max_dots2 - len(h[1]))) -
                        (len(is_installed) - color_len))
@@ -100,10 +101,32 @@ def debug_report():
         print(name, "." * (max_dots - len(name)), value)
 
 
-def main():
-    op_report()
+def parse_arguments():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--hide_operator_status',
+        action='store_true',
+        help=
+        'Suppress display of installation and compatiblity statuses of DeepSpeed operators. '
+    )
+    parser.add_argument('--hide_errors_and_warnings',
+                        action='store_true',
+                        help='Suppress warning and error messages.')
+    args = parser.parse_args()
+    return args
+
+
+def main(hide_operator_status=False, hide_errors_and_warnings=False):
+    if not hide_operator_status:
+        op_report(verbose=not hide_errors_and_warnings)
     debug_report()
 
 
+def cli_main():
+    args = parse_arguments()
+    main(hide_operator_status=args.hide_operator_status,
+         hide_errors_and_warnings=args.hide_errors_and_warnings)
+
+
 if __name__ == "__main__":
     main()

op_builder/async_io.py

@@ -84,14 +84,14 @@ def check_for_libaio_pkg(self):
                 break
         return found
 
-    def is_compatible(self):
+    def is_compatible(self, verbose=True):
         # Check for the existence of libaio by using distutils
         # to compile and link a test program that calls io_submit,
         # which is a function provided by libaio that is used in the async_io op.
         # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS
         # respectively to specify the directories for libaio.h and libaio.so.
         aio_compatible = self.has_function('io_submit', ('aio', ))
-        if not aio_compatible:
+        if verbose and not aio_compatible:
             self.warning(
                 f"{self.NAME} requires the dev libaio .so object and headers but these were not found."
             )
@@ -103,4 +103,4 @@ def is_compatible(self):
             self.warning(
                 "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found."
             )
-        return super().is_compatible() and aio_compatible
+        return super().is_compatible(verbose) and aio_compatible

op_builder/builder.py

@@ -153,7 +153,7 @@ def cxx_args(self):
         '''
         return []
 
-    def is_compatible(self):
+    def is_compatible(self, verbose=True):
         '''
         Check if all non-python dependencies are satisfied to build this op
         '''
@@ -370,7 +370,7 @@ def load(self, verbose=True):
             return self.jit_load(verbose)
 
     def jit_load(self, verbose=True):
-        if not self.is_compatible():
+        if not self.is_compatible(verbose):
             raise RuntimeError(
                 f"Unable to JIT load the {self.name} op due to it not being compatible due to hardware/software issue."
             )
@@ -482,8 +482,8 @@ def version_dependent_macros(self):
             version_ge_1_5 = ['-DVERSION_GE_1_5']
         return version_ge_1_1 + version_ge_1_3 + version_ge_1_5
 
-    def is_compatible(self):
-        return super().is_compatible()
+    def is_compatible(self, verbose=True):
+        return super().is_compatible(verbose)
 
     def builder(self):
         from torch.utils.cpp_extension import CUDAExtension

op_builder/cpu_adagrad.py

@@ -14,7 +14,7 @@ class CPUAdagradBuilder(CUDAOpBuilder):
     def __init__(self):
         super().__init__(name=self.NAME)
 
-    def is_compatible(self):
+    def is_compatible(self, verbose=True):
         # Disable on Windows.
         return sys.platform != "win32"
 

op_builder/cpu_adam.py

@@ -14,7 +14,7 @@ class CPUAdamBuilder(CUDAOpBuilder):
     def __init__(self):
         super().__init__(name=self.NAME)
 
-    def is_compatible(self):
+    def is_compatible(self, verbose=True):
         # Disable on Windows.
         return sys.platform != "win32"
 

op_builder/sparse_attn.py

@@ -21,7 +21,7 @@ def sources(self):
     def cxx_args(self):
         return ['-O2', '-fopenmp']
 
-    def is_compatible(self):
+    def is_compatible(self, verbose=True):
         # Check to see if llvm and cmake are installed since they are dependencies
         #required_commands = ['llvm-config|llvm-config-9', 'cmake']
         #command_status = list(map(self.command_exists, required_commands))
@@ -52,4 +52,4 @@ def is_compatible(self):
                 f'{self.NAME} requires a torch version >= 1.5 but detected {TORCH_MAJOR}.{TORCH_MINOR}'
             )
 
-        return super().is_compatible() and torch_compatible and cuda_compatible
+        return super().is_compatible(verbose) and torch_compatible and cuda_compatible