From 581de320b97a305d60a04247fd50c326aefbb81a Mon Sep 17 00:00:00 2001 From: Nikolay Edigaryev Date: Thu, 6 Feb 2025 00:50:01 +0400 Subject: [PATCH] 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 | 17 +++++- internal/command/worker/run.go | 7 +++ internal/controller/api_vms.go | 6 -- internal/controller/api_workers.go | 2 + internal/controller/scheduler/scheduler.go | 28 +++++++++- internal/tests/implicit_cpu_memory_test.go | 64 ++++++++++++++++++++++ internal/worker/option.go | 7 +++ internal/worker/worker.go | 11 +++- pkg/resource/v1/v1.go | 13 +++++ pkg/resource/v1/worker.go | 7 +++ 10 files changed, 150 insertions(+), 12 deletions(-) create mode 100644 internal/tests/implicit_cpu_memory_test.go diff --git a/internal/command/get/vm.go b/internal/command/get/vm.go index b7800d6..712e94a 100644 --- a/internal/command/get/vm.go +++ b/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) diff --git a/internal/command/worker/run.go b/internal/command/worker/run.go index d7cee80..bc1b59e 100644 --- a/internal/command/worker/run.go +++ b/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 { diff --git a/internal/controller/api_vms.go b/internal/controller/api_vms.go index e35616e..85d1cfc 100644 --- a/internal/controller/api_vms.go +++ b/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() diff --git a/internal/controller/api_workers.go b/internal/controller/api_workers.go index adaf3ce..f5d43c4 100644 --- a/internal/controller/api_workers.go +++ b/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) diff --git a/internal/controller/scheduler/scheduler.go b/internal/controller/scheduler/scheduler.go index 0bcfa9a..5412a5b 100644 --- a/internal/controller/scheduler/scheduler.go +++ b/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{} diff --git a/internal/tests/implicit_cpu_memory_test.go b/internal/tests/implicit_cpu_memory_test.go new file mode 100644 index 0000000..ab6c0c5 --- /dev/null +++ b/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) +} diff --git a/internal/worker/option.go b/internal/worker/option.go index d16cb5b..7e87056 100644 --- a/internal/worker/option.go +++ b/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() diff --git a/internal/worker/worker.go b/internal/worker/worker.go index 2e0b15a..54db6e2 100644 --- a/internal/worker/worker.go +++ b/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 diff --git a/pkg/resource/v1/v1.go b/pkg/resource/v1/v1.go index 5425411..bcdcede 100644 --- a/pkg/resource/v1/v1.go +++ b/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"` diff --git a/pkg/resource/v1/worker.go b/pkg/resource/v1/worker.go index 577b332..02d4bff 100644 --- a/pkg/resource/v1/worker.go +++ b/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 }