diff --git a/core/gallery/models.go b/core/gallery/models.go index 6b20ad7b2b40..953ef767ffed 100644 --- a/core/gallery/models.go +++ b/core/gallery/models.go @@ -6,6 +6,7 @@ import ( "fmt" "os" "path/filepath" + "slices" "strings" "dario.cat/mergo" @@ -293,13 +294,24 @@ func GetLocalModelConfiguration(basePath string, name string) (*ModelConfig, err return ReadConfigFile[ModelConfig](galleryFile) } -func DeleteModelFromSystem(systemState *system.SystemState, name string) error { - additionalFiles := []string{} +func listModelFiles(systemState *system.SystemState, name string) ([]string, error) { configFile := filepath.Join(systemState.Model.ModelsPath, fmt.Sprintf("%s.yaml", name)) if err := utils.VerifyPath(configFile, systemState.Model.ModelsPath); err != nil { - return fmt.Errorf("failed to verify path %s: %w", configFile, err) + return nil, fmt.Errorf("failed to verify path %s: %w", configFile, err) + } + + // os.PathSeparator is not allowed in model names. Replace them with "__" to avoid conflicts with file paths. + name = strings.ReplaceAll(name, string(os.PathSeparator), "__") + + galleryFile := filepath.Join(systemState.Model.ModelsPath, galleryFileName(name)) + if err := utils.VerifyPath(galleryFile, systemState.Model.ModelsPath); err != nil { + return nil, fmt.Errorf("failed to verify path %s: %w", galleryFile, err) } + + additionalFiles := []string{} + allFiles := []string{} + // Galleryname is the name of the model in this case dat, err := os.ReadFile(configFile) if err == nil { @@ -307,7 +319,7 @@ func DeleteModelFromSystem(systemState *system.SystemState, name string) error { err = yaml.Unmarshal(dat, &modelConfig) if err != nil { - return err + return nil, err } if modelConfig.Model != "" { additionalFiles = append(additionalFiles, modelConfig.ModelFileName()) @@ -318,26 +330,15 @@ func DeleteModelFromSystem(systemState *system.SystemState, name string) error { } } - // os.PathSeparator is not allowed in model names. Replace them with "__" to avoid conflicts with file paths. - name = strings.ReplaceAll(name, string(os.PathSeparator), "__") - - galleryFile := filepath.Join(systemState.Model.ModelsPath, galleryFileName(name)) - if err := utils.VerifyPath(galleryFile, systemState.Model.ModelsPath); err != nil { - return fmt.Errorf("failed to verify path %s: %w", galleryFile, err) - } - - var filesToRemove []string - - // Delete all the files associated to the model // read the model config galleryconfig, err := ReadConfigFile[ModelConfig](galleryFile) if err == nil && galleryconfig != nil { for _, f := range galleryconfig.Files { fullPath := filepath.Join(systemState.Model.ModelsPath, f.Filename) if err := utils.VerifyPath(fullPath, systemState.Model.ModelsPath); err != nil { - return fmt.Errorf("failed to verify path %s: %w", fullPath, err) + return allFiles, fmt.Errorf("failed to verify path %s: %w", fullPath, err) } - filesToRemove = append(filesToRemove, fullPath) + allFiles = append(allFiles, fullPath) } } else { log.Error().Err(err).Msgf("failed to read gallery file %s", configFile) @@ -346,18 +347,68 @@ func DeleteModelFromSystem(systemState *system.SystemState, name string) error { for _, f := range additionalFiles { fullPath := filepath.Join(filepath.Join(systemState.Model.ModelsPath, f)) if err := utils.VerifyPath(fullPath, systemState.Model.ModelsPath); err != nil { - return fmt.Errorf("failed to verify path %s: %w", fullPath, err) + return allFiles, fmt.Errorf("failed to verify path %s: %w", fullPath, err) } - filesToRemove = append(filesToRemove, fullPath) + allFiles = append(allFiles, fullPath) } - filesToRemove = append(filesToRemove, galleryFile) + allFiles = append(allFiles, galleryFile) // skip duplicates - filesToRemove = utils.Unique(filesToRemove) + allFiles = utils.Unique(allFiles) + + return allFiles, nil +} + +func DeleteModelFromSystem(systemState *system.SystemState, name string) error { + configFile := filepath.Join(systemState.Model.ModelsPath, fmt.Sprintf("%s.yaml", name)) + + filesToRemove, err := listModelFiles(systemState, name) + if err != nil { + return err + } + + allOtherFiles := []string{} + // Get all files of all other models + fi, err := os.ReadDir(systemState.Model.ModelsPath) + if err != nil { + return err + } + for _, f := range fi { + if f.IsDir() { + continue + } + if strings.HasPrefix(f.Name(), "._gallery_") { + continue + } + if !strings.HasSuffix(f.Name(), ".yaml") && !strings.HasSuffix(f.Name(), ".yml") { + continue + } + if f.Name() == fmt.Sprintf("%s.yaml", name) || f.Name() == fmt.Sprintf("%s.yml", name) { + continue + } + + name := strings.TrimSuffix(f.Name(), ".yaml") + name = strings.TrimSuffix(name, ".yml") + + log.Debug().Msgf("Checking file %s", f.Name()) + files, err := listModelFiles(systemState, name) + if err != nil { + log.Debug().Err(err).Msgf("failed to list files for model %s", f.Name()) + continue + } + allOtherFiles = append(allOtherFiles, files...) + } + + log.Debug().Msgf("Files to remove: %+v", filesToRemove) + log.Debug().Msgf("All other files: %+v", allOtherFiles) // Removing files for _, f := range filesToRemove { + if slices.Contains(allOtherFiles, f) { + log.Debug().Msgf("Skipping file %s because it is part of another model", f) + continue + } if e := os.Remove(f); e != nil { log.Error().Err(e).Msgf("failed to remove file %s", f) } diff --git a/core/gallery/models_test.go b/core/gallery/models_test.go index df0bee06ce8e..39b8a8427155 100644 --- a/core/gallery/models_test.go +++ b/core/gallery/models_test.go @@ -183,5 +183,98 @@ var _ = Describe("Model test", func() { _, err = InstallModel(context.TODO(), systemState, "../../../foo", c, map[string]interface{}{}, func(string, string, string, float64) {}, true) Expect(err).To(HaveOccurred()) }) + + It("does not delete shared model files when one config is deleted", func() { + tempdir, err := os.MkdirTemp("", "test") + Expect(err).ToNot(HaveOccurred()) + defer os.RemoveAll(tempdir) + + systemState, err := system.GetSystemState( + system.WithModelPath(tempdir), + ) + Expect(err).ToNot(HaveOccurred()) + + // Create a shared model file + sharedModelFile := filepath.Join(tempdir, "shared_model.bin") + err = os.WriteFile(sharedModelFile, []byte("fake model content"), 0600) + Expect(err).ToNot(HaveOccurred()) + + // Create first model configuration + config1 := `name: model1 +model: shared_model.bin` + err = os.WriteFile(filepath.Join(tempdir, "model1.yaml"), []byte(config1), 0600) + Expect(err).ToNot(HaveOccurred()) + + // Create first model's gallery file + galleryConfig1 := ModelConfig{ + Name: "model1", + Files: []File{ + {Filename: "shared_model.bin"}, + }, + } + galleryData1, err := yaml.Marshal(galleryConfig1) + Expect(err).ToNot(HaveOccurred()) + err = os.WriteFile(filepath.Join(tempdir, "._gallery_model1.yaml"), galleryData1, 0600) + Expect(err).ToNot(HaveOccurred()) + + // Create second model configuration sharing the same model file + config2 := `name: model2 +model: shared_model.bin` + err = os.WriteFile(filepath.Join(tempdir, "model2.yaml"), []byte(config2), 0600) + Expect(err).ToNot(HaveOccurred()) + + // Create second model's gallery file + galleryConfig2 := ModelConfig{ + Name: "model2", + Files: []File{ + {Filename: "shared_model.bin"}, + }, + } + galleryData2, err := yaml.Marshal(galleryConfig2) + Expect(err).ToNot(HaveOccurred()) + err = os.WriteFile(filepath.Join(tempdir, "._gallery_model2.yaml"), galleryData2, 0600) + Expect(err).ToNot(HaveOccurred()) + + // Verify both configurations exist + _, err = os.Stat(filepath.Join(tempdir, "model1.yaml")) + Expect(err).ToNot(HaveOccurred()) + _, err = os.Stat(filepath.Join(tempdir, "model2.yaml")) + Expect(err).ToNot(HaveOccurred()) + + // Verify the shared model file exists + _, err = os.Stat(sharedModelFile) + Expect(err).ToNot(HaveOccurred()) + + // Delete the first model + err = DeleteModelFromSystem(systemState, "model1") + Expect(err).ToNot(HaveOccurred()) + + // Verify the first configuration is deleted + _, err = os.Stat(filepath.Join(tempdir, "model1.yaml")) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, os.ErrNotExist)).To(BeTrue()) + + // Verify the shared model file still exists (not deleted because model2 still uses it) + _, err = os.Stat(sharedModelFile) + Expect(err).ToNot(HaveOccurred(), "shared model file should not be deleted when used by other configs") + + // Verify the second configuration still exists + _, err = os.Stat(filepath.Join(tempdir, "model2.yaml")) + Expect(err).ToNot(HaveOccurred()) + + // Now delete the second model + err = DeleteModelFromSystem(systemState, "model2") + Expect(err).ToNot(HaveOccurred()) + + // Verify the second configuration is deleted + _, err = os.Stat(filepath.Join(tempdir, "model2.yaml")) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, os.ErrNotExist)).To(BeTrue()) + + // Verify the shared model file is now deleted (no more references) + _, err = os.Stat(sharedModelFile) + Expect(err).To(HaveOccurred(), "shared model file should be deleted when no configs reference it") + Expect(errors.Is(err, os.ErrNotExist)).To(BeTrue()) + }) }) })