mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
047e68235f
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/64733 The previous implementation was wrong when CPU scheduling affinity is set. In fact, it is still wrong if Ninja is not being used. When there is CPU scheduling affinity set, the number of processors available on the system likely exceeds the number of processors that are usable to the build. We ought to use `len(os.sched_getaffinity(0))` to determine the effective parallelism. This change is more minimal and instead just delegates to Ninja (which handles this correctly) when it is used. Test Plan: I verified this worked as correctly using Ninja on a 96-core machine with 24 cores available for scheduling by checking: * the cmake command did not specify "-j" * the number of top-level jobs in top/pstree never exceeded 26 (24 + 2) And I verified we get the legacy behavior by specifying USE_NINJA=0 on the build. Reviewed By: jbschlosser, driazati Differential Revision: D30968796 Pulled By: dagitses fbshipit-source-id: 29547dd378fea793957bcc2f7d52d5def1ecace2
88 lines
3.3 KiB
Python
88 lines
3.3 KiB
Python
import contextlib
|
|
import os
|
|
import typing
|
|
from typing import Iterator, Optional, Sequence
|
|
import unittest
|
|
import unittest.mock
|
|
|
|
import tools.setup_helpers.env # noqa: F401 unused but resolves circular import
|
|
import tools.setup_helpers.cmake
|
|
|
|
|
|
T = typing.TypeVar('T')
|
|
|
|
|
|
class TestCMake(unittest.TestCase):
|
|
|
|
@unittest.mock.patch('multiprocessing.cpu_count')
|
|
def test_build_jobs(self, mock_cpu_count: unittest.mock.MagicMock) -> None:
|
|
"""Tests that the number of build jobs comes out correctly."""
|
|
mock_cpu_count.return_value = 13
|
|
cases = [
|
|
# MAX_JOBS, USE_NINJA, IS_WINDOWS, want
|
|
(( '8', True, False), ['-j', '8']), # noqa: E201,E241
|
|
(( None, True, False), None), # noqa: E201,E241
|
|
(( None, True, True), None), # noqa: E201,E241
|
|
(( None, False, True), ['/p:CL_MPCount=13']), # noqa: E201,E241
|
|
]
|
|
for (max_jobs, use_ninja, is_windows), want in cases:
|
|
with self.subTest(MAX_JOBS=max_jobs, USE_NINJA=use_ninja, IS_WINDOWS=is_windows):
|
|
with contextlib.ExitStack() as stack:
|
|
stack.enter_context(env_var('MAX_JOBS', max_jobs))
|
|
stack.enter_context(unittest.mock.patch.object(tools.setup_helpers.cmake, 'USE_NINJA', use_ninja))
|
|
stack.enter_context(unittest.mock.patch.object(tools.setup_helpers.cmake, 'IS_WINDOWS', is_windows))
|
|
|
|
cmake = tools.setup_helpers.cmake.CMake()
|
|
|
|
with unittest.mock.patch.object(cmake, 'run') as cmake_run:
|
|
cmake.build({})
|
|
|
|
cmake_run.assert_called_once()
|
|
call, = cmake_run.mock_calls
|
|
build_args, _ = call.args
|
|
|
|
if want is None:
|
|
self.assertNotIn('-j', build_args)
|
|
else:
|
|
self.assert_contains_sequence(build_args, want)
|
|
|
|
@staticmethod
|
|
def assert_contains_sequence(sequence: Sequence[T], subsequence: Sequence[T]) -> None:
|
|
"""Raises an assertion if the subsequence is not contained in the sequence."""
|
|
if len(subsequence) == 0:
|
|
return # all sequences contain the empty subsequence
|
|
|
|
# Iterate over all windows of len(subsequence). Stop if the
|
|
# window matches.
|
|
for i in range(len(sequence) - len(subsequence) + 1):
|
|
candidate = sequence[i : i + len(subsequence)]
|
|
assert len(candidate) == len(subsequence) # sanity check
|
|
if candidate == subsequence:
|
|
return # found it
|
|
raise AssertionError(f'{subsequence} not found in {sequence}')
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def env_var(key: str, value: Optional[str]) -> Iterator[None]:
|
|
"""Sets/clears an environment variable within a Python context."""
|
|
# Get the previous value and then override it.
|
|
previous_value = os.environ.get(key)
|
|
set_env_var(key, value)
|
|
try:
|
|
yield
|
|
finally:
|
|
# Restore to previous value.
|
|
set_env_var(key, previous_value)
|
|
|
|
|
|
def set_env_var(key: str, value: Optional[str]) -> None:
|
|
"""Sets/clears an environment variable."""
|
|
if value is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = value
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|