1318 lines
49 KiB
Python
1318 lines
49 KiB
Python
# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# 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.
|
|
|
|
import os
|
|
import shlex
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace, SimpleNamespace as ModuleSimpleNamespace
|
|
from unittest import mock
|
|
|
|
from setuptools import Distribution
|
|
|
|
from paddle.utils.cpp_extension.cpp_extension import (
|
|
_as_command_list,
|
|
_get_num_workers,
|
|
_is_ninja_available,
|
|
_join_ninja_shell_list,
|
|
_ninja_escape_path,
|
|
_run_ninja_build,
|
|
_write_ninja_file,
|
|
)
|
|
from paddle.utils.cpp_extension.extension_utils import (
|
|
_write_setup_file,
|
|
)
|
|
|
|
|
|
class TestNinjaHelperFunctions(unittest.TestCase):
|
|
def test_is_ninja_available(self):
|
|
result = _is_ninja_available()
|
|
self.assertIsInstance(result, bool)
|
|
if sys.platform.startswith('linux'):
|
|
self.assertTrue(result, "ninja should be available on Linux CI")
|
|
|
|
def test_is_ninja_available_handles_probe_failure(self):
|
|
with mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.subprocess.check_output',
|
|
side_effect=OSError("ninja missing"),
|
|
):
|
|
self.assertFalse(_is_ninja_available())
|
|
|
|
def test_ninja_escape_path(self):
|
|
self.assertEqual(_ninja_escape_path("/path/to/file"), "/path/to/file")
|
|
self.assertEqual(
|
|
_ninja_escape_path("/path with space/file"),
|
|
"/path$ with$ space/file",
|
|
)
|
|
self.assertEqual(
|
|
_ninja_escape_path("C:/path/to/file"), "C$:/path/to/file"
|
|
)
|
|
self.assertEqual(
|
|
_ninja_escape_path("/path/$var/file"), "/path/$$var/file"
|
|
)
|
|
self.assertEqual(
|
|
_ninja_escape_path("C:/path with $var/file"),
|
|
"C$:/path$ with$ $$var/file",
|
|
)
|
|
|
|
def test_join_ninja_shell_list(self):
|
|
self.assertEqual(
|
|
_join_ninja_shell_list("simple string"), "simple string"
|
|
)
|
|
self.assertEqual(_join_ninja_shell_list([]), "")
|
|
result = _join_ninja_shell_list(['-c', '-O2', '-I/usr/include'])
|
|
self.assertIn('-c', result)
|
|
self.assertIn('-O2', result)
|
|
result = _join_ninja_shell_list(['/path with space', '-c'])
|
|
self.assertIn('/path with space', result)
|
|
|
|
def test_as_command_list(self):
|
|
self.assertEqual(_as_command_list("gcc"), ["gcc"])
|
|
self.assertEqual(_as_command_list(["gcc", "-c"]), ["gcc", "-c"])
|
|
self.assertEqual(_as_command_list(("gcc", "-c")), ["gcc", "-c"])
|
|
|
|
def test_write_ninja_file(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
ninja_path = os.path.join(tmpdir, "build.ninja")
|
|
_write_ninja_file(ninja_path, "ninja_required_version = 1.5")
|
|
|
|
with open(ninja_path, 'r', encoding='utf-8') as f:
|
|
written_content = f.read()
|
|
self.assertTrue(written_content.endswith('\n'))
|
|
self.assertIn("ninja_required_version = 1.5", written_content)
|
|
|
|
content_with_newline = "rule compile\n command = $cxx -c $in\n"
|
|
_write_ninja_file(ninja_path, content_with_newline)
|
|
with open(ninja_path, 'r', encoding='utf-8') as f:
|
|
written_content = f.read()
|
|
self.assertTrue(written_content.endswith('\n'))
|
|
|
|
nested_path = os.path.join(tmpdir, "nested", "dir", "build.ninja")
|
|
_write_ninja_file(nested_path, "test content")
|
|
self.assertTrue(os.path.exists(nested_path))
|
|
|
|
def test_get_num_workers_with_max_jobs_env(self):
|
|
with mock.patch.dict(os.environ, {'MAX_JOBS': '4'}, clear=False):
|
|
result = _get_num_workers(verbose=False)
|
|
self.assertEqual(result, 4)
|
|
|
|
def test_run_ninja_build_windows_vc_env(self):
|
|
if sys.platform != 'win32':
|
|
self.skipTest("Windows-only test")
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
ninja_path = os.path.join(tmpdir, "build.ninja")
|
|
_write_ninja_file(ninja_path, "rule cc\n command = echo hello\n")
|
|
|
|
vc_env_mock = {'PATH': '/vc/bin', 'INCLUDE': '/vc/include'}
|
|
|
|
with (
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.distutils.util.get_platform',
|
|
return_value='win-amd64',
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._get_vc_env',
|
|
return_value=vc_env_mock,
|
|
),
|
|
mock.patch('subprocess.run') as mock_run,
|
|
):
|
|
_run_ninja_build(tmpdir, verbose=True, error_prefix="Test")
|
|
mock_run.assert_called_once()
|
|
|
|
def test_run_ninja_build_with_work_directory(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
build_dir = os.path.join(tmpdir, 'build')
|
|
work_dir = os.path.join(tmpdir, 'src')
|
|
os.makedirs(build_dir)
|
|
os.makedirs(work_dir)
|
|
|
|
with (
|
|
mock.patch.dict(
|
|
os.environ, {'VSCMD_ARG_TGT_ARCH': 'x64'}, clear=False
|
|
),
|
|
mock.patch('subprocess.run') as mock_run,
|
|
):
|
|
_run_ninja_build(
|
|
build_dir,
|
|
verbose=True,
|
|
error_prefix='Test',
|
|
work_directory=work_dir,
|
|
)
|
|
|
|
command = mock_run.call_args.args[0]
|
|
self.assertIn('-f', command)
|
|
self.assertIn(os.path.join(build_dir, 'build.ninja'), command)
|
|
self.assertEqual(mock_run.call_args.kwargs['cwd'], work_dir)
|
|
|
|
def test_run_ninja_build_windows_with_vscmd_env(self):
|
|
if sys.platform != 'win32':
|
|
self.skipTest("Windows-only test")
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
ninja_path = os.path.join(tmpdir, "build.ninja")
|
|
_write_ninja_file(ninja_path, "rule cc\n command = echo hello\n")
|
|
|
|
with (
|
|
mock.patch.dict(
|
|
os.environ, {'VSCMD_ARG_TGT_ARCH': 'x64'}, clear=False
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.IS_WINDOWS', True
|
|
),
|
|
mock.patch('subprocess.run') as mock_run,
|
|
):
|
|
_run_ninja_build(tmpdir, verbose=True, error_prefix="Test")
|
|
mock_run.assert_called_once()
|
|
|
|
|
|
class _FakeUnixCompiler:
|
|
compiler_type = 'unix'
|
|
|
|
def __init__(self, objects, build_map):
|
|
self.src_extensions = ['.c', '.cc', '.cpp']
|
|
self.compiler_so = ['g++']
|
|
self.compiler_cxx = ['g++']
|
|
self._objects = objects
|
|
self._build_map = build_map
|
|
self._compile_calls = []
|
|
self.verbose = False
|
|
|
|
def compile(self, *args, **kwargs):
|
|
raise AssertionError("original unix compile should be replaced")
|
|
|
|
def link_shared_object(self, *args, **kwargs):
|
|
return None
|
|
|
|
def object_filenames(self, *args, **kwargs):
|
|
return self._objects
|
|
|
|
def _setup_compile(
|
|
self,
|
|
output_dir,
|
|
macros,
|
|
include_dirs,
|
|
sources,
|
|
depends,
|
|
extra_postargs,
|
|
):
|
|
pp_opts = [f'-I{inc}' for inc in (include_dirs or [])]
|
|
return macros, self._objects, extra_postargs, pp_opts, self._build_map
|
|
|
|
def _get_cc_args(self, pp_opts, debug, extra_preargs):
|
|
return list(extra_preargs or []) + list(pp_opts)
|
|
|
|
def _compile(self, obj, src, ext, cc_args, cflags, pp_opts):
|
|
self._compile_calls.append(
|
|
{
|
|
'obj': obj,
|
|
'src': src,
|
|
'ext': ext,
|
|
'cc_args': list(cc_args),
|
|
'cflags': list(cflags),
|
|
'pp_opts': list(pp_opts),
|
|
'compiler_so': list(self.compiler_so)
|
|
if isinstance(self.compiler_so, list)
|
|
else self.compiler_so,
|
|
}
|
|
)
|
|
|
|
def set_executable(self, key, value):
|
|
setattr(self, key, value)
|
|
|
|
def set_executables(self, **kwargs):
|
|
return None
|
|
|
|
|
|
class _FakeMsvcCompiler:
|
|
compiler_type = 'msvc'
|
|
|
|
def __init__(self, objects, build_map):
|
|
self.src_extensions = ['.c', '.cc', '.cpp']
|
|
self._cpp_extensions = ['.c', '.cc', '.cpp']
|
|
self.cc = ['cl.exe']
|
|
self.compile_options = ['/nologo', '/O2', '/W3', '/MD']
|
|
self.compile_options_debug = [
|
|
'/nologo',
|
|
'/Od',
|
|
'/MDd',
|
|
'/Zi',
|
|
'/W3',
|
|
'/D_DEBUG',
|
|
]
|
|
self.initialized = True
|
|
self._objects = objects
|
|
self._build_map = build_map
|
|
self._spawned_cmds = []
|
|
self.verbose = False
|
|
self.spawn = self._spawn
|
|
|
|
def compile(self, *args, **kwargs):
|
|
macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
|
|
kwargs.get('output_dir'),
|
|
kwargs.get('macros'),
|
|
kwargs.get('include_dirs'),
|
|
args[0],
|
|
kwargs.get('depends'),
|
|
kwargs.get('extra_postargs'),
|
|
)
|
|
del macros, extra_postargs
|
|
for obj in objects:
|
|
src, ext = build[obj]
|
|
cmd = [
|
|
'cl.exe',
|
|
'/Tp' + src,
|
|
'/Fo' + obj,
|
|
*pp_opts,
|
|
]
|
|
if ext == '.c':
|
|
cmd[1] = '/Tc' + src
|
|
self.spawn(cmd)
|
|
return objects
|
|
|
|
def initialize(self):
|
|
self.initialized = True
|
|
|
|
def object_filenames(self, *args, **kwargs):
|
|
return self._objects
|
|
|
|
def _setup_compile(
|
|
self,
|
|
output_dir,
|
|
macros,
|
|
include_dirs,
|
|
sources,
|
|
depends,
|
|
extra_postargs,
|
|
):
|
|
pp_opts = [f'/I{inc}' for inc in (include_dirs or [])]
|
|
return macros, self._objects, extra_postargs, pp_opts, self._build_map
|
|
|
|
def _spawn(self, cmd):
|
|
self._spawned_cmds.append(list(cmd))
|
|
|
|
|
|
class TestBuildExtension(unittest.TestCase):
|
|
def _build_extension(self, **kwargs):
|
|
from paddle.utils.cpp_extension.cpp_extension import BuildExtension
|
|
|
|
return BuildExtension(dist=Distribution(), **kwargs)
|
|
|
|
def _run_build_with_fake_compiler(
|
|
self,
|
|
compiler,
|
|
sources,
|
|
extra_compile_args,
|
|
include_dirs=None,
|
|
):
|
|
ext = SimpleNamespace(
|
|
name='fake_extension',
|
|
_full_name='fake_extension',
|
|
sources=sources,
|
|
extra_compile_args=extra_compile_args,
|
|
)
|
|
cmd = self._build_extension(use_ninja=True)
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
cmd.extensions = [ext]
|
|
cmd.compiler = compiler
|
|
cmd.build_temp = tmpdir
|
|
cmd.build_lib = tmpdir
|
|
cmd.verbose = False
|
|
cmd.contain_cuda_file = any(
|
|
source.endswith('.cu') for source in sources
|
|
)
|
|
|
|
captured = {}
|
|
|
|
def fake_write_ninja_file(path, content):
|
|
captured['ninja_path'] = path
|
|
captured['ninja_content'] = content
|
|
|
|
def fake_run_ninja_build(
|
|
build_directory, verbose, error_prefix, work_directory=None
|
|
):
|
|
captured['ninja_build_directory'] = build_directory
|
|
captured['ninja_work_directory'] = work_directory
|
|
|
|
def fake_build_extensions(_):
|
|
_.compiler.compile(
|
|
sources,
|
|
output_dir=_.build_temp,
|
|
macros=[],
|
|
include_dirs=include_dirs or [],
|
|
debug=False,
|
|
extra_preargs=[],
|
|
extra_postargs=extra_compile_args,
|
|
depends=None,
|
|
)
|
|
|
|
patches = [
|
|
mock.patch.object(cmd, '_check_abi', return_value=None),
|
|
mock.patch.object(cmd, '_record_op_info', return_value=None),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.clean_object_if_change_cflags',
|
|
return_value=None,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.define_paddle_extension_name',
|
|
return_value=None,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._reset_so_rpath',
|
|
return_value=None,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._write_ninja_file',
|
|
side_effect=fake_write_ninja_file,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._run_ninja_build',
|
|
side_effect=fake_run_ninja_build,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.build_ext.build_extensions',
|
|
side_effect=fake_build_extensions,
|
|
),
|
|
]
|
|
|
|
with (
|
|
patches[0],
|
|
patches[1],
|
|
patches[2],
|
|
patches[3],
|
|
patches[4],
|
|
patches[5],
|
|
patches[6],
|
|
patches[7],
|
|
):
|
|
cmd.build_extensions()
|
|
|
|
return captured
|
|
|
|
def _patch_build_extension_common(self, cmd):
|
|
return [
|
|
mock.patch.object(cmd, '_check_abi', return_value=None),
|
|
mock.patch.object(cmd, '_record_op_info', return_value=None),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.clean_object_if_change_cflags',
|
|
return_value=None,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.define_paddle_extension_name',
|
|
return_value=None,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._reset_so_rpath',
|
|
return_value=None,
|
|
),
|
|
]
|
|
|
|
@staticmethod
|
|
def _expected_nvcc_path(cuda_home='/opt/cuda'):
|
|
return os.path.join(cuda_home, 'bin', 'nvcc')
|
|
|
|
def test_unix_ninja_build_file_contains_multiple_sources(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
sources = [
|
|
str(test_dir / "custom_extension.cc"),
|
|
str(test_dir / "custom_sub.cc"),
|
|
]
|
|
objects = [
|
|
str(test_dir / "build" / "custom_extension.o"),
|
|
str(test_dir / "build" / "custom_sub.o"),
|
|
]
|
|
build_map = {
|
|
objects[0]: (sources[0], '.cc'),
|
|
objects[1]: (sources[1], '.cc'),
|
|
}
|
|
compiler = _FakeUnixCompiler(objects, build_map)
|
|
|
|
captured = self._run_build_with_fake_compiler(
|
|
compiler,
|
|
sources,
|
|
extra_compile_args={'cxx': ['-w', '-g'], 'nvcc': []},
|
|
include_dirs=[str(test_dir / "include")],
|
|
)
|
|
|
|
content = captured['ninja_content']
|
|
self.assertIn('rule compile', content)
|
|
self.assertIn('deps = gcc', content)
|
|
self.assertIn(_ninja_escape_path(os.path.abspath(objects[0])), content)
|
|
self.assertIn(_ninja_escape_path(os.path.abspath(objects[1])), content)
|
|
self.assertIn(_ninja_escape_path(os.path.abspath(sources[0])), content)
|
|
self.assertIn(_ninja_escape_path(os.path.abspath(sources[1])), content)
|
|
self.assertIn('post_cflags = -w -g', content)
|
|
self.assertIn(
|
|
f'default {_ninja_escape_path(os.path.abspath(objects[0]))} {_ninja_escape_path(os.path.abspath(objects[1]))}',
|
|
content,
|
|
)
|
|
self.assertTrue(captured['ninja_path'].endswith('build.ninja'))
|
|
|
|
def test_windows_ninja_build_file_contains_multiple_sources(self):
|
|
if sys.platform != 'win32':
|
|
self.skipTest("Windows-only test")
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
sources = [
|
|
str(test_dir / "custom_extension.cc"),
|
|
str(test_dir / "custom_sub.cc"),
|
|
]
|
|
objects = [
|
|
str(test_dir / "build" / "custom_extension.obj"),
|
|
str(test_dir / "build" / "custom_sub.obj"),
|
|
]
|
|
build_map = {
|
|
objects[0]: (sources[0], '.cc'),
|
|
objects[1]: (sources[1], '.cc'),
|
|
}
|
|
compiler = _FakeMsvcCompiler(objects, build_map)
|
|
|
|
captured = self._run_build_with_fake_compiler(
|
|
compiler,
|
|
sources,
|
|
extra_compile_args={'cxx': ['/wd4244'], 'nvcc': []},
|
|
include_dirs=[str(test_dir / "include dir")],
|
|
)
|
|
|
|
content = captured['ninja_content']
|
|
self.assertIn('rule compile', content)
|
|
self.assertIn('deps = msvc', content)
|
|
self.assertIn('command = $cxx /showIncludes', content)
|
|
self.assertIn('cl.exe', content)
|
|
self.assertIn('/wd4244', content)
|
|
self.assertIn(_ninja_escape_path(os.path.abspath(objects[0])), content)
|
|
self.assertIn(_ninja_escape_path(os.path.abspath(objects[1])), content)
|
|
self.assertIn(_ninja_escape_path(os.path.abspath(sources[0])), content)
|
|
self.assertIn(_ninja_escape_path(os.path.abspath(sources[1])), content)
|
|
self.assertIn('/I', content)
|
|
self.assertTrue(captured['ninja_path'].endswith('build.ninja'))
|
|
|
|
def test_windows_ninja_quotes_flags_once(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
include_dir = test_dir / "include dir"
|
|
sources = [
|
|
str(test_dir / "custom_extension.cc"),
|
|
str(test_dir / "custom_kernel.cu"),
|
|
]
|
|
objects = [
|
|
str(test_dir / "build" / "custom_extension.obj"),
|
|
str(test_dir / "build" / "custom_kernel.obj"),
|
|
]
|
|
build_map = {
|
|
objects[0]: (sources[0], '.cc'),
|
|
objects[1]: (sources[1], '.cu'),
|
|
}
|
|
compiler = _FakeMsvcCompiler(objects, build_map)
|
|
|
|
with (
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.IS_WINDOWS',
|
|
True,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
|
|
'/opt/cuda',
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.prepare_win_cudaflags',
|
|
side_effect=lambda flags: ['--prepared', *flags],
|
|
),
|
|
):
|
|
captured = self._run_build_with_fake_compiler(
|
|
compiler,
|
|
sources,
|
|
extra_compile_args={
|
|
'cxx': ['/DNAME=value with space'],
|
|
'nvcc': ['--compiler-options=/Wall /WX'],
|
|
},
|
|
include_dirs=[str(include_dir)],
|
|
)
|
|
|
|
content = captured['ninja_content']
|
|
self.assertIn(subprocess.list2cmdline([f'/I{include_dir}']), content)
|
|
self.assertIn(
|
|
subprocess.list2cmdline(['/DNAME=value with space']), content
|
|
)
|
|
self.assertIn(
|
|
subprocess.list2cmdline(['--compiler-options=/Wall /WX']),
|
|
content,
|
|
)
|
|
self.assertNotIn('\\"/I', content)
|
|
self.assertNotIn('\\"/DNAME', content)
|
|
|
|
def test_unix_ninja_build_file_contains_cuda_sources(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
sources = [
|
|
str(test_dir / "custom_extension.cc"),
|
|
str(test_dir / "custom_kernel.cu"),
|
|
]
|
|
objects = [
|
|
str(test_dir / "build" / "custom_extension.o"),
|
|
str(test_dir / "build" / "custom_kernel.cu.o"),
|
|
]
|
|
build_map = {
|
|
objects[0]: (sources[0], '.cc'),
|
|
objects[1]: (sources[1], '.cu'),
|
|
}
|
|
compiler = _FakeUnixCompiler(objects, build_map)
|
|
|
|
with (
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
|
|
'/opt/cuda',
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._get_ccache_home',
|
|
return_value='/usr/bin/ccache',
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.prepare_unix_cudaflags',
|
|
side_effect=lambda flags: ['--prepared', *flags],
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_rocm',
|
|
return_value=False,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_custom_device',
|
|
return_value=False,
|
|
),
|
|
):
|
|
captured = self._run_build_with_fake_compiler(
|
|
compiler,
|
|
sources,
|
|
extra_compile_args={
|
|
'cxx': ['-w'],
|
|
'nvcc': ['--gpu-flag'],
|
|
},
|
|
include_dirs=[str(test_dir / "include")],
|
|
)
|
|
|
|
content = captured['ninja_content']
|
|
self.assertIn('cxx = /usr/bin/ccache g++', content)
|
|
self.assertIn(
|
|
f'nvcc = /usr/bin/ccache {self._expected_nvcc_path()}',
|
|
content,
|
|
)
|
|
self.assertIn('rule cuda_compile', content)
|
|
self.assertIn('cuda_post_cflags = --prepared --gpu-flag', content)
|
|
self.assertIn('post_cflags = -w -D_GLIBCXX_USE_CXX11_ABI=1', content)
|
|
self.assertIn('-DPADDLE_WITH_CUDA', content)
|
|
|
|
def test_unix_ninja_quotes_flags_once(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
include_dir = test_dir / "include dir"
|
|
sources = [
|
|
str(test_dir / "custom_extension.cc"),
|
|
str(test_dir / "custom_kernel.cu"),
|
|
]
|
|
objects = [
|
|
str(test_dir / "build" / "custom_extension.o"),
|
|
str(test_dir / "build" / "custom_kernel.cu.o"),
|
|
]
|
|
build_map = {
|
|
objects[0]: (sources[0], '.cc'),
|
|
objects[1]: (sources[1], '.cu'),
|
|
}
|
|
compiler = _FakeUnixCompiler(objects, build_map)
|
|
|
|
with (
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.IS_WINDOWS',
|
|
False,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
|
|
'/opt/cuda',
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._get_ccache_home',
|
|
return_value=None,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.prepare_unix_cudaflags',
|
|
side_effect=lambda flags: ['--prepared', *flags],
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_rocm',
|
|
return_value=False,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_custom_device',
|
|
return_value=False,
|
|
),
|
|
):
|
|
captured = self._run_build_with_fake_compiler(
|
|
compiler,
|
|
sources,
|
|
extra_compile_args={
|
|
'cxx': ['-DNAME=value with space'],
|
|
'nvcc': ['--compiler-options=-Wall -Wextra'],
|
|
},
|
|
include_dirs=[str(include_dir)],
|
|
)
|
|
|
|
content = captured['ninja_content']
|
|
self.assertIn(shlex.join([f'-I{include_dir}']), content)
|
|
self.assertIn(shlex.join(['-DNAME=value with space']), content)
|
|
self.assertIn(shlex.join(['--compiler-options=-Wall -Wextra']), content)
|
|
self.assertNotIn('"\'"\'"', content)
|
|
|
|
def test_unix_ninja_preserves_relative_include_flags(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
sources = [
|
|
str(test_dir / "custom_extension.cc"),
|
|
str(test_dir / "custom_kernel.cu"),
|
|
]
|
|
objects = [
|
|
str(test_dir / "build" / "custom_extension.o"),
|
|
str(test_dir / "build" / "custom_kernel.cu.o"),
|
|
]
|
|
build_map = {
|
|
objects[0]: (sources[0], '.cc'),
|
|
objects[1]: (sources[1], '.cu'),
|
|
}
|
|
compiler = _FakeUnixCompiler(objects, build_map)
|
|
|
|
with (
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.os.getcwd',
|
|
return_value=str(test_dir),
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
|
|
'/opt/cuda',
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._get_ccache_home',
|
|
return_value=None,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.prepare_unix_cudaflags',
|
|
side_effect=lambda flags: ['--prepared', *flags],
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_rocm',
|
|
return_value=False,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_custom_device',
|
|
return_value=False,
|
|
),
|
|
):
|
|
captured = self._run_build_with_fake_compiler(
|
|
compiler,
|
|
sources,
|
|
extra_compile_args={
|
|
'cxx': ['-Ithird_party/cxx', '-isystem', 'system/cxx'],
|
|
'nvcc': ['-Ithird_party/cuda', '-I', 'cuda_split'],
|
|
},
|
|
include_dirs=['public/include'],
|
|
)
|
|
|
|
content = captured['ninja_content']
|
|
self.assertIn('-Ipublic/include', content)
|
|
self.assertIn('-Ithird_party/cxx', content)
|
|
self.assertIn('-isystem system/cxx', content)
|
|
self.assertIn('-Ithird_party/cuda', content)
|
|
self.assertIn('-I cuda_split', content)
|
|
self.assertEqual(captured['ninja_work_directory'], str(test_dir))
|
|
|
|
def test_unix_ninja_uses_hipcc_for_rocm_sources(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
sources = [
|
|
str(test_dir / "custom_extension.cc"),
|
|
str(test_dir / "custom_kernel.cu"),
|
|
]
|
|
objects = [
|
|
str(test_dir / "build" / "custom_extension.o"),
|
|
str(test_dir / "build" / "custom_kernel.cu.o"),
|
|
]
|
|
build_map = {
|
|
objects[0]: (sources[0], '.cc'),
|
|
objects[1]: (sources[1], '.cu'),
|
|
}
|
|
compiler = _FakeUnixCompiler(objects, build_map)
|
|
|
|
with (
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.ROCM_HOME',
|
|
'/opt/rocm',
|
|
create=True,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._get_ccache_home',
|
|
return_value=None,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.prepare_unix_cudaflags',
|
|
side_effect=lambda flags: ['--prepared', *flags],
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_rocm',
|
|
return_value=True,
|
|
),
|
|
):
|
|
captured = self._run_build_with_fake_compiler(
|
|
compiler,
|
|
sources,
|
|
extra_compile_args={
|
|
'cxx': ['-w'],
|
|
'hipcc': ['--hip-flag'],
|
|
},
|
|
include_dirs=[str(test_dir / "include")],
|
|
)
|
|
|
|
content = captured['ninja_content']
|
|
self.assertIn('nvcc = ', content)
|
|
self.assertIn('hipcc', content)
|
|
self.assertIn('cuda_post_cflags = --prepared --hip-flag', content)
|
|
self.assertIn('-D__HIP_PLATFORM_HCC__', content)
|
|
self.assertIn(
|
|
'-DTHRUST_DEVICE_SYSTEM=THRUST_DEVICE_SYSTEM_HIP', content
|
|
)
|
|
self.assertIn('-DPADDLE_WITH_HIP', content)
|
|
|
|
def test_unix_ninja_reports_missing_corex_compiler(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
sources = [str(test_dir / "custom_kernel.cu")]
|
|
objects = [str(test_dir / "build" / "custom_kernel.cu.o")]
|
|
build_map = {objects[0]: (sources[0], '.cu')}
|
|
compiler = _FakeUnixCompiler(objects, build_map)
|
|
|
|
with (
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_rocm',
|
|
return_value=False,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_custom_device',
|
|
return_value=True,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.os.path.isfile',
|
|
return_value=False,
|
|
),
|
|
self.assertRaisesRegex(ValueError, 'Corex compiler'),
|
|
):
|
|
self._run_build_with_fake_compiler(
|
|
compiler,
|
|
sources,
|
|
extra_compile_args={'nvcc': ['--gpu-flag']},
|
|
)
|
|
|
|
def test_use_ninja_attribute_default(self):
|
|
build_ext = self._build_extension()
|
|
self.assertEqual(build_ext.use_ninja, _is_ninja_available())
|
|
|
|
def test_use_ninja_attribute_explicit_false(self):
|
|
build_ext = self._build_extension(use_ninja=False)
|
|
self.assertFalse(build_ext.use_ninja)
|
|
|
|
def test_use_ninja_attribute_explicit_true(self):
|
|
build_ext = self._build_extension(use_ninja=True)
|
|
self.assertEqual(build_ext.use_ninja, _is_ninja_available())
|
|
|
|
def test_use_ninja_falls_back_when_ninja_missing(self):
|
|
with (
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._is_ninja_available',
|
|
return_value=False,
|
|
),
|
|
mock.patch('builtins.print') as mock_print,
|
|
):
|
|
build_ext = self._build_extension(use_ninja=True)
|
|
self.assertFalse(build_ext.use_ninja)
|
|
mock_print.assert_any_call(
|
|
"Ninja is not available, falling back to the distutils backend."
|
|
)
|
|
|
|
def test_unix_compiler_with_use_ninja_false_executes_compile(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
sources = [
|
|
str(test_dir / "custom_extension.cc"),
|
|
str(test_dir / "custom_kernel.cu"),
|
|
]
|
|
objects = [
|
|
str(test_dir / "build" / "custom_extension.o"),
|
|
str(test_dir / "build" / "custom_kernel.o"),
|
|
]
|
|
build_map = {
|
|
objects[0]: (sources[0], '.cc'),
|
|
objects[1]: (sources[1], '.cu'),
|
|
}
|
|
compiler = _FakeUnixCompiler(objects, build_map)
|
|
|
|
ext = SimpleNamespace(
|
|
name='fake_extension',
|
|
_full_name='fake_extension',
|
|
sources=sources,
|
|
extra_compile_args={'cxx': ['-w'], 'nvcc': ['--gpu-flag']},
|
|
)
|
|
cmd = self._build_extension(use_ninja=False)
|
|
|
|
cmd.extensions = [ext]
|
|
cmd.compiler = compiler
|
|
cmd.build_temp = tmpdir
|
|
cmd.build_lib = tmpdir
|
|
cmd.verbose = False
|
|
cmd.contain_cuda_file = True
|
|
|
|
patches = self._patch_build_extension_common(cmd)
|
|
|
|
with (
|
|
patches[0],
|
|
patches[1],
|
|
patches[2],
|
|
patches[3],
|
|
patches[4],
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
|
|
'/opt/cuda',
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.prepare_unix_cudaflags',
|
|
side_effect=lambda flags: ['--prepared', *flags],
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._get_ccache_home',
|
|
return_value=None,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._get_num_workers',
|
|
return_value=4,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.os.cpu_count',
|
|
return_value=8,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_rocm',
|
|
return_value=False,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_custom_device',
|
|
return_value=False,
|
|
),
|
|
mock.patch('builtins.print') as mock_print,
|
|
):
|
|
|
|
def fake_build_extensions(_):
|
|
compiled = _.compiler.compile(
|
|
sources,
|
|
output_dir=_.build_temp,
|
|
macros=[],
|
|
include_dirs=[str(test_dir / "include")],
|
|
debug=False,
|
|
extra_preargs=[],
|
|
extra_postargs=ext.extra_compile_args,
|
|
depends=None,
|
|
)
|
|
self.assertEqual(compiled, objects)
|
|
|
|
with mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.build_ext.build_extensions',
|
|
side_effect=fake_build_extensions,
|
|
):
|
|
cmd.build_extensions()
|
|
|
|
self.assertEqual(len(compiler._compile_calls), 2)
|
|
cc_call = next(
|
|
call
|
|
for call in compiler._compile_calls
|
|
if call['src'].endswith('.cc')
|
|
)
|
|
cu_call = next(
|
|
call
|
|
for call in compiler._compile_calls
|
|
if call['src'].endswith('.cu')
|
|
)
|
|
self.assertIn('-DPADDLE_WITH_CUDA', cc_call['cflags'])
|
|
self.assertIn('-D_GLIBCXX_USE_CXX11_ABI=1', cc_call['cflags'])
|
|
self.assertIn('--prepared', cu_call['cflags'])
|
|
self.assertEqual(cu_call['compiler_so'], self._expected_nvcc_path())
|
|
mock_print.assert_any_call(
|
|
"Using 2 workers for compilation. HINT: export MAX_JOBS=n to set the number of workers"
|
|
)
|
|
|
|
def test_msvc_compiler_with_use_ninja_false_executes_compile(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
sources = [
|
|
str(test_dir / "custom_extension.cc"),
|
|
str(test_dir / "custom_kernel.cu"),
|
|
]
|
|
objects = [
|
|
str(test_dir / "build" / "custom_extension.obj"),
|
|
str(test_dir / "build" / "custom_kernel.obj"),
|
|
]
|
|
build_map = {
|
|
objects[0]: (sources[0], '.cc'),
|
|
objects[1]: (sources[1], '.cu'),
|
|
}
|
|
compiler = _FakeMsvcCompiler(objects, build_map)
|
|
|
|
ext = SimpleNamespace(
|
|
name='fake_extension',
|
|
_full_name='fake_extension',
|
|
sources=sources,
|
|
extra_compile_args={
|
|
'cxx': ['/wd4244'],
|
|
'nvcc': ['--gpu-flag'],
|
|
},
|
|
)
|
|
cmd = self._build_extension(use_ninja=False)
|
|
|
|
cmd.extensions = [ext]
|
|
cmd.compiler = compiler
|
|
cmd.build_temp = tmpdir
|
|
cmd.build_lib = tmpdir
|
|
cmd.verbose = False
|
|
cmd.contain_cuda_file = True
|
|
|
|
patches = self._patch_build_extension_common(cmd)
|
|
|
|
with (
|
|
patches[0],
|
|
patches[1],
|
|
patches[2],
|
|
patches[3],
|
|
patches[4],
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
|
|
'/opt/cuda',
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.prepare_win_cudaflags',
|
|
side_effect=lambda flags: ['--prepared', *flags],
|
|
),
|
|
):
|
|
|
|
def fake_build_extensions(_):
|
|
_.compiler.compile(
|
|
sources,
|
|
output_dir=_.build_temp,
|
|
macros=[],
|
|
include_dirs=[str(test_dir / "include dir")],
|
|
debug=False,
|
|
extra_preargs=[],
|
|
extra_postargs=ext.extra_compile_args,
|
|
depends=None,
|
|
)
|
|
|
|
with mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.build_ext.build_extensions',
|
|
side_effect=fake_build_extensions,
|
|
):
|
|
cmd.build_extensions()
|
|
|
|
self.assertEqual(len(compiler._spawned_cmds), 2)
|
|
cc_cmd = next(
|
|
cmd
|
|
for cmd in compiler._spawned_cmds
|
|
if any(arg.endswith('.cc') for arg in cmd)
|
|
)
|
|
cu_cmd = next(
|
|
cmd
|
|
for cmd in compiler._spawned_cmds
|
|
if any(arg.endswith('.cu') for arg in cmd)
|
|
)
|
|
self.assertIn('-DPADDLE_WITH_CUDA', cc_cmd)
|
|
self.assertIn('/wd4244', cc_cmd)
|
|
self.assertEqual(cu_cmd[0], self._expected_nvcc_path())
|
|
self.assertIn('--prepared', cu_cmd)
|
|
self.assertIn('--use-local-env', cu_cmd)
|
|
|
|
def test_windows_ninja_build_file_contains_cuda_sources(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
sources = [
|
|
str(test_dir / "custom_extension.cc"),
|
|
str(test_dir / "custom_kernel.cu"),
|
|
]
|
|
objects = [
|
|
str(test_dir / "build" / "custom_extension.obj"),
|
|
str(test_dir / "build" / "custom_kernel.obj"),
|
|
]
|
|
build_map = {
|
|
objects[0]: (sources[0], '.cc'),
|
|
objects[1]: (sources[1], '.cu'),
|
|
}
|
|
compiler = _FakeMsvcCompiler(objects, build_map)
|
|
|
|
with (
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
|
|
'/opt/cuda',
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.prepare_win_cudaflags',
|
|
side_effect=lambda flags: ['--prepared', *flags],
|
|
),
|
|
):
|
|
captured = self._run_build_with_fake_compiler(
|
|
compiler,
|
|
sources,
|
|
extra_compile_args={
|
|
'cxx': ['/wd4244'],
|
|
'nvcc': ['--gpu-flag'],
|
|
},
|
|
include_dirs=[str(test_dir / "include dir")],
|
|
)
|
|
|
|
content = captured['ninja_content']
|
|
self.assertIn(f'nvcc = {self._expected_nvcc_path()}', content)
|
|
self.assertIn('rule cuda_compile', content)
|
|
self.assertIn('-Xcompiler /EHsc', content)
|
|
self.assertIn('--prepared', content)
|
|
self.assertIn('/DPADDLE_WITH_CUDA', content)
|
|
self.assertIn('deps = msvc', content)
|
|
|
|
def test_windows_ninja_preserves_relative_include_flags(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
sources = [
|
|
str(test_dir / "custom_extension.cc"),
|
|
str(test_dir / "custom_kernel.cu"),
|
|
]
|
|
objects = [
|
|
str(test_dir / "build" / "custom_extension.obj"),
|
|
str(test_dir / "build" / "custom_kernel.obj"),
|
|
]
|
|
build_map = {
|
|
objects[0]: (sources[0], '.cc'),
|
|
objects[1]: (sources[1], '.cu'),
|
|
}
|
|
compiler = _FakeMsvcCompiler(objects, build_map)
|
|
|
|
with (
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.os.getcwd',
|
|
return_value=str(test_dir),
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
|
|
'/opt/cuda',
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.prepare_win_cudaflags',
|
|
side_effect=lambda flags: ['--prepared', *flags],
|
|
),
|
|
):
|
|
captured = self._run_build_with_fake_compiler(
|
|
compiler,
|
|
sources,
|
|
extra_compile_args={
|
|
'cxx': ['/Ithird_party/cxx'],
|
|
'nvcc': ['/Ithird_party/cuda'],
|
|
},
|
|
include_dirs=['public/include'],
|
|
)
|
|
|
|
content = captured['ninja_content']
|
|
self.assertIn('/Ipublic/include', content)
|
|
self.assertIn('/Ithird_party/cxx', content)
|
|
self.assertIn('/Ithird_party/cuda', content)
|
|
self.assertEqual(captured['ninja_work_directory'], str(test_dir))
|
|
|
|
def test_windows_ninja_ignores_invalid_extra_compile_args(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
test_dir = Path(tmpdir)
|
|
sources = [
|
|
str(test_dir / "custom_extension.cc"),
|
|
str(test_dir / "custom_kernel.cu"),
|
|
]
|
|
objects = [
|
|
str(test_dir / "build" / "custom_extension.obj"),
|
|
str(test_dir / "build" / "custom_kernel.obj"),
|
|
]
|
|
build_map = {
|
|
objects[0]: (sources[0], '.cc'),
|
|
objects[1]: (sources[1], '.cu'),
|
|
}
|
|
compiler = _FakeMsvcCompiler(objects, build_map)
|
|
invalid_extra_args = 'not-a-list'
|
|
|
|
with (
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
|
|
'/opt/cuda',
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.prepare_win_cudaflags',
|
|
side_effect=lambda flags: ['--prepared', *flags],
|
|
),
|
|
):
|
|
captured = self._run_build_with_fake_compiler(
|
|
compiler,
|
|
sources,
|
|
extra_compile_args=invalid_extra_args,
|
|
include_dirs=[str(test_dir / "include dir")],
|
|
)
|
|
|
|
content = captured['ninja_content']
|
|
self.assertIn('rule cuda_compile', content)
|
|
self.assertIn('cuda_post_cflags = ', content)
|
|
self.assertIn('-Xcompiler /EHsc', content)
|
|
self.assertIn('--prepared --use-local-env', content)
|
|
self.assertIn('/DPADDLE_WITH_CUDA', content)
|
|
self.assertNotIn(invalid_extra_args, content)
|
|
|
|
|
|
class TestNinjaGeneratedSetupFile(unittest.TestCase):
|
|
def test_load_setup_file_uses_default_build_extension_options(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
setup_path = os.path.join(tmpdir, "demo_setup.py")
|
|
build_dir = os.path.join(tmpdir, "build")
|
|
sources = ['custom_extension.cc', 'custom_sub.cc']
|
|
|
|
_write_setup_file(
|
|
'demo_extension',
|
|
sources,
|
|
setup_path,
|
|
build_dir,
|
|
['include_dir'],
|
|
['library_dir'],
|
|
['-w', '-g'],
|
|
[],
|
|
[],
|
|
)
|
|
|
|
content = Path(setup_path).read_text(encoding='utf-8')
|
|
|
|
self.assertIn('BuildExtension.with_options(', content)
|
|
self.assertIn("output_dir=r'", content)
|
|
self.assertIn('no_python_abi_suffix=True', content)
|
|
self.assertIn(
|
|
"sources=['custom_extension.cc', 'custom_sub.cc']", content
|
|
)
|
|
self.assertNotIn('use_ninja=', content)
|
|
|
|
|
|
class TestRunNinjaBuild(unittest.TestCase):
|
|
def test_run_ninja_build_uses_verbose_subprocess_streams(self):
|
|
with (
|
|
mock.patch.dict(
|
|
os.environ, {'VSCMD_ARG_TGT_ARCH': 'x64'}, clear=False
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.subprocess.run'
|
|
) as mock_run,
|
|
):
|
|
_run_ninja_build('/tmp/build', verbose=True, error_prefix='Test')
|
|
|
|
_, kwargs = mock_run.call_args
|
|
self.assertEqual(mock_run.call_args.args[0], ['ninja', '-v'])
|
|
self.assertEqual(kwargs['cwd'], '/tmp/build')
|
|
self.assertIsNone(kwargs['stdout'])
|
|
self.assertIsNone(kwargs['stderr'])
|
|
|
|
def test_run_ninja_build_uses_num_workers_and_quiet_streams(self):
|
|
with (
|
|
mock.patch.dict(
|
|
os.environ,
|
|
{'MAX_JOBS': '3', 'VSCMD_ARG_TGT_ARCH': 'x64'},
|
|
clear=False,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.subprocess.run'
|
|
) as mock_run,
|
|
):
|
|
_run_ninja_build('/tmp/build', verbose=False, error_prefix='Test')
|
|
|
|
_, kwargs = mock_run.call_args
|
|
self.assertEqual(mock_run.call_args.args[0], ['ninja', '-v', '-j', '3'])
|
|
self.assertEqual(kwargs['stdout'], subprocess.PIPE)
|
|
self.assertEqual(kwargs['stderr'], subprocess.STDOUT)
|
|
self.assertTrue(kwargs['text'])
|
|
|
|
def test_run_ninja_build_windows_merges_vc_env(self):
|
|
vc_env = {'path': 'vc/bin', 'include': 'vc/include'}
|
|
with (
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.IS_WINDOWS', True
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.PLAT_TO_VCVARS',
|
|
{'win-amd64': 'x86_amd64'},
|
|
create=True,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.distutils',
|
|
ModuleSimpleNamespace(
|
|
util=ModuleSimpleNamespace(get_platform=lambda: 'win-amd64')
|
|
),
|
|
create=True,
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension._get_vc_env',
|
|
return_value=vc_env,
|
|
create=True,
|
|
),
|
|
mock.patch.dict(
|
|
os.environ, {'Path': 'user/bin', 'TMP': 'tmp'}, clear=True
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.subprocess.run'
|
|
) as mock_run,
|
|
):
|
|
_run_ninja_build('/tmp/build', verbose=False, error_prefix='Test')
|
|
|
|
env = mock_run.call_args.kwargs['env']
|
|
self.assertEqual(env['PATH'], 'vc/bin')
|
|
self.assertEqual(env['INCLUDE'], 'vc/include')
|
|
self.assertEqual(env['TMP'], 'tmp')
|
|
|
|
def test_run_ninja_build_raises_runtime_error(self):
|
|
with (
|
|
mock.patch.dict(
|
|
os.environ, {'VSCMD_ARG_TGT_ARCH': 'x64'}, clear=False
|
|
),
|
|
mock.patch(
|
|
'paddle.utils.cpp_extension.cpp_extension.subprocess.run',
|
|
side_effect=subprocess.CalledProcessError(1, ['ninja', '-v']),
|
|
),
|
|
self.assertRaisesRegex(RuntimeError, 'Prefix'),
|
|
):
|
|
_run_ninja_build('/tmp/build', verbose=False, error_prefix='Prefix')
|
|
|
|
def test_get_num_workers_verbose_without_env(self):
|
|
with (
|
|
mock.patch.dict(os.environ, {}, clear=True),
|
|
mock.patch('sys.stderr') as mock_stderr,
|
|
mock.patch('builtins.print') as mock_print,
|
|
):
|
|
result = _get_num_workers(verbose=True)
|
|
|
|
self.assertIsNone(result)
|
|
mock_print.assert_called_with(
|
|
'Allowing ninja to set a default number of workers... '
|
|
'(overridable by setting the environment variable MAX_JOBS=N)',
|
|
file=mock_stderr,
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|