Skip to content

Commit

Permalink
Add dask config
Browse files Browse the repository at this point in the history
  • Loading branch information
jacobtomlinson committed Jan 25, 2024
1 parent bacaae7 commit d00ace7
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 2 deletions.
1 change: 1 addition & 0 deletions dask_ctl/ctl.yaml
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
ctl:
disable_discovery: []
cluster-spec: null
7 changes: 6 additions & 1 deletion dask_ctl/lifecycle.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import importlib
from typing import List

import dask.config
from dask.widgets import get_template
from dask.utils import typename
from distributed.deploy import LocalCluster
Expand All @@ -12,7 +13,7 @@


def create_cluster(
spec_path: str = "dask-cluster.yaml",
spec_path: str = None,
local_fallback: bool = False,
asynchronous: bool = False,
) -> Cluster:
Expand Down Expand Up @@ -47,6 +48,10 @@ def create_cluster(
LocalCluster(b3973c71, 'tcp://127.0.0.1:8786', workers=4, threads=12, memory=17.18 GB)
"""
spec_path = (
dask.config.get("ctl.cluster-spec", None, override_with=spec_path)
or "dask-cluster.yaml"
)

async def _create_cluster():
try:
Expand Down
7 changes: 6 additions & 1 deletion dask_ctl/tests/test_lifecycle.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import pytest
import ast

import dask.config
from dask.distributed import LocalCluster

from dask_ctl.lifecycle import create_cluster, get_snippet
Expand All @@ -14,9 +15,13 @@ def test_create_cluster(simple_spec_path):


def test_create_cluster_fallback():
with pytest.raises(DaskClusterConfigNotFound):
with pytest.raises(DaskClusterConfigNotFound, match="dask-cluster.yaml"):
cluster = create_cluster()

with dask.config.set({"ctl.cluster-spec": "foo.yaml"}):
with pytest.raises(DaskClusterConfigNotFound, match="foo.yaml"):
cluster = create_cluster()

cluster = create_cluster(local_fallback=True)
assert isinstance(cluster, LocalCluster)

Expand Down

0 comments on commit d00ace7

Please sign in to comment.