executorlib/tests/executor/test_api.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
import os
import shutil
import unittest

from executorlib import get_cache_data
from executorlib.api import TestClusterExecutor
from executorlib.task_scheduler.interactive.dependency_plot import generate_nodes_and_edges_for_plotting
from executorlib.standalone.serialize import cloudpickle_register

try:
    import h5py

    skip_h5py_test = False
except ImportError:
    skip_h5py_test = True


def add_function(parameter_1, parameter_2):
    return parameter_1 + parameter_2


def foo(x):
    return x + 1


@unittest.skipIf(
    skip_h5py_test, "h5py is not installed, so the h5io tests are skipped."
)
class TestTestClusterExecutor(unittest.TestCase):
    def test_cache_dir(self):
        with TestClusterExecutor(cache_directory="not_this_dir", resource_dict={}) as exe:
            cloudpickle_register(ind=1)
            future = exe.submit(
                foo,
                1,
                resource_dict={
                    "cache_directory": "rather_this_dir",
                    "cache_key": "foo",
                },
            )
            self.assertEqual(future.result(), 2)
        self.assertFalse(os.path.exists("not_this_dir"))
        cache_lst = get_cache_data(cache_directory="not_this_dir")
        self.assertEqual(len(cache_lst), 0)
        self.assertTrue(os.path.exists("rather_this_dir"))
        cache_lst = get_cache_data(cache_directory="rather_this_dir")
        self.assertEqual(len(cache_lst), 1)
        with TestClusterExecutor(cache_directory="not_this_dir", resource_dict={}) as exe:
            cloudpickle_register(ind=1)
            future = exe.submit(
                foo,
                1,
                resource_dict={
                    "cache_directory": "rather_this_dir",
                    "cache_key": "foo",
                },
            )
            self.assertEqual(future.result(), 2)
        self.assertFalse(os.path.exists("not_this_dir"))
        cache_lst = get_cache_data(cache_directory="not_this_dir")
        self.assertEqual(len(cache_lst), 0)
        self.assertTrue(os.path.exists("rather_this_dir"))
        cache_lst = get_cache_data(cache_directory="rather_this_dir")
        self.assertEqual(len(cache_lst), 1)

    def test_empty(self):
        with TestClusterExecutor(cache_directory="rather_this_dir") as exe:
            cloudpickle_register(ind=1)
            future = exe.submit(foo,1)
            self.assertEqual(future.result(), 2)
        self.assertTrue(os.path.exists("rather_this_dir"))
        cache_lst = get_cache_data(cache_directory="rather_this_dir")
        self.assertEqual(len(cache_lst), 1)

    def test_executor_dependency_plot(self):
        with TestClusterExecutor(
            plot_dependency_graph=True,
        ) as exe:
            cloudpickle_register(ind=1)
            future_1 = exe.submit(add_function, 1, parameter_2=2)
            future_2 = exe.submit(add_function, 1, parameter_2=future_1)
            self.assertTrue(future_1.done())
            self.assertTrue(future_2.done())
            self.assertEqual(len(exe._task_scheduler._future_hash_dict), 2)
            self.assertEqual(len(exe._task_scheduler._task_hash_dict), 2)
            nodes, edges = generate_nodes_and_edges_for_plotting(
                task_hash_dict=exe._task_scheduler._task_hash_dict,
                future_hash_inverse_dict={
                    v: k for k, v in exe._task_scheduler._future_hash_dict.items()
                },
            )
            self.assertEqual(len(nodes), 4)
            self.assertEqual(len(edges), 4)

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