Allow creating VMs with implicit CPU and memory (#243)

* Allow creating VMs with implicit CPU and memory

* Clarify why cpu/memory can be 0 a bit better

* Controller(API): don't forget to update DefaultCPU and DefaultMemory

* Add an integration test for implicit CPU and memory

internal/command/get/vm.go

@@ -65,8 +65,21 @@ func runGetVM(cmd *cobra.Command, args []string) error {
 	table.AddRow("Created", createdAtInfo)
 	table.AddRow("Image", vm.Image)
 	table.AddRow("Image pull policy", vm.ImagePullPolicy)
-	table.AddRow("CPU", vm.CPU)
-	table.AddRow("Memory", vm.Memory)
+
+	cpu := vm.CPU
+	if cpu == 0 {
+		// Implicit CPU assignment, CPU will always be 0
+		cpu = vm.AssignedCPU
+	}
+	table.AddRow("CPU", cpu)
+
+	memory := vm.Memory
+	if memory == 0 {
+		// Implicit memory assignment, memory will always be 0
+		memory = vm.AssignedMemory
+	}
+	table.AddRow("Memory", memory)
+
 	table.AddRow("Softnet enabled", vm.NetSoftnet)
 	table.AddRow("Bridged networking interface", nonEmptyOrNone(vm.NetBridged))
 	table.AddRow("Headless mode", vm.Headless)

internal/command/worker/run.go

@@ -29,6 +29,8 @@ var bootstrapTokenStdin bool
 var logFilePath string
 var stringToStringResources map[string]string
 var noPKI bool
+var defaultCPU uint64
+var defaultMemory uint64
 var debug bool
 
 func newRunCommand() *cobra.Command {
@@ -53,6 +55,10 @@ func newRunCommand() *cobra.Command {
 		"do not use the host's root CA set and instead validate the Controller's presented "+
 			"certificate using a bootstrap token (or manually via fingerprint, "+
 			"if no bootstrap token is provided)")
+	cmd.PersistentFlags().Uint64Var(&defaultCPU, "default-cpu", 4, "number of CPUs to use for VMs "+
+		"that do not explicitly specify a value")
+	cmd.PersistentFlags().Uint64Var(&defaultMemory, "default-memory", 8*1024, "megabytes of memory "+
+		"to use for VMs that do not explicitly specify a value")
 	cmd.PersistentFlags().BoolVar(&debug, "debug", false, "enable debug logging")
 
 	return cmd
@@ -116,6 +122,7 @@ func runWorker(cmd *cobra.Command, args []string) (err error) {
 		controllerClient,
 		worker.WithName(name),
 		worker.WithResources(resources),
+		worker.WithDefaultCPUAndMemory(defaultCPU, defaultMemory),
 		worker.WithLogger(logger),
 	)
 	if err != nil {

internal/controller/api_vms.go

@@ -34,12 +34,6 @@ func (controller *Controller) createVM(ctx *gin.Context) responder.Responder {
 	if vm.Image == "" {
 		return responder.JSON(http.StatusPreconditionFailed, NewErrorResponse("VM image is empty"))
 	}
-	if vm.CPU == 0 {
-		return responder.JSON(http.StatusPreconditionFailed, NewErrorResponse("VM CPU is zero"))
-	}
-	if vm.Memory == 0 {
-		return responder.JSON(http.StatusPreconditionFailed, NewErrorResponse("VM memory is zero"))
-	}
 
 	vm.Status = v1.VMStatusPending
 	vm.CreatedAt = time.Now()

internal/controller/api_workers.go

@@ -98,6 +98,8 @@ func (controller *Controller) createWorker(ctx *gin.Context) responder.Responder
 
 		dbWorker.LastSeen = worker.LastSeen
 		dbWorker.Resources = worker.Resources
+		dbWorker.DefaultCPU = worker.DefaultCPU
+		dbWorker.DefaultMemory = worker.DefaultMemory
 
 		if err := txn.SetWorker(*dbWorker); err != nil {
 			return responder.Error(err)

internal/controller/scheduler/scheduler.go

@@ -134,7 +134,7 @@ func (scheduler *Scheduler) RequestScheduling() {
 	}
 }
 
-//nolint:gocognit // this logic could be said to be considered even more complex if split into multiple functions
+//nolint:gocognit,gocyclo // this logic could be seen as even more complex if split into multiple functions
 func (scheduler *Scheduler) schedulingLoopIteration() error {
 	affectedWorkers := mapset.NewSet[string]()
 
@@ -289,6 +289,30 @@ NextVM:
 				unscheduledVM.Worker = worker.Name
 				unscheduledVM.ScheduledAt = time.Now()
 
+				// Fill out the actual CPU allocation
+				if unscheduledVM.CPU == 0 {
+					// Provide defaults for VMs with implicit CPU specification
+					if worker.DefaultCPU != 0 {
+						unscheduledVM.AssignedCPU = worker.DefaultCPU
+					} else {
+						unscheduledVM.AssignedCPU = 4
+					}
+				} else {
+					unscheduledVM.AssignedCPU = unscheduledVM.CPU
+				}
+
+				// Fill out the actual memory allocation
+				if unscheduledVM.Memory == 0 {
+					// Provide defaults for VMs with implicit memory specification
+					if worker.DefaultMemory != 0 {
+						unscheduledVM.AssignedMemory = worker.DefaultMemory
+					} else {
+						unscheduledVM.AssignedMemory = 8192
+					}
+				} else {
+					unscheduledVM.AssignedMemory = unscheduledVM.Memory
+				}
+
 				if err := txn.SetVM(unscheduledVM); err != nil {
 					return err
 				}
@@ -434,6 +458,8 @@ func (scheduler *Scheduler) healthCheckVM(txn storepkg.Transaction, vm v1.VM) er
 		vm.Status = v1.VMStatusPending
 		vm.StatusMessage = ""
 		vm.Worker = ""
+		vm.AssignedCPU = 0
+		vm.AssignedMemory = 0
 		vm.RestartedAt = time.Now()
 		vm.RestartCount++
 		vm.ScheduledAt = time.Time{}

internal/tests/implicit_cpu_memory_test.go

@@ -0,0 +1,64 @@
+package tests_test
+
+import (
+	"context"
+	"github.com/cirruslabs/orchard/internal/tests/devcontroller"
+	"github.com/cirruslabs/orchard/internal/tests/wait"
+	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
+	"github.com/stretchr/testify/require"
+	"testing"
+	"time"
+)
+
+func TestImplicitCPUMemory(t *testing.T) {
+	ctx := context.Background()
+
+	// Create a development environment
+	devClient, _, _ := devcontroller.StartIntegrationTestEnvironmentWithAdditionalOpts(t,
+		false, nil,
+		true, nil,
+	)
+
+	// Create a worker with default CPU and memory values
+	const workerName = "worker"
+
+	_, err := devClient.Workers().Create(ctx, v1.Worker{
+		Meta: v1.Meta{
+			Name: workerName,
+		},
+		Resources: map[string]uint64{
+			v1.ResourceTartVMs: 2,
+		},
+		DefaultCPU:    12,
+		DefaultMemory: 3456,
+	})
+	require.NoError(t, err)
+
+	// Create a VM with implicit CPU and memory
+	vmName := "test-vm"
+
+	require.NoError(t, devClient.VMs().Create(ctx, &v1.VM{
+		Meta: v1.Meta{
+			Name: vmName,
+		},
+		Image:  "example.com/doesnt/matter:latest",
+		Status: v1.VMStatusPending,
+	}))
+
+	// Wait for the VM to be assigned
+	require.True(t, wait.Wait(2*time.Minute, func() bool {
+		vm, err := devClient.VMs().Get(context.Background(), vmName)
+		require.NoError(t, err)
+
+		t.Logf("Waiting for the VM %s to be assigned to a worker", vmName)
+
+		return vm.Worker == workerName
+	}), "VM was %s expected to be assigned to the worker %q, but was assigned to the worker %q",
+		vmName, workerName)
+
+	// Ensure that the VM is using default CPU and memory values from the worker
+	vm, err := devClient.VMs().Get(context.Background(), vmName)
+	require.NoError(t, err)
+	require.EqualValues(t, 12, vm.AssignedCPU)
+	require.EqualValues(t, 3456, vm.AssignedMemory)
+}

internal/worker/option.go

@@ -19,6 +19,13 @@ func WithResources(resources v1.Resources) Option {
 	}
 }
 
+func WithDefaultCPUAndMemory(defaultCPU uint64, defaultMemory uint64) Option {
+	return func(worker *Worker) {
+		worker.defaultCPU = defaultCPU
+		worker.defaultMemory = defaultMemory
+	}
+}
+
 func WithLogger(logger *zap.Logger) Option {
 	return func(worker *Worker) {
 		worker.logger = logger.Sugar()

internal/worker/worker.go

@@ -33,6 +33,9 @@ type Worker struct {
 	pollTicker    *time.Ticker
 	resources     v1.Resources
 
+	defaultCPU    uint64
+	defaultMemory uint64
+
 	vmPullTimeHistogram metric.Float64Histogram
 
 	logger *zap.SugaredLogger
@@ -191,9 +194,11 @@ func (worker *Worker) registerWorker(ctx context.Context) error {
 		Meta: v1.Meta{
 			Name: worker.name,
 		},
-		Resources: worker.resources,
-		LastSeen:  time.Now(),
-		MachineID: platformUUID,
+		Resources:     worker.resources,
+		LastSeen:      time.Now(),
+		MachineID:     platformUUID,
+		DefaultCPU:    worker.defaultCPU,
+		DefaultMemory: worker.defaultMemory,
 	})
 	if err != nil {
 		return err

pkg/resource/v1/v1.go

@@ -34,6 +34,19 @@ type VM struct {
 	// Worker field is set by the Controller to assign this VM to a specific Worker.
 	Worker string `json:"worker,omitempty"`
 
+	// AssignedCPU is set by the Controller when the VM is scheduled.
+	//
+	// It's set to CPU when CPU non-zero, otherwise the value is taken from
+	// Worker's DefaultCPU field. If Worker's DefaultCPU field is zero, it defaults
+	// to 4.
+	AssignedCPU uint64 `json:"assignedCPU,omitempty"`
+	// AssignedMemory is set by the Controller
+	//
+	// It's set to Memory when Memory non-zero, otherwise the value is taken from
+	// Worker's DefaultCPU field. If Worker's DefaultCPU field is zero, it defaults
+	// to 8192.
+	AssignedMemory uint64 `json:"assignedMemory,omitempty"`
+
 	Username      string    `json:"username,omitempty"`
 	Password      string    `json:"password,omitempty"`
 	StartupScript *VMScript `json:"startup_script,omitempty"`

pkg/resource/v1/worker.go

@@ -14,6 +14,13 @@ type Worker struct {
 	// Resources available on this Worker.
 	Resources Resources `json:"resources,omitempty"`
 
+	// DefaultCPU is the amount of CPUs to assign to a VM
+	// when it doesn't explicitly request a specific amount.
+	DefaultCPU uint64 `json:"defaultCPU,omitempty"`
+	// DefaultMemory is the amount of memory to assign to a VM
+	// when it doesn't explicitly request a specific amount.
+	DefaultMemory uint64 `json:"defaultMemory,omitempty"`
+
 	Meta
 }