-
Notifications
You must be signed in to change notification settings - Fork 82
De-duplicate code #387
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
De-duplicate code #387
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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) | ||
| } | ||
| }() | ||
|
Comment on lines
181
to
188
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The call to To avoid duplicate operations and improve clarity, the explicit cleanup call should be removed. If go func() {
vllmErr := common.ProcessExitHandler(v.log, vllmSandbox, tailBuf, serverLogStream, socket)
vllmErrors <- vllmErr
close(vllmErrors)
}() |
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,96 @@ | ||
| package common | ||
|
|
||
| import ( | ||
| "errors" | ||
| "io" | ||
| "os" | ||
| "strings" | ||
|
|
||
| "github.com/docker/model-runner/pkg/internal/utils" | ||
| "github.com/docker/model-runner/pkg/logging" | ||
| "github.com/docker/model-runner/pkg/sandbox" | ||
| ) | ||
|
|
||
| // SanitizedArgsLog logs command arguments with sanitization for safe logging | ||
| func SanitizedArgsLog(log logging.Logger, label string, args []string) { | ||
| sanitizedArgs := make([]string, len(args)) | ||
| for i, arg := range args { | ||
| sanitizedArgs[i] = utils.SanitizeForLog(arg) | ||
| } | ||
| log.Infof("%s: %v", label, sanitizedArgs) | ||
| } | ||
|
|
||
| // HandleSocketCleanup removes the socket file at the given path, ignoring if it doesn't exist | ||
| func HandleSocketCleanup(socket string) error { | ||
| if err := os.RemoveAll(socket); err != nil && !errors.Is(err, os.ErrNotExist) { | ||
| return err | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| // ProcessExitHandler handles the exit of a backend process and captures its output | ||
| func ProcessExitHandler( | ||
| log logging.Logger, | ||
| sandboxInstance sandbox.Sandbox, | ||
| tailBuf io.ReadWriter, | ||
| serverLogStream io.Closer, | ||
| socket string, | ||
| ) error { | ||
| serverLogStream.Close() | ||
|
|
||
| errOutput := new(strings.Builder) | ||
| if _, err := io.Copy(errOutput, tailBuf); err != nil { | ||
| log.Warnf("failed to read server output tail: %v", err) | ||
| } | ||
|
|
||
| cmdErr := sandboxInstance.Command().Wait() | ||
| outputStr := errOutput.String() | ||
| if len(outputStr) != 0 { | ||
| return &BackendExitError{ | ||
| Err: cmdErr, | ||
| Output: outputStr, | ||
| } | ||
| } else { | ||
| return cmdErr | ||
| } | ||
| } | ||
|
|
||
| // BackendExitError represents an error when a backend process exits | ||
| type BackendExitError struct { | ||
| Err error | ||
| Output string | ||
| } | ||
|
|
||
| func (e *BackendExitError) Error() string { | ||
| if e.Output != "" { | ||
| return e.Err.Error() + "\nwith output: " + e.Output | ||
| } | ||
| return e.Err.Error() | ||
| } | ||
|
|
||
| // 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 | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The call to
common.HandleSocketCleanup(socket)on lines 210-212 appears to be redundant, as thesocketvariable is already passed tocommon.ProcessExitHandleron line 207. This implies thatProcessExitHandleris responsible for the socket cleanup.To avoid duplicate operations and improve clarity, the explicit cleanup call should be removed. If
ProcessExitHandlerdoes not handle cleanup, its signature should be changed to not accept asocketparameter.