Fixes for Adaptive

dask_jobqueue/core.py

@@ -1,20 +1,70 @@
+from __future__ import absolute_import, division, print_function
+
 import logging
+import math
 import shlex
 import socket
 import subprocess
 import sys
+from collections import OrderedDict
 from contextlib import contextmanager
 
 import dask
 import docrep
 from distributed import LocalCluster
 from distributed.deploy import Cluster
-from distributed.utils import get_ip_interface, ignoring, parse_bytes, tmpfile
+from distributed.utils import get_ip_interface, parse_bytes, tmpfile
+from distributed.diagnostics.plugin import SchedulerPlugin
 
 logger = logging.getLogger(__name__)
 docstrings = docrep.DocstringProcessor()
 
 
+def _job_id_from_worker_name(name):
+    ''' utility to parse the job ID from the worker name
+
+    template: 'prefix-jobid[-proc]'
+    '''
+    pieces = name.split('-')
+    print(name, pieces)
+    if len(pieces) == 2:
+        return pieces[-1]
+    else:
+        return pieces[-2]
+
+
+class JobQueuePlugin(SchedulerPlugin):
+    def __init__(self):
+        self.pending_jobs = OrderedDict()
+        self.running_jobs = OrderedDict()
+        self.finished_jobs = OrderedDict()
+        self.all_workers = {}
+
+    def add_worker(self, scheduler, worker=None, name=None, **kwargs):
+        ''' Run when a new worker enters the cluster'''
+        w = scheduler.workers[worker]
+        job_id = _job_id_from_worker_name(w.name)
+        self.all_workers[worker] = (w.name, job_id)
+
+        # if this is the first worker for this job, move job to running
+        if job_id not in self.running_jobs:
+            self.running_jobs[job_id] = self.pending_jobs.pop(job_id)
+
+        # add worker to dict of workers in this job
+        self.running_jobs[job_id][w.name] = w
+
+    def remove_worker(self, scheduler=None, worker=None, **kwargs):
+        ''' Run when a worker leaves the cluster'''
+        name, job_id = self.all_workers[worker]
+
+        # remove worker from this job
+        del self.running_jobs[job_id][name]
+
+        # once there are no more workers, move this job to finished_jobs
+        if not self.running_jobs[job_id]:
+            self.finished_jobs[job_id] = self.running_jobs.pop(job_id)
+
+
 @docstrings.get_sectionsf('JobQueueCluster')
 class JobQueueCluster(Cluster):
     """ Base class to launch Dask Clusters for Job queues
@@ -75,6 +125,8 @@ class JobQueueCluster(Cluster):
     submit_command = None
     cancel_command = None
     scheduler_name = ''
+    _adaptive_options = {
+        'worker_key': lambda ws: _job_id_from_worker_name(ws.name)}
 
     def __init__(self,
                  name=None,
@@ -117,6 +169,10 @@ def __init__(self,
         if env_extra is None:
             env_extra = dask.config.get('jobqueue.%s.env-extra' % self.scheduler_name)
 
+        if '-' in name:
+            raise ValueError(
+                'name (%s) can not include the `-` character' % name)
+
         #This attribute should be overriden
         self.job_header = None
 
@@ -128,15 +184,17 @@ def __init__(self,
 
         self.cluster = LocalCluster(n_workers=0, ip=host, **kwargs)
 
+        # plugin for tracking job status
+        self._scheduler_plugin = JobQueuePlugin()
+        self.cluster.scheduler.add_plugin(self._scheduler_plugin)
+
         # Keep information on process, threads and memory, for use in
         # subclasses
         self.worker_memory = parse_bytes(memory) if memory is not None else None
         self.worker_processes = processes
         self.worker_threads = threads
         self.name = name
 
-        self.jobs = dict()
-        self.n = 0
         self._adaptive = None
 
         self._env_header = '\n'.join(env_extra)
@@ -152,22 +210,36 @@ def __init__(self,
         if memory is not None:
             self._command_template += " --memory-limit %s" % memory
         if name is not None:
-            self._command_template += " --name %s" % name
-            self._command_template += "-%(n)d" # Keep %(n) to be replaced later
+            # worker names follow this template: {NAME}-{JOB_ID}
+            self._command_template += " --name %s-${JOB_ID}" % name
         if death_timeout is not None:
             self._command_template += " --death-timeout %s" % death_timeout
         if local_directory is not None:
             self._command_template += " --local-directory %s" % local_directory
         if extra is not None:
             self._command_template += extra
 
+    @property
+    def pending_jobs(self):
+        """ Jobs pending in the queue """
+        return self._scheduler_plugin.pending_jobs
+
+    @property
+    def running_jobs(self):
+        """ Jobs with currenly active workers """
+        return self._scheduler_plugin.running_jobs
+
+    @property
+    def finished_jobs(self):
+        """ Jobs that have finished """
+        return self._scheduler_plugin.finished_jobs
+
     def job_script(self):
         """ Construct a job submission script """
-        self.n += 1
-        template = self._command_template % {'n': self.n}
-        return self._script_template % {'job_header': self.job_header,
-                                        'env_header': self._env_header,
-                                        'worker_command': template}
+        pieces = {'job_header': self.job_header,
+                  'env_header': self._env_header,
+                  'worker_command': self._command_template}
+        return self._script_template % pieces
 
     @contextmanager
     def job_file(self):
@@ -179,14 +251,12 @@ def job_file(self):
 
     def start_workers(self, n=1):
         """ Start workers and point them to our local scheduler """
-        workers = []
-        for _ in range(n):
+        num_jobs = math.ceil(n / self.worker_processes)
+        for _ in range(num_jobs):
             with self.job_file() as fn:
                 out = self._call(shlex.split(self.submit_command) + [fn])
                 job = self._job_id_from_submit_output(out.decode())
-                self.jobs[self.n] = job
-                workers.append(self.n)
-        return workers
+                self._scheduler_plugin.pending_jobs[job] = {}
 
     @property
     def scheduler(self):
@@ -215,12 +285,12 @@ def _calls(self, cmds):
         Also logs any stderr information
         """
         logger.debug("Submitting the following calls to command line")
+        procs = []
         for cmd in cmds:
             logger.debug(' '.join(cmd))
-        procs = [subprocess.Popen(cmd,
-                                  stdout=subprocess.PIPE,
-                                  stderr=subprocess.PIPE)
-                 for cmd in cmds]
+            procs.append(subprocess.Popen(cmd,
+                                          stdout=subprocess.PIPE,
+                                          stderr=subprocess.PIPE))
 
         result = []
         for proc in procs:
@@ -238,29 +308,43 @@ def stop_workers(self, workers):
         """ Stop a list of workers"""
         if not workers:
             return
-        workers = list(map(int, workers))
-        jobs = [self.jobs[w] for w in workers]
-        self._call([self.cancel_command] + list(jobs))
+        jobs = []
         for w in workers:
-            with ignoring(KeyError):
-                del self.jobs[w]
+            if isinstance(w, dict):
+                jobs.append(_job_id_from_worker_name(w['name']))
+            else:
+                jobs.append(_job_id_from_worker_name(w.name))
+        self.stop_jobs(jobs)
+
+    def stop_jobs(self, jobs):
+        """ Stop a list of jobs"""
+        if jobs:
+            self._call([self.cancel_command] + list(jobs))
 
     def scale_up(self, n, **kwargs):
         """ Brings total worker count up to ``n`` """
-        return self.start_workers(n - len(self.jobs))
+        active_and_pending = sum([len(j.workers) for j in
+                                  self.running_jobs.values()])
+        active_and_pending += self.worker_processes * len(self.pending_jobs)
+        self.start_workers(n - active_and_pending)
 
     def scale_down(self, workers):
         ''' Close the workers with the given addresses '''
-        if isinstance(workers, dict):
-            names = {v['name'] for v in workers.values()}
-            job_ids = {name.split('-')[-2] for name in names}
-            self.stop_workers(job_ids)
+        worker_states = []
+        for w in workers:
+            try:
+                # Get the actual WorkerState
+                worker_states.append(self.scheduler.workers[w])
+            except KeyError:
+                logger.debug('worker %s is already gone' % w)
+        self.stop_workers(worker_states)
 
     def __enter__(self):
         return self
 
     def __exit__(self, type, value, traceback):
-        self.stop_workers(self.jobs)
+        jobs = list(self.pending_jobs.keys()) + list(self.running_jobs.keys())
+        self.stop_jobs(jobs)
         self.cluster.__exit__(type, value, traceback)
 
     def _job_id_from_submit_output(self, out):

dask_jobqueue/jobqueue.yaml

@@ -1,6 +1,6 @@
 jobqueue:
   pbs:
-    name: dask-worker
+    name: dask_worker
 
     # Dask worker options
     threads: 2
@@ -20,7 +20,7 @@ jobqueue:
     job-extra: []
 
   sge:
-    name: dask-worker
+    name: dask_worker
 
     # Dask worker options
     threads: 2
@@ -40,7 +40,7 @@ jobqueue:
     resource-spec: null
 
   slurm:
-    name: dask-worker
+    name: dask_worker
 
     # Dask worker options
     threads: 2
@@ -61,7 +61,7 @@ jobqueue:
     job-extra: {}
 
   moab:
-    name: dask-worker
+    name: dask_worker
 
     # Dask worker options
     threads: 2

dask_jobqueue/pbs.py

@@ -75,7 +75,10 @@ def __init__(self, queue=None, project=None, resource_spec=None, walltime=None,
         # Instantiate args and parameters from parent abstract class
         super(PBSCluster, self).__init__(**kwargs)
 
-        header_lines = []
+        # Try to find a project name from environment variable
+        project = project or os.environ.get('PBS_ACCOUNT')
+
+        header_lines = ['#!/usr/bin/env bash']
         # PBS header build
         if self.name is not None:
             header_lines.append('#PBS -N %s' % self.name)
@@ -95,6 +98,7 @@ def __init__(self, queue=None, project=None, resource_spec=None, walltime=None,
         if walltime is not None:
             header_lines.append('#PBS -l walltime=%s' % walltime)
         header_lines.extend(['#PBS %s' % arg for arg in job_extra])
+        header_lines.append('JOB_ID=${PBS_JOBID%.*}')
 
         # Declare class attribute that shall be overriden
         self.job_header = '\n'.join(header_lines)

dask_jobqueue/sge.py

@@ -1,3 +1,5 @@
+from __future__ import absolute_import, division, print_function
+
 import logging
 
 import dask
@@ -56,8 +58,7 @@ def __init__(self, queue=None, project=None, resource_spec=None, walltime=None,
 
         super(SGECluster, self).__init__(**kwargs)
 
-        header_lines = ['#!/bin/bash']
-
+        header_lines = ['#!/usr/bin/env bash']
         if self.name is not None:
             header_lines.append('#$ -N %(name)s')
         if queue is not None:

dask_jobqueue/slurm.py

@@ -74,7 +74,7 @@ def __init__(self, queue=None, project=None, walltime=None,
         super(SLURMCluster, self).__init__(**kwargs)
 
         # Always ask for only one task
-        header_lines = []
+        header_lines = ['#!/usr/bin/env bash']
         # SLURM header build
         if self.name is not None:
             header_lines.append('#SBATCH -J %s' % self.name)
@@ -102,6 +102,7 @@ def __init__(self, queue=None, project=None, walltime=None,
 
         if walltime is not None:
             header_lines.append('#SBATCH -t %s' % walltime)
+        header_lines.append('JOB_ID=${SLURM_JOB_ID%;*}')
         header_lines.extend(['#SBATCH %s' % arg for arg in job_extra])
 
         # Declare class attribute that shall be overriden

dask_jobqueue/tests/__init__.py

@@ -0,0 +1,2 @@
+
+QUEUE_WAIT = 60  # seconds

dask_jobqueue/tests/test_jobqueue_core.py

@@ -3,11 +3,6 @@
 from dask_jobqueue import JobQueueCluster
 
 
-def test_jq_core_placeholder():
-    # to test that CI is working
-    pass
-
-
 def test_errors():
     with pytest.raises(NotImplementedError) as info:
         JobQueueCluster()