def _test_multi_mpi_wrap(self):
if self.run_tests:
controller._close()
# Create the cluster
nodes = 1
fork_cluster = self.cluster(name="multifork_cluster",
fork_mpi=True,
nodes=nodes,
maximum_jobs=2,
**self.common_kwargs)
# Create the function that wraps tasks for this cluster
@on_cluster(cluster=fork_cluster)
@mpi_task(cluster_id=fork_cluster.name)
def mpi_wrap_task(**kwargs):
return mpi_wrap(**kwargs)
def test_function(script_path, return_wrapped_command=False):
t = mpi_wrap_task(
executable=self.executable,
exec_args=script_path,
return_wrapped_command=return_wrapped_command,
)
return t
self._multi_mpi_wrap_assert(test_function, nodes)
controller._close()
else:
pass
def _test_minimum_jobs(self):
if self.run_tests:
controller._close()
# Create the cluster
nodes = 1
fork_cluster = self.cluster(name="fork_cluster",
fork_mpi=True,
nodes=nodes,
minimum_jobs=2,
**self.common_kwargs)
# Create the function that wraps tasks for this cluster
@on_cluster(cluster=fork_cluster)
@mpi_task(cluster_id=fork_cluster.name)
def mpi_wrap_task(**kwargs):
return mpi_wrap(**kwargs)
def test_function(script_path, return_wrapped_command=False):
t = mpi_wrap_task(
executable=self.executable,
exec_args=script_path,
return_wrapped_command=return_wrapped_command,
)
return t
# Wait for a few seconds and the workers will start
time.sleep(5)
self.assertEqual(
len(fork_cluster.client.scheduler_info()["workers"]), 2)
controller._close()
else:
pass
def setUp(self):
# Kill any existing clusters
controller._close()
self.local_cluster = LocalCluster(name="test")
self.executable = "python"
self.script_path = os.path.abspath(
os.path.join(os.path.dirname(__file__), "resources",
"helloworld.py"))
self.number_of_processes = 4
@mpi_task(cluster_id="test")
def mpi_wrap_task(**kwargs):
return mpi_wrap(**kwargs)
@on_cluster(cluster=self.local_cluster, cluster_id="test")
def test_function(
script_path,
mpi_launcher=MPIEXEC,
launcher_args=None,
nodes=1,
ntasks_per_node=4,
cpus_per_task=1,
return_wrapped_command=False,
):
mpi_tasks = ntasks_per_node * nodes
t = mpi_wrap_task(
executable=self.executable,
exec_args=script_path,
mpi_launcher=mpi_launcher,
launcher_args=launcher_args,
mpi_tasks=mpi_tasks,
cpus_per_task=cpus_per_task,
ntasks_per_node=ntasks_per_node,
nodes=nodes,
return_wrapped_command=return_wrapped_command,
)
result = t.result()
return result
self.test_function = test_function
def mpi_task1(task_name):
comm = get_task_mpi_comm()
size = comm.Get_size()
# Since it is a return value it will only get printed by root
return "Running %d tasks of type %s." % (size, task_name)
self.mpi_task1 = mpi_task1
def string_task(string, kwarg_string=None):
return " ".join([s for s in [string, kwarg_string] if s])
self.string_task = string_task
def _test_single_mpi_tasks(self):
# Assume here we have srun support
if self.run_tests:
controller._close()
nodes = 2
custom_cluster = self.cluster(name="mpiCluster",
nodes=nodes,
**self.common_kwargs)
@on_cluster(cluster=custom_cluster)
@mpi_task(cluster_id=custom_cluster.name)
def task1(task_name):
from mpi4py import MPI
comm = get_task_mpi_comm()
size = comm.Get_size()
name = MPI.Get_processor_name()
all_nodes = comm.gather(name, root=0)
if all_nodes:
all_nodes = list(set(all_nodes))
all_nodes.sort()
else:
all_nodes = []
# Since it is a return value it will only get printed by root
return_string = "Running %d tasks of type %s on nodes %s." % (
size,
task_name,
all_nodes,
)
return return_string
@on_cluster(cluster=custom_cluster)
@mpi_task(cluster_id=custom_cluster.name)
def task2(name, task_name="default"):
comm = get_task_mpi_comm()
rank = comm.Get_rank()
# This only appears in the slurm job output
return_string = "Hi %s, my rank is %d for task of type %s" % (
name,
rank,
task_name,
)
return return_string
self._single_mpi_tasks_assert(task1, task2, nodes)
controller._close()
else:
pass
def setUp(self):
# Kill any existing clusters
controller._close()
self.number_of_processes_per_node = 2
self.run_tests = False
if which(self.default_mpi_launcher["launcher"]) is not None:
print(
"Found {} so assuming we have {}, running MPI test (with {})".
format(
self.default_mpi_launcher["launcher"],
self.cluster_name,
self.mpi_launcher,
))
self.run_tests = True
self.common_kwargs = {
"interface":
None,
"walltime":
"00:04:00",
"cores_per_node":
2,
"minimum_cores":
2,
"hyperthreading_factor":
1,
"ntasks_per_node":
2,
"memory":
self.memory,
"mpi_mode":
True,
"env_extra": [
"export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1",
"export OMPI_ALLOW_RUN_AS_ROOT=1",
],
"mpi_launcher":
self.mpi_launcher,
"local_directory":
"/tmp",
"queue":
self.queue_name,
}
self.executable = "python"
self.script_path = os.path.abspath(
os.path.join(os.path.dirname(__file__), "helloworld.py"))
def _test_multi_mpi_tasks(self):
#
# Assume here we have srun support
if self.run_tests:
controller._close()
# We only have 2 worker nodes so to have multiple jobs we need one worker
# per node
custom_cluster = self.cluster(name="mpiMultiCluster",
nodes=1,
maximum_jobs=2,
**self.common_kwargs)
@on_cluster(cluster=custom_cluster)
@mpi_task(cluster_id=custom_cluster.name)
def task(task_name):
import time
from mpi4py import MPI
comm = get_task_mpi_comm()
size = comm.Get_size()
name = MPI.Get_processor_name()
all_nodes = comm.gather(name, root=0)
if all_nodes:
all_nodes = list(set(all_nodes))
all_nodes.sort()
else:
all_nodes = []
# Since it is a return value it will only get printed by root
return_string = "Running %d tasks of type %s on nodes %s." % (
size,
task_name,
all_nodes,
)
# Add a sleep to make the task substantial enough to require scaling
time.sleep(1)
return return_string
self._multi_mpi_tasks_assert(task)
controller._close()
else:
pass
def tearDown(self):
# Kill any existing clusters
controller._close()
def tearDown(self):
controller._close()
def setUp(self):
# Kill any existing clusters
controller._close()
self.cluster_class = CustomPBSCluster
def tearDown(self):
clusters_controller_singleton._close()