Refactor imports (#153)

* Refactor submodule structure and imports

* Update docs

* Install all deps in CI

.circleci/config.yml

@@ -38,6 +38,9 @@ jobs:
           key: miniconda-v1-{{ checksum "ci/environment-3.7.yml" }}
           paths:
             - "/home/circleci/miniconda"
+      - run:
+          command: |
+            /home/circleci/miniconda/envs/dask-cloudprovider-test/bin/pip install -e .[all]
       - run:
           command: |
             /home/circleci/miniconda/envs/dask-cloudprovider-test/bin/py.test dask_cloudprovider

ci/environment-3.7.yml

@@ -28,7 +28,5 @@ dependencies:
   - tornado >=5
   - zict >=0.1.3
   - pip:
-      - aiobotocore
       - git+https://github.com/dask/dask
       - git+https://github.com/dask/distributed
-      - azureml-sdk >=1.1.5

dask_cloudprovider/__init__.py

@@ -1,29 +1,27 @@
 from . import config
 
-try:
-    from .providers.aws.ecs import ECSCluster, FargateCluster
-    from .providers.aws.ec2 import EC2Cluster
-except ImportError:
-    pass
-try:
-    from .providers.azure.azureml import AzureMLCluster
-except ImportError:
-    pass
-try:
-    from .providers.digitalocean.droplet import DropletCluster
-except ImportError:
-    pass
-
-__all__ = [
-    "ECSCluster",
-    "EC2Cluster",
-    "FargateCluster",
-    "AzureMLCluster",
-    "DropletCluster",
-]
-
 from ._version import get_versions
 
 __version__ = get_versions()["version"]
 
 del get_versions
+
+
+def __getattr__(name):
+    if name in ["EC2Cluster", "ECSCluster", "FargateCluster"]:
+        raise ImportError(
+            "AWS cluster managers have been moved into the aws subpackage. "
+            f"Please import dask_cloudprovider.aws.{name}"
+        )
+
+    if name in ["AzureMLCluster"]:
+        raise ImportError(
+            "Azure cluster managers have been moved into the azure subpackage. "
+            f"Please import dask_cloudprovider.azure.{name}"
+        )
+
+    if name in ["DropletCluster"]:
+        raise ImportError(
+            "DigitalOcean cluster managers have been moved into the digitalocean subpackage. "
+            f"Please import dask_cloudprovider.digitalocean.{name}"
+        )

dask_cloudprovider/aws/__init__.py

@@ -0,0 +1,2 @@
+from .ec2 import EC2Cluster
+from .ecs import ECSCluster, FargateCluster

dask_cloudprovider/providers/aws/ec2.py → dask_cloudprovider/aws/ec2.py

@@ -1,13 +1,13 @@
 import asyncio
 
 import dask
-from dask_cloudprovider.providers.generic.vmcluster import (
+from dask_cloudprovider.generic.vmcluster import (
     VMCluster,
     VMInterface,
     SchedulerMixin,
     WorkerMixin,
 )
-from dask_cloudprovider.providers.aws.helper import (
+from dask_cloudprovider.aws.helper import (
     get_latest_ami_id,
     get_default_vpc,
     get_vpc_subnets,

dask_cloudprovider/providers/aws/ecs.py → dask_cloudprovider/aws/ecs.py

@@ -9,7 +9,7 @@
 
 from dask_cloudprovider.utils.logs import Log, Logs
 from dask_cloudprovider.utils.timeout import Timeout
-from dask_cloudprovider.providers.aws.helper import (
+from dask_cloudprovider.aws.helper import (
     dict_to_aws,
     aws_to_dict,
     get_sleep_duration,
@@ -617,15 +617,15 @@ class ECSCluster(SpecCluster):
     Examples
     --------
 
-    >>> from dask_cloudprovider import ECSCluster
+    >>> from dask_cloudprovider.aws import ECSCluster
     >>> cluster = ECSCluster(cluster_arn="arn:aws:ecs:<region>:<acctid>:cluster/<clustername>")
 
     There is also support in ``ECSCluster`` for GPU aware Dask clusters. To do
     this you need to create an ECS cluster with GPU capable instances (from the
     ``g3``, ``p3`` or ``p3dn`` families) and specify the number of GPUs each worker task
     should have.
 
-    >>> from dask_cloudprovider import ECSCluster
+    >>> from dask_cloudprovider.aws import ECSCluster
     >>> cluster = ECSCluster(
     ...     cluster_arn="arn:aws:ecs:<region>:<acctid>:cluster/<gpuclustername>",
     ...     worker_gpu=1)
@@ -1230,13 +1230,13 @@ class FargateCluster(ECSCluster):
     The ``FargateCluster`` will create a new Fargate ECS cluster by default along
     with all the IAM roles, security groups, and so on that it needs to function.
 
-    >>> from dask_cloudprovider import FargateCluster
+    >>> from dask_cloudprovider.aws import FargateCluster
     >>> cluster = FargateCluster()
 
     Note that in many cases you will want to specify a custom Docker image to ``FargateCluster`` so that Dask has the
     packages it needs to execute your workflow.
 
-    >>> from dask_cloudprovider import FargateCluster
+    >>> from dask_cloudprovider.aws import FargateCluster
     >>> cluster = FargateCluster(image="<hub-user>/<repo-name>[:<tag>]")
 
     One strategy to ensure that package versions match between your custom environment and the Docker container is to

dask_cloudprovider/providers/aws/tests/test_ec2.py → dask_cloudprovider/aws/tests/test_ec2.py

@@ -2,8 +2,8 @@
 
 aiobotocore = pytest.importorskip("aiobotocore")
 
-from dask_cloudprovider.providers.aws.ec2 import EC2Cluster
-from dask_cloudprovider.providers.aws.helper import get_latest_ami_id
+from dask_cloudprovider.aws.ec2 import EC2Cluster
+from dask_cloudprovider.aws.helper import get_latest_ami_id
 from dask.distributed import Client
 from distributed.core import Status
 

dask_cloudprovider/aws/tests/test_ecs.py

@@ -0,0 +1,8 @@
+import pytest
+
+aiobotocore = pytest.importorskip("aiobotocore")
+
+
+def test_import():
+    from dask_cloudprovider.aws import ECSCluster  # noqa
+    from dask_cloudprovider.aws import FargateCluster  # noqa

dask_cloudprovider/providers/aws/tests/test_helper.py → dask_cloudprovider/aws/tests/test_helper.py

@@ -1,5 +1,5 @@
 def test_aws_to_dict_and_back():
-    from dask_cloudprovider.providers.aws.helper import aws_to_dict, dict_to_aws
+    from dask_cloudprovider.aws.helper import aws_to_dict, dict_to_aws
 
     aws_dict = [{"key": "hello", "value": "world"}]
     aws_upper_dict = [{"Key": "hello", "Value": "world"}]
@@ -16,7 +16,7 @@ def test_aws_to_dict_and_back():
 
 
 def test_get_sleep_duration_first_try():
-    from dask_cloudprovider.providers.aws.helper import get_sleep_duration
+    from dask_cloudprovider.aws.helper import get_sleep_duration
 
     duration = get_sleep_duration(
         current_try=0, min_sleep_millis=10, max_sleep_millis=5000
@@ -25,7 +25,7 @@ def test_get_sleep_duration_first_try():
 
 
 def test_get_sleep_duration_max():
-    from dask_cloudprovider.providers.aws.helper import get_sleep_duration
+    from dask_cloudprovider.aws.helper import get_sleep_duration
 
     duration = get_sleep_duration(
         current_try=23, min_sleep_millis=10, max_sleep_millis=5000
@@ -34,7 +34,7 @@ def test_get_sleep_duration_max():
 
 
 def test_get_sleep_duration_negative_try():
-    from dask_cloudprovider.providers.aws.helper import get_sleep_duration
+    from dask_cloudprovider.aws.helper import get_sleep_duration
 
     duration = get_sleep_duration(
         current_try=-1, min_sleep_millis=10, max_sleep_millis=5000

dask_cloudprovider/azure/__init__.py

@@ -0,0 +1 @@
+from .azureml import AzureMLCluster

dask_cloudprovider/providers/azure/azureml.py → dask_cloudprovider/azure/azureml.py

@@ -152,7 +152,7 @@ class AzureMLCluster(Cluster):
     First, import all necessary modules.
 
     >>> from azureml.core import Workspace
-    >>> from dask_cloudprovider import AzureMLCluster
+    >>> from dask_cloudprovider.azure import AzureMLCluster
 
     Next, create the ``Workspace`` object given your AzureML ``Workspace`` parameters. Check
     more in the AzureML documentation for

dask_cloudprovider/azure/tests/test_azureml_gpu.py

@@ -0,0 +1,7 @@
+import pytest
+
+pytest.importorskip("azureml")
+
+
+def test_aml():
+    from dask_cloudprovider.azure import AzureMLCluster  # noqa

dask_cloudprovider/cli/ecs.py

@@ -6,7 +6,7 @@
 from distributed.cli.utils import check_python_3, install_signal_handlers
 from tornado.ioloop import IOLoop, TimeoutError
 
-from dask_cloudprovider import ECSCluster
+from dask_cloudprovider.aws import ECSCluster
 
 
 logger = logging.getLogger(__name__)

dask_cloudprovider/digitalocean/__init__.py

@@ -0,0 +1 @@
+from .droplet import DropletCluster

dask_cloudprovider/providers/digitalocean/droplet.py → dask_cloudprovider/digitalocean/droplet.py

@@ -1,7 +1,7 @@
 import asyncio
 
 import dask
-from dask_cloudprovider.providers.generic.vmcluster import (
+from dask_cloudprovider.generic.vmcluster import (
     VMCluster,
     VMInterface,
     SchedulerMixin,
@@ -142,7 +142,7 @@ class DropletCluster(VMCluster):
 
     Create the cluster.
 
-    >>> from dask_cloudprovider import DropletCluster
+    >>> from dask_cloudprovider.digitalocean import DropletCluster
     >>> cluster = DropletCluster(n_workers=1)
     Creating scheduler instance
     Created droplet dask-38b817c1-scheduler

dask_cloudprovider/providers/digitalocean/tests/test_droplet.py → dask_cloudprovider/digitalocean/tests/test_droplet.py

@@ -4,7 +4,7 @@
 
 digitalocean = pytest.importorskip("digitalocean")
 
-from dask_cloudprovider.providers.digitalocean.droplet import DropletCluster
+from dask_cloudprovider.digitalocean.droplet import DropletCluster
 from dask.distributed import Client
 from distributed.core import Status
 

dask_cloudprovider/providers/generic/tests/test_vmcluster.py → dask_cloudprovider/generic/tests/test_vmcluster.py

@@ -1,6 +1,6 @@
 import pytest
 
-from dask_cloudprovider.providers.generic.vmcluster import VMCluster
+from dask_cloudprovider.generic.vmcluster import VMCluster
 from distributed.core import Status
 
 

dask_cloudprovider/tests/test_imports.py

@@ -0,0 +1,22 @@
+import pytest
+
+
+def test_imports():
+    from dask_cloudprovider.aws import EC2Cluster  # noqa
+    from dask_cloudprovider.aws import ECSCluster  # noqa
+    from dask_cloudprovider.aws import FargateCluster  # noqa
+    from dask_cloudprovider.azure import AzureMLCluster  # noqa
+    from dask_cloudprovider.digitalocean import DropletCluster  # noqa
+
+
+def test_import_exceptions():
+    with pytest.raises(ImportError):
+        from dask_cloudprovider import EC2Cluster  # noqa
+    with pytest.raises(ImportError):
+        from dask_cloudprovider import ECSCluster  # noqa
+    with pytest.raises(ImportError):
+        from dask_cloudprovider import FargateCluster  # noqa
+    with pytest.raises(ImportError):
+        from dask_cloudprovider import AzureMLCluster  # noqa
+    with pytest.raises(ImportError):
+        from dask_cloudprovider import DropletCluster  # noqa

doc/source/aws.rst

@@ -1,7 +1,7 @@
 Amazon Web Services (AWS)
 =========================
 
-.. currentmodule:: dask_cloudprovider
+.. currentmodule:: dask_cloudprovider.aws
 
 .. autosummary::
    EC2Cluster

doc/source/azure.rst

@@ -1,7 +1,7 @@
 Microsoft Azure
 ===============
 
-.. currentmodule:: dask_cloudprovider
+.. currentmodule:: dask_cloudprovider.azure
 
 .. autosummary::
    AzureMLCluster

doc/source/config.rst

@@ -29,7 +29,7 @@ to give the scheduler in megabytes. This can be configured in the following ways
 
 .. code-block:: python
 
-   from dask_cloudprovider import FargateCluster
+   from dask_cloudprovider.aws import FargateCluster
 
    cluster = FargateCluster(
        scheduler_mem=8192

doc/source/digitalocean.rst

@@ -1,7 +1,7 @@
 DigitalOcean
 ============
 
-.. currentmodule:: dask_cloudprovider
+.. currentmodule:: dask_cloudprovider.digitalocean
 
 .. autosummary::
    DropletCluster

doc/source/gpus.rst

@@ -12,7 +12,7 @@ Each cluster manager handles this differently but generally you will need to con
 - Ensure the OS/Docker image has the NVIDIA drivers. For Docker images it is recommended to use the [RAPIDS images](https://hub.docker.com/r/rapidsai/rapidsai/).
 - Set the ``worker_module`` config option to ``dask_cuda.cli.dask_cuda_worker`` or ``worker_command`` option to ``dask-cuda-worker``.
 
-In the following AWS :class:`EC2Cluster` example we set the ``ami`` to be a Deep Learning AMI with NVIDIA drivers, the ``docker_image`` to RAPIDS, the ``instance_type``
+In the following AWS :class:`dask_cloudprovider.aws.EC2Cluster` example we set the ``ami`` to be a Deep Learning AMI with NVIDIA drivers, the ``docker_image`` to RAPIDS, the ``instance_type``
 to ``p3.2xlarge`` which has one NVIDIA Tesla V100 and the ``worker_module`` to ``dask_cuda.cli.dask_cuda_worker``.
 
 .. code-block:: python

doc/source/index.rst

@@ -11,7 +11,7 @@ will result in cloud resources being created for you.
 
 .. code-block:: python
 
-    from dask_cloudprovider import FargateCluster
+    from dask_cloudprovider.aws import FargateCluster
     cluster = FargateCluster(
         # Cluster manager specific config kwargs
     )
@@ -43,7 +43,7 @@ this code.
 
 .. code-block:: python
 
-    from dask_cloudprovider import FargateCluster
+    from dask_cloudprovider.aws import FargateCluster
     from dask.distributed import Client
 
     with FargateCluster(...) as cluster:

doc/source/packer.rst

@@ -58,7 +58,7 @@ returns the cloud-init file that would be generated.
 
 .. code-block:: python
 
-    from dask_cloudprovider import EC2Cluster
+    from dask_cloudprovider.aws import EC2Cluster
 
     cloud_init_config = EC2Cluster.get_cloud_init(
         # Pass any kwargs here you would normally pass to ``EC2Cluster``
@@ -196,7 +196,7 @@ Then to use our new image we can create an ``EC2Cluster`` specifying the AMI and
 .. code-block:: python
 
     from dask.distributed import Client
-    from dask_cloudprovider import EC2Cluster
+    from dask_cloudprovider.aws import EC2Cluster
 
     cluster = EC2Cluster(
         ami="ami-064f8db7634d19647",  # AMI ID provided by Packer
@@ -214,7 +214,7 @@ To launch `RAPIDS <https://rapids.ai/>`_ on AWS EC2 we can select a GPU instance
 
 .. code-block:: python
 
-    from dask_cloudprovider import EC2Cluster
+    from dask_cloudprovider.aws import EC2Cluster
 
     cluster = EC2Cluster(
         ami="ami-0c7c7d78f752f8f17",  # Deep Learning AMI (this ID varies by region so find yours in the AWS Console)
@@ -311,7 +311,7 @@ We can then run our code snippet again but this time it will take less than 5 mi
 .. code-block:: python
 
     from dask.distributed import Client
-    from dask_cloudprovider import EC2Cluster
+    from dask_cloudprovider.aws import EC2Cluster
 
     cluster = EC2Cluster(
         ami="ami-04e5539cb82859e69",  # AMI ID provided by Packer