def test_multiple_overlapping_clusters(loop): with DRMAACluster(scheduler_port=0) as cluster_1: cluster_1.start_workers(1) with Client(cluster_1, loop=loop) as client_1: with DRMAACluster(scheduler_port=0) as cluster_2: cluster_2.start_workers(1) with Client(cluster_2, loop=loop) as client_2: future_1 = client_1.submit(inc, 1) future_2 = client_2.submit(inc, 2) assert future_1.result() == 2 assert future_2.result() == 3
def main(nworkers): cluster = DRMAACluster(silence_logs=logging.INFO) cluster.start_workers(nworkers) def handle_signal(sig, frame): cluster.close() sys.exit(0) signal.signal(signal.SIGINT, handle_signal) signal.signal(signal.SIGTERM, handle_signal) while True: sleep(1)
def get_drmaa_cluster(): """ Instatiate a DRMAACluster for use with the LSF scheduler on the Janelia Research Campus compute cluster. This is a wrapper for dask_drmaa.DRMMACluster that uses reasonable default settings for the dask workers. Specifically, this ensures that dask workers use the /scratch/$USER directory for temporary files and also that each worker runs on a single core. This wrapper also directs the $WORKER.err and $WORKER.log files to /scratch/$USER. """ from dask_drmaa import DRMAACluster import os # we need these on each worker to prevent multithreaded numerical operations pre_exec = ( "export NUM_MKL_THREADS=1", "export OPENBLAS_NUM_THREADS=1", "export OPENMP_NUM_THREADS=1", "export OMP_NUM_THREADS=1", ) local_directory = "/scratch/" + os.environ["USER"] output_path = ":" + local_directory error_path = output_path cluster_kwargs_pass = {} cluster_kwargs_pass.setdefault( "template", { "args": ["--nthreads", "1", "--local-directory", local_directory], "jobEnvironment": os.environ, "outputPath": output_path, "errorPath": error_path, }, ) cluster_kwargs_pass["preexec_commands"] = pre_exec cluster = DRMAACluster(**cluster_kwargs_pass) return cluster
def test_pythonpath(): tmpdir = tempfile.mkdtemp(prefix='test_drmaa_pythonpath_', dir='.') try: with open(os.path.join(tmpdir, "bzzz_unlikely_module_name.py"), "w") as f: f.write("""if 1: def f(): return 5 """) def func(): import bzzz_unlikely_module_name return bzzz_unlikely_module_name.f() with DRMAACluster( scheduler_port=0, preexec_commands=['export PYTHONPATH=%s:PYTHONPATH' % tmpdir], ) as cluster: with Client(cluster) as client: cluster.start_workers(2) x = client.submit(func) assert x.result() == 5 finally: shutil.rmtree(tmpdir)
def test_str(loop): with DRMAACluster(scheduler_port=0) as cluster: cluster.start_workers(2) assert 'DRMAACluster' in str(cluster) assert 'DRMAACluster' in repr(cluster) assert '2' in str(cluster) assert '2' in repr(cluster) 1 + 1
def test_no_workers(loop): with DRMAACluster(scheduler_port=0) as cluster: with Client(cluster, loop=loop) as client: cluster.start_workers(0) assert not cluster.workers cluster.stop_workers([]) assert not os.path.exists(cluster.script)
def test_job_name_as_name(loop): with DRMAACluster(scheduler_port=0) as cluster: cluster.start_workers(2) while len(cluster.scheduler.workers) < 2: sleep(0.1) names = {cluster.scheduler.worker_info[w]['name'] for w in cluster.scheduler.workers} assert names == set(cluster.workers)
def test_logs(loop): with DRMAACluster(scheduler_port=0) as cluster: cluster.start_workers(2) while len(cluster.scheduler.workers) < 2: sleep(0.1) for w in cluster.workers: fn = 'worker.%s.err' % w assert os.path.exists(fn) with open(fn) as f: assert "worker" in f.read()
def test_simple(loop): with DRMAACluster(scheduler_port=0) as cluster: with Client(cluster, loop=loop) as client: cluster.start_workers(2) future = client.submit(lambda x: x + 1, 1) assert future.result() == 2 cluster.stop_workers(cluster.workers) start = time() while client.ncores(): sleep(0.2) assert time() < start + 60
def test_passed_script(loop): with tmpfile(extension='sh') as fn: with open(fn, 'w') as f: f.write(make_job_script(executable=worker_bin_path, name='foo')) os.chmod(fn, 0o777) with DRMAACluster(scheduler_port=0, script=fn) as cluster: tmp_script_location = cluster.script assert cluster.script.split(os.path.sep)[-1] == fn.split( os.path.sep)[-1] job = cluster.start_workers(1) with Client(cluster, loop=loop) as client: assert client.submit(lambda x: x + 1, 10).result() == 11 assert os.path.exists(fn) # doesn't cleanup provided script assert not os.path.exists(tmp_script_location)
def test_stop_workers_politely(loop): with DRMAACluster(scheduler_port=0) as cluster: with Client(cluster, loop=loop) as client: cluster.start_workers(2) while len(client.ncores()) < 2: sleep(0.1) futures = client.scatter(list(range(10))) a, b = cluster.workers cluster.stop_workers(a) data = client.gather(futures) assert data == list(range(10))
def setup_cluster(): from dask_drmaa import DRMAACluster import os cluster_kwargs_pass = {} cluster_kwargs_pass.setdefault( 'template', { 'args': [ '--nthreads', '1', '--local-directory', '/scratch/' + os.environ['USER']], 'jobEnvironment': os.environ } ) cluster = DRMAACluster(**cluster_kwargs_pass) return cluster
def test_stop_single_worker(loop): with DRMAACluster(scheduler_port=0) as cluster: with Client(cluster, loop=loop) as client: cluster.start_workers(2) future = client.submit(lambda x: x + 1, 1) assert future.result() == 2 while len(client.ncores()) < 2: sleep(0.1) a, b = cluster.workers local_dir = client.run(lambda dask_worker: dask_worker.local_dir, workers=[a])[a] assert os.path.exists(local_dir) cluster.stop_workers(a) start = time() while len(client.ncores()) != 1: sleep(0.2) assert time() < start + 60 assert not os.path.exists(local_dir)
def test_stdout_in_worker(): """ stdout and stderr should be redirected and line-buffered in workers. """ def inc_and_print(x): print("stdout: inc_and_print(%s)" % (x, )) print("stderr: inc_and_print(%s)" % (x, ), file=sys.stderr) return x + 1 def get_lines(fn): with open(fn) as f: return [line.strip() for line in f] with DRMAACluster(scheduler_port=0, diagnostics_port=None) as cluster: with Client(cluster) as client: cluster.start_workers(1) future = client.submit(inc_and_print, 1) assert future.result() == 2 w, = cluster.workers.values() assert "stdout: inc_and_print(1)" in get_lines(w.stdout) assert "stderr: inc_and_print(1)" in get_lines(w.stderr)
def test_ip(): ip = socket.gethostbyname(socket.gethostname()) with DRMAACluster(ip=ip, scheduler_port=0, diagnostics_port=None) as cluster: assert cluster.local_cluster.scheduler.ip == ip