def test_configuration_full_specification(conda_env, tmpdir):
spec = _make_specification(environment=conda_env,
worker_memory='512 MiB',
scheduler_memory='512 MiB',
name='test-configuration-full-specification')
fn = os.path.join(str(tmpdir), 'spec.yaml')
with open(fn, 'w') as f:
f.write(spec.to_yaml())
# path to full specification
with dask.config.set({'yarn': {'specification': fn}}):
spec = _make_specification()
assert spec == spec
# full specification inlined
with dask.config.set({'yarn': {'specification': spec.to_dict()}}):
spec = _make_specification()
assert spec == spec
# when overrides specified, ignores full specification
override = {'specification': spec.to_dict()}
override.update(dask.config.get('yarn'))
override['name'] = 'config-name'
with dask.config.set({'yarn': override}):
spec = _make_specification(environment='foo',
name='test-name',
n_workers=4)
assert spec.services['dask.worker'].instances == 4
assert spec.name == 'test-name'
def test_configuration_full_specification(conda_env, tmpdir):
spec = _make_specification(
environment=conda_env,
worker_memory="512 MiB",
scheduler_memory="512 MiB",
name="test-configuration-full-specification",
)
fn = os.path.join(str(tmpdir), "spec.yaml")
with open(fn, "w") as f:
f.write(spec.to_yaml())
# path to full specification
with dask.config.set({"yarn": {"specification": fn}}):
spec = _make_specification()
assert spec == spec
# full specification inlined
with dask.config.set({"yarn": {"specification": spec.to_dict()}}):
spec = _make_specification()
assert spec == spec
# when overrides specified, ignores full specification
override = {"specification": spec.to_dict()}
override.update(dask.config.get("yarn"))
override["name"] = "config-name"
with dask.config.set({"yarn": override}):
spec = _make_specification(environment="foo",
name="test-name",
n_workers=4)
assert spec.services["dask.worker"].instances == 4
assert spec.name == "test-name"
def test_configuration():
config = {
'yarn': {
'environment': 'myenv.tar.gz',
'queue': 'myqueue',
'name': 'dask-yarn-tests',
'tags': ['a', 'b', 'c'],
'specification': None,
'worker': {
'memory': '1234MB',
'count': 1,
'vcores': 1,
'restarts': -1,
'env': {
'foo': 'bar'
}
},
'scheduler': {
'memory': '1234MB',
'vcores': 1
}
}
}
with dask.config.set(config):
spec = _make_specification()
assert spec.name == 'dask-yarn-tests'
assert spec.queue == 'myqueue'
assert spec.tags == {'a', 'b', 'c'}
assert spec.services['dask.worker'].resources.memory == 1234
assert spec.services['dask.worker'].env == {'foo': 'bar'}
assert spec.services['dask.scheduler'].resources.memory == 1234
def test_configuration(deploy_mode):
config = {
"yarn": {
"environment": "myenv.tar.gz",
"queue": "myqueue",
"name": "dask-yarn-tests",
"user": "******",
"tags": ["a", "b", "c"],
"specification": None,
"deploy-mode": deploy_mode,
"worker": {
"memory": "1234 MiB",
"count": 1,
"vcores": 1,
"restarts": -1,
"gpus": 2,
"env": {
"foo": "bar"
},
"worker_class": "abc",
"worker_options": {
"FOO": "BAZ"
},
},
"scheduler": {
"memory": "1234 MiB",
"vcores": 1,
"gpus": 0
},
"host": "0.0.0.0",
"port": 8786,
"dashboard_address": ":8787",
}
}
with dask.config.set(config):
spec = _make_specification()
assert spec.name == "dask-yarn-tests"
assert spec.user == "alice"
assert spec.queue == "myqueue"
assert spec.tags == {"a", "b", "c"}
assert spec.services["dask.worker"].resources.memory == 1234
assert spec.services["dask.worker"].resources.gpus == 2
assert spec.services["dask.worker"].env == {
"foo": "bar",
"worker_class": "abc",
"worker_options": "eyJGT08iOiAiQkFaIn0=",
}
if deploy_mode == "remote":
assert spec.services["dask.scheduler"].resources.memory == 1234
assert spec.services["dask.scheduler"].resources.gpus == 0
else:
assert "dask.scheduler" not in spec.services
kwargs = _make_scheduler_kwargs()
assert kwargs["host"] == "0.0.0.0"
assert kwargs["port"] == 8786
assert kwargs["dashboard_address"] == ":8787"
def test_from_specification(skein_client, conda_env, tmpdir, loop):
spec = _make_specification(environment=conda_env,
worker_memory='512 MB',
scheduler_memory='512 MB',
name=APPNAME)
fn = os.path.join(str(tmpdir), 'spec.yaml')
with open(fn, 'w') as f:
f.write(spec.to_yaml())
with YarnCluster.from_specification(fn, skein_client=skein_client) as cluster:
with Client(cluster, loop=loop):
pass
check_is_shutdown(skein_client, cluster.app_id)
def test_make_specification_errors():
with dask.config.set({"yarn.environment": None}):
with pytest.raises(ValueError) as info:
_make_specification()
assert "conda-pack" in str(info.value)
with pytest.raises(ValueError) as info:
_make_specification(environment="foo.tar.gz", deploy_mode="unknown")
assert "deploy_mode" in str(info.value)
with pytest.raises(ValueError) as info:
_make_specification(environment="foo.tar.gz", worker_memory=1234)
assert "1234 MiB" in str(info.value)
with pytest.raises(ValueError) as info:
_make_specification(environment="foo.tar.gz", scheduler_memory=1234)
assert "1234 MiB" in str(info.value)
def test_make_specification_errors():
with dask.config.set({'yarn.environment': None}):
with pytest.raises(ValueError) as info:
_make_specification()
assert 'conda-pack' in str(info.value)
with pytest.raises(ValueError) as info:
_make_specification(environment='foo.tar.gz', deploy_mode='unknown')
assert 'deploy_mode' in str(info.value)
with pytest.raises(ValueError) as info:
_make_specification(environment='foo.tar.gz', worker_memory=1234)
assert '1234 MiB' in str(info.value)
with pytest.raises(ValueError) as info:
_make_specification(environment='foo.tar.gz', scheduler_memory=1234)
assert '1234 MiB' in str(info.value)
def test_from_specification(skein_client, conda_env, tmpdir):
spec = _make_specification(
environment=conda_env,
worker_memory="512 MiB",
scheduler_memory="512 MiB",
name="dask-yarn-test-from-specification",
)
fn = os.path.join(str(tmpdir), "spec.yaml")
with open(fn, "w") as f:
f.write(spec.to_yaml())
with YarnCluster.from_specification(fn,
skein_client=skein_client) as cluster:
with Client(cluster):
pass
check_is_shutdown(skein_client, cluster.app_id)
def test_environment_relative_paths():
relpath = "path/to/foo.tar.gz"
abspath = os.path.abspath(relpath)
a = _make_specification(environment=relpath).to_dict()
b = _make_specification(environment=abspath).to_dict()
assert a == b
def test_environment_relative_paths(conda_env):
assert (_make_specification(
environment=conda_env).to_dict() == _make_specification(
environment=os.path.relpath(conda_env)).to_dict())
