Fix an import issue introduced by aio-libs/aiobotocore#877 (#307)

dask · Aug 25, 2021 · fcc6cf3 · fcc6cf3
1 parent a344aec
commit fcc6cf3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/dask_cloudprovider/aws/ec2.py b/dask_cloudprovider/aws/ec2.py
@@ -17,7 +17,7 @@
 from dask_cloudprovider.utils.timeout import Timeout
 
 try:
-    import aiobotocore
+    from aiobotocore.session import get_session
     import botocore.exceptions
 except ImportError as e:
     msg = (
@@ -407,7 +407,7 @@ def __init__(
         debug=False,
         **kwargs,
     ):
-        self.boto_session = aiobotocore.get_session()
+        self.boto_session = get_session()
         self.config = dask.config.get("cloudprovider.ec2", {})
         self.scheduler_class = EC2Scheduler
         self.worker_class = EC2Worker

diff --git a/dask_cloudprovider/aws/ecs.py b/dask_cloudprovider/aws/ecs.py
@@ -24,7 +24,7 @@
 
 try:
     from botocore.exceptions import ClientError
-    import aiobotocore
+    from aiobotocore.session import get_session
 except ImportError as e:
     msg = (
         "Dask Cloud Provider AWS requirements are not installed.\n\n"
@@ -729,7 +729,7 @@ def __init__(
         self._region_name = region_name
         self._platform_version = platform_version
         self._lock = asyncio.Lock()
-        self.session = aiobotocore.get_session()
+        self.session = get_session()
         super().__init__(**kwargs)
 
     def _client(self, name: str):
@@ -1390,7 +1390,7 @@ async def _cleanup_stale_resources():
 
     """
     # Clean up clusters (clusters with no running tasks)
-    session = aiobotocore.get_session()
+    session = get_session()
     async with session.create_client("ecs") as ecs:
         active_clusters = []
         clusters_to_delete = []