De-deduplicate code

pkg/inference/backends/diffusers/diffusers.go

@@ -6,8 +6,6 @@
 	"fmt"
 	"net/http"
 	"os"
-	"os/exec"
-	"path/filepath"
 	"strings"
 
 	"github.com/docker/model-runner/pkg/diskusage"
@@ -28,7 +26,6 @@
 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 @@
 		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"
@@ -225,13 +207,4 @@
 		VRAM: 6 * 1024 * 1024 * 1024, // 6GB VRAM (average estimate)
 	}, 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...)
-}

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

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...)
+}

pkg/inference/backends/sglang/sglang.go

@@ -6,8 +6,6 @@
 	"fmt"
 	"net/http"
 	"os"
-	"os/exec"
-	"path/filepath"
 	"strings"
 
 	"github.com/docker/model-runner/pkg/diskusage"
@@ -27,7 +25,6 @@
 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 @@
 		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"
@@ -203,13 +185,4 @@
 		VRAM: 1,
 	}, 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...)
-}