executorlib/tests/executor/test_slurm_cluster.py at 7086c4e2458dbfc177d4e5a4c18ae8a2ba509a7a · pyiron/executorlib · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import os
import importlib
import unittest
import shutil

from executorlib import SlurmClusterExecutor
from executorlib.standalone.serialize import cloudpickle_register

if shutil.which("srun") is not None:
    skip_slurm_test = False
else:
    skip_slurm_test = True

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

try:
    from executorlib.standalone.hdf import dump

    skip_h5py_test = False
except ImportError:
    skip_h5py_test = True

try:
    import pysqa

    skip_pysqa_test = False
except ImportError:
    skip_pysqa_test = True

submission_template = """\
#!/bin/bash
#SBATCH --output=time.out
#SBATCH --job-name={{job_name}}
#SBATCH --chdir={{working_directory}}
#SBATCH --get-user-env=L
#SBATCH --ntasks={{cores}}

{{command}}
"""


def mpi_funct(i):
    from mpi4py import MPI

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


@unittest.skipIf(
    skip_slurm_test or skip_mpi4py_test or skip_h5py_test,
    "h5py or mpi4py or SLRUM are not installed, so the h5py, slurm and mpi4py tests are skipped.",
)
class TestCacheExecutorPysqa(unittest.TestCase):
    def test_executor(self):
        with SlurmClusterExecutor(
            resource_dict={"cores": 2, "cwd": "executorlib_cache", "submission_template": submission_template},
            block_allocation=False,
            cache_directory="executorlib_cache",
            pmi_mode="pmi2",
        ) as exe:
            cloudpickle_register(ind=1)
            fs1 = exe.submit(mpi_funct, 1)
            self.assertFalse(fs1.done())
            self.assertEqual(fs1.result(), [(1, 2, 0), (1, 2, 1)])
            self.assertEqual(len(os.listdir("executorlib_cache")), 3)
            self.assertTrue(fs1.done())

    def test_executor_no_cwd(self):
        with SlurmClusterExecutor(
            resource_dict={"cores": 2, "submission_template": submission_template},
            block_allocation=False,
            cache_directory="executorlib_cache",
            pmi_mode="pmi2",
        ) as exe:
            cloudpickle_register(ind=1)
            fs1 = exe.submit(mpi_funct, 1)
            self.assertFalse(fs1.done())
            self.assertEqual(fs1.result(), [(1, 2, 0), (1, 2, 1)])
            self.assertEqual(len(os.listdir("executorlib_cache")), 2)
            self.assertTrue(fs1.done())

    def test_executor_existing_files(self):
        with SlurmClusterExecutor(
            resource_dict={"cores": 2, "cwd": "executorlib_cache", "submission_template": submission_template},
            block_allocation=False,
            cache_directory="executorlib_cache",
            pmi_mode="pmi2",
        ) as exe:
            cloudpickle_register(ind=1)
            fs1 = exe.submit(mpi_funct, 1)
            self.assertFalse(fs1.done())
            self.assertEqual(fs1.result(), [(1, 2, 0), (1, 2, 1)])
            self.assertTrue(fs1.done())
            self.assertEqual(len(os.listdir("executorlib_cache")), 3)
            for file_name in os.listdir("executorlib_cache"):
                file_path = os.path.join("executorlib_cache", file_name )
                os.remove(file_path)
                if ".h5" in file_path:
                    task_key = file_path[:-5] + "_i.h5"
                    dump(file_name=task_key, data_dict={"a": 1})

        with SlurmClusterExecutor(
            resource_dict={"cores": 2, "cwd": "executorlib_cache", "submission_template": submission_template},
            block_allocation=False,
            cache_directory="executorlib_cache",
            pmi_mode="pmi2",
        ) as exe:
            cloudpickle_register(ind=1)
            fs1 = exe.submit(mpi_funct, 1)
            self.assertFalse(fs1.done())
            self.assertEqual(fs1.result(), [(1, 2, 0), (1, 2, 1)])
            self.assertTrue(fs1.done())
            self.assertEqual(len(os.listdir("executorlib_cache")), 3)

    def tearDown(self):
        shutil.rmtree("executorlib_cache", ignore_errors=True)


@unittest.skipIf(skip_pysqa_test, "pysqa is not installed, so the pysqa tests are skipped.")
class TestSlurmClusterInit(unittest.TestCase):
    def test_slurm_cluster_block_allocation(self):
        with self.assertRaises(ValueError):
            SlurmClusterExecutor(block_allocation=True)

    def test_slurm_cluster_file(self):
        self.assertTrue(SlurmClusterExecutor(block_allocation=False))