De-duplicate code

In various places

Signed-off-by: Eric Curtin <[email protected]>

Author: ericcurtin

main.go

@@ -8,14 +8,14 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
-	"strings"
 	"syscall"
 
 	"github.com/docker/model-runner/pkg/distribution/transport/resumable"
 	"github.com/docker/model-runner/pkg/gpuinfo"
 	"github.com/docker/model-runner/pkg/inference"
 	"github.com/docker/model-runner/pkg/inference/backends/llamacpp"
 	"github.com/docker/model-runner/pkg/inference/backends/vllm"
+	"github.com/docker/model-runner/pkg/inference/common"
 	"github.com/docker/model-runner/pkg/inference/config"
 	"github.com/docker/model-runner/pkg/inference/memory"
 	"github.com/docker/model-runner/pkg/inference/models"
@@ -255,7 +255,7 @@ func createLlamaCppConfigFromEnv() config.BackendConfig {
 	}
 
 	// Split the string by spaces, respecting quoted arguments
-	args := splitArgs(argsStr)
+	args := common.SplitArgs(argsStr)
 
 	// Check for disallowed arguments
 	disallowedArgs := []string{"--model", "--host", "--embeddings", "--mmproj"}
@@ -273,29 +273,4 @@ func createLlamaCppConfigFromEnv() config.BackendConfig {
 	}
 }
 
-// splitArgs splits a string into arguments, respecting quoted arguments
-func splitArgs(s string) []string {
-	var args []string
-	var currentArg strings.Builder
-	inQuotes := false
-
-	for _, r := range s {
-		switch {
-		case r == '"' || r == '\'':
-			inQuotes = !inQuotes
-		case r == ' ' && !inQuotes:
-			if currentArg.Len() > 0 {
-				args = append(args, currentArg.String())
-				currentArg.Reset()
-			}
-		default:
-			currentArg.WriteRune(r)
-		}
-	}
 
-	if currentArg.Len() > 0 {
-		args = append(args, currentArg.String())
-	}
-
-	return args
-}

pkg/inference/backends/llamacpp/llamacpp.go

@@ -7,7 +7,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"io/fs"
 	"net/http"
 	"os"
 	"os/exec"
@@ -23,9 +22,9 @@ import (
 
 	"github.com/docker/model-runner/pkg/diskusage"
 	"github.com/docker/model-runner/pkg/inference"
+	"github.com/docker/model-runner/pkg/inference/common"
 	"github.com/docker/model-runner/pkg/inference/config"
 	"github.com/docker/model-runner/pkg/inference/models"
-	"github.com/docker/model-runner/pkg/internal/utils"
 	"github.com/docker/model-runner/pkg/logging"
 	"github.com/docker/model-runner/pkg/sandbox"
 	"github.com/docker/model-runner/pkg/tailbuffer"
@@ -149,8 +148,8 @@ func (l *llamaCpp) Run(ctx context.Context, socket, model string, _ string, mode
 		}
 	}
 
-	if err := os.RemoveAll(socket); err != nil && !errors.Is(err, fs.ErrNotExist) {
-		l.log.Warnf("failed to remove socket file %s: %w\n", socket, err)
+	if err := common.HandleSocketCleanup(socket); err != nil {
+		l.log.Warnf("failed to remove socket file %s: %v\n", socket, err)
 		l.log.Warnln("llama.cpp may not be able to start")
 	}
 
@@ -177,12 +176,7 @@ func (l *llamaCpp) Run(ctx context.Context, socket, model string, _ string, mode
 		}
 	}
 
-	// Sanitize args for safe logging
-	sanitizedArgs := make([]string, len(args))
-	for i, arg := range args {
-		sanitizedArgs[i] = utils.SanitizeForLog(arg)
-	}
-	l.log.Infof("llamaCppArgs: %v", sanitizedArgs)
+	common.SanitizedArgsLog(l.log, "llamaCppArgs", args)
 	tailBuf := tailbuffer.NewTailBuffer(1024)
 	serverLogStream := l.serverLog.Writer()
 	out := io.MultiWriter(serverLogStream, tailBuf)
@@ -210,24 +204,11 @@ func (l *llamaCpp) Run(ctx context.Context, socket, model string, _ string, mode
 
 	llamaCppErrors := make(chan error, 1)
 	go func() {
-		llamaCppErr := llamaCppSandbox.Command().Wait()
-		serverLogStream.Close()
-
-		errOutput := new(strings.Builder)
-		if _, err := io.Copy(errOutput, tailBuf); err != nil {
-			l.log.Warnf("failed to read server output tail: %w", err)
-		}
-
-		if len(errOutput.String()) != 0 {
-			llamaCppErr = fmt.Errorf("llama.cpp exit status: %w\nwith output: %s", llamaCppErr, errOutput.String())
-		} else {
-			llamaCppErr = fmt.Errorf("llama.cpp exit status: %w", llamaCppErr)
-		}
-
+		llamaCppErr := common.ProcessExitHandler(l.log, llamaCppSandbox, tailBuf, serverLogStream, socket)
 		llamaCppErrors <- llamaCppErr
 		close(llamaCppErrors)
-		if err := os.Remove(socket); err != nil && !errors.Is(err, fs.ErrNotExist) {
-			l.log.Warnf("failed to remove socket file %s on exit: %w\n", socket, err)
+		if err := common.HandleSocketCleanup(socket); err != nil {
+			l.log.Warnf("failed to remove socket file %s on exit: %v\n", socket, err)
 		}
 	}()
 	defer func() {

pkg/inference/backends/vllm/vllm.go

@@ -17,9 +17,9 @@ import (
 	"github.com/docker/model-runner/pkg/diskusage"
 	"github.com/docker/model-runner/pkg/distribution/types"
 	"github.com/docker/model-runner/pkg/inference"
+	"github.com/docker/model-runner/pkg/inference/common"
 	"github.com/docker/model-runner/pkg/inference/models"
 	"github.com/docker/model-runner/pkg/inference/platform"
-	"github.com/docker/model-runner/pkg/internal/utils"
 	"github.com/docker/model-runner/pkg/logging"
 	"github.com/docker/model-runner/pkg/sandbox"
 	"github.com/docker/model-runner/pkg/tailbuffer"
@@ -118,7 +118,7 @@ func (v *vLLM) Run(ctx context.Context, socket, model string, modelRef string, m
 		}
 	}
 
-	if err := os.RemoveAll(socket); err != nil && !errors.Is(err, fs.ErrNotExist) {
+	if err := common.HandleSocketCleanup(socket); err != nil {
 		v.log.Warnf("failed to remove socket file %s: %v\n", socket, err)
 		v.log.Warnln("vLLM may not be able to start")
 	}
@@ -151,12 +151,7 @@ func (v *vLLM) Run(ctx context.Context, socket, model string, modelRef string, m
 
 	args = append(args, "--served-model-name", model, modelRef)
 
-	// Sanitize args for safe logging
-	sanitizedArgs := make([]string, len(args))
-	for i, arg := range args {
-		sanitizedArgs[i] = utils.SanitizeForLog(arg)
-	}
-	v.log.Infof("vLLM args: %v", sanitizedArgs)
+	common.SanitizedArgsLog(v.log, "vLLM args", args)
 	tailBuf := tailbuffer.NewTailBuffer(1024)
 	serverLogStream := v.serverLog.Writer()
 	out := io.MultiWriter(serverLogStream, tailBuf)
@@ -184,23 +179,10 @@ func (v *vLLM) Run(ctx context.Context, socket, model string, modelRef string, m
 
 	vllmErrors := make(chan error, 1)
 	go func() {
-		vllmErr := vllmSandbox.Command().Wait()
-		serverLogStream.Close()
-
-		errOutput := new(strings.Builder)
-		if _, err := io.Copy(errOutput, tailBuf); err != nil {
-			v.log.Warnf("failed to read server output tail: %v", err)
-		}
-
-		if len(errOutput.String()) != 0 {
-			vllmErr = fmt.Errorf("vLLM exit status: %w\nwith output: %s", vllmErr, errOutput.String())
-		} else {
-			vllmErr = fmt.Errorf("vLLM exit status: %w", vllmErr)
-		}
-
+		vllmErr := common.ProcessExitHandler(v.log, vllmSandbox, tailBuf, serverLogStream, socket)
 		vllmErrors <- vllmErr
 		close(vllmErrors)
-		if err := os.Remove(socket); err != nil && !errors.Is(err, fs.ErrNotExist) {
+		if err := common.HandleSocketCleanup(socket); err != nil {
 			v.log.Warnf("failed to remove socket file %s on exit: %v\n", socket, err)
 		}
 	}()