Add on_host_maintenance option in config (#253)

* add on_host_maintenance option in config * lint Co-authored-by: Jacob Tomlinson <jacobtomlinson@users.noreply.github.com>
dask · Feb 23, 2021 · adca69c · adca69c
1 parent ce0d5e5
commit adca69c
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/dask_cloudprovider/cloudprovider.yaml b/dask_cloudprovider/cloudprovider.yaml
@@ -96,6 +96,7 @@ cloudprovider:
     zone: "us-east1-c" # the zone of of where to launch the instances
     network: "default" # the network/subnetwork in GCP to use
     projectid: "" # name of the google cloud project
+    on_host_maintenance: "TERMINATE"
     machine_type: "n1-standard-1" # size of the machine type to use
     filesystem_size: 50 # amount in GBs of hard drive space to allocate
     ngpus: "" # number of GPUs to use

diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py
@@ -52,6 +52,7 @@ def __init__(
         projectid=None,
         machine_type=None,
         filesystem_size=None,
+        on_host_maintenance=None,
         source_image=None,
         docker_image=None,
         network=None,
@@ -68,6 +69,9 @@ def __init__(
 
         self.cluster = cluster
         self.config = config
+        self.on_host_maintenance = on_host_maintenance or self.config.get(
+            "on_host_maintenance"
+        )
         self.projectid = projectid or self.config.get("projectid")
         self.zone = zone or self.config.get("zone")
 
@@ -149,7 +153,7 @@ def create_gcp_config(self):
             "labels": {"container-vm": "dask-cloudprovider"},
             "scheduling": {
                 "preemptible": ("true" if self.preemptible else "false"),
-                "onHostMaintenance": "TERMINATE",
+                "onHostMaintenance": self.on_host_maintenance.upper(),
                 "automaticRestart": ("false" if self.preemptible else "true"),
                 "nodeAffinities": [],
             },
@@ -425,6 +429,8 @@ class GCPCluster(VMCluster):
         You can see a list of GPUs available in each zone with ``gcloud compute accelerator-types list``.
     filesystem_size: int (optional)
         The VM filesystem size in GB. Defaults to ``50``.
+    on_host_maintenance: str (optional)
+        The Host Maintenance GCP option.  Defaults to ``TERMINATE``.
     n_workers: int (optional)
         Number of workers to initialise the cluster with. Defaults to ``0``.
     bootstrap: bool (optional)
@@ -506,7 +512,7 @@ class GCPCluster(VMCluster):
     Source Image: projects/ubuntu-os-cloud/global/images/ubuntu-minimal-1804-bionic-v20201014
     Docker Image: daskdev/dask:latest
     Machine Type: n1-standard-1
-    Filesytsem Size: 50
+    Filesystem Size: 50
     N-GPU Type:
     Zone: us-east1-c
     Creating scheduler instance
@@ -531,6 +537,7 @@ def __init__(
         zone=None,
         network=None,
         machine_type=None,
+        on_host_maintenance=None,
         source_image=None,
         docker_image=None,
         ngpus=None,
@@ -564,6 +571,8 @@ def __init__(
             "source_image": source_image or self.config.get("source_image"),
             "docker_image": docker_image or self.config.get("docker_image"),
             "filesystem_size": filesystem_size or self.config.get("filesystem_size"),
+            "on_host_maintenance": on_host_maintenance
+            or self.config.get("on_host_maintenance"),
             "zone": zone or self.config.get("zone"),
             "machine_type": self.machine_type,
             "ngpus": ngpus or self.config.get("ngpus"),