Blobfs integration into worker (#352)

- Add blobfs cache path to worker
- Integrate into images client

go.mod

@@ -15,7 +15,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/ec2 v1.144.0
 	github.com/aws/aws-sdk-go-v2/service/ecr v1.24.4
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.48.1
-	github.com/beam-cloud/blobcache-v2 v0.0.0-20240620173035-e5940dda6ffa
+	github.com/beam-cloud/blobcache-v2 v0.0.0-20240707233531-8f1d070bac0b
 	github.com/beam-cloud/clip v0.0.0-20240606013643-20d5cfc38ac5
 	github.com/beam-cloud/go-runc v0.0.0-20231222221338-b89899f33170
 	github.com/bsm/redislock v0.9.4

go.sum

@@ -94,6 +94,8 @@ github.com/aws/smithy-go v1.20.0/go.mod h1:uo5RKksAl4PzhqaAbjd4rLgFoq5koTsQKYuGe
 github.com/aybabtme/rgbterm v0.0.0-20170906152045-cc83f3b3ce59/go.mod h1:q/89r3U2H7sSsE2t6Kca0lfwTK8JdoNGS/yzM/4iH5I=
 github.com/beam-cloud/blobcache-v2 v0.0.0-20240620173035-e5940dda6ffa h1:9ZlxYBetjhYyhXBNjd5p05Lla7nLNfrnMh0rb6ynyTc=
 github.com/beam-cloud/blobcache-v2 v0.0.0-20240620173035-e5940dda6ffa/go.mod h1:+2CZ7oUTp7W42Sc8Z3SBsqW6DWX4/CTwAM/nvlttcVw=
+github.com/beam-cloud/blobcache-v2 v0.0.0-20240707233531-8f1d070bac0b h1:kDS7D3hfi6XRPggsqLt8qZHqGr57HVe93HlnY5pOVkk=
+github.com/beam-cloud/blobcache-v2 v0.0.0-20240707233531-8f1d070bac0b/go.mod h1:rZaHc09LperUvxb/tuXlgmtH5P4gRs+uvj2DdXEjIf4=
 github.com/beam-cloud/clip v0.0.0-20240606013643-20d5cfc38ac5 h1:iHTcl1bR/osK/21xvF7stY47MlI2Z1yrtqaXAI+P3Tw=
 github.com/beam-cloud/clip v0.0.0-20240606013643-20d5cfc38ac5/go.mod h1:FO7taXHUAqgx33PjeB6LbSLpKob3Ceyo9Po64nq5TR0=
 github.com/beam-cloud/go-runc v0.0.0-20231222221338-b89899f33170 h1:KYVz18kobBGU8URM9Srn++2tcL9e0PcwYyH0Z4GYicM=

pkg/providers/remote_config.go

@@ -55,7 +55,17 @@ func GetRemoteConfig(baseConfig types.AppConfig, tailscale *network.Tailscale) (
 		if err != nil {
 			return nil, err
 		}
+
 		remoteConfig.ImageService.BlobCache.Metadata.RedisAddr = fmt.Sprintf("%s:%d", blobcacheRedisHostname, 6379)
+
+		if remoteConfig.ImageService.BlobCache.BlobFs.Enabled {
+			for idx, sourceConfig := range remoteConfig.ImageService.BlobCache.BlobFs.Sources {
+				if sourceConfig.Mode == "juicefs" {
+					remoteConfig.ImageService.BlobCache.BlobFs.Sources[idx].JuiceFS.RedisURI = remoteConfig.Storage.JuiceFS.RedisURI
+				}
+			}
+		}
+
 	}
 
 	return &remoteConfig, nil

pkg/scheduler/pool.go

@@ -17,10 +17,12 @@ const (
 	tmpVolumeName               string = "beta9-tmp"
 	logVolumeName               string = "beta9-logs"
 	imagesVolumeName            string = "beta9-images"
+	cacheVolumeName             string = "beta9-cache"
 	defaultContainerName        string = "worker"
 	defaultWorkerEntrypoint     string = "/usr/local/bin/worker"
 	defaultWorkerLogPath        string = "/var/log/worker"
 	defaultImagesPath           string = "/images"
+	defaultCachePath            string = "/cache"
 )
 
 type WorkerPoolController interface {

pkg/scheduler/pool_external.go

@@ -459,6 +459,15 @@ func (wpc *ExternalWorkerPoolController) getWorkerVolumes(workerMemory int64) []
 				},
 			},
 		},
+		{
+			Name: cacheVolumeName,
+			VolumeSource: corev1.VolumeSource{
+				HostPath: &corev1.HostPathVolumeSource{
+					Path: defaultCachePath,
+					Type: &hostPathType,
+				},
+			},
+		},
 	}
 }
 
@@ -474,6 +483,11 @@ func (wpc *ExternalWorkerPoolController) getWorkerVolumeMounts() []corev1.Volume
 			MountPath: "/images",
 			ReadOnly:  false,
 		},
+		{
+			Name:      cacheVolumeName,
+			MountPath: "/cache",
+			ReadOnly:  true,
+		},
 		{
 			Name:      logVolumeName,
 			MountPath: defaultWorkerLogPath,

pkg/worker/image.go

@@ -39,6 +39,7 @@ const (
 var (
 	baseImageCachePath string = "/images/cache"
 	baseImageMountPath string = "/images/mnt/%s"
+	baseBlobFsPath     string = "/cache"
 )
 
 var requiredContainerDirectories []string = []string{"/workspace", "/volumes"}
@@ -129,6 +130,32 @@ func (c *ImageClient) PullLazy(request *types.ContainerRequest) error {
 		localCachePath = ""
 	}
 
+	if c.config.BlobCache.BlobFs.Enabled {
+		sourcePath := fmt.Sprintf("images/%s.clip", imageId)
+		sourceOffset := int64(0)
+
+		// If the image archive is already cached in blobcache, then we can use that as the local cache path
+		baseBlobFsContentPath := fmt.Sprintf("%s/%s", baseBlobFsPath, sourcePath)
+		if _, err := os.Stat(baseBlobFsContentPath); err == nil {
+			localCachePath = baseBlobFsContentPath
+		} else {
+
+			log.Printf("<%s> - blobfs cache entry not found for image<%s>, storing content nearby\n", request.ContainerId, imageId)
+
+			startTime := time.Now()
+			// Otherwise, lets cache it in a nearby blobcache host
+			_, err := c.cacheClient.StoreContentFromSource(sourcePath, sourceOffset)
+			if err == nil {
+				localCachePath = baseBlobFsContentPath
+			}
+
+			elapsed := time.Since(startTime)
+
+			log.Printf("<%s> - blobfs cache took %v\n", request.ContainerId, elapsed)
+
+		}
+	}
+
 	remoteArchivePath := fmt.Sprintf("%s/%s.%s", c.imageCachePath, imageId, c.registry.ImageFileExtension)
 	var err error = nil