azure: Add support for using an external availability set and resource group

cmd/kola/options.go

@@ -124,6 +124,8 @@ func init() {
 	bv(&kola.AzureOptions.UsePrivateIPs, "azure-use-private-ips", false, "Assume nodes are reachable using private IP addresses")
 	bv(&kola.AzureOptions.UseIdentity, "azure-identity", false, "Use VM managed identity for authentication (default false)")
 	sv(&kola.AzureOptions.DiskController, "azure-disk-controller", "default", "Use a specific disk-controller for storage (default \"default\", also \"nvme\" and \"scsi\")")
+	sv(&kola.AzureOptions.ResourceGroup, "azure-resource-group", "", "Deploy resources in an existing resource group")
+	sv(&kola.AzureOptions.AvailabilitySet, "azure-availability-set", "", "Deploy instances with an existing availibity set")
 
 	// do-specific options
 	sv(&kola.DOOptions.ConfigPath, "do-config-file", "", "DigitalOcean config file (default \"~/"+auth.DOConfigPath+"\")")

platform/api/azure/api.go

@@ -154,6 +154,10 @@ func New(opts *Options) (*API, error) {
 		client = management.NewAnonymousClient()
 	}
 
+	if opts.AvailabilitySet != "" && opts.ResourceGroup == "" {
+		return nil, fmt.Errorf("ResourceGroup must match AvailabilitySet")
+	}
+
 	api := &API{
 		client: client,
 		Opts:   opts,

platform/api/azure/instance.go

@@ -39,6 +39,21 @@ type Machine struct {
 	PublicIPName     string
 }
 
+func (a *API) getAvset() string {
+	if a.Opts.AvailabilitySet == "" {
+		return ""
+	}
+	return fmt.Sprintf("/subscriptions/%s/resourceGroups/%s/providers/Microsoft.Compute/availabilitySets/%s", a.Opts.SubscriptionID, a.Opts.ResourceGroup, a.Opts.AvailabilitySet)
+}
+
+func (a *API) getVMRG(rg string) string {
+	vmrg := rg
+	if a.Opts.ResourceGroup != "" {
+		vmrg = a.Opts.ResourceGroup
+	}
+	return vmrg
+}
+
 func (a *API) getVMParameters(name, userdata, sshkey, storageAccountURI string, ip *network.PublicIPAddress, nic *network.Interface) compute.VirtualMachine {
 	osProfile := compute.OSProfile{
 		AdminUsername: util.StrToPtr("core"),
@@ -156,10 +171,18 @@ func (a *API) getVMParameters(name, userdata, sshkey, storageAccountURI string,
 		}
 	}
 
+	availabilitySetID := a.getAvset()
+	if availabilitySetID != "" {
+		vm.VirtualMachineProperties.AvailabilitySet = &compute.SubResource{ID: &availabilitySetID}
+	}
+
 	return vm
 }
 
 func (a *API) CreateInstance(name, userdata, sshkey, resourceGroup, storageAccount string, network Network) (*Machine, error) {
+	// only VMs are created in the user supplied resource group, kola still manages a resource group
+	// for the gallery and storage account.
+	vmResourceGroup := a.getVMRG(resourceGroup)
 	subnet := network.subnet
 
 	ip, err := a.createPublicIP(resourceGroup)
@@ -181,22 +204,31 @@ func (a *API) CreateInstance(name, userdata, sshkey, resourceGroup, storageAccou
 	vmParams := a.getVMParameters(name, userdata, sshkey, fmt.Sprintf("https://%s.blob.core.windows.net/", storageAccount), ip, nic)
 	plog.Infof("Creating Instance %s", name)
 
-	future, err := a.compClient.CreateOrUpdate(context.TODO(), resourceGroup, name, vmParams)
+	clean := func() {
+		_, _ = a.compClient.Delete(context.TODO(), vmResourceGroup, name, &forceDelete)
+		_, _ = a.intClient.Delete(context.TODO(), resourceGroup, *nic.Name)
+		_, _ = a.ipClient.Delete(context.TODO(), resourceGroup, *ip.Name)
+	}
+
+	future, err := a.compClient.CreateOrUpdate(context.TODO(), vmResourceGroup, name, vmParams)
 	if err != nil {
+		clean()
 		return nil, err
 	}
 	err = future.WaitForCompletionRef(context.TODO(), a.compClient.Client)
 	if err != nil {
+		clean()
 		return nil, err
 	}
 	_, err = future.Result(a.compClient)
 	if err != nil {
+		clean()
 		return nil, err
 	}
 	plog.Infof("Instance %s created", name)
 
 	err = util.WaitUntilReady(5*time.Minute, 10*time.Second, func() (bool, error) {
-		vm, err := a.compClient.Get(context.TODO(), resourceGroup, name, "")
+		vm, err := a.compClient.Get(context.TODO(), vmResourceGroup, name, "")
 		if err != nil {
 			return false, err
 		}
@@ -209,13 +241,11 @@ func (a *API) CreateInstance(name, userdata, sshkey, resourceGroup, storageAccou
 	})
 	plog.Infof("Instance %s ready", name)
 	if err != nil {
-		_, _ = a.compClient.Delete(context.TODO(), resourceGroup, name, &forceDelete)
-		_, _ = a.intClient.Delete(context.TODO(), resourceGroup, *nic.Name)
-		_, _ = a.ipClient.Delete(context.TODO(), resourceGroup, *ip.Name)
+		clean()
 		return nil, fmt.Errorf("waiting for machine to become active: %v", err)
 	}
 
-	vm, err := a.compClient.Get(context.TODO(), resourceGroup, name, "")
+	vm, err := a.compClient.Get(context.TODO(), vmResourceGroup, name, "")
 	if err != nil {
 		return nil, err
 	}
@@ -245,6 +275,7 @@ func (a *API) CreateInstance(name, userdata, sshkey, resourceGroup, storageAccou
 // TerminateInstance deletes a VM created by CreateInstance. Public IP, NIC and
 // OS disk are deleted automatically together with the VM.
 func (a *API) TerminateInstance(machine *Machine, resourceGroup string) error {
+	resourceGroup = a.getVMRG(resourceGroup)
 	future, err := a.compClient.Delete(context.TODO(), resourceGroup, machine.ID, &forceDelete)
 	if err != nil {
 		return err
@@ -272,7 +303,8 @@ func (a *API) GetConsoleOutput(name, resourceGroup, storageAccount string) ([]by
 	k := *kr.Keys
 	key := *k[0].Value
 
-	vm, err := a.compClient.Get(context.TODO(), resourceGroup, name, compute.InstanceViewTypesInstanceView)
+	vmResourceGroup := a.getVMRG(resourceGroup)
+	vm, err := a.compClient.Get(context.TODO(), vmResourceGroup, name, compute.InstanceViewTypesInstanceView)
 	if err != nil {
 		return nil, fmt.Errorf("could not get VM: %v", err)
 	}

platform/api/azure/options.go

@@ -51,6 +51,10 @@ type Options struct {
 
 	// Azure Storage API endpoint suffix. If unset, the Azure SDK default will be used.
 	StorageEndpointSuffix string
-	// UseUserData can be use to enable custom data only or user-data only.
+	// UseUserData can be used to enable custom data only or user-data only.
 	UseUserData bool
+	// ResourceGroup is an existing resource group to deploy resources in.
+	ResourceGroup string
+	// AvailabilitySet is an existing availability set to deploy the instance in.
+	AvailabilitySet string
 }