def test_init_both_cluster_and_id(self):
_id = "test1"
original_cluster = CustomSLURMCluster()
with self.assertRaises(AssertionError):
@on_cluster(cluster_id=_id, cluster=original_cluster)
def f():
pass
controller.delete_cluster(original_cluster.name)
original_cluster = CustomSLURMCluster()
@on_cluster(cluster_id=original_cluster.name, cluster=original_cluster)
def f():
pass
cluster, client = controller.get_cluster(id_=original_cluster.name)
self.assertEqual(original_cluster, cluster)
self.assertEqual(cluster.name, original_cluster.name)
self.assertIsInstance(client, Client)
controller.delete_cluster(original_cluster.name)
original_cluster = CustomSLURMCluster(name=_id)
@on_cluster(cluster_id=_id, cluster=original_cluster)
def f():
pass
cluster, client = controller.get_cluster(id_=_id)
self.assertEqual(original_cluster, cluster)
self.assertEqual(cluster.name, _id)
self.assertIsInstance(client, Client)
controller.delete_cluster(_id)
def test_init_only_cluster(self):
original_cluster = CustomSLURMCluster()
@on_cluster(cluster=original_cluster)
def f():
pass
cluster, client = controller.get_cluster(original_cluster.name)
self.assertEqual(original_cluster, cluster)
self.assertIsInstance(client, Client)
controller.delete_cluster(cluster.name)
_id = "test1"
original_cluster = CustomSLURMCluster(name=_id)
@on_cluster(cluster=original_cluster)
def f():
pass
cluster, client = controller.get_cluster(_id)
self.assertEqual(original_cluster, cluster)
self.assertIsInstance(client, Client)
controller.delete_cluster(cluster.name)
def test_init_only_id(self):
_id = "test1"
original_cluster = CustomSLURMCluster(name=_id)
with self.assertRaises(ClusterException) as ctx:
controller.add_cluster(cluster=original_cluster)
self.assertEqual('Cluster "{}" already exists!'.format(_id),
str(ctx.exception))
@on_cluster(cluster_id=_id)
def f():
pass
cluster, client = controller.get_cluster(id_=_id)
self.assertEqual(original_cluster, cluster)
self.assertEqual(cluster.name, _id)
self.assertIsInstance(client, Client)
controller.delete_cluster(_id)
GPU = True
set_default_cluster(CustomSLURMCluster)
if GPU:
GROMACS_gpu_cluster = CustomSLURMCluster(
name="GROMACS_gpu_cluster",
walltime="00:15:00",
nodes=2,
mpi_mode=True,
fork_mpi=True,
queue_type="gpus",
maximum_scale=5,
env_extra=[
"module --force purge",
"module use /usr/local/software/jureca/OtherStages",
"module load Stages/Devel-2019a",
"module load Intel",
"module load ParaStationMPI",
"module load GROMACS",
"module load GPUtil", # Only required for our hello_world2.py example
"module load dask",
"module load jobqueue_features",
],
)
if KNL:
GROMACS_knl_cluster = CustomSLURMCluster(
name="GROMACS_knl_cluster",
walltime="00:15:00",
nodes=4,
# logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.DEBUG)
def inc(x):
time.sleep(1)
return x + 1
def mul(x):
time.sleep(1)
return x * 2
@on_cluster(
cluster_id="inc-cluster",
cluster=CustomSLURMCluster(name="inc-cluster", walltime="00:04:00"),
scale=2,
)
@task(cluster_id="inc-cluster")
def increment_task(x):
return inc(x)
@on_cluster(
cluster_id="mul-cluster",
cluster=CustomSLURMCluster(name="mul-cluster", walltime="00:04:00"),
scale=2,
)
@task(cluster_id="mul-cluster")
def multiplication_task(x):
return mul(x)
from __future__ import print_function
import sys
from jobqueue_features.clusters import CustomSLURMCluster
from jobqueue_features.decorators import on_cluster, mpi_task
from jobqueue_features.mpi_wrapper import SRUN
# import logging
# logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.DEBUG)
custom_cluster = CustomSLURMCluster(name="mpiCluster",
walltime="00:03:00",
nodes=2,
mpi_mode=True,
mpi_launcher=SRUN)
@mpi_task(cluster_id="mpiCluster")
def task1(task_name):
from mpi4py import MPI
comm = MPI.COMM_WORLD
size = comm.Get_size()
name = MPI.Get_processor_name()
all_nodes = comm.gather(name, root=0)
if all_nodes:
all_nodes = set(all_nodes)
else:
all_nodes = []
# Since it is a return value it will only get printed by root
# import logging
# logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.DEBUG)
# set_default_cluster(LocalCluster) # set LocalCluster as default cluster type
set_default_cluster(
CustomSLURMCluster) # set CustomSLURMCluster as default cluster type
@on_cluster()
@task()
def square(x):
return x**2
@on_cluster(cluster_id="other",
cluster=CustomSLURMCluster(name="other", walltime="00:04:00"))
@task(cluster_id="other")
def inc(x):
return x + 1
def main():
sq_tasks = list(map(square, range(1, 11)))
inc_tasks = list(map(inc, range(1, 11)))
print([t.result() for t in sq_tasks])
print([t.result() for t in inc_tasks])
if __name__ == "__main__":
start = time.time()
main()