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

Usage Features (#863)

Browse source
This commit is contained in:
Dave 2023-08-18 15:23:14 -04:00 • committed by GitHub
parent 2bacd0180d
commit 8cb1061c11
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
40 changed files with 1222 additions and 317 deletions
Download patch file Download diff file Expand all files Collapse all files

32
pkg/model/initializers.go
View file

@ -6,6 +6,7 @@ import (
"os"
"os/signal"
"path/filepath"
"strconv"
"strings"
"syscall"
"time"
@ -64,10 +65,33 @@ var AutoLoadBackends []string = []string{
PiperBackend,
}
func (ml *ModelLoader) StopGRPC() {
for _, p := range ml.grpcProcesses {
p.Stop()
func (ml *ModelLoader) GetGRPCPID(id string) (int, error) {
p, exists := ml.grpcProcesses[id]
if !exists {
return -1, fmt.Errorf("no grpc backend found for %s", id)
}
return strconv.Atoi(p.PID)
}
type GRPCProcessFilter = func(p *process.Process) bool
func includeAllProcesses(_ *process.Process) bool {
return true
}
func (ml *ModelLoader) StopGRPC(filter GRPCProcessFilter) {
for _, p := range ml.grpcProcesses {
if filter(p) {
p.Stop()
}
}
}
func (ml *ModelLoader) StopAllGRPC() {
ml.StopGRPC(includeAllProcesses)
// for _, p := range ml.grpcProcesses {
// p.Stop()
// }
}
func (ml *ModelLoader) startProcess(grpcProcess, id string, serverAddress string) error {
@ -252,7 +276,7 @@ func (ml *ModelLoader) GreedyLoader(opts ...Option) (*grpc.Client, error) {
// Is this really needed? BackendLoader already does this
ml.mu.Lock()
if m := ml.checkIsLoaded(o.model); m != nil {
if m := ml.CheckIsLoaded(o.model); m != nil {
log.Debug().Msgf("Model '%s' already loaded", o.model)
ml.mu.Unlock()
return m, nil

4
pkg/model/loader.go
View file

@ -103,7 +103,7 @@ func (ml *ModelLoader) LoadModel(modelName string, loader func(string, string) (
defer ml.mu.Unlock()
// Check if we already have a loaded model
if model := ml.checkIsLoaded(modelName); model != nil {
if model := ml.CheckIsLoaded(modelName); model != nil {
return model, nil
}
@ -128,7 +128,7 @@ func (ml *ModelLoader) LoadModel(modelName string, loader func(string, string) (
return model, nil
}
func (ml *ModelLoader) checkIsLoaded(s string) *grpc.Client {
func (ml *ModelLoader) CheckIsLoaded(s string) *grpc.Client {
if m, ok := ml.models[s]; ok {
log.Debug().Msgf("Model already loaded in memory: %s", s)