Function test_cpu_affinity_psutil

dask/tests/test_system.py:26–40

Test that cpu_count() respects CPU affinity set by psutil

(affinity)

@pytest.mark.skipif(MACOS, reason="No CPU affinity in psutil")
@pytest.mark.parametrize(("affinity"), [{0}, {1}, {0, 1}, {0, 2}])
def test_cpu_affinity_psutil(affinity):
    """Test that cpu_count() respects CPU affinity set by psutil"""
    psutil = pytest.importorskip("psutil")
    proc = psutil.Process()
    prev = proc.cpu_affinity()
    if prev is None:
        pytest.skip("No support for CPU affinity")
    if not affinity.issubset(set(prev)):
        pytest.skip("Not enough CPUs")  # pragma: no cover

    proc.cpu_affinity(affinity)
    try:
        assert cpu_count() == len(affinity)
    finally:
        proc.cpu_affinity(prev)


@pytest.mark.skipif(not LINUX, reason="No command line API for CPU affinity")