def test_local_cluster_adapt(self):
adapt_kwargs = {"minimum": 1, "maximum": 1}
called_with = None
class MyCluster(distributed.LocalCluster):
def adapt(self, **kwargs):
nonlocal called_with
called_with = kwargs
super().adapt(**kwargs)
executor = DaskExecutor(
cluster_class=MyCluster,
cluster_kwargs={
"processes": False,
"n_workers": 0
},
adapt_kwargs=adapt_kwargs,
)
assert executor.adapt_kwargs == adapt_kwargs
with executor.start():
res = executor.wait(executor.submit(lambda x: x + 1, 1))
assert res == 2
assert called_with == adapt_kwargs
def test_debug_is_converted_to_silence_logs(self, monkeypatch):
client = MagicMock()
monkeypatch.setattr(prefect.engine.executors.dask, "Client", client)
# debug False
executor = DaskExecutor(debug=False)
with executor.start():
pass
assert client.called
assert client.call_args[-1]["silence_logs"] == logging.CRITICAL
# debug True
executor = DaskExecutor(debug=True)
with executor.start():
pass
assert client.called
assert client.call_args[-1]["silence_logs"] == logging.WARNING
def test_init_kwargs_are_passed_to_init(self, monkeypatch):
client = MagicMock()
monkeypatch.setattr(prefect.engine.executors.dask, "Client", client)
executor = DaskExecutor(test_kwarg="test_value")
with executor.start():
pass
assert client.called
assert client.call_args[-1]["test_kwarg"] == "test_value"
def test_context_tags_are_passed_to_map(self, monkeypatch):
client = MagicMock()
monkeypatch.setattr(prefect.engine.executors.dask, "Client", client)
executor = DaskExecutor()
with executor.start():
with prefect.context(task_tags=["dask-resource:GPU=1"]):
executor.map(lambda: None, [1, 2])
kwargs = client.return_value.__enter__.return_value.map.call_args[1]
assert kwargs["resources"] == {"GPU": 1.0}
def test_task_names_are_passed_to_map(self, monkeypatch):
client = MagicMock()
monkeypatch.setattr(prefect.engine.executors.dask, "Client", client)
executor = DaskExecutor()
with executor.start():
with prefect.context(task_full_name="FISH![0]"):
executor.map(lambda: None, [1, 2])
kwargs = client.return_value.__enter__.return_value.map.call_args[1]
assert kwargs["key"].startswith("FISH![0]")
def test_context_tags_are_passed_to_submit(self, monkeypatch):
client = MagicMock()
monkeypatch.setattr(distributed, "Client", client)
executor = DaskExecutor()
with executor.start():
with prefect.context(task_tags=["dask-resource:GPU=1"]):
executor.submit(lambda: None)
kwargs = client.return_value.__enter__.return_value.submit.call_args[1]
assert kwargs["resources"] == {"GPU": 1.0}
def test_task_names_are_passed_to_submit(self, monkeypatch):
client = MagicMock()
monkeypatch.setattr(distributed, "Client", client)
executor = DaskExecutor()
with executor.start():
with prefect.context(task_full_name="FISH!"):
executor.submit(lambda: None)
kwargs = client.return_value.__enter__.return_value.submit.call_args[1]
assert kwargs["key"].startswith("FISH!")
def test_start_local_cluster(self):
executor = DaskExecutor(cluster_kwargs={"processes": False})
assert executor.cluster_class == distributed.LocalCluster
assert executor.cluster_kwargs == {
"processes": False,
"silence_logs": logging.CRITICAL,
}
with executor.start():
res = executor.wait(executor.submit(lambda x: x + 1, 1))
assert res == 2
def test_connect_to_running_cluster(self):
with distributed.Client(processes=False, set_as_default=False) as client:
executor = DaskExecutor(address=client.scheduler.address)
assert executor.address == client.scheduler.address
assert executor.cluster_class is None
assert executor.cluster_kwargs is None
assert executor.client_kwargs == {"set_as_default": False}
with executor.start():
res = executor.wait(executor.submit(lambda x: x + 1, 1))
assert res == 2
def test_exit_early_with_external_or_inproc_cluster_waits_for_pending_futures(
self, kind, monkeypatch
):
key = "TESTING_%s" % uuid.uuid4().hex
monkeypatch.setenv(key, "initial")
def slow():
time.sleep(0.5)
os.environ[key] = "changed"
def pending(x):
# This function shouldn't ever start, since it's pending when the
# shutdown signal is received
os.environ[key] = "changed more"
if kind == "external":
with distributed.Client(processes=False, set_as_default=False) as client:
executor = DaskExecutor(address=client.scheduler.address)
with executor.start():
fut = executor.submit(slow)
fut2 = executor.submit(pending, fut) # noqa
time.sleep(0.2)
assert os.environ[key] == "changed"
elif kind == "inproc":
executor = DaskExecutor(cluster_kwargs={"processes": False})
with executor.start():
fut = executor.submit(slow)
fut2 = executor.submit(pending, fut) # noqa
time.sleep(0.2)
assert os.environ[key] == "changed"
assert executor.client is None
assert executor._futures is None
assert executor._should_run_event is None
def test_executor_logs_worker_events(self, caplog):
caplog.set_level(logging.DEBUG, logger="prefect")
with distributed.Client(
n_workers=1, processes=False, set_as_default=False
) as client:
executor = DaskExecutor(address=client.scheduler.address)
with executor.start():
client.cluster.scale(4)
while len(client.scheduler_info()["workers"]) < 4:
time.sleep(0.1)
client.cluster.scale(1)
while len(client.scheduler_info()["workers"]) > 1:
time.sleep(0.1)
assert any("Worker %s added" == rec.msg for rec in caplog.records)
assert any("Worker %s removed" == rec.msg for rec in caplog.records)
def test_temporary_cluster_forcefully_cancels_pending_tasks(self, tmpdir):
filname = tmpdir.join("signal")
def slow():
time.sleep(10)
with open(filname, "w") as f:
f.write("Got here")
executor = DaskExecutor()
with executor.start():
start = time.time()
fut = executor.submit(slow) # noqa
time.sleep(0.1)
stop = time.time()
# Cluster shutdown before task could complete
assert stop - start < 5
assert not os.path.exists(filname)
