Merge pull request #1003 from mat007/command-modifier

Add an optional CommandModifier for backend runner processes

Author: doringeman

pkg/inference/backends/diffusers/diffusers.go

@@ -56,11 +56,13 @@ type diffusers struct {
 	installDir string
 	// registryMirrors is the list of registry mirrors to try before registry-1.docker.io.
 	registryMirrors []string
+	// commandModifier, if non-nil, is applied to the server process before it starts.
+	commandModifier func(*exec.Cmd)
 }
 
 // New creates a new diffusers-based backend for image generation.
 // 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, conf *Config, customPythonPath string, registryMirrors []string) (inference.Backend, error) {
+func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, conf *Config, customPythonPath string, registryMirrors []string, commandModifier func(*exec.Cmd)) (inference.Backend, error) {
 	// If no config is provided, use the default configuration
 	if conf == nil {
 		conf = NewDefaultConfig()
@@ -81,6 +83,7 @@ func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Log
 		customPythonPath: customPythonPath,
 		installDir:       installDir,
 		registryMirrors:  registryMirrors,
+		commandModifier:  commandModifier,
 	}, nil
 }
 
@@ -265,6 +268,7 @@ func (d *diffusers) Run(ctx context.Context, socket, model string, modelRef stri
 		Logger:           d.log,
 		ServerLogWriter:  logging.NewWriter(d.serverLog),
 		ErrorTransformer: ExtractPythonError,
+		CommandModifier:  d.commandModifier,
 	})
 }
 

pkg/inference/backends/llamacpp/llamacpp.go

@@ -53,6 +53,8 @@ type llamaCpp struct {
 	gpuSupported bool
 	// registryMirrors is the list of registry mirrors to try before registry-1.docker.io.
 	registryMirrors []string
+	// commandModifier, if non-nil, is applied to the server process before it starts.
+	commandModifier func(*exec.Cmd)
 }
 
 // New creates a new llama.cpp-based backend. installDir is the directory that
@@ -65,6 +67,7 @@ func New(
 	installDir string,
 	conf config.BackendConfig,
 	registryMirrors []string,
+	commandModifier func(*exec.Cmd),
 ) (inference.Backend, error) {
 	// If no config is provided, use the default configuration
 	if conf == nil {
@@ -87,6 +90,7 @@ func New(
 		status:          inference.FormatNotInstalled(""),
 		config:          conf,
 		registryMirrors: registryMirrors,
+		commandModifier: commandModifier,
 	}, nil
 }
 
@@ -228,6 +232,7 @@ func (l *llamaCpp) Run(ctx context.Context, socket, model string, _ string, mode
 		Logger:           l.log,
 		ServerLogWriter:  logging.NewWriter(l.serverLog),
 		ErrorTransformer: ExtractLlamaCppError,
+		CommandModifier:  l.commandModifier,
 	})
 }
 

pkg/inference/backends/mlx/mlx.go

@@ -39,11 +39,13 @@ type mlx struct {
 	pythonPath string
 	// customPythonPath is an optional custom path to the python3 binary.
 	customPythonPath string
+	// commandModifier, if non-nil, is applied to the server process before it starts.
+	commandModifier func(*exec.Cmd)
 }
 
 // New creates a new MLX-based backend.
 // customPythonPath is an optional path to a custom python3 binary; if empty, the default path is used.
-func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, conf *Config, customPythonPath string) (inference.Backend, error) {
+func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, conf *Config, customPythonPath string, commandModifier func(*exec.Cmd)) (inference.Backend, error) {
 	// If no config is provided, use the default configuration
 	if conf == nil {
 		conf = NewDefaultMLXConfig()
@@ -56,6 +58,7 @@ func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Log
 		config:           conf,
 		status:           inference.FormatNotInstalled(""),
 		customPythonPath: customPythonPath,
+		commandModifier:  commandModifier,
 	}, nil
 }
 
@@ -145,6 +148,7 @@ func (m *mlx) Run(ctx context.Context, socket, model string, modelRef string, mo
 		Args:            args,
 		Logger:          m.log,
 		ServerLogWriter: logging.NewWriter(m.serverLog),
+		CommandModifier: m.commandModifier,
 	})
 }
 

pkg/inference/backends/runner.go

@@ -47,6 +47,12 @@ type RunnerConfig struct {
 	// process environment. If nil or empty, the backend inherits the parent
 	// env as-is (an empty non-nil slice is treated the same as nil).
 	Env []string
+	// CommandModifier, if non-nil, is invoked on the backend command immediately
+	// before it is started, after the internal setup (cancellation, stdio, env).
+	// Embedders use it to customize process attributes such as credentials
+	// (SysProcAttr), environment, or working directory. It composes with, and
+	// runs after, the internal modifier.
+	CommandModifier func(*exec.Cmd)
 }
 
 // Logger interface for backend logging
@@ -113,6 +119,15 @@ func RunBackend(ctx context.Context, config RunnerConfig) error {
 			if len(config.Env) > 0 {
 				command.Env = append(os.Environ(), config.Env...)
 			}
+			if config.CommandModifier != nil {
+				// Materialize the inherited environment first: a nil Env means
+				// "inherit the parent", but a modifier that appends to it would
+				// otherwise start from an empty slice and drop the parent env.
+				if command.Env == nil {
+					command.Env = os.Environ()
+				}
+				config.CommandModifier(command)
+			}
 		},
 		config.SandboxPath,
 		config.BinaryPath,

pkg/inference/backends/sglang/sglang.go

@@ -47,11 +47,13 @@ type sglang struct {
 	pythonPath string
 	// customPythonPath is an optional custom path to the python3 binary.
 	customPythonPath string
+	// commandModifier, if non-nil, is applied to the server process before it starts.
+	commandModifier func(*exec.Cmd)
 }
 
 // New creates a new SGLang-based backend.
 // customPythonPath is an optional path to a custom python3 binary; if empty, the default path is used.
-func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, conf *Config, customPythonPath string) (inference.Backend, error) {
+func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, conf *Config, customPythonPath string, commandModifier func(*exec.Cmd)) (inference.Backend, error) {
 	// If no config is provided, use the default configuration
 	if conf == nil {
 		conf = NewDefaultSGLangConfig()
@@ -64,6 +66,7 @@ func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Log
 		config:           conf,
 		status:           inference.FormatNotInstalled(""),
 		customPythonPath: customPythonPath,
+		commandModifier:  commandModifier,
 	}, nil
 }
 
@@ -174,6 +177,7 @@ func (s *sglang) Run(ctx context.Context, socket, model string, modelRef string,
 		Args:            args,
 		Logger:          s.log,
 		ServerLogWriter: logging.NewWriter(s.serverLog),
+		CommandModifier: s.commandModifier,
 	})
 }
 

pkg/inference/backends/vllm/vllm.go

@@ -8,6 +8,7 @@ import (
 	"io/fs"
 	"net/http"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"strings"
 
@@ -45,14 +46,17 @@ type vLLM struct {
 	customBinaryPath string
 	// registryMirrors is the list of registry mirrors to try before registry-1.docker.io.
 	registryMirrors []string
+	// commandModifier, if non-nil, is applied to the server process before it starts.
+	commandModifier func(*exec.Cmd)
 }
 
 // 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
-	RegistryMirrors []string // registry mirrors tried before registry-1.docker.io
+	Config          *Config         // Linux-only: extra vllm args (nil = defaults)
+	LinuxBinaryPath string          // Linux: custom vllm binary path
+	MetalPythonPath string          // macOS ARM64: custom python path
+	RegistryMirrors []string        // registry mirrors tried before registry-1.docker.io
+	CommandModifier func(*exec.Cmd) // applied to the server process before it starts
 }
 
 // New creates the appropriate vLLM backend for the current platform.
@@ -61,9 +65,9 @@ type Options struct {
 // 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, opts.RegistryMirrors)
+		return newMetal(log, modelManager, serverLog, opts.MetalPythonPath, opts.RegistryMirrors, opts.CommandModifier)
 	}
-	return newLinux(log, modelManager, serverLog, opts.Config, opts.LinuxBinaryPath, opts.RegistryMirrors)
+	return newLinux(log, modelManager, serverLog, opts.Config, opts.LinuxBinaryPath, opts.RegistryMirrors, opts.CommandModifier)
 }
 
 // NeedsDeferredInstall reports whether vllm on the current platform
@@ -74,7 +78,7 @@ func NeedsDeferredInstall() bool {
 
 // 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 newLinux(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, conf *Config, customBinaryPath string, registryMirrors []string) (inference.Backend, error) {
+func newLinux(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, conf *Config, customBinaryPath string, registryMirrors []string, commandModifier func(*exec.Cmd)) (inference.Backend, error) {
 	// If no config is provided, use the default configuration
 	if conf == nil {
 		conf = NewDefaultVLLMConfig()
@@ -88,6 +92,7 @@ func newLinux(log logging.Logger, modelManager *models.Manager, serverLog loggin
 		status:           inference.FormatNotInstalled(""),
 		customBinaryPath: customBinaryPath,
 		registryMirrors:  registryMirrors,
+		commandModifier:  commandModifier,
 	}, nil
 }
 
@@ -189,6 +194,7 @@ func (v *vLLM) Run(ctx context.Context, socket, model string, modelRef string, m
 		Args:            args,
 		Logger:          v.log,
 		ServerLogWriter: logging.NewWriter(v.serverLog),
+		CommandModifier: v.commandModifier,
 	})
 }
 

pkg/inference/backends/vllm/vllm_metal.go

@@ -53,11 +53,13 @@ type vllmMetal struct {
 	status string
 	// registryMirrors is the list of registry mirrors to try before registry-1.docker.io.
 	registryMirrors []string
+	// commandModifier, if non-nil, is applied to the server process before it starts.
+	commandModifier func(*exec.Cmd)
 }
 
 // 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 newMetal(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, customPythonPath string, registryMirrors []string) (inference.Backend, error) {
+func newMetal(log logging.Logger, modelManager *models.Manager, serverLog logging.Logger, customPythonPath string, registryMirrors []string, commandModifier func(*exec.Cmd)) (inference.Backend, error) {
 	homeDir, err := os.UserHomeDir()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get user home directory: %w", err)
@@ -72,6 +74,7 @@ func newMetal(log logging.Logger, modelManager *models.Manager, serverLog loggin
 		installDir:       installDir,
 		status:           inference.FormatNotInstalled(""),
 		registryMirrors:  registryMirrors,
+		commandModifier:  commandModifier,
 	}, nil
 }
 
@@ -250,6 +253,7 @@ func (v *vllmMetal) Run(ctx context.Context, socket, model string, modelRef stri
 		Logger:          v.log,
 		ServerLogWriter: logging.NewWriter(v.serverLog),
 		Env:             []string{"VLLM_HOST_IP=127.0.0.1"},
+		CommandModifier: v.commandModifier,
 	})
 }
 

pkg/routing/backends.go

@@ -1,6 +1,8 @@
 package routing
 
 import (
+	"os/exec"
+
 	"github.com/docker/model-runner/pkg/inference"
 	"github.com/docker/model-runner/pkg/inference/backends/diffusers"
 	"github.com/docker/model-runner/pkg/inference/backends/llamacpp"
@@ -39,6 +41,12 @@ type BackendsConfig struct {
 	// when pulling backend images. Populated from MODEL_RUNNER_REGISTRY_MIRRORS or
 	// injected by Docker Desktop from daemon.json registry-mirrors.
 	RegistryMirrors []string
+
+	// CommandModifier, if non-nil, is applied to every backend runner process
+	// immediately before it starts (see backends.RunnerConfig.CommandModifier).
+	// Embedders use it to customize process attributes such as credentials or
+	// environment; nil leaves the process unchanged.
+	CommandModifier func(*exec.Cmd)
 }
 
 // DefaultBackendDefs returns BackendDef entries for the configured backends.
@@ -54,13 +62,13 @@ func DefaultBackendDefs(cfg BackendsConfig) []BackendDef {
 
 	defs := []BackendDef{
 		{Name: llamacpp.Name, Deferred: llamacpp.NeedsDeferredInstall(), Init: func(mm *models.Manager) (inference.Backend, error) {
-			return llamacpp.New(cfg.Log, mm, sl(llamacpp.Name), cfg.LlamaCppPath, cfg.LlamaCppConfig, cfg.RegistryMirrors)
+			return llamacpp.New(cfg.Log, mm, sl(llamacpp.Name), cfg.LlamaCppPath, cfg.LlamaCppConfig, cfg.RegistryMirrors, cfg.CommandModifier)
 		}},
 	}
 
 	if cfg.IncludeMLX {
 		defs = append(defs, BackendDef{Name: mlx.Name, Init: func(mm *models.Manager) (inference.Backend, error) {
-			return mlx.New(cfg.Log, mm, sl(mlx.Name), nil, cfg.MLXPath)
+			return mlx.New(cfg.Log, mm, sl(mlx.Name), nil, cfg.MLXPath, cfg.CommandModifier)
 		}})
 	}
 
@@ -73,6 +81,7 @@ func DefaultBackendDefs(cfg BackendsConfig) []BackendDef {
 					LinuxBinaryPath: cfg.VLLMPath,
 					MetalPythonPath: cfg.VLLMMetalPath,
 					RegistryMirrors: cfg.RegistryMirrors,
+					CommandModifier: cfg.CommandModifier,
 				})
 			},
 		})
@@ -83,7 +92,7 @@ func DefaultBackendDefs(cfg BackendsConfig) []BackendDef {
 			Name:     diffusers.Name,
 			Deferred: true,
 			Init: func(mm *models.Manager) (inference.Backend, error) {
-				return diffusers.New(cfg.Log, mm, sl(diffusers.Name), nil, cfg.DiffusersPath, cfg.RegistryMirrors)
+				return diffusers.New(cfg.Log, mm, sl(diffusers.Name), nil, cfg.DiffusersPath, cfg.RegistryMirrors, cfg.CommandModifier)
 			},
 		})
 	}

pkg/server/server.go

@@ -179,7 +179,7 @@ func Run(ctx context.Context, cfg Config) error {
 				RegistryMirrors:  envconfig.RegistryMirrors(),
 			}),
 			routing.BackendDef{Name: sglang.Name, Init: func(mm *models.Manager) (inference.Backend, error) {
-				return sglang.New(log, mm, log.With("component", sglang.Name), nil, sglangServerPath)
+				return sglang.New(log, mm, log.With("component", sglang.Name), nil, sglangServerPath, nil)
 			}},
 		),
 		OnBackendError: func(name string, err error) {