From b8a1045ffd8f10593b43f33a2faa381f768835af Mon Sep 17 00:00:00 2001 From: Luke Marsden Date: Wed, 1 Nov 2023 21:49:08 +0000 Subject: [PATCH 1/7] wip --- api/cmd/helix/root.go | 1 + api/pkg/runner/controller.go | 29 +++++++++++++++------ api/pkg/runner/server.go | 50 ++++++++++++++++++++++++++++++++++++ go.mod | 1 + 4 files changed, 73 insertions(+), 8 deletions(-) diff --git a/api/cmd/helix/root.go b/api/cmd/helix/root.go index d4ae7c7315..99830fdd3f 100644 --- a/api/cmd/helix/root.go +++ b/api/cmd/helix/root.go @@ -21,6 +21,7 @@ func NewRootCmd() *cobra.Command { } RootCmd.AddCommand(newServeCmd()) RootCmd.AddCommand(newRunnerCmd()) + RootCmd.AddCommand(newRunCmd()) return RootCmd } diff --git a/api/pkg/runner/controller.go b/api/pkg/runner/controller.go index 2a73d705cd..3653d2c32b 100644 --- a/api/pkg/runner/controller.go +++ b/api/pkg/runner/controller.go @@ -72,14 +72,15 @@ func NewRunner( ctx context.Context, options RunnerOptions, ) (*Runner, error) { - if options.ID == "" { - return nil, fmt.Errorf("id is required") - } - if options.ApiHost == "" { - return nil, fmt.Errorf("api url is required") - } - if options.ApiToken == "" { - return nil, fmt.Errorf("api token is required") + if options.ApiHost != "" { + // these are only required if api-host is specified, we can also run in + // a purely local mode + if options.ID == "" { + return nil, fmt.Errorf("id is required") + } + if options.ApiToken == "" { + return nil, fmt.Errorf("api token is required") + } } if options.MemoryString != "" { bytes, err := bytesize.Parse(options.MemoryString) @@ -320,6 +321,14 @@ func (r *Runner) getNextSession(ctx context.Context, queryParams url.Values) (*t // we pass that free memory back to the master API - it will filter out any tasks // for models that would require more memory than we have available func (r *Runner) getNextGlobalSession(ctx context.Context) (*types.Session, error) { + if r.httpClientOptions.Host == "" { + // we are in local only mode... the next session will be injected into + // us rather than queried from the control server + // TODO: it would be nice to still support memory limits and other + // smarts for local tasks + return nil, nil + } + freeMemory := r.getFreeMemory() if freeMemory < r.lowestMemoryRequirement { @@ -454,6 +463,10 @@ func (r *Runner) handleTaskResponse(ctx context.Context, instanceID string, task } func (r *Runner) uploadWorkerResponse(res *types.WorkerTaskResponse, session *types.Session) error { + if r.httpClientOptions.Host == "" { + // no upstream server configured, skip uploading + return nil + } if len(res.Files) > 0 { // create a new multipart form body := &bytes.Buffer{} diff --git a/api/pkg/runner/server.go b/api/pkg/runner/server.go index ed0a71ea1d..6e648cdf29 100644 --- a/api/pkg/runner/server.go +++ b/api/pkg/runner/server.go @@ -5,12 +5,14 @@ import ( "encoding/json" "fmt" "net/http" + "sync" "time" "github.com/gorilla/mux" "github.com/lukemarsden/helix/api/pkg/server" "github.com/lukemarsden/helix/api/pkg/system" "github.com/lukemarsden/helix/api/pkg/types" + "gopkg.in/yaml.v3" ) type RunnerServerOptions struct { @@ -21,6 +23,9 @@ type RunnerServerOptions struct { type RunnerServer struct { Options RunnerServerOptions Controller *Runner + // in-memory state to record status that would normally be posted up as a result + State map[string]types.WorkerTaskResponse + StateMtx sync.Mutex } func NewRunnerServer( @@ -41,6 +46,13 @@ func (runnerServer *RunnerServer) ListenAndServe(ctx context.Context, cm *system subrouter := router.PathPrefix(server.API_SUB_PATH).Subrouter() + // TODO: record worker response state locally, _in memory_ if we are in "local only mode" + // an endpoint to add our next session + subrouter.HandleFunc("/worker/session", server.Wrapper(runnerServer.setNextGlobalSession)).Methods("POST") + + // an endpoint to query the local state + subrouter.HandleFunc("/worker/state", server.Wrapper(runnerServer.state)).Methods("GET") + // pull the next task for an already running wrapper subrouter.HandleFunc("/worker/task/{instanceid}", server.WrapperWithConfig(runnerServer.getWorkerTask, server.WrapperConfig{ SilenceErrors: true, @@ -76,9 +88,47 @@ func (runnerServer *RunnerServer) respondWorkerTask(res http.ResponseWriter, req if err != nil { return nil, err } + taskResponse, err = runnerServer.Controller.handleTaskResponse(req.Context(), vars["instanceid"], taskResponse) if err != nil { return nil, err } + + // record in-memory for any local clients who want to query us + runnerServer.State[vars["instanceid"]] = *taskResponse + + stateYAML, err := yaml.Marshal(runnerServer.State) + if err != nil { + return nil, err + } + fmt.Println("==========================================") + fmt.Println(" LOCAL STATE") + fmt.Println("==========================================") + fmt.Println(string(stateYAML)) + fmt.Println("==========================================") + return taskResponse, nil } + +func (runnerServer *RunnerServer) state(res http.ResponseWriter, req *http.Request) (map[string]types.WorkerTaskResponse, error) { + return runnerServer.State, nil +} + +func (runnerServer *RunnerServer) setNextGlobalSession(res http.ResponseWriter, req *http.Request) (*types.WorkerTask, error) { + session := &types.Session{} + err := json.NewDecoder(req.Body).Decode(session) + if err != nil { + return nil, err + } + + // just start it instantly for now... + // TODO: what's the distinction between session and task? + // + // why does getNextGlobalSession always immediately start a new model + // instance? shouldn't it assign it to an existing one potentially? + runnerServer.Controller.createModelInstance(req.Context(), session) + + // TODO: Implement the logic to set the next global session + + return nil, nil +} diff --git a/go.mod b/go.mod index ba675442cf..67667dc597 100644 --- a/go.mod +++ b/go.mod @@ -59,4 +59,5 @@ require ( google.golang.org/genproto/googleapis/rpc v0.0.0-20231009173412-8bfb1ae86b6c // indirect google.golang.org/grpc v1.58.2 // indirect google.golang.org/protobuf v1.31.0 // indirect + gopkg.in/yaml.v3 v3.0.1 ) From b47c9ebe2685614a240fd86dec18cfc9d1e302f5 Mon Sep 17 00:00:00 2001 From: Luke Marsden Date: Thu, 2 Nov 2023 11:28:40 +0000 Subject: [PATCH 2/7] what to do next --- api/pkg/runner/server.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/api/pkg/runner/server.go b/api/pkg/runner/server.go index 6e648cdf29..2bc791acbb 100644 --- a/api/pkg/runner/server.go +++ b/api/pkg/runner/server.go @@ -126,6 +126,11 @@ func (runnerServer *RunnerServer) setNextGlobalSession(res http.ResponseWriter, // // why does getNextGlobalSession always immediately start a new model // instance? shouldn't it assign it to an existing one potentially? + + // what to do next: try running this code, get it working, then figure out + // how to make 'helix run' reuse an existing session (which i'm pretty sure + // this won't do) - also figure out how to write out results to disk, etc + runnerServer.Controller.createModelInstance(req.Context(), session) // TODO: Implement the logic to set the next global session From 777ba289a44d537ecf254080631b7031ff77711f Mon Sep 17 00:00:00 2001 From: Luke Marsden Date: Fri, 3 Nov 2023 17:07:05 +0000 Subject: [PATCH 3/7] wrap err --- api/pkg/runner/server.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/pkg/runner/server.go b/api/pkg/runner/server.go index 2bc791acbb..c00b81ae8b 100644 --- a/api/pkg/runner/server.go +++ b/api/pkg/runner/server.go @@ -118,7 +118,7 @@ func (runnerServer *RunnerServer) setNextGlobalSession(res http.ResponseWriter, session := &types.Session{} err := json.NewDecoder(req.Body).Decode(session) if err != nil { - return nil, err + return nil, fmt.Errorf("error decoding session as post body: %s", err) } // just start it instantly for now... From 143868f0f4608ff8c571d5245e33bd40cf89d9ac Mon Sep 17 00:00:00 2001 From: Luke Marsden Date: Sat, 4 Nov 2023 09:40:44 +0000 Subject: [PATCH 4/7] implement local queue for locally injected sessions, ensuring we reject ones that won't fit in GPU memory. What to do next: why is it hanging?? --- api/pkg/runner/controller.go | 94 +++++++++++++++++++++++++++++++----- api/pkg/runner/server.go | 6 ++- 2 files changed, 88 insertions(+), 12 deletions(-) diff --git a/api/pkg/runner/controller.go b/api/pkg/runner/controller.go index e839bab988..48004d92d9 100644 --- a/api/pkg/runner/controller.go +++ b/api/pkg/runner/controller.go @@ -64,6 +64,8 @@ type Runner struct { // no point asking for more top level tasks // we get this on boot by asking the model package lowestMemoryRequirement uint64 + + localQueue []*types.Session } func NewRunner( @@ -107,6 +109,44 @@ func NewRunner( return runner, nil } +func modelInstanceMatchesSession(modelInstance *ModelInstance, session *types.Session) bool { + return modelInstance.filter.Mode == session.Mode && + modelInstance.filter.Type == session.Type && + (modelInstance.filter.FinetuneFile == session.FinetuneFile || + (modelInstance.filter.FinetuneFile == "none" && session.FinetuneFile == "")) +} + +func (r *Runner) AddToLocalQueue(ctx context.Context, session *types.Session) error { + // iterate over model instances to see if one exists and if it doesn't, create it. + // then add session to localQueue + + // Check if a model instance exists for the session's model ID + found := false + + // loop over r.activeModelInstances, checking whether the filters on the + // model instance match the session mode, type and finetune + for _, modelInstance := range r.activeModelInstances { + if modelInstanceMatchesSession(modelInstance, session) { + // no need to create another one, because there's already one which will match the session + log.Printf("🟠 Found modelInstance %+v which matches session %+v", modelInstance, session) + found = true + break + } + } + if !found { + // Create a new model instance because it doesn't exist + log.Printf("🟠 No currently running modelInstance for session %+v, starting a new one", session) + err := r.createModelInstance(ctx, session) + if err != nil { + return err + } + } + + // Add the session to the local queue + r.localQueue = append(r.localQueue, session) + return nil +} + // this should be run in a go-routine func (r *Runner) StartLooping() { for { @@ -222,6 +262,7 @@ func (r *Runner) getNextGlobalSession(ctx context.Context) (*types.Session, erro func (r *Runner) createModelInstance(ctx context.Context, session *types.Session) error { r.modelMutex.Lock() defer r.modelMutex.Unlock() + modelInstance, err := NewModelInstance( r.Ctx, session, @@ -241,6 +282,15 @@ func (r *Runner) createModelInstance(ctx context.Context, session *types.Session if err != nil { return err } + + // belt and braces in remote case and reject jobs that won't fit in local case + modelMem := modelInstance.model.GetMemoryRequirements(session.Mode) + freeMem := r.getFreeMemory() + if modelMem > freeMem { + // refuse to start or record the model instance, it will just get GC'd at this point + return fmt.Errorf("cannot fit model requiring gpu memory %d into available gpu memory %d", modelMem, freeMem) + } + log.Debug(). Msgf("🔵 runner started model instance: %s", modelInstance.id) @@ -287,7 +337,27 @@ func (r *Runner) getNextTask(ctx context.Context, instanceID string) (*types.Wor var session *types.Session - if modelInstance.nextSession != nil { + foundLocalQueuedSession := false + for i, sess := range r.localQueue { + if modelInstanceMatchesSession(modelInstance, sess) { + foundLocalQueuedSession = true + // remove it from the local queue + r.localQueue = append(r.localQueue[:i], r.localQueue[i+1:]...) + session = sess + break + } + + } + // as the first check, we need to ask if there's a session in localQueue + // that matches this model instance. if there is, we've got a local + // session and it takes precedence over remote work + + // if there is, call modelInstance.queueSession on it + + if foundLocalQueuedSession { + // queue it, and fall thru below to assign + go modelInstance.queueSession(session) + } else if modelInstance.nextSession != nil { // if there is a session in the nextSession cache then we return it immediately session = modelInstance.nextSession modelInstance.nextSession = nil @@ -298,19 +368,21 @@ func (r *Runner) getNextTask(ctx context.Context, instanceID string) (*types.Wor // ask the upstream api server if there is another task // if there is - then assign it to the queuedSession // and call "pre" - queryParams := url.Values{} + if r.httpClientOptions.Host != "" { + queryParams := url.Values{} - queryParams.Add("model_name", string(modelInstance.filter.ModelName)) - queryParams.Add("mode", string(modelInstance.filter.Mode)) - queryParams.Add("finetune_file", string(modelInstance.filter.FinetuneFile)) + queryParams.Add("model_name", string(modelInstance.filter.ModelName)) + queryParams.Add("mode", string(modelInstance.filter.Mode)) + queryParams.Add("finetune_file", string(modelInstance.filter.FinetuneFile)) - apiSession, err := r.getNextSession(ctx, queryParams) - if err != nil { - return nil, err - } + apiSession, err := r.getNextSession(ctx, queryParams) + if err != nil { + return nil, err + } - if apiSession != nil { - go modelInstance.queueSession(apiSession) + if apiSession != nil { + go modelInstance.queueSession(apiSession) + } } } diff --git a/api/pkg/runner/server.go b/api/pkg/runner/server.go index c00b81ae8b..edfd6d605b 100644 --- a/api/pkg/runner/server.go +++ b/api/pkg/runner/server.go @@ -131,7 +131,11 @@ func (runnerServer *RunnerServer) setNextGlobalSession(res http.ResponseWriter, // how to make 'helix run' reuse an existing session (which i'm pretty sure // this won't do) - also figure out how to write out results to disk, etc - runnerServer.Controller.createModelInstance(req.Context(), session) + err = runnerServer.Controller.AddToLocalQueue(req.Context(), session) + // err = runnerServer.Controller.createModelInstance(req.Context(), session) + if err != nil { + return nil, err + } // TODO: Implement the logic to set the next global session From a7d6816a6026c3edd937741a24b7f70f53ae03ab Mon Sep 17 00:00:00 2001 From: Luke Marsden Date: Sun, 5 Nov 2023 08:40:03 +0000 Subject: [PATCH 5/7] switch to xsync.MapOf --- api/pkg/runner/controller.go | 73 +++++++++++++++++------------------- go.mod | 1 + go.sum | 2 + 3 files changed, 38 insertions(+), 38 deletions(-) diff --git a/api/pkg/runner/controller.go b/api/pkg/runner/controller.go index 48004d92d9..5bac0c12bb 100644 --- a/api/pkg/runner/controller.go +++ b/api/pkg/runner/controller.go @@ -11,7 +11,6 @@ import ( "net/url" "os" "runtime/debug" - "sync" "time" "github.com/davecgh/go-spew/spew" @@ -20,6 +19,7 @@ import ( "github.com/lukemarsden/helix/api/pkg/model" "github.com/lukemarsden/helix/api/pkg/server" "github.com/lukemarsden/helix/api/pkg/types" + "github.com/puzpuzpuz/xsync/v3" "github.com/rs/zerolog/log" ) @@ -54,10 +54,9 @@ type Runner struct { httpClientOptions server.ClientOptions - modelMutex sync.RWMutex // the map of model instances that we have loaded // and are currently running - activeModelInstances map[string]*ModelInstance + activeModelInstances *xsync.MapOf[string, *ModelInstance] // the lowest amount of memory that something can run with // if we have less than this amount of memory then there is @@ -65,7 +64,8 @@ type Runner struct { // we get this on boot by asking the model package lowestMemoryRequirement uint64 - localQueue []*types.Session + // local sessions, which will be executed in no particular order + localSessions *xsync.MapOf[string, *types.Session] } func NewRunner( @@ -104,7 +104,7 @@ func NewRunner( Host: options.ApiHost, Token: options.ApiToken, }, - activeModelInstances: map[string]*ModelInstance{}, + activeModelInstances: xsync.NewMapOf[string, *ModelInstance](), } return runner, nil } @@ -125,14 +125,15 @@ func (r *Runner) AddToLocalQueue(ctx context.Context, session *types.Session) er // loop over r.activeModelInstances, checking whether the filters on the // model instance match the session mode, type and finetune - for _, modelInstance := range r.activeModelInstances { + r.activeModelInstances.Range(func(key string, modelInstance *ModelInstance) bool { if modelInstanceMatchesSession(modelInstance, session) { // no need to create another one, because there's already one which will match the session log.Printf("🟠 Found modelInstance %+v which matches session %+v", modelInstance, session) found = true - break + return false } - } + return true + }) if !found { // Create a new model instance because it doesn't exist log.Printf("🟠 No currently running modelInstance for session %+v, starting a new one", session) @@ -143,7 +144,7 @@ func (r *Runner) AddToLocalQueue(ctx context.Context, session *types.Session) er } // Add the session to the local queue - r.localQueue = append(r.localQueue, session) + r.localSessions.Store(session.ID, session) return nil } @@ -194,22 +195,21 @@ func (r *Runner) loop(ctx context.Context) error { // loop over the active model instances and stop any that have not processed a job // in the last timeout seconds func (r *Runner) checkForStaleModelInstances(ctx context.Context, timeout time.Duration) error { - r.modelMutex.Lock() - defer r.modelMutex.Unlock() - for _, activeModelInstance := range r.activeModelInstances { + r.activeModelInstances.Range(func(key string, activeModelInstance *ModelInstance) bool { // this means we are booting so let's leave it alone to boot if activeModelInstance.lastActivityTimestamp == 0 { - continue + return true } if activeModelInstance.lastActivityTimestamp+int64(timeout.Seconds()) < time.Now().Unix() { log.Info().Msgf("Killing stale model instance %s", activeModelInstance.id) err := activeModelInstance.stopProcess() if err != nil { log.Error().Msgf("error stopping model instance %s: %s", activeModelInstance.id, err.Error()) - continue + return true } } - } + return true + }) return nil } @@ -245,9 +245,10 @@ func (r *Runner) getNextGlobalSession(ctx context.Context) (*types.Session, erro // (i.e. if we get 100 text inferences then the chance is we boot 100 model instances) // before trying to run another type of model - for _, modelInstance := range r.activeModelInstances { + r.activeModelInstances.Range(func(key string, modelInstance *ModelInstance) bool { queryParams.Add("reject", fmt.Sprintf("%s:%s", modelInstance.filter.ModelName, modelInstance.filter.Mode)) - } + return true + }) return r.getNextSession(ctx, queryParams) } @@ -260,9 +261,6 @@ func (r *Runner) getNextGlobalSession(ctx context.Context) (*types.Session, erro // and will add the de-prioritise filter to the next request // so that we get a different job type func (r *Runner) createModelInstance(ctx context.Context, session *types.Session) error { - r.modelMutex.Lock() - defer r.modelMutex.Unlock() - modelInstance, err := NewModelInstance( r.Ctx, session, @@ -284,12 +282,13 @@ func (r *Runner) createModelInstance(ctx context.Context, session *types.Session } // belt and braces in remote case and reject jobs that won't fit in local case - modelMem := modelInstance.model.GetMemoryRequirements(session.Mode) - freeMem := r.getFreeMemory() + modelMem := float32(modelInstance.model.GetMemoryRequirements(session.Mode)) / 1024 / 1024 / 1024 + freeMem := float32(r.getFreeMemory()) / 1024 / 1024 / 1024 if modelMem > freeMem { // refuse to start or record the model instance, it will just get GC'd at this point - return fmt.Errorf("cannot fit model requiring gpu memory %d into available gpu memory %d", modelMem, freeMem) + return fmt.Errorf("cannot fit model requiring gpu memory %.2f into available gpu memory %.2f", modelMem, freeMem) } + log.Printf("🟠 Fitting model requiring gpu memory %.2f into available gpu memory %.2f", modelMem, freeMem) log.Debug(). Msgf("🔵 runner started model instance: %s", modelInstance.id) @@ -307,14 +306,13 @@ func (r *Runner) createModelInstance(ctx context.Context, session *types.Session if err != nil { return err } - r.activeModelInstances[modelInstance.id] = modelInstance + + r.activeModelInstances.Store(modelInstance.id, modelInstance) go func() { <-modelInstance.finishChan - r.modelMutex.Lock() - defer r.modelMutex.Unlock() log.Debug(). Msgf("🔵 runner stop model instance: %s", modelInstance.id) - delete(r.activeModelInstances, modelInstance.id) + r.activeModelInstances.Delete(modelInstance.id) }() return nil } @@ -330,7 +328,7 @@ func (r *Runner) getNextTask(ctx context.Context, instanceID string) (*types.Wor if instanceID == "" { return nil, fmt.Errorf("instanceid is required") } - modelInstance, ok := r.activeModelInstances[instanceID] + modelInstance, ok := r.activeModelInstances.Load(instanceID) if !ok { return nil, fmt.Errorf("instance not found: %s", instanceID) } @@ -338,16 +336,16 @@ func (r *Runner) getNextTask(ctx context.Context, instanceID string) (*types.Wor var session *types.Session foundLocalQueuedSession := false - for i, sess := range r.localQueue { + r.localSessions.Range(func(i string, sess *types.Session) bool { if modelInstanceMatchesSession(modelInstance, sess) { foundLocalQueuedSession = true // remove it from the local queue - r.localQueue = append(r.localQueue[:i], r.localQueue[i+1:]...) + r.localSessions.Delete(i) session = sess - break + return false } - - } + return true + }) // as the first check, we need to ask if there's a session in localQueue // that matches this model instance. if there is, we've got a local // session and it takes precedence over remote work @@ -406,7 +404,7 @@ func (r *Runner) handleTaskResponse(ctx context.Context, instanceID string, task if taskResponse == nil { return nil, fmt.Errorf("task response is required") } - modelInstance, ok := r.activeModelInstances[instanceID] + modelInstance, ok := r.activeModelInstances.Load(instanceID) if !ok { return nil, fmt.Errorf("instance not found: %s", instanceID) } @@ -475,12 +473,11 @@ func (r *Runner) getNextSession(ctx context.Context, queryParams url.Values) (*t } func (r *Runner) getUsedMemory() uint64 { - r.modelMutex.RLock() - defer r.modelMutex.RUnlock() memoryUsed := uint64(0) - for _, modelInstance := range r.activeModelInstances { + r.activeModelInstances.Range(func(i string, modelInstance *ModelInstance) bool { memoryUsed += modelInstance.model.GetMemoryRequirements(modelInstance.filter.Mode) - } + return true + }) return memoryUsed } diff --git a/go.mod b/go.mod index 67667dc597..1d16adbee7 100644 --- a/go.mod +++ b/go.mod @@ -41,6 +41,7 @@ require ( github.com/mattn/go-isatty v0.0.19 // indirect github.com/opentracing/opentracing-go v1.2.0 // indirect github.com/pkg/errors v0.9.1 // indirect + github.com/puzpuzpuz/xsync/v3 v3.0.1 github.com/segmentio/ksuid v1.0.4 // indirect github.com/spf13/pflag v1.0.5 // indirect go.opencensus.io v0.24.0 // indirect diff --git a/go.sum b/go.sum index 4771f7ae78..ad3968f12b 100644 --- a/go.sum +++ b/go.sum @@ -117,6 +117,8 @@ github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA= +github.com/puzpuzpuz/xsync/v3 v3.0.1 h1:yhTYnDJlgIYp/3Bb14b43VfUPrk/QNJ1HrLYEZ8r2AE= +github.com/puzpuzpuz/xsync/v3 v3.0.1/go.mod h1:VjzYrABPabuM4KyBh1Ftq6u8nhwY5tBPKP9jpmh0nnA= github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg= github.com/rs/zerolog v1.31.0 h1:FcTR3NnLWW+NnTwwhFWiJSZr4ECLpqCm6QsEnyvbV4A= github.com/rs/zerolog v1.31.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss= From 66d4840521a8ff69ddff3094902aecaee9aaa120 Mon Sep 17 00:00:00 2001 From: Luke Marsden Date: Sun, 5 Nov 2023 08:52:44 +0000 Subject: [PATCH 6/7] initialize properly --- api/pkg/runner/controller.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/api/pkg/runner/controller.go b/api/pkg/runner/controller.go index 5bac0c12bb..c71d09ec93 100644 --- a/api/pkg/runner/controller.go +++ b/api/pkg/runner/controller.go @@ -65,6 +65,7 @@ type Runner struct { lowestMemoryRequirement uint64 // local sessions, which will be executed in no particular order + // TODO: maybe preserve insertion order localSessions *xsync.MapOf[string, *types.Session] } @@ -105,6 +106,7 @@ func NewRunner( Token: options.ApiToken, }, activeModelInstances: xsync.NewMapOf[string, *ModelInstance](), + localSessions: xsync.NewMapOf[string, *types.Session](), } return runner, nil } From f8e813a811bc3bcd4535a073e9881e01ec5fbf3a Mon Sep 17 00:00:00 2001 From: Luke Marsden Date: Sun, 5 Nov 2023 09:04:14 +0000 Subject: [PATCH 7/7] Delete after killing --- api/pkg/runner/controller.go | 1 + 1 file changed, 1 insertion(+) diff --git a/api/pkg/runner/controller.go b/api/pkg/runner/controller.go index c71d09ec93..77f2b80962 100644 --- a/api/pkg/runner/controller.go +++ b/api/pkg/runner/controller.go @@ -209,6 +209,7 @@ func (r *Runner) checkForStaleModelInstances(ctx context.Context, timeout time.D log.Error().Msgf("error stopping model instance %s: %s", activeModelInstance.id, err.Error()) return true } + r.activeModelInstances.Delete(activeModelInstance.id) } return true })