Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion pkg/inference/backends/diffusers/diffusers.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -81,6 +83,7 @@ func New(log logging.Logger, modelManager *models.Manager, serverLog logging.Log
customPythonPath: customPythonPath,
installDir: installDir,
registryMirrors: registryMirrors,
commandModifier: commandModifier,
}, nil
}

Expand Down Expand Up @@ -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,
})
}

Expand Down
5 changes: 5 additions & 0 deletions pkg/inference/backends/llamacpp/llamacpp.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 {
Expand All @@ -87,6 +90,7 @@ func New(
status: inference.FormatNotInstalled(""),
config: conf,
registryMirrors: registryMirrors,
commandModifier: commandModifier,
}, nil
}

Expand Down Expand Up @@ -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,
})
}

Expand Down
6 changes: 5 additions & 1 deletion pkg/inference/backends/mlx/mlx.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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
}

Expand Down Expand Up @@ -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,
})
}

Expand Down
15 changes: 15 additions & 0 deletions pkg/inference/backends/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +50 to 58

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion: Clarify and possibly adjust environment handling semantics before CommandModifier runs.

With the new logic, both nil and empty Env are materialized to os.Environ() before CommandModifier runs. This matches the Env comment but prevents CommandModifier from distinguishing between "inherit parent" vs. "start from empty and selectively add vars." Embedders who want a minimal env must now overwrite command.Env inside CommandModifier. Please either document this explicitly in RunnerConfig.CommandModifier or change the flow so CommandModifier sees the raw Env configuration and can decide how to materialize or replace it.

Suggested implementation:

	// 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).
	//
	// When Runner constructs the *exec.Cmd, it materializes Env into cmd.Env
	// based on this policy before CommandModifier runs; embedders will not
	// be able to distinguish "inherit parent" from "start from empty and
	// selectively add vars" via Env alone.
	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.
	//
	// Note: by the time CommandModifier runs, cmd.Env has already been
	// materialized from RunnerConfig.Env according to the inheritance policy
	// above. Embedders that require a minimal or custom environment should
	// overwrite cmd.Env inside CommandModifier rather than relying on a
	// nil or empty RunnerConfig.Env.
	CommandModifier func(*exec.Cmd)
}
			if len(config.Env) > 0 {
				command.Env = append(os.Environ(), config.Env...)
			}
			// At this point, command.Env reflects the materialized environment
			// based on RunnerConfig.Env (inherit + overlay). CommandModifier
			// may further adjust or replace command.Env if a different policy
			// is desired by the embedder.
			if config.CommandModifier != nil {

If there are other constructors or helpers that build *exec.Cmd using RunnerConfig.Env and CommandModifier, mirror the same documentation and behavior there so embedders consistently see a materialized environment before CommandModifier runs.

Expand Down Expand Up @@ -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)
}
Comment thread
mat007 marked this conversation as resolved.
},
config.SandboxPath,
config.BinaryPath,
Expand Down
6 changes: 5 additions & 1 deletion pkg/inference/backends/sglang/sglang.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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
}

Expand Down Expand Up @@ -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,
})
}

Expand Down
20 changes: 13 additions & 7 deletions pkg/inference/backends/vllm/vllm.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"io/fs"
"net/http"
"os"
"os/exec"
"path/filepath"
"strings"

Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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()
Expand All @@ -88,6 +92,7 @@ func newLinux(log logging.Logger, modelManager *models.Manager, serverLog loggin
status: inference.FormatNotInstalled(""),
customBinaryPath: customBinaryPath,
registryMirrors: registryMirrors,
commandModifier: commandModifier,
}, nil
}

Expand Down Expand Up @@ -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,
})
}

Expand Down
6 changes: 5 additions & 1 deletion pkg/inference/backends/vllm/vllm_metal.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -72,6 +74,7 @@ func newMetal(log logging.Logger, modelManager *models.Manager, serverLog loggin
installDir: installDir,
status: inference.FormatNotInstalled(""),
registryMirrors: registryMirrors,
commandModifier: commandModifier,
}, nil
}

Expand Down Expand Up @@ -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,
})
}

Expand Down
15 changes: 12 additions & 3 deletions pkg/routing/backends.go
Original file line number Diff line number Diff line change
@@ -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"
Expand Down Expand Up @@ -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.
Expand All @@ -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)
}})
}

Expand All @@ -73,6 +81,7 @@ func DefaultBackendDefs(cfg BackendsConfig) []BackendDef {
LinuxBinaryPath: cfg.VLLMPath,
MetalPythonPath: cfg.VLLMMetalPath,
RegistryMirrors: cfg.RegistryMirrors,
CommandModifier: cfg.CommandModifier,
})
},
})
Expand All @@ -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)
},
})
}
Expand Down
2 changes: 1 addition & 1 deletion pkg/server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Loading