Check for model folder when creating container

worker/docker.go

@@ -3,7 +3,10 @@ package worker
 import (
 	"context"
 	"errors"
+	"fmt"
 	"log/slog"
+	"os"
+	"path/filepath"
 	"strings"
 	"sync"
 	"time"
@@ -150,6 +153,17 @@ func (m *DockerManager) HasCapacity(ctx context.Context, pipeline, modelID strin
 	return err == nil
 }
 
+// ModelExists checks that the models directory and specific model folder exists before loading a container or processing a request
+func (m *DockerManager) ModelExists(modelID string) bool {
+	modelPathCheck := filepath.Join(m.modelDir, "models--"+strings.ReplaceAll(modelID, "/", "--"))
+	if _, err := os.Stat(modelPathCheck); os.IsNotExist(err) {
+		slog.Error(fmt.Sprintf("model %s does not exist at %s", modelID, modelPathCheck))
+		return false
+	} else {
+		return true
+	}
+}
+
 func (m *DockerManager) createContainer(ctx context.Context, pipeline string, modelID string, keepWarm bool, optimizationFlags OptimizationFlags) (*RunnerContainer, error) {
 	containerName := dockerContainerName(pipeline, modelID)
 
@@ -158,6 +172,10 @@ func (m *DockerManager) createContainer(ctx context.Context, pipeline string, mo
 		return nil, err
 	}
 
+	if !m.ModelExists(modelID) {
+		return nil, fmt.Errorf("model %s does not exist", modelID)
+	}
+
 	slog.Info("Starting managed container", slog.String("gpu", gpu), slog.String("name", containerName), slog.String("modelID", modelID))
 
 	// Add optimization flags as environment variables.

worker/worker.go

@@ -5,8 +5,12 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"log/slog"
+	"os"
+	"path/filepath"
 	"strconv"
+	"strings"
 	"sync"
 )
 
@@ -362,6 +366,22 @@ func (w *Worker) HasCapacity(pipeline, modelID string) bool {
 	return ok
 }
 
+// HasCapacity checks if an unused managed container exists or if a GPU is available for a new container.
+func (m *Worker) ModelExists(modelID string) bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	//Check if the model folder exists
+	modelPathCheck := filepath.Join(m.manager.modelDir, "models--"+strings.ReplaceAll(modelID, "/", "--"))
+	if _, err := os.Stat(modelPathCheck); os.IsNotExist(err) {
+		slog.Error(fmt.Sprintf("model %s does not exist at %s", modelID, modelPathCheck))
+		return false
+	} else {
+		slog.Info("Model found: ", slog.String("path", modelPathCheck))
+		return true
+	}
+}
+
 func (w *Worker) borrowContainer(ctx context.Context, pipeline, modelID string) (*RunnerContainer, error) {
 	w.mu.Lock()