feat: add K8s scheduling options, image pull secrets, and dedicated model downloader images

Author: 照微

cmd/llmctl/autobenchmark/dashboard.go

@@ -50,11 +50,12 @@ const (
 )
 
 type dashboardOptions struct {
-	cf         *genericclioptions.ConfigFlags
-	configFile string
-	name       string
-	image      string
-	localPort  int
+	cf               *genericclioptions.ConfigFlags
+	configFile       string
+	name             string
+	image            string
+	imagePullSecrets []string
+	localPort        int
 }
 
 func newDashboardCmd(cf *genericclioptions.ConfigFlags) *cobra.Command {
@@ -83,6 +84,7 @@ If [name] is provided, it overrides the experiment name from the config file.`,
 
 	cmd.Flags().StringVarP(&opts.configFile, "config", "f", "", "Path to auto-benchmark config file (required)")
 	cmd.Flags().StringVar(&opts.image, "image", dashboardImage, "Dashboard container image")
+	cmd.Flags().StringArrayVar(&opts.imagePullSecrets, "image-pull-secret", nil, "Image pull secret names for private registries (can be specified multiple times)")
 	cmd.Flags().IntVarP(&opts.localPort, "port", "p", defaultLocalPort, "Local port for port-forward")
 	_ = cmd.MarkFlagRequired("config")
 
@@ -180,6 +182,13 @@ func (o *dashboardOptions) run(ctx context.Context) error {
 		},
 	}
 
+	// Apply image pull secrets
+	if len(o.imagePullSecrets) > 0 {
+		for _, name := range o.imagePullSecrets {
+			deploy.Spec.Template.Spec.ImagePullSecrets = append(deploy.Spec.Template.Spec.ImagePullSecrets, corev1.LocalObjectReference{Name: name})
+		}
+	}
+
 	// Delete existing if present, then create
 	deployClient := clientset.AppsV1().Deployments(namespace)
 	_, err = deployClient.Get(ctx, deployName, metav1.GetOptions{})

cmd/llmctl/autobenchmark/run.go

@@ -37,11 +37,12 @@ import (
 )
 
 type runOptions struct {
-	cf             *genericclioptions.ConfigFlags
-	configFile     string
-	name           string
-	image          string
-	serviceAccount string
+	cf               *genericclioptions.ConfigFlags
+	configFile       string
+	name             string
+	image            string
+	serviceAccount   string
+	imagePullSecrets []string
 }
 
 func newRunCmd(cf *genericclioptions.ConfigFlags) *cobra.Command {
@@ -67,6 +68,7 @@ See the guide for RBAC setup instructions.`,
 	cmd.Flags().StringVar(&opts.name, "name", "", "Experiment name (defaults to auto-generated name from config)")
 	cmd.Flags().StringVar(&opts.image, "image", controllerImage, "Controller image")
 	cmd.Flags().StringVar(&opts.serviceAccount, "service-account", "", "ServiceAccount for the controller Job (required)")
+	cmd.Flags().StringArrayVar(&opts.imagePullSecrets, "image-pull-secret", nil, "Image pull secret names for private registries (can be specified multiple times)")
 	_ = cmd.MarkFlagRequired("config")
 	_ = cmd.MarkFlagRequired("service-account")
 
@@ -187,6 +189,7 @@ func (o *runOptions) run(ctx context.Context) error {
 			Template: corev1.PodTemplateSpec{
 				Spec: corev1.PodSpec{
 					ServiceAccountName: o.serviceAccount,
+					ImagePullSecrets:   toImagePullSecrets(o.imagePullSecrets),
 					RestartPolicy:      corev1.RestartPolicyNever,
 					Containers: []corev1.Container{
 						{
@@ -266,6 +269,17 @@ func validateExpName(name string) error {
 	return nil
 }
 
+func toImagePullSecrets(names []string) []corev1.LocalObjectReference {
+	if len(names) == 0 {
+		return nil
+	}
+	refs := make([]corev1.LocalObjectReference, len(names))
+	for i, name := range names {
+		refs[i] = corev1.LocalObjectReference{Name: name}
+	}
+	return refs
+}
+
 func sanitizeCMKey(name string) string {
 	name = strings.ToLower(name)
 	name = invalidCMKeyChars.ReplaceAllString(name, "-")

cmd/llmctl/benchmark/benchmark.go

@@ -70,11 +70,12 @@ type BenchmarkOptions struct {
 
 	extraArgs map[string]string
 
-	image         string
-	cpuRequest    string
-	cpuLimit      string
-	memoryRequest string
-	memoryLimit   string
+	image            string
+	imagePullSecrets []string
+	cpuRequest       string
+	cpuLimit         string
+	memoryRequest    string
+	memoryLimit      string
 
 	wait bool
 }
@@ -103,11 +104,12 @@ type BenchmarkConfig struct {
 
 	ExtraArgs map[string]string `json:"extraArgs,omitempty"`
 
-	Image         string `json:"image,omitempty"`
-	CPURequest    string `json:"cpuRequest,omitempty"`
-	CPULimit      string `json:"cpuLimit,omitempty"`
-	MemoryRequest string `json:"memoryRequest,omitempty"`
-	MemoryLimit   string `json:"memoryLimit,omitempty"`
+	Image            string   `json:"image,omitempty"`
+	ImagePullSecrets []string `json:"imagePullSecrets,omitempty"`
+	CPURequest       string   `json:"cpuRequest,omitempty"`
+	CPULimit         string   `json:"cpuLimit,omitempty"`
+	MemoryRequest    string   `json:"memoryRequest,omitempty"`
+	MemoryLimit      string   `json:"memoryLimit,omitempty"`
 
 	Wait *bool `json:"wait,omitempty"`
 }
@@ -195,6 +197,7 @@ version management.`,
 			"Defaults to the generated job name if not specified.")
 
 	cmd.Flags().StringVar(&benchmarkOpts.image, "image", benchmarkOpts.image, "Container image used for benchmark job")
+	cmd.Flags().StringArrayVar(&benchmarkOpts.imagePullSecrets, "image-pull-secret", nil, "Image pull secret names for private registries (can be specified multiple times)")
 	cmd.Flags().StringVar(&benchmarkOpts.cpuRequest, "cpu-request", benchmarkOpts.cpuRequest, "CPU request for benchmark pod")
 	cmd.Flags().StringVar(&benchmarkOpts.cpuLimit, "cpu-limit", benchmarkOpts.cpuLimit, "CPU limit for benchmark pod")
 	cmd.Flags().StringVar(&benchmarkOpts.memoryRequest, "memory-request", benchmarkOpts.memoryRequest, "Memory request for benchmark pod")
@@ -262,7 +265,7 @@ func getConflictingFlags(cmd *cobra.Command) []string {
 		"api-backend", "api-base", "api-port", "api-key", "api-model-name",
 		"model-tokenizer",
 		"experiment-base-dir", "experiment-folder-name",
-		"image", "cpu-request", "cpu-limit", "memory-request", "memory-limit",
+		"image", "image-pull-secret", "cpu-request", "cpu-limit", "memory-request", "memory-limit",
 		"wait", "extra-args",
 	}
 	var conflicting []string
@@ -291,6 +294,7 @@ func applyConfigToOptions(cfg *BenchmarkConfig, cmd *cobra.Command) {
 	setStringOpt(&benchmarkOpts.experimentBaseDir, cfg.ExperimentBaseDir, "experiment-base-dir", cmd)
 	setStringOpt(&benchmarkOpts.experimentFolderName, cfg.ExperimentFolderName, "experiment-folder-name", cmd)
 	setStringOpt(&benchmarkOpts.image, cfg.Image, "image", cmd)
+	setStringSliceOpt(&benchmarkOpts.imagePullSecrets, cfg.ImagePullSecrets, "image-pull-secret", cmd)
 	setStringOpt(&benchmarkOpts.cpuRequest, cfg.CPURequest, "cpu-request", cmd)
 	setStringOpt(&benchmarkOpts.cpuLimit, cfg.CPULimit, "cpu-limit", cmd)
 	setStringOpt(&benchmarkOpts.memoryRequest, cfg.MemoryRequest, "memory-request", cmd)

cmd/llmctl/benchmark/dashboard_cmd.go

@@ -53,6 +53,7 @@ type DashboardOptions struct {
 	localPort         int
 	openBrowser       bool
 	image             string
+	imagePullSecrets  []string
 }
 
 var dashboardOpts DashboardOptions
@@ -99,6 +100,8 @@ Press Ctrl+C to stop the dashboard and clean up resources.`,
 		"Automatically open browser after dashboard is ready")
 	cmd.Flags().StringVar(&dashboardOpts.image, "image", defaultDashboardImage,
 		"Container image for the benchmark dashboard")
+	cmd.Flags().StringArrayVar(&dashboardOpts.imagePullSecrets, "image-pull-secret", nil,
+		"Image pull secret names for private registries (can be specified multiple times)")
 
 	return cmd
 }

cmd/llmctl/benchmark/dashboard_pod.go

@@ -38,7 +38,7 @@ func buildDashboardPod(namespace string, pvcComponents *PVCComponents) *corev1.P
 		volumeMount.SubPath = pvcComponents.SubPath
 	}
 
-	return &corev1.Pod{
+	pod := &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
 			GenerateName: "rbg-benchmark-dashboard-",
 			Namespace:    namespace,
@@ -94,4 +94,12 @@ func buildDashboardPod(namespace string, pvcComponents *PVCComponents) *corev1.P
 			RestartPolicy: corev1.RestartPolicyNever,
 		},
 	}
+
+	if len(dashboardOpts.imagePullSecrets) > 0 {
+		for _, name := range dashboardOpts.imagePullSecrets {
+			pod.Spec.ImagePullSecrets = append(pod.Spec.ImagePullSecrets, corev1.LocalObjectReference{Name: name})
+		}
+	}
+
+	return pod
 }

cmd/llmctl/benchmark/job.go

@@ -132,6 +132,12 @@ func buildBenchmarkJob(namespace, rbgName string) (*batchv1.Job, error) {
 		RestartPolicy: corev1.RestartPolicyNever,
 	}
 
+	if len(benchmarkOpts.imagePullSecrets) > 0 {
+		for _, name := range benchmarkOpts.imagePullSecrets {
+			podSpec.ImagePullSecrets = append(podSpec.ImagePullSecrets, corev1.LocalObjectReference{Name: name})
+		}
+	}
+
 	backoffLimit := int32(0)
 	ttlSecondsAfterFinished := int32(604800) // 7 days
 
@@ -178,6 +184,7 @@ func buildBenchmarkAnnotations() (map[string]string, error) {
 		ExperimentFolderName: benchmarkOpts.experimentFolderName,
 		ExtraArgs:            benchmarkOpts.extraArgs,
 		Image:                benchmarkOpts.image,
+		ImagePullSecrets:     benchmarkOpts.imagePullSecrets,
 		CPURequest:           benchmarkOpts.cpuRequest,
 		CPULimit:             benchmarkOpts.cpuLimit,
 		MemoryRequest:        benchmarkOpts.memoryRequest,

cmd/llmctl/model/list.go

@@ -36,7 +36,9 @@ import (
 
 func newListCmd(cf *genericclioptions.ConfigFlags) *cobra.Command {
 	var (
-		storage string
+		storage          string
+		image            string
+		imagePullSecrets []string
 	)
 
 	cmd := &cobra.Command{
@@ -94,7 +96,7 @@ Examples:
 			}
 
 			// Create a job to scan storage and list models
-			job := buildListModelsJob(storageplugin.DefaultMountPath)
+			job := buildListModelsJob(storageplugin.DefaultMountPath, image, imagePullSecrets)
 
 			// Mount storage (provisions resources for OSS and adds volumes/mounts)
 			ctrlClient, err := util.GetControllerRuntimeClient(cf)
@@ -141,12 +143,16 @@ Examples:
 	}
 
 	cmd.Flags().StringVar(&storage, "storage", "", "Storage to use (overrides default)")
+	cmd.Flags().StringVar(&image, "image", defaultListImage, "Container image for the list job")
+	cmd.Flags().StringArrayVar(&imagePullSecrets, "image-pull-secret", nil, "Image pull secret names for private registries (can be specified multiple times)")
 
 	return cmd
 }
 
+const defaultListImage = "alpine:latest"
+
 // buildListModelsJob creates a Job to scan storage and list models
-func buildListModelsJob(mountPath string) *batchv1.Job {
+func buildListModelsJob(mountPath, image string, imagePullSecrets []string) *batchv1.Job {
 	timestamp := time.Now().Unix()
 	jobName := fmt.Sprintf("list-models-%d", timestamp)
 
@@ -198,7 +204,7 @@ echo "]" >> "$OUTPUT_FILE"
 cat "$OUTPUT_FILE"
 `, mountPath)
 
-	return &batchv1.Job{
+	job := &batchv1.Job{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:   jobName,
 			Labels: labels,
@@ -215,7 +221,7 @@ cat "$OUTPUT_FILE"
 					Containers: []corev1.Container{
 						{
 							Name:    "scanner",
-							Image:   "alpine:latest",
+							Image:   image,
 							Command: []string{"/bin/sh", "-c"},
 							Args:    []string{scanScript},
 						},
@@ -225,6 +231,14 @@ cat "$OUTPUT_FILE"
 			},
 		},
 	}
+
+	if len(imagePullSecrets) > 0 {
+		for _, name := range imagePullSecrets {
+			job.Spec.Template.Spec.ImagePullSecrets = append(job.Spec.Template.Spec.ImagePullSecrets, corev1.LocalObjectReference{Name: name})
+		}
+	}
+
+	return job
 }
 
 // getJobOutput retrieves the output from the job's pod

cmd/llmctl/model/list_test.go

@@ -45,38 +45,38 @@ func TestInt64Ptr(t *testing.T) {
 // --- buildListModelsJob ---
 
 func TestBuildListModelsJob_NamePrefix(t *testing.T) {
-	job := buildListModelsJob("/models")
+	job := buildListModelsJob("/models", defaultListImage, nil)
 	assert.True(t, strings.HasPrefix(job.Name, "list-models-"), "job name should start with list-models-")
 }
 
 func TestBuildListModelsJob_Labels(t *testing.T) {
-	job := buildListModelsJob("/models")
+	job := buildListModelsJob("/models", defaultListImage, nil)
 	assert.Equal(t, "true", job.Labels["rbg-list-job"])
 }
 
 func TestBuildListModelsJob_ContainerSpec(t *testing.T) {
-	job := buildListModelsJob("/data/models")
+	job := buildListModelsJob("/models", defaultListImage, nil)
 	require.Len(t, job.Spec.Template.Spec.Containers, 1)
 	c := job.Spec.Template.Spec.Containers[0]
 	assert.Equal(t, "scanner", c.Name)
-	assert.Equal(t, "alpine:latest", c.Image)
+	assert.Equal(t, defaultListImage, c.Image)
 	assert.Equal(t, []string{"/bin/sh", "-c"}, c.Command)
 }
 
 func TestBuildListModelsJob_ScriptContainsMountPath(t *testing.T) {
-	job := buildListModelsJob("/custom/path")
+	job := buildListModelsJob("/custom/path", defaultListImage, nil)
 	args := job.Spec.Template.Spec.Containers[0].Args
 	require.Len(t, args, 1)
 	assert.Contains(t, args[0], "/custom/path")
 }
 
 func TestBuildListModelsJob_RestartPolicy(t *testing.T) {
-	job := buildListModelsJob("/models")
+	job := buildListModelsJob("/models", defaultListImage, nil)
 	assert.Equal(t, corev1.RestartPolicyNever, job.Spec.Template.Spec.RestartPolicy)
 }
 
 func TestBuildListModelsJob_Limits(t *testing.T) {
-	job := buildListModelsJob("/models")
+	job := buildListModelsJob("/models", defaultListImage, nil)
 	assert.Equal(t, int32(2), *job.Spec.BackoffLimit)
 	assert.Equal(t, int64(300), *job.Spec.ActiveDeadlineSeconds)
 	assert.Equal(t, int32(60), *job.Spec.TTLSecondsAfterFinished)

cmd/llmctl/model/pull.go

@@ -41,12 +41,14 @@ import (
 
 func newPullCmd(cf *genericclioptions.ConfigFlags) *cobra.Command {
 	var (
-		revision string
-		source   string
-		storage  string
-		wait     bool
-		memory   string
-		cpu      string
+		revision         string
+		source           string
+		storage          string
+		wait             bool
+		memory           string
+		cpu              string
+		image            string
+		imagePullSecrets []string
 	)
 
 	cmd := &cobra.Command{
@@ -144,6 +146,10 @@ Examples:
 				return fmt.Errorf("failed to generate download template: %w", err)
 			}
 
+			if image != "" && len(podTemplate.Spec.Containers) > 0 {
+				podTemplate.Spec.Containers[0].Image = image
+			}
+
 			// Inject metadata saving logic - wraps the original command to save model info after download
 			injectMetadataSave(podTemplate, modelID, revision, modelPath)
 
@@ -164,6 +170,13 @@ Examples:
 			// Apply resource limits to the download container
 			applyPullResources(podTemplate, cpu, memory)
 
+			// Apply image pull secrets
+			if len(imagePullSecrets) > 0 {
+				for _, name := range imagePullSecrets {
+					podTemplate.Spec.ImagePullSecrets = append(podTemplate.Spec.ImagePullSecrets, corev1.LocalObjectReference{Name: name})
+				}
+			}
+
 			// Create the Job
 			job := buildPullJob(modelID, podTemplate)
 
@@ -208,6 +221,8 @@ Examples:
 	cmd.Flags().BoolVar(&wait, "wait", true, "Wait for the pull job to complete and stream logs")
 	cmd.Flags().StringVar(&memory, "memory", defaultPullMemory, "Memory request/limit for the download container (e.g. 8Gi, 16Gi)")
 	cmd.Flags().StringVar(&cpu, "cpu", defaultPullCPU, "CPU request/limit for the download container (e.g. 2, 4)")
+	cmd.Flags().StringVar(&image, "image", "", "Container image override for the download job (default: from source plugin)")
+	cmd.Flags().StringArrayVar(&imagePullSecrets, "image-pull-secret", nil, "Image pull secret names for private registries (can be specified multiple times)")
 
 	return cmd
 }