Ability to pre-pull VM images by targeting workers with labels

This commit is contained in:
Nikolay Edigaryev
2025-12-10 20:43:34 +01:00
parent 173a34b083
commit aa36ad4cb6
29 changed files with 1463 additions and 8 deletions
+31
View File
@@ -183,6 +183,37 @@ func (controller *Controller) initAPI() *gin.Engine {
controller.appendVMEvents(c).Respond(c)
})
// Image pulls
v1.POST("/imagepulls", func(c *gin.Context) {
controller.createImagePull(c).Respond(c)
})
v1.PUT("/imagepulls/:name/state", func(c *gin.Context) {
controller.updateImagePullState(c).Respond(c)
})
v1.GET("/imagepulls/:name", func(c *gin.Context) {
controller.getImagePull(c).Respond(c)
})
v1.GET("/imagepulls", func(c *gin.Context) {
controller.listImagePulls(c).Respond(c)
})
v1.DELETE("/imagepulls/:name", func(c *gin.Context) {
controller.deleteImagePull(c).Respond(c)
})
// Image pull jobs
v1.POST("/imagepulljobs", func(c *gin.Context) {
controller.createImagePullJob(c).Respond(c)
})
v1.GET("/imagepulljobs/:name", func(c *gin.Context) {
controller.getImagePullJob(c).Respond(c)
})
v1.GET("/imagepulljobs", func(c *gin.Context) {
controller.listImagePullJobs(c).Respond(c)
})
v1.DELETE("/imagepulljobs/:name", func(c *gin.Context) {
controller.deleteImagePullJob(c).Respond(c)
})
return ginEngine
}
+1
View File
@@ -19,6 +19,7 @@ func (controller *Controller) controllerInfo(ctx *gin.Context) responder.Respond
capabilities := []v1pkg.ControllerCapability{
v1pkg.ControllerCapabilityRPCV1,
v1pkg.ControllerCapabilityVMStateEndpoint,
v1pkg.ControllerCapabilityImagePullResource,
}
if controller.experimentalRPCV2 {
+124
View File
@@ -0,0 +1,124 @@
package controller
import (
"errors"
"net/http"
"time"
storepkg "github.com/cirruslabs/orchard/internal/controller/store"
"github.com/cirruslabs/orchard/internal/responder"
"github.com/cirruslabs/orchard/internal/simplename"
v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
)
func (controller *Controller) createImagePullJob(ctx *gin.Context) responder.Responder {
// Auth
if responder := controller.authorize(ctx, v1.ServiceAccountRoleComputeWrite); responder != nil {
return responder
}
// Parse user input
var userPullJob v1.ImagePullJob
if err := ctx.ShouldBindJSON(&userPullJob); err != nil {
return responder.JSON(http.StatusBadRequest, NewErrorResponse("invalid JSON was provided"))
}
// Validate user input
if userPullJob.Name == "" {
return responder.JSON(http.StatusPreconditionFailed, NewErrorResponse("name field cannot be empty"))
} else if err := simplename.Validate(userPullJob.Name); err != nil {
return responder.JSON(http.StatusPreconditionFailed,
NewErrorResponse("name field %v", err))
}
if userPullJob.Image == "" {
return responder.JSON(http.StatusPreconditionFailed, NewErrorResponse("image field cannot be empty"))
}
// Provide defaults
userPullJob.CreatedAt = time.Now()
userPullJob.UID = uuid.NewString()
response := controller.storeUpdate(func(txn storepkg.Transaction) responder.Responder {
// Does this resource already exists?
_, err := txn.GetImagePullJob(userPullJob.Name)
if err != nil && !errors.Is(err, storepkg.ErrNotFound) {
controller.logger.Errorf("failed to check if the image pull job exists in the DB: %v", err)
return responder.Code(http.StatusInternalServerError)
}
if err == nil {
return responder.JSON(http.StatusConflict, NewErrorResponse("image pull job with this name "+
"already exists"))
}
if err := txn.SetImagePullJob(userPullJob); err != nil {
controller.logger.Errorf("failed to create image pull job in the DB: %v", err)
return responder.Code(http.StatusInternalServerError)
}
return responder.JSON(http.StatusOK, &userPullJob)
})
return response
}
func (controller *Controller) getImagePullJob(ctx *gin.Context) responder.Responder {
// Auth
if responder := controller.authorize(ctx, v1.ServiceAccountRoleComputeRead); responder != nil {
return responder
}
name := ctx.Param("name")
return controller.storeView(func(txn storepkg.Transaction) responder.Responder {
dbPullJob, err := txn.GetImagePullJob(name)
if err != nil {
return responder.Error(err)
}
return responder.JSON(http.StatusOK, dbPullJob)
})
}
func (controller *Controller) listImagePullJobs(ctx *gin.Context) responder.Responder {
// Auth
if responder := controller.authorize(ctx, v1.ServiceAccountRoleComputeRead); responder != nil {
return responder
}
return controller.storeView(func(txn storepkg.Transaction) responder.Responder {
dbPullJobs, err := txn.ListImagePullJobs()
if err != nil {
return responder.Error(err)
}
return responder.JSON(http.StatusOK, dbPullJobs)
})
}
func (controller *Controller) deleteImagePullJob(ctx *gin.Context) responder.Responder {
// Auth
if responder := controller.authorize(ctx, v1.ServiceAccountRoleComputeWrite); responder != nil {
return responder
}
name := ctx.Param("name")
return controller.storeUpdate(func(txn storepkg.Transaction) responder.Responder {
_, err := txn.GetImagePullJob(name)
if err != nil {
return responder.Error(err)
}
err = txn.DeleteImagePullJob(name)
if err != nil {
return responder.Error(err)
}
return responder.Code(http.StatusOK)
})
}
+159
View File
@@ -0,0 +1,159 @@
package controller
import (
"errors"
"net/http"
"time"
storepkg "github.com/cirruslabs/orchard/internal/controller/store"
"github.com/cirruslabs/orchard/internal/responder"
"github.com/cirruslabs/orchard/internal/simplename"
v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
)
func (controller *Controller) createImagePull(ctx *gin.Context) responder.Responder {
// Auth
if responder := controller.authorize(ctx, v1.ServiceAccountRoleComputeWrite); responder != nil {
return responder
}
// Parse user input
var userPull v1.ImagePull
if err := ctx.ShouldBindJSON(&userPull); err != nil {
return responder.JSON(http.StatusBadRequest, NewErrorResponse("invalid JSON was provided"))
}
// Validate user input
if userPull.Name == "" {
return responder.JSON(http.StatusPreconditionFailed, NewErrorResponse("name field cannot is empty"))
} else if err := simplename.Validate(userPull.Name); err != nil {
return responder.JSON(http.StatusPreconditionFailed,
NewErrorResponse("name field %v", err))
}
if userPull.Image == "" {
return responder.JSON(http.StatusPreconditionFailed, NewErrorResponse("image field cannot be empty"))
}
if userPull.Worker == "" {
return responder.JSON(http.StatusPreconditionFailed, NewErrorResponse("worker field cannot be empty"))
}
// Provide defaults
userPull.CreatedAt = time.Now()
userPull.UID = uuid.NewString()
response := controller.storeUpdate(func(txn storepkg.Transaction) responder.Responder {
// Does this resource already exists?
_, err := txn.GetImagePull(userPull.Name)
if err != nil && !errors.Is(err, storepkg.ErrNotFound) {
controller.logger.Errorf("failed to check if the image pull exists in the DB: %v", err)
return responder.Code(http.StatusInternalServerError)
}
if err == nil {
return responder.JSON(http.StatusConflict, NewErrorResponse("image pull with this name already exists"))
}
if err := txn.SetImagePull(userPull); err != nil {
controller.logger.Errorf("failed to create image pull in the DB: %v", err)
return responder.Code(http.StatusInternalServerError)
}
return responder.JSON(http.StatusOK, &userPull)
})
return response
}
func (controller *Controller) updateImagePullState(ctx *gin.Context) responder.Responder {
// Auth
if responder := controller.authorize(ctx, v1.ServiceAccountRoleComputeWrite); responder != nil {
return responder
}
// Parse user input
var userPull v1.ImagePull
if err := ctx.ShouldBindJSON(&userPull); err != nil {
return responder.JSON(http.StatusBadRequest, NewErrorResponse("invalid JSON was provided"))
}
name := ctx.Param("name")
return controller.storeUpdate(func(txn storepkg.Transaction) responder.Responder {
dbPull, err := txn.GetImagePull(name)
if err != nil {
return responder.Error(err)
}
dbPull.PullState = userPull.PullState
if err := txn.SetImagePull(*dbPull); err != nil {
controller.logger.Errorf("failed to update image pull in the DB: %v", err)
return responder.Code(http.StatusInternalServerError)
}
return responder.JSON(http.StatusOK, dbPull)
})
}
func (controller *Controller) getImagePull(ctx *gin.Context) responder.Responder {
// Auth
if responder := controller.authorize(ctx, v1.ServiceAccountRoleComputeRead); responder != nil {
return responder
}
name := ctx.Param("name")
return controller.storeView(func(txn storepkg.Transaction) responder.Responder {
dbPull, err := txn.GetImagePull(name)
if err != nil {
return responder.Error(err)
}
return responder.JSON(http.StatusOK, dbPull)
})
}
func (controller *Controller) listImagePulls(ctx *gin.Context) responder.Responder {
// Auth
if responder := controller.authorize(ctx, v1.ServiceAccountRoleComputeRead); responder != nil {
return responder
}
return controller.storeView(func(txn storepkg.Transaction) responder.Responder {
dbPulls, err := txn.ListImagePulls()
if err != nil {
return responder.Error(err)
}
return responder.JSON(http.StatusOK, dbPulls)
})
}
func (controller *Controller) deleteImagePull(ctx *gin.Context) responder.Responder {
// Auth
if responder := controller.authorize(ctx, v1.ServiceAccountRoleComputeWrite); responder != nil {
return responder
}
name := ctx.Param("name")
return controller.storeUpdate(func(txn storepkg.Transaction) responder.Responder {
_, err := txn.GetImagePull(name)
if err != nil {
return responder.Error(err)
}
err = txn.DeleteImagePull(name)
if err != nil {
return responder.Error(err)
}
return responder.Code(http.StatusOK)
})
}
+259
View File
@@ -0,0 +1,259 @@
package scheduler
import (
"errors"
"fmt"
"time"
storepkg "github.com/cirruslabs/orchard/internal/controller/store"
v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
"github.com/google/go-cmp/cmp"
"github.com/google/uuid"
"github.com/samber/lo"
)
func (scheduler *Scheduler) imagePullLoopIteration() error {
// Get a lagging view of image pulls, image pull jobs and workers
var imagePulls []v1.ImagePull
var imagePullJobs []v1.ImagePullJob
var workers []v1.Worker
if err := scheduler.store.View(func(txn storepkg.Transaction) error {
var err error
imagePulls, err = txn.ListImagePulls()
if err != nil {
return err
}
imagePullJobs, err = txn.ListImagePullJobs()
if err != nil {
return err
}
workers, err = txn.ListWorkers()
if err != nil {
return err
}
return nil
}); err != nil {
return err
}
scheduler.logger.Debugf("processing %d image pull jobs, %d image pulls and %d workers",
len(imagePullJobs), len(imagePulls), len(workers))
// Schedule new image pulls and update image pull job states
imagePullJobIndex := map[string]v1.ImagePullJob{}
for _, imagePullJob := range imagePullJobs {
imagePullJobIndex[imagePullJob.Name] = imagePullJob
// Schedule new image pulls
existingImagePulls := lo.Filter(imagePulls, func(imagePull v1.ImagePull, _ int) bool {
return lo.ContainsBy(imagePull.OwnerReferences, func(ownerReference v1.OwnerReference) bool {
return ownerReference == imagePullJob.OwnerReference()
})
})
for _, worker := range workers {
// Should we create an image pull for this worker?
if !worker.Labels.Contains(imagePullJob.Labels) {
continue
}
// Have we already created an image pull for this worker?
if _, ok := lo.Find(existingImagePulls, func(imagePull v1.ImagePull) bool {
return imagePull.Worker == worker.Name
}); ok {
continue
}
// Create an image pull for this worker
scheduler.logger.Debugf("creating image pull for job %s and worker %s",
imagePullJob.Name, worker.Name)
newImagePull, err := scheduler.createImagePull(imagePullJob, worker)
if err != nil {
return err
}
existingImagePulls = append(existingImagePulls, *newImagePull)
}
// Craft the current image pull job state
newImagePullJobState := v1.ImagePullJobState{
Total: int64(len(existingImagePulls)),
Progressing: int64(lo.CountBy(existingImagePulls, func(imagePull v1.ImagePull) bool {
return v1.ConditionIsTrue(imagePull.Conditions, v1.ConditionTypeProgressing)
})),
Succeeded: int64(lo.CountBy(existingImagePulls, func(imagePull v1.ImagePull) bool {
return v1.ConditionIsTrue(imagePull.Conditions, v1.ConditionTypeCompleted)
})),
Failed: int64(lo.CountBy(existingImagePulls, func(imagePull v1.ImagePull) bool {
return v1.ConditionIsTrue(imagePull.Conditions, v1.ConditionTypeFailed)
})),
}
if newImagePullJobState.Progressing != 0 {
newImagePullJobState.Conditions = append(newImagePullJobState.Conditions, v1.Condition{
Type: v1.ConditionTypeProgressing,
State: v1.ConditionStateTrue,
})
}
if newImagePullJobState.Failed != 0 {
newImagePullJobState.Conditions = append(newImagePullJobState.Conditions, v1.Condition{
Type: v1.ConditionTypeFailed,
State: v1.ConditionStateTrue,
})
} else if (newImagePullJobState.Progressing + newImagePullJobState.Succeeded + newImagePullJobState.Failed) == newImagePullJobState.Total {
newImagePullJobState.Conditions = append(newImagePullJobState.Conditions, v1.Condition{
Type: v1.ConditionTypeCompleted,
State: v1.ConditionStateTrue,
})
}
// Is the current image pull job state any different?
if cmp.Equal(imagePullJob.ImagePullJobState, newImagePullJobState) {
continue
}
// Update image pull job state
if err := scheduler.store.Update(func(txn storepkg.Transaction) error {
dbPullJob, err := txn.GetImagePullJob(imagePullJob.Name)
if err != nil {
// Is the image pull job still exists?
if errors.Is(err, storepkg.ErrNotFound) {
return nil
}
return err
}
// Is it the same image pull job?
if dbPullJob.UID != imagePullJob.UID {
return nil
}
dbPullJob.ImagePullJobState = newImagePullJobState
return txn.SetImagePullJob(*dbPullJob)
}); err != nil {
return err
}
}
// Garbage collect orphaned image pulls
for _, imagePull := range imagePulls {
if lo.ContainsBy(imagePull.OwnerReferences, func(ownerReference v1.OwnerReference) bool {
imagePullJob, ok := imagePullJobIndex[ownerReference.Name]
if !ok {
return false
}
return ownerReference == imagePullJob.OwnerReference()
}) {
// This image pull is still controlled by an existing image pull job
continue
}
scheduler.logger.Debugf("removing image pull %s with non-existent owner reference", imagePull.Name)
if err := scheduler.store.Update(func(txn storepkg.Transaction) error {
// Is this image pull still exists?
dbImagePull, err := txn.GetImagePull(imagePull.Name)
if err != nil {
if errors.Is(err, storepkg.ErrNotFound) {
scheduler.logger.Warnf("image pull %s is gone, perhaps the user "+
"manually deleted it?", imagePull.Name)
return nil
}
return err
}
// Is this the same image pull?
if imagePull.UID != dbImagePull.UID {
return nil
}
return txn.DeleteImagePull(imagePull.Name)
}); err != nil {
return err
}
}
return nil
}
func (scheduler *Scheduler) createImagePull(imagePullJob v1.ImagePullJob, worker v1.Worker) (*v1.ImagePull, error) {
var imagePull v1.ImagePull
if err := scheduler.store.Update(func(txn storepkg.Transaction) error {
dbPullJob, err := txn.GetImagePullJob(imagePullJob.Name)
if err != nil {
// Is the image pull job still exists?
if errors.Is(err, storepkg.ErrNotFound) {
return nil
}
return err
}
// Is it the same pull job?
if dbPullJob.UID != imagePullJob.UID {
return nil
}
dbWorker, err := txn.GetWorker(worker.Name)
if err != nil {
// Is the worker still exists?
if errors.Is(err, storepkg.ErrNotFound) {
return nil
}
return err
}
// Do the worker labels still match?
if !dbWorker.Labels.Contains(dbPullJob.Labels) {
return nil
}
imagePull = v1.ImagePull{
Meta: v1.Meta{
Name: fmt.Sprintf("%s-%s", dbPullJob.Name, dbWorker.Name),
CreatedAt: time.Now(),
},
UID: uuid.NewString(),
OwnerReferences: []v1.OwnerReference{
dbPullJob.OwnerReference(),
},
Image: dbPullJob.Image,
Worker: worker.Name,
}
_, err = txn.GetImagePull(imagePull.Name)
if err != nil && !errors.Is(err, storepkg.ErrNotFound) {
return err
}
if !errors.Is(err, storepkg.ErrNotFound) {
scheduler.logger.Warnf("image pull %s already exists, perhaps the user "+
"manually created it?", imagePull.Name)
return nil
}
if err := txn.SetImagePull(imagePull); err != nil {
return err
}
return nil
}); err != nil {
return nil, err
}
return &imagePull, nil
}
+10 -2
View File
@@ -99,16 +99,24 @@ func (scheduler *Scheduler) Run() {
scheduler.logger.Errorf("Failed to health-check VMs: %v", err)
}
imagePullLoopIterationStart := time.Now()
err = scheduler.imagePullLoopIteration()
imagePullLoopIterationEnd := time.Now()
if err != nil {
scheduler.logger.Errorf("Failed to process image pulls and image pull jobs: %v", err)
}
schedulingLoopIterationStart := time.Now()
numWorkersScheduling, numVMsScheduling, err := scheduler.schedulingLoopIteration()
schedulingLoopIterationEnd := time.Now()
scheduler.logger.Debugf("Health checking loop iteration for %d workers and %d VMs took %v, "+
"scheduling loop iteration for %d workers and %d VMs took %v",
"scheduling loop iteration for %d workers and %d VMs took %v, image pull loop iteration took %v",
numWorkersHealth, numVMsHealth,
healthCheckingLoopIterationEnd.Sub(healthCheckingLoopIterationStart),
numWorkersScheduling, numVMsScheduling,
schedulingLoopIterationEnd.Sub(schedulingLoopIterationStart))
schedulingLoopIterationEnd.Sub(schedulingLoopIterationStart),
imagePullLoopIterationEnd.Sub(imagePullLoopIterationStart))
if err != nil {
scheduler.logger.Errorf("Failed to schedule VMs: %v", err)
@@ -0,0 +1,30 @@
//nolint:dupl // maybe we'll figure out how to make DB resource accessors generic in the future
package badger
import (
"path"
"github.com/cirruslabs/orchard/pkg/resource/v1"
)
const SpaceImagePulls = "/imagepulls"
func ImagePullKey(name string) []byte {
return []byte(path.Join(SpaceImagePulls, name))
}
func (txn *Transaction) GetImagePull(name string) (*v1.ImagePull, error) {
return genericGet[v1.ImagePull](txn, ImagePullKey(name))
}
func (txn *Transaction) SetImagePull(pull v1.ImagePull) error {
return genericSet[v1.ImagePull](txn, ImagePullKey(pull.Name), pull)
}
func (txn *Transaction) DeleteImagePull(name string) error {
return genericDelete(txn, ImagePullKey(name))
}
func (txn *Transaction) ListImagePulls() ([]v1.ImagePull, error) {
return genericList[v1.ImagePull](txn, []byte(SpaceImagePulls))
}
@@ -0,0 +1,30 @@
//nolint:dupl // maybe we'll figure out how to make DB resource accessors generic in the future
package badger
import (
"path"
"github.com/cirruslabs/orchard/pkg/resource/v1"
)
const SpaceImagePullJobs = "/imagepulljobs"
func ImagePullJobKey(name string) []byte {
return []byte(path.Join(SpaceImagePullJobs, name))
}
func (txn *Transaction) GetImagePullJob(name string) (*v1.ImagePullJob, error) {
return genericGet[v1.ImagePullJob](txn, ImagePullJobKey(name))
}
func (txn *Transaction) SetImagePullJob(pull v1.ImagePullJob) error {
return genericSet[v1.ImagePullJob](txn, ImagePullJobKey(pull.Name), pull)
}
func (txn *Transaction) DeleteImagePullJob(name string) error {
return genericDelete(txn, ImagePullJobKey(name))
}
func (txn *Transaction) ListImagePullJobs() ([]v1.ImagePullJob, error) {
return genericList[v1.ImagePullJob](txn, []byte(SpaceImagePullJobs))
}
+10
View File
@@ -47,4 +47,14 @@ type Transaction interface {
GetClusterSettings() (*v1.ClusterSettings, error)
SetClusterSettings(clusterSettings v1.ClusterSettings) error
GetImagePull(name string) (result *v1.ImagePull, err error)
SetImagePull(pull v1.ImagePull) (err error)
DeleteImagePull(name string) (err error)
ListImagePulls() (result []v1.ImagePull, err error)
GetImagePullJob(name string) (result *v1.ImagePullJob, err error)
SetImagePullJob(pull v1.ImagePullJob) (err error)
DeleteImagePullJob(name string) (err error)
ListImagePullJobs() (result []v1.ImagePullJob, err error)
}