async def test_cluster_options():
config = Config()
config.InProcessBackend.cluster_config_class = MyClusterConfig
config.InProcessBackend.cluster_options = options.Options(
options.Integer("option_one", default=1, min=1, max=4, target="option_one_b"),
options.Select("option_two", options=[("small", 1.5), ("large", 15)]),
)
async with temp_gateway(config=config) as g:
async with g.gateway_client() as gateway:
# Create with no parameters
async with gateway.new_cluster() as cluster:
report = await gateway.get_cluster(cluster.name)
assert report.options == {"option_one": 1, "option_two": "small"}
# Create with parameters
async with gateway.new_cluster(option_two="large") as cluster:
report = await gateway.get_cluster(cluster.name)
assert report.options == {"option_one": 1, "option_two": "large"}
# With options object
opts = await gateway.cluster_options()
opts.option_one = 2
async with gateway.new_cluster(opts, option_two="large") as cluster:
report = await gateway.get_cluster(cluster.name)
assert report.options == {"option_one": 2, "option_two": "large"}
# Bad parameters
with pytest.raises(TypeError):
await gateway.new_cluster(cluster_options=10)
with pytest.raises(ValueError) as exc:
await gateway.new_cluster(option_two="medium")
assert "option_two" in str(exc.value)
async def test_cluster_options_client_config(monkeypatch):
monkeypatch.setenv("TEST_OPTION_TWO", "large")
config = Config()
config.InProcessBackend.cluster_config_class = MyClusterConfig
config.InProcessBackend.cluster_options = options.Options(
options.Integer("option_one",
default=1,
min=1,
max=4,
target="option_one_b"),
options.Select("option_two", options=[("small", 1.5), ("large", 15)]),
)
async with temp_gateway(config=config) as g:
async with g.gateway_client() as gateway:
with dask.config.set(gateway__cluster__options={"option_one": 2}):
# Local config can override
opts = await gateway.cluster_options()
assert opts.option_one == 2
assert opts.option_two == "small"
# Without local config, uses server-side defaults
opts = await gateway.cluster_options(use_local_defaults=False)
assert opts.option_one == 1
assert opts.option_two == "small"
with dask.config.set(gateway__cluster__options={
"option_two": "{TEST_OPTION_TWO}"
}):
# Values can format from environment variables
opts = await gateway.cluster_options()
assert opts.option_one == 1
assert opts.option_two == "large"
with dask.config.set(gateway__cluster__options={
"option_two": "{TEST_OPTION_TWO}",
"option_one": 3,
}):
# Defaults are merged with kwargs to new_cluster
async with gateway.new_cluster(option_one=2) as cluster:
report = await gateway.get_cluster(cluster.name)
assert report.options == {
"option_one": 2,
"option_two": "large"
}
# If passing `cluster_options`, defaults are assumed already applied
opts = await gateway.cluster_options()
opts.option_two = "small"
async with gateway.new_cluster(opts) as cluster:
report = await gateway.get_cluster(cluster.name)
assert report.options == {
"option_one": 3,
"option_two": "small"
}
async def test_cluster_manager_options(tmpdir):
async with temp_gateway(
cluster_manager_class=ClusterOptionsManager,
cluster_manager_options=options.Options(
options.Integer(
"option_one", default=1, min=1, max=4, target="option_one_b"
),
options.Select("option_two", options=[("small", 1.5), ("large", 15)]),
),
temp_dir=str(tmpdir.join("dask-gateway")),
) as gateway_proc:
async with Gateway(
address=gateway_proc.public_urls.connect_url,
proxy_address=gateway_proc.gateway_urls.connect_url,
asynchronous=True,
) as gateway:
# Create with no parameters
cluster = await gateway.new_cluster()
cluster_obj = gateway_proc.db.cluster_from_name(cluster.name)
assert cluster_obj.manager.option_one_b == 1
assert cluster_obj.manager.option_two == 1.5
assert cluster_obj.options == {"option_one": 1, "option_two": "small"}
await cluster.shutdown()
# Create with parameters
cluster = await gateway.new_cluster(option_two="large")
cluster_obj = gateway_proc.db.cluster_from_name(cluster.name)
assert cluster_obj.manager.option_one_b == 1
assert cluster_obj.manager.option_two == 15
assert cluster_obj.options == {"option_one": 1, "option_two": "large"}
await cluster.shutdown()
# With options object
opts = await gateway.cluster_options()
opts.option_one = 2
cluster = await gateway.new_cluster(opts, option_two="large")
cluster_obj = gateway_proc.db.cluster_from_name(cluster.name)
assert cluster_obj.manager.option_one_b == 2
assert cluster_obj.manager.option_two == 15
assert cluster_obj.options == {"option_one": 2, "option_two": "large"}
await cluster.shutdown()
# Bad parameters
with pytest.raises(TypeError):
await gateway.new_cluster(cluster_options=10)
with pytest.raises(ValueError) as exc:
await gateway.new_cluster(option_two="medium")
assert "option_two" in str(exc.value)
def server_opts():
return server_options.Options(
server_options.Integer("integer_field", 1, target="integer_field2"),
server_options.Float("float_field", 2.5, min=2, max=4),
server_options.Select("select_field", [("a", 5), "b", ("c", 10)]),
)
def test_select():
field = "field_name"
default = "banana"
options = ["apple", "banana", ("orange", 1)]
target = "target_name"
label = "Field Name"
# Default values
s_opt = server_options.Select(field, options)
assert s_opt.default == "apple"
assert s_opt.label == field
assert s_opt.target == field
# Specified values propogate
s_opt = server_options.Select(field,
options,
default=default,
target=target,
label=label)
assert s_opt.default == default
assert s_opt.target == target
assert s_opt.label == label
def check(opt):
assert opt.validate("apple") == "apple"
assert opt.validate("orange") == "orange"
with pytest.raises(TypeError):
opt.validate(1)
with pytest.raises(ValueError):
opt.validate("grape")
# Server-side validation
check(s_opt)
# Bad server-side constructors
with pytest.raises(TypeError):
server_options.Select(field, 1)
with pytest.raises(TypeError):
server_options.Select(field, [1, 2, 3])
with pytest.raises(ValueError):
server_options.Select(field, [])
# Serialization works
spec = s_opt.json_spec()
c_opt = client_options.Field._from_spec(spec)
assert isinstance(c_opt, client_options.Select)
assert c_opt.value == default
assert c_opt.options == ("apple", "banana", "orange")
assert c_opt.label == label
# Client-side validation
check(c_opt)
# Bad client-side constructors
with pytest.raises(TypeError):
client_options.Select(field, default, options=1)
with pytest.raises(TypeError):
client_options.Select(field, default, options=[1, 2, 3])
with pytest.raises(ValueError):
client_options.Select(field, default, options=[])
async def test_cluster_manager_options_client_config(tmpdir, monkeypatch):
monkeypatch.setenv("TEST_OPTION_TWO", "large")
async with temp_gateway(
cluster_manager_class=ClusterOptionsManager,
cluster_manager_options=options.Options(
options.Integer("option_one",
default=1,
min=1,
max=4,
target="option_one_b"),
options.Select("option_two",
options=[("small", 1.5), ("large", 15)]),
),
temp_dir=str(tmpdir),
) as gateway_proc:
async with Gateway(
address=gateway_proc.public_urls.connect_url,
proxy_address=gateway_proc.gateway_urls.connect_url,
asynchronous=True,
) as gateway:
with dask.config.set(gateway__cluster__options={"option_one": 2}):
# Local config can override
opts = await gateway.cluster_options()
assert opts.option_one == 2
assert opts.option_two == "small"
# Without local config, uses server-side defaults
opts = await gateway.cluster_options(use_local_defaults=False)
assert opts.option_one == 1
assert opts.option_two == "small"
with dask.config.set(gateway__cluster__options={
"option_two": "{TEST_OPTION_TWO}"
}):
# Values can format from environment variables
opts = await gateway.cluster_options()
assert opts.option_one == 1
assert opts.option_two == "large"
with dask.config.set(gateway__cluster__options={
"option_two": "{TEST_OPTION_TWO}",
"option_one": 3,
}):
# Defaults are merged with kwargs to new_cluster
cluster = await gateway.new_cluster(option_one=2)
cluster_obj = gateway_proc.db.cluster_from_name(cluster.name)
assert cluster_obj.options == {
"option_one": 2,
"option_two": "large"
}
await cluster.shutdown()
# If passing `cluster_options`, defaults are assumed already applied
opts = await gateway.cluster_options()
opts.option_two = "small"
cluster = await gateway.new_cluster(opts)
cluster_obj = gateway_proc.db.cluster_from_name(cluster.name)
assert cluster_obj.options == {
"option_one": 3,
"option_two": "small"
}
await cluster.shutdown()
