Skip to content

Commit 16297b0

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

File tree

17 files changed

+112
-166
lines changed

17 files changed

+112
-166
lines changed

‎Makefile‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -142,38 +142,38 @@ install: build
142142
sudo cp bin/workloadmanager /usr/local/bin/
143143

144144
# Docker image variables
145-
APISERVER_IMAGE ?= workloadmanager:latest
145+
WORKLOAD_MANAGER_IMAGE ?= workloadmanager:latest
146146
IMAGE_REGISTRY ?= ""
147147

148148
# Docker and Kubernetes targets
149149
docker-build:
150150
@echo "Building Docker image..."
151-
docker build -t $(APISERVER_IMAGE) .
151+
docker build -t $(WORKLOAD_MANAGER_IMAGE) .
152152

153153
# Multi-architecture build (supports amd64, arm64)
154154
docker-buildx:
155155
@echo "Building multi-architecture Docker image..."
156-
docker buildx build --platform linux/amd64,linux/arm64 -t $(APISERVER_IMAGE) .
156+
docker buildx build --platform linux/amd64,linux/arm64 -t $(WORKLOAD_MANAGER_IMAGE) .
157157

158158
# Multi-architecture build and push
159159
docker-buildx-push:
160160
@if [ -z "$(IMAGE_REGISTRY)" ]; then \
161161
echo "Error: IMAGE_REGISTRY not set. Usage: make docker-buildx-push IMAGE_REGISTRY=your-registry.com"; \
162162
exit 1; \
163163
fi
164-
@echo "Building and pushing multi-architecture Docker image to $(IMAGE_REGISTRY)/$(APISERVER_IMAGE)..."
164+
@echo "Building and pushing multi-architecture Docker image to $(IMAGE_REGISTRY)/$(WORKLOAD_MANAGER_IMAGE)..."
165165
docker buildx build --platform linux/amd64,linux/arm64 \
166-
-t $(IMAGE_REGISTRY)/$(APISERVER_IMAGE) \
166+
-t $(IMAGE_REGISTRY)/$(WORKLOAD_MANAGER_IMAGE) \
167167
--push .
168168

169169
docker-push: docker-build
170170
@if [ -z "$(IMAGE_REGISTRY)" ]; then \
171171
echo "Error: IMAGE_REGISTRY not set. Usage: make docker-push IMAGE_REGISTRY=your-registry.com"; \
172172
exit 1; \
173173
fi
174-
@echo "Tagging and pushing Docker image to $(IMAGE_REGISTRY)/$(APISERVER_IMAGE)..."
175-
docker tag $(APISERVER_IMAGE) $(IMAGE_REGISTRY)/$(APISERVER_IMAGE)
176-
docker push $(IMAGE_REGISTRY)/$(APISERVER_IMAGE)
174+
@echo "Tagging and pushing Docker image to $(IMAGE_REGISTRY)/$(WORKLOAD_MANAGER_IMAGE)..."
175+
docker tag $(WORKLOAD_MANAGER_IMAGE) $(IMAGE_REGISTRY)/$(WORKLOAD_MANAGER_IMAGE)
176+
docker push $(IMAGE_REGISTRY)/$(WORKLOAD_MANAGER_IMAGE)
177177

178178
k8s-deploy:
179179
@echo "Deploying to Kubernetes..."
@@ -190,7 +190,7 @@ k8s-logs:
190190
# Load image to kind cluster
191191
kind-load:
192192
@echo "Loading image to kind..."
193-
kind load docker-image $(APISERVER_IMAGE)
193+
kind load docker-image $(WORKLOAD_MANAGER_IMAGE)
194194

195195
# Sandbox image targets
196196
SANDBOX_IMAGE ?= sandbox:latest

‎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: 9 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"log"
77
"time"
88

9+
"github.com/volcano-sh/agentcube/pkg/common/types"
910
"github.com/volcano-sh/agentcube/pkg/redis"
1011
"k8s.io/apimachinery/pkg/api/errors"
1112
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -22,13 +23,6 @@ type garbageCollector struct {
2223
redisClient redis.Client
2324
}
2425

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

9866
if len(gcSandboxes) > 0 {
99-
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))
10068
}
10169

10270
errs := make([]error, 0, len(gcSandboxes))
10371
// delete sandboxes
10472
for _, gcSandbox := range gcSandboxes {
105-
if gcSandbox.Kind == "SandboxClaim" {
106-
err = gc.deleteSandboxClaim(ctx, gcSandbox.Namespace, gcSandbox.Name)
73+
if gcSandbox.Kind == types.SandboxClaimsKind {
74+
err = gc.deleteSandboxClaim(ctx, gcSandbox.SandboxNamespace, gcSandbox.Name)
10775
} else {
108-
err = gc.deleteSandbox(ctx, gcSandbox.Namespace, gcSandbox.Name)
76+
err = gc.deleteSandbox(ctx, gcSandbox.SandboxNamespace, gcSandbox.Name)
10977
}
11078
if err != nil {
11179
errs = append(errs, err)
11280
continue
11381
}
114-
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)
11583
err = gc.redisClient.DeleteSandboxBySessionIDTx(ctx, gcSandbox.SessionID)
11684
if err != nil {
11785
errs = append(errs, err)

‎pkg/workloadmanager/handlers.go‎

Lines changed: 20 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"github.com/volcano-sh/agentcube/pkg/common/types"
1515
"github.com/volcano-sh/agentcube/pkg/redis"
1616
sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
17+
"sigs.k8s.io/agent-sandbox/controllers"
1718
extensionsv1alpha1 "sigs.k8s.io/agent-sandbox/extensions/api/v1alpha1"
1819
)
1920

@@ -75,7 +76,7 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
7576
namespace := sandbox.Namespace
7677

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

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

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

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

219219
// Call sandbox init endpoint with JWT-signed request
@@ -223,6 +223,7 @@ func (s *Server) handleCreateSandbox(c *gin.Context) {
223223
sandbox.Labels[SessionIdLabelKey],
224224
createAgentRequest.PublicKey,
225225
createAgentRequest.Metadata,
226+
createAgentRequest.InitTimeoutSeconds,
226227
)
227228

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

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

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

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

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

0 commit comments

Comments
 (0)