diff --git a/pkg/inference/backends/diffusers/diffusers.go b/pkg/inference/backends/diffusers/diffusers.go index a966c6678..47bd27405 100644 --- a/pkg/inference/backends/diffusers/diffusers.go +++ b/pkg/inference/backends/diffusers/diffusers.go @@ -6,8 +6,6 @@ import ( "fmt" "net/http" "os" - "os/exec" - "path/filepath" "strings" "github.com/docker/model-runner/pkg/diskusage" @@ -28,7 +26,6 @@ const ( var ( ErrNotImplemented = errors.New("not implemented") ErrDiffusersNotFound = errors.New("diffusers package not installed") - ErrPythonNotFound = errors.New("python3 not found in PATH") ErrNoDDUFFile = errors.New("no DDUF file found in model bundle") ) @@ -91,37 +88,22 @@ func (d *diffusers) Install(_ context.Context, _ *http.Client) error { return ErrNotImplemented } - var pythonPath string - - // Use custom python path if specified - if d.customPythonPath != "" { - pythonPath = d.customPythonPath - } else { - venvPython := filepath.Join(diffusersDir, "bin", "python3") - pythonPath = venvPython - - if _, err := os.Stat(venvPython); err != nil { - // Fall back to system Python - systemPython, err := exec.LookPath("python3") - if err != nil { - d.status = ErrPythonNotFound.Error() - return ErrPythonNotFound - } - pythonPath = systemPython - } + pythonPath, err := backends.FindPythonPath(d.customPythonPath, diffusersDir) + if err != nil { + d.status = err.Error() + return err } - d.pythonPath = pythonPath // Check if diffusers is installed - if err := d.pythonCmd("-c", "import diffusers").Run(); err != nil { + if err := backends.NewPythonCmd(d.pythonPath, "-c", "import diffusers").Run(); err != nil { d.status = "diffusers package not installed" d.log.Warnf("diffusers package not found. Install with: uv pip install diffusers torch") return ErrDiffusersNotFound } // Get version - output, err := d.pythonCmd("-c", "import diffusers; print(diffusers.__version__)").Output() + output, err := backends.NewPythonCmd(d.pythonPath, "-c", "import diffusers; print(diffusers.__version__)").Output() if err != nil { d.log.Warnf("could not get diffusers version: %v", err) d.status = "running diffusers version: unknown" @@ -226,12 +208,3 @@ func (d *diffusers) GetRequiredMemoryForModel(_ context.Context, _ string, _ *in }, nil } -// pythonCmd creates an exec.Cmd that runs python with the given arguments. -// It uses the configured pythonPath if available, otherwise falls back to "python3". -func (d *diffusers) pythonCmd(args ...string) *exec.Cmd { - pythonBinary := "python3" - if d.pythonPath != "" { - pythonBinary = d.pythonPath - } - return exec.Command(pythonBinary, args...) -} diff --git a/pkg/inference/backends/mlx/mlx.go b/pkg/inference/backends/mlx/mlx.go index 6ba3683ed..ef787aa3a 100644 --- a/pkg/inference/backends/mlx/mlx.go +++ b/pkg/inference/backends/mlx/mlx.go @@ -80,22 +80,11 @@ func (m *mlx) Install(ctx context.Context, httpClient *http.Client) error { return errors.New("MLX is only available on macOS ARM64") } - var pythonPath string - - // Use custom python path if specified - if m.customPythonPath != "" { - pythonPath = m.customPythonPath - } else { - // Check if Python 3 is available - var err error - pythonPath, err = exec.LookPath("python3") - if err != nil { - m.status = ErrStatusNotFound.Error() - return ErrStatusNotFound - } + pythonPath, err := backends.FindPythonPath(m.customPythonPath, "") + if err != nil { + m.status = ErrStatusNotFound.Error() + return ErrStatusNotFound } - - // Store the python path for later use m.pythonPath = pythonPath // Check if mlx-lm package is installed by attempting to import it diff --git a/pkg/inference/backends/python.go b/pkg/inference/backends/python.go new file mode 100644 index 000000000..7b6a8d853 --- /dev/null +++ b/pkg/inference/backends/python.go @@ -0,0 +1,42 @@ +package backends + +import ( + "errors" + "os" + "os/exec" + "path/filepath" +) + +// ErrPythonNotFound is returned when python3 cannot be found. +var ErrPythonNotFound = errors.New("python3 not found in PATH") + +// FindPythonPath returns the path to a python3 binary. +// If customPath is non-empty, it is returned directly. +// If envDir is non-empty and envDir/bin/python3 exists, that is returned. +// Otherwise, python3 is looked up in PATH. Returns ErrPythonNotFound if none found. +func FindPythonPath(customPath, envDir string) (string, error) { + if customPath != "" { + return customPath, nil + } + if envDir != "" { + venvPython := filepath.Join(envDir, "bin", "python3") + if _, err := os.Stat(venvPython); err == nil { + return venvPython, nil + } + } + systemPython, err := exec.LookPath("python3") + if err != nil { + return "", ErrPythonNotFound + } + return systemPython, nil +} + +// NewPythonCmd creates an exec.Cmd that runs python3 with the given arguments. +// If pythonPath is empty, "python3" is used. +func NewPythonCmd(pythonPath string, args ...string) *exec.Cmd { + binary := "python3" + if pythonPath != "" { + binary = pythonPath + } + return exec.Command(binary, args...) +} diff --git a/pkg/inference/backends/sglang/sglang.go b/pkg/inference/backends/sglang/sglang.go index 802c48802..de727218e 100644 --- a/pkg/inference/backends/sglang/sglang.go +++ b/pkg/inference/backends/sglang/sglang.go @@ -6,8 +6,6 @@ import ( "fmt" "net/http" "os" - "os/exec" - "path/filepath" "strings" "github.com/docker/model-runner/pkg/diskusage" @@ -27,7 +25,6 @@ const ( var ( ErrNotImplemented = errors.New("not implemented") ErrSGLangNotFound = errors.New("sglang package not installed") - ErrPythonNotFound = errors.New("python3 not found in PATH") ) // sglang is the SGLang-based backend implementation. @@ -86,37 +83,22 @@ func (s *sglang) Install(_ context.Context, _ *http.Client) error { return ErrNotImplemented } - var pythonPath string - - // Use custom python path if specified - if s.customPythonPath != "" { - pythonPath = s.customPythonPath - } else { - venvPython := filepath.Join(sglangDir, "bin", "python3") - pythonPath = venvPython - - if _, err := os.Stat(venvPython); err != nil { - // Fall back to system Python - systemPython, err := exec.LookPath("python3") - if err != nil { - s.status = ErrPythonNotFound.Error() - return ErrPythonNotFound - } - pythonPath = systemPython - } + pythonPath, err := backends.FindPythonPath(s.customPythonPath, sglangDir) + if err != nil { + s.status = err.Error() + return err } - s.pythonPath = pythonPath // Check if sglang is installed - if err := s.pythonCmd("-c", "import sglang").Run(); err != nil { + if err := backends.NewPythonCmd(s.pythonPath, "-c", "import sglang").Run(); err != nil { s.status = "sglang package not installed" s.log.Warnf("sglang package not found. Install with: uv pip install sglang") return ErrSGLangNotFound } // Get version - output, err := s.pythonCmd("-c", "import sglang; print(sglang.__version__)").Output() + output, err := backends.NewPythonCmd(s.pythonPath, "-c", "import sglang; print(sglang.__version__)").Output() if err != nil { s.log.Warnf("could not get sglang version: %v", err) s.status = "running sglang version: unknown" @@ -204,12 +186,3 @@ func (s *sglang) GetRequiredMemoryForModel(_ context.Context, _ string, _ *infer }, nil } -// pythonCmd creates an exec.Cmd that runs python with the given arguments. -// It uses the configured pythonPath if available, otherwise falls back to "python3". -func (s *sglang) pythonCmd(args ...string) *exec.Cmd { - pythonBinary := "python3" - if s.pythonPath != "" { - pythonBinary = s.pythonPath - } - return exec.Command(pythonBinary, args...) -}