[FR] Config per registered model
khari998 opened this issue · comments
Willingness to contribute
Yes. I would be willing to contribute this feature with guidance from the MLflow community.
Proposal Summary
Add an optional config parameter when registering a model where we can add a JSON serializable configuration per each model that can be retrieved from the API when querying the model registry. This would allow us to use more sophisticated configurations that model specific tags are not currently suited for.
Motivation
Right now we are able to set tags per registered models which satisfies most cases for most people, however for the project I am working on, my models rely on a rather large JSON configuration that is used to dynamically generate the required inputs for the model. I tried storing this json as a string on a tag but ran into a character limitation.
My current implementation is to store this JSON as an artifact but this only exists for each experiment run. This results in a very clunky way of retrieving the config when trying to run inference, where I need to first search the model registry for the model I need, store a tag on the model for which run was associated with it in order to then find the artifact, use the run id to find the artifact directory with the JSON, deserialize the JSON, and then use that to draw inference from my model. At scale with multiple models, this becomes a clear bottleneck for performance that can be easily avoided.
The easiest way to solve this problem is to just allow us to store a dictionary on the registered model directly. I can't see why we can't have a config parameter in the API that allows us to register a model or log a registered model. That way when we query the model registry via the API, this dictionary can just come up under the config key on the API response. If there are certain things that might break allowing us to add dictionaries with any datatypes as values, we can just ensure that the dictionary is JSON serializable. We don't need to do anything fancy like indexing the keys of this config dictionary to be searchable or anything via a filter query; tags already solve this problem. This is only for allowing advanced configurations to be made available for each registered model so that they can be used for inference more efficiently if needed.
I don't think this would be a hard feature to implement but I don't know the codebase for ml flow enough to know where to start implementing this.
Details
No response
What component(s) does this affect?
-
area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
cc @smurching: What are your thoughts on supporting expanded registry metadata by exposing run artifact data through an API?
Would love any input to this. @smurching
If this is not an easy integration I will try to dive deeper into the codebase myself to see if there is a solution I can come up with
@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.
@khari998 thanks for filing, sorry for the slow response! @BenWilson2 @khari998 my first question would be whether this use case could be solved by passing metadata when logging the model (see example API doc). This metadata gets persisted as part of the MLmodel file / model artifacts, so can support larger JSON configs etc, and can also be loaded back using the mlflow.models.Model.load API (see docs) e.g.
import mlflow
mlmodel = mlflow.models.Model.load(f"models:/{registered_model_name}/{model_version}")
loaded_metadata = mlmodel.metadata
I feel this API is a bit hard to discover/find so we may want to make it more discoverable at least in docs, if it sounds like this approach can work
I believe this should actually solve this problem @smurching ! My apologies, I did not see this in the documentation. Everything I've seen so far has recommended storing everything as an artifact.
I also learned that I have to wrap my model in a custom model wrapper in order to use the load_context method in order to retrieve this during inference do logic with it so that the process is automatic. I believe the combination of these two things should allow me to do this all in one REST API call and reduce the logic that is needed by a ton.