inference: Return memory requirement in estimation error

Signed-off-by: Piotr Stankiewicz <[email protected]>

Author: Piotr Stankiewicz

pkg/inference/memory/estimator.go

@@ -11,7 +11,7 @@ import (
 type MemoryEstimator interface {
 	SetDefaultBackend(MemoryEstimatorBackend)
 	GetRequiredMemoryForModel(context.Context, string, *inference.BackendConfiguration) (*inference.RequiredMemory, error)
-	HaveSufficientMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (bool, error)
+	HaveSufficientMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (bool, *inference.RequiredMemory, *inference.RequiredMemory, error)
 }
 
 type MemoryEstimatorBackend interface {
@@ -39,10 +39,11 @@ func (m *memoryEstimator) GetRequiredMemoryForModel(ctx context.Context, model s
 	return m.defaultBackend.GetRequiredMemoryForModel(ctx, model, config)
 }
 
-func (m *memoryEstimator) HaveSufficientMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (bool, error) {
+func (m *memoryEstimator) HaveSufficientMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (bool, *inference.RequiredMemory, *inference.RequiredMemory, error) {
 	req, err := m.GetRequiredMemoryForModel(ctx, model, config)
 	if err != nil {
-		return false, fmt.Errorf("estimating required memory for model: %w", err)
+		return false, nil, nil, fmt.Errorf("estimating required memory for model: %w", err)
 	}
-	return m.systemMemoryInfo.HaveSufficientMemory(*req), nil
+	totalMemory := m.systemMemoryInfo.GetTotalMemory()
+	return m.systemMemoryInfo.HaveSufficientMemory(*req), req, &totalMemory, nil
 }

pkg/inference/memory/system.go

@@ -40,6 +40,8 @@ func NewSystemMemoryInfo(log logging.Logger, gpuInfo *gpuinfo.GPUInfo) (SystemMe
 			log.Infof("Running on system with %d MB RAM", ramSize/1024/1024)
 		}
 	}
+	ramSize = 1024 * 1024 * 1024
+	vramSize = 1024 * 1024 * 1024
 	return &systemMemoryInfo{
 		log:         log,
 		totalMemory: inference.RequiredMemory{RAM: ramSize, VRAM: vramSize},

pkg/inference/models/manager.go

@@ -163,15 +163,16 @@ func (m *Manager) handleCreateModel(w http.ResponseWriter, r *http.Request) {
 	// besides pulling (such as model building).
 	if !request.IgnoreRuntimeMemoryCheck {
 		m.log.Infof("Will estimate memory required for %q", request.From)
-		proceed, err := m.memoryEstimator.HaveSufficientMemoryForModel(r.Context(), request.From, nil)
+		proceed, req, totalMem, err := m.memoryEstimator.HaveSufficientMemoryForModel(r.Context(), request.From, nil)
 		if err != nil {
 			m.log.Warnf("Failed to calculate memory required for model %q: %s", request.From, err)
 			// Prefer staying functional in case of unexpected estimation errors.
 			proceed = true
 		}
 		if !proceed {
-			m.log.Warnf("Runtime memory requirement for model %q exceeds total system memory", request.From)
-			http.Error(w, "Runtime memory requirement for model exceeds total system memory", http.StatusInsufficientStorage)
+			errstr := fmt.Sprintf("Runtime memory requirement for model %q exceeds total system memory: required %d RAM %d VRAM, system %d RAM %d VRAM", request.From, req.RAM, req.VRAM, totalMem.RAM, totalMem.VRAM)
+			m.log.Warnf(errstr)
+			http.Error(w, errstr, http.StatusInsufficientStorage)
 			return
 		}
 	}

pkg/inference/models/manager_test.go

@@ -30,8 +30,8 @@ func (me *mockMemoryEstimator) GetRequiredMemoryForModel(_ context.Context, _ st
 	return &inference.RequiredMemory{RAM: 0, VRAM: 0}, nil
 }
 
-func (me *mockMemoryEstimator) HaveSufficientMemoryForModel(_ context.Context, _ string, _ *inference.BackendConfiguration) (bool, error) {
-	return true, nil
+func (me *mockMemoryEstimator) HaveSufficientMemoryForModel(_ context.Context, _ string, _ *inference.BackendConfiguration) (bool, *inference.RequiredMemory, *inference.RequiredMemory, error) {
+	return true, nil, nil, nil
 }
 
 // getProjectRoot returns the absolute path to the project root directory