diff --git a/charts/ome-resources/templates/model-agent-daemonset/daemonset.yaml b/charts/ome-resources/templates/model-agent-daemonset/daemonset.yaml index 827fa376d..32ff6b2b2 100644 --- a/charts/ome-resources/templates/model-agent-daemonset/daemonset.yaml +++ b/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: diff --git a/charts/ome-resources/values.yaml b/charts/ome-resources/values.yaml index 1ec40172e..5793ef7ec 100644 --- a/charts/ome-resources/values.yaml +++ b/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: diff --git a/cmd/model-agent/main.go b/cmd/model-agent/main.go index 160c1a7ca..880d63830 100644 --- a/cmd/model-agent/main.go +++ b/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) diff --git a/cmd/model-agent/main_test.go b/cmd/model-agent/main_test.go index f156f99b9..2b8715a74 100644 --- a/cmd/model-agent/main_test.go +++ b/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) diff --git a/pkg/alfred/policy/defrag/scoring.go b/pkg/alfred/policy/defrag/scoring.go index 03a5268f8..1da6eb112 100644 --- a/pkg/alfred/policy/defrag/scoring.go +++ b/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 diff --git a/pkg/modelagent/gopher.go b/pkg/modelagent/gopher.go index 867493974..11c5afeb4 100644 --- a/pkg/modelagent/gopher.go +++ b/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() diff --git a/pkg/modelagent/gopher_test.go b/pkg/modelagent/gopher_test.go index 4fca0cc4e..d5c310c04 100644 --- a/pkg/modelagent/gopher_test.go +++ b/pkg/modelagent/gopher_test.go @@ -7,6 +7,8 @@ import ( "fmt" "os" "path/filepath" + "sync" + "sync/atomic" "testing" "time" @@ -27,6 +29,7 @@ import ( "sigs.k8s.io/ome/pkg/apis/ome/v1beta1" omev1beta1lister "sigs.k8s.io/ome/pkg/client/listers/ome/v1beta1" "sigs.k8s.io/ome/pkg/constants" + "sigs.k8s.io/ome/pkg/ociobjectstore" "sigs.k8s.io/ome/pkg/utils/storage" ) @@ -1131,6 +1134,87 @@ func TestEnqueueTaskClassifiesStartupReadyLocalPathAsRevalidation(t *testing.T) assert.Equal(t, 0, g.taskQueue.len()) } +func TestWithModelVerificationConcurrencyCreatesSharedLimiter(t *testing.T) { + g := &Gopher{} + + WithModelVerificationConcurrency(8)(g) + + assert.Equal(t, 8, g.modelVerificationConcurrency) + require.NotNil(t, g.modelVerificationLimiter) + assert.Equal(t, 8, g.modelVerificationLimiter.limit()) + + WithModelVerificationConcurrency(0)(g) + assert.Equal(t, 1, g.modelVerificationConcurrency) + require.NotNil(t, g.modelVerificationLimiter) + assert.Equal(t, 1, g.modelVerificationLimiter.limit()) +} + +func TestVerificationConcurrencyIsSharedAcrossModels(t *testing.T) { + g := &Gopher{} + WithModelVerificationConcurrency(3)(g) + + uris := make([]ociobjectstore.ObjectURI, 12) + for i := range uris { + uris[i] = ociobjectstore.ObjectURI{ObjectName: fmt.Sprintf("model/file-%02d", i), Prefix: "model"} + } + + var active atomic.Int32 + var maxActive atomic.Int32 + validate := func(_ ociobjectstore.ObjectURI, _ string) (bool, error) { + current := active.Add(1) + for { + observed := maxActive.Load() + if current <= observed || maxActive.CompareAndSwap(observed, current) { + break + } + } + time.Sleep(10 * time.Millisecond) + active.Add(-1) + return true, nil + } + + start := make(chan struct{}) + var runs sync.WaitGroup + runs.Add(2) + for range 2 { + go func() { + defer runs.Done() + <-start + errs := g.verifyDownloadedFilesWithValidator(uris, "/models", validate) + assert.Empty(t, errs) + }() + } + close(start) + runs.Wait() + + assert.Equal(t, int32(3), maxActive.Load()) +} + +func TestVerifyDownloadedFilesWithValidatorReportsFailures(t *testing.T) { + g := &Gopher{} + WithModelVerificationConcurrency(2)(g) + uris := []ociobjectstore.ObjectURI{ + {ObjectName: "model/valid", Prefix: "model"}, + {ObjectName: "model/mismatch", Prefix: "model"}, + {ObjectName: "model/error", Prefix: "model"}, + } + + errs := g.verifyDownloadedFilesWithValidator(uris, "/models", func(obj ociobjectstore.ObjectURI, _ string) (bool, error) { + switch obj.ObjectName { + case "model/mismatch": + return false, nil + case "model/error": + return false, errors.New("read failed") + default: + return true, nil + } + }) + + require.Len(t, errs, 2) + assert.ErrorContains(t, errs["model/mismatch"], "MD5 or size mismatch") + assert.EqualError(t, errs["model/error"], "read failed") +} + func TestCaptureStartupReadyModelsCapturesOnlyReadyEntries(t *testing.T) { readyKey := constants.GetModelConfigMapKey("service-ns", "ready-model", false) updatingKey := constants.GetModelConfigMapKey("service-ns", "updating-model", false) diff --git a/pkg/modelagent/verification_limiter.go b/pkg/modelagent/verification_limiter.go new file mode 100644 index 000000000..a1c610a37 --- /dev/null +++ b/pkg/modelagent/verification_limiter.go @@ -0,0 +1,26 @@ +package modelagent + +// verificationLimiter bounds concurrent file integrity checks across all +// model downloads handled by one model-agent process. +type verificationLimiter struct { + permits chan struct{} +} + +func newVerificationLimiter(concurrency int) *verificationLimiter { + if concurrency < 1 { + concurrency = 1 + } + return &verificationLimiter{permits: make(chan struct{}, concurrency)} +} + +func (l *verificationLimiter) acquire() { + l.permits <- struct{}{} +} + +func (l *verificationLimiter) release() { + <-l.permits +} + +func (l *verificationLimiter) limit() int { + return cap(l.permits) +}