[Core] Parallelize OCI object storage model verification (MD5)

charts/ome-resources/templates/model-agent-daemonset/daemonset.yaml

@@ -61,6 +61,8 @@ spec:
         - '2'
         - --num-high-priority-worker
         - '1'
+        - --model-verification-concurrency
+        - {{ .Values.modelAgent.modelVerificationConcurrency | quote }}
         - --same-path-wait-timeout
         - 30m
         env:

charts/ome-resources/values.yaml

@@ -125,6 +125,7 @@ modelAgent:
   # Installs whose runtimes fetch models themselves can leave this off.
   enabled: false
   hostPath: /mnt/data/models
+  modelVerificationConcurrency: 1
   priorityClassName: system-node-critical
   serviceAccountName: ome-model-agent
   image:

cmd/model-agent/main.go

@@ -32,20 +32,21 @@ import (
 
 // config holds all configuration parameters for the model agent
 type config struct {
-	port                  int
-	modelsRootDir         string
-	modelsRootDirOnHost   string
-	nodeName              string
-	nodeLabelRetry        int
-	concurrency           int
-	multipartConcurrency  int
-	downloadRetry         int
-	downloadAuthType      string
-	numDownloadWorker     int
-	numHighPriorityWorker int
-	samePathWaitTimeout   time.Duration
-	namespace             string
-	logLevel              string
+	port                         int
+	modelsRootDir                string
+	modelsRootDirOnHost          string
+	nodeName                     string
+	nodeLabelRetry               int
+	concurrency                  int
+	multipartConcurrency         int
+	modelVerificationConcurrency int
+	downloadRetry                int
+	downloadAuthType             string
+	numDownloadWorker            int
+	numHighPriorityWorker        int
+	samePathWaitTimeout          time.Duration
+	namespace                    string
+	logLevel                     string
 }
 
 // Logger type alias for zap.SugaredLogger
@@ -73,6 +74,7 @@ func init() {
 	rootCmd.PersistentFlags().IntVar(&cfg.downloadRetry, "download-retry", 3, "Number of retries for downloading")
 	rootCmd.PersistentFlags().IntVar(&cfg.concurrency, "concurrency", 4, "Number of concurrent download workers per gopher")
 	rootCmd.PersistentFlags().IntVar(&cfg.multipartConcurrency, "multipart-concurrency", 4, "Number of concurrent multipart download workers per gopher")
+	rootCmd.PersistentFlags().IntVar(&cfg.modelVerificationConcurrency, "model-verification-concurrency", 1, "Maximum concurrent OCI model file integrity checks across the model-agent pod")
 	rootCmd.PersistentFlags().IntVar(&cfg.numDownloadWorker, "num-download-worker", 5, "Number of download workers")
 	rootCmd.PersistentFlags().IntVar(&cfg.numHighPriorityWorker, "num-high-priority-worker", 1, "Number of high-priority workers for delete and same-path reuse tasks")
 	rootCmd.PersistentFlags().DurationVar(&cfg.samePathWaitTimeout, "same-path-wait-timeout", 30*time.Minute, "Maximum time to wait for same-path model reuse before falling back to normal download")
@@ -274,6 +276,7 @@ func initializeComponents(
 		logger,
 		baseModelInformer.Lister(),
 		clusterBaseModelInformer.Lister(),
+		modelagent.WithModelVerificationConcurrency(cfg.modelVerificationConcurrency),
 	)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to create gopher: %w", err)

cmd/model-agent/main_test.go

@@ -113,6 +113,7 @@ func TestDefaultConfig(t *testing.T) {
 	testCmd.Flags().IntVar(&cfg.downloadRetry, "download-retry", 3, "retry times for model download")
 	testCmd.Flags().StringVar(&cfg.downloadAuthType, "download-auth-type", "instance-principal", "authentication method for model download")
 	testCmd.Flags().IntVar(&cfg.numDownloadWorker, "num-download-worker", 3, "number of download workers")
+	testCmd.Flags().IntVar(&cfg.modelVerificationConcurrency, "model-verification-concurrency", 1, "model verification concurrency")
 	testCmd.Flags().IntVar(&cfg.numHighPriorityWorker, "num-high-priority-worker", 1, "number of high-priority workers")
 	testCmd.Flags().DurationVar(&cfg.samePathWaitTimeout, "same-path-wait-timeout", 30*time.Minute, "same-path wait timeout")
 	testCmd.Flags().StringVar(&cfg.namespace, "namespace", "ome", "the namespace of the ome model agents daemon set")
@@ -129,6 +130,7 @@ func TestDefaultConfig(t *testing.T) {
 	assert.Equal(t, 3, cfg.downloadRetry)
 	assert.Equal(t, "instance-principal", cfg.downloadAuthType)
 	assert.Equal(t, 3, cfg.numDownloadWorker)
+	assert.Equal(t, 1, cfg.modelVerificationConcurrency)
 	assert.Equal(t, 1, cfg.numHighPriorityWorker)
 	assert.Equal(t, 30*time.Minute, cfg.samePathWaitTimeout)
 	assert.Equal(t, "ome", cfg.namespace)

pkg/alfred/policy/defrag/scoring.go

@@ -27,7 +27,7 @@ type SizeFrag struct {
 }
 
 // PoolScore is the full scoring breakdown for one hardware pool
-// (Node.GPUPool). 
+// (Node.GPUPool).
 type PoolScore struct {
 	Pool string
 

pkg/modelagent/gopher.go

@@ -55,21 +55,23 @@ type activeDownload struct {
 }
 
 type Gopher struct {
-	modelConfigParser      *modelparser.ModelConfigParser
-	configMapReconciler    *ConfigMapReconciler
-	downloadRetry          int
-	concurrency            int
-	multipartConcurrency   int
-	modelRootDir           string
-	xetConfig              *xet.Config
-	kubeClient             kubernetes.Interface
-	gopherChan             chan *GopherTask
-	nodeLabelReconciler    *NodeLabelReconciler
-	metrics                *Metrics
-	logger                 *zap.SugaredLogger
-	configMapMutex         sync.Mutex // Mutex to coordinate ConfigMap access
-	baseModelLister        omev1beta1lister.BaseModelLister
-	clusterBaseModelLister omev1beta1lister.ClusterBaseModelLister
+	modelConfigParser            *modelparser.ModelConfigParser
+	configMapReconciler          *ConfigMapReconciler
+	downloadRetry                int
+	concurrency                  int
+	multipartConcurrency         int
+	modelVerificationConcurrency int
+	modelVerificationLimiter     *verificationLimiter
+	modelRootDir                 string
+	xetConfig                    *xet.Config
+	kubeClient                   kubernetes.Interface
+	gopherChan                   chan *GopherTask
+	nodeLabelReconciler          *NodeLabelReconciler
+	metrics                      *Metrics
+	logger                       *zap.SugaredLogger
+	configMapMutex               sync.Mutex // Mutex to coordinate ConfigMap access
+	baseModelLister              omev1beta1lister.BaseModelLister
+	clusterBaseModelLister       omev1beta1lister.ClusterBaseModelLister
 
 	// Track active downloads for cancellation
 	activeDownloads      map[string]activeDownload // key: model UID
@@ -82,6 +84,20 @@ type Gopher struct {
 	startupReadyModelKeys map[string]struct{}
 }
 
+type GopherOption func(*Gopher)
+
+// WithModelVerificationConcurrency bounds concurrent OCI model file integrity
+// checks across all downloads handled by this model-agent process.
+func WithModelVerificationConcurrency(concurrency int) GopherOption {
+	return func(gopher *Gopher) {
+		if concurrency < 1 {
+			concurrency = 1
+		}
+		gopher.modelVerificationConcurrency = concurrency
+		gopher.modelVerificationLimiter = newVerificationLimiter(concurrency)
+	}
+}
+
 const (
 	BigFileSizeInMB = 200
 
@@ -105,7 +121,8 @@ func NewGopher(
 	metrics *Metrics,
 	logger *zap.SugaredLogger,
 	baseModelLister omev1beta1lister.BaseModelLister,
-	clusterBaseModelLister omev1beta1lister.ClusterBaseModelLister) (*Gopher, error) {
+	clusterBaseModelLister omev1beta1lister.ClusterBaseModelLister,
+	options ...GopherOption) (*Gopher, error) {
 
 	if xetConfig == nil {
 		return nil, fmt.Errorf("xet hugging face config cannot be nil")
@@ -114,26 +131,34 @@ func NewGopher(
 		samePathWaitTimeout = defaultSamePathWaitTimeout
 	}
 
-	return &Gopher{
-		modelConfigParser:      modelConfigParser,
-		configMapReconciler:    configMapReconciler,
-		downloadRetry:          downloadRetry,
-		concurrency:            concurrency,
-		multipartConcurrency:   multipartConcurrency,
-		modelRootDir:           modelRootDir,
-		xetConfig:              xetConfig,
-		kubeClient:             kubeClient,
-		gopherChan:             gopherChan,
-		nodeLabelReconciler:    nodeLabelReconciler,
-		metrics:                metrics,
-		logger:                 logger,
-		activeDownloads:        make(map[string]activeDownload),
-		baseModelLister:        baseModelLister,
-		clusterBaseModelLister: clusterBaseModelLister,
-		taskQueue:              newGopherTaskQueue(),
-		samePathWaitDelay:      defaultSamePathWaitDelay,
-		samePathWaitTimeout:    samePathWaitTimeout,
-	}, nil
+	gopher := &Gopher{
+		modelConfigParser:            modelConfigParser,
+		configMapReconciler:          configMapReconciler,
+		downloadRetry:                downloadRetry,
+		concurrency:                  concurrency,
+		multipartConcurrency:         multipartConcurrency,
+		modelVerificationConcurrency: 1,
+		modelVerificationLimiter:     newVerificationLimiter(1),
+		modelRootDir:                 modelRootDir,
+		xetConfig:                    xetConfig,
+		kubeClient:                   kubeClient,
+		gopherChan:                   gopherChan,
+		nodeLabelReconciler:          nodeLabelReconciler,
+		metrics:                      metrics,
+		logger:                       logger,
+		activeDownloads:              make(map[string]activeDownload),
+		baseModelLister:              baseModelLister,
+		clusterBaseModelLister:       clusterBaseModelLister,
+		taskQueue:                    newGopherTaskQueue(),
+		samePathWaitDelay:            defaultSamePathWaitDelay,
+		samePathWaitTimeout:          samePathWaitTimeout,
+	}
+	for _, option := range options {
+		if option != nil {
+			option(gopher)
+		}
+	}
+	return gopher, nil
 }
 
 func (s *Gopher) Run(stopCh <-chan struct{}, numWorker int, numHighPriorityWorker int) {
@@ -1291,7 +1316,8 @@ func (s *Gopher) downloadModel(ctx context.Context, uri *ociobjectstore.ObjectUR
 	}
 
 	// Perform final verification of all downloaded files
-	s.logger.Info("Performing final integrity verification of all downloaded files...")
+	s.logger.Infof("Performing final integrity verification of %d downloaded files with pod-wide concurrency %d...",
+		len(objectUris), s.effectiveModelVerificationConcurrency())
 	verificationStartTime := time.Now()
 	verificationErrors := s.verifyDownloadedFiles(ociOSDataStore, objectUris, destPath, task)
 	verificationDuration := time.Since(verificationStartTime)
@@ -1324,23 +1350,7 @@ func (s *Gopher) downloadModel(ctx context.Context, uri *ociobjectstore.ObjectUR
 }
 
 func (s *Gopher) verifyDownloadedFiles(ociOSDataStore *ociobjectstore.OCIOSDataStore, uris []ociobjectstore.ObjectURI, destPath string, task *GopherTask) map[string]error {
-	errors := make(map[string]error)
-	for _, obj := range uris {
-		relativeName := filepath.Join(destPath, ociobjectstore.TrimObjectPrefix(obj.ObjectName, obj.Prefix))
-		// Fallback: if relativeName is empty, use the object name directly
-		if relativeName == "" {
-			relativeName = obj.ObjectName
-		}
-
-		valid, err := ociOSDataStore.IsLocalCopyValid(obj, relativeName)
-		if err != nil {
-			errors[obj.ObjectName] = err
-			continue
-		}
-		if !valid {
-			errors[obj.ObjectName] = fmt.Errorf("MD5 or size mismatch for %s", obj.ObjectName)
-		}
-	}
+	errors := s.verifyDownloadedFilesWithValidator(uris, destPath, ociOSDataStore.IsLocalCopyValid)
 
 	// Record verification result in metrics
 	modelType, namespace, name := GetModelTypeNamespaceAndName(task)
@@ -1349,6 +1359,71 @@ func (s *Gopher) verifyDownloadedFiles(ociOSDataStore *ociobjectstore.OCIOSDataS
 	return errors
 }
 
+type localCopyValidator func(ociobjectstore.ObjectURI, string) (bool, error)
+
+type verificationResult struct {
+	objectName string
+	err        error
+}
+
+func (s *Gopher) verifyDownloadedFilesWithValidator(uris []ociobjectstore.ObjectURI, destPath string, validate localCopyValidator) map[string]error {
+	errors := make(map[string]error)
+	if len(uris) == 0 {
+		return errors
+	}
+
+	workerCount := min(s.effectiveModelVerificationConcurrency(), len(uris))
+	limiter := s.modelVerificationLimiter
+	if limiter == nil {
+		limiter = newVerificationLimiter(workerCount)
+	}
+	jobs := make(chan ociobjectstore.ObjectURI)
+	results := make(chan verificationResult, len(uris))
+	var workers sync.WaitGroup
+	workers.Add(workerCount)
+	for range workerCount {
+		go func() {
+			defer workers.Done()
+			for obj := range jobs {
+				limiter.acquire()
+				relativeName := filepath.Join(destPath, ociobjectstore.TrimObjectPrefix(obj.ObjectName, obj.Prefix))
+				if relativeName == "" {
+					relativeName = obj.ObjectName
+				}
+				valid, err := validate(obj, relativeName)
+				limiter.release()
+				if err == nil && !valid {
+					err = fmt.Errorf("MD5 or size mismatch for %s", obj.ObjectName)
+				}
+				results <- verificationResult{objectName: obj.ObjectName, err: err}
+			}
+		}()
+	}
+
+	go func() {
+		for _, obj := range uris {
+			jobs <- obj
+		}
+		close(jobs)
+		workers.Wait()
+		close(results)
+	}()
+
+	for result := range results {
+		if result.err != nil {
+			errors[result.objectName] = result.err
+		}
+	}
+	return errors
+}
+
+func (s *Gopher) effectiveModelVerificationConcurrency() int {
+	if s.modelVerificationConcurrency < 1 {
+		return 1
+	}
+	return s.modelVerificationConcurrency
+}
+
 func (s *Gopher) deleteModel(destPath string, task *GopherTask) error {
 	startTime := time.Now()