ttl more keys

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,
         }

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),

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
 	}

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)
+}

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 {

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
 		}

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)
+}

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")
 	}

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",

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
 }
 

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)