Merge pull request #200 from mrocklin/status

Use Status enum

dask_cloudprovider/aws/ecs.py

@@ -20,6 +20,7 @@
 
 from distributed.deploy.spec import SpecCluster
 from distributed.utils import warn_on_duration
+from distributed.core import Status
 
 try:
     from botocore.exceptions import ClientError
@@ -153,7 +154,7 @@ def __init__(
         self._fargate_use_private_ip = fargate_use_private_ip
         self.kwargs = kwargs
         self.task_kwargs = task_kwargs
-        self.status = "created"
+        self.status = Status.created
 
     def __await__(self):
         async def _():
@@ -301,7 +302,7 @@ async def start(self):
             self.public_ip = interface["Association"]["PublicIp"]
         self.private_ip = interface["PrivateIpAddresses"][0]["PrivateIpAddress"]
         await self._set_address_from_logs()
-        self.status = "running"
+        self.status = Status.running
 
     async def close(self, **kwargs):
         if self.task:
@@ -311,7 +312,7 @@ async def close(self, **kwargs):
             while self.task["lastStatus"] in ["RUNNING"]:
                 await asyncio.sleep(1)
                 await self._update_task()
-        self.status = "closed"
+        self.status = Status.closed
 
     @property
     def task_id(self):
@@ -737,11 +738,11 @@ def _client(self, name: str):
     async def _start(
         self,
     ):
-        while self.status == "starting":
+        while self.status == Status.starting:
             await asyncio.sleep(0.01)
-        if self.status == "running":
+        if self.status == Status.running:
             return
-        if self.status == "closed":
+        if self.status == Status.closed:
             raise ValueError("Cluster is closed")
 
         self.config = dask.config.get("cloudprovider.ecs", {})

dask_cloudprovider/azureml/azureml.py

@@ -8,7 +8,7 @@
 import dask
 from contextlib import suppress
 from distributed.deploy.cluster import Cluster
-from distributed.core import rpc
+from distributed.core import rpc, Status
 from distributed.utils import LoopRunner, log_errors, format_bytes
 from tornado.ioloop import PeriodicCallback
 
@@ -1053,7 +1053,7 @@ def scale_down(self, workers=1):
 
     # close cluster
     async def _close(self):
-        if self.status == "closed":
+        if self.status == Status.closed:
             return
 
         while self.workers_list:
@@ -1065,7 +1065,7 @@ async def _close(self):
             self.run.complete()
             self.run.cancel()
 
-        self.status = "closed"
+        self.status = Status.closed
         self.__print_message("Scheduler and workers are disconnected.")
 
         if self.portforward_proc is not None:

dask_cloudprovider/cli/ecs.py

@@ -4,6 +4,7 @@
 
 import click
 from distributed.cli.utils import check_python_3, install_signal_handlers
+from distributed.core import Status
 from tornado.ioloop import IOLoop, TimeoutError
 
 from dask_cloudprovider.aws import ECSCluster
@@ -214,7 +215,7 @@ def main(
 
     async def run():
         logger.info("Ready")
-        while cluster.status != "closed":
+        while cluster.status != Status.closed:
             await sleep(0.2)
 
     def on_signal(signum):

requirements.txt

@@ -1,3 +1,3 @@
-dask>=2.2.0
+dask>=2.30
 distributed>=2.30
 jinja2