From 5fe0a80c83993bd7d32f272da8475ccb8da455c8 Mon Sep 17 00:00:00 2001 From: Luke Lombardi <33990301+luke-lombardi@users.noreply.github.com> Date: Fri, 5 Jun 2026 10:55:42 -0400 Subject: [PATCH] ttl more keys --- benchmarks/b9bench/cache_suite.py | 13 +++- benchmarks/sandbox_startup_report.py | 2 +- pkg/abstractions/endpoint/buffer.go | 15 +++- pkg/abstractions/endpoint/buffer_test.go | 42 ++++++++++ pkg/abstractions/endpoint/instance.go | 4 +- pkg/abstractions/pod/proxy.go | 2 +- pkg/abstractions/pod/proxy_test.go | 39 ++++++++++ pkg/api/v1/stub.go | 15 ++++ pkg/api/v1/stub_test.go | 16 ++++ pkg/repository/worker_redis.go | 5 ++ pkg/repository/worker_redis_test.go | 7 ++ pkg/worker/criu.go | 99 +++++++++++++++++++++--- pkg/worker/lifecycle_test.go | 67 ++++++++++++++++ 13 files changed, 306 insertions(+), 20 deletions(-) create mode 100644 pkg/abstractions/endpoint/buffer_test.go create mode 100644 pkg/abstractions/pod/proxy_test.go diff --git a/benchmarks/b9bench/cache_suite.py b/benchmarks/b9bench/cache_suite.py index c0ef11110..7b25eb467 100644 --- a/benchmarks/b9bench/cache_suite.py +++ b/benchmarks/b9bench/cache_suite.py @@ -225,6 +225,8 @@ class BenchmarkConfig: direct_cache_disk_probe_size_mb: int sandbox_cpu: str sandbox_memory: str + sandbox_gpu: str + sandbox_gpu_count: int sandbox_keep_warm_seconds: int sandbox_create_retries: int sandbox_ready_timeout_seconds: float @@ -364,6 +366,12 @@ def parse(cls) -> "BenchmarkConfig": ) add("--sandbox-cpu", default=os.getenv("BENCH_CACHE_SANDBOX_CPU", "4")) add("--sandbox-memory", default=os.getenv("BENCH_CACHE_SANDBOX_MEMORY", "4096")) + add("--sandbox-gpu", default=os.getenv("BENCH_CACHE_SANDBOX_GPU", "")) + add( + "--sandbox-gpu-count", + type=int, + default=env_int("BENCH_CACHE_SANDBOX_GPU_COUNT", 0), + ) add( "--sandbox-keep-warm-seconds", type=int, @@ -2378,7 +2386,8 @@ def prepare_runtime(self): image=image, cpu=parse_sdk_cpu(self.config.sandbox_cpu), memory=parse_sdk_memory(self.config.sandbox_memory), - gpu="", + gpu=self.config.sandbox_gpu, + gpu_count=self.config.sandbox_gpu_count, keep_warm_seconds=self.config.sandbox_keep_warm_seconds, authorized=False, volumes=[volume], @@ -3876,6 +3885,8 @@ def _report_config(self, volume_id: str) -> dict[str, Any]: "readMethod": self.config.read_method, "sandboxCpu": self.config.sandbox_cpu, "sandboxMemory": self.config.sandbox_memory, + "sandboxGpu": self.config.sandbox_gpu, + "sandboxGpuCount": self.config.sandbox_gpu_count, "forceNewImage": self.config.force_new_image, "forceReadThroughAfterPrepare": self.config.force_read_through_after_prepare, } diff --git a/benchmarks/sandbox_startup_report.py b/benchmarks/sandbox_startup_report.py index 2797919f6..768c246bf 100644 --- a/benchmarks/sandbox_startup_report.py +++ b/benchmarks/sandbox_startup_report.py @@ -101,7 +101,7 @@ ("network_create_namespace_ms", "network.create_namespace", "Create namespace", False), ("network_configure_namespace_ms", "network.configure_namespace", "Configure namespace", False), ("network_ip_lock_ms", "network.ip_lock", "Acquire IP lock", False), - ("network_ip_scan_ms", "network.ip_scan", "Scan allocated IPs", False), + ("network_ip_load_ms", "network.ip_load", "Load allocated IPs", False), ("network_ip_assign_ms", "network.ip_assign", "Assign container IP", False), ("network_set_container_ip_ms", "network.set_container_ip", "Persist container IP", False), ("network_restrictions_ms", "network.restrictions", "Network restrictions", False), diff --git a/pkg/abstractions/endpoint/buffer.go b/pkg/abstractions/endpoint/buffer.go index 221a0e8fe..8c89f8ff0 100644 --- a/pkg/abstractions/endpoint/buffer.go +++ b/pkg/abstractions/endpoint/buffer.go @@ -301,7 +301,7 @@ func (rb *RequestBuffer) requestTokens(containerId string) (int, error) { if err != nil && err != redis.Nil { return 0, err } else if err == redis.Nil { - created, err := rb.rdb.SetNX(rb.ctx, tokenKey, rb.maxTokens, 0).Result() + created, err := rb.rdb.SetNX(rb.ctx, tokenKey, rb.maxTokens, rb.requestTokenTTL()).Result() if err != nil { return 0, err } @@ -321,6 +321,15 @@ func (rb *RequestBuffer) requestTokens(containerId string) (int, error) { return val, nil } +func (rb *RequestBuffer) requestTokenTTL() time.Duration { + ttl := time.Duration(rb.stubConfig.TaskPolicy.Timeout) * time.Second + if ttl <= 0 { + return time.Duration(DefaultEndpointRequestTimeoutS) * time.Second + } + + return ttl +} + func (rb *RequestBuffer) acquireRequestToken(containerId string) error { tokenKey := Keys.endpointRequestTokens(rb.workspace.Name, rb.stubId, containerId) tokenCount, err := rb.rdb.Decr(rb.ctx, tokenKey).Result() @@ -336,7 +345,7 @@ func (rb *RequestBuffer) acquireRequestToken(containerId string) error { return errors.New("too many in-flight requests") } - err = rb.rdb.Expire(rb.ctx, tokenKey, time.Duration(rb.stubConfig.TaskPolicy.Timeout)*time.Second).Err() + err = rb.rdb.Expire(rb.ctx, tokenKey, rb.requestTokenTTL()).Err() if err != nil { return err } @@ -352,7 +361,7 @@ func (rb *RequestBuffer) releaseRequestToken(containerId, taskId string) error { return err } - err = rb.rdb.Expire(rb.ctx, tokenKey, time.Duration(rb.stubConfig.TaskPolicy.Timeout)*time.Second).Err() + err = rb.rdb.Expire(rb.ctx, tokenKey, rb.requestTokenTTL()).Err() if err != nil { return err } diff --git a/pkg/abstractions/endpoint/buffer_test.go b/pkg/abstractions/endpoint/buffer_test.go new file mode 100644 index 000000000..9ca49d7c5 --- /dev/null +++ b/pkg/abstractions/endpoint/buffer_test.go @@ -0,0 +1,42 @@ +package endpoint + +import ( + "context" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/beam-cloud/beta9/pkg/common" + "github.com/beam-cloud/beta9/pkg/types" + "github.com/tj/assert" +) + +func TestRequestTokensSetsTTLWhenCreated(t *testing.T) { + server, err := miniredis.Run() + assert.Nil(t, err) + defer server.Close() + + rdb, err := common.NewRedisClient(types.RedisConfig{ + Addrs: []string{server.Addr()}, + Mode: types.RedisModeSingle, + }) + assert.Nil(t, err) + + rb := &RequestBuffer{ + ctx: context.Background(), + rdb: rdb, + workspace: &types.Workspace{Name: "workspace"}, + stubId: "stub", + stubConfig: &types.StubConfigV1{TaskPolicy: types.TaskPolicy{Timeout: 120}}, + maxTokens: 3, + } + + tokens, err := rb.requestTokens("container") + assert.Nil(t, err) + assert.Equal(t, 3, tokens) + + ttl, err := rdb.TTL(context.Background(), Keys.endpointRequestTokens("workspace", "stub", "container")).Result() + assert.Nil(t, err) + assert.True(t, ttl > 0) + assert.True(t, ttl <= 120*time.Second) +} diff --git a/pkg/abstractions/endpoint/instance.go b/pkg/abstractions/endpoint/instance.go index dbc79234e..ba2a6581d 100644 --- a/pkg/abstractions/endpoint/instance.go +++ b/pkg/abstractions/endpoint/instance.go @@ -33,7 +33,7 @@ type endpointInstance struct { } func (i *endpointInstance) startContainers(containersToRun int) error { - secrets, err := abstractions.ConfigureContainerRequestSecrets(i.Workspace, *i.buffer.stubConfig) + secrets, err := abstractions.ConfigureContainerRequestSecrets(i.Workspace, *i.StubConfig) if err != nil { return err } @@ -79,7 +79,7 @@ func (i *endpointInstance) startContainers(containersToRun int) error { containerId, i.Stub.Object.ExternalId, i.Workspace, - *i.buffer.stubConfig, + *i.StubConfig, i.Stub.ExternalId, ) if err != nil { diff --git a/pkg/abstractions/pod/proxy.go b/pkg/abstractions/pod/proxy.go index aebd46da4..6068c6500 100644 --- a/pkg/abstractions/pod/proxy.go +++ b/pkg/abstractions/pod/proxy.go @@ -503,7 +503,7 @@ func (pb *PodProxyBuffer) containerConnections(containerId string) (int, error) if err != nil && err != redis.Nil { return 0, err } else if err == redis.Nil { - created, err := pb.rdb.SetNX(pb.ctx, tokenKey, 0, 0).Result() + created, err := pb.rdb.SetNX(pb.ctx, tokenKey, 0, podContainerConnectionTimeout).Result() if err != nil { return 0, err } diff --git a/pkg/abstractions/pod/proxy_test.go b/pkg/abstractions/pod/proxy_test.go new file mode 100644 index 000000000..e93930ff5 --- /dev/null +++ b/pkg/abstractions/pod/proxy_test.go @@ -0,0 +1,39 @@ +package pod + +import ( + "context" + "testing" + + "github.com/alicebob/miniredis/v2" + "github.com/beam-cloud/beta9/pkg/common" + "github.com/beam-cloud/beta9/pkg/types" + "github.com/tj/assert" +) + +func TestContainerConnectionsSetsTTLWhenCreated(t *testing.T) { + server, err := miniredis.Run() + assert.Nil(t, err) + defer server.Close() + + rdb, err := common.NewRedisClient(types.RedisConfig{ + Addrs: []string{server.Addr()}, + Mode: types.RedisModeSingle, + }) + assert.Nil(t, err) + + pb := &PodProxyBuffer{ + ctx: context.Background(), + rdb: rdb, + workspace: &types.Workspace{Name: "workspace"}, + stubId: "stub", + } + + connections, err := pb.containerConnections("container") + assert.Nil(t, err) + assert.Equal(t, 0, connections) + + ttl, err := rdb.TTL(context.Background(), Keys.podContainerConnections("workspace", "stub", "container")).Result() + assert.Nil(t, err) + assert.True(t, ttl > 0) + assert.True(t, ttl <= podContainerConnectionTimeout) +} diff --git a/pkg/api/v1/stub.go b/pkg/api/v1/stub.go index 3f3e24d9d..f6b788ea5 100644 --- a/pkg/api/v1/stub.go +++ b/pkg/api/v1/stub.go @@ -615,6 +615,21 @@ func (g *StubGroup) UpdateConfig(ctx echo.Context) error { return HTTPBadRequest(errorMsg) } + if stubConfig.CheckpointEnabled { + if stub.Type.IsServe() { + return HTTPBadRequest("Checkpoints are not supported for serve stubs") + } + if stubConfig.Runtime.GpuCount > 1 { + return HTTPBadRequest("Checkpoints are yet not supported for multi-GPU") + } + if len(stubConfig.Runtime.Gpus) > 1 { + return HTTPBadRequest("Checkpoints are yet not supported between multiple GPUs") + } + if !stub.Workspace.StorageAvailable() { + return HTTPBadRequest("workspace storage is required for checkpoints") + } + } + if err := g.backendRepo.UpdateStubConfig(ctx.Request().Context(), stub.Id, &stubConfig); err != nil { return HTTPInternalServerError("Failed to update stub config") } diff --git a/pkg/api/v1/stub_test.go b/pkg/api/v1/stub_test.go index 0ea1e799e..b52bc844d 100644 --- a/pkg/api/v1/stub_test.go +++ b/pkg/api/v1/stub_test.go @@ -248,6 +248,22 @@ func TestUpdateConfig(t *testing.T) { expectedStatus: http.StatusOK, expectedFields: []string{"runtime.cpu", "runtime.gpu"}, }, + { + name: "Test checkpoint update requires workspace storage", + stubID: "test-stub-checkpoint", + workspaceID: 1, + requestBody: UpdateConfigRequest{ + Fields: map[string]interface{}{ + "checkpoint_enabled": true, + }, + }, + setupMock: func(mock *sqlmock.Sqlmock) { + rows := generateMockStubRows(1, "test-stub-checkpoint", "Test Stub", `{"runtime":{"cpu":1000,"memory":1000},"checkpoint_enabled":false}`, 1) + (*mock).ExpectQuery("SELECT").WithArgs("test-stub-checkpoint").WillReturnRows(rows) + }, + expectedStatus: http.StatusBadRequest, + expectedError: true, + }, { name: "Test stub not found", stubID: "non-existent-stub", diff --git a/pkg/repository/worker_redis.go b/pkg/repository/worker_redis.go index e84fcf071..8367ca21b 100644 --- a/pkg/repository/worker_redis.go +++ b/pkg/repository/worker_redis.go @@ -190,6 +190,11 @@ func (r *WorkerRedisRepository) RemoveWorker(workerId string) error { return err } + err = r.rdb.Del(context.TODO(), common.RedisKeys.SchedulerContainerWorkerIndex(workerId)).Err() + if err != nil { + return err + } + return nil } diff --git a/pkg/repository/worker_redis_test.go b/pkg/repository/worker_redis_test.go index 6ab8fcfea..fdf0be4b2 100644 --- a/pkg/repository/worker_redis_test.go +++ b/pkg/repository/worker_redis_test.go @@ -39,6 +39,9 @@ func TestAddAndRemoveWorker(t *testing.T) { err = repo.AddWorker(newWorker) assert.Nil(t, err) + err = rdb.SAdd(context.TODO(), common.RedisKeys.SchedulerContainerWorkerIndex(newWorker.Id), "scheduler:container:state:stale").Err() + assert.Nil(t, err) + worker, err := repo.GetWorkerById(newWorker.Id) assert.Nil(t, err) assert.Equal(t, newWorker.FreeCpu, worker.FreeCpu) @@ -49,6 +52,10 @@ func TestAddAndRemoveWorker(t *testing.T) { err = repo.RemoveWorker(newWorker.Id) assert.Nil(t, err) + exists, err := rdb.Exists(context.TODO(), common.RedisKeys.SchedulerContainerWorkerIndex(newWorker.Id)).Result() + assert.Nil(t, err) + assert.Equal(t, int64(0), exists) + err = repo.RemoveWorker(newWorker.Id) assert.Error(t, err) diff --git a/pkg/worker/criu.go b/pkg/worker/criu.go index 01f95767b..1702f06f6 100644 --- a/pkg/worker/criu.go +++ b/pkg/worker/criu.go @@ -28,6 +28,7 @@ import ( const ( gpuCntEnvKey = "GPU_COUNT" defaultCheckpointDeadline = 10 * time.Minute + checkpointStateRPCTimeout = 5 * time.Second readyLogRate = 10 checkpointFsDir = "filesystem" checkpointArchiveExtension = ".tar" @@ -173,7 +174,7 @@ type CreateCheckpointOpts struct { // Waits for the container to be ready to checkpoint at the desired point in execution, ie. // after all processes within a container have reached a checkpointable state -func (s *Worker) createCheckpoint(ctx context.Context, opts *CreateCheckpointOpts) error { +func (s *Worker) createCheckpoint(ctx context.Context, opts *CreateCheckpointOpts) (err error) { instance, exists := s.containerInstances.Get(opts.Request.ContainerId) if !exists { return fmt.Errorf("container instance not found") @@ -185,17 +186,53 @@ func (s *Worker) createCheckpoint(ctx context.Context, opts *CreateCheckpointOpt log.Info().Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msg("creating checkpoint") + readyForCheckpoint := false + checkpointComplete := false + writeCheckpointComplete := func() error { + if !opts.WaitForSignal || !readyForCheckpoint || checkpointComplete { + return nil + } + f, err := os.Create(filepath.Join(checkpointSignalDir(opts.Request.ContainerId), checkpointCompleteFileName)) + if err != nil { + log.Error().Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msgf("failed to create checkpoint complete file: %v", err) + return err + } + if err := f.Close(); err != nil { + log.Error().Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msgf("failed to close checkpoint complete file: %v", err) + return err + } + checkpointComplete = true + log.Info().Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msg("checkpoint complete signal created") + return nil + } + defer func() { + if err == nil || !opts.WaitForSignal || !readyForCheckpoint || checkpointComplete { + return + } + if signalErr := writeCheckpointComplete(); signalErr != nil { + log.Error().Err(signalErr).Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msg("failed to unblock checkpoint waiter") + } + }() + if opts.WaitForSignal { err := s.waitForCheckpointSignal(ctx, opts.Request, opts.OutputLogger) if err != nil { log.Error().Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msgf("failed to wait for checkpoint signal: %v", err) return err } + readyForCheckpoint = true } + checkpointCtx, cancel := context.WithTimeout(ctx, defaultCheckpointDeadline) + defer cancel() + // Proceed to create the checkpoint - checkpointPath, err := s.criuManager.CreateCheckpoint(ctx, instance.Runtime, opts.CheckpointId, opts.Request) + phaseStart := time.Now() + log.Info().Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msg("runtime checkpoint starting") + checkpointPath, err := s.criuManager.CreateCheckpoint(checkpointCtx, instance.Runtime, opts.CheckpointId, opts.Request) if err != nil { + log.Error().Err(err).Dur("duration", time.Since(phaseStart)).Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msg("runtime checkpoint failed") + _ = writeCheckpointComplete() if stateErr := s.createCheckpointState(opts.CheckpointId, opts.Request, types.CheckpointStatusCheckpointFailed, opts.ContainerIp); stateErr != nil { log.Error().Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msgf("failed to create checkpoint state: %v", stateErr) } @@ -206,28 +243,39 @@ func (s *Worker) createCheckpoint(ctx context.Context, opts *CreateCheckpointOpt return err } + log.Info().Dur("duration", time.Since(phaseStart)).Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Str("checkpoint_path", checkpointPath).Msg("runtime checkpoint completed") parentDir := filepath.Dir(instance.Overlay.TopLayerPath()) upperDir := path.Join(parentDir, "upper") + phaseStart = time.Now() + log.Info().Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Str("upper_dir", upperDir).Msg("checkpoint filesystem copy starting") err = copyDirectory(upperDir, path.Join(checkpointPath, checkpointFsDir), []string{"config.json", "outputs", "snapshot"}) if err != nil { log.Error().Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msgf("failed to copy upper directory: %v", err) + if opts.OutputLogger != nil { + opts.OutputLogger.Error("Failed to copy checkpoint filesystem") + } return err } + log.Info().Dur("duration", time.Since(phaseStart)).Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msg("checkpoint filesystem copy completed") - metadata, err := s.persistCheckpoint(ctx, opts.Request, opts.CheckpointId, checkpointPath) + phaseStart = time.Now() + log.Info().Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msg("checkpoint persistence starting") + metadata, err := s.persistCheckpoint(checkpointCtx, opts.Request, opts.CheckpointId, checkpointPath) if err != nil { + _ = writeCheckpointComplete() _ = s.createCheckpointState(opts.CheckpointId, opts.Request, types.CheckpointStatusCheckpointFailed, opts.ContainerIp) log.Error().Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msgf("failed to persist checkpoint: %v", err) + if opts.OutputLogger != nil { + opts.OutputLogger.Error("Failed to persist checkpoint") + } return err } + log.Info().Dur("duration", time.Since(phaseStart)).Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Str("hash", metadata.hash).Int64("size_bytes", metadata.sizeBytes).Msg("checkpoint persistence completed") if opts.WaitForSignal { - // Create a file accessible to the container to indicate that the checkpoint has been captured - _, err = os.Create(filepath.Join(checkpointSignalDir(opts.Request.ContainerId), checkpointCompleteFileName)) - if err != nil { - log.Error().Str("container_id", opts.Request.ContainerId).Str("checkpoint_id", opts.CheckpointId).Msgf("failed to create checkpoint complete file: %v", err) + if err := writeCheckpointComplete(); err != nil { return err } } @@ -287,13 +335,25 @@ func (s *Worker) persistCheckpoint(ctx context.Context, request *types.Container _ = os.Remove(archivePath) defer os.Remove(archivePath) + phaseStart := time.Now() + log.Info().Str("checkpoint_id", checkpointId).Str("archive_path", archivePath).Msg("checkpoint archive creation starting") if err := createTar(checkpointPath, archivePath); err != nil { return nil, err } + archiveInfo, _ := os.Stat(archivePath) + archiveSize := int64(0) + if archiveInfo != nil { + archiveSize = archiveInfo.Size() + } + log.Info().Dur("duration", time.Since(phaseStart)).Str("checkpoint_id", checkpointId).Int64("size_bytes", archiveSize).Msg("checkpoint archive creation completed") + + phaseStart = time.Now() + log.Info().Str("checkpoint_id", checkpointId).Str("archive_path", archivePath).Msg("checkpoint archive hash starting") hash, size, err := fileSHA256(archivePath) if err != nil { return nil, err } + log.Info().Dur("duration", time.Since(phaseStart)).Str("checkpoint_id", checkpointId).Str("hash", hash).Int64("size_bytes", size).Msg("checkpoint archive hash completed") originKey := checkpointOriginKey(checkpointId) storageClient, err := clients.NewWorkspaceStorageClient(ctx, request.Workspace.Name, request.Workspace.Storage) @@ -304,6 +364,8 @@ func (s *Worker) persistCheckpoint(ctx context.Context, request *types.Container if err != nil { return nil, err } + phaseStart = time.Now() + log.Info().Str("checkpoint_id", checkpointId).Str("origin_key", originKey).Int64("size_bytes", size).Msg("checkpoint workspace upload starting") if err := storageClient.UploadWithReader(ctx, originKey, f); err != nil { _ = f.Close() return nil, err @@ -311,12 +373,17 @@ func (s *Worker) persistCheckpoint(ctx context.Context, request *types.Container if err := f.Close(); err != nil { return nil, err } + log.Info().Dur("duration", time.Since(phaseStart)).Str("checkpoint_id", checkpointId).Str("origin_key", originKey).Int64("size_bytes", size).Msg("checkpoint workspace upload completed") + phaseStart = time.Now() + log.Info().Str("checkpoint_id", checkpointId).Str("hash", hash).Int64("size_bytes", size).Msg("checkpoint cache store starting") if _, err := s.cacheManager.client.StoreContentFromLocalFile(cache.LocalContentSource{ Path: archivePath, CachePath: originKey, }, cache.StoreContentOptions{RoutingKey: hash, Lock: true}); err != nil { - log.Warn().Err(err).Str("checkpoint_id", checkpointId).Str("hash", hash).Msg("failed to store checkpoint archive in cache") + log.Warn().Err(err).Dur("duration", time.Since(phaseStart)).Str("checkpoint_id", checkpointId).Str("hash", hash).Msg("failed to store checkpoint archive in cache") + } else { + log.Info().Dur("duration", time.Since(phaseStart)).Str("checkpoint_id", checkpointId).Str("hash", hash).Msg("checkpoint cache store completed") } return &checkpointCacheMetadata{ @@ -588,7 +655,9 @@ func (s *Worker) IsCRIUAvailable(gpuCount uint32) bool { func (s *Worker) createCheckpointState(checkpointId string, request *types.ContainerRequest, status types.CheckpointStatus, containerIp string) error { req := checkpointStateRequest(checkpointId, request, status, containerIp) - _, err := handleGRPCResponse(s.backendRepoClient.CreateCheckpoint(context.Background(), req)) + ctx, cancel := context.WithTimeout(context.Background(), checkpointStateRPCTimeout) + defer cancel() + _, err := handleGRPCResponse(s.backendRepoClient.CreateCheckpoint(ctx, req)) return err } @@ -600,7 +669,9 @@ func (s *Worker) createAvailableCheckpointState(checkpointId string, request *ty req.OriginKey = metadata.originKey req.Locality = metadata.locality req.Accelerator = metadata.accelerator - _, err := handleGRPCResponse(s.backendRepoClient.CreateCheckpoint(context.Background(), req)) + ctx, cancel := context.WithTimeout(context.Background(), checkpointStateRPCTimeout) + defer cancel() + _, err := handleGRPCResponse(s.backendRepoClient.CreateCheckpoint(ctx, req)) return err } @@ -618,7 +689,9 @@ func checkpointStateRequest(checkpointId string, request *types.ContainerRequest } func (s *Worker) updateCheckpointState(checkpointId string, request *types.ContainerRequest, status types.CheckpointStatus) error { - _, err := handleGRPCResponse(s.backendRepoClient.UpdateCheckpoint(context.Background(), &pb.UpdateCheckpointRequest{ + ctx, cancel := context.WithTimeout(context.Background(), checkpointStateRPCTimeout) + defer cancel() + _, err := handleGRPCResponse(s.backendRepoClient.UpdateCheckpoint(ctx, &pb.UpdateCheckpointRequest{ CheckpointId: checkpointId, Status: string(status), })) @@ -627,7 +700,9 @@ func (s *Worker) updateCheckpointState(checkpointId string, request *types.Conta } func (s *Worker) updateCheckpointRestored(checkpointId string) error { - _, err := handleGRPCResponse(s.backendRepoClient.UpdateCheckpoint(context.Background(), &pb.UpdateCheckpointRequest{ + ctx, cancel := context.WithTimeout(context.Background(), checkpointStateRPCTimeout) + defer cancel() + _, err := handleGRPCResponse(s.backendRepoClient.UpdateCheckpoint(ctx, &pb.UpdateCheckpointRequest{ CheckpointId: checkpointId, LastRestoredAt: timestamppb.Now(), })) diff --git a/pkg/worker/lifecycle_test.go b/pkg/worker/lifecycle_test.go index b172d2d02..c21f00565 100644 --- a/pkg/worker/lifecycle_test.go +++ b/pkg/worker/lifecycle_test.go @@ -550,6 +550,54 @@ func TestAttemptRestoreCheckpointTreatsGenericErrorAsRestoreFailure(t *testing.T require.Nil(t, backendRepoClient.lastUpdate.LastRestoredAt) } +func TestCreateCheckpointUnblocksRunnerOnPersistenceFailure(t *testing.T) { + tmpDir := t.TempDir() + containerID := "checkpoint-unblock" + signalDir := checkpointSignalDir(containerID) + require.NoError(t, os.RemoveAll(signalDir)) + require.NoError(t, os.MkdirAll(signalDir, 0755)) + t.Cleanup(func() { _ = os.RemoveAll(signalDir) }) + require.NoError(t, os.WriteFile(filepath.Join(signalDir, checkpointSignalFileName), nil, 0644)) + + layerDir := filepath.Join(tmpDir, "layer-0") + mergedDir := filepath.Join(layerDir, "merged") + upperDir := filepath.Join(layerDir, "upper") + require.NoError(t, os.MkdirAll(mergedDir, 0755)) + require.NoError(t, os.MkdirAll(upperDir, 0755)) + require.NoError(t, os.WriteFile(filepath.Join(upperDir, "ready.txt"), []byte("ok"), 0644)) + + request := &types.ContainerRequest{ + ContainerId: containerID, + Stub: types.StubWithRelated{Stub: types.Stub{ExternalId: "stub"}}, + Workspace: types.Workspace{Name: "workspace"}, + } + worker := &Worker{ + containerInstances: common.NewSafeMap[*ContainerInstance](), + criuManager: &checkpointPathCRIUManager{path: filepath.Join(tmpDir, "checkpoint")}, + backendRepoClient: &fakeBackendRepoClient{}, + } + worker.containerInstances.Set(containerID, &ContainerInstance{ + Id: containerID, + Runtime: &mockRuntime{}, + Overlay: common.NewContainerOverlay( + request, + mergedDir, + filepath.Join(tmpDir, "overlay"), + ), + }) + + err := worker.createCheckpoint(context.Background(), &CreateCheckpointOpts{ + Request: request, + CheckpointId: "checkpoint", + ContainerIp: "127.0.0.1", + OutputLogger: slog.New(slog.NewTextHandler(io.Discard, nil)), + WaitForSignal: true, + }) + + require.Error(t, err) + require.FileExists(t, filepath.Join(signalDir, checkpointCompleteFileName)) +} + func TestAddRequestMountsBuildsVolumeCacheMap(t *testing.T) { localPath := filepath.Join(t.TempDir(), "volume") spec := getTestBaseSpec() @@ -1083,6 +1131,25 @@ func (m *startedCRIUManager) RestoreCheckpoint(ctx context.Context, rt runtime.R return 0, nil } +type checkpointPathCRIUManager struct { + path string +} + +func (m *checkpointPathCRIUManager) Available() bool { + return true +} + +func (m *checkpointPathCRIUManager) CreateCheckpoint(ctx context.Context, rt runtime.Runtime, checkpointId string, request *types.ContainerRequest) (string, error) { + if err := os.MkdirAll(m.path, 0755); err != nil { + return "", err + } + return m.path, nil +} + +func (m *checkpointPathCRIUManager) RestoreCheckpoint(ctx context.Context, rt runtime.Runtime, opts *RestoreOpts) (int, error) { + return 0, nil +} + type restoreErrorCRIUManager struct { exitCode int err error