RFC: Refactor the ModelManager

[bitgamma]

Proposal size

Major feature (spans multiple well-scoped PRs)

Updates

No response

User story

The goal is to restructure the model manager to be able to easily add data sources, catalogs and storage options. Another goal is to split this big components into more manageable and hopefully easier to unit-test subcomponents.

High-level design

ModelManager does a lot of stuff and is one of the core components of lemonade. The things it handles are:

1. Model catalogs (server, user models, cloud models, HF and modelscope lookup, extra directories)
2. Checkpoint resolution: mapping a checkpoint name to a file on the disk
3. Model indexer: mapping which models from the catalog can be resolved to actual checkpoints
4. Checkpoint downloader: allows downloading checkpoints from different data sources.

All these concerns are inherently related but currently there is no clear separation, which result in model_manager.cpp being over 6.7k LoC.

My proposal is to refactor this into a modular system, where each concern is handled through interfaces. Let's call them:

• Catalog: gives way to discover models. Some catalogs might have multiple instance with different parameters, so we could allow downloading/sharing remote catalogs in the same json format as the user/server models. Some catalogs can be listed (locally stored ones) while other only searched (HF, ModelScope, etc)
• CheckpointResolver: Locates checkpoints on disk. The concrete instances now would be HuggingFaceResolver (models stored as a huggingface hub dir) and PlainDirectoryResolver (the extra model dir). Again, multiple instances of each could be supported for maximum flexibility.
• DownloadSource: where we actually download checkpoints from.

There will be a coordination layer to wire these layers and the model definitions must be enriched with the download source (possibly listing alternative sources for the same file) to make this wiring possible.

The ModelIndex keeps track of which models are actually available for download and running, with a filtering layer (following the current filtering rules) and an aliasing layer (which relies on the existing AliasManager. I am not sure if the indexer should run completely at startup (like it does now) or if should have a persistent index to make things easier.

Breaking changes

A migration will be needed for the data structures and current model-related json files. The model manager API will likely somewhat change, but the goal is to restructure its internal API rather than what it exposes

Maintenance plan

this being a core component it will be inevitably touched many times. Hopefully the ability to test components in isolation should make the code more robust.

Risks

there is a lot that could go wrong during implementation and very careful testing will be needed. It is also likely we will observe some regressions that will need to be addressed

[jeremyfowers]

Thanks for starting this discussion! I need to read and understand it in depth another time, but we should absolutely not have 6.7k LoC files so clearly something should be done.

Something I've discussed with @ramkrishna2910 is that there seems to be some smart-router-specific functionality burried deep in ModelManager as well. Curious to get your take @bitgamma on whether that should come out as well, as I don't see the smart router mentioned in this RFC yet.

[bitgamma]

I honestly missed that. I have kept the analysis high level to understand the general sentiment on the topic (considering that it comes with risks). If there is interest in pursuing that I'll expand the RFC into a more concrete plan.