feat(catalog): remove orphaned models (kubeflow#1962)

* feat(catalog): remove orphaned models

Remove models from the database that are no longer referenced. This can
happen because a model is removed or excluded from a source, or the
source can be deleted or disabled.

Signed-off-by: Paul Boyd <paul@pboyd.io>

* fix(catalog): race condition causing empty models list

Signed-off-by: Paul Boyd <paul@pboyd.io>

---------

Signed-off-by: Paul Boyd <paul@pboyd.io>

Author: pboyd

catalog/internal/catalog/catalog_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"reflect"
 	"sort"
+	"sync"
 	"testing"
 	"time"
 
@@ -376,6 +377,9 @@ func TestLoadCatalogSourcesWithMockRepositories(t *testing.T) {
 	// Wait a bit for the goroutine to process
 	time.Sleep(100 * time.Millisecond)
 
+	mockModelRepo.mu.RLock()
+	defer mockModelRepo.mu.RUnlock()
+
 	// Verify that the model was saved
 	if len(mockModelRepo.SavedModels) != 1 {
 		t.Errorf("Expected 1 model to be saved, got %d", len(mockModelRepo.SavedModels))
@@ -390,11 +394,17 @@ func TestLoadCatalogSourcesWithMockRepositories(t *testing.T) {
 		}
 	}
 
+	mockModelArtifactRepo.mu.RLock()
+	defer mockModelArtifactRepo.mu.RUnlock()
+
 	// Verify that artifacts were saved
 	if len(mockModelArtifactRepo.SavedArtifacts) != 1 {
 		t.Errorf("Expected 1 model artifact to be saved, got %d", len(mockModelArtifactRepo.SavedArtifacts))
 	}
 
+	mockMetricsArtifactRepo.mu.RLock()
+	defer mockMetricsArtifactRepo.mu.RUnlock()
+
 	if len(mockMetricsArtifactRepo.SavedMetrics) != 1 {
 		t.Errorf("Expected 1 metrics artifact to be saved, got %d", len(mockMetricsArtifactRepo.SavedMetrics))
 	}
@@ -470,6 +480,9 @@ func TestLoadCatalogSourcesWithRepositoryErrors(t *testing.T) {
 	// Wait for processing
 	time.Sleep(100 * time.Millisecond)
 
+	mockModelRepo.mu.RLock()
+	defer mockModelRepo.mu.RUnlock()
+
 	// Verify that no models were saved due to the error
 	if len(mockModelRepo.SavedModels) != 0 {
 		t.Errorf("Expected 0 models to be saved due to error, got %d", len(mockModelRepo.SavedModels))
@@ -638,11 +651,14 @@ func (m *MockCatalogModelRepositoryWithErrors) Save(model dbmodels.CatalogModel)
 
 // MockCatalogModelRepository mocks the CatalogModelRepository interface.
 type MockCatalogModelRepository struct {
+	mu          sync.RWMutex
 	SavedModels []dbmodels.CatalogModel
 	NextID      int32
 }
 
 func (m *MockCatalogModelRepository) GetByID(id int32) (dbmodels.CatalogModel, error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	for _, model := range m.SavedModels {
 		if model.GetID() != nil && *model.GetID() == id {
 			return model, nil
@@ -652,6 +668,8 @@ func (m *MockCatalogModelRepository) GetByID(id int32) (dbmodels.CatalogModel, e
 }
 
 func (m *MockCatalogModelRepository) List(listOptions dbmodels.CatalogModelListOptions) (*mrmodels.ListWrapper[dbmodels.CatalogModel], error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	return &mrmodels.ListWrapper[dbmodels.CatalogModel]{
 		Items:    m.SavedModels,
 		PageSize: int32(len(m.SavedModels)),
@@ -660,6 +678,8 @@ func (m *MockCatalogModelRepository) List(listOptions dbmodels.CatalogModelListO
 }
 
 func (m *MockCatalogModelRepository) GetByName(name string) (dbmodels.CatalogModel, error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	for _, model := range m.SavedModels {
 		if model.GetAttributes() != nil && model.GetAttributes().Name != nil && *model.GetAttributes().Name == name {
 			return model, nil
@@ -669,6 +689,9 @@ func (m *MockCatalogModelRepository) GetByName(name string) (dbmodels.CatalogMod
 }
 
 func (m *MockCatalogModelRepository) Save(model dbmodels.CatalogModel) (dbmodels.CatalogModel, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	m.NextID++
 	id := m.NextID
 
@@ -685,13 +708,31 @@ func (m *MockCatalogModelRepository) Save(model dbmodels.CatalogModel) (dbmodels
 	return savedModel, nil
 }
 
+func (m *MockCatalogModelRepository) DeleteBySource(sourceID string) error {
+	// Mock implementation - no-op for testing
+	return nil
+}
+
+func (m *MockCatalogModelRepository) DeleteByID(id int32) error {
+	// Mock implementation - no-op for testing
+	return nil
+}
+
+func (m *MockCatalogModelRepository) GetDistinctSourceIDs() ([]string, error) {
+	// Mock implementation - return empty list by default
+	return []string{}, nil
+}
+
 // MockCatalogModelArtifactRepository mocks the CatalogModelArtifactRepository interface.
 type MockCatalogModelArtifactRepository struct {
+	mu             sync.RWMutex
 	SavedArtifacts []dbmodels.CatalogModelArtifact
 	NextID         int32
 }
 
 func (m *MockCatalogModelArtifactRepository) GetByID(id int32) (dbmodels.CatalogModelArtifact, error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	for _, artifact := range m.SavedArtifacts {
 		if artifact.GetID() != nil && *artifact.GetID() == id {
 			return artifact, nil
@@ -701,6 +742,8 @@ func (m *MockCatalogModelArtifactRepository) GetByID(id int32) (dbmodels.Catalog
 }
 
 func (m *MockCatalogModelArtifactRepository) List(listOptions dbmodels.CatalogModelArtifactListOptions) (*mrmodels.ListWrapper[dbmodels.CatalogModelArtifact], error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	return &mrmodels.ListWrapper[dbmodels.CatalogModelArtifact]{
 		Items:    m.SavedArtifacts,
 		PageSize: int32(len(m.SavedArtifacts)),
@@ -709,6 +752,9 @@ func (m *MockCatalogModelArtifactRepository) List(listOptions dbmodels.CatalogMo
 }
 
 func (m *MockCatalogModelArtifactRepository) Save(modelArtifact dbmodels.CatalogModelArtifact, parentResourceID *int32) (dbmodels.CatalogModelArtifact, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	m.NextID++
 	id := m.NextID
 
@@ -727,11 +773,15 @@ func (m *MockCatalogModelArtifactRepository) Save(modelArtifact dbmodels.Catalog
 
 // MockCatalogMetricsArtifactRepository mocks the CatalogMetricsArtifactRepository interface.
 type MockCatalogMetricsArtifactRepository struct {
+	mu           sync.RWMutex
 	SavedMetrics []dbmodels.CatalogMetricsArtifact
 	NextID       int32
 }
 
 func (m *MockCatalogMetricsArtifactRepository) GetByID(id int32) (dbmodels.CatalogMetricsArtifact, error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
 	for _, metrics := range m.SavedMetrics {
 		if metrics.GetID() != nil && *metrics.GetID() == id {
 			return metrics, nil
@@ -741,6 +791,9 @@ func (m *MockCatalogMetricsArtifactRepository) GetByID(id int32) (dbmodels.Catal
 }
 
 func (m *MockCatalogMetricsArtifactRepository) List(listOptions dbmodels.CatalogMetricsArtifactListOptions) (*mrmodels.ListWrapper[dbmodels.CatalogMetricsArtifact], error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
 	return &mrmodels.ListWrapper[dbmodels.CatalogMetricsArtifact]{
 		Items:    m.SavedMetrics,
 		PageSize: int32(len(m.SavedMetrics)),
@@ -749,6 +802,9 @@ func (m *MockCatalogMetricsArtifactRepository) List(listOptions dbmodels.Catalog
 }
 
 func (m *MockCatalogMetricsArtifactRepository) Save(metricsArtifact dbmodels.CatalogMetricsArtifact, parentResourceID *int32) (dbmodels.CatalogMetricsArtifact, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	m.NextID++
 	id := m.NextID
 
@@ -766,6 +822,9 @@ func (m *MockCatalogMetricsArtifactRepository) Save(metricsArtifact dbmodels.Cat
 }
 
 func (m *MockCatalogMetricsArtifactRepository) BatchSave(metricsArtifacts []dbmodels.CatalogMetricsArtifact, parentResourceID *int32) ([]dbmodels.CatalogMetricsArtifact, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	savedArtifacts := make([]dbmodels.CatalogMetricsArtifact, len(metricsArtifacts))
 
 	for i, metricsArtifact := range metricsArtifacts {
@@ -790,11 +849,14 @@ func (m *MockCatalogMetricsArtifactRepository) BatchSave(metricsArtifacts []dbmo
 
 // MockCatalogArtifactRepository mocks the CatalogArtifactRepository interface.
 type MockCatalogArtifactRepository struct {
+	mu             sync.RWMutex
 	SavedArtifacts []dbmodels.CatalogArtifact
 	NextID         int32
 }
 
 func (m *MockCatalogArtifactRepository) GetByID(id int32) (dbmodels.CatalogArtifact, error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	for _, artifact := range m.SavedArtifacts {
 		// Check both model and metrics artifacts for the ID
 		if artifact.CatalogModelArtifact != nil && artifact.CatalogModelArtifact.GetID() != nil && *artifact.CatalogModelArtifact.GetID() == id {
@@ -808,6 +870,8 @@ func (m *MockCatalogArtifactRepository) GetByID(id int32) (dbmodels.CatalogArtif
 }
 
 func (m *MockCatalogArtifactRepository) List(listOptions dbmodels.CatalogArtifactListOptions) (*mrmodels.ListWrapper[dbmodels.CatalogArtifact], error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	return &mrmodels.ListWrapper[dbmodels.CatalogArtifact]{
 		Items:    m.SavedArtifacts,
 		PageSize: int32(len(m.SavedArtifacts)),

catalog/internal/catalog/hf_catalog.go

@@ -577,6 +577,12 @@ func (p *hfModelProvider) emit(ctx context.Context, models []ModelProviderRecord
 			return
 		}
 	}
+
+	// Send an empty record to indicate that we're done with the batch.
+	select {
+	case out <- ModelProviderRecord{}:
+	case <-done:
+	}
 }
 
 // validateCredentials checks if the HuggingFace API key credentials are valid

catalog/internal/catalog/loader.go

@@ -8,6 +8,7 @@ import (
 	"path/filepath"
 	"sync"
 
+	mapset "github.com/deckarep/golang-set/v2"
 	"github.com/golang/glog"
 	dbmodels "github.com/kubeflow/model-registry/catalog/internal/db/models"
 	"github.com/kubeflow/model-registry/catalog/internal/db/service"
@@ -26,6 +27,9 @@ type ModelProviderRecord struct {
 // expected to spawn a goroutine and return immediately. The returned channel must
 // close when the goroutine ends. The goroutine should end when the context is
 // canceled, but may end sooner.
+//
+// The function may emit a record with a nil Model to indicate that the
+// complete set of models has been sent.
 type ModelProviderFunc func(ctx context.Context, source *Source, reldir string) (<-chan ModelProviderRecord, error)
 
 var registeredModelProviders = map[string]ModelProviderFunc{}
@@ -120,8 +124,14 @@ func (l *Loader) Start(ctx context.Context) error {
 		}
 	}
 
+	// Delete models from unknown or disabled sources
+	err := l.removeModelsFromMissingSources()
+	if err != nil {
+		return fmt.Errorf("faied to remove models from missing sources: %w", err)
+	}
+
 	// Phase 2: Load models from merged sources (once, after all merging is complete)
-	err := l.loadAllModels(ctx)
+	err = l.loadAllModels(ctx)
 	if err != nil {
 		return err
 	}
@@ -269,6 +279,9 @@ func (l *Loader) updateDatabase(ctx context.Context) error {
 
 	go func() {
 		for record := range records {
+			if record.Model == nil {
+				continue
+			}
 			attr := record.Model.GetAttributes()
 			if attr == nil || attr.Name == nil {
 				continue
@@ -377,7 +390,30 @@ func (l *Loader) readProviderRecords(ctx context.Context) <-chan ModelProviderRe
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
+
+			modelNames := []string{}
+
 			for r := range records {
+				if r.Model == nil {
+					glog.V(2).Infof("%s: trigger cleanup", source.Id)
+
+					// Copy the list of model names, then clear it.
+					modelNameSet := mapset.NewSet(modelNames...)
+					modelNames = modelNames[:0]
+
+					go func() {
+						err := l.removeOrphanedModelsFromSource(source.Id, modelNameSet)
+						if err != nil {
+							glog.Errorf("error removing orphaned models: %v", err)
+						}
+					}()
+					continue
+				}
+
+				if attr := r.Model.GetAttributes(); attr != nil && attr.Name != nil {
+					modelNames = append(modelNames, *attr.Name)
+				}
+
 				// Set source_id on every returned model.
 				l.setModelSourceID(r.Model, source.Id)
 
@@ -424,3 +460,57 @@ func (l *Loader) setModelSourceID(model dbmodels.CatalogModel, sourceID string)
 
 	*props = append(*props, mrmodels.NewStringProperty("source_id", sourceID, false))
 }
+
+func (l *Loader) removeModelsFromMissingSources() error {
+	enabledSourceIDs := mapset.NewSet[string]()
+	for id, source := range l.Sources.AllSources() {
+		if source.Enabled == nil || *source.Enabled {
+			enabledSourceIDs.Add(id)
+		}
+	}
+
+	existingSourceIDs, err := l.services.CatalogModelRepository.GetDistinctSourceIDs()
+	if err != nil {
+		return fmt.Errorf("unable to retrieve existing source IDs: %w", err)
+	}
+
+	for oldSource := range mapset.NewSet(existingSourceIDs...).Difference(enabledSourceIDs).Iter() {
+		glog.Infof("Removing models from source %s", oldSource)
+
+		err = l.services.CatalogModelRepository.DeleteBySource(oldSource)
+		if err != nil {
+			return fmt.Errorf("unable to remove models from source %q: %w", oldSource, err)
+		}
+	}
+
+	return nil
+}
+
+func (l *Loader) removeOrphanedModelsFromSource(sourceID string, valid mapset.Set[string]) error {
+	list, err := l.services.CatalogModelRepository.List(dbmodels.CatalogModelListOptions{
+		SourceIDs: &[]string{sourceID},
+	})
+	if err != nil {
+		return fmt.Errorf("unable to list models from source %q: %w", sourceID, err)
+	}
+
+	for _, model := range list.Items {
+		attr := model.GetAttributes()
+		if attr == nil || attr.Name == nil || model.GetID() == nil {
+			continue
+		}
+
+		if valid.Contains(*attr.Name) {
+			continue
+		}
+
+		glog.Infof("Removing %s model %s", sourceID, *attr.Name)
+
+		err = l.services.CatalogModelRepository.DeleteByID(*model.GetID())
+		if err != nil {
+			return fmt.Errorf("unable to remove model %d (%s from source %s): %w", *model.GetID(), *attr.Name, sourceID, err)
+		}
+	}
+
+	return nil
+}