Move overlay to /dev/shm for image builds (#116)

* move overlay to /dev/shm for builds * add in configurable build container settings * add function * wip * wip --------- Co-authored-by: Luke Lombardi <luke@beam.cloud>
beam-cloud · Mar 27, 2024 · 1b646e6 · 1b646e6
1 parent 677c5fc
commit 1b646e6
Show file tree

Hide file tree

Showing 7 changed files with 53 additions and 26 deletions.
diff --git a/internal/abstractions/image/build.go b/internal/abstractions/image/build.go
@@ -126,11 +126,23 @@ func (b *Builder) Build(ctx context.Context, opts *BuildOpts, outputChan chan co
 	sourceImage := fmt.Sprintf("%s/%s:%s", opts.BaseImageRegistry, opts.BaseImageName, opts.BaseImageTag)
 	containerId := b.genContainerId()
 
+	// Allow config to override default build container settings
+	cpu := defaultBuildContainerCpu
+	memory := defaultBuildContainerMemory
+
+	if b.config.ImageService.BuildContainerCpu > 0 {
+		cpu = b.config.ImageService.BuildContainerCpu
+	}
+
+	if b.config.ImageService.BuildContainerMemory > 0 {
+		memory = b.config.ImageService.BuildContainerMemory
+	}
+
 	err = b.scheduler.Run(&types.ContainerRequest{
 		ContainerId: containerId,
 		Env:         []string{},
-		Cpu:         defaultBuildContainerCpu,
-		Memory:      defaultBuildContainerMemory,
+		Cpu:         cpu,
+		Memory:      memory,
 		ImageId:     baseImageId,
 		SourceImage: &sourceImage,
 		WorkspaceId: authInfo.Workspace.ExternalId,

diff --git a/internal/abstractions/image/image.go b/internal/abstractions/image/image.go
@@ -21,9 +21,8 @@ type ImageService interface {
 
 type RuncImageService struct {
 	pb.UnimplementedImageServiceServer
-	builder   *Builder
-	scheduler *scheduler.Scheduler
-	config    types.AppConfig
+	builder *Builder
+	config  types.AppConfig
 }
 
 type ImageServiceOpts struct {

diff --git a/internal/common/overlay.go b/internal/common/overlay.go
@@ -11,7 +11,6 @@ import (
 
 type ContainerOverlay struct {
 	containerId string
-	bundlePath  string
 	layers      []ContainerOverlayLayer
 	root        string
 	overlayPath string
@@ -25,14 +24,9 @@ type ContainerOverlayLayer struct {
 	merged string
 }
 
-func NewContainerOverlay(containerId string, bundlePath string, overlayPath string, rootPath string) *ContainerOverlay {
-	if rootPath == "" {
-		rootPath = bundlePath
-	}
-
+func NewContainerOverlay(containerId string, rootPath string, overlayPath string) *ContainerOverlay {
 	return &ContainerOverlay{
 		containerId: containerId,
-		bundlePath:  bundlePath,
 		layers:      []ContainerOverlayLayer{},
 		root:        rootPath,
 		overlayPath: overlayPath,

diff --git a/internal/types/config.go b/internal/types/config.go
@@ -76,6 +76,8 @@ type ImageServiceConfig struct {
 	RegistryCredentialProviderName string                `key:"registryCredentialProvider" json:"registry_credential_provider_name"`
 	Registries                     ImageRegistriesConfig `key:"registries" json:"registries"`
 	EnableTLS                      bool                  `key:"enableTLS" json:"enable_tls"`
+	BuildContainerCpu              int64                 `key:"buildContainerCpu" json:"build_container_cpu"`
+	BuildContainerMemory           int64                 `key:"buildContainerMemory" json:"build_container_memory"`
 	Runner                         RunnerConfig          `key:"runner" json:"runner"`
 }
 

diff --git a/internal/worker/image.go b/internal/worker/image.go
@@ -29,13 +29,11 @@ import (
 )
 
 const (
-	imagePullCommand       string = "skopeo"
-	imageCachePath         string = "/dev/shm/images"
-	imageAvailableFilename string = "IMAGE_AVAILABLE"
-	imageMountLockFilename string = "IMAGE_MOUNT_LOCK"
+	imagePullCommand string = "skopeo"
+	imageCachePath   string = "/dev/shm/images"
 )
 
-var requiredContainerDirectories []string = []string{"/workspace", "/volumes", "/snapshot"}
+var requiredContainerDirectories []string = []string{"/workspace", "/volumes"}
 
 type ImageClient struct {
 	registry       *common.ImageRegistry

diff --git a/internal/worker/runc_server.go b/internal/worker/runc_server.go
@@ -71,11 +71,11 @@ func (s *RunCServer) Start() error {
 }
 
 func (s *RunCServer) RunCKill(ctx context.Context, in *pb.RunCKillRequest) (*pb.RunCKillResponse, error) {
-	err := s.runcHandle.Kill(ctx, in.ContainerId, int(syscall.SIGTERM), &runc.KillOpts{
+	_ = s.runcHandle.Kill(ctx, in.ContainerId, int(syscall.SIGTERM), &runc.KillOpts{
 		All: true,
 	})
 
-	err = s.runcHandle.Delete(ctx, in.ContainerId, &runc.DeleteOpts{
+	err := s.runcHandle.Delete(ctx, in.ContainerId, &runc.DeleteOpts{
 		Force: true,
 	})
 

diff --git a/internal/worker/worker.go b/internal/worker/worker.go
@@ -252,6 +252,7 @@ func (s *Worker) shouldShutDown(lastContainerRequest time.Time) bool {
 // Spawn a single container and stream output to stdout/stderr
 func (s *Worker) RunContainer(request *types.ContainerRequest) error {
 	containerID := request.ContainerId
+
 	bundlePath := filepath.Join(s.userImagePath, request.ImageId)
 
 	hostname := fmt.Sprintf("%s:%d", s.podAddr, defaultWorkerServerPort)
@@ -440,6 +441,31 @@ func (s *Worker) clearContainer(containerId string, request *types.ContainerRequ
 	}()
 }
 
+// isBuildRequest checks if the sourceImage field is not-nil, which means the container request is for a build container
+func (s *Worker) isBuildRequest(request *types.ContainerRequest) bool {
+	if request.SourceImage != nil {
+		return true
+	}
+
+	return false
+}
+
+func (s *Worker) createOverlay(request *types.ContainerRequest, bundlePath string) *common.ContainerOverlay {
+	// For images that have a rootfs, set that as the root path
+	// otherwise, assume runc config files are in the rootfs themselves
+	rootPath := filepath.Join(bundlePath, "rootfs")
+	if _, err := os.Stat(rootPath); os.IsNotExist(err) {
+		rootPath = bundlePath
+	}
+
+	overlayPath := baseConfigPath
+	if s.isBuildRequest(request) {
+		overlayPath = "/dev/shm"
+	}
+
+	return common.NewContainerOverlay(request.ContainerId, rootPath, overlayPath)
+}
+
 // spawn a container using runc binary
 func (s *Worker) spawn(request *types.ContainerRequest, bundlePath string, spec *specs.Spec, outputChan chan common.OutputMsg) {
 	s.workerRepo.AddContainerRequestToWorker(s.workerId, request.ContainerId, request)
@@ -455,19 +481,15 @@ func (s *Worker) spawn(request *types.ContainerRequest, bundlePath string, spec
 		s.terminateContainer(containerId, request, &exitCode, &containerErr)
 	}()
 
-	// For images that have a rootfs, set that as the root path
-	// otherwise, assume runc config files are in the rootfs themselves
-	rootPath := filepath.Join(bundlePath, "rootfs")
-	if _, err := os.Stat(rootPath); os.IsNotExist(err) {
-		rootPath = bundlePath
-	}
+	// Create overlayfs for container
+	overlay := s.createOverlay(request, bundlePath)
 
 	// Add the container instance to the runningContainers map
 	containerInstance := &ContainerInstance{
 		Id:         containerId,
 		StubId:     request.StubId,
 		BundlePath: bundlePath,
-		Overlay:    common.NewContainerOverlay(containerId, bundlePath, baseConfigPath, rootPath),
+		Overlay:    overlay,
 		Spec:       spec,
 		ExitCode:   -1,
 		OutputWriter: common.NewOutputWriter(func(s string) {