diff --git a/worker/worker.go b/worker/worker.go index 07e4f29c..48b55235 100644 --- a/worker/worker.go +++ b/worker/worker.go @@ -40,22 +40,24 @@ func (sb EnvValue) String() string { type OptimizationFlags map[string]EnvValue type Worker struct { - manager *DockerManager + manager *DockerManager + noManagedContainers bool externalContainers map[string]*RunnerContainer mu *sync.Mutex } -func NewWorker(containerImageID string, gpus []string, modelDir string) (*Worker, error) { +func NewWorker(containerImageID string, gpus []string, modelDir string, noManagedContainers bool) (*Worker, error) { manager, err := NewDockerManager(containerImageID, gpus, modelDir) if err != nil { return nil, err } return &Worker{ - manager: manager, - externalContainers: make(map[string]*RunnerContainer), - mu: &sync.Mutex{}, + manager: manager, + noManagedContainers: noManagedContainers, + externalContainers: make(map[string]*RunnerContainer), + mu: &sync.Mutex{}, }, nil } @@ -329,6 +331,10 @@ func (w *Worker) borrowContainer(ctx context.Context, pipeline, modelID string) w.mu.Unlock() + if w.noManagedContainers { + return nil, errors.New("no runners available") + } + return w.manager.Borrow(ctx, pipeline, modelID) }