From 4f3c8575295d2dbad9ae84654277ca153ab42a88 Mon Sep 17 00:00:00 2001
From: Tom White <tom.e.white@gmail.com>
Date: Tue, 9 Feb 2021 11:20:37 +0000
Subject: [PATCH 1/3] Add option for preemptible instances on GCP

---
 dask_cloudprovider/gcp/instances.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py
index af8f1371..0b87cd56 100644
--- a/dask_cloudprovider/gcp/instances.py
+++ b/dask_cloudprovider/gcp/instances.py
@@ -61,6 +61,7 @@ def __init__(
         bootstrap=None,
         gpu_instance=None,
         auto_shutdown=None,
+        preemptible=None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -84,6 +85,7 @@ def __init__(
         self.gpu_instance = gpu_instance
         self.bootstrap = bootstrap
         self.auto_shutdown = auto_shutdown
+        self.preemptible = preemptible
 
         self.general_zone = "-".join(self.zone.split("-")[:2])  # us-east1-c -> us-east1
 
@@ -146,9 +148,9 @@ def create_gcp_config(self):
             },
             "labels": {"container-vm": "dask-cloudprovider"},
             "scheduling": {
-                "preemptible": "false",
+                "preemptible": ("true" if self.preemptible else "false"),
                 "onHostMaintenance": "TERMINATE",
-                "automaticRestart": "true",
+                "automaticRestart": ("false" if self.preemptible else "true"),
                 "nodeAffinities": [],
             },
             "shieldedInstanceConfig": {
@@ -447,6 +449,8 @@ class GCPCluster(VMCluster):
         Configures communication security in this cluster. Can be a security
         object, or True. If True, temporary self-signed credentials will
         be created automatically. Default is ``True``.
+    preemptible: bool (optional)
+        Whether to use preemptible instances in this cluster. Defaults to ``False``.
 
     Examples
     --------
@@ -533,6 +537,7 @@ def __init__(
         filesystem_size=None,
         auto_shutdown=None,
         bootstrap=True,
+        preemptible=None,
         **kwargs,
     ):
 
@@ -566,6 +571,7 @@ def __init__(
             "gpu_instance": self.gpu_instance,
             "bootstrap": self.bootstrap,
             "auto_shutdown": self.auto_shutdown,
+            "preemptible": preemptible if preemptible is not None else self.config.get("preemptible"),
         }
         self.scheduler_options = {**self.options}
         self.worker_options = {**self.options}

From 1828882053d623e967d5c150eda01c3433a03dec Mon Sep 17 00:00:00 2001
From: Tom White <tom.e.white@gmail.com>
Date: Tue, 9 Feb 2021 15:31:31 +0000
Subject: [PATCH 2/3] Change default and fix formatting

---
 dask_cloudprovider/gcp/instances.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py
index 0b87cd56..922f091d 100644
--- a/dask_cloudprovider/gcp/instances.py
+++ b/dask_cloudprovider/gcp/instances.py
@@ -61,7 +61,7 @@ def __init__(
         bootstrap=None,
         gpu_instance=None,
         auto_shutdown=None,
-        preemptible=None,
+        preemptible=False,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -571,7 +571,9 @@ def __init__(
             "gpu_instance": self.gpu_instance,
             "bootstrap": self.bootstrap,
             "auto_shutdown": self.auto_shutdown,
-            "preemptible": preemptible if preemptible is not None else self.config.get("preemptible"),
+            "preemptible": preemptible
+            if preemptible is not None
+            else self.config.get("preemptible"),
         }
         self.scheduler_options = {**self.options}
         self.worker_options = {**self.options}

From acf7846898bd71331d7a7f992f45472dd41b32bc Mon Sep 17 00:00:00 2001
From: Tom White <tom.e.white@gmail.com>
Date: Mon, 15 Feb 2021 12:54:03 +0000
Subject: [PATCH 3/3] Scheduler instances are not preemptible

---
 dask_cloudprovider/gcp/instances.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py
index 922f091d..7ce9113b 100644
--- a/dask_cloudprovider/gcp/instances.py
+++ b/dask_cloudprovider/gcp/instances.py
@@ -270,6 +270,7 @@ class GCPScheduler(SchedulerMixin, GCPInstance):
     """Scheduler running in a GCP instance."""
 
     def __init__(self, *args, **kwargs):
+        kwargs.pop("preemptible", None)  # scheduler instances are not preemptible
         super().__init__(*args, **kwargs)
 
     async def start(self):
@@ -450,7 +451,7 @@ class GCPCluster(VMCluster):
         object, or True. If True, temporary self-signed credentials will
         be created automatically. Default is ``True``.
     preemptible: bool (optional)
-        Whether to use preemptible instances in this cluster. Defaults to ``False``.
+        Whether to use preemptible instances for workers in this cluster. Defaults to ``False``.
 
     Examples
     --------