Extract provider specific requirements into extras_require in setup.py (

#125)

* Extract provider specific requirements into extras_require in setup.py

* Update documentation

* Reformat using black

* Add `all` group to extras_require to enable installing all extras

* Add documentation for the `all` extra

* Reformat using black

dask_cloudprovider/__init__.py

@@ -1,6 +1,13 @@
 from . import config
-from .providers.aws.ecs import ECSCluster, FargateCluster
-from .providers.azure.azureml import AzureMLCluster
+
+try:
+    from .providers.aws.ecs import ECSCluster, FargateCluster
+except ImportError:
+    pass
+try:
+    from .providers.azure.azureml import AzureMLCluster
+except ImportError:
+    pass
 
 __all__ = ["ECSCluster", "FargateCluster", "AzureMLCluster"]
 

dask_cloudprovider/providers/aws/ecs.py

@@ -1,24 +1,33 @@
-import asyncio
-import logging
-import uuid
-import warnings
-import weakref
-from typing import List
-
-from botocore.exceptions import ClientError
-import aiobotocore
-import dask
-
-from dask_cloudprovider.utils.logs import Log, Logs
-from dask_cloudprovider.utils.timeout import Timeout
-from dask_cloudprovider.providers.aws.helper import (
-    dict_to_aws,
-    aws_to_dict,
-    get_sleep_duration,
-)
-
-from distributed.deploy.spec import SpecCluster
-from distributed.utils import warn_on_duration
+try:
+    import asyncio
+    import logging
+    import uuid
+    import warnings
+    import weakref
+    from typing import List
+
+    from botocore.exceptions import ClientError
+    import aiobotocore
+    import dask
+
+    from dask_cloudprovider.utils.logs import Log, Logs
+    from dask_cloudprovider.utils.timeout import Timeout
+    from dask_cloudprovider.providers.aws.helper import (
+        dict_to_aws,
+        aws_to_dict,
+        get_sleep_duration,
+    )
+
+    from distributed.deploy.spec import SpecCluster
+    from distributed.utils import warn_on_duration
+except ImportError as e:
+    msg = (
+        "Dask Cloud Provider AWS requirements are not installed.\n\n"
+        "Please either conda or pip install as follows:\n\n"
+        "  conda install dask-cloudprovider                           # either conda install\n"
+        '  python -m pip install "dask-cloudprovider[aws]" --upgrade  # or python -m pip install'
+    )
+    raise ImportError(msg) from e
 
 logger = logging.getLogger(__name__)
 

dask_cloudprovider/providers/azure/azureml.py

@@ -1,21 +1,30 @@
-from azureml.core import Experiment, RunConfiguration, ScriptRunConfig
-from azureml.core.compute import AmlCompute, ComputeTarget
-from azureml.core.compute_target import ComputeTargetException
-from azureml.train.estimator import Estimator
-from azureml.core.runconfig import MpiConfiguration
-
-import time, os, socket, subprocess, logging
-import pathlib
-import threading
-
-from contextlib import suppress
-
-import dask
-from contextlib import suppress
-from distributed.deploy.cluster import Cluster
-from distributed.core import rpc
-from distributed.utils import LoopRunner, log_errors, format_bytes
-from tornado.ioloop import PeriodicCallback
+try:
+    from azureml.core import Experiment, RunConfiguration, ScriptRunConfig
+    from azureml.core.compute import AmlCompute, ComputeTarget
+    from azureml.core.compute_target import ComputeTargetException
+    from azureml.train.estimator import Estimator
+    from azureml.core.runconfig import MpiConfiguration
+
+    import time, os, socket, subprocess, logging
+    import pathlib
+    import threading
+
+    from contextlib import suppress
+
+    import dask
+    from contextlib import suppress
+    from distributed.deploy.cluster import Cluster
+    from distributed.core import rpc
+    from distributed.utils import LoopRunner, log_errors, format_bytes
+    from tornado.ioloop import PeriodicCallback
+except ImportError as e:
+    msg = (
+        "Dask Cloud Provider Azure requirements are not installed.\n\n"
+        "Please either conda or pip install as follows:\n\n"
+        "  conda install dask-cloudprovider                             # either conda install\n"
+        '  python -m pip install "dask-cloudprovider[azure]" --upgrade  # or python -m pip install'
+    )
+    raise ImportError(msg) from e
 
 logger = logging.getLogger(__name__)
 

doc/source/index.rst

@@ -15,7 +15,14 @@ Pip
 
 .. code-block:: console
 
-   $ pip install dask-cloudprovider
+   $ pip install dask-cloudprovider[all]
+
+You can also restrict your install to a specific cloud provider by giving their name instead of ``all``.
+
+.. code-block:: console
+
+   $ pip install dask-cloudprovider[aws]  # or
+   $ pip install dask-cloudprovider[azure]
 
 Conda
 ^^^^^

requirements.txt

@@ -1,4 +1,2 @@
-aiobotocore>=0.10.2
 dask>=2.2.0
 distributed>=2.3.1
-azureml-sdk>=1.0.83

setup.py

@@ -5,6 +5,12 @@
 
 import versioneer
 
+extras_require = {
+    "aws": ["aiobotocore>=0.10.2"],
+    "azure": ["azureml-sdk>=1.0.83"],
+}
+extras_require["all"] = set(pkg for pkgs in extras_require.values() for pkg in pkgs)
+
 setup(
     name="dask-cloudprovider",
     cmdclass=versioneer.get_cmdclass(),
@@ -18,6 +24,7 @@
     long_description=(open("README.md").read() if exists("README.md") else ""),
     zip_safe=False,
     install_requires=list(open("requirements.txt").read().strip().split("\n")),
+    extras_require=extras_require,
     entry_points="""
     [console_scripts]
     dask-ecs=dask_cloudprovider.cli.ecs:go