query for libaio package using known package managers (#1250)

* aio: test for libaio with various package managers

* aio: note typical tool used to install libaio package

* setup: abort with error if cannot build requested op

* setup: define op_envvar to return op build environment variable

* setup: call is_compatible once for each op

* setup: only print suggestion to disable op when its envvar not set

* setup: add method to abort from fatal error

* Revert "setup: add method to abort from fatal error"

This reverts commit 0e4cde6b0a650591c3fafface7e27b4efd9aad4f.

* setup: add method to abort from fatal error
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

op_builder/async_io.py

 """
 Copyright 2020 The Microsoft DeepSpeed Team
 """
+import distutils.spawn
+import subprocess
+
 from .builder import OpBuilder
 
 
@@ -50,6 +53,37 @@ class AsyncIOBuilder(OpBuilder):
     def extra_ldflags(self):
         return ['-laio']
 
+    def check_for_libaio_pkg(self):
+        libs = dict(
+            dpkg=["-l",
+                  "libaio-dev",
+                  "apt"],
+            pacman=["-Q",
+                    "libaio",
+                    "pacman"],
+            rpm=["-q",
+                 "libaio-devel",
+                 "yum"],
+        )
+
+        found = False
+        for pkgmgr, data in libs.items():
+            flag, lib, tool = data
+            path = distutils.spawn.find_executable(pkgmgr)
+            if path is not None:
+                cmd = f"{pkgmgr} {flag} {lib}"
+                result = subprocess.Popen(cmd,
+                                          stdout=subprocess.PIPE,
+                                          stderr=subprocess.PIPE,
+                                          shell=True)
+                if result.wait() == 0:
+                    found = True
+                else:
+                    self.warning(
+                        f"{self.NAME}: please install the {lib} package with {tool}")
+                break
+        return found
+
     def is_compatible(self):
         # Check for the existence of libaio by using distutils
         # to compile and link a test program that calls io_submit,
@@ -59,6 +93,14 @@ class AsyncIOBuilder(OpBuilder):
         aio_compatible = self.has_function('io_submit', ('aio', ))
         if not aio_compatible:
             self.warning(
-                f"{self.NAME} requires libaio but it is missing. Can be fixed by: `apt install libaio-dev`."
+                f"{self.NAME} requires the dev libaio .so object and headers but these were not found."
+            )
+
+            # Check for the libaio package via known package managers
+            # to print suggestions on which package to install.
+            self.check_for_libaio_pkg()
+
+            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

op_builder/builder.py

@@ -7,7 +7,13 @@ import time
 import importlib
 from pathlib import Path
 import subprocess
+import shlex
+import shutil
+import tempfile
 import distutils.ccompiler
+import distutils.log
+import distutils.sysconfig
+from distutils.errors import CompileError, LinkError
 from abc import ABC, abstractmethod
 
 YELLOW = '\033[93m'
@@ -161,9 +167,84 @@ class OpBuilder(ABC):
             valid = valid or result.wait() == 0
         return valid
 
-    def has_function(self, funcname, libraries):
-        compiler = distutils.ccompiler.new_compiler()
-        return compiler.has_function(funcname, libraries=libraries)
+    def has_function(self, funcname, libraries, verbose=False):
+        '''
+        Test for existence of a function within a tuple of libraries.
+
+        This is used as a smoke test to check whether a certain library is avaiable.
+        As a test, this creates a simple C program that calls the specified function,
+        and then distutils is used to compile that program and link it with the specified libraries.
+        Returns True if both the compile and link are successful, False otherwise.
+        '''
+        tempdir = None  # we create a temporary directory to hold various files
+        filestderr = None  # handle to open file to which we redirect stderr
+        oldstderr = None  # file descriptor for stderr
+        try:
+            # Echo compile and link commands that are used.
+            if verbose:
+                distutils.log.set_verbosity(1)
+
+            # Create a compiler object.
+            compiler = distutils.ccompiler.new_compiler(verbose=verbose)
+
+            # Configure compiler and linker to build according to Python install.
+            distutils.sysconfig.customize_compiler(compiler)
+
+            # Create a temporary directory to hold test files.
+            tempdir = tempfile.mkdtemp()
+
+            # Define a simple C program that calls the function in question
+            prog = "void %s(void); int main(int argc, char** argv) { %s(); return 0; }" % (
+                funcname,
+                funcname)
+
+            # Write the test program to a file.
+            filename = os.path.join(tempdir, 'test.c')
+            with open(filename, 'w') as f:
+                f.write(prog)
+
+            # Redirect stderr file descriptor to a file to silence compile/link warnings.
+            if not verbose:
+                filestderr = open(os.path.join(tempdir, 'stderr.txt'), 'w')
+                oldstderr = os.dup(sys.stderr.fileno())
+                os.dup2(filestderr.fileno(), sys.stderr.fileno())
+
+            # Attempt to compile the C program into an object file.
+            cflags = shlex.split(os.environ.get('CFLAGS', ""))
+            objs = compiler.compile([filename],
+                                    extra_preargs=self.strip_empty_entries(cflags))
+
+            # Attempt to link the object file into an executable.
+            # Be sure to tack on any libraries that have been specified.
+            ldflags = shlex.split(os.environ.get('LDFLAGS', ""))
+            compiler.link_executable(objs,
+                                     os.path.join(tempdir,
+                                                  'a.out'),
+                                     extra_preargs=self.strip_empty_entries(ldflags),
+                                     libraries=libraries)
+
+            # Compile and link succeeded
+            return True
+
+        except CompileError:
+            return False
+
+        except LinkError:
+            return False
+
+        except:
+            return False
+
+        finally:
+            # Restore stderr file descriptor and close the stderr redirect file.
+            if oldstderr is not None:
+                os.dup2(oldstderr, sys.stderr.fileno())
+            if filestderr is not None:
+                filestderr.close()
+
+            # Delete the temporary directory holding the test program and stderr files.
+            if tempdir is not None:
+                shutil.rmtree(tempdir)
 
     def strip_empty_entries(self, args):
         '''

setup.py

@@ -33,6 +33,15 @@ except ImportError:
 
 from op_builder import ALL_OPS, get_default_compute_capatabilities
 
+RED_START = '\033[31m'
+RED_END = '\033[0m'
+ERROR = f"{RED_START} [ERROR] {RED_END}"
+
+
+def abort(msg):
+    print(f"{ERROR} {msg}")
+    assert False, msg
+
 
 def fetch_requirements(path):
     with open(path, 'r') as fd:
@@ -101,16 +110,29 @@ def command_exists(cmd):
         return result.wait() == 0
 
 
-def op_enabled(op_name):
+def op_envvar(op_name):
     assert hasattr(ALL_OPS[op_name], 'BUILD_VAR'), \
         f"{op_name} is missing BUILD_VAR field"
-    env_var = ALL_OPS[op_name].BUILD_VAR
+    return ALL_OPS[op_name].BUILD_VAR
+
+
+def op_enabled(op_name):
+    env_var = op_envvar(op_name)
     return int(os.environ.get(env_var, BUILD_OP_DEFAULT))
 
 
+compatible_ops = dict.fromkeys(ALL_OPS.keys(), False)
 install_ops = dict.fromkeys(ALL_OPS.keys(), False)
 for op_name, builder in ALL_OPS.items():
     op_compatible = builder.is_compatible()
+    compatible_ops[op_name] = op_compatible
+
+    # If op is requested but not available, throw an error
+    if op_enabled(op_name) and not op_compatible:
+        env_var = op_envvar(op_name)
+        if env_var not in os.environ:
+            builder.warning(f"One can disable {op_name} with {env_var}=0")
+        abort(f"Unable to pre-compile {op_name}")
 
     # If op is compatible update install reqs so it can potentially build/run later
     if op_compatible:
@@ -123,8 +145,6 @@ for op_name, builder in ALL_OPS.items():
         install_ops[op_name] = op_enabled(op_name)
         ext_modules.append(builder.builder())
 
-compatible_ops = {op_name: op.is_compatible() for (op_name, op) in ALL_OPS.items()}
-
 print(f'Install Ops={install_ops}')
 
 # Write out version/git info