Add MachineMaintenancePolicy utility

This is a helper which recommends a maintenance policy as a function of
the given machine type.

For example, most accelerator and metal machines require the `TERMINATE`
policy, while most other machines default to `MIGRATE`.

PiperOrigin-RevId: 934421171

Author: a-r-n

fixtures.go

@@ -20,6 +20,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"regexp"
 	"strings"
 
 	daisy "github.com/GoogleCloudPlatform/compute-daisy"
@@ -53,6 +54,24 @@ const (
 	DefaultMachineType = "n1-standard-1"
 )
 
+var (
+	// knownTerminateOnlyMachineRegexes is a list of regexes for machine types that must use the
+	// TERMINATE maintenance policy.
+	knownTerminateOnlyMachineRegexes = []*regexp.Regexp{
+		// Accelerator families.
+		regexp.MustCompile(`^a\d+[x]?-.*`),
+		regexp.MustCompile(`^g\d+-.*`),
+
+		// TPU families.
+		regexp.MustCompile(`^ct\d.*`),
+		regexp.MustCompile(`^tpu\d.*`),
+
+		// Special cases.
+		regexp.MustCompile(`^h4d-.*`),
+		regexp.MustCompile(`^z3-.*`),
+	}
+)
+
 // TestVM is a test VM.
 type TestVM struct {
 	name         string
@@ -147,6 +166,23 @@ func DiskTypeNeeded(machineType string) string {
 	return HyperdiskBalanced
 }
 
+// MachineMaintenancePolicy returns the recommended maintenance policy for the given machine type.
+// For instance, metal and accelerator machines must use TERMINATE, while most others can use
+// MIGRATE.
+func MachineMaintenancePolicy(machineType string) string {
+	if IsMetal(machineType) {
+		return "TERMINATE"
+	}
+
+	for _, regex := range knownTerminateOnlyMachineRegexes {
+		if regex.MatchString(machineType) {
+			return "TERMINATE"
+		}
+	}
+
+	return "MIGRATE"
+}
+
 // CreateTestVM adds the necessary steps to create a VM with the specified
 // name to the workflow.
 func (t *TestWorkflow) CreateTestVM(name string) (*TestVM, error) {

fixtures_test.go

@@ -979,3 +979,50 @@ func TestIsMetal(t *testing.T) {
 		})
 	}
 }
+
+func TestMachineMaintenancePolicy(t *testing.T) {
+	cases := []struct {
+		name        string
+		machineType string
+		want        string
+	}{
+		{
+			name:        "regular machine",
+			machineType: "c3-highcpu-176",
+			want:        "MIGRATE",
+		},
+		{
+			name:        "metal machine",
+			machineType: "c3-highcpu-192-metal",
+			want:        "TERMINATE",
+		},
+		{
+			name:        "accelerator machine",
+			machineType: "a2-highgpu-1g",
+			want:        "TERMINATE",
+		},
+		{
+			name:        "h4d special case",
+			machineType: "h4d-standard-4",
+			want:        "TERMINATE",
+		},
+		{
+			name:        "z3 special case",
+			machineType: "z3-highmem-176-standardlssd",
+			want:        "TERMINATE",
+		},
+		{
+			name:        "unknown machine",
+			machineType: "unknown-machine",
+			want:        "MIGRATE",
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.machineType, func(t *testing.T) {
+			got := MachineMaintenancePolicy(tc.machineType)
+			if got != tc.want {
+				t.Errorf("MachineMaintenancePolicy(%q) = %v, want %v", tc.machineType, got, tc.want)
+			}
+		})
+	}
+}