first commit

deafrica_tools/dask.py

@@ -0,0 +1,106 @@
+"""
+Functions for simplifying the creation of a local dask cluster.
+"""
+
+from importlib.util import find_spec
+import os
+import dask
+from aiohttp import ClientConnectionError
+from datacube.utils.dask import start_local_dask
+from datacube.utils.rio import configure_s3_access
+
+_HAVE_PROXY = bool(find_spec('jupyter_server_proxy'))
+_IS_AWS = ('AWS_ACCESS_KEY_ID' in os.environ or
+           'AWS_DEFAULT_REGION' in os.environ)
+
+
+def create_local_dask_cluster(spare_mem='3Gb', display_client=True, return_client=False):
+    """
+    Using the datacube utils function `start_local_dask`, generate
+    a local dask cluster. Automatically detects if on AWS or NCI.
+
+    Parameters
+    ----------
+    spare_mem : String, optional
+        The amount of memory, in Gb, to leave for the notebook to run.
+        This memory will not be used by the cluster. e.g '3Gb'
+    display_client : Bool, optional
+        An optional boolean indicating whether to display a summary of
+        the dask client, including a link to monitor progress of the
+        analysis. Set to False to hide this display.
+    return_client : Bool, optional
+        An optional boolean indicating whether to return the dask client
+        object.
+
+    """
+
+    if _HAVE_PROXY:
+        # Configure dashboard link to go over proxy
+        prefix = os.environ.get('JUPYTERHUB_SERVICE_PREFIX', '/')
+        dask.config.set({"distributed.dashboard.link":
+                         prefix + "proxy/{port}/status"})
+
+    # Start up a local cluster
+    client = start_local_dask(mem_safety_margin=spare_mem)
+
+    if _IS_AWS:
+        # Configure GDAL for s3 access
+        configure_s3_access(aws_unsigned=True,
+                            client=client)
+
+    # Show the dask cluster settings
+    if display_client:
+        from IPython.display import display
+        display(client)
+    
+    # return the client as an object
+    if return_client:
+        return client
+   
+
+try:
+    from dask_gateway import Gateway
+
+    def create_dask_gateway_cluster(profile='r5_L', workers=2):
+        """
+        Create a cluster in our internal dask cluster.
+
+        Parameters
+        ----------
+        profile : str
+            Possible values are:
+                - r5_L (2 cores, 15GB memory)
+                - r5_XL (4 cores, 31GB memory)
+                - r5_2XL (8 cores, 63GB memory)
+                - r5_4XL (16 cores, 127GB memory)
+
+        workers : int
+            Number of workers in the cluster.
+        """
+        try:
+            gateway = Gateway()
+            
+            # Close any existing clusters
+            cluster_names = gateway.list_clusters()
+            if len(cluster_names) > 0:
+                print("Cluster(s) still running:", cluster_names)
+                for n in cluster_names:
+                    cluster = gateway.connect(n.name)
+                    cluster.shutdown()            
+            
+            options = gateway.cluster_options()
+            options['profile'] = profile
+
+            # limit username to alphanumeric characters
+            # kubernetes pods won't launch if labels contain anything other than [a-Z, -, _]
+            options['jupyterhub_user'] = ''.join(c if c.isalnum() else '-' for c in os.getenv('JUPYTERHUB_USER'))
+
+            cluster = gateway.new_cluster(options)
+            cluster.scale(workers)
+            return cluster
+        except ClientConnectionError:
+            raise ConnectionError("access to dask gateway cluster unauthorized")
+
+except ImportError:
+    def create_dask_gateway_cluster(*args, **kwargs):
+        raise NotImplementedError