feat: preload models on configure and in interactive mode #629
Conversation
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 3 issues, and left some high level feedback:
- The
preloadOnlyKeyused withcontext.WithValueis a plain string, which risks collisions with other context values; consider defining a dedicated unexported type (e.g.type preloadOnlyKeyType struct{}) and using a value of that type as the key. - In the
Configurehandler’s preload goroutine, errors fromjson.Marshalandhttp.NewRequestWithContextare ignored; even if the preload is best-effort, it would be helpful to at least log these failures so issues with model preloading are diagnosable.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The `preloadOnlyKey` used with `context.WithValue` is a plain string, which risks collisions with other context values; consider defining a dedicated unexported type (e.g. `type preloadOnlyKeyType struct{}`) and using a value of that type as the key.
- In the `Configure` handler’s preload goroutine, errors from `json.Marshal` and `http.NewRequestWithContext` are ignored; even if the preload is best-effort, it would be helpful to at least log these failures so issues with model preloading are diagnosable.
## Individual Comments
### Comment 1
<location> `pkg/inference/scheduling/http_handler.go:24-28` </location>
<code_context>
"github.com/docker/model-runner/pkg/middleware"
)
+const preloadOnlyKey string = "preloadOnly"
+
// HTTPHandler handles HTTP requests for the scheduler.
</code_context>
<issue_to_address>
**issue (bug_risk):** Use a dedicated context key type instead of a plain string for `preloadOnlyKey`.
A plain string context key can collide with keys from other packages (including stdlib/middleware). Define a private key type instead (e.g. `type preloadOnlyKeyType struct{}` with `var preloadOnlyKey preloadOnlyKeyType`) to make this key uniquely identifiable while still supporting the `X-Preload-Only` header path.
</issue_to_address>
### Comment 2
<location> `pkg/inference/scheduling/http_handler.go:382-391` </location>
<code_context>
+ // Preload the model in the background by calling handleOpenAIInference with preload-only context.
+ // This makes Compose preload the model as well as it calls `configure` by default.
+ go func() {
+ preloadBody, _ := json.Marshal(OpenAIInferenceRequest{Model: configureRequest.Model})
+ ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+ defer cancel()
+ preloadReq, _ := http.NewRequestWithContext(
+ context.WithValue(ctx, preloadOnlyKey, true),
+ http.MethodPost,
+ inference.InferencePrefix+"/v1/chat/completions",
+ bytes.NewReader(preloadBody),
+ )
+ preloadReq.Header.Set("User-Agent", r.UserAgent())
+ if backend != nil {
+ preloadReq.SetPathValue("backend", backend.Name())
+ }
+ h.handleOpenAIInference(httptest.NewRecorder(), preloadReq)
+ }()
+
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Avoid silently discarding errors when building the preload request in the goroutine.
`json.Marshal` and `http.NewRequestWithContext` both ignore their errors, so any failure just skips the preload with no signal. Even for best-effort behavior, we should at least log these errors (or return early) so malformed config or context issues are visible during startup.
</issue_to_address>
### Comment 3
<location> `cmd/cli/commands/run.go:763-765` </location>
<code_context>
return nil
}
+ // For interactive mode, eagerly load the model in the background
+ // while the user types their first query
+ go func() {
+ _ = desktopClient.Preload(context.Background(), model)
+ }()
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Consider using `cmd.Context()` instead of `context.Background()` for the interactive preload goroutine.
Using `context.Background()` lets the preload continue even after the CLI command exits (for example, if the user cancels quickly). Passing `cmd.Context()` (or a derived context) into `Preload` would ensure the preload is cancelled with the command lifecycle instead.
Suggested implementation:
```golang
// For interactive mode, eagerly load the model in the background
// while the user types their first query
go func() {
_ = desktopClient.Preload(cmd.Context(), model)
}()
```
If `context` is only imported for this call site elsewhere in the file, you can remove the unused `context` import after this change. Otherwise, no additional code changes are required.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Summary of ChangesHello @doringeman, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a mechanism to preload models into memory, addressing the issue of high first-request latency for backends like vLLM. It achieves this by adding a Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a model preloading feature to improve first-request latency, which is a valuable enhancement for user experience. The implementation is well-structured, adding a Preload method to the client and using a special header/context key on the backend to short-circuit inference after model loading. My review focuses on improving error handling in a few areas, particularly for background tasks where failures could otherwise go unnoticed. Overall, the changes are solid and well-thought-out.
Add preload support to reduce first-request latency for slow-starting backends like vLLM. It works by having `configure`, which is also used by Compose, call `handleOpenAIInference` with a `preloadOnly` context key in order to stop right after loading the model. The CLI side calls `handleOpenAIInference` directly with `X-Preload-Only` set. Fixes docker#623. Signed-off-by: Dorin Geman <dorin.geman@docker.com>
doringeman
force-pushed
the
preload-optimization
branch
from
bbacaad to
bdf158b
Compare
1 participant
Add preload support to reduce first-request latency for slow-starting backends like vLLM.
It works by having
configure, which is also used by Compose, callhandleOpenAIInferencewith apreloadOnlycontext key in order to stop right after loading the model. The CLI side callshandleOpenAIInferencedirectly withX-Preload-Onlyset.Fixes #623.
To test: