[Core] Prioritize real downloads over validation replay

pkg/modelagent/configmap_reconciler.go

@@ -712,6 +712,10 @@ Returns:
 func (c *ConfigMapReconciler) getConfigMap(ctx context.Context) (*corev1.ConfigMap, error) {
 	existingConfigMap, err := c.kubeClient.CoreV1().ConfigMaps(c.namespace).Get(ctx, c.nodeName, metav1.GetOptions{})
 	if err != nil {
+		if errors.IsNotFound(err) {
+			c.logger.Infof("Node %s configmap does not exist yet", c.nodeName)
+			return nil, err
+		}
 		c.logger.Errorf("Failed retrieve node %s configmap: %v", c.nodeName, err)
 		return nil, err
 	}

pkg/modelagent/gopher.go

@@ -46,6 +46,7 @@ type GopherTask struct {
 	TensorRTLLMShapeFilter *TensorRTLLMShapeFilter
 	SamePathWaitStartedAt  time.Time
 	NormalPriorityOnly     bool
+	RevalidationReplay     bool
 }
 
 type activeDownload struct {
@@ -77,13 +78,16 @@ type Gopher struct {
 	taskQueue           *gopherTaskQueue
 	samePathWaitDelay   time.Duration
 	samePathWaitTimeout time.Duration
+
+	startupReadyModelKeys map[string]struct{}
 }
 
 const (
 	BigFileSizeInMB = 200
 
-	defaultSamePathWaitDelay   = 30 * time.Second
-	defaultSamePathWaitTimeout = 30 * time.Minute
+	defaultSamePathWaitDelay           = 30 * time.Second
+	defaultSamePathWaitTimeout         = 30 * time.Minute
+	defaultStartupReadySnapshotTimeout = 5 * time.Second
 )
 
 func NewGopher(
@@ -133,6 +137,10 @@ func NewGopher(
 }
 
 func (s *Gopher) Run(stopCh <-chan struct{}, numWorker int, numHighPriorityWorker int) {
+	startupSnapshotCtx, cancelStartupSnapshot := context.WithTimeout(context.Background(), defaultStartupReadySnapshotTimeout)
+	defer cancelStartupSnapshot()
+	s.captureStartupReadyModels(startupSnapshotCtx)
+
 	// Start the ConfigMap reconciliation service
 	s.configMapReconciler.StartReconciliation()
 	s.logger.Info("Started ConfigMap reconciliation service")
@@ -192,6 +200,8 @@ func (s *Gopher) enqueueTask(task *GopherTask) {
 	}
 	if task.TaskType == Delete {
 		s.cancelActiveDownload(task)
+	} else {
+		s.classifyStartupRevalidation(task)
 	}
 	s.taskQueue.enqueue(task)
 }
@@ -679,10 +689,78 @@ func (s *Gopher) demoteToNormalPriority(task *GopherTask) {
 		return
 	}
 	task.NormalPriorityOnly = true
+	s.classifyStartupRevalidation(task)
 	s.logger.Infof("Demoting %s to normal priority for fallback download/validation", getModelInfoForLogging(task))
 	s.enqueueTask(task)
 }
 
+func (s *Gopher) classifyStartupRevalidation(task *GopherTask) bool {
+	if task == nil || task.RevalidationReplay || task.TaskType != Download {
+		return false
+	}
+	if !s.isStartupRevalidation(task) {
+		return false
+	}
+	task.NormalPriorityOnly = true
+	task.RevalidationReplay = true
+	return true
+}
+
+func (s *Gopher) captureStartupReadyModels(ctx context.Context) {
+	if s.configMapReconciler == nil {
+		return
+	}
+	configMap, err := s.configMapReconciler.getConfigMap(ctx)
+	if err != nil {
+		if apierrors.IsNotFound(err) {
+			s.logger.Infof("No startup Ready model snapshot because node ConfigMap does not exist yet")
+			s.startupReadyModelKeys = map[string]struct{}{}
+			return
+		}
+		s.logger.Warnf("Cannot capture startup Ready model snapshot: %v", err)
+		s.startupReadyModelKeys = map[string]struct{}{}
+		return
+	}
+	readyModelKeys := make(map[string]struct{})
+	for key, data := range configMap.Data {
+		if hasModelEntryStatus(data, ModelStatusReady) {
+			readyModelKeys[key] = struct{}{}
+		}
+	}
+	s.startupReadyModelKeys = readyModelKeys
+	s.logger.Infof("Captured %d Ready models from startup ConfigMap snapshot", len(readyModelKeys))
+}
+
+func (s *Gopher) isStartupRevalidation(task *GopherTask) bool {
+	if len(s.startupReadyModelKeys) == 0 {
+		return false
+	}
+	modelKey := getModelID(task.BaseModel, task.ClusterBaseModel)
+	if _, wasReady := s.startupReadyModelKeys[modelKey]; !wasReady {
+		return false
+	}
+
+	var baseModelSpec v1beta1.BaseModelSpec
+	if task.BaseModel != nil {
+		baseModelSpec = task.BaseModel.Spec
+	} else if task.ClusterBaseModel != nil {
+		baseModelSpec = task.ClusterBaseModel.Spec
+	} else {
+		return false
+	}
+	if baseModelSpec.Storage == nil || baseModelSpec.Storage.StorageUri == nil || baseModelSpec.Storage.Path == nil || *baseModelSpec.Storage.Path == "" {
+		return false
+	}
+	storageType, err := storage.GetStorageType(*baseModelSpec.Storage.StorageUri)
+	if err != nil || storageType != storage.StorageTypeOCI {
+		return false
+	}
+
+	destPath := getDestPath(&baseModelSpec, s.modelRootDir)
+	fileInfo, err := os.Stat(destPath)
+	return err == nil && fileInfo.IsDir()
+}
+
 func shouldUseSamePathObjectStorageReuse(task *GopherTask) bool {
 	return task != nil && task.TaskType == Download
 }
@@ -1101,6 +1179,21 @@ func sameModelStoragePath(currentStorage *v1beta1.StorageSpec, candidateStorage
 	return getDestPath(&candidateSpec, modelRootDir) == destPath
 }
 
+func filterObjectStorageObjectsForTask(objects []objectstorage.ObjectSummary, task *GopherTask) []objectstorage.ObjectSummary {
+	if task == nil || task.TensorRTLLMShapeFilter == nil ||
+		!task.TensorRTLLMShapeFilter.IsTensorrtLLMModel ||
+		task.TensorRTLLMShapeFilter.ModelType != string(constants.ServingBaseModel) {
+		return objects
+	}
+	shapeFilteredObjects := make([]objectstorage.ObjectSummary, 0)
+	for _, object := range objects {
+		if object.Name != nil && strings.Contains(*object.Name, fmt.Sprintf("/%s/", task.TensorRTLLMShapeFilter.ShapeAlias)) {
+			shapeFilteredObjects = append(shapeFilteredObjects, object)
+		}
+	}
+	return shapeFilteredObjects
+}
+
 func (s *Gopher) downloadModel(ctx context.Context, uri *ociobjectstore.ObjectURI, destPath string, task *GopherTask) error {
 	startTime := time.Now()
 	defer func() {
@@ -1146,15 +1239,7 @@ func (s *Gopher) downloadModel(ctx context.Context, uri *ociobjectstore.ObjectUR
 	// Shape filtering for TensorRTLLM
 	if task.TensorRTLLMShapeFilter != nil && task.TensorRTLLMShapeFilter.IsTensorrtLLMModel && task.TensorRTLLMShapeFilter.ModelType == string(constants.ServingBaseModel) {
 		s.logger.Infof("TensorRTLLM Serving model detected. Start filtering model files that doesn't belong to the node shape %s in model bucket folder", task.TensorRTLLMShapeFilter.ShapeAlias)
-		shapeFilteredObjects := make([]objectstorage.ObjectSummary, 0)
-		for _, object := range objects {
-			if object.Name != nil {
-				if strings.Contains(*object.Name, fmt.Sprintf("/%s/", task.TensorRTLLMShapeFilter.ShapeAlias)) {
-					shapeFilteredObjects = append(shapeFilteredObjects, object)
-				}
-			}
-		}
-		objects = shapeFilteredObjects
+		objects = filterObjectStorageObjectsForTask(objects, task)
 
 		if len(objects) == 0 {
 			return fmt.Errorf("no suitable objects found for shape %s", task.TensorRTLLMShapeFilter.ShapeAlias)

pkg/modelagent/gopher_task_queue.go

@@ -8,11 +8,12 @@ import (
 )
 
 type gopherTaskQueue struct {
-	mutex  sync.Mutex
-	cond   *sync.Cond
-	high   []*GopherTask
-	normal []*GopherTask
-	closed bool
+	mutex              sync.Mutex
+	cond               *sync.Cond
+	high               []*GopherTask
+	normalDownload     []*GopherTask
+	normalRevalidation []*GopherTask
+	closed             bool
 }
 
 func newGopherTaskQueue() *gopherTaskQueue {
@@ -34,25 +35,33 @@ func (q *gopherTaskQueue) enqueue(task *GopherTask) {
 		// Delete preempts pending work for the same model and should run before
 		// reuse-wait tasks, so it is the only non-FIFO insertion.
 		q.high = removeSupersededTasks(q.high, task)
-		q.normal = removeSupersededTasks(q.normal, task)
+		q.normalDownload = removeSupersededTasks(q.normalDownload, task)
+		q.normalRevalidation = removeSupersededTasks(q.normalRevalidation, task)
 		q.high = append([]*GopherTask{task}, q.high...)
 	} else if shouldUseHighPriorityQueue(task) {
 		q.high = append(q.high, task)
+	} else if task.RevalidationReplay {
+		q.normalRevalidation = append(q.normalRevalidation, task)
 	} else {
-		q.normal = append(q.normal, task)
+		q.normalDownload = append(q.normalDownload, task)
 	}
 	q.cond.Broadcast()
 }
 
 func (q *gopherTaskQueue) popNormal() (*GopherTask, bool) {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
-	for len(q.normal) == 0 && !q.closed {
+	for len(q.normalDownload) == 0 && len(q.normalRevalidation) == 0 && !q.closed {
 		q.cond.Wait()
 	}
-	if len(q.normal) > 0 {
-		task := q.normal[0]
-		q.normal = q.normal[1:]
+	if len(q.normalDownload) > 0 {
+		task := q.normalDownload[0]
+		q.normalDownload = q.normalDownload[1:]
+		return task, true
+	}
+	if len(q.normalRevalidation) > 0 {
+		task := q.normalRevalidation[0]
+		q.normalRevalidation = q.normalRevalidation[1:]
 		return task, true
 	}
 	return nil, false
@@ -82,15 +91,15 @@ func (q *gopherTaskQueue) close() {
 func (q *gopherTaskQueue) len() int {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
-	return len(q.high) + len(q.normal)
+	return len(q.high) + len(q.normalDownload) + len(q.normalRevalidation)
 }
 
 func shouldUseHighPriorityQueue(task *GopherTask) bool {
 	return task.TaskType == Delete || isObjectStorageDownloadTask(task) || (!task.NormalPriorityOnly && !task.SamePathWaitStartedAt.IsZero())
 }
 
 func isObjectStorageDownloadTask(task *GopherTask) bool {
-	if task == nil || task.TaskType != Download || task.NormalPriorityOnly {
+	if task == nil || task.TaskType != Download || task.NormalPriorityOnly || task.RevalidationReplay {
 		return false
 	}
 	var storageSpec *v1beta1.StorageSpec

pkg/modelagent/gopher_task_queue_test.go

@@ -103,6 +103,35 @@ func TestGopherTaskQueueKeepsDownloadOverrideNormal(t *testing.T) {
 	assert.Equal(t, "oci-model", queued.BaseModel.Name)
 }
 
+func TestGopherTaskQueuePrioritizesNormalDownloadBeforeRevalidationReplay(t *testing.T) {
+	queue := newGopherTaskQueue()
+	validation := &GopherTask{
+		TaskType:           Download,
+		NormalPriorityOnly: true,
+		RevalidationReplay: true,
+		BaseModel: &v1beta1.BaseModel{
+			ObjectMeta: metav1.ObjectMeta{Name: "validation", Namespace: "service-ns", UID: "validation-uid"},
+		},
+	}
+	download := &GopherTask{
+		TaskType:           Download,
+		NormalPriorityOnly: true,
+		BaseModel: &v1beta1.BaseModel{
+			ObjectMeta: metav1.ObjectMeta{Name: "download", Namespace: "service-ns", UID: "download-uid"},
+		},
+	}
+
+	queue.enqueue(validation)
+	queue.enqueue(download)
+
+	task, ok := queue.popNormal()
+	require.True(t, ok)
+	assert.Equal(t, "download", task.BaseModel.Name)
+	task, ok = queue.popNormal()
+	require.True(t, ok)
+	assert.Equal(t, "validation", task.BaseModel.Name)
+}
+
 func TestGopherTaskQueueDeleteSupersedesPendingDownloadsForSameModel(t *testing.T) {
 	queue := newGopherTaskQueue()
 	model := &v1beta1.BaseModel{
@@ -203,6 +232,26 @@ func TestGopherTaskQueueDemotedSamePathWaitUsesNormalQueue(t *testing.T) {
 	assert.Equal(t, "demoted", task.BaseModel.Name)
 }
 
+func TestGopherTaskQueueDeleteSupersedesPendingRevalidationReplayForSameModel(t *testing.T) {
+	queue := newGopherTaskQueue()
+	model := &v1beta1.BaseModel{
+		ObjectMeta: metav1.ObjectMeta{Name: "model", Namespace: "service-ns", UID: "model-uid"},
+	}
+
+	queue.enqueue(&GopherTask{
+		TaskType:           Download,
+		BaseModel:          model,
+		NormalPriorityOnly: true,
+		RevalidationReplay: true,
+	})
+	queue.enqueue(&GopherTask{TaskType: Delete, BaseModel: model})
+
+	task, ok := queue.popHighPriority()
+	require.True(t, ok)
+	assert.Equal(t, Delete, task.TaskType)
+	assert.Equal(t, 0, queue.len())
+}
+
 func TestGopherTaskQueueEnqueueWakesMatchingBlockedWorker(t *testing.T) {
 	queue := newGopherTaskQueue()
 	normalDone := make(chan struct{})