inference: Return memory requirement in estimation error

p1-0tr · p1-0tr · commit 1f8d94b11784 · 2025-09-04T16:23:09.000+02:00
Signed-off-by: Piotr Stankiewicz &lt;piotr.stankiewicz@docker.com&gt;
diff --git a/pkg/inference/backend.go b/pkg/inference/backend.go
@@ -91,5 +91,5 @@ type Backend interface {
 	GetDiskUsage() (int64, error)
 	// GetRequiredMemoryForModel returns the required working memory for a given
 	// model.
-	GetRequiredMemoryForModel(ctx context.Context, model string, config *BackendConfiguration) (*RequiredMemory, error)
+	GetRequiredMemoryForModel(ctx context.Context, model string, config *BackendConfiguration) (RequiredMemory, error)
 }
diff --git a/pkg/inference/backends/llamacpp/llamacpp.go b/pkg/inference/backends/llamacpp/llamacpp.go
@@ -230,22 +230,22 @@ func (l *llamaCpp) GetDiskUsage() (int64, error) {
 	return size, nil
 }
 
-func (l *llamaCpp) GetRequiredMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (*inference.RequiredMemory, error) {
+func (l *llamaCpp) GetRequiredMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (inference.RequiredMemory, error) {
 	var mdlGguf *parser.GGUFFile
 	var mdlConfig types.Config
 	inStore, err := l.modelManager.IsModelInStore(model)
 	if err != nil {
-		return nil, fmt.Errorf("checking if model is in local store: %w", err)
+		return inference.RequiredMemory{}, fmt.Errorf("checking if model is in local store: %w", err)
 	}
 	if inStore {
 		mdlGguf, mdlConfig, err = l.parseLocalModel(model)
 		if err != nil {
-			return nil, &inference.ErrGGUFParse{Err: err}
+			return inference.RequiredMemory{}, &inference.ErrGGUFParse{Err: err}
 		}
 	} else {
 		mdlGguf, mdlConfig, err = l.parseRemoteModel(ctx, model)
 		if err != nil {
-			return nil, &inference.ErrGGUFParse{Err: err}
+			return inference.RequiredMemory{}, &inference.ErrGGUFParse{Err: err}
 		}
 	}
 
@@ -278,7 +278,7 @@ func (l *llamaCpp) GetRequiredMemoryForModel(ctx context.Context, model string,
 		vram = 1
 	}
 
-	return &inference.RequiredMemory{
+	return inference.RequiredMemory{
 		RAM:  ram,
 		VRAM: vram,
 	}, nil
diff --git a/pkg/inference/backends/mlx/mlx.go b/pkg/inference/backends/mlx/mlx.go
@@ -63,6 +63,6 @@ func (m *mlx) GetDiskUsage() (int64, error) {
 	return 0, nil
 }
 
-func (m *mlx) GetRequiredMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (*inference.RequiredMemory, error) {
-	return nil, errors.New("not implemented")
+func (m *mlx) GetRequiredMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (inference.RequiredMemory, error) {
+	return inference.RequiredMemory{}, errors.New("not implemented")
 }
diff --git a/pkg/inference/backends/vllm/vllm.go b/pkg/inference/backends/vllm/vllm.go
@@ -63,6 +63,6 @@ func (v *vLLM) GetDiskUsage() (int64, error) {
 	return 0, nil
 }
 
-func (v *vLLM) GetRequiredMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (*inference.RequiredMemory, error) {
-	return nil, errors.New("not implemented")
+func (v *vLLM) GetRequiredMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (inference.RequiredMemory, error) {
+	return inference.RequiredMemory{}, errors.New("not implemented")
 }
diff --git a/pkg/inference/memory/estimator.go b/pkg/inference/memory/estimator.go
@@ -10,12 +10,12 @@ 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)
+	GetRequiredMemoryForModel(context.Context, string, *inference.BackendConfiguration) (inference.RequiredMemory, error)
+	HaveSufficientMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (bool, inference.RequiredMemory, inference.RequiredMemory, error)
 }
 
 type MemoryEstimatorBackend interface {
-	GetRequiredMemoryForModel(context.Context, string, *inference.BackendConfiguration) (*inference.RequiredMemory, error)
+	GetRequiredMemoryForModel(context.Context, string, *inference.BackendConfiguration) (inference.RequiredMemory, error)
 }
 
 type memoryEstimator struct {
@@ -31,18 +31,18 @@ func (m *memoryEstimator) SetDefaultBackend(backend MemoryEstimatorBackend) {
 	m.defaultBackend = backend
 }
 
-func (m *memoryEstimator) GetRequiredMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (*inference.RequiredMemory, error) {
+func (m *memoryEstimator) GetRequiredMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (inference.RequiredMemory, error) {
 	if m.defaultBackend == nil {
-		return nil, errors.New("default backend not configured")
+		return inference.RequiredMemory{}, errors.New("default backend not configured")
 	}
 
 	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, inference.RequiredMemory{}, inference.RequiredMemory{}, fmt.Errorf("estimating required memory for model: %w", err)
 	}
-	return m.systemMemoryInfo.HaveSufficientMemory(*req), nil
+	return m.systemMemoryInfo.HaveSufficientMemory(req), req, m.systemMemoryInfo.GetTotalMemory(), nil
 }
diff --git a/pkg/inference/models/manager.go b/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
 		}
 	}
diff --git a/pkg/inference/models/manager_test.go b/pkg/inference/models/manager_test.go
@@ -26,12 +26,12 @@ type mockMemoryEstimator struct{}
 
 func (me *mockMemoryEstimator) SetDefaultBackend(_ memory.MemoryEstimatorBackend) {}
 
-func (me *mockMemoryEstimator) GetRequiredMemoryForModel(_ context.Context, _ string, _ *inference.BackendConfiguration) (*inference.RequiredMemory, error) {
-	return &inference.RequiredMemory{RAM: 0, VRAM: 0}, nil
+func (me *mockMemoryEstimator) GetRequiredMemoryForModel(_ context.Context, _ string, _ *inference.BackendConfiguration) (inference.RequiredMemory, error) {
+	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, inference.RequiredMemory{}, inference.RequiredMemory{}, nil
 }
 
 // getProjectRoot returns the absolute path to the project root directory
diff --git a/pkg/inference/scheduling/loader.go b/pkg/inference/scheduling/loader.go
@@ -401,7 +401,7 @@ func (l *loader) load(ctx context.Context, backendName, modelID, modelRef string
 		// e.g. model is too new for gguf-parser-go to know. We should provide a cleaner
 		// way to bypass these checks.
 		l.log.Warnf("Could not parse model(%s), memory checks will be ignored for it. Error: %s", modelID, parseErr)
-		memory = &inference.RequiredMemory{
+		memory = inference.RequiredMemory{
 			RAM:  0,
 			VRAM: 0,
 		}

Original file line number	Diff line number	Diff line change
`@@ -91,5 +91,5 @@ type Backend interface {`
`91`	`91`	`GetDiskUsage() (int64, error)`
`92`	`92`	`// GetRequiredMemoryForModel returns the required working memory for a given`
`93`	`93`	`// model.`
`94`		`- GetRequiredMemoryForModel(ctx context.Context, model string, config BackendConfiguration) (RequiredMemory, error)`
	`94`	`+ GetRequiredMemoryForModel(ctx context.Context, model string, config *BackendConfiguration) (RequiredMemory, error)`
`95`	`95`	`}`
Original file line number	Diff line number	Diff line change
`@@ -230,22 +230,22 @@ func (l *llamaCpp) GetDiskUsage() (int64, error) {`
`230`	`230`	`return size, nil`
`231`	`231`	`}`
`232`	`232`
`233`		`-func (l llamaCpp) GetRequiredMemoryForModel(ctx context.Context, model string, config inference.BackendConfiguration) (*inference.RequiredMemory, error) {`
	`233`	`+func (l llamaCpp) GetRequiredMemoryForModel(ctx context.Context, model string, config inference.BackendConfiguration) (inference.RequiredMemory, error) {`
`234`	`234`	`var mdlGguf *parser.GGUFFile`
`235`	`235`	`var mdlConfig types.Config`
`236`	`236`	`inStore, err := l.modelManager.IsModelInStore(model)`
`237`	`237`	`if err != nil {`
`238`		`- return nil, fmt.Errorf("checking if model is in local store: %w", err)`
	`238`	`+ return inference.RequiredMemory{}, fmt.Errorf("checking if model is in local store: %w", err)`
`239`	`239`	`}`
`240`	`240`	`if inStore {`
`241`	`241`	`mdlGguf, mdlConfig, err = l.parseLocalModel(model)`
`242`	`242`	`if err != nil {`
`243`		`- return nil, &inference.ErrGGUFParse{Err: err}`
	`243`	`+ return inference.RequiredMemory{}, &inference.ErrGGUFParse{Err: err}`
`244`	`244`	`}`
`245`	`245`	`} else {`
`246`	`246`	`mdlGguf, mdlConfig, err = l.parseRemoteModel(ctx, model)`
`247`	`247`	`if err != nil {`
`248`		`- return nil, &inference.ErrGGUFParse{Err: err}`
	`248`	`+ return inference.RequiredMemory{}, &inference.ErrGGUFParse{Err: err}`
`249`	`249`	`}`
`250`	`250`	`}`
`251`	`251`
`@@ -278,7 +278,7 @@ func (l *llamaCpp) GetRequiredMemoryForModel(ctx context.Context, model string,`
`278`	`278`	`vram = 1`
`279`	`279`	`}`
`280`	`280`
`281`		`- return &inference.RequiredMemory{`
	`281`	`+ return inference.RequiredMemory{`
`282`	`282`	`RAM: ram,`
`283`	`283`	`VRAM: vram,`
`284`	`284`	`}, nil`
Original file line number	Diff line number	Diff line change
`@@ -63,6 +63,6 @@ func (m *mlx) GetDiskUsage() (int64, error) {`
`63`	`63`	`return 0, nil`
`64`	`64`	`}`
`65`	`65`
`66`		`-func (m mlx) GetRequiredMemoryForModel(ctx context.Context, model string, config inference.BackendConfiguration) (*inference.RequiredMemory, error) {`
`67`		`- return nil, errors.New("not implemented")`
	`66`	`+func (m mlx) GetRequiredMemoryForModel(ctx context.Context, model string, config inference.BackendConfiguration) (inference.RequiredMemory, error) {`
	`67`	`+ return inference.RequiredMemory{}, errors.New("not implemented")`
`68`	`68`	`}`
Original file line number	Diff line number	Diff line change
`@@ -63,6 +63,6 @@ func (v *vLLM) GetDiskUsage() (int64, error) {`
`63`	`63`	`return 0, nil`
`64`	`64`	`}`
`65`	`65`
`66`		`-func (v vLLM) GetRequiredMemoryForModel(ctx context.Context, model string, config inference.BackendConfiguration) (*inference.RequiredMemory, error) {`
`67`		`- return nil, errors.New("not implemented")`
	`66`	`+func (v vLLM) GetRequiredMemoryForModel(ctx context.Context, model string, config inference.BackendConfiguration) (inference.RequiredMemory, error) {`
	`67`	`+ return inference.RequiredMemory{}, errors.New("not implemented")`
`68`	`68`	`}`
Original file line number	Diff line number	Diff line change
`@@ -10,12 +10,12 @@ import (`
`10`	`10`
`11`	`11`	`type MemoryEstimator interface {`
`12`	`12`	`SetDefaultBackend(MemoryEstimatorBackend)`
`13`		`- GetRequiredMemoryForModel(context.Context, string, inference.BackendConfiguration) (inference.RequiredMemory, error)`
`14`		`- HaveSufficientMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (bool, error)`
	`13`	`+ GetRequiredMemoryForModel(context.Context, string, *inference.BackendConfiguration) (inference.RequiredMemory, error)`
	`14`	`+ HaveSufficientMemoryForModel(ctx context.Context, model string, config *inference.BackendConfiguration) (bool, inference.RequiredMemory, inference.RequiredMemory, error)`
`15`	`15`	`}`
`16`	`16`
`17`	`17`	`type MemoryEstimatorBackend interface {`
`18`		`- GetRequiredMemoryForModel(context.Context, string, inference.BackendConfiguration) (inference.RequiredMemory, error)`
	`18`	`+ GetRequiredMemoryForModel(context.Context, string, *inference.BackendConfiguration) (inference.RequiredMemory, error)`
`19`	`19`	`}`
`20`	`20`
`21`	`21`	`type memoryEstimator struct {`
`@@ -31,18 +31,18 @@ func (m *memoryEstimator) SetDefaultBackend(backend MemoryEstimatorBackend) {`
`31`	`31`	`m.defaultBackend = backend`
`32`	`32`	`}`
`33`	`33`
`34`		`-func (m memoryEstimator) GetRequiredMemoryForModel(ctx context.Context, model string, config inference.BackendConfiguration) (*inference.RequiredMemory, error) {`
	`34`	`+func (m memoryEstimator) GetRequiredMemoryForModel(ctx context.Context, model string, config inference.BackendConfiguration) (inference.RequiredMemory, error) {`
`35`	`35`	`if m.defaultBackend == nil {`
`36`		`- return nil, errors.New("default backend not configured")`
	`36`	`+ return inference.RequiredMemory{}, errors.New("default backend not configured")`
`37`	`37`	`}`
`38`	`38`
`39`	`39`	`return m.defaultBackend.GetRequiredMemoryForModel(ctx, model, config)`
`40`	`40`	`}`
`41`	`41`
`42`		`-func (m memoryEstimator) HaveSufficientMemoryForModel(ctx context.Context, model string, config inference.BackendConfiguration) (bool, error) {`
	`42`	`+func (m memoryEstimator) HaveSufficientMemoryForModel(ctx context.Context, model string, config inference.BackendConfiguration) (bool, inference.RequiredMemory, inference.RequiredMemory, error) {`
`43`	`43`	`req, err := m.GetRequiredMemoryForModel(ctx, model, config)`
`44`	`44`	`if err != nil {`
`45`		`- return false, fmt.Errorf("estimating required memory for model: %w", err)`
	`45`	`+ return false, inference.RequiredMemory{}, inference.RequiredMemory{}, fmt.Errorf("estimating required memory for model: %w", err)`
`46`	`46`	`}`
`47`		`- return m.systemMemoryInfo.HaveSufficientMemory(*req), nil`
	`47`	`+ return m.systemMemoryInfo.HaveSufficientMemory(req), req, m.systemMemoryInfo.GetTotalMemory(), nil`
`48`	`48`	`}`
Original file line number	Diff line number	Diff line change
`@@ -26,12 +26,12 @@ type mockMemoryEstimator struct{}`
`26`	`26`
`27`	`27`	`func (me *mockMemoryEstimator) SetDefaultBackend(_ memory.MemoryEstimatorBackend) {}`
`28`	`28`
`29`		`-func (me mockMemoryEstimator) GetRequiredMemoryForModel(_ context.Context, _ string, _ inference.BackendConfiguration) (*inference.RequiredMemory, error) {`
`30`		`- return &inference.RequiredMemory{RAM: 0, VRAM: 0}, nil`
	`29`	`+func (me mockMemoryEstimator) GetRequiredMemoryForModel(_ context.Context, _ string, _ inference.BackendConfiguration) (inference.RequiredMemory, error) {`
	`30`	`+ return inference.RequiredMemory{RAM: 0, VRAM: 0}, nil`
`31`	`31`	`}`
`32`	`32`
`33`		`-func (me mockMemoryEstimator) HaveSufficientMemoryForModel(_ context.Context, _ string, _ inference.BackendConfiguration) (bool, error) {`
`34`		`- return true, nil`
	`33`	`+func (me mockMemoryEstimator) HaveSufficientMemoryForModel(_ context.Context, _ string, _ inference.BackendConfiguration) (bool, inference.RequiredMemory, inference.RequiredMemory, error) {`
	`34`	`+ return true, inference.RequiredMemory{}, inference.RequiredMemory{}, nil`
`35`	`35`	`}`
`36`	`36`
`37`	`37`	`// getProjectRoot returns the absolute path to the project root directory`
Original file line number	Diff line number	Diff line change
`@@ -401,7 +401,7 @@ func (l *loader) load(ctx context.Context, backendName, modelID, modelRef string`
`401`	`401`	`// e.g. model is too new for gguf-parser-go to know. We should provide a cleaner`
`402`	`402`	`// way to bypass these checks.`
`403`	`403`	`l.log.Warnf("Could not parse model(%s), memory checks will be ignored for it. Error: %s", modelID, parseErr)`
`404`		`- memory = &inference.RequiredMemory{`
	`404`	`+ memory = inference.RequiredMemory{`
`405`	`405`	`RAM: 0,`
`406`	`406`	`VRAM: 0,`
`407`	`407`	`}`