dearwolf/LocalAI
1
0
Fork 0
mirror of https://github.com/mudler/LocalAI.git synced 2025-05-25 13:04:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 13

feat: initial watchdog implementation (#1341)

Browse source 
* feat: initial watchdog implementation

Signed-off-by: Ettore Di Giacinto <mudler@users.noreply.github.com>

* fiuxups

* Add more output

* wip: idletime checker

* wire idle watchdog checks

* enlarge watchdog time window

* small fixes

* Use stopmodel

* Always delete process

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

---------

Signed-off-by: Ettore Di Giacinto <mudler@users.noreply.github.com>
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto 2023-11-26 18:36:23 +01:00 • committed by GitHub
parent 9482acfdfc
commit 824612f1b4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 341 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

26
pkg/model/loader.go
View file

@ -63,12 +63,17 @@ type ModelLoader struct {
models map[string]ModelAddress
grpcProcesses map[string]*process.Process
templates map[TemplateType]map[string]*template.Template
wd *WatchDog
}
type ModelAddress string
func (m ModelAddress) GRPC(parallel bool) *grpc.Client {
return grpc.NewClient(string(m), parallel)
func (m ModelAddress) GRPC(parallel bool, wd *WatchDog) *grpc.Client {
enableWD := false
if wd != nil {
enableWD = true
}
return grpc.NewClient(string(m), parallel, wd, enableWD)
}
func NewModelLoader(modelPath string) *ModelLoader {
@ -79,10 +84,15 @@ func NewModelLoader(modelPath string) *ModelLoader {
templates: make(map[TemplateType]map[string]*template.Template),
grpcProcesses: make(map[string]*process.Process),
}
nml.initializeTemplateMap()
return nml
}
func (ml *ModelLoader) SetWatchDog(wd *WatchDog) {
ml.wd = wd
}
func (ml *ModelLoader) ExistsInModelPath(s string) bool {
return existsInPath(ml.ModelPath, s)
}
@ -139,11 +149,17 @@ func (ml *ModelLoader) LoadModel(modelName string, loader func(string, string) (
func (ml *ModelLoader) ShutdownModel(modelName string) error {
ml.mu.Lock()
defer ml.mu.Unlock()
return ml.StopModel(modelName)
}
func (ml *ModelLoader) StopModel(modelName string) error {
defer ml.deleteProcess(modelName)
if _, ok := ml.models[modelName]; !ok {
return fmt.Errorf("model %s not found", modelName)
}
return ml.deleteProcess(modelName)
return nil
//return ml.deleteProcess(modelName)
}
func (ml *ModelLoader) CheckIsLoaded(s string) ModelAddress {
@ -153,7 +169,7 @@ func (ml *ModelLoader) CheckIsLoaded(s string) ModelAddress {
if c, ok := ml.grpcClients[s]; ok {
client = c
} else {
client = m.GRPC(false)
client = m.GRPC(false, ml.wd)
}
if !client.HealthCheck(context.Background()) {