feat: support dragonfly endpoint for model fetch #25
Conversation
Summary of ChangesHello @imeoer, 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 primarily focuses on enhancing the model fetching mechanism and improving progress reporting within the system. It includes a significant update of Go module dependencies to leverage newer features and maintain security. Key functional changes involve modifying model artifact handling to provide total layer counts, which is then utilized by an improved progress reporting system. Additionally, it introduces support for a Dragonfly endpoint, enabling more efficient and distributed model fetching. 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. 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 support for a Dragonfly endpoint for model fetching and significantly refactors the model artifact handling and progress reporting logic. The dependency updates in go.mod and go.sum bring in newer versions of many packages, which is good for security and compatibility. The changes to pkg/service/model.go, pkg/service/puller.go, and pkg/status/hook.go aim to provide more accurate progress tracking. My review focuses on ensuring the new progress reporting logic is robust. I've identified a few areas for improvement to correctly handle progress totals for filtered downloads and to account for potential duplicate layers in model artifacts, which could affect the accuracy of the progress display.
pkg/service/puller.go
Outdated
| logger.WithContext(ctx).Infof( | ||
| "fetching model without weights: %s, file patterns: %s", | ||
| reference, strings.Join(patterns, ", "), | ||
| "fetching model without weights: %s, file patterns: %s, total files: %d", | ||
| reference, strings.Join(patterns, ", "), total, | ||
| ) | ||
|
|
||
| p.hook.SetTotal(total) |
There was a problem hiding this comment.
The total variable holds the total number of layers in the artifact, but when excludeModelWeights is true, only a subset of layers (patterns) is fetched. By setting the progress hook's total to the unfiltered count (p.hook.SetTotal(total)), the progress bar will not show accurate progress (e.g., it will show 5/20 when all 5 requested files are complete, instead of 5/5). The total for the progress hook should be the number of files that are actually being fetched.
Additionally, the log message is misleading. It would be clearer to state both the number of files being fetched and the total files in the model.
| logger.WithContext(ctx).Infof( | |
| "fetching model without weights: %s, file patterns: %s", | |
| reference, strings.Join(patterns, ", "), | |
| "fetching model without weights: %s, file patterns: %s, total files: %d", | |
| reference, strings.Join(patterns, ", "), total, | |
| ) | |
| p.hook.SetTotal(total) | |
| logger.WithContext(ctx).Infof( | |
| "fetching %d of %d model files without weights: %s, file patterns: %s", | |
| len(patterns), total, reference, strings.Join(patterns, ", "), | |
| ) | |
| p.hook.SetTotal(len(patterns)) |
| } | ||
|
|
||
| return layers, nil | ||
| return layers, len(m.artifact.Layers), nil |
There was a problem hiding this comment.
The function returns len(m.artifact.Layers) as the total count, which doesn't account for potential duplicate layers in the artifact. This can lead to an incorrect total being reported. The GetSize function in this file correctly deduplicates layers by digest, and a similar approach should be used here to provide an accurate count of unique layers.
| return layers, len(m.artifact.Layers), nil | |
| digestMap := make(map[string]bool) | |
| for idx := range m.artifact.Layers { | |
| digestMap[m.artifact.Layers[idx].Digest] = true | |
| } | |
| return layers, len(digestMap), nil |
| } | ||
|
|
||
| if h.manifest != nil { | ||
| return len(h.manifest.Layers) |
There was a problem hiding this comment.
The previous implementation for calculating the total number of layers correctly handled potential duplicates by using a map of digests. This new implementation len(h.manifest.Layers) is a regression as it does not account for duplicate layers in the manifest, which could lead to an incorrect total for progress reporting. It would be safer to restore the logic of counting unique layers.
digestMap := make(map[digest.Digest]bool)
for _, layer := range h.manifest.Layers {
digestMap[layer.Digest] = true
}
return len(digestMap)
Signed-off-by: imeoer <yansong.ys@antgroup.com>
1 participant
This pull request updates the
go.moddependencies and enhances the model artifact handling and progress reporting logic. The most significant changes include updating multiple dependencies, improving the way model artifact layers and patterns are retrieved, and refining progress tracking during model pulls.Dependency updates:
go.mod, includinggithub.com/modelpack/modctl,github.com/prometheus/client_golang,go.opentelemetry.io/otel,google.golang.org/grpc, and many others. These upgrades improve security, compatibility, and feature support. [1] [2] [3] [4]Model artifact handling:
ModelArtifact.getLayersto return both the filtered layers and the total number of artifact layers, allowing downstream consumers to know the total file count regardless of filtering. [1] [2]ModelArtifact.GetPatternsto return the total layer count along with the filtered patterns, and adjusted its callers accordingly. [1] [2] [3]Progress reporting improvements:
Hookstruct and related logic to track and report the total number of files/layers more accurately, using a newtotalfield that can be set externally. [1] [2] [3]