diff --git a/pkg/inference/backends/diffusers/diffusers.go b/pkg/inference/backends/diffusers/diffusers.go index 1cda1967c..4b6517bce 100644 --- a/pkg/inference/backends/diffusers/diffusers.go +++ b/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, }) } diff --git a/pkg/inference/backends/llamacpp/llamacpp.go b/pkg/inference/backends/llamacpp/llamacpp.go index 16f1ff912..2b25d79db 100644 --- a/pkg/inference/backends/llamacpp/llamacpp.go +++ b/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, }) } diff --git a/pkg/inference/backends/mlx/mlx.go b/pkg/inference/backends/mlx/mlx.go index cd9dd8694..5431836bf 100644 --- a/pkg/inference/backends/mlx/mlx.go +++ b/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, }) } diff --git a/pkg/inference/backends/runner.go b/pkg/inference/backends/runner.go index 8e0608053..811ffd6d4 100644 --- a/pkg/inference/backends/runner.go +++ b/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, diff --git a/pkg/inference/backends/sglang/sglang.go b/pkg/inference/backends/sglang/sglang.go index 2e083ba99..f3338ae18 100644 --- a/pkg/inference/backends/sglang/sglang.go +++ b/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, }) } diff --git a/pkg/inference/backends/vllm/vllm.go b/pkg/inference/backends/vllm/vllm.go index 38f4152a9..cdfb7a469 100644 --- a/pkg/inference/backends/vllm/vllm.go +++ b/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, }) } diff --git a/pkg/inference/backends/vllm/vllm_metal.go b/pkg/inference/backends/vllm/vllm_metal.go index ddf752688..4caa8ec27 100644 --- a/pkg/inference/backends/vllm/vllm_metal.go +++ b/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, }) } diff --git a/pkg/routing/backends.go b/pkg/routing/backends.go index 79c9d6466..451e7484f 100644 --- a/pkg/routing/backends.go +++ b/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) }, }) } diff --git a/pkg/server/server.go b/pkg/server/server.go index eeff6b2ad..b47fbd275 100644 --- a/pkg/server/server.go +++ b/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) {