Skip to content

Commit 1c41bc6

Browse files
tjucodertjucoder
committed
fix workload manager review comment
Signed-off-by: tjucoder <chinesecoder@foxmail.com>
1 parent 1dbaac4 commit 1c41bc6

File tree

17 files changed

+113
-168
lines changed

17 files changed

+113
-168
lines changed

‎Makefile‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -153,39 +153,39 @@ install: build
153153
sudo cp bin/workloadmanager /usr/local/bin/
154154

155155
# Docker image variables
156-
APISERVER_IMAGE ?= workloadmanager:latest
156+
WORKLOAD_MANAGER_IMAGE ?= workloadmanager:latest
157157
ROUTER_IMAGE ?= agentcube-router:latest
158158
IMAGE_REGISTRY ?= ""
159159

160160
# Docker and Kubernetes targets
161161
docker-build:
162162
@echo "Building Docker image..."
163-
docker build -t $(APISERVER_IMAGE) .
163+
docker build -t $(WORKLOAD_MANAGER_IMAGE) .
164164

165165
# Multi-architecture build (supports amd64, arm64)
166166
docker-buildx:
167167
@echo "Building multi-architecture Docker image..."
168-
docker buildx build --platform linux/amd64,linux/arm64 -t $(APISERVER_IMAGE) .
168+
docker buildx build --platform linux/amd64,linux/arm64 -t $(WORKLOAD_MANAGER_IMAGE) .
169169

170170
# Multi-architecture build and push
171171
docker-buildx-push:
172172
@if [ -z "$(IMAGE_REGISTRY)" ]; then \
173173
echo "Error: IMAGE_REGISTRY not set. Usage: make docker-buildx-push IMAGE_REGISTRY=your-registry.com"; \
174174
exit 1; \
175175
fi
176-
@echo "Building and pushing multi-architecture Docker image to $(IMAGE_REGISTRY)/$(APISERVER_IMAGE)..."
176+
@echo "Building and pushing multi-architecture Docker image to $(IMAGE_REGISTRY)/$(WORKLOAD_MANAGER_IMAGE)..."
177177
docker buildx build --platform linux/amd64,linux/arm64 \
178-
-t $(IMAGE_REGISTRY)/$(APISERVER_IMAGE) \
178+
-t $(IMAGE_REGISTRY)/$(WORKLOAD_MANAGER_IMAGE) \
179179
--push .
180180

181181
docker-push: docker-build
182182
@if [ -z "$(IMAGE_REGISTRY)" ]; then \
183183
echo "Error: IMAGE_REGISTRY not set. Usage: make docker-push IMAGE_REGISTRY=your-registry.com"; \
184184
exit 1; \
185185
fi
186-
@echo "Tagging and pushing Docker image to $(IMAGE_REGISTRY)/$(APISERVER_IMAGE)..."
187-
docker tag $(APISERVER_IMAGE) $(IMAGE_REGISTRY)/$(APISERVER_IMAGE)
188-
docker push $(IMAGE_REGISTRY)/$(APISERVER_IMAGE)
186+
@echo "Tagging and pushing Docker image to $(IMAGE_REGISTRY)/$(WORKLOAD_MANAGER_IMAGE)..."
187+
docker tag $(WORKLOAD_MANAGER_IMAGE) $(IMAGE_REGISTRY)/$(WORKLOAD_MANAGER_IMAGE)
188+
docker push $(IMAGE_REGISTRY)/$(WORKLOAD_MANAGER_IMAGE)
189189

190190
k8s-deploy:
191191
@echo "Deploying to Kubernetes..."
@@ -202,7 +202,7 @@ k8s-logs:
202202
# Load image to kind cluster
203203
kind-load:
204204
@echo "Loading image to kind..."
205-
kind load docker-image $(APISERVER_IMAGE)
205+
kind load docker-image $(WORKLOAD_MANAGER_IMAGE)
206206

207207
# Router Docker targets
208208
docker-build-router:

‎cmd/workload-manager/main.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ func main() {
4141
enableTLS = flag.Bool("enable-tls", false, "Enable TLS (HTTPS)")
4242
tlsCert = flag.String("tls-cert", "", "Path to TLS certificate file")
4343
tlsKey = flag.String("tls-key", "", "Path to TLS key file")
44+
enableAuth = flag.Bool("enable-auth", false, "Enable Authentication")
4445
)
4546

4647
// Parse command line flags
@@ -83,6 +84,7 @@ func main() {
8384
EnableTLS: *enableTLS,
8485
TLSCert: *tlsCert,
8586
TLSKey: *tlsKey,
87+
EnableAuth: *enableAuth,
8688
}
8789

8890
// Create and initialize API server

‎pkg/common/types/sandbox.go‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,10 @@ import (
66
)
77

88
type SandboxRedis struct {
9+
Kind string `json:"kind"`
910
SandboxID string `json:"sandboxId"`
1011
SandboxNamespace string `json:"sandboxNamespace"`
11-
SandboxName string `json:"sandboxName"`
12-
SandboxClaimName string `json:"sandboxClaimName"`
12+
Name string `json:"name"`
1313
EntryPoints []SandboxEntryPoints `json:"entryPoints"`
1414
SessionID string `json:"sessionId"`
1515
CreatedAt time.Time `json:"createdAt"`
@@ -26,12 +26,13 @@ type SandboxEntryPoints struct {
2626
}
2727

2828
type CreateSandboxRequest struct {
29-
Kind string `json:"kind"`
30-
Name string `json:"name"`
31-
Namespace string `json:"namespace"`
32-
Auth Auth `json:"auth"`
33-
Metadata map[string]string `json:"metadata"`
34-
PublicKey string `json:"publicKey,omitempty"`
29+
Kind string `json:"kind"`
30+
Name string `json:"name"`
31+
Namespace string `json:"namespace"`
32+
Auth Auth `json:"auth"`
33+
Metadata map[string]string `json:"metadata"`
34+
PublicKey string `json:"publicKey,omitempty"`
35+
InitTimeoutSeconds int `json:"initTimeoutSeconds,omitempty"`
3536
}
3637

3738
type Auth struct {

‎pkg/common/types/types.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,4 +3,6 @@ package types
33
const (
44
AgentRuntimeKind = "AgentRuntime"
55
CodeInterpreterKind = "CodeInterpreter"
6+
SandboxKind = "Sandbox"
7+
SandboxClaimsKind = "SandboxClaim"
68
)

‎pkg/redis/client_test.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ func newTestClient(t *testing.T) (*client, *miniredis.Miniredis) {
3131
func newTestSandbox(id string, sessionID string, expiresAt time.Time) *types.SandboxRedis {
3232
return &types.SandboxRedis{
3333
SandboxID: id,
34-
SandboxName: "test-sandbox-" + id,
34+
Name: "test-sandbox-" + id,
3535
EntryPoints: nil,
3636
SessionID: sessionID,
3737
CreatedAt: time.Now().UTC(),
@@ -55,7 +55,7 @@ func TestClient_StoreSandbox(t *testing.T) {
5555
sandboxRedis := &types.SandboxRedis{
5656
SessionID: "session-id-TestClient_StoreSandbox-01",
5757
SandboxNamespace: "agent-cube",
58-
SandboxName: "fake-sandbox-01",
58+
Name: "fake-sandbox-01",
5959
ExpiresAt: time.Now(),
6060
}
6161
err := c.StoreSandbox(ctx, sandboxRedis, time.Hour)
@@ -72,7 +72,7 @@ func TestClient_UpdateSandbox(t *testing.T) {
7272
sandboxRedis := &types.SandboxRedis{
7373
SessionID: "session-id-TestClient_StoreSandbox-02",
7474
SandboxNamespace: "agent-cube",
75-
SandboxName: "fake-sandbox-01",
75+
Name: "fake-sandbox-01",
7676
ExpiresAt: time.Now(),
7777
}
7878
err := c.UpdateSandbox(ctx, sandboxRedis, time.Hour)

‎pkg/workloadmanager/auth.go‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,12 +63,7 @@ const (
6363

6464
// authMiddleware provides service account token authentication middleware
6565
func (s *Server) authMiddleware(c *gin.Context) {
66-
if !s.enableAuth {
67-
c.Next()
68-
return
69-
}
70-
// Skip authentication for health check endpoint
71-
if c.Request.URL.Path == "/health" {
66+
if !s.config.EnableAuth {
7267
c.Next()
7368
return
7469
}

‎pkg/workloadmanager/codeinterpreter_controller.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -293,7 +293,7 @@ func (r *CodeInterpreterReconciler) convertToPodTemplate(template *runtimev1alph
293293
Resources: template.Resources,
294294
VolumeMounts: []corev1.VolumeMount{
295295
{
296-
Name: "jwt-public-key",
296+
Name: JWTPublicKeyVolumeName,
297297
MountPath: "/etc/picod",
298298
ReadOnly: true,
299299
},
@@ -303,10 +303,10 @@ func (r *CodeInterpreterReconciler) convertToPodTemplate(template *runtimev1alph
303303
RuntimeClassName: template.RuntimeClassName,
304304
Volumes: []corev1.Volume{
305305
{
306-
Name: "jwt-public-key",
306+
Name: JWTPublicKeyVolumeName,
307307
VolumeSource: corev1.VolumeSource{
308308
Secret: &corev1.SecretVolumeSource{
309-
SecretName: "agentcube-jwt-public-key",
309+
SecretName: JWTPublicKeySecretName,
310310
},
311311
},
312312
},

‎pkg/workloadmanager/garbage_collection.go‎

Lines changed: 10 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,11 @@ import (
66
"log"
77
"time"
88

9+
"github.com/volcano-sh/agentcube/pkg/common/types"
10+
"github.com/volcano-sh/agentcube/pkg/redis"
911
"k8s.io/apimachinery/pkg/api/errors"
1012
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
1113
utilerrors "k8s.io/apimachinery/pkg/util/errors"
12-
13-
"github.com/volcano-sh/agentcube/pkg/redis"
1414
)
1515

1616
const (
@@ -23,13 +23,6 @@ type garbageCollector struct {
2323
redisClient redis.Client
2424
}
2525

26-
type garbageCollectorSandbox struct {
27-
Kind string // Sandbox or SandboxClaim
28-
Name string
29-
Namespace string
30-
SessionID string
31-
}
32-
3326
func newGarbageCollector(k8sClient *K8sClient, redisClient redis.Client, interval time.Duration) *garbageCollector {
3427
return &garbageCollector{
3528
k8sClient: k8sClient,
@@ -66,53 +59,27 @@ func (gc *garbageCollector) once() {
6659
if err != nil {
6760
log.Printf("garbage collector error listing expired sandboxes: %v", err)
6861
}
69-
gcSandboxes := make([]garbageCollectorSandbox, 0, len(inactiveSandboxes)+len(expiredSandboxes))
70-
for _, inactive := range inactiveSandboxes {
71-
gcSandboxObj := garbageCollectorSandbox{
72-
Namespace: inactive.SandboxNamespace,
73-
SessionID: inactive.SessionID,
74-
}
75-
if inactive.SandboxClaimName != "" {
76-
gcSandboxObj.Kind = "SandboxClaim"
77-
gcSandboxObj.Name = inactive.SandboxClaimName
78-
} else {
79-
gcSandboxObj.Kind = "Sandbox"
80-
gcSandboxObj.Name = inactive.SandboxName
81-
}
82-
gcSandboxes = append(gcSandboxes, gcSandboxObj)
83-
}
84-
for _, expired := range expiredSandboxes {
85-
gcSandboxObj := garbageCollectorSandbox{
86-
Namespace: expired.SandboxNamespace,
87-
SessionID: expired.SessionID,
88-
}
89-
if expired.SandboxClaimName != "" {
90-
gcSandboxObj.Kind = "SandboxClaim"
91-
gcSandboxObj.Name = expired.SandboxClaimName
92-
} else {
93-
gcSandboxObj.Kind = "Sandbox"
94-
gcSandboxObj.Name = expired.SandboxName
95-
}
96-
gcSandboxes = append(gcSandboxes, gcSandboxObj)
97-
}
62+
gcSandboxes := make([]*types.SandboxRedis, 0, len(inactiveSandboxes)+len(expiredSandboxes))
63+
gcSandboxes = append(gcSandboxes, inactiveSandboxes...)
64+
gcSandboxes = append(gcSandboxes, expiredSandboxes...)
9865

9966
if len(gcSandboxes) > 0 {
100-
log.Printf("garbage collector found %d sandboxes to be delete", len(gcSandboxes))
67+
log.Printf("garbage collector found %d sandboxes to be deleted", len(gcSandboxes))
10168
}
10269

10370
errs := make([]error, 0, len(gcSandboxes))
10471
// delete sandboxes
10572
for _, gcSandbox := range gcSandboxes {
106-
if gcSandbox.Kind == "SandboxClaim" {
107-
err = gc.deleteSandboxClaim(ctx, gcSandbox.Namespace, gcSandbox.Name)
73+
if gcSandbox.Kind == types.SandboxClaimsKind {
74+
err = gc.deleteSandboxClaim(ctx, gcSandbox.SandboxNamespace, gcSandbox.Name)
10875
} else {
109-
err = gc.deleteSandbox(ctx, gcSandbox.Namespace, gcSandbox.Name)
76+
err = gc.deleteSandbox(ctx, gcSandbox.SandboxNamespace, gcSandbox.Name)
11077
}
11178
if err != nil {
11279
errs = append(errs, err)
11380
continue
11481
}
115-
log.Printf("garbage collector %s %s/%s session %s deleted", gcSandbox.Kind, gcSandbox.Namespace, gcSandbox.Name, gcSandbox.SessionID)
82+
log.Printf("garbage collector %s %s/%s session %s deleted", gcSandbox.Kind, gcSandbox.SandboxNamespace, gcSandbox.Name, gcSandbox.SessionID)
11683
err = gc.redisClient.DeleteSandboxBySessionIDTx(ctx, gcSandbox.SessionID)
11784
if err != nil {
11885
errs = append(errs, err)

‎pkg/workloadmanager/handlers.go‎

Lines changed: 20 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"github.com/gin-gonic/gin"
1313
redisv9 "github.com/redis/go-redis/v9"
1414
sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
15+
"sigs.k8s.io/agent-sandbox/controllers"
1516
extensionsv1alpha1 "sigs.k8s.io/agent-sandbox/extensions/api/v1alpha1"
1617

1718
"github.com/volcano-sh/agentcube/pkg/common/types"
@@ -76,7 +77,7 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
7677
namespace := sandbox.Namespace
7778

7879
dynamicClient := s.k8sClient.dynamicClient
79-
if s.enableAuth {
80+
if s.config.EnableAuth {
8081
// Extract user information from context
8182
userToken, userNamespace, _, serviceAccountName := extractUserInfo(c)
8283
if userToken == "" || userNamespace == "" || serviceAccountName == "" {
@@ -162,7 +163,7 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
162163
}()
163164

164165
sandboxPodName := ""
165-
if podName, exists := createdSandbox.Annotations["agents.x-k8s.io/pod-name"]; exists {
166+
if podName, exists := createdSandbox.Annotations[controllers.SanboxPodNameAnnotation]; exists {
166167
sandboxPodName = podName
167168
}
168169
podIP, err := s.k8sClient.GetSandboxPodIP(c.Request.Context(), namespace, sandboxName, sandboxPodName)
@@ -194,6 +195,12 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
194195
return
195196
}
196197

198+
if len(redisCacheInfo.EntryPoints) == 0 {
199+
respondError(c, http.StatusInternalServerError, "SANDBOX_INIT_FAILED",
200+
"No access endpoint found for sandbox initialization")
201+
return
202+
}
203+
197204
// Code Interpreter sandbox created, init code interpreter
198205
// Find the /init endpoint from entryPoints
199206
var initEndpoint string
@@ -204,17 +211,10 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
204211
}
205212
}
206213

207-
// If no /init path found, use the first entryPoint endpoint with /init appended
214+
// If no /init path found, use the first entryPoint endpoint fallback
208215
if initEndpoint == "" {
209-
if len(redisCacheInfo.EntryPoints) > 0 {
210-
initEndpoint = fmt.Sprintf("%s://%s",
211-
redisCacheInfo.EntryPoints[0].Protocol,
212-
redisCacheInfo.EntryPoints[0].Endpoint)
213-
} else {
214-
respondError(c, http.StatusInternalServerError, "SANDBOX_INIT_FAILED",
215-
"No access endpoint found for sandbox initialization")
216-
return
217-
}
216+
initEndpoint = fmt.Sprintf("%s://%s", redisCacheInfo.EntryPoints[0].Protocol,
217+
redisCacheInfo.EntryPoints[0].Endpoint)
218218
}
219219

220220
// Call sandbox init endpoint with JWT-signed request
@@ -224,6 +224,7 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
224224
sandbox.Labels[SessionIdLabelKey],
225225
createAgentRequest.PublicKey,
226226
createAgentRequest.Metadata,
227+
createAgentRequest.InitTimeoutSeconds,
227228
)
228229

229230
if err != nil {
@@ -263,7 +264,7 @@ func (s *Server) handleDeleteSandbox(c *gin.Context) {
263264
}
264265

265266
dynamicClient := s.k8sClient.dynamicClient
266-
if s.enableAuth {
267+
if s.config.EnableAuth {
267268
// Extract user information from context
268269
userToken, userNamespace, _, serviceAccountName := extractUserInfo(c)
269270

@@ -282,20 +283,18 @@ func (s *Server) handleDeleteSandbox(c *gin.Context) {
282283
dynamicClient = userClient.dynamicClient
283284
}
284285

285-
if sandbox.SandboxClaimName != "" {
286-
// SandboxClaimName is not empty, we should delete SandboxClaim
287-
err = deleteSandboxClaim(c.Request.Context(), dynamicClient, sandbox.SandboxNamespace, sandbox.SandboxClaimName)
286+
if sandbox.Kind == types.SandboxClaimsKind {
287+
err = deleteSandboxClaim(c.Request.Context(), dynamicClient, sandbox.SandboxNamespace, sandbox.Name)
288288
if err != nil {
289-
log.Printf("failed to delete sandbox claim %s/%s: %v", sandbox.SandboxNamespace, sandbox.SandboxClaimName, err)
289+
log.Printf("failed to delete sandbox claim %s/%s: %v", sandbox.SandboxNamespace, sandbox.Name, err)
290290
respondError(c, http.StatusForbidden, "SANDBOX_CLAIM_DELETE_FAILED",
291291
fmt.Sprintf("Failed to delete sandbox claim (namespace: %s): %v", sandbox.SandboxNamespace, err))
292292
return
293293
}
294294
} else {
295-
// SandboxClaimName is empty, we should delete Sandbox directly
296-
err = deleteSandbox(c.Request.Context(), dynamicClient, sandbox.SandboxNamespace, sandbox.SandboxName)
295+
err = deleteSandbox(c.Request.Context(), dynamicClient, sandbox.SandboxNamespace, sandbox.Name)
297296
if err != nil {
298-
log.Printf("failed to delete sandbox claim %s/%s: %v", sandbox.SandboxNamespace, sandbox.SandboxName, err)
297+
log.Printf("failed to delete sandbox %s/%s: %v", sandbox.SandboxNamespace, sandbox.Name, err)
299298
respondError(c, http.StatusForbidden, "SANDBOX_DELETE_FAILED",
300299
fmt.Sprintf("Failed to delete sandbox (namespace: %s): %v", sandbox.SandboxNamespace, err))
301300
return
@@ -309,12 +308,7 @@ func (s *Server) handleDeleteSandbox(c *gin.Context) {
309308
return
310309
}
311310

312-
objectType := SandboxGVR.Resource
313-
objectName := sandbox.SandboxName
314-
if sandbox.SandboxClaimName != "" {
315-
objectName = sandbox.SandboxClaimName
316-
}
317-
log.Printf("delete %s %s/%s successfully, sessionID: %v ", objectType, sandbox.SandboxNamespace, objectName, sandbox.SessionID)
311+
log.Printf("delete %s %s/%s successfully, sessionID: %v ", sandbox.Kind, sandbox.SandboxNamespace, sandbox.Name, sandbox.SessionID)
318312
respondJSON(c, http.StatusOK, map[string]string{
319313
"message": "Sandbox deleted successfully",
320314
})

0 commit comments

Comments
 (0)