diff --git a/MANIFEST.in b/MANIFEST.in index 08f025a4..61e0ef4f 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,3 +1,6 @@ +recursive-include dask_jobqueue *.py +recursive-include dask_jobqueue *.yaml + include LICENSE.txt include README.rst diff --git a/ci/none.sh b/ci/none.sh index 333f18d7..1c67925e 100644 --- a/ci/none.sh +++ b/ci/none.sh @@ -7,6 +7,8 @@ function jobqueue_before_install { ./ci/conda_setup.sh export PATH="$HOME/miniconda/bin:$PATH" conda install --yes -c conda-forge python=$TRAVIS_PYTHON_VERSION dask distributed flake8 pytest docrep + pip install --no-cache-dir git+https://github.com/dask/dask.git --upgrade --no-deps + pip install --no-cache-dir git+https://github.com/dask/distributed.git --upgrade --no-deps } function jobqueue_install { diff --git a/ci/pbs/Dockerfile b/ci/pbs/Dockerfile index 59e74ebd..60d3011d 100644 --- a/ci/pbs/Dockerfile +++ b/ci/pbs/Dockerfile @@ -31,6 +31,8 @@ RUN curl -o miniconda.sh https://repo.continuum.io/miniconda/Miniconda3-latest-L /opt/anaconda/bin/conda clean -tipy && \ rm -f miniconda.sh RUN conda install --yes -c conda-forge python=3.6 dask distributed flake8 pytest docrep +RUN pip install --no-cache-dir git+https://github.com/dask/dask.git --upgrade --no-deps +RUN pip install --no-cache-dir git+https://github.com/dask/distributed.git --upgrade --no-deps # Copy entrypoint and other needed scripts COPY ./*.sh / diff --git a/ci/sge/Dockerfile-master b/ci/sge/Dockerfile-master index 515b59ae..b5789514 100644 --- a/ci/sge/Dockerfile-master +++ b/ci/sge/Dockerfile-master @@ -12,6 +12,8 @@ ENV PATH /opt/anaconda/bin:$PATH RUN conda install -n root conda=4.4.11 && conda clean -tipy RUN conda install -c conda-forge dask distributed blas pytest mock ipython pip psutil && conda clean -tipy RUN pip install --no-cache-dir drmaa +RUN pip install --no-cache-dir git+https://github.com/dask/dask.git --upgrade --no-deps +RUN pip install --no-cache-dir git+https://github.com/dask/distributed.git --upgrade --no-deps COPY ./*.sh / COPY ./*.txt / diff --git a/ci/sge/Dockerfile-slave b/ci/sge/Dockerfile-slave index 516eb940..587b1adb 100644 --- a/ci/sge/Dockerfile-slave +++ b/ci/sge/Dockerfile-slave @@ -12,6 +12,8 @@ ENV PATH /opt/anaconda/bin:$PATH RUN conda install -n root conda=4.4.11 && conda clean -tipy RUN conda install -c conda-forge dask distributed blas pytest mock ipython pip psutil && conda clean -tipy RUN pip install --no-cache-dir drmaa +RUN pip install --no-cache-dir git+https://github.com/dask/dask.git --upgrade --no-deps +RUN pip install --no-cache-dir git+https://github.com/dask/distributed.git --upgrade --no-deps COPY ./setup-slave.sh / COPY ./*.sh / diff --git a/ci/slurm/Dockerfile b/ci/slurm/Dockerfile index 814cf792..be0f22d0 100644 --- a/ci/slurm/Dockerfile +++ b/ci/slurm/Dockerfile @@ -6,6 +6,8 @@ RUN curl -o miniconda.sh https://repo.continuum.io/miniconda/Miniconda3-latest-L rm -f miniconda.sh ENV PATH /opt/anaconda/bin:$PATH RUN conda install --yes -c conda-forge python=3.6 dask distributed flake8 pytest docrep +RUN pip install --no-cache-dir git+https://github.com/dask/dask.git --upgrade --no-deps +RUN pip install --no-cache-dir git+https://github.com/dask/distributed.git --upgrade --no-deps ENV LC_ALL en_US.UTF-8 diff --git a/dask_jobqueue/__init__.py b/dask_jobqueue/__init__.py index 996e9c12..12f870c1 100644 --- a/dask_jobqueue/__init__.py +++ b/dask_jobqueue/__init__.py @@ -1,4 +1,5 @@ # flake8: noqa +from . import config from .core import JobQueueCluster from .pbs import PBSCluster from .slurm import SLURMCluster diff --git a/dask_jobqueue/config.py b/dask_jobqueue/config.py new file mode 100644 index 00000000..7ab63bb3 --- /dev/null +++ b/dask_jobqueue/config.py @@ -0,0 +1,15 @@ +from __future__ import print_function, division, absolute_import + +import os + +import dask +import yaml + + +fn = os.path.join(os.path.dirname(__file__), 'jobqueue.yaml') +dask.config.ensure_file(source=fn) + +with open(fn) as f: + defaults = yaml.load(f) + +dask.config.update(dask.config.config, defaults, priority='old') diff --git a/dask_jobqueue/core.py b/dask_jobqueue/core.py index 37c1a67f..d2e6a53d 100644 --- a/dask_jobqueue/core.py +++ b/dask_jobqueue/core.py @@ -6,6 +6,7 @@ import sys from contextlib import contextmanager +import dask import docrep from distributed import LocalCluster from distributed.deploy import Cluster @@ -78,15 +79,15 @@ class JobQueueCluster(Cluster): cancel_command = None def __init__(self, - name='dask-worker', - threads=2, - processes=4, - memory='8GB', - interface=None, - death_timeout=60, - local_directory=None, - extra='', - env_extra=[], + name=dask.config.get('jobqueue.name'), + threads=dask.config.get('jobqueue.threads'), + processes=dask.config.get('jobqueue.processes'), + memory=dask.config.get('jobqueue.memory'), + interface=dask.config.get('jobqueue.interface'), + death_timeout=dask.config.get('jobqueue.death-timeout'), + local_directory=dask.config.get('jobqueue.local-directory'), + extra=dask.config.get('jobqueue.extra'), + env_extra=dask.config.get('jobqueue.env-extra'), **kwargs ): """ """ diff --git a/dask_jobqueue/jobqueue.yaml b/dask_jobqueue/jobqueue.yaml new file mode 100644 index 00000000..018936f4 --- /dev/null +++ b/dask_jobqueue/jobqueue.yaml @@ -0,0 +1,26 @@ +jobqueue: + name: dask-worker + threads: 2 + processes: 4 + memory: 8GB + interface: null + death-timeout: 60 + local-directory: null + extra: "" + env-extra: [] + + queue: null + project: null + walltime: '00:30:00' + + pbs: + resource-spec: null + job-extra: [] + + sge: + resource-spec: null + + slurm: + job-cpu: null + job-mem: null + job-extra: {} diff --git a/dask_jobqueue/pbs.py b/dask_jobqueue/pbs.py index f7f490cf..350547ad 100644 --- a/dask_jobqueue/pbs.py +++ b/dask_jobqueue/pbs.py @@ -4,6 +4,8 @@ import math import os +import dask + from .core import JobQueueCluster, docstrings logger = logging.getLogger(__name__) @@ -58,11 +60,11 @@ class PBSCluster(JobQueueCluster): cancel_command = 'qdel' def __init__(self, - queue=None, - project=None, - resource_spec=None, - walltime='00:30:00', - job_extra=[], + queue=dask.config.get('jobqueue.queue'), + project=dask.config.get('jobqueue.project'), + resource_spec=dask.config.get('jobqueue.pbs.resource-spec'), + walltime=dask.config.get('jobqueue.walltime'), + job_extra=dask.config.get('jobqueue.pbs.job-extra'), **kwargs): # Instantiate args and parameters from parent abstract class diff --git a/dask_jobqueue/sge.py b/dask_jobqueue/sge.py index 489a75ff..e6a2a841 100644 --- a/dask_jobqueue/sge.py +++ b/dask_jobqueue/sge.py @@ -1,5 +1,7 @@ import logging +import dask + from .core import JobQueueCluster, docstrings logger = logging.getLogger(__name__) @@ -42,10 +44,10 @@ class SGECluster(JobQueueCluster): cancel_command = 'qdel' def __init__(self, - queue=None, - project=None, - resource_spec=None, - walltime='0:30:00', + queue=dask.config.get('jobqueue.queue'), + project=dask.config.get('jobqueue.project'), + resource_spec=dask.config.get('jobqueue.sge.resource-spec'), + walltime=dask.config.get('jobqueue.walltime'), **kwargs): super(SGECluster, self).__init__(**kwargs) diff --git a/dask_jobqueue/slurm.py b/dask_jobqueue/slurm.py index 13f1ad19..d4878315 100644 --- a/dask_jobqueue/slurm.py +++ b/dask_jobqueue/slurm.py @@ -5,6 +5,8 @@ import os import sys +import dask + from .core import JobQueueCluster, docstrings logger = logging.getLogger(__name__) @@ -58,12 +60,12 @@ class SLURMCluster(JobQueueCluster): cancel_command = 'scancel' def __init__(self, - queue=None, - project=None, - walltime='00:30:00', - job_cpu=None, - job_mem=None, - job_extra=[], + queue=dask.config.get('jobqueue.queue'), + project=dask.config.get('jobqueue.project'), + walltime=dask.config.get('jobqueue.walltime'), + job_cpu=dask.config.get('jobqueue.slurm.job-cpu'), + job_mem=dask.config.get('jobqueue.slurm.job-mem'), + job_extra=dask.config.get('jobqueue.slurm.job-extra'), **kwargs): super(SLURMCluster, self).__init__(**kwargs) diff --git a/setup.py b/setup.py index 3747abf5..cfc19ddf 100755 --- a/setup.py +++ b/setup.py @@ -22,6 +22,7 @@ url='https://github.com/dask/dask-jobqueue', license='BSD 3-Clause', packages=['dask_jobqueue'], + include_package_data=True, install_requires=install_requires, tests_require=['pytest >= 2.7.1'], long_description=long_description,