deepspeed

1
# Copyright (c) Microsoft Corporation.
2
# SPDX-License-Identifier: Apache-2.0
3

4
# DeepSpeed Team
5

6
import distutils.spawn
7
import subprocess
8

9
from .builder import OpBuilder
10

11

12
class AsyncIOBuilder(OpBuilder):
13
    BUILD_VAR = "DS_BUILD_AIO"
14
    NAME = "async_io"
15

16
    def __init__(self):
17
        super().__init__(name=self.NAME)
18

19
    def absolute_name(self):
20
        return f'deepspeed.ops.aio.{self.NAME}_op'
21

22
    def sources(self):
23
        return [
24
            'csrc/aio/py_lib/deepspeed_py_copy.cpp', 'csrc/aio/py_lib/py_ds_aio.cpp',
25
            'csrc/aio/py_lib/deepspeed_py_aio.cpp', 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp',
26
            'csrc/aio/py_lib/deepspeed_aio_thread.cpp', 'csrc/aio/common/deepspeed_aio_utils.cpp',
27
            'csrc/aio/common/deepspeed_aio_common.cpp', 'csrc/aio/common/deepspeed_aio_types.cpp',
28
            'csrc/aio/py_lib/deepspeed_pin_tensor.cpp'
29
        ]
30

31
    def include_paths(self):
32
        return ['csrc/aio/py_lib', 'csrc/aio/common']
33

34
    def cxx_args(self):
35
        # -O0 for improved debugging, since performance is bound by I/O
36
        CPU_ARCH = self.cpu_arch()
37
        SIMD_WIDTH = self.simd_width()
38
        import torch  # Keep this import here to avoid errors when building DeepSpeed wheel without torch installed
39
        TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2])
40
        if TORCH_MAJOR >= 2 and TORCH_MINOR >= 1:
41
            CPP_STD = '-std=c++17'
42
        else:
43
            CPP_STD = '-std=c++14'
44
        return [
45
            '-g',
46
            '-Wall',
47
            '-O0',
48
            CPP_STD,
49
            '-shared',
50
            '-fPIC',
51
            '-Wno-reorder',
52
            CPU_ARCH,
53
            '-fopenmp',
54
            SIMD_WIDTH,
55
            '-laio',
56
        ]
57

58
    def extra_ldflags(self):
59
        return ['-laio']
60

61
    def check_for_libaio_pkg(self):
62
        libs = dict(
63
            dpkg=["-l", "libaio-dev", "apt"],
64
            pacman=["-Q", "libaio", "pacman"],
65
            rpm=["-q", "libaio-devel", "yum"],
66
        )
67

68
        found = False
69
        for pkgmgr, data in libs.items():
70
            flag, lib, tool = data
71
            path = distutils.spawn.find_executable(pkgmgr)
72
            if path is not None:
73
                cmd = f"{pkgmgr} {flag} {lib}"
74
                result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
75
                if result.wait() == 0:
76
                    found = True
77
                else:
78
                    self.warning(f"{self.NAME}: please install the {lib} package with {tool}")
79
                break
80
        return found
81

82
    def is_compatible(self, verbose=True):
83
        # Check for the existence of libaio by using distutils
84
        # to compile and link a test program that calls io_submit,
85
        # which is a function provided by libaio that is used in the async_io op.
86
        # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS
87
        # respectively to specify the directories for libaio.h and libaio.so.
88
        aio_compatible = self.has_function('io_pgetevents', ('aio', ))
89
        if verbose and not aio_compatible:
90
            self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.")
91

92
            # Check for the libaio package via known package managers
93
            # to print suggestions on which package to install.
94
            self.check_for_libaio_pkg()
95

96
            self.warning(
97
                "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found."
98
            )
99
        return super().is_compatible(verbose) and aio_compatible
100