[draft_model] Fix OCI mount options for draft_model#2225
Conversation
xbezdick
requested review from
bmahabirbu,
cgruver,
engelmi,
jhjaggars,
maxamillion,
rhatdan and
swarajpande5
as code owners
Reviewer's GuideRefactors model identifier extraction to accept an optional model argument and fixes draft model OCI mount behavior by mounting individual blob files instead of a single bind mount, while skipping incompatible multimodal-specific files for draft models. File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
Summary of ChangesHello @xbezdick, 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 significantly improves the handling of draft models within the system by refining how they are mounted in OCI containers. The core purpose is to enable robust support for draft models that consist of multiple files and to ensure that the container mount options are configured accurately. It also introduces a mechanism to intelligently ignore specific file types, such as 'mmproj' and 'chat_template', when processing draft models, which is crucial for preventing conflicts and facilitating advanced features like speculative decoding. Highlights
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
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey there - I've reviewed your changes - here's some feedback:
- The draft model files are all mounted directly under
MNT_DIR/{file.name}, which will collide with the main model mounts when filenames overlap (e.g., tokenizer/config/weights for 440B + 30B); consider mounting the draft model into its own subdirectory (e.g.,MNT_DIR/draft/{file.name}) and updating consumers accordingly. - Now that
extract_model_identifierstakes an optionalmodelargument in the base class, you may want to add an explicit type hint and default consistently across all overrides (and consider a docstring update) to make the new calling pattern clearer and avoid accidental misuse.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The draft model files are all mounted directly under `MNT_DIR/{file.name}`, which will collide with the main model mounts when filenames overlap (e.g., tokenizer/config/weights for 440B + 30B); consider mounting the draft model into its own subdirectory (e.g., `MNT_DIR/draft/{file.name}`) and updating consumers accordingly.
- Now that `extract_model_identifiers` takes an optional `model` argument in the base class, you may want to add an explicit type hint and default consistently across all overrides (and consider a docstring update) to make the new calling pattern clearer and avoid accidental misuse.Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
![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 correctly refactors extract_model_identifiers to be reusable for draft models and fixes the mounting logic for draft models, especially for multi-file models. The changes are logical and address the described issues. However, I've found a critical syntax error in ramalama/transports/base.py due to incorrect indentation that will prevent the code from running. I've also provided a suggestion to improve code quality by using an enum for type checking. After addressing the critical issue, this PR should be good to merge.
|
Thanks @xbezdick @olliewalsh @ieaves @engelmi PTAL |
There was a problem hiding this comment.
This looks really good. One of the AI reviews identified
The draft model files are all mounted directly under MNT_DIR/{file.name}, which will collide with the main model mounts when filenames overlap
This looks like a real issue and namespacing the mount directory to something like {MNT_DIR}/drafts/{file.name} seems like a reasonable approach.
I think you'll need to modify RamalamaModelContext.draft_model_path to reflect the namespaced mount path and it would be amazing to add a test for this scenario as well.
xbezdick
force-pushed
the
main
branch
4 times, most recently
from
55d3869 to
b31ccca
Compare
|
A friendly reminder that this PR had no activity for 30 days. |
|
@xbezdick still working on this? |
|
@xbezdick is it possible to add unit or e2e test coverage for this? Can run the tests on your fork if enable github actions (ensure main branch is up to date first) |
xbezdick
force-pushed
the
main
branch
3 times, most recently
from
a65e914 to
817ccb0
Compare
When supplying draft model the old code pretty much tried to bindmount model with both src and destination in the end translating to same location. Also there was no support for mulitple file models so when one tried to use 440B as main model and 30B as draft model it would not be possible. To prevent file collisions for example for mmproj and chat_template we use namespaced directory. As of now draft_models won't work with multimodal models but we can still do speculative decoding when disabling vision by passing --no-mmproj. Signed-off-by: Lukas Bezdicka <lbezdick@redhat.com>
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe changes refactor draft model mounting in the transport layer by replacing a single mount point with per-file mounts under a Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@ramalama/transports/base.py`:
- Around line 450-452: The draft model's tag is being parsed through the primary
transport by calling extract_model_identifiers(model=self.draft_model.model);
change this to resolve the tag using the draft model object itself (pass
self.draft_model into extract_model_identifiers) so the draft transport's rules
are used, then call self.draft_model.model_store.get_ref_file(draft_tag); update
any variable usages (draft_tag) accordingly to avoid NoRefFileFound when main
and draft transports differ.
In `@ramalama/transports/url.py`:
- Around line 62-70: In extract_model_identifiers, the file:// branch ignores
the model override and always uses self.model; change it to use the provided
model parameter when present (falling back to self.model if model is None) so
that extract_model_identifiers(model=...) returns identifiers for the explicit
file path; update the file-branch expressions (Path(...).name and
Path(...).parent) to use model (or the fallback) and preserve the existing
normalization call normalize_host_path_for_container(...).removeprefix("/") and
the returned tuple format.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: d03b8511-ef6b-427a-bc8a-2efa58a6819c
📒 Files selected for processing (4)
ramalama/transports/base.pyramalama/transports/huggingface.pyramalama/transports/ollama.pyramalama/transports/url.py
| # Get the model tag for draft model | ||
| _, draft_tag, _ = self.extract_model_identifiers(model=self.draft_model.model) | ||
| ref_file_draft = self.draft_model.model_store.get_ref_file(draft_tag) |
There was a problem hiding this comment.
Resolve the draft tag with self.draft_model, not self.
This reparses self.draft_model.model through the primary model's transport rules. If the main model and draft model use different transports, you can derive the wrong tag and fail with NoRefFileFound even though the draft snapshot exists.
Suggested fix
- _, draft_tag, _ = self.extract_model_identifiers(model=self.draft_model.model)
- ref_file_draft = self.draft_model.model_store.get_ref_file(draft_tag)
+ draft_tag = self.draft_model.model_tag
+ ref_file_draft = self.draft_model.model_store.get_ref_file(draft_tag)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@ramalama/transports/base.py` around lines 450 - 452, The draft model's tag is
being parsed through the primary transport by calling
extract_model_identifiers(model=self.draft_model.model); change this to resolve
the tag using the draft model object itself (pass self.draft_model into
extract_model_identifiers) so the draft transport's rules are used, then call
self.draft_model.model_store.get_ref_file(draft_tag); update any variable usages
(draft_tag) accordingly to avoid NoRefFileFound when main and draft transports
differ.
| def extract_model_identifiers(self, model=None): | ||
| if self.type == "file": | ||
| return ( | ||
| Path(self.model).name, | ||
| "latest", | ||
| normalize_host_path_for_container(str(Path(self.model).parent)).removeprefix("/"), | ||
| ) | ||
|
|
||
| model_name, model_tag, model_organization = super().extract_model_identifiers() | ||
| model_name, model_tag, model_organization = super().extract_model_identifiers(model) |
There was a problem hiding this comment.
Honor the model override in the file:// branch.
Lines 63-68 still read self.model, so extract_model_identifiers(model=...) returns identifiers for the bound instance instead of the explicit file path. That breaks the new reuse path for alternate or draft file models.
Suggested fix
def extract_model_identifiers(self, model=None):
+ if model is None:
+ model = self.model
if self.type == "file":
return (
- Path(self.model).name,
+ Path(model).name,
"latest",
- normalize_host_path_for_container(str(Path(self.model).parent)).removeprefix("/"),
+ normalize_host_path_for_container(str(Path(model).parent)).removeprefix("/"),
)
model_name, model_tag, model_organization = super().extract_model_identifiers(model)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@ramalama/transports/url.py` around lines 62 - 70, In
extract_model_identifiers, the file:// branch ignores the model override and
always uses self.model; change it to use the provided model parameter when
present (falling back to self.model if model is None) so that
extract_model_identifiers(model=...) returns identifiers for the explicit file
path; update the file-branch expressions (Path(...).name and Path(...).parent)
to use model (or the fallback) and preserve the existing normalization call
normalize_host_path_for_container(...).removeprefix("/") and the returned tuple
format.
xbezdick
temporarily deployed
to
macos-installer
GitHub Actions
Inactive
4 participants
When supplying draft model the old code pretty much tried to bindmount model with both src and destination in the end translating to same location. Also there was no support for mulitple file models so when one tried to use 440B as main model and 30B as draft model it would not be possible.
Lastly mmproj and chat_template would translate to the same location as main model so we ignore them. As of now draft_models won't work with multimodal models but we can still do speculative decoding when disabling vision by passing --no-mmproj.
Summary by Sourcery
Support correct identification and mounting of draft models, including multi-file models, while reusing the existing model identifier parsing across transports.
Bug Fixes:
Enhancements: