Fix deletion of model in cases of ModeFailed and ScheduleFailed with server disconnects

scheduler/pkg/agent/server.go

@@ -465,6 +465,14 @@ func (s *Server) removeServerReplicaImpl(serverName string, serverReplicaIdx int
 			s.logger.Debugf("Failed to reschedule model %s when server %s replica %d disconnected", modelName, serverName, serverReplicaIdx)
 		}
 	}
+	// retry failed models
+	// this is perhaps counterintuitive, but we want to retry failed models on other servers
+	// specifically in the case of model state `LoadFailed` and the server replica disconnects, we want to reconcile
+	// the model state with th new set of active servers
+	// note that this will also retry `ScheduleFailed`, which is a side effect of calling `ScheduleFailedModels`
+	if _, err := s.scheduler.ScheduleFailedModels(); err != nil {
+		s.logger.WithError(err).Errorf("Failed to reschedule failed models when server %s replica %d disconnected", serverName, serverReplicaIdx)
+	}
 }
 
 func (s *Server) drainServerReplicaImpl(serverName string, serverReplicaIdx int) {

scheduler/pkg/scheduler/scheduler.go

@@ -134,7 +134,7 @@ func (s *SimpleScheduler) scheduleToServer(modelName string) error {
 		logger.Debugf("Model %s is deleted ensuring removed", modelName)
 		err = s.store.UpdateLoadedModels(modelName, latestModel.GetVersion(), server, []*store.ServerReplica{})
 		if err != nil {
-			logger.Warnf("Failed to unschedule model replicas for model %s on server %s", modelName, server)
+			logger.WithError(err).Warnf("Failed to unschedule model replicas for model %s on server %s", modelName, server)
 		}
 
 	} else {

scheduler/pkg/store/memory.go

@@ -511,7 +511,7 @@ func (m *MemoryStore) updateModelStateImpl(
 			modelKey, version, serverKey, replicaIdx, desiredState.String(),
 		)
 
-		// Update models loaded onto replica if loaded or unloaded is state
+		// Update models loaded onto replica for relevant state
 		if desiredState == Loaded || desiredState == Loading || desiredState == Unloaded || desiredState == LoadFailed {
 			server, ok := m.store.servers[serverKey]
 			if ok {

scheduler/pkg/store/memory_status.go

@@ -118,7 +118,8 @@ func (m *MemoryStore) FailedScheduling(modelVersion *ModelVersion, reason string
 		AvailableReplicas:   modelVersion.state.AvailableReplicas,
 		UnavailableReplicas: modelVersion.GetModel().GetDeploymentSpec().GetReplicas() - modelVersion.state.AvailableReplicas,
 	}
-
+	// make sure we reset server
+	modelVersion.server = ""
 	m.eventHub.PublishModelEvent(
 		modelFailureEventSource,
 		coordinator.ModelEventMsg{