def test_db_encrypt_keys_from_env(monkeypatch):
keys = [Fernet.generate_key(), Fernet.generate_key()]
val = b";".join(keys).decode()
monkeypatch.setenv("DASK_GATEWAY_ENCRYPT_KEYS", val)
gateway = DaskGateway()
gateway.initialize([])
assert gateway.backend.db_encrypt_keys == keys
def test_resume_clusters_forbid_in_memory_db():
c = Config()
c.DBBackendBase.db_url = "sqlite://"
c.DBBackendBase.stop_clusters_on_shutdown = False
with pytest.raises(ValueError) as exc:
gateway = DaskGateway(config=c)
gateway.initialize([])
assert "stop_clusters_on_shutdown" in str(exc.value)
def test_db_encrypt_keys_invalid(tmpdir):
c = Config()
c.DBBackendBase.db_url = "sqlite:///%s" % tmpdir.join("dask_gateway.sqlite")
c.DBBackendBase.db_encrypt_keys = ["abc"]
with pytest.raises(ValueError) as exc:
gateway = DaskGateway(config=c)
gateway.initialize([])
assert "DASK_GATEWAY_ENCRYPT_KEYS" in str(exc.value)
def test_db_encrypt_keys_required(tmpdir, capsys):
c = Config()
c.DBBackendBase.db_url = "sqlite:///%s" % tmpdir.join("dask_gateway.sqlite")
with pytest.raises(SystemExit) as exc:
gateway = DaskGateway(config=c)
gateway.initialize([])
gateway.start()
assert exc.value.code == 1
captured = capsys.readouterr()
assert "DASK_GATEWAY_ENCRYPT_KEYS" in captured.err
def test_shutdown_on_startup_error(tmpdir, capsys):
# A configuration that will cause a failure at runtime (not init time)
c = Config()
c.Proxy.tls_cert = str(tmpdir.join("tls_cert.pem"))
gateway = DaskGateway(config=c)
with pytest.raises(SystemExit) as exc:
gateway.initialize([])
gateway.start()
assert exc.value.code == 1
captured = capsys.readouterr()
assert "tls_cert" in captured.err
def test_db_encrypt_keys_required(tmpdir):
with pytest.raises(ValueError) as exc:
gateway = DaskGateway(
gateway_url="tls://127.0.0.1:0",
private_url="http://127.0.0.1:0",
public_url="http://127.0.0.1:0",
temp_dir=str(tmpdir),
db_url="sqlite:///%s" % tmpdir.join("dask_gateway.sqlite"),
authenticator_class="dask_gateway_server.auth.DummyAuthenticator",
)
gateway.initialize([])
assert "DASK_GATEWAY_ENCRYPT_KEYS" in str(exc.value)
async def test_shutdown_on_startup_error(tmpdir):
# A configuration that will cause a failure at runtime (not init time)
gateway = DaskGateway(
gateway_url="tls://127.0.0.1:0",
private_url="http://127.0.0.1:0",
public_url="http://127.0.0.1:0",
temp_dir=str(tmpdir),
tls_cert=str(tmpdir.join("tls_cert.pem")),
authenticator_class="dask_gateway_server.auth.DummyAuthenticator",
)
with pytest.raises(SystemExit) as exc:
gateway.initialize([])
await gateway.start_or_exit()
assert exc.value.code == 1
class temp_gateway:
def __init__(self, **kwargs):
c = Config()
c.DaskGateway.backend_class = InProcessBackend
config2 = kwargs.pop("config", None)
c.DaskGateway.address = "127.0.0.1:0"
c.Proxy.address = "127.0.0.1:0"
c.DaskGateway.authenticator_class = (
"dask_gateway_server.auth.SimpleAuthenticator")
c.DaskGateway.update(kwargs)
if config2:
c.merge(config2)
self.config = c
async def __aenter__(self):
self.gateway = DaskGateway(config=self.config)
self.gateway.initialize([])
await self.gateway.setup()
await self.gateway.backend.proxy._proxy_contacted
# Awaiting _proxy_contacted isn't failsafe and can lead to "ValueError,
# 404 NOT FOUND" if not complemented with a sufficiently long sleep
# following it. See https://github.com/dask/dask-gateway/pull/529 for a
# discussion on this.
await asyncio.sleep(0.25)
self.address = f"http://{self.gateway.backend.proxy.address}"
self.proxy_address = f"gateway://{self.gateway.backend.proxy.tcp_address}"
return self
async def __aexit__(self, *args):
await self.gateway.cleanup()
def gateway_client(self, **kwargs):
defaults = {
"address": self.address,
"proxy_address": self.proxy_address,
"asynchronous": True,
}
defaults.update(kwargs)
return Gateway(**defaults)
async def local_dask_gateway_server(
local_dask_gateway_server_config: traitlets.config.Config,
) -> AsyncIterator[DaskGatewayServer]:
print("--> creating local dask gateway server")
dask_gateway_server = DaskGateway(config=local_dask_gateway_server_config)
dask_gateway_server.initialize([]) # that is a shitty one!
print("--> local dask gateway server initialized")
await dask_gateway_server.setup()
await dask_gateway_server.backend.proxy._proxy_contacted # pylint: disable=protected-access
print("--> local dask gateway server setup completed")
yield DaskGatewayServer(
f"http://{dask_gateway_server.backend.proxy.address}",
f"gateway://{dask_gateway_server.backend.proxy.tcp_address}",
local_dask_gateway_server_config.SimpleAuthenticator.
password, # type: ignore
dask_gateway_server,
)
print("--> local dask gateway server switching off...")
await dask_gateway_server.cleanup()
print("...done")
class temp_gateway(object):
def __init__(self, **kwargs):
c = Config()
c.DaskGateway.backend_class = InProcessBackend
config2 = kwargs.pop("config", None)
c.DaskGateway.address = "127.0.0.1:0"
c.Proxy.address = "127.0.0.1:0"
c.DaskGateway.authenticator_class = (
"dask_gateway_server.auth.SimpleAuthenticator"
)
c.DaskGateway.update(kwargs)
if config2:
c.merge(config2)
self.config = c
async def __aenter__(self):
self.gateway = DaskGateway(config=self.config)
self.gateway.initialize([])
await self.gateway.setup()
await self.gateway.backend.proxy._proxy_contacted
self.address = f"http://{self.gateway.backend.proxy.address}"
self.proxy_address = f"gateway://{self.gateway.backend.proxy.tcp_address}"
return self
async def __aexit__(self, *args):
await self.gateway.cleanup()
def gateway_client(self, **kwargs):
defaults = {
"address": self.address,
"proxy_address": self.proxy_address,
"asynchronous": True,
}
defaults.update(kwargs)
return Gateway(**defaults)
async def local_dask_gateway_server(
cluster_id_resource_name: str,
) -> AsyncIterator[DaskGatewayServer]:
c = traitlets.config.Config()
c.DaskGateway.backend_class = UnsafeLocalBackend # type: ignore
c.DaskGateway.address = "127.0.0.1:0" # type: ignore
c.Proxy.address = "127.0.0.1:0" # type: ignore
c.DaskGateway.authenticator_class = "dask_gateway_server.auth.SimpleAuthenticator" # type: ignore
c.SimpleAuthenticator.password = "******" # type: ignore
c.ClusterConfig.worker_cmd = [ # type: ignore
"dask-worker",
"--resources",
f"CPU=12,GPU=1,MPI=1,RAM={16e9},{cluster_id_resource_name}=1",
]
# NOTE: This must be set such that the local unsafe backend creates a worker with enough cores/memory
c.ClusterConfig.worker_cores = 12 # type: ignore
c.ClusterConfig.worker_memory = "16G" # type: ignore
c.DaskGateway.log_level = "DEBUG" # type: ignore
print("--> creating local dask gateway server")
dask_gateway_server = DaskGateway(config=c)
dask_gateway_server.initialize([]) # that is a shitty one!
print("--> local dask gateway server initialized")
await dask_gateway_server.setup()
await dask_gateway_server.backend.proxy._proxy_contacted # pylint: disable=protected-access
print("--> local dask gateway server setup completed")
yield DaskGatewayServer(
f"http://{dask_gateway_server.backend.proxy.address}",
f"gateway://{dask_gateway_server.backend.proxy.tcp_address}",
c.SimpleAuthenticator.password, # type: ignore
dask_gateway_server,
)
print("--> local dask gateway server switching off...")
await dask_gateway_server.cleanup()
print("...done")