Merge pull request #700 from doringeman/unify-vllm

doringeman · web-flow · commit 9e75f0499c6a · 2026-02-19T16:33:46.000+02:00
refactor: unify vllm and vllm-metal into a single "vllm" backend
diff --git a/cmd/cli/commands/install-runner.go b/cmd/cli/commands/install-runner.go
@@ -16,7 +16,7 @@ import (
 	"github.com/docker/model-runner/pkg/inference/backends/diffusers"
 	"github.com/docker/model-runner/pkg/inference/backends/llamacpp"
 	"github.com/docker/model-runner/pkg/inference/backends/vllm"
-	"github.com/docker/model-runner/pkg/inference/backends/vllmmetal"
+	"github.com/docker/model-runner/pkg/inference/platform"
 	"github.com/moby/moby/api/types/container"
 	"github.com/spf13/cobra"
 )
@@ -29,7 +29,7 @@ const (
 	// installation will try to reach the model runner while waiting for it to
 	// be ready.
 	installWaitRetryInterval = 500 * time.Millisecond
-	backendUsage             = "Specify backend (" + llamacpp.Name + "|" + vllm.Name + "|" + diffusers.Name + "|" + vllmmetal.Name + "). Default: " + llamacpp.Name
+	backendUsage             = "Specify backend (" + llamacpp.Name + "|" + vllm.Name + "|" + diffusers.Name + "). Default: " + llamacpp.Name
 )
 
 // waitForStandaloneRunnerAfterInstall waits for a standalone model runner
@@ -242,14 +242,14 @@ type runnerOptions struct {
 
 // runInstallOrStart is shared logic for install-runner and start-runner commands
 func runInstallOrStart(cmd *cobra.Command, opts runnerOptions, debug bool) error {
-	// vllm-metal is installed on-demand via the running model runner,
-	// not as a standalone container. This applies to all engine kinds.
-	if opts.backend == vllmmetal.Name {
-		cmd.Println("Installing vllm-metal backend...")
-		if err := desktopClient.InstallBackend(vllmmetal.Name); err != nil {
-			return fmt.Errorf("failed to install vllm-metal backend: %w", err)
+	// On macOS ARM64, the vllm backend requires deferred installation
+	// (on-demand via the running model runner), not as a standalone container.
+	if opts.backend == vllm.Name && platform.SupportsVLLMMetal() {
+		cmd.Println("Installing vllm backend...")
+		if err := desktopClient.InstallBackend(vllm.Name); err != nil {
+			return fmt.Errorf("failed to install vllm backend: %w", err)
 		}
-		cmd.Println("vllm-metal backend installed successfully")
+		cmd.Println("vllm backend installed successfully")
 		return nil
 	}
 
@@ -340,7 +340,7 @@ func runInstallOrStart(cmd *cobra.Command, opts runnerOptions, debug bool) error
 	}
 
 	// Validate backend selection
-	validBackends := []string{llamacpp.Name, vllm.Name, diffusers.Name, vllmmetal.Name}
+	validBackends := []string{llamacpp.Name, vllm.Name, diffusers.Name}
 	if opts.backend != "" {
 		isValid := false
 		for _, valid := range validBackends {
@@ -354,8 +354,8 @@ func runInstallOrStart(cmd *cobra.Command, opts runnerOptions, debug bool) error
 		}
 	}
 
-	// Validate backend-GPU compatibility
-	if opts.backend == vllm.Name && gpu != gpupkg.GPUSupportCUDA {
+	// Validate backend-GPU compatibility (only on Linux; macOS ARM64 uses Metal)
+	if opts.backend == vllm.Name && !platform.SupportsVLLMMetal() && gpu != gpupkg.GPUSupportCUDA {
 		return fmt.Errorf("--backend vllm requires CUDA GPU support (--gpu=cuda or auto-detected CUDA)")
 	}
 
diff --git a/cmd/cli/docs/reference/docker_model_install-runner.yaml b/cmd/cli/docs/reference/docker_model_install-runner.yaml
@@ -8,8 +8,7 @@ plink: docker_model.yaml
 options:
     - option: backend
       value_type: string
-      description: |
-        Specify backend (llama.cpp|vllm|diffusers|vllm-metal). Default: llama.cpp
+      description: 'Specify backend (llama.cpp|vllm|diffusers). Default: llama.cpp'
       deprecated: false
       hidden: false
       experimental: false
diff --git a/cmd/cli/docs/reference/docker_model_reinstall-runner.yaml b/cmd/cli/docs/reference/docker_model_reinstall-runner.yaml
@@ -8,8 +8,7 @@ plink: docker_model.yaml
 options:
     - option: backend
       value_type: string
-      description: |
-        Specify backend (llama.cpp|vllm|diffusers|vllm-metal). Default: llama.cpp
+      description: 'Specify backend (llama.cpp|vllm|diffusers). Default: llama.cpp'
       deprecated: false
       hidden: false
       experimental: false
diff --git a/cmd/cli/docs/reference/docker_model_start-runner.yaml b/cmd/cli/docs/reference/docker_model_start-runner.yaml
@@ -10,8 +10,7 @@ plink: docker_model.yaml
 options:
     - option: backend
       value_type: string
-      description: |
-        Specify backend (llama.cpp|vllm|diffusers|vllm-metal). Default: llama.cpp
+      description: 'Specify backend (llama.cpp|vllm|diffusers). Default: llama.cpp'
       deprecated: false
       hidden: false
       experimental: false
diff --git a/cmd/cli/docs/reference/model_install-runner.md b/cmd/cli/docs/reference/model_install-runner.md
@@ -7,7 +7,7 @@ Install Docker Model Runner (Docker Engine only)
 
 | Name             | Type     | Default     | Description                                                                                            |
 |:-----------------|:---------|:------------|:-------------------------------------------------------------------------------------------------------|
-| `--backend`      | `string` |             | Specify backend (llama.cpp\|vllm\|diffusers\|vllm-metal). Default: llama.cpp                           |
+| `--backend`      | `string` |             | Specify backend (llama.cpp\|vllm\|diffusers). Default: llama.cpp                                       |
 | `--debug`        | `bool`   |             | Enable debug logging                                                                                   |
 | `--do-not-track` | `bool`   |             | Do not track models usage in Docker Model Runner                                                       |
 | `--gpu`          | `string` | `auto`      | Specify GPU support (none\|auto\|cuda\|rocm\|musa\|cann)                                               |
diff --git a/cmd/cli/docs/reference/model_reinstall-runner.md b/cmd/cli/docs/reference/model_reinstall-runner.md
@@ -7,7 +7,7 @@ Reinstall Docker Model Runner (Docker Engine only)
 
 | Name             | Type     | Default     | Description                                                                                            |
 |:-----------------|:---------|:------------|:-------------------------------------------------------------------------------------------------------|
-| `--backend`      | `string` |             | Specify backend (llama.cpp\|vllm\|diffusers\|vllm-metal). Default: llama.cpp                           |
+| `--backend`      | `string` |             | Specify backend (llama.cpp\|vllm\|diffusers). Default: llama.cpp                                       |
 | `--debug`        | `bool`   |             | Enable debug logging                                                                                   |
 | `--do-not-track` | `bool`   |             | Do not track models usage in Docker Model Runner                                                       |
 | `--gpu`          | `string` | `auto`      | Specify GPU support (none\|auto\|cuda\|rocm\|musa\|cann)                                               |
diff --git a/cmd/cli/docs/reference/model_start-runner.md b/cmd/cli/docs/reference/model_start-runner.md
@@ -7,7 +7,7 @@ Start Docker Model Runner (Docker Engine only)
 
 | Name             | Type     | Default     | Description                                                                                            |
 |:-----------------|:---------|:------------|:-------------------------------------------------------------------------------------------------------|
-| `--backend`      | `string` |             | Specify backend (llama.cpp\|vllm\|diffusers\|vllm-metal). Default: llama.cpp                           |
+| `--backend`      | `string` |             | Specify backend (llama.cpp\|vllm\|diffusers). Default: llama.cpp                                       |
 | `--debug`        | `bool`   |             | Enable debug logging                                                                                   |
 | `--do-not-track` | `bool`   |             | Do not track models usage in Docker Model Runner                                                       |
 | `--gpu`          | `string` | `auto`      | Specify GPU support (none\|auto\|cuda\|rocm\|musa\|cann)                                               |
diff --git a/main.go b/main.go
@@ -122,28 +122,30 @@ func main() {
 			Logger:        log.WithFields(logrus.Fields{"component": "model-manager"}),
 			Transport:     baseTransport,
 		},
-		Backends: append(append(
+		Backends: append(
 			routing.DefaultBackendDefs(routing.BackendsConfig{
 				Log:                  log,
 				LlamaCppVendoredPath: llamaServerPath,
 				LlamaCppUpdatedPath:  updatedServerPath,
 				LlamaCppConfig:       llamaCppConfig,
 				IncludeMLX:           true,
 				MLXPath:              mlxServerPath,
+				IncludeVLLM:          includeVLLM,
+				VLLMPath:             vllmServerPath,
+				VLLMMetalPath:        vllmMetalServerPath,
 			}),
 			routing.BackendDef{Name: sglang.Name, Init: func(mm *models.Manager) (inference.Backend, error) {
 				return sglang.New(log, mm, log.WithFields(logrus.Fields{"component": sglang.Name}), nil, sglangServerPath)
 			}},
 			routing.BackendDef{Name: diffusers.Name, Init: func(mm *models.Manager) (inference.Backend, error) {
 				return diffusers.New(log, mm, log.WithFields(logrus.Fields{"component": diffusers.Name}), nil, diffusersServerPath)
 			}},
-		), vllmBackendDefs(log, vllmServerPath)...),
+		),
 		OnBackendError: func(name string, err error) {
 			log.Fatalf("unable to initialize %s backend: %v", name, err)
 		},
-		DefaultBackendName:  llamacpp.Name,
-		VLLMMetalServerPath: vllmMetalServerPath,
-		HTTPClient:          http.DefaultClient,
+		DefaultBackendName: llamacpp.Name,
+		HTTPClient:         http.DefaultClient,
 		MetricsTracker: metrics.NewTracker(
 			http.DefaultClient,
 			log.WithField("component", "metrics"),
diff --git a/pkg/inference/backends/vllm/vllm.go b/pkg/inference/backends/vllm/vllm.go
@@ -44,9 +44,33 @@ type vLLM struct {
 	customBinaryPath string
 }
 
-// New creates a new vLLM-based backend.
+// Options holds the configuration for the unified vLLM backend constructor.
+type Options struct {
+	Config          *Config // Linux-only: extra vllm args (nil = defaults)
+	LinuxBinaryPath string  // Linux: custom vllm binary path
+	MetalPythonPath string  // macOS ARM64: custom python path
+}
+
+// New creates the appropriate vLLM backend for the current platform.
+// On macOS ARM64, it returns the vllm-metal backend; on Linux, the standard
+// vLLM backend. On unsupported platforms, the returned backend's Install/Run
+// methods return errors.
+func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, opts Options) (inference.Backend, error) {
+	if platform.SupportsVLLMMetal() {
+		return newMetal(log, modelManager, serverLog, opts.MetalPythonPath)
+	}
+	return newLinux(log, modelManager, serverLog, opts.Config, opts.LinuxBinaryPath)
+}
+
+// NeedsDeferredInstall reports whether vllm on the current platform
+// requires deferred (on-demand) installation.
+func NeedsDeferredInstall() bool {
+	return platform.SupportsVLLMMetal()
+}
+
+// newLinux creates a new Linux vLLM-based backend.
 // customBinaryPath is an optional path to a custom vllm binary; if empty, the default path is used.
-func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, conf *Config, customBinaryPath string) (inference.Backend, error) {
+func newLinux(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, conf *Config, customBinaryPath string) (inference.Backend, error) {
 	// If no config is provided, use the default configuration
 	if conf == nil {
 		conf = NewDefaultVLLMConfig()
diff --git a/pkg/inference/backends/vllm/vllm_metal.go b/pkg/inference/backends/vllm/vllm_metal.go
@@ -1,4 +1,4 @@
-package vllmmetal
+package vllm
 
 import (
 	"context"
@@ -19,12 +19,9 @@ import (
 	"github.com/docker/model-runner/pkg/inference/platform"
 	"github.com/docker/model-runner/pkg/internal/dockerhub"
 	"github.com/docker/model-runner/pkg/logging"
-	"github.com/sirupsen/logrus"
 )
 
 const (
-	// Name is the backend name.
-	Name              = "vllm-metal"
 	defaultInstallDir = ".docker/model-runner/vllm-metal"
 	// vllmMetalVersion is the vllm-metal release tag to download from Docker Hub.
 	vllmMetalVersion = "v0.1.0-20260126-121650"
@@ -53,9 +50,9 @@ type vllmMetal struct {
 	status string
 }
 
-// New creates a new vllm-metal backend.
+// newMetal creates a new vllm-metal backend.
 // customPythonPath is an optional path to a custom python3 binary; if empty, the default installation is used.
-func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, customPythonPath string) (inference.Backend, error) {
+func newMetal(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, customPythonPath string) (inference.Backend, error) {
 	homeDir, err := os.UserHomeDir()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get user home directory: %w", err)
@@ -72,22 +69,6 @@ func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Log
 	}, nil
 }
 
-// TryRegister initializes the vllm-metal backend if the platform supports it
-// and registers it in the provided backends map. It returns the backend names
-// whose installation should be deferred until explicitly requested.
-func TryRegister(log logging.Logger, modelManager *models.Manager, backends map[string]inference.Backend, serverPath string) []string {
-	if !platform.SupportsVLLMMetal() {
-		return nil
-	}
-	backend, err := New(log, modelManager, log.WithFields(logrus.Fields{"component": Name}), serverPath)
-	if err != nil {
-		log.Warnf("Failed to initialize vllm-metal backend: %v", err)
-		return nil
-	}
-	backends[Name] = backend
-	return []string{Name}
-}
-
 // Name implements inference.Backend.Name.
 func (v *vllmMetal) Name() string {
 	return Name
diff --git a/pkg/inference/scheduling/scheduler.go b/pkg/inference/scheduling/scheduler.go
@@ -14,7 +14,6 @@ import (
 	"github.com/docker/model-runner/pkg/inference/backends/mlx"
 	"github.com/docker/model-runner/pkg/inference/backends/sglang"
 	"github.com/docker/model-runner/pkg/inference/backends/vllm"
-	"github.com/docker/model-runner/pkg/inference/backends/vllmmetal"
 	"github.com/docker/model-runner/pkg/inference/models"
 	"github.com/docker/model-runner/pkg/inference/platform"
 	"github.com/docker/model-runner/pkg/internal/utils"
@@ -29,15 +28,17 @@ import (
 type PlatformSupport interface {
 	SupportsMLX() bool
 	SupportsVLLM() bool
+	SupportsVLLMMetal() bool
 	SupportsSGLang() bool
 }
 
 // defaultPlatformSupport delegates to the platform package.
 type defaultPlatformSupport struct{}
 
-func (defaultPlatformSupport) SupportsMLX() bool    { return platform.SupportsMLX() }
-func (defaultPlatformSupport) SupportsVLLM() bool   { return platform.SupportsVLLM() }
-func (defaultPlatformSupport) SupportsSGLang() bool { return platform.SupportsSGLang() }
+func (defaultPlatformSupport) SupportsMLX() bool       { return platform.SupportsMLX() }
+func (defaultPlatformSupport) SupportsVLLM() bool      { return platform.SupportsVLLM() }
+func (defaultPlatformSupport) SupportsVLLMMetal() bool { return platform.SupportsVLLMMetal() }
+func (defaultPlatformSupport) SupportsSGLang() bool    { return platform.SupportsSGLang() }
 
 // Scheduler is used to coordinate inference scheduling across multiple backends
 // and models.
@@ -121,8 +122,9 @@ func (s *Scheduler) Run(ctx context.Context) error {
 
 // selectBackendForModel selects the appropriate backend for a model based on its format.
 // If the model is in safetensors format, it will prefer the best available backend:
-// - On macOS: vllm-metal > MLX
-// - On Linux: vLLM > SGLang
+// - vLLM (handles platform dispatch internally: vllm-metal on macOS ARM64, standard vLLM on Linux)
+// - MLX on macOS
+// - SGLang on Linux
 func (s *Scheduler) selectBackendForModel(model types.Model, backend inference.Backend, modelRef string) inference.Backend {
 	config, err := model.Config()
 	if err != nil {
@@ -131,29 +133,19 @@ func (s *Scheduler) selectBackendForModel(model types.Model, backend inference.B
 	}
 
 	if config.GetFormat() == types.FormatSafetensors {
-		// Prefer vllm-metal for safetensors models on macOS (most feature-rich for Metal),
-		// but only if it has been installed.
-		if vllmMetalBackend, ok := s.backends[vllmmetal.Name]; ok && vllmMetalBackend != nil {
-			if s.installer.isInstalled(vllmmetal.Name) {
-				return vllmMetalBackend
+		// Prefer vLLM for safetensors models (handles platform dispatch internally)
+		if s.platformSupport.SupportsVLLM() || s.platformSupport.SupportsVLLMMetal() {
+			if vllmBackend, ok := s.backends[vllm.Name]; ok && vllmBackend != nil {
+				return vllmBackend
 			}
-			s.log.Infof("vllm-metal backend is available but not installed. "+
-				"To install, run: docker model install-runner --backend %s", vllmmetal.Name)
-			return vllmMetalBackend
 		}
 		// Fall back to MLX on macOS
 		if s.platformSupport.SupportsMLX() {
 			if mlxBackend, ok := s.backends[mlx.Name]; ok && mlxBackend != nil {
 				return mlxBackend
 			}
 		}
-		// Prefer vLLM for safetensors models on Linux
-		if s.platformSupport.SupportsVLLM() {
-			if vllmBackend, ok := s.backends[vllm.Name]; ok && vllmBackend != nil {
-				return vllmBackend
-			}
-		}
-		// Fall back to SGLang if vLLM is not available
+		// Fall back to SGLang on Linux
 		if s.platformSupport.SupportsSGLang() {
 			if sglangBackend, ok := s.backends[sglang.Name]; ok && sglangBackend != nil {
 				return sglangBackend
diff --git a/pkg/inference/scheduling/select_backend_test.go b/pkg/inference/scheduling/select_backend_test.go
@@ -14,14 +14,16 @@ import (
 
 // mockPlatformSupport allows tests to control platform capability checks.
 type mockPlatformSupport struct {
-	mlx    bool
-	vllm   bool
-	sglang bool
+	mlx       bool
+	vllm      bool
+	vllmMetal bool
+	sglang    bool
 }
 
-func (m mockPlatformSupport) SupportsMLX() bool    { return m.mlx }
-func (m mockPlatformSupport) SupportsVLLM() bool   { return m.vllm }
-func (m mockPlatformSupport) SupportsSGLang() bool { return m.sglang }
+func (m mockPlatformSupport) SupportsMLX() bool       { return m.mlx }
+func (m mockPlatformSupport) SupportsVLLM() bool      { return m.vllm }
+func (m mockPlatformSupport) SupportsVLLMMetal() bool { return m.vllmMetal }
+func (m mockPlatformSupport) SupportsSGLang() bool    { return m.sglang }
 
 // mockModel is a minimal Model implementation for testing.
 type mockModel struct {
@@ -82,7 +84,7 @@ func TestSelectBackendForModel(t *testing.T) {
 			expectedBackend: vllm.Name,
 		},
 		{
-			name: "macOS with MLX and vLLM registered selects MLX for safetensors",
+			name: "macOS without vllm-metal support falls back to MLX for safetensors",
 			backends: map[string]inference.Backend{
 				"llamacpp":  llamacppBackend,
 				mlx.Name:    mlxBackend,
@@ -94,6 +96,19 @@ func TestSelectBackendForModel(t *testing.T) {
 			model:           safetensorsModel,
 			expectedBackend: mlx.Name,
 		},
+		{
+			name: "macOS ARM64 with vllm-metal support selects unified vllm for safetensors",
+			backends: map[string]inference.Backend{
+				"llamacpp":  llamacppBackend,
+				mlx.Name:    mlxBackend,
+				vllm.Name:   vllmBackend,
+				sglang.Name: sglangBackend,
+			},
+			defaultBackend:  llamacppBackend,
+			platform:        mockPlatformSupport{mlx: true, vllmMetal: true},
+			model:           safetensorsModel,
+			expectedBackend: vllm.Name,
+		},
 		{
 			name: "Linux with only SGLang selects SGLang for safetensors",
 			backends: map[string]inference.Backend{
diff --git a/pkg/routing/backends.go b/pkg/routing/backends.go
diff --git a/pkg/routing/service.go b/pkg/routing/service.go
diff --git a/vllm_backend.go b/vllm_backend.go
diff --git a/vllm_backend_stub.go b/vllm_backend_stub.go
