Refactor: get_instance method saved in state#200
Conversation
Resolves skops-dev#197 Description Currently, during the dispatch of the get_instance functions, the class stored in the state is being loaded to determine which function to dispatch to. This is bad because module loading can be dangerous. We will add auditing but it is intended to be on the level of get_instance itself, not for the dispatch mechanism. In this PR, the state returned by get_state functions is augmented with the name of the get_instance method required to load the object. This way, we can look up the correct method based on the state and don't need to use the modified singledispatch mechanism, thus avoiding loading modules during dispatching. Implementation Whereas for get_state, we still rely in singledispatch, for get_instance we now use a simple dictionary that looks up the function based on its name (which in turn is stored in the state). The dictionary, going by the name of GET_INSTANCE_MAPPING, is populated similarly to how the get_instance functions were registered previously with singledispatch. There was an issue with circular imports (e.g. get_instance > GET_INSTANCE_MAPPING > ndarray_get_instance > get_instance), hence the get_instance function was moved to its own module, _dispatch.py (other name suggestions welcome). For some types, we now need extra get_state functions because they have specific get_instance methods. So e.g. sgd_loss_get_state just wraps reduce_get_state and adds sgd_loss_get_instance as its loader. Coincidental changes Since we no longer have to inspect the contents of state to determine the function to dispatch to for get_instance, we can fall back to the Python implementation of singledispatch instead of rolling our own. This side effect is a big win. The function Tree_get_instance was renamed to tree_get_instance for consistency. In the debug_dispatch_functions, there was some code from a time when the state was allowed not to be a dict (json-serializable objects). Now we always have a dict, so this dead code was removed. Also, this fixture was elevated to module-level scope, since it only needs to run once.
|
@skops-dev/maintainers ready for review |
Tests pass without it.
|
Addendum: I noticed (thanks codecov) that I had accidentally not registered I'm not exactly sure which change we did that resulted in this, since I'm pretty sure it used to be necessary, but it's not anymore. Therefore, I removed the corresponding code. |
| from typing import Any, Callable | ||
| from zipfile import ZipFile | ||
|
|
||
| GET_INSTANCE_MAPPING: dict[str, Callable[[dict[str, Any], ZipFile], Any]] = {} |
There was a problem hiding this comment.
can we not have the type here? we're gonna change signature later and we'd have to keep these in sync.
| get_instance_func = GET_INSTANCE_MAPPING[state["__loader__"]] | ||
| except KeyError: | ||
| raise TypeError( | ||
| f"Creating an instance of type {type(state)} is not supported yet" |
There was a problem hiding this comment.
type of state is a dict. We should say something like:
f" Can't find loader {state["__loader__"]} for type {state[__module__]}.{state["__class__"]}."
| SquaredLoss, | ||
| ) | ||
| from sklearn.tree._tree import Tree | ||
| from sklearn.utils import Bunch |
There was a problem hiding this comment.
this is because we load and create the class of the object whatever it was now in dict_get_instance instead of creating a dict.
- Remove type annotation for GET_INSTANCE_MAPPING - Better error message if loader not found Also changed: - Misleading var names in get_state
|
@adrinjalali I addressed your comments, also added a test for the error message. |
adrinjalali
left a comment
adrinjalali
left a comment
There was a problem hiding this comment.
I'm mostly happy about the code we're removing here lol
2 participants
Resolves #197
Description
Currently, during the dispatch of the
get_instancefunctions, the class stored in the state is being loaded to determine which function to dispatch to. This is bad because module loading can be dangerous. We will add auditing but it is intended to be on the level ofget_instanceitself, not for the dispatch mechanism.In this PR, the state returned by
get_statefunctions is augmented with the name of the get_instance method required to load the object. This way, we can look up the correct method based on the state and don't need to use the modifiedsingledispatchmechanism, thus avoiding loading modules during dispatching.Implementation
Whereas for
get_state, we still rely insingledispatch, forget_instancewe now use a simple dictionary that looks up the function based on its name (which in turn is stored in the state). The dictionary, going by the name ofGET_INSTANCE_MAPPING, is populated similarly to how theget_instancefunctions were registered previously withsingledispatch.There was an issue with circular imports (e.g.
get_instance > GET_INSTANCE_MAPPING > ndarray_get_instance > get_instance), hence theget_instancefunction was moved to its own module,_dispatch.py(better name suggestions are welcome).For some types, we now need extra get_state functions because they have specific
get_instancemethods. So e.g.sgd_loss_get_statejust wrapsreduce_get_stateand addssgd_loss_get_instanceas its loader.Coincidental changes
Since we no longer have to inspect the contents of state to determine the function to dispatch to for
get_instance, we can fall back to the Python implementation ofsingledispatchinstead of rolling our own. This side effect is a big win.The function
Tree_get_instancewas renamed totree_get_instancefor consistency.In the
debug_dispatch_functions, there was some code from a time when the state was allowed not to be a dict (json-serializable objects). Now we always have a dict, so this dead code was removed.Also, this fixture was elevated to module-level scope, since it only needs to run once.