From 4be8cbc7841009d01923f8248e0053a9c36ab796 Mon Sep 17 00:00:00 2001
From: jonathanmetzman <31354670+jonathanmetzman@users.noreply.github.com>
Date: Mon, 25 Nov 2024 14:07:40 -0500
Subject: [PATCH] Schedule preemptible batch jobs properly. (#4433)

---
 src/clusterfuzz/_internal/google_cloud_utils/batch.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/src/clusterfuzz/_internal/google_cloud_utils/batch.py b/src/clusterfuzz/_internal/google_cloud_utils/batch.py
index 520460c4a7..8fb8b8b639 100644
--- a/src/clusterfuzz/_internal/google_cloud_utils/batch.py
+++ b/src/clusterfuzz/_internal/google_cloud_utils/batch.py
@@ -144,6 +144,16 @@ def _get_task_spec(batch_workload_spec):
   return task_spec
 
 
+def _set_preemptible(instance_policy,
+                     batch_workload_spec: BatchWorkloadSpec) -> None:
+  if batch_workload_spec.preemptible:
+    instance_policy.provisioning_model = (
+        batch.AllocationPolicy.ProvisioningModel.PREEMPTIBLE)
+  else:
+    instance_policy.provisioning_model = (
+        batch.AllocationPolicy.ProvisioningModel.STANDARD)
+
+
 def _get_allocation_policy(spec):
   """Returns the allocation policy for a BatchWorkloadSpec."""
   disk = batch.AllocationPolicy.Disk()
@@ -153,6 +163,7 @@ def _get_allocation_policy(spec):
   instance_policy = batch.AllocationPolicy.InstancePolicy()
   instance_policy.boot_disk = disk
   instance_policy.machine_type = spec.machine_type
+  _set_preemptible(instance_policy, spec)
   instances = batch.AllocationPolicy.InstancePolicyOrTemplate()
   instances.policy = instance_policy