✨ Enable to override params at predict time in KedroPipeline…

kedro_mlflow/framework/hooks/mlflow_hook.py

@@ -389,7 +389,18 @@ def after_pipeline_run(
                     if isinstance(model_signature, str):
                         if model_signature == "auto":
                             input_data = catalog.load(pipeline.input_name)
-                            model_signature = infer_signature(model_input=input_data)
+
+                            # all pipeline params will be overridable at predict time: https://mlflow.org/docs/latest/model/signatures.html#model-signatures-with-inference-params
+                            # I add the special "runner" parameter to be able to choose it at runtime
+                            pipeline_params = {
+                                ds_name[7:]: catalog.load(ds_name)
+                                for ds_name in pipeline.inference.inputs()
+                                if ds_name.startswith("params:")
+                            } | {"runner": "SequentialRunner"}
+                            model_signature = infer_signature(
+                                model_input=input_data,
+                                params=pipeline_params,
+                            )
 
                     mlflow.pyfunc.log_model(
                         python_model=kedro_pipeline_model,

kedro_mlflow/mlflow/kedro_pipeline_model.py

@@ -196,17 +196,38 @@ def load_context(self, context):
             updated_catalog._datasets[name]._filepath = Path(uri)
             self.loaded_catalog.save(name=name, data=updated_catalog.load(name))
 
-    def predict(self, context, model_input):
+    def predict(self, context, model_input, params=None):
         # we create an empty hook manager but do NOT register hooks
         # because we want this model be executable outside of a kedro project
+
+        # params can pass
+        # TODO globals
+        # TODO runtime
+        # TODO parameters -> I'd prefer not have them, but it would require catalog to be able to not be fully resolved if we want to pass runtime and globals
+        # TODO hooks
+        # TODO runner
+
+        params = params or {}
+        runner_class = params.pop("runner", "SequentialRunner")
+        runner = (
+            self.runner
+        )  # runner="build it dynamically from runner class" or self.runner
+
         hook_manager = _create_hook_manager()
+        # _register_hooks(hook_manager, predict_params.hooks)
+
+        for name, value in params.items():
+            # no need to check if params are ni the catalog, because mlflow already checks that the params mathc the signature
+            param = f"params:{name}"
+            self._logger.info(f"Using {param}={value} for the prediction")
+            self.loaded_catalog.save(name=param, data=value)
 
         self.loaded_catalog.save(
             name=self.input_name,
             data=model_input,
         )
 
-        run_output = self.runner.run(
+        run_output = runner.run(
             pipeline=self.pipeline,
             catalog=self.loaded_catalog,
             hook_manager=hook_manager,