chore(refactor): drop duplicated shutdown logics (#3589)

* chore(refactor): drop duplicated shutdown logics

- Handle locking in Shutdown and CheckModelIsLoaded in a more go-idiomatic way
- Drop duplicated code and re-organize shutdown code

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix: drop leftover

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* chore: improve logging and add missing locks

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto 2024-09-17 16:51:40 +02:00 committed by GitHub
parent eee1fb2c75
commit a53392f919
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 32 additions and 38 deletions

View file

@ -118,9 +118,6 @@ func (ml *ModelLoader) ListModels() []*Model {
}
func (ml *ModelLoader) LoadModel(modelName string, loader func(string, string) (*Model, error)) (*Model, error) {
ml.mu.Lock()
defer ml.mu.Unlock()
// Check if we already have a loaded model
if model := ml.CheckIsLoaded(modelName); model != nil {
return model, nil
@ -139,6 +136,8 @@ func (ml *ModelLoader) LoadModel(modelName string, loader func(string, string) (
return nil, fmt.Errorf("loader didn't return a model")
}
ml.mu.Lock()
defer ml.mu.Unlock()
ml.models[modelName] = model
return model, nil
@ -168,6 +167,8 @@ func (ml *ModelLoader) ShutdownModel(modelName string) error {
}
func (ml *ModelLoader) CheckIsLoaded(s string) *Model {
ml.mu.Lock()
defer ml.mu.Unlock()
m, ok := ml.models[s]
if !ok {
return nil