test_spawner.py

from concurrent.futures import CancelledError, Future
import importlib.util
from queue import Queue
from time import sleep
import shutil
import os
from sys import platform
import unittest

import numpy as np

from executorlib.task_scheduler.base import TaskSchedulerBase
from executorlib.standalone.interactive.spawner import MpiExecSpawner, set_current_directory_in_environment
from executorlib.task_scheduler.interactive.blockallocation import BlockAllocationTaskScheduler, _execute_multiple_tasks
from executorlib.task_scheduler.interactive.onetoone import OneProcessTaskScheduler
from executorlib.standalone.interactive.backend import call_funct
from executorlib.standalone.serialize import cloudpickle_register

try:
    import h5py

    skip_h5py_test = False
except ImportError:
    skip_h5py_test = True

skip_mpi4py_test = importlib.util.find_spec("mpi4py") is None


def calc(i):
    return i


def calc_array(i):
    return np.array(i**2)


def echo_funct(i):
    return i


def get_global(memory=None):
    return memory


def set_global():
    return {"memory": np.array([5])}


def mpi_funct(i):
    from mpi4py import MPI

    size = MPI.COMM_WORLD.Get_size()
    rank = MPI.COMM_WORLD.Get_rank()
    return i, size, rank


def raise_error():
    raise RuntimeError


def sleep_one(i):
    sleep(1)
    return i


class TestBlockAllocationTaskSchedulerSerial(unittest.TestCase):
    def test_two_workers_submit_serial_tasks(self):
        with BlockAllocationTaskScheduler(
            max_workers=2,
            executor_kwargs={},
            spawner=MpiExecSpawner,
        ) as exe:
            cloudpickle_register(ind=1)
            fs_1 = exe.submit(calc, 1)
            fs_2 = exe.submit(calc, 2)
            self.assertEqual(fs_1.result(), 1)
            self.assertEqual(fs_2.result(), 2)
            self.assertTrue(fs_1.done())
            self.assertTrue(fs_2.done())

    def test_max_workers(self):
        with BlockAllocationTaskScheduler(
            max_workers=2,
            executor_kwargs={},
            spawner=MpiExecSpawner,
        ) as exe:
            self.assertEqual(exe.max_workers, 2)

    def test_one_worker_submit_serial_tasks(self):
        with BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={},
            spawner=MpiExecSpawner,
        ) as exe:
            cloudpickle_register(ind=1)
            fs_1 = exe.submit(calc, 1)
            fs_2 = exe.submit(calc, 2)
            self.assertEqual(fs_1.result(), 1)
            self.assertEqual(fs_2.result(), 2)
            self.assertTrue(fs_1.done())
            self.assertTrue(fs_2.done())


class TestOneProcessTaskSchedulerSerial(unittest.TestCase):
    def test_two_workers_submit_serial_tasks(self):
        with OneProcessTaskScheduler(
            max_cores=2,
            executor_kwargs={},
            spawner=MpiExecSpawner,
        ) as exe:
            cloudpickle_register(ind=1)
            fs_1 = exe.submit(calc, 1)
            fs_2 = exe.submit(calc, 2)
            self.assertEqual(fs_1.result(), 1)
            self.assertEqual(fs_2.result(), 2)
            self.assertTrue(fs_1.done())
            self.assertTrue(fs_2.done())

    def test_max_workers(self):
        with OneProcessTaskScheduler(
            max_workers=2,
            executor_kwargs={},
            spawner=MpiExecSpawner,
        ) as exe:
            self.assertEqual(exe.max_workers, 2)

    def test_one_worker_submit_serial_tasks(self):
        with OneProcessTaskScheduler(
            max_cores=1,
            executor_kwargs={},
            spawner=MpiExecSpawner,
        ) as exe:
            cloudpickle_register(ind=1)
            fs_1 = exe.submit(calc, 1)
            fs_2 = exe.submit(calc, 2)
            self.assertEqual(fs_1.result(), 1)
            self.assertEqual(fs_2.result(), 2)
            self.assertTrue(fs_1.done())
            self.assertTrue(fs_2.done())


@unittest.skipIf(
    skip_mpi4py_test, "mpi4py is not installed, so the mpi4py tests are skipped."
)
class TestBlockAllocationTaskSchedulerMPI(unittest.TestCase):
    def test_block_allocation_mpi_two_cores(self):
        with BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={"cores": 2},
            spawner=MpiExecSpawner,
        ) as exe:
            cloudpickle_register(ind=1)
            fs_1 = exe.submit(mpi_funct, 1)
            self.assertEqual(fs_1.result(), [(1, 2, 0), (1, 2, 1)])
            self.assertTrue(fs_1.done())

    def test_block_allocation_mpi_multiple_submissions(self):
        with BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={"cores": 2},
            spawner=MpiExecSpawner,
        ) as p:
            cloudpickle_register(ind=1)
            fs1 = p.submit(mpi_funct, 1)
            fs2 = p.submit(mpi_funct, 2)
            fs3 = p.submit(mpi_funct, 3)
            output = [
                fs1.result(),
                fs2.result(),
                fs3.result(),
            ]
        self.assertEqual(
            output,
            [[(1, 2, 0), (1, 2, 1)], [(2, 2, 0), (2, 2, 1)], [(3, 2, 0), (3, 2, 1)]],
        )

    def test_block_allocation_mpi_echo_broadcast(self):
        with BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={"cores": 2},
            spawner=MpiExecSpawner,
        ) as p:
            cloudpickle_register(ind=1)
            output = p.submit(echo_funct, 2).result()
        self.assertEqual(output, [2, 2])


@unittest.skipIf(
    skip_mpi4py_test, "mpi4py is not installed, so the mpi4py tests are skipped."
)
class TestOneProcessTaskSchedulerMPI(unittest.TestCase):
    def test_one_process_mpi_two_cores(self):
        with OneProcessTaskScheduler(
            max_cores=2,
            executor_kwargs={"cores": 2},
            spawner=MpiExecSpawner,
        ) as exe:
            cloudpickle_register(ind=1)
            fs_1 = exe.submit(mpi_funct, 1)
            self.assertEqual(fs_1.result(), [(1, 2, 0), (1, 2, 1)])
            self.assertTrue(fs_1.done())

    def test_one_process_mpi_multiple_submissions(self):
        with OneProcessTaskScheduler(
            max_cores=2,
            executor_kwargs={"cores": 2},
            spawner=MpiExecSpawner,
        ) as p:
            cloudpickle_register(ind=1)
            fs1 = p.submit(mpi_funct, 1)
            fs2 = p.submit(mpi_funct, 2)
            fs3 = p.submit(mpi_funct, 3)
            output = [
                fs1.result(),
                fs2.result(),
                fs3.result(),
            ]
        self.assertEqual(
            output,
            [[(1, 2, 0), (1, 2, 1)], [(2, 2, 0), (2, 2, 1)], [(3, 2, 0), (3, 2, 1)]],
        )

    def test_one_process_mpi_echo_broadcast(self):
        with OneProcessTaskScheduler(
            max_cores=2,
            executor_kwargs={"cores": 2},
            spawner=MpiExecSpawner,
        ) as p:
            cloudpickle_register(ind=1)
            output = p.submit(echo_funct, 2).result()
        self.assertEqual(output, [2, 2])


class TestBlockAllocationTaskSchedulerInitFunction(unittest.TestCase):
    def test_internal_memory(self):
        with BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={
                "cores": 1,
                "init_function": set_global,
            },
            spawner=MpiExecSpawner,
        ) as p:
            f = p.submit(get_global)
            self.assertFalse(f.done())
            self.assertEqual(f.result(), np.array([5]))
            self.assertTrue(f.done())

    def test_call_funct(self):
        self.assertEqual(
            call_funct(
                input_dict={"fn": get_global, "args": (), "kwargs": {}},
                memory={"memory": 4},
            ),
            4,
        )

    def test_execute_task(self):
        f = Future()
        q = Queue()
        q.put({"fn": get_global, "args": (), "kwargs": {}, "future": f})
        q.put({"shutdown": True, "wait": True})
        cloudpickle_register(ind=1)
        _execute_multiple_tasks(
            future_queue=q,
            cores=1,
            openmpi_oversubscribe=False,
            spawner=MpiExecSpawner,
            init_function=set_global,
        )
        self.assertEqual(f.result(), np.array([5]))
        q.join()

    @unittest.skipIf(
        skip_mpi4py_test, "mpi4py is not installed, so the mpi4py tests are skipped."
    )
    def test_internal_memory_mpi(self):
        with BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={
                "cores": 2,
                "init_function": set_global,
            },
            spawner=MpiExecSpawner,
        ) as p:
            cloudpickle_register(ind=1)
            f = p.submit(get_global)
            result = f.result()
            self.assertEqual(len(result), 2)
            np.testing.assert_array_equal(result[0], np.array([5]))
            np.testing.assert_array_equal(result[1], np.array([5]))


class TestBlockAllocationTaskScheduler(unittest.TestCase):
    def test_submit_tracks_future_state(self):
        with BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={"cores": 1},
            spawner=MpiExecSpawner,
        ) as p:
            output = p.submit(calc_array, i=2)
            self.assertEqual(len(p), 1)
            self.assertTrue(isinstance(output, Future))
            self.assertFalse(output.done())
            sleep(1)
            self.assertTrue(output.done())
            self.assertEqual(len(p), 0)
        self.assertEqual(output.result(), np.array(4))

    def test_executor_multi_submission(self):
        with BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={"cores": 1},
            spawner=MpiExecSpawner,
        ) as p:
            fs_1 = p.submit(calc_array, i=2)
            fs_2 = p.submit(calc_array, i=2)
            self.assertEqual(fs_1.result(), np.array(4))
            self.assertEqual(fs_2.result(), np.array(4))
            self.assertTrue(fs_1.done())
            self.assertTrue(fs_2.done())

    @unittest.skipIf(platform == "darwin", "Skipping test on macOS due to known issues")
    def test_shutdown(self):
        p = BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={"cores": 1},
            spawner=MpiExecSpawner,
        )
        fs1 = p.submit(sleep_one, i=2)
        fs2 = p.submit(sleep_one, i=4)
        sleep(1)
        p.shutdown(wait=True, cancel_futures=True)
        self.assertTrue(fs1.done())
        self.assertTrue(fs2.done())
        self.assertEqual(fs1.result(), 2)
        with self.assertRaises(CancelledError):
            fs2.result()

    def test_map_returns_array_results(self):
        with BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={"cores": 1},
            spawner=MpiExecSpawner,
        ) as p:
            output = list(p.map(calc_array, [1, 2, 3]))
        self.assertEqual(output, [np.array(1), np.array(4), np.array(9)])

    def test_executor_exception(self):
        with self.assertRaises(RuntimeError):
            with BlockAllocationTaskScheduler(
                max_workers=1,
                executor_kwargs={"cores": 1},
                spawner=MpiExecSpawner,
            ) as p:
                fs = p.submit(raise_error)
                fs.result()

    def test_executor_exception_future(self):
        with self.assertRaises(RuntimeError):
            with BlockAllocationTaskScheduler(
                max_workers=1,
                executor_kwargs={"cores": 1},
                spawner=MpiExecSpawner,
            ) as p:
                fs = p.submit(raise_error)
                fs.result()

    @unittest.skipIf(
        skip_mpi4py_test, "mpi4py is not installed, so the mpi4py tests are skipped."
    )
    def test_block_allocation_task_scheduler_info(self):
        meta_data_exe_dict = {
            "cores": 2,
            "spawner": "<class 'executorlib.standalone.interactive.spawner.MpiExecSpawner'>",
            "hostname_localhost": True,
            "init_function": None,
            "cwd": None,
            "openmpi_oversubscribe": False,
            "max_workers": 1,
        }
        with BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={
                "cores": 2,
                "hostname_localhost": True,
                "init_function": None,
                "cwd": None,
                "openmpi_oversubscribe": False,
            },
            spawner=MpiExecSpawner,
        ) as exe:
            for k, v in meta_data_exe_dict.items():
                if k != "spawner":
                    self.assertEqual(exe.info[k], v)
                else:
                    self.assertEqual(str(exe.info[k]), v)
        with TaskSchedulerBase() as exe:
            self.assertIsNone(exe.info)

    def test_one_process_task_scheduler_info(self):
        meta_data_exe_dict = {
            "cores": 2,
            "spawner": "<class 'executorlib.standalone.interactive.spawner.MpiExecSpawner'>",
            "hostname_localhost": True,
            "cwd": None,
            "openmpi_oversubscribe": False,
            "max_cores": 2,
        }
        with OneProcessTaskScheduler(
            max_cores=2,
            executor_kwargs={
                "cores": 2,
                "hostname_localhost": True,
                "cwd": None,
                "openmpi_oversubscribe": False,
            },
            spawner=MpiExecSpawner,
        ) as exe:
            for k, v in meta_data_exe_dict.items():
                if k != "spawner":
                    self.assertEqual(exe.info[k], v)
                else:
                    self.assertEqual(str(exe.info[k]), v)

    @unittest.skipIf(
        skip_mpi4py_test, "mpi4py is not installed, so the mpi4py tests are skipped."
    )
    def test_submit_mpi_task_tracks_future_state(self):
        with BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={"cores": 2},
            spawner=MpiExecSpawner,
        ) as p:
            output = p.submit(mpi_funct, i=2)
            self.assertEqual(len(p), 1)
            self.assertTrue(isinstance(output, Future))
            self.assertFalse(output.done())
            sleep(2)
            self.assertTrue(output.done())
            self.assertEqual(len(p), 0)
        self.assertEqual(output.result(), [(2, 2, 0), (2, 2, 1)])

    @unittest.skipIf(
        skip_mpi4py_test, "mpi4py is not installed, so the mpi4py tests are skipped."
    )
    def test_map_mpi_tasks(self):
        with BlockAllocationTaskScheduler(
            max_workers=1,
            executor_kwargs={"cores": 2},
            spawner=MpiExecSpawner,
        ) as p:
            output = list(p.map(mpi_funct, [1, 2, 3]))
        self.assertEqual(
            output,
            [[(1, 2, 0), (1, 2, 1)], [(2, 2, 0), (2, 2, 1)], [(3, 2, 0), (3, 2, 1)]],
        )

    def test_execute_task_failed_no_argument(self):
        f = Future()
        q = Queue()
        q.put({"fn": calc_array, "args": (), "kwargs": {}, "future": f})
        q.put({"shutdown": True, "wait": True})
        cloudpickle_register(ind=1)
        _execute_multiple_tasks(
            future_queue=q,
            cores=1,
            openmpi_oversubscribe=False,
            spawner=MpiExecSpawner,
        )
        with self.assertRaises(TypeError):
            f.result()
        q.join()

    def test_execute_task_failed_wrong_argument(self):
        f = Future()
        q = Queue()
        q.put({"fn": calc_array, "args": (), "kwargs": {"j": 4}, "future": f})
        q.put({"shutdown": True, "wait": True})
        cloudpickle_register(ind=1)
        _execute_multiple_tasks(
            future_queue=q,
            cores=1,
            openmpi_oversubscribe=False,
            spawner=MpiExecSpawner,
        )
        with self.assertRaises(TypeError):
            f.result()
        q.join()

    def test_execute_task(self):
        f = Future()
        q = Queue()
        q.put({"fn": calc_array, "args": (), "kwargs": {"i": 2}, "future": f})
        q.put({"shutdown": True, "wait": True})
        cloudpickle_register(ind=1)
        _execute_multiple_tasks(
            future_queue=q,
            cores=1,
            openmpi_oversubscribe=False,
            spawner=MpiExecSpawner,
        )
        self.assertEqual(f.result(), np.array(4))
        q.join()

    @unittest.skipIf(
        skip_mpi4py_test, "mpi4py is not installed, so the mpi4py tests are skipped."
    )
    def test_execute_task_parallel(self):
        f = Future()
        q = Queue()
        q.put({"fn": calc_array, "args": (), "kwargs": {"i": 2}, "future": f})
        q.put({"shutdown": True, "wait": True})
        cloudpickle_register(ind=1)
        _execute_multiple_tasks(
            future_queue=q,
            cores=2,
            openmpi_oversubscribe=False,
            spawner=MpiExecSpawner,
        )
        self.assertEqual(f.result(), [np.array(4), np.array(4)])
        q.join()


class TestBlockAllocationTaskSchedulerCache(unittest.TestCase):
    def tearDown(self):
        shutil.rmtree("executorlib_cache", ignore_errors=True)

    @unittest.skipIf(
        skip_h5py_test, "h5py is not installed, so the h5py tests are skipped."
    )
    def test_execute_task_cache(self):
        f = Future()
        q = Queue()
        q.put({"fn": calc, "args": (), "kwargs": {"i": 1}, "future": f})
        q.put({"shutdown": True, "wait": True})
        cloudpickle_register(ind=1)
        _execute_multiple_tasks(
            future_queue=q,
            cores=1,
            openmpi_oversubscribe=False,
            spawner=MpiExecSpawner,
            cache_directory="executorlib_cache",
        )
        self.assertEqual(f.result(), 1)
        q.join()

    @unittest.skipIf(
        skip_h5py_test, "h5py is not installed, so the h5py tests are skipped."
    )
    def test_execute_task_cache_failed_no_argument(self):
        f = Future()
        q = Queue()
        q.put({"fn": calc_array, "args": (), "kwargs": {}, "future": f})
        q.put({"shutdown": True, "wait": True})
        cloudpickle_register(ind=1)
        _execute_multiple_tasks(
            future_queue=q,
            cores=1,
            openmpi_oversubscribe=False,
            spawner=MpiExecSpawner,
            cache_directory="executorlib_cache",
        )
        with self.assertRaises(TypeError):
            f.result()
        q.join()


class TestEnvManipulation(unittest.TestCase):
    def test_set_current_directory_in_environment(self):
        env = os.environ
        if "PYTHONPATH" in env:
            python_path = env["PYTHONPATH"]
            del env["PYTHONPATH"]
        else:
            python_path = None
        self.assertFalse("PYTHONPATH" in env)
        set_current_directory_in_environment()
        self.assertTrue("PYTHONPATH" in env)
        self.assertEqual(env["PYTHONPATH"], os.getcwd())
        env["PYTHONPATH"] = "/my/special/path"
        set_current_directory_in_environment()
        self.assertEqual(env["PYTHONPATH"], os.getcwd() + ":/my/special/path")
        if python_path is not None:
            env["PYTHONPATH"] = python_path