Repository: RhodiumGroup/rhg_compute_tools

Issue:
KubeCluster.adapt() does not work
we get rapid increases and decreases of the number of workers. Also, this doesn't play nice with multiple threads

Discussion:

Patch:
diff --git a/requirements.txt b/requirements.txt
index 6603199..d19ea9f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
-matplotlib==3.0.3
 google-cloud-storage==1.16.1
 dask-kubernetes==0.8.0
+matplotlib==3.0.3
+numpy>=1.14
diff --git a/rhg_compute_tools/gcs.py b/rhg_compute_tools/gcs.py
index 1a45bf7..d8a610a 100644
--- a/rhg_compute_tools/gcs.py
+++ b/rhg_compute_tools/gcs.py
@@ -33,29 +33,30 @@ def get_bucket(cred_path):
     return bucket
 
 
-def _get_path_types(src,dest):
+def _get_path_types(src, dest):
     src_gs = src
     dest_gs = dest
     dest_gcs = dest
     if src_gs.startswith('/gcs/'):
-        src_gs = src.replace('/gcs/','gs://')
+        src_gs = src.replace('/gcs/', 'gs://')
     if dest_gs.startswith('/gcs/'):
         dest_gs = dest.replace('/gcs/', 'gs://')
     if dest_gcs.startswith('gs://'):
-        dest_gcs = dest.replace('gs://', '/gcs/')  
-        
+        dest_gcs = dest.replace('gs://', '/gcs/')
+
     return src_gs, dest_gs, dest_gcs
 
 
 def cp_to_gcs(*args, **kwargs):
     '''Deprecated. Use `cp_gcs`.'''
     return cp_gcs(*args, **kwargs)
-    
+
+
 def cp_gcs(src, dest, cp_flags=[]):
     '''Copy a file or recursively copy a directory from local
     path to GCS or vice versa. Must have already authenticated to use.
-    Notebook servers are automatically authenticated, but workers 
-    need to pass the path to the authentication json file to the 
+    Notebook servers are automatically authenticated, but workers
+    need to pass the path to the authentication json file to the
     GCLOUD_DEFAULT_TOKEN_FILE env var.
     This is done automatically for rhg-data.json when using the get_worker
     wrapper.
@@ -63,10 +64,10 @@ def cp_gcs(src, dest, cp_flags=[]):
     Parameters
     ----------
     src, dest : str
-        The paths to the source and destination file or directory. 
+        The paths to the source and destination file or directory.
         If on GCS, either the `/gcs` or `gs:/` prefix will work.
     cp_flags : list of str, optional
-        String of flags to add to the gsutil cp command. e.g. 
+        String of flags to add to the gsutil cp command. e.g.
         `cp_flags=['r']` will run the command `gsutil -m cp -r...`
         (recursive copy)
 
@@ -85,8 +86,8 @@ def cp_gcs(src, dest, cp_flags=[]):
     # make sure we're using URL
     # if /gcs or gs:/ are not in src or not in dest
     # then these won't change anything
-    src_gs, dest_gs, dest_gcs = _get_path_types(src,dest)
-    
+    src_gs, dest_gs, dest_gcs = _get_path_types(src, dest)
+
     # if directory already existed cp would put src into dest_gcs
     if exists(dest_gcs):
         dest_base = join(dest_gcs, basename(src))
@@ -94,7 +95,11 @@ def cp_gcs(src, dest, cp_flags=[]):
     else:
         dest_base = dest_gcs
 
-    cmd = 'gsutil -m cp ' + ' '.join(['-'+f for f in cp_flags]) + ' {} {}'.format(src_gs, dest_gs)
+    cmd = (
+        'gsutil -m cp '
+        + ' '.join(['-' + f for f in cp_flags])
+        + ' {} {}'.format(src_gs, dest_gs))
+
     cmd = shlex.split(cmd)
     p = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE)
@@ -111,11 +116,13 @@ def cp_gcs(src, dest, cp_flags=[]):
 
     return stdout, stderr, end_time - st_time
 
+
 def sync_to_gcs(*args, **kwargs):
     '''Deprecated. Use sync_gcs'''
     return sync_gcs(*args, **kwargs)
 
-def sync_gcs(src, dest, sync_flags=['r','d']):
+
+def sync_gcs(src, dest, sync_flags=['r', 'd']):
     '''Sync a directory from local to GCS or vice versa. Uses `gsutil rsync`.
     Must have already authenticated to use. Notebook servers
     are automatically authenticated, but workers need to pass the path
@@ -126,10 +133,10 @@ def sync_gcs(src, dest, sync_flags=['r','d']):
     Parameters
     ----------
     src, dest : str
-        The paths to the source and destination file or directory. 
+        The paths to the source and destination file or directory.
         If on GCS, either the `/gcs` or `gs:/` prefix will work.
     sync_flags : list of str, optional
-        String of flags to add to the gsutil cp command. e.g. 
+        String of flags to add to the gsutil cp command. e.g.
         `sync_flags=['r','d']` will run the command `gsutil -m cp -r -d...`
         (recursive copy, delete any files on dest that are not on src).
         This is the default set of flags.
@@ -149,20 +156,24 @@ def sync_gcs(src, dest, sync_flags=['r','d']):
     # remove trailing /'s
     src = src.rstrip('/')
     dest = dest.rstrip('/')
-    
+
     # make sure we're using URL
     # if /gcs or gs:/ are not in src or not in dest
     # then these won't change anything
-    src_gs, dest_gs, dest_gcs = _get_path_types(src,dest)
+    src_gs, dest_gs, dest_gcs = _get_path_types(src, dest)
+
+    cmd = (
+        'gsutil -m rsync '
+        + ' '.join(['-' + f for f in sync_flags])
+        + ' {} {}'.format(src_gs, dest_gs))
 
-    cmd = 'gsutil -m rsync ' + ' '.join(['-'+f for f in sync_flags]) + ' {} {}'.format(src_gs, dest_gs)
     cmd = shlex.split(cmd)
     p = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE)
     stdout, stderr = p.communicate()
 
-    # need to add directories if you were recursively copying a directory TO gcs
-    # now make directory blobs on gcs so that gcsfuse recognizes it
+    # need to add directories if you were recursively copying a directory TO
+    # gcs now make directory blobs on gcs so that gcsfuse recognizes it
     if dest_gcs.startswith('/gcs/'):
         dirs_to_make = [x[0].replace(src, dest_gcs) for x in os.walk(src)]
         for d in dirs_to_make:
@@ -170,4 +181,4 @@ def sync_gcs(src, dest, sync_flags=['r','d']):
 
     end_time = dt.now()
 
-    return stdout, stderr, end_time - st_time
\ No newline at end of file
+    return stdout, stderr, end_time - st_time
diff --git a/rhg_compute_tools/kubernetes.py b/rhg_compute_tools/kubernetes.py
index d9515dc..dad65ae 100644
--- a/rhg_compute_tools/kubernetes.py
+++ b/rhg_compute_tools/kubernetes.py
@@ -1,220 +1,8 @@
 # -*- coding: utf-8 -*-
+"""Tools for interacting with kubernetes."""
 
-import dask
-import dask.distributed
-
-from distributed.deploy.adaptive import Adaptive
-from distributed.deploy.cluster import Cluster
-from distributed.deploy.local import LocalCluster
 from dask_kubernetes import KubeCluster
-from toolz import valmap, second, compose, groupby
-from distributed.utils import log_errors, PeriodicCallback, parse_timedelta
-
-import atexit
-from datetime import timedelta
-import logging
-import math
-from time import sleep
-import weakref
-import toolz
-
-from tornado import gen
-
-from distributed.core import CommClosedError
-from distributed.utils import (sync, ignoring, All, silence_logging, LoopRunner,
-        log_errors, thread_state)
-from distributed.nanny import Nanny
-from distributed.scheduler import Scheduler
-from distributed.worker import Worker, _ncores
-
-import logging
-
-
-logger = logging.getLogger('distributed.deploy.adaptive')
-logger.setLevel(logging.DEBUG)
-
-
-class MyAdaptive(Adaptive):
-    def needs_cpu(self):
-        total_occupancy = self.scheduler.total_occupancy
-        total_cores = sum([ws.ncores for ws in self.scheduler.workers.values()])
-
-        if total_occupancy / (total_cores + 1e-9) > self.startup_cost * 2:
-            logger.info("CPU limit exceeded [%d occupancy / %d cores]",
-                        total_occupancy, total_cores)
-
-            num_workers = len(self.scheduler.workers)
-
-            tasks_processing = 0
-
-            for w in self.scheduler.workers.values():
-                tasks_processing += len(w.processing)
-
-                if tasks_processing > num_workers:
-                    logger.info(
-                        "pending tasks exceed number of workers "
-                        "[%d tasks / %d workers]",
-                        tasks_processing, num_workers)
-
-                    return True
-
-        return False
-
-
-class MyCluster(Cluster):
-    def adapt(self, **kwargs):
-        with ignoring(AttributeError):
-            self._adaptive.stop()
-        if not hasattr(self, '_adaptive_options'):
-            self._adaptive_options = {}
-        self._adaptive_options.update(kwargs)
-        self._adaptive = MyAdaptive(self.scheduler, self, **self._adaptive_options)
-        return self._adaptive
-
-
-class MyLocalCluster(MyCluster, LocalCluster):
-    pass
-
-
-'''
-subclass dask_kubernetes.KubeCluster
-'''
-
-import getpass
-import logging
-import os
-import socket
-import string
-import time
-from urllib.parse import urlparse
-import uuid
-from weakref import finalize
-
-try:
-    import yaml
-except ImportError:
-    yaml = False
-
 import dask
-from distributed.deploy import LocalCluster, Cluster
-from distributed.comm.utils import offload
-import kubernetes
-from tornado import gen
-
-from dask_kubernetes.objects import make_pod_from_dict, clean_pod_template
-
-
-class MyKubeCluster(MyCluster, KubeCluster):
-    def __init__(
-            self,
-            pod_template=None,
-            name=None,
-            namespace=None,
-            n_workers=None,
-            host=None,
-            port=None,
-            env=None,
-            **kwargs
-    ):
-        name = name or dask.config.get('kubernetes.name')
-        namespace = namespace or dask.config.get('kubernetes.namespace')
-        n_workers = n_workers if n_workers is not None else dask.config.get('kubernetes.count.start')
-        host = host or dask.config.get('kubernetes.host')
-        port = port if port is not None else dask.config.get('kubernetes.port')
-        env = env if env is not None else dask.config.get('kubernetes.env')
-
-        if not pod_template and dask.config.get('kubernetes.worker-template', None):
-            d = dask.config.get('kubernetes.worker-template')
-            d = dask.config.expand_environment_variables(d)
-            pod_template = make_pod_from_dict(d)
-
-        if not pod_template and dask.config.get('kubernetes.worker-template-path', None):
-            import yaml
-            fn = dask.config.get('kubernetes.worker-template-path')
-            fn = fn.format(**os.environ)
-            with open(fn) as f:
-                d = yaml.safe_load(f)
-            d = dask.config.expand_environment_variables(d)
-            pod_template = make_pod_from_dict(d)
-
-        if not pod_template:
-            msg = ("Worker pod specification not provided. See KubeCluster "
-                   "docstring for ways to specify workers")
-            raise ValueError(msg)
-
-        self.cluster = MyLocalCluster(ip=host or socket.gethostname(),
-                                    scheduler_port=port,
-                                    n_workers=0, **kwargs)
-        try:
-            kubernetes.config.load_incluster_config()
-        except kubernetes.config.ConfigException:
-            kubernetes.config.load_kube_config()
-
-        self.core_api = kubernetes.client.CoreV1Api()
-
-        if namespace is None:
-            namespace = _namespace_default()
-
-        name = name.format(user=getpass.getuser(),
-                           uuid=str(uuid.uuid4())[:10],
-                           **os.environ)
-        name = escape(name)
-
-        self.pod_template = clean_pod_template(pod_template)
-        # Default labels that can't be overwritten
-        self.pod_template.metadata.labels['dask.org/cluster-name'] = name
-        self.pod_template.metadata.labels['user'] = escape(getpass.getuser())
-        self.pod_template.metadata.labels['app'] = 'dask'
-        self.pod_template.metadata.labels['component'] = 'dask-worker'
-        self.pod_template.metadata.namespace = namespace
-
-        self.pod_template.spec.containers[0].env.append(
-            kubernetes.client.V1EnvVar(name='DASK_SCHEDULER_ADDRESS',
-                                       value=self.scheduler_address)
-        )
-        if env:
-            self.pod_template.spec.containers[0].env.extend([
-                kubernetes.client.V1EnvVar(name=k, value=str(v))
-                for k, v in env.items()
-            ])
-        self.pod_template.metadata.generate_name = name
-
-        finalize(self, _cleanup_pods, self.namespace, self.pod_template.metadata.labels)
-
-        if n_workers:
-            self.scale(n_workers)
-
-def _cleanup_pods(namespace, labels):
-    """ Remove all pods with these labels in this namespace """
-    api = kubernetes.client.CoreV1Api()
-    pods = api.list_namespaced_pod(namespace, label_selector=format_labels(labels))
-    for pod in pods.items:
-        try:
-            api.delete_namespaced_pod(pod.metadata.name, namespace,
-                                      kubernetes.client.V1DeleteOptions())
-            logger.info('Deleted pod: %s', pod.metadata.name)
-        except kubernetes.client.rest.ApiException as e:
-            # ignore error if pod is already removed
-            if e.status != 404:
-                raise
-
-def _namespace_default():
-    ns_path = '/var/run/secrets/kubernetes.io/serviceaccount/namespace'
-    if os.path.exists(ns_path):
-        with open(ns_path) as f:
-            return f.read().strip()
-    return 'default'
-
-
-valid_characters = string.ascii_letters + string.digits + '_-.'
-
-
-def escape(s):
-    return ''.join(c for c in s if c in valid_characters)
-
-
-
-"""Tools for interacting with kubernetes."""
 import dask.distributed as dd
 import yaml as yml
 import traceback as tb
@@ -243,6 +31,7 @@ def _append_docstring(func_with_docstring):
         return func
     return decorator
 
+
 def get_cluster(
         name=None,
         extra_pip_packages=None,
@@ -255,7 +44,8 @@ def get_cluster(
         env_items=None,
         scaling_factor=1,
         dask_config_dict={},
-        template_path='~/worker-template.yml'):
+        template_path='~/worker-template.yml',
+        **kwargs):
     """
     Start dask.kubernetes cluster and dask.distributed client
 
@@ -333,10 +123,9 @@ def get_cluster(
 
     """
 
-    ## update dask settings
+    # update dask settings
     dask.config.set(dask_config_dict)
 
-
     template_path = os.path.expanduser(template_path)
 
     with open(template_path, 'r') as f:
@@ -362,7 +151,7 @@ def get_cluster(
         container['env'].append({
             'name': 'GCLOUD_DEFAULT_TOKEN_FILE',
             'value': cred_path})
-        
+
     elif cred_name is not None:
         container['env'].append({
             'name': 'GCLOUD_DEFAULT_TOKEN_FILE',
@@ -414,7 +203,7 @@ def get_cluster(
     requests['cpu'] = format_request(float(cpus) * scaling_factor)
 
     # start cluster and client and return
-    cluster = MyKubeCluster.from_dict(template)
+    cluster = KubeCluster.from_dict(template)
 
     client = dd.Client(cluster)
 
diff --git a/rhg_compute_tools/utils.py b/rhg_compute_tools/utils.py
index f2f130c..5a45528 100644
--- a/rhg_compute_tools/utils.py
+++ b/rhg_compute_tools/utils.py
@@ -4,6 +4,7 @@ import json
 import numpy as np
 import os
 
+
 def expand(func):
     '''
     Decorator to expand an (args, kwargs) tuple in function calls
@@ -157,7 +158,7 @@ class NumpyEncoder(json.JSONEncoder):
     '''
     Helper class for json.dumps to coerce numpy objects to native python
     '''
-    
+
     def default(self, obj):
         if isinstance(obj, np.ndarray):
             return obj.tolist()
@@ -172,23 +173,30 @@ class NumpyEncoder(json.JSONEncoder):
         return json.JSONEncoder.default(self, obj)
 
 
-def checkpoint(jobs, futures, job_name, log_dir='.', extra_pending=None, extra_errors=None, extra_others=None):
+def checkpoint(
+        jobs,
+        futures,
+        job_name,
+        log_dir='.',
+        extra_pending=None,
+        extra_errors=None,
+        extra_others=None):
     '''
     checkpoint and save a job state to disk
     '''
-    
+
     err_msg = (
         "lengths do not match: jobs [{}] != futures [{}]"
         .format(len(jobs), len(futures)))
 
     assert len(jobs) == len(futures), err_msg
-    
+
     if extra_pending is None:
         extra_pending = []
-    
+
     if extra_errors is None:
         extra_errors = []
-        
+
     if extra_others is None:
         extra_others = {}
 
@@ -199,18 +207,18 @@ def checkpoint(jobs, futures, job_name, log_dir='.', extra_pending=None, extra_e
     errored_jobs = (
         [jobs[i] for i, f in enumerate(futures) if f.status == 'error']
         + extra_errors)
-    
+
     other_jobs = {}
-    
+
     for i, f in enumerate(futures):
         if f.status in ['pending', 'error', 'finished']:
             continue
 
         if f.status not in other_jobs:
             other_jobs[f.status] = []
-        
+
         other_jobs[f.status].append(jobs[i])
-        
+
     for k, v in extra_others.items():
         if k not in other_jobs:
             other_jobs[k] = []
@@ -252,5 +260,5 @@ def recover(job_name, log_dir='.'):
             other = {}
         else:
             other = json.loads(content)
-        
+
     return pending, errored, other
diff --git a/rhg_compute_tools/xarray.py b/rhg_compute_tools/xarray.py
index 8246549..5f01a67 100644
--- a/rhg_compute_tools/xarray.py
+++ b/rhg_compute_tools/xarray.py
@@ -42,7 +42,7 @@ def dataarrays_from_delayed(futures, client=None):
         >>> fut = client.map(build_arr, range(3))
         >>> arrs = dataarrays_from_delayed(fut)
         >>> arrs[-1]  # doctest: +ELLIPSIS
-        <xarray.DataArray 'from-value-...' (x: 2)>
+        <xarray.DataArray ...(x: 2)>
         dask.array<shape=(2,), dtype=int64, chunksize=(2,)>
         Coordinates:
           * x        (x) <U1 'a' 'b'
@@ -52,12 +52,13 @@ def dataarrays_from_delayed(futures, client=None):
     .. code-block:: python
 
         >>> xr.concat(arrs, dim='simulation') # doctest: +ELLIPSIS
-        <xarray.DataArray 'from-value-...' (simulation: 3, x: 2)>
+        <xarray.DataArray ...(simulation: 3, x: 2)>
         dask.array<shape=(3, 2), dtype=int64, chunksize=(1, 2)>
         Coordinates:
           * x        (x) <U1 'a' 'b'
         Dimensions without coordinates: simulation
 
+        >>> client.close()
     '''
 
     if client is None:
@@ -130,11 +131,13 @@ def dataarray_from_delayed(futures, dim=None, client=None):
         ...
 
         >>> da  # doctest: +ELLIPSIS
-        <xarray.DataArray 'from-value-...' (simulation: 3, x: 2)>
+        <xarray.DataArray ...(simulation: 3, x: 2)>
         dask.array<shape=(3, 2), dtype=int64, chunksize=(1, 2)>
         Coordinates:
           * x           (x) <U1 'a' 'b'
           * simulation  (simulation) int64 0 1 2
+
+        >>> client.close()
     '''
 
     data_arrays = dataarrays_from_delayed(futures, client=client)
@@ -202,6 +205,8 @@ def datasets_from_delayed(futures, client=None):
         Dimensions without coordinates: y
         Data variables:
             var1     (y, x) int64 dask.array<shape=(3, 2), chunksize=(1, 2)>
+
+        >>> client.close()
     '''
 
     if client is None:
@@ -313,6 +318,8 @@ def dataset_from_delayed(futures, dim=None, client=None):
           * y        (y) int64 0 1 2
         Data variables:
             var1     (y, x) int64 dask.array<shape=(3, 2), chunksize=(1, 2)>
+
+        >>> client.close()
     '''
 
     datasets = datasets_from_delayed(futures, client=client)


Test patch:
diff --git a/tests/test_rhg_compute_tools.py b/tests/test_rhg_compute_tools.py
index 853defd..acc44db 100644
--- a/tests/test_rhg_compute_tools.py
+++ b/tests/test_rhg_compute_tools.py
@@ -4,14 +4,13 @@
 """Tests for `rhg_compute_tools` package."""
 
 import pytest
-import inspect
 
-from rhg_compute_tools import gcs, kubernetes, utils
+from rhg_compute_tools import kubernetes
 
 
 def monkeypatch_cluster(func):
     def inner(monkeypatch, mem, cpu, scale):
-        def mock_KubeCluster_from_dict(dict_argument):
+        def mock_KubeCluster_from_dict(dict_argument, **kwargs):
             return dict_argument
 
         def mock_dask_Client(cluster):