bikes.io.registries

Savers, loaders, and registers for model registries.

View Source

  1"""Savers, loaders, and registers for model registries."""
  2
  3# %% IMPORTS
  4
  5import abc
  6import typing as T
  7
  8import mlflow
  9import pydantic as pdt
 10
 11from bikes.core import models, schemas
 12from bikes.utils import signers
 13
 14# %% TYPES
 15
 16# Results of model registry operations
 17Info: T.TypeAlias = mlflow.models.model.ModelInfo
 18Alias: T.TypeAlias = mlflow.entities.model_registry.ModelVersion
 19Version: T.TypeAlias = mlflow.entities.model_registry.ModelVersion
 20
 21# %% HELPERS
 22
 23
 24def uri_for_model_alias(name: str, alias: str) -> str:
 25    """Create a model URI from a model name and an alias.
 26
 27    Args:
 28        name (str): name of the mlflow registered model.
 29        alias (str): alias of the registered model.
 30
 31    Returns:
 32        str: model URI as "models:/name@alias".
 33    """
 34    return f"models:/{name}@{alias}"
 35
 36
 37def uri_for_model_version(name: str, version: int) -> str:
 38    """Create a model URI from a model name and a version.
 39
 40    Args:
 41        name (str): name of the mlflow registered model.
 42        version (int): version of the registered model.
 43
 44    Returns:
 45        str: model URI as "models:/name/version."
 46    """
 47    return f"models:/{name}/{version}"
 48
 49
 50def uri_for_model_alias_or_version(name: str, alias_or_version: str | int) -> str:
 51    """Create a model URi from a model name and an alias or version.
 52
 53    Args:
 54        name (str): name of the mlflow registered model.
 55        alias_or_version (str | int): alias or version of the registered model.
 56
 57    Returns:
 58        str: model URI as "models:/name@alias" or "models:/name/version" based on input.
 59    """
 60    if isinstance(alias_or_version, int):
 61        return uri_for_model_version(name=name, version=alias_or_version)
 62    else:
 63        return uri_for_model_alias(name=name, alias=alias_or_version)
 64
 65
 66# %% SAVERS
 67
 68
 69class Saver(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
 70    """Base class for saving models in registry.
 71
 72    Separate model definition from serialization.
 73    e.g., to switch between serialization flavors.
 74
 75    Parameters:
 76        path (str): model path inside the Mlflow store.
 77    """
 78
 79    KIND: str
 80
 81    path: str = "model"
 82
 83    @abc.abstractmethod
 84    def save(
 85        self, model: models.Model, signature: signers.Signature, input_example: schemas.Inputs
 86    ) -> Info:
 87        """Save a model in the model registry.
 88
 89        Args:
 90            model (models.Model): project model to save.
 91            signature (signers.Signature): model signature.
 92            input_example (schemas.Inputs): sample of inputs.
 93
 94        Returns:
 95            Info: model saving information.
 96        """
 97
 98
 99class CustomSaver(Saver):
100    """Saver for project models using the Mlflow PyFunc module.
101
102    https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html
103    """
104
105    KIND: T.Literal["CustomSaver"] = "CustomSaver"
106
107    class Adapter(mlflow.pyfunc.PythonModel):  # type: ignore[misc]
108        """Adapt a custom model to the Mlflow PyFunc flavor for saving operations.
109
110        https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html?#mlflow.pyfunc.PythonModel
111        """
112
113        def __init__(self, model: models.Model):
114            """Initialize the custom saver adapter.
115
116            Args:
117                model (models.Model): project model.
118            """
119            self.model = model
120
121        def predict(
122            self,
123            context: mlflow.pyfunc.PythonModelContext,
124            model_input: schemas.Inputs,
125            params: dict[str, T.Any] | None = None,
126        ) -> schemas.Outputs:
127            """Generate predictions with a custom model for the given inputs.
128
129            Args:
130                context (mlflow.pyfunc.PythonModelContext): mlflow context.
131                model_input (schemas.Inputs): inputs for the mlflow model.
132                params (dict[str, T.Any] | None): additional parameters.
133
134            Returns:
135                schemas.Outputs: validated outputs of the project model.
136            """
137            return self.model.predict(inputs=model_input)
138
139    @T.override
140    def save(
141        self, model: models.Model, signature: signers.Signature, input_example: schemas.Inputs
142    ) -> Info:
143        adapter = CustomSaver.Adapter(model=model)
144        return mlflow.pyfunc.log_model(
145            python_model=adapter,
146            signature=signature,
147            artifact_path=self.path,
148            input_example=input_example,
149        )
150
151
152class BuiltinSaver(Saver):
153    """Saver for built-in models using an Mlflow flavor module.
154
155    https://mlflow.org/docs/latest/models.html#built-in-model-flavors
156
157    Parameters:
158        flavor (str): Mlflow flavor module to use for the serialization.
159    """
160
161    KIND: T.Literal["BuiltinSaver"] = "BuiltinSaver"
162
163    flavor: str
164
165    @T.override
166    def save(
167        self,
168        model: models.Model,
169        signature: signers.Signature,
170        input_example: schemas.Inputs | None = None,
171    ) -> mlflow.entities.model_registry.ModelVersion:
172        builtin_model = model.get_internal_model()
173        module = getattr(mlflow, self.flavor)
174        return module.log_model(
175            builtin_model, artifact_path=self.path, signature=signature, input_example=input_example
176        )
177
178
179SaverKind = CustomSaver | BuiltinSaver
180
181# %% LOADERS
182
183
184class Loader(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
185    """Base class for loading models from registry.
186
187    Separate model definition from deserialization.
188    e.g., to switch between deserialization flavors.
189    """
190
191    KIND: str
192
193    class Adapter(abc.ABC):
194        """Adapt any model for the project inference."""
195
196        @abc.abstractmethod
197        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
198            """Generate predictions with the internal model for the given inputs.
199
200            Args:
201                inputs (schemas.Inputs): validated inputs for the project model.
202
203            Returns:
204                schemas.Outputs: validated outputs of the project model.
205            """
206
207    @abc.abstractmethod
208    def load(self, uri: str) -> "Loader.Adapter":
209        """Load a model from the model registry.
210
211        Args:
212            uri (str): URI of a model to load.
213
214        Returns:
215            Loader.Adapter: model loaded.
216        """
217
218
219class CustomLoader(Loader):
220    """Loader for custom models using the Mlflow PyFunc module.
221
222    https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html
223    """
224
225    KIND: T.Literal["CustomLoader"] = "CustomLoader"
226
227    class Adapter(Loader.Adapter):
228        """Adapt a custom model for the project inference."""
229
230        def __init__(self, model: mlflow.pyfunc.PyFuncModel) -> None:
231            """Initialize the adapter from an mlflow pyfunc model.
232
233            Args:
234                model (mlflow.pyfunc.PyFuncModel): mlflow pyfunc model.
235            """
236            self.model = model
237
238        @T.override
239        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
240            # model validation is already done in predict
241            outputs = self.model.predict(data=inputs)
242            return T.cast(schemas.Outputs, outputs)
243
244    @T.override
245    def load(self, uri: str) -> "CustomLoader.Adapter":
246        model = mlflow.pyfunc.load_model(model_uri=uri)
247        adapter = CustomLoader.Adapter(model=model)
248        return adapter
249
250
251class BuiltinLoader(Loader):
252    """Loader for built-in models using the Mlflow PyFunc module.
253
254    Note: use Mlflow PyFunc instead of flavors to use standard API.
255
256    https://mlflow.org/docs/latest/models.html#built-in-model-flavors
257    """
258
259    KIND: T.Literal["BuiltinLoader"] = "BuiltinLoader"
260
261    class Adapter(Loader.Adapter):
262        """Adapt a builtin model for the project inference."""
263
264        def __init__(self, model: mlflow.pyfunc.PyFuncModel) -> None:
265            """Initialize the adapter from an mlflow pyfunc model.
266
267            Args:
268                model (mlflow.pyfunc.PyFuncModel): mlflow pyfunc model.
269            """
270            self.model = model
271
272        @T.override
273        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
274            columns = list(schemas.OutputsSchema.to_schema().columns)
275            outputs = self.model.predict(data=inputs)  # unchecked data!
276            return schemas.Outputs(outputs, columns=columns, index=inputs.index)
277
278    @T.override
279    def load(self, uri: str) -> "BuiltinLoader.Adapter":
280        model = mlflow.pyfunc.load_model(model_uri=uri)
281        adapter = BuiltinLoader.Adapter(model=model)
282        return adapter
283
284
285LoaderKind = CustomLoader | BuiltinLoader
286
287# %% REGISTERS
288
289
290class Register(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
291    """Base class for registring models to a location.
292
293    Separate model definition from its registration.
294    e.g., to change the model registry backend.
295
296    Parameters:
297        tags (dict[str, T.Any]): tags for the model.
298    """
299
300    KIND: str
301
302    tags: dict[str, T.Any] = {}
303
304    @abc.abstractmethod
305    def register(self, name: str, model_uri: str) -> Version:
306        """Register a model given its name and URI.
307
308        Args:
309            name (str): name of the model to register.
310            model_uri (str): URI of a model to register.
311
312        Returns:
313            Version: information about the registered model.
314        """
315
316
317class MlflowRegister(Register):
318    """Register for models in the Mlflow Model Registry.
319
320    https://mlflow.org/docs/latest/model-registry.html
321    """
322
323    KIND: T.Literal["MlflowRegister"] = "MlflowRegister"
324
325    @T.override
326    def register(self, name: str, model_uri: str) -> Version:
327        return mlflow.register_model(name=name, model_uri=model_uri, tags=self.tags)
328
329
330RegisterKind = MlflowRegister

Info: TypeAlias = mlflow.models.model.ModelInfo

Alias: TypeAlias = mlflow.entities.model_registry.model_version.ModelVersion

Version: TypeAlias = mlflow.entities.model_registry.model_version.ModelVersion

def uri_for_model_alias(name: str, alias: str) -> str: View Source

25def uri_for_model_alias(name: str, alias: str) -> str:
26    """Create a model URI from a model name and an alias.
27
28    Args:
29        name (str): name of the mlflow registered model.
30        alias (str): alias of the registered model.
31
32    Returns:
33        str: model URI as "models:/name@alias".
34    """
35    return f"models:/{name}@{alias}"

Create a model URI from a model name and an alias.

Arguments:

name (str): name of the mlflow registered model.
alias (str): alias of the registered model.

Returns:

str: model URI as "models:/name@alias".

def uri_for_model_version(name: str, version: int) -> str: View Source

38def uri_for_model_version(name: str, version: int) -> str:
39    """Create a model URI from a model name and a version.
40
41    Args:
42        name (str): name of the mlflow registered model.
43        version (int): version of the registered model.
44
45    Returns:
46        str: model URI as "models:/name/version."
47    """
48    return f"models:/{name}/{version}"

Create a model URI from a model name and a version.

Arguments:

name (str): name of the mlflow registered model.
version (int): version of the registered model.

Returns:

str: model URI as "models:/name/version."

def uri_for_model_alias_or_version(name: str, alias_or_version: str | int) -> str: View Source

51def uri_for_model_alias_or_version(name: str, alias_or_version: str | int) -> str:
52    """Create a model URi from a model name and an alias or version.
53
54    Args:
55        name (str): name of the mlflow registered model.
56        alias_or_version (str | int): alias or version of the registered model.
57
58    Returns:
59        str: model URI as "models:/name@alias" or "models:/name/version" based on input.
60    """
61    if isinstance(alias_or_version, int):
62        return uri_for_model_version(name=name, version=alias_or_version)
63    else:
64        return uri_for_model_alias(name=name, alias=alias_or_version)

Create a model URi from a model name and an alias or version.

Arguments:

name (str): name of the mlflow registered model.
alias_or_version (str | int): alias or version of the registered model.

Returns:

str: model URI as "models:/name@alias" or "models:/name/version" based on input.

class Saver(abc.ABC, pydantic.main.BaseModel): View Source

70class Saver(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
71    """Base class for saving models in registry.
72
73    Separate model definition from serialization.
74    e.g., to switch between serialization flavors.
75
76    Parameters:
77        path (str): model path inside the Mlflow store.
78    """
79
80    KIND: str
81
82    path: str = "model"
83
84    @abc.abstractmethod
85    def save(
86        self, model: models.Model, signature: signers.Signature, input_example: schemas.Inputs
87    ) -> Info:
88        """Save a model in the model registry.
89
90        Args:
91            model (models.Model): project model to save.
92            signature (signers.Signature): model signature.
93            input_example (schemas.Inputs): sample of inputs.
94
95        Returns:
96            Info: model saving information.
97        """

Base class for saving models in registry.

Separate model definition from serialization. e.g., to switch between serialization flavors.

Arguments:

path (str): model path inside the Mlflow store.

KIND: str

path: str

@abc.abstractmethod

def save( self, model: bikes.core.models.Model, signature: mlflow.models.signature.ModelSignature, input_example: pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema]) -> mlflow.models.model.ModelInfo: View Source

84    @abc.abstractmethod
85    def save(
86        self, model: models.Model, signature: signers.Signature, input_example: schemas.Inputs
87    ) -> Info:
88        """Save a model in the model registry.
89
90        Args:
91            model (models.Model): project model to save.
92            signature (signers.Signature): model signature.
93            input_example (schemas.Inputs): sample of inputs.
94
95        Returns:
96            Info: model saving information.
97        """

Save a model in the model registry.

Arguments:

model (models.Model): project model to save.
signature (signers.Signature): model signature.
input_example (schemas.Inputs): sample of inputs.

Returns:

Info: model saving information.

model_config = {'strict': True, 'frozen': True, 'extra': 'forbid'}

model_fields = {'KIND': FieldInfo(annotation=str, required=True), 'path': FieldInfo(annotation=str, required=False, default='model')}

model_computed_fields = {}

Inherited Members

pydantic.main.BaseModel: BaseModel; model_extra; model_fields_set; model_construct; model_copy; model_dump; model_dump_json; model_json_schema; model_parametrized_name; model_post_init; model_rebuild; model_validate; model_validate_json; model_validate_strings; dict; json; parse_obj; parse_raw; parse_file; from_orm; construct; copy; schema; schema_json; validate; update_forward_refs

class CustomSaver(Saver): View Source

100class CustomSaver(Saver):
101    """Saver for project models using the Mlflow PyFunc module.
102
103    https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html
104    """
105
106    KIND: T.Literal["CustomSaver"] = "CustomSaver"
107
108    class Adapter(mlflow.pyfunc.PythonModel):  # type: ignore[misc]
109        """Adapt a custom model to the Mlflow PyFunc flavor for saving operations.
110
111        https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html?#mlflow.pyfunc.PythonModel
112        """
113
114        def __init__(self, model: models.Model):
115            """Initialize the custom saver adapter.
116
117            Args:
118                model (models.Model): project model.
119            """
120            self.model = model
121
122        def predict(
123            self,
124            context: mlflow.pyfunc.PythonModelContext,
125            model_input: schemas.Inputs,
126            params: dict[str, T.Any] | None = None,
127        ) -> schemas.Outputs:
128            """Generate predictions with a custom model for the given inputs.
129
130            Args:
131                context (mlflow.pyfunc.PythonModelContext): mlflow context.
132                model_input (schemas.Inputs): inputs for the mlflow model.
133                params (dict[str, T.Any] | None): additional parameters.
134
135            Returns:
136                schemas.Outputs: validated outputs of the project model.
137            """
138            return self.model.predict(inputs=model_input)
139
140    @T.override
141    def save(
142        self, model: models.Model, signature: signers.Signature, input_example: schemas.Inputs
143    ) -> Info:
144        adapter = CustomSaver.Adapter(model=model)
145        return mlflow.pyfunc.log_model(
146            python_model=adapter,
147            signature=signature,
148            artifact_path=self.path,
149            input_example=input_example,
150        )

Saver for project models using the Mlflow PyFunc module.

https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html

KIND: Literal['CustomSaver']

@T.override

140    @T.override
141    def save(
142        self, model: models.Model, signature: signers.Signature, input_example: schemas.Inputs
143    ) -> Info:
144        adapter = CustomSaver.Adapter(model=model)
145        return mlflow.pyfunc.log_model(
146            python_model=adapter,
147            signature=signature,
148            artifact_path=self.path,
149            input_example=input_example,
150        )

Save a model in the model registry.

Arguments:

model (models.Model): project model to save.
signature (signers.Signature): model signature.
input_example (schemas.Inputs): sample of inputs.

Returns:

Info: model saving information.

model_config = {'strict': True, 'frozen': True, 'extra': 'forbid'}

model_fields = {'KIND': FieldInfo(annotation=Literal['CustomSaver'], required=False, default='CustomSaver'), 'path': FieldInfo(annotation=str, required=False, default='model')}

model_computed_fields = {}

Inherited Members

pydantic.main.BaseModel: BaseModel; model_extra; model_fields_set; model_construct; model_copy; model_dump; model_dump_json; model_json_schema; model_parametrized_name; model_post_init; model_rebuild; model_validate; model_validate_json; model_validate_strings; dict; json; parse_obj; parse_raw; parse_file; from_orm; construct; copy; schema; schema_json; validate; update_forward_refs
Saver: path

class CustomSaver.Adapter(mlflow.pyfunc.model.PythonModel): View Source

108    class Adapter(mlflow.pyfunc.PythonModel):  # type: ignore[misc]
109        """Adapt a custom model to the Mlflow PyFunc flavor for saving operations.
110
111        https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html?#mlflow.pyfunc.PythonModel
112        """
113
114        def __init__(self, model: models.Model):
115            """Initialize the custom saver adapter.
116
117            Args:
118                model (models.Model): project model.
119            """
120            self.model = model
121
122        def predict(
123            self,
124            context: mlflow.pyfunc.PythonModelContext,
125            model_input: schemas.Inputs,
126            params: dict[str, T.Any] | None = None,
127        ) -> schemas.Outputs:
128            """Generate predictions with a custom model for the given inputs.
129
130            Args:
131                context (mlflow.pyfunc.PythonModelContext): mlflow context.
132                model_input (schemas.Inputs): inputs for the mlflow model.
133                params (dict[str, T.Any] | None): additional parameters.
134
135            Returns:
136                schemas.Outputs: validated outputs of the project model.
137            """
138            return self.model.predict(inputs=model_input)

Adapt a custom model to the Mlflow PyFunc flavor for saving operations.

https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html?#mlflow.pyfunc.PythonModel

CustomSaver.Adapter(model: bikes.core.models.Model) View Source

114        def __init__(self, model: models.Model):
115            """Initialize the custom saver adapter.
116
117            Args:
118                model (models.Model): project model.
119            """
120            self.model = model

Initialize the custom saver adapter.

Arguments:

model (models.Model): project model.

model

def predict( self, context: mlflow.pyfunc.model.PythonModelContext, model_input: pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema], params: dict[str, typing.Any] | None = None) -> pandera.typing.pandas.DataFrame[bikes.core.schemas.OutputsSchema]: View Source

122        def predict(
123            self,
124            context: mlflow.pyfunc.PythonModelContext,
125            model_input: schemas.Inputs,
126            params: dict[str, T.Any] | None = None,
127        ) -> schemas.Outputs:
128            """Generate predictions with a custom model for the given inputs.
129
130            Args:
131                context (mlflow.pyfunc.PythonModelContext): mlflow context.
132                model_input (schemas.Inputs): inputs for the mlflow model.
133                params (dict[str, T.Any] | None): additional parameters.
134
135            Returns:
136                schemas.Outputs: validated outputs of the project model.
137            """
138            return self.model.predict(inputs=model_input)

Generate predictions with a custom model for the given inputs.

Arguments:

context (mlflow.pyfunc.PythonModelContext): mlflow context.
model_input (schemas.Inputs): inputs for the mlflow model.
params (dict[str, T.Any] | None): additional parameters.

Returns:

schemas.Outputs: validated outputs of the project model.

Inherited Members

mlflow.pyfunc.model.PythonModel: load_context; predict_stream

class BuiltinSaver(Saver): View Source

153class BuiltinSaver(Saver):
154    """Saver for built-in models using an Mlflow flavor module.
155
156    https://mlflow.org/docs/latest/models.html#built-in-model-flavors
157
158    Parameters:
159        flavor (str): Mlflow flavor module to use for the serialization.
160    """
161
162    KIND: T.Literal["BuiltinSaver"] = "BuiltinSaver"
163
164    flavor: str
165
166    @T.override
167    def save(
168        self,
169        model: models.Model,
170        signature: signers.Signature,
171        input_example: schemas.Inputs | None = None,
172    ) -> mlflow.entities.model_registry.ModelVersion:
173        builtin_model = model.get_internal_model()
174        module = getattr(mlflow, self.flavor)
175        return module.log_model(
176            builtin_model, artifact_path=self.path, signature=signature, input_example=input_example
177        )

Saver for built-in models using an Mlflow flavor module.

https://mlflow.org/docs/latest/models.html#built-in-model-flavors

Arguments:

flavor (str): Mlflow flavor module to use for the serialization.

KIND: Literal['BuiltinSaver']

flavor: str

@T.override

def save( self, model: bikes.core.models.Model, signature: mlflow.models.signature.ModelSignature, input_example: Optional[pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema]] = None) -> mlflow.entities.model_registry.model_version.ModelVersion: View Source

166    @T.override
167    def save(
168        self,
169        model: models.Model,
170        signature: signers.Signature,
171        input_example: schemas.Inputs | None = None,
172    ) -> mlflow.entities.model_registry.ModelVersion:
173        builtin_model = model.get_internal_model()
174        module = getattr(mlflow, self.flavor)
175        return module.log_model(
176            builtin_model, artifact_path=self.path, signature=signature, input_example=input_example
177        )

Save a model in the model registry.

Arguments:

model (models.Model): project model to save.
signature (signers.Signature): model signature.
input_example (schemas.Inputs): sample of inputs.

Returns:

Info: model saving information.

model_config = {'strict': True, 'frozen': True, 'extra': 'forbid'}

model_fields = {'KIND': FieldInfo(annotation=Literal['BuiltinSaver'], required=False, default='BuiltinSaver'), 'path': FieldInfo(annotation=str, required=False, default='model'), 'flavor': FieldInfo(annotation=str, required=True)}

model_computed_fields = {}

Inherited Members

pydantic.main.BaseModel: BaseModel; model_extra; model_fields_set; model_construct; model_copy; model_dump; model_dump_json; model_json_schema; model_parametrized_name; model_post_init; model_rebuild; model_validate; model_validate_json; model_validate_strings; dict; json; parse_obj; parse_raw; parse_file; from_orm; construct; copy; schema; schema_json; validate; update_forward_refs
Saver: path

SaverKind = CustomSaver | BuiltinSaver

class Loader(abc.ABC, pydantic.main.BaseModel): View Source

185class Loader(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
186    """Base class for loading models from registry.
187
188    Separate model definition from deserialization.
189    e.g., to switch between deserialization flavors.
190    """
191
192    KIND: str
193
194    class Adapter(abc.ABC):
195        """Adapt any model for the project inference."""
196
197        @abc.abstractmethod
198        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
199            """Generate predictions with the internal model for the given inputs.
200
201            Args:
202                inputs (schemas.Inputs): validated inputs for the project model.
203
204            Returns:
205                schemas.Outputs: validated outputs of the project model.
206            """
207
208    @abc.abstractmethod
209    def load(self, uri: str) -> "Loader.Adapter":
210        """Load a model from the model registry.
211
212        Args:
213            uri (str): URI of a model to load.
214
215        Returns:
216            Loader.Adapter: model loaded.
217        """

Base class for loading models from registry.

Separate model definition from deserialization. e.g., to switch between deserialization flavors.

KIND: str

@abc.abstractmethod

def load(self, uri: str) -> Loader.Adapter: View Source

208    @abc.abstractmethod
209    def load(self, uri: str) -> "Loader.Adapter":
210        """Load a model from the model registry.
211
212        Args:
213            uri (str): URI of a model to load.
214
215        Returns:
216            Loader.Adapter: model loaded.
217        """

Load a model from the model registry.

Arguments:

uri (str): URI of a model to load.

Returns:

Loader.Adapter: model loaded.

model_config = {'strict': True, 'frozen': True, 'extra': 'forbid'}

model_fields = {'KIND': FieldInfo(annotation=str, required=True)}

model_computed_fields = {}

Inherited Members

pydantic.main.BaseModel: BaseModel; model_extra; model_fields_set; model_construct; model_copy; model_dump; model_dump_json; model_json_schema; model_parametrized_name; model_post_init; model_rebuild; model_validate; model_validate_json; model_validate_strings; dict; json; parse_obj; parse_raw; parse_file; from_orm; construct; copy; schema; schema_json; validate; update_forward_refs

class Loader.Adapter(abc.ABC): View Source

194    class Adapter(abc.ABC):
195        """Adapt any model for the project inference."""
196
197        @abc.abstractmethod
198        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
199            """Generate predictions with the internal model for the given inputs.
200
201            Args:
202                inputs (schemas.Inputs): validated inputs for the project model.
203
204            Returns:
205                schemas.Outputs: validated outputs of the project model.
206            """

Adapt any model for the project inference.

@abc.abstractmethod

def predict( self, inputs: pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema]) -> pandera.typing.pandas.DataFrame[bikes.core.schemas.OutputsSchema]: View Source

197        @abc.abstractmethod
198        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
199            """Generate predictions with the internal model for the given inputs.
200
201            Args:
202                inputs (schemas.Inputs): validated inputs for the project model.
203
204            Returns:
205                schemas.Outputs: validated outputs of the project model.
206            """

Generate predictions with the internal model for the given inputs.

Arguments:

inputs (schemas.Inputs): validated inputs for the project model.

Returns:

schemas.Outputs: validated outputs of the project model.

class CustomLoader(Loader): View Source

220class CustomLoader(Loader):
221    """Loader for custom models using the Mlflow PyFunc module.
222
223    https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html
224    """
225
226    KIND: T.Literal["CustomLoader"] = "CustomLoader"
227
228    class Adapter(Loader.Adapter):
229        """Adapt a custom model for the project inference."""
230
231        def __init__(self, model: mlflow.pyfunc.PyFuncModel) -> None:
232            """Initialize the adapter from an mlflow pyfunc model.
233
234            Args:
235                model (mlflow.pyfunc.PyFuncModel): mlflow pyfunc model.
236            """
237            self.model = model
238
239        @T.override
240        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
241            # model validation is already done in predict
242            outputs = self.model.predict(data=inputs)
243            return T.cast(schemas.Outputs, outputs)
244
245    @T.override
246    def load(self, uri: str) -> "CustomLoader.Adapter":
247        model = mlflow.pyfunc.load_model(model_uri=uri)
248        adapter = CustomLoader.Adapter(model=model)
249        return adapter

Loader for custom models using the Mlflow PyFunc module.

https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html

KIND: Literal['CustomLoader']

@T.override

def load(self, uri: str) -> CustomLoader.Adapter: View Source

245    @T.override
246    def load(self, uri: str) -> "CustomLoader.Adapter":
247        model = mlflow.pyfunc.load_model(model_uri=uri)
248        adapter = CustomLoader.Adapter(model=model)
249        return adapter

Load a model from the model registry.

Arguments:

uri (str): URI of a model to load.

Returns:

Loader.Adapter: model loaded.

model_config = {'strict': True, 'frozen': True, 'extra': 'forbid'}

model_fields = {'KIND': FieldInfo(annotation=Literal['CustomLoader'], required=False, default='CustomLoader')}

model_computed_fields = {}

Inherited Members

pydantic.main.BaseModel: BaseModel; model_extra; model_fields_set; model_construct; model_copy; model_dump; model_dump_json; model_json_schema; model_parametrized_name; model_post_init; model_rebuild; model_validate; model_validate_json; model_validate_strings; dict; json; parse_obj; parse_raw; parse_file; from_orm; construct; copy; schema; schema_json; validate; update_forward_refs

class CustomLoader.Adapter(Loader.Adapter): View Source

228    class Adapter(Loader.Adapter):
229        """Adapt a custom model for the project inference."""
230
231        def __init__(self, model: mlflow.pyfunc.PyFuncModel) -> None:
232            """Initialize the adapter from an mlflow pyfunc model.
233
234            Args:
235                model (mlflow.pyfunc.PyFuncModel): mlflow pyfunc model.
236            """
237            self.model = model
238
239        @T.override
240        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
241            # model validation is already done in predict
242            outputs = self.model.predict(data=inputs)
243            return T.cast(schemas.Outputs, outputs)

Adapt a custom model for the project inference.

CustomLoader.Adapter(model: mlflow.pyfunc.PyFuncModel) View Source

231        def __init__(self, model: mlflow.pyfunc.PyFuncModel) -> None:
232            """Initialize the adapter from an mlflow pyfunc model.
233
234            Args:
235                model (mlflow.pyfunc.PyFuncModel): mlflow pyfunc model.
236            """
237            self.model = model

Initialize the adapter from an mlflow pyfunc model.

Arguments:

model (mlflow.pyfunc.PyFuncModel): mlflow pyfunc model.

model

@T.override

def predict( self, inputs: pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema]) -> pandera.typing.pandas.DataFrame[bikes.core.schemas.OutputsSchema]: View Source

239        @T.override
240        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
241            # model validation is already done in predict
242            outputs = self.model.predict(data=inputs)
243            return T.cast(schemas.Outputs, outputs)

Generate predictions with the internal model for the given inputs.

Arguments:

inputs (schemas.Inputs): validated inputs for the project model.

Returns:

schemas.Outputs: validated outputs of the project model.

class BuiltinLoader(Loader): View Source

252class BuiltinLoader(Loader):
253    """Loader for built-in models using the Mlflow PyFunc module.
254
255    Note: use Mlflow PyFunc instead of flavors to use standard API.
256
257    https://mlflow.org/docs/latest/models.html#built-in-model-flavors
258    """
259
260    KIND: T.Literal["BuiltinLoader"] = "BuiltinLoader"
261
262    class Adapter(Loader.Adapter):
263        """Adapt a builtin model for the project inference."""
264
265        def __init__(self, model: mlflow.pyfunc.PyFuncModel) -> None:
266            """Initialize the adapter from an mlflow pyfunc model.
267
268            Args:
269                model (mlflow.pyfunc.PyFuncModel): mlflow pyfunc model.
270            """
271            self.model = model
272
273        @T.override
274        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
275            columns = list(schemas.OutputsSchema.to_schema().columns)
276            outputs = self.model.predict(data=inputs)  # unchecked data!
277            return schemas.Outputs(outputs, columns=columns, index=inputs.index)
278
279    @T.override
280    def load(self, uri: str) -> "BuiltinLoader.Adapter":
281        model = mlflow.pyfunc.load_model(model_uri=uri)
282        adapter = BuiltinLoader.Adapter(model=model)
283        return adapter

Loader for built-in models using the Mlflow PyFunc module.

Note: use Mlflow PyFunc instead of flavors to use standard API.

https://mlflow.org/docs/latest/models.html#built-in-model-flavors

KIND: Literal['BuiltinLoader']

@T.override

def load(self, uri: str) -> BuiltinLoader.Adapter: View Source

279    @T.override
280    def load(self, uri: str) -> "BuiltinLoader.Adapter":
281        model = mlflow.pyfunc.load_model(model_uri=uri)
282        adapter = BuiltinLoader.Adapter(model=model)
283        return adapter

Load a model from the model registry.

Arguments:

uri (str): URI of a model to load.

Returns:

Loader.Adapter: model loaded.

model_config = {'strict': True, 'frozen': True, 'extra': 'forbid'}

model_fields = {'KIND': FieldInfo(annotation=Literal['BuiltinLoader'], required=False, default='BuiltinLoader')}

model_computed_fields = {}

Inherited Members

pydantic.main.BaseModel: BaseModel; model_extra; model_fields_set; model_construct; model_copy; model_dump; model_dump_json; model_json_schema; model_parametrized_name; model_post_init; model_rebuild; model_validate; model_validate_json; model_validate_strings; dict; json; parse_obj; parse_raw; parse_file; from_orm; construct; copy; schema; schema_json; validate; update_forward_refs

class BuiltinLoader.Adapter(Loader.Adapter): View Source

262    class Adapter(Loader.Adapter):
263        """Adapt a builtin model for the project inference."""
264
265        def __init__(self, model: mlflow.pyfunc.PyFuncModel) -> None:
266            """Initialize the adapter from an mlflow pyfunc model.
267
268            Args:
269                model (mlflow.pyfunc.PyFuncModel): mlflow pyfunc model.
270            """
271            self.model = model
272
273        @T.override
274        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
275            columns = list(schemas.OutputsSchema.to_schema().columns)
276            outputs = self.model.predict(data=inputs)  # unchecked data!
277            return schemas.Outputs(outputs, columns=columns, index=inputs.index)

Adapt a builtin model for the project inference.

BuiltinLoader.Adapter(model: mlflow.pyfunc.PyFuncModel) View Source

265        def __init__(self, model: mlflow.pyfunc.PyFuncModel) -> None:
266            """Initialize the adapter from an mlflow pyfunc model.
267
268            Args:
269                model (mlflow.pyfunc.PyFuncModel): mlflow pyfunc model.
270            """
271            self.model = model

Initialize the adapter from an mlflow pyfunc model.

Arguments:

model (mlflow.pyfunc.PyFuncModel): mlflow pyfunc model.

model

@T.override

def predict( self, inputs: pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema]) -> pandera.typing.pandas.DataFrame[bikes.core.schemas.OutputsSchema]: View Source

273        @T.override
274        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
275            columns = list(schemas.OutputsSchema.to_schema().columns)
276            outputs = self.model.predict(data=inputs)  # unchecked data!
277            return schemas.Outputs(outputs, columns=columns, index=inputs.index)

Generate predictions with the internal model for the given inputs.

Arguments:

inputs (schemas.Inputs): validated inputs for the project model.

Returns:

schemas.Outputs: validated outputs of the project model.

LoaderKind = CustomLoader | BuiltinLoader

class Register(abc.ABC, pydantic.main.BaseModel): View Source

291class Register(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
292    """Base class for registring models to a location.
293
294    Separate model definition from its registration.
295    e.g., to change the model registry backend.
296
297    Parameters:
298        tags (dict[str, T.Any]): tags for the model.
299    """
300
301    KIND: str
302
303    tags: dict[str, T.Any] = {}
304
305    @abc.abstractmethod
306    def register(self, name: str, model_uri: str) -> Version:
307        """Register a model given its name and URI.
308
309        Args:
310            name (str): name of the model to register.
311            model_uri (str): URI of a model to register.
312
313        Returns:
314            Version: information about the registered model.
315        """

Base class for registring models to a location.

Separate model definition from its registration. e.g., to change the model registry backend.

Arguments:

tags (dict[str, T.Any]): tags for the model.

KIND: str

tags: dict[str, typing.Any]

@abc.abstractmethod

def register( self, name: str, model_uri: str) -> mlflow.entities.model_registry.model_version.ModelVersion: View Source

305    @abc.abstractmethod
306    def register(self, name: str, model_uri: str) -> Version:
307        """Register a model given its name and URI.
308
309        Args:
310            name (str): name of the model to register.
311            model_uri (str): URI of a model to register.
312
313        Returns:
314            Version: information about the registered model.
315        """

Arguments:

name (str): name of the model to register.
model_uri (str): URI of a model to register.

Returns:

Version: information about the registered model.

model_config = {'strict': True, 'frozen': True, 'extra': 'forbid'}

model_fields = {'KIND': FieldInfo(annotation=str, required=True), 'tags': FieldInfo(annotation=dict[str, Any], required=False, default={})}

model_computed_fields = {}

Inherited Members

pydantic.main.BaseModel: BaseModel; model_extra; model_fields_set; model_construct; model_copy; model_dump; model_dump_json; model_json_schema; model_parametrized_name; model_post_init; model_rebuild; model_validate; model_validate_json; model_validate_strings; dict; json; parse_obj; parse_raw; parse_file; from_orm; construct; copy; schema; schema_json; validate; update_forward_refs

class MlflowRegister(Register): View Source

318class MlflowRegister(Register):
319    """Register for models in the Mlflow Model Registry.
320
321    https://mlflow.org/docs/latest/model-registry.html
322    """
323
324    KIND: T.Literal["MlflowRegister"] = "MlflowRegister"
325
326    @T.override
327    def register(self, name: str, model_uri: str) -> Version:
328        return mlflow.register_model(name=name, model_uri=model_uri, tags=self.tags)

https://mlflow.org/docs/latest/model-registry.html

KIND: Literal['MlflowRegister']

@T.override

def register( self, name: str, model_uri: str) -> mlflow.entities.model_registry.model_version.ModelVersion: View Source

326    @T.override
327    def register(self, name: str, model_uri: str) -> Version:
328        return mlflow.register_model(name=name, model_uri=model_uri, tags=self.tags)

Arguments:

name (str): name of the model to register.
model_uri (str): URI of a model to register.

Returns:

Version: information about the registered model.

model_config = {'strict': True, 'frozen': True, 'extra': 'forbid'}

model_fields = {'KIND': FieldInfo(annotation=Literal['MlflowRegister'], required=False, default='MlflowRegister'), 'tags': FieldInfo(annotation=dict[str, Any], required=False, default={})}

model_computed_fields = {}

Inherited Members

pydantic.main.BaseModel: BaseModel; model_extra; model_fields_set; model_construct; model_copy; model_dump; model_dump_json; model_json_schema; model_parametrized_name; model_post_init; model_rebuild; model_validate; model_validate_json; model_validate_strings; dict; json; parse_obj; parse_raw; parse_file; from_orm; construct; copy; schema; schema_json; validate; update_forward_refs
Register: tags

RegisterKind = <class 'MlflowRegister'>