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
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
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_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_default_urls_from_config():
hostname = socket.gethostname()
gateway = DaskGateway()
gateway.init_server_urls()
assert gateway.public_urls.bind_host == ""
assert gateway.public_urls.connect.hostname == hostname
assert gateway.gateway_urls.bind.scheme == "tls"
assert gateway.gateway_urls.bind_host == ""
assert gateway.gateway_urls.connect.hostname == hostname
assert gateway.private_urls.bind_port != 0
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 __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 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
def test_generate_config(tmpdir, capfd):
cfg_file = str(tmpdir.join("dask_gateway_config.py"))
orig_text = "c.foo = 'bar'"
with open(cfg_file, "w") as f:
f.write(orig_text)
with pytest.raises(SystemExit) as exc:
DaskGateway.launch_instance(["generate-config", "--output", cfg_file])
DaskGateway.clear_instance()
assert "already exists" in exc.value.code
out, err = capfd.readouterr()
assert not out
assert not err
assert os.path.exists(cfg_file)
with open(cfg_file) as f:
cfg_text = f.read()
assert cfg_text == orig_text
DaskGateway.launch_instance(
["generate-config", "--force", "--output", cfg_file])
DaskGateway.clear_instance()
out, err = capfd.readouterr()
assert cfg_file in out
assert not err
with open(cfg_file) as f:
cfg_text = f.read()
assert "DaskGateway.cluster_manager_class" in cfg_text
assert "ClusterManager.worker_start_timeout" in cfg_text
def test_proxy_cli(kind, tmpdir, monkeypatch):
cfg_file = str(tmpdir.join("dask_gateway_config.py"))
text = ("c.DaskGateway.gateway_url = 'tls://127.0.0.1:8786'\n"
"c.DaskGateway.public_url = 'http://127.0.0.1:8888'\n"
"c.SchedulerProxy.api_url = 'http://127.0.0.1:8866'\n"
"c.WebProxy.api_url = 'http://127.0.0.1:8867'\n"
"c.ProxyBase.log_level = 'debug'\n"
"c.ProxyBase.auth_token = 'abcde'")
with open(cfg_file, "w") as f:
f.write(text)
if kind == "web":
address = "127.0.0.1:8888"
api_address = "127.0.0.1:8867"
else:
address = "127.0.0.1:8786"
api_address = "127.0.0.1:8866"
called_with = []
def mock_execle(*args):
called_with.extend(args)
monkeypatch.setattr(os, "execle", mock_execle)
DaskGateway.launch_instance([kind + "-proxy", "-f", cfg_file])
DaskGateway.clear_instance()
SchedulerProxyApp.clear_instance()
WebProxyApp.clear_instance()
assert called_with
env = called_with.pop()
assert called_with == [
_PROXY_EXE,
"dask-gateway-proxy",
kind,
"-address",
address,
"-api-address",
api_address,
"-log-level",
"debug",
]
assert "DASK_GATEWAY_PROXY_TOKEN" in env
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")
def test_resume_clusters_forbid_in_memory_db(tmpdir):
with pytest.raises(ValueError) as exc:
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://",
stop_clusters_on_shutdown=False,
authenticator_class="dask_gateway_server.auth.DummyAuthenticator",
)
assert "stop_clusters_on_shutdown" in str(exc.value)
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)
def test_proxy_cli(tmpdir, monkeypatch):
cfg_file = str(tmpdir.join("dask_gateway_config.py"))
text = ("c.DaskGateway.address = '127.0.0.1:8888'\n"
"c.Proxy.address = '127.0.0.1:8866'\n"
"c.Proxy.tcp_address = '127.0.0.1:8867'\n"
"c.Proxy.log_level = 'debug'\n"
"c.Proxy.api_token = 'abcde'")
with open(cfg_file, "w") as f:
f.write(text)
called_with = []
def mock_execle(*args):
called_with.extend(args)
monkeypatch.setattr(os, "execle", mock_execle)
DaskGateway.launch_instance(
["proxy", "-f", cfg_file, "--log-level", "warn"])
DaskGateway.clear_instance()
ProxyApp.clear_instance()
assert called_with
env = called_with.pop()
assert called_with == [
_PROXY_EXE,
"dask-gateway-proxy",
"-address",
"127.0.0.1:8866",
"-tcp-address",
"127.0.0.1:8867",
"-api-url",
"http://127.0.0.1:8888/api/v1/routes",
"-log-level",
"warn",
]
assert "DASK_GATEWAY_PROXY_TOKEN" in env
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")
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
async def __aenter__(self):
self.gateway = DaskGateway.instance(config=self.config)
self.gateway.initialize([])
await self.gateway.start_async()
return self.gateway
async def __aexit__(self, *args):
await self.gateway.stop_async(stop_event_loop=False)
DaskGateway.clear_instance()