bikes.io.registries

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

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

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

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

106class CustomSaver(Saver):
107    """Saver for project models using the Mlflow PyFunc module.
108
109    Nothing is pickled: the adapter is logged as code (MLflow "models from code"), the
110    model params as its `model_config`, and the fitted internal model as a skops file.
111
112    https://mlflow.org/docs/latest/ml/model/models-from-code/
113    """
114
115    KIND: T.Literal["CustomSaver"] = "CustomSaver"
116
117    # Code file that MLflow copies with the model and runs to rebuild the adapter.
118    CODE: T.ClassVar[Path] = Path(__file__).with_name("pyfunc_model.py")
119
120    class Adapter(PythonModel):  # type: ignore[misc]
121        """Adapt a custom model to the Mlflow PyFunc flavor.
122
123        https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html?#mlflow.pyfunc.PythonModel
124        """
125
126        model: models.Model
127
128        def load_context(self, context: PythonModelContext) -> None:
129            """Rebuild the project model from its params and its internal model file.
130
131            Args:
132                context (mlflow.PythonModelContext): mlflow context with config and artifacts.
133            """
134            model = pdt.TypeAdapter(models.ModelKind).validate_python(context.model_config)
135            self.model = model.load_internal_model(path=context.artifacts["model"])
136
137        def predict(
138            self,
139            context: PythonModelContext,  # noqa: ARG002  # required by mlflow PythonModel.predict
140            model_input: pd.DataFrame,
141            params: dict[str, T.Any] | None = None,  # noqa: ARG002  # required by mlflow PythonModel.predict
142        ) -> schemas.Outputs:
143            """Generate predictions with a custom model for the given inputs.
144
145            Args:
146                context (mlflow.PythonModelContext): mlflow context.
147                model_input (pd.DataFrame): inputs to validate for the mlflow model.
148                params (dict[str, T.Any] | None): additional parameters.
149
150            Returns:
151                schemas.Outputs: validated outputs of the project model.
152            """
153            # MLflow accepts pandas here and validates the logged signature;
154            # Pandera enforces the project's constraints before the typed model call.
155            inputs = schemas.InputsSchema.check(data=model_input)
156            return self.model.predict(inputs=inputs)
157
158    @T.override
159    def save(
160        self,
161        model: models.Model,
162        signature: signers.Signature,
163        input_example: schemas.Inputs,
164    ) -> Info:
165        with tempfile.TemporaryDirectory() as directory:
166            path = Path(directory) / "model.skops"
167            model.save_internal_model(path=str(path))
168            return mlflow.pyfunc.log_model(
169                name=self.path,
170                python_model=str(self.CODE),
171                artifacts={"model": str(path)},
172                model_config=model.model_dump(),
173                signature=signature,
174                input_example=input_example,
175            )

120    class Adapter(PythonModel):  # type: ignore[misc]
121        """Adapt a custom model to the Mlflow PyFunc flavor.
122
123        https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html?#mlflow.pyfunc.PythonModel
124        """
125
126        model: models.Model
127
128        def load_context(self, context: PythonModelContext) -> None:
129            """Rebuild the project model from its params and its internal model file.
130
131            Args:
132                context (mlflow.PythonModelContext): mlflow context with config and artifacts.
133            """
134            model = pdt.TypeAdapter(models.ModelKind).validate_python(context.model_config)
135            self.model = model.load_internal_model(path=context.artifacts["model"])
136
137        def predict(
138            self,
139            context: PythonModelContext,  # noqa: ARG002  # required by mlflow PythonModel.predict
140            model_input: pd.DataFrame,
141            params: dict[str, T.Any] | None = None,  # noqa: ARG002  # required by mlflow PythonModel.predict
142        ) -> schemas.Outputs:
143            """Generate predictions with a custom model for the given inputs.
144
145            Args:
146                context (mlflow.PythonModelContext): mlflow context.
147                model_input (pd.DataFrame): inputs to validate for the mlflow model.
148                params (dict[str, T.Any] | None): additional parameters.
149
150            Returns:
151                schemas.Outputs: validated outputs of the project model.
152            """
153            # MLflow accepts pandas here and validates the logged signature;
154            # Pandera enforces the project's constraints before the typed model call.
155            inputs = schemas.InputsSchema.check(data=model_input)
156            return self.model.predict(inputs=inputs)

178class BuiltinSaver(Saver):
179    """Saver for built-in models using an Mlflow flavor module.
180
181    https://mlflow.org/docs/latest/models.html#built-in-model-flavors
182
183    Parameters:
184        flavor (str): Mlflow flavor module to use for the serialization.
185        kwargs (dict[str, T.Any]): flavor-specific options for `log_model`.
186            e.g., sklearn saves with skops, which refuses to reload tree models
187            unless their types are trusted: `{"skops_trusted_types": [...]}`.
188    """
189
190    KIND: T.Literal["BuiltinSaver"] = "BuiltinSaver"
191
192    flavor: str
193    kwargs: dict[str, T.Any] = {}
194
195    @T.override
196    def save(
197        self,
198        model: models.Model,
199        signature: signers.Signature,
200        input_example: schemas.Inputs,
201    ) -> Info:
202        builtin_model = model.get_internal_model()
203        module = getattr(mlflow, self.flavor)
204        return module.log_model(
205            builtin_model,
206            name=self.path,
207            signature=signature,
208            input_example=input_example,
209            **self.kwargs,
210        )

218class Loader(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
219    """Base class for loading models from registry.
220
221    Separate model definition from deserialization.
222    e.g., to switch between deserialization flavors.
223    """
224
225    KIND: str
226
227    class Adapter(abc.ABC):
228        """Adapt any model for the project inference."""
229
230        @abc.abstractmethod
231        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
232            """Generate predictions with the internal model for the given inputs.
233
234            Args:
235                inputs (schemas.Inputs): validated inputs for the project model.
236
237            Returns:
238                schemas.Outputs: validated outputs of the project model.
239            """
240
241    @abc.abstractmethod
242    def load(self, uri: str) -> Loader.Adapter:
243        """Load a model from the model registry.
244
245        Args:
246            uri (str): URI of a model to load.
247
248        Returns:
249            Loader.Adapter: model loaded.
250        """

227    class Adapter(abc.ABC):
228        """Adapt any model for the project inference."""
229
230        @abc.abstractmethod
231        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
232            """Generate predictions with the internal model for the given inputs.
233
234            Args:
235                inputs (schemas.Inputs): validated inputs for the project model.
236
237            Returns:
238                schemas.Outputs: validated outputs of the project model.
239            """

253class CustomLoader(Loader):
254    """Loader for custom models using the Mlflow PyFunc module.
255
256    https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html
257    """
258
259    KIND: T.Literal["CustomLoader"] = "CustomLoader"
260
261    class Adapter(Loader.Adapter):
262        """Adapt a custom model for the project inference."""
263
264        def __init__(self, model: PyFuncModel) -> None:
265            """Initialize the adapter from an mlflow pyfunc model.
266
267            Args:
268                model (PyFuncModel): mlflow pyfunc model.
269            """
270            self.model = model
271
272        @T.override
273        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
274            # model validation is already done in predict
275            outputs = self.model.predict(data=inputs)
276            return T.cast(schemas.Outputs, outputs)
277
278    @T.override
279    def load(self, uri: str) -> CustomLoader.Adapter:
280        model = mlflow.pyfunc.load_model(model_uri=uri)
281        return CustomLoader.Adapter(model=model)

261    class Adapter(Loader.Adapter):
262        """Adapt a custom model for the project inference."""
263
264        def __init__(self, model: PyFuncModel) -> None:
265            """Initialize the adapter from an mlflow pyfunc model.
266
267            Args:
268                model (PyFuncModel): mlflow pyfunc model.
269            """
270            self.model = model
271
272        @T.override
273        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
274            # model validation is already done in predict
275            outputs = self.model.predict(data=inputs)
276            return T.cast(schemas.Outputs, outputs)

284class BuiltinLoader(Loader):
285    """Loader for built-in models using the Mlflow PyFunc module.
286
287    Note: use Mlflow PyFunc instead of flavors to use standard API.
288
289    https://mlflow.org/docs/latest/models.html#built-in-model-flavors
290    """
291
292    KIND: T.Literal["BuiltinLoader"] = "BuiltinLoader"
293
294    class Adapter(Loader.Adapter):
295        """Adapt a builtin model for the project inference."""
296
297        def __init__(self, model: PyFuncModel) -> None:
298            """Initialize the adapter from an mlflow pyfunc model.
299
300            Args:
301                model (PyFuncModel): mlflow pyfunc model.
302            """
303            self.model = model
304
305        @T.override
306        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
307            columns = list(schemas.OutputsSchema.to_schema().columns)
308            outputs = self.model.predict(data=inputs)  # unchecked data!
309            return schemas.Outputs(outputs, columns=columns, index=inputs.index)
310
311    @T.override
312    def load(self, uri: str) -> BuiltinLoader.Adapter:
313        model = mlflow.pyfunc.load_model(model_uri=uri)
314        return BuiltinLoader.Adapter(model=model)

294    class Adapter(Loader.Adapter):
295        """Adapt a builtin model for the project inference."""
296
297        def __init__(self, model: PyFuncModel) -> None:
298            """Initialize the adapter from an mlflow pyfunc model.
299
300            Args:
301                model (PyFuncModel): mlflow pyfunc model.
302            """
303            self.model = model
304
305        @T.override
306        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
307            columns = list(schemas.OutputsSchema.to_schema().columns)
308            outputs = self.model.predict(data=inputs)  # unchecked data!
309            return schemas.Outputs(outputs, columns=columns, index=inputs.index)

322class Register(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
323    """Base class for registring models to a location.
324
325    Separate model definition from its registration.
326    e.g., to change the model registry backend.
327
328    Parameters:
329        tags (dict[str, T.Any]): tags for the model.
330    """
331
332    KIND: str
333
334    tags: dict[str, T.Any] = {}
335
336    @abc.abstractmethod
337    def register(self, name: str, model_uri: str) -> Version:
338        """Register a model given its name and URI.
339
340        Args:
341            name (str): name of the model to register.
342            model_uri (str): URI of a model to register.
343
344        Returns:
345            Version: information about the registered model.
346        """

349class MlflowRegister(Register):
350    """Register for models in the Mlflow Model Registry.
351
352    https://mlflow.org/docs/latest/model-registry.html
353    """
354
355    KIND: T.Literal["MlflowRegister"] = "MlflowRegister"
356
357    @T.override
358    def register(self, name: str, model_uri: str) -> Version:
359        return mlflow.register_model(name=name, model_uri=model_uri, tags=self.tags)