From a90e1f9b14f87b4b67883726cd94b5611b5d9b4b Mon Sep 17 00:00:00 2001
From: Thomas Carmet <8408330+tcarmet@users.noreply.github.com>
Date: Thu, 5 Oct 2023 00:00:48 -0700
Subject: [PATCH] PTFE-1041 setup capability to spawn GCP spot instances (#455)

---
 runner_manager/models/backend.py | 15 +++++++++++++++
 tests/unit/backend/test_gcp.py   | 11 +++++++++++
 2 files changed, 26 insertions(+)

diff --git a/runner_manager/models/backend.py b/runner_manager/models/backend.py
index 278f7b13..586321cd 100644
--- a/runner_manager/models/backend.py
+++ b/runner_manager/models/backend.py
@@ -10,6 +10,7 @@
     Items,
     Metadata,
     NetworkInterface,
+    Scheduling,
 )
 from mypy_boto3_ec2.literals import InstanceTypeType, VolumeTypeType
 from mypy_boto3_ec2.type_defs import (
@@ -110,6 +111,7 @@ class GCPInstanceConfig(InstanceConfig):
     image_family: str = "ubuntu-2004-lts"
     image_project: str = "ubuntu-os-cloud"
     machine_type: str = "e2-small"
+    spot: bool = False
     network: str = "global/networks/default"
     enable_nested_virtualization: bool = True
     labels: Optional[Dict[str, str]] = {}
@@ -124,6 +126,18 @@ class GCPInstanceConfig(InstanceConfig):
     class Config:
         arbitrary_types_allowed = True
 
+    @property
+    def scheduling(self) -> Scheduling:
+        """Configure scheduling."""
+        if self.spot:
+            return Scheduling(
+                provisioning_model="SPOT", instance_termination_action="DELETE"
+            )
+        else:
+            return Scheduling(
+                provisioning_model="STANDARD", instance_termination_action="DEFAULT"
+            )
+
     def configure_metadata(self, runner: Runner) -> Metadata:
         items: List[Items] = []
         env: RunnerEnv = self.runner_env(runner)
@@ -148,6 +162,7 @@ def configure_instance(self, runner: Runner) -> Instance:
             advanced_machine_features=AdvancedMachineFeatures(
                 enable_nested_virtualization=self.enable_nested_virtualization
             ),
+            scheduling=self.scheduling,
         )
 
 
diff --git a/tests/unit/backend/test_gcp.py b/tests/unit/backend/test_gcp.py
index bb8fa1c2..031c9a8c 100644
--- a/tests/unit/backend/test_gcp.py
+++ b/tests/unit/backend/test_gcp.py
@@ -62,6 +62,17 @@ def test_gcp_instance(runner: Runner):
     assert startup is True
 
 
+def test_gcp_spot_config(runner: Runner):
+    gcp_instance: GCPInstanceConfig = GCPInstanceConfig(spot=True)
+    instance = gcp_instance.configure_instance(runner)
+    assert instance.scheduling.provisioning_model == "SPOT"
+    assert instance.scheduling.instance_termination_action == "DELETE"
+    gcp_instance: GCPInstanceConfig = GCPInstanceConfig(spot=False)
+    instance = gcp_instance.configure_instance(runner)
+    assert instance.scheduling.provisioning_model == "STANDARD"
+    assert instance.scheduling.instance_termination_action == "DEFAULT"
+
+
 @mark.skipif(
     not os.getenv("GOOGLE_APPLICATION_CREDENTIALS"), reason="GCP credentials not found"
 )