中的Try语句(用于mpi4py)

Question

是否有一种方法可以使cimport在Cython中具有相当于Python语句的功能？

就像这样：

try:
    cimport something
except ImportError:
    pass

我需要它来编写一个可以使用或不使用mpi4py编译的Cython扩展。在编译语言中，这是非常标准的，mpi命令可以放在#ifdef和#endif预处理器指令之间。我们如何在Cython中获得相同的结果？

我试过了，但没有用：

try:
    from mpi4py import MPI
    from mpi4py cimport MPI
    from mpi4py.mpi_c cimport *
except ImportError:
    rank = 0
    nb_proc = 1

# solve a incompatibility between openmpi and mpi4py versions
cdef extern from 'mpi-compat.h': pass

does_it_work = 'Not yet'

实际上，如果正确安装了mpi4py，它就工作得很好，但是如果import mpi4py引发ImportError，则Cython文件不会编译，我得到了错误：

Error compiling Cython file:
------------------------------------------------------------
...

try:
    from mpi4py import MPI
    from mpi4py cimport MPI
   ^
------------------------------------------------------------

mod.pyx:4:4: 'mpi4py.pxd' not found

文件setup.py

from setuptools import setup, Extension
from Cython.Distutils import build_ext

import os
here = os.path.abspath(os.path.dirname(__file__))

include_dirs = [here]

try:
    import mpi4py
except ImportError:
    pass
else:
    INCLUDE_MPI = '/usr/lib/openmpi/include'
    include_dirs.extend([
        INCLUDE_MPI,
        mpi4py.get_include()])

name = 'mod'
ext = Extension(
    name,
    include_dirs=include_dirs,
    sources=['mod.pyx'])

setup(name=name,
      cmdclass={"build_ext": build_ext},
      ext_modules=[ext])

Answer 1

以这种方式使用try-catch块是您无法做到的事情。您正在创建的扩展模块必须静态编译，并根据它使用cimport在C级别加载的内容进行链接。尝试捕获块是在导入模块时执行的，而不是在编译时执行的。

另一方面，从理论上讲，您应该能够获得使用Cython的对条件编译的支持所需要的效果。在您的setup.py文件中，您可以检查是否定义了所需的模块，然后定义要传递给Cython编译器的环境变量，这反过来取决于是否存在所需的模块。

在Cython试验之一中有一个如何做到这一点的例子。在那里，它们将包含所需环境变量的字典作为关键字参数pyrex_compile_time_env传递给Cython的Extension类的构造函数，该关键字参数是重命名 to cython_compile_time_env，Cython.Build.Dependencies.cythonize称为compile_time_env)。

Answer 2

谢谢您非常有用的回答@IanH。我举了一个例子来说明它给出了什么。

文件setup.py

from setuptools import setup
from Cython.Distutils.extension import Extension
from Cython.Distutils import build_ext

import os
here = os.path.abspath(os.path.dirname(__file__))

import numpy as np
include_dirs = [here, np.get_include()]

try:
    import mpi4py
except ImportError:
    MPI4PY = False
else:
    MPI4PY = True
    INCLUDE_MPI = '/usr/lib/openmpi/include'
    include_dirs.extend([
        INCLUDE_MPI,
        mpi4py.get_include()])

name = 'mod'
ext = Extension(
    name,
    include_dirs=include_dirs,
    cython_compile_time_env={'MPI4PY': MPI4PY},
    sources=['mod.pyx'])

setup(name=name,
      cmdclass={"build_ext": build_ext},
      ext_modules=[ext])

if not MPI4PY:
    print('Warning: since importing mpi4py raises an ImportError,\n'
          '         the extensions are compiled without mpi and \n'
          '         will work only in sequencial.')

还有文件mod.pyx，带有一些真正的mpi命令：

import numpy as np
cimport numpy as np

try:
    from mpi4py import MPI
except ImportError:
    nb_proc = 1
    rank = 0
else:
    comm = MPI.COMM_WORLD
    nb_proc = comm.size
    rank = comm.Get_rank()

IF MPI4PY:
    from mpi4py cimport MPI
    from mpi4py.mpi_c cimport *

    # solve an incompatibility between openmpi and mpi4py versions
    cdef extern from 'mpi-compat.h': pass

    print('mpi4py ok')
ELSE:
    print('no mpi4py')

n = 8
if n % nb_proc != 0:
    raise ValueError('The number of processes is incorrect.')

if rank == 0:
    data_seq = np.ones([n], dtype=np.int32)
    s_seq = data_seq.sum()
else:
    data_seq = np.zeros([n], dtype=np.int32)

if nb_proc > 1:
    data_local = np.zeros([n/nb_proc], dtype=np.int32)
    comm.Scatter(data_seq, data_local, root=0)
else:
    data_local = data_seq

s = data_local.sum()
if nb_proc > 1:
    s = comm.allreduce(s, op=MPI.SUM)

if rank == 0:
    print('s: {}; s_seq: {}'.format(s, s_seq))
    assert s == s_seq

使用python setup.py build_ext --inplace构建，用python -c "import mod"和mpirun -np 4 python -c "import mod"进行测试。如果没有安装mpi4py，仍然可以构建模块并按顺序使用它。