Adding implementation

docker-compose.yaml

@@ -31,6 +31,7 @@ services:
       - ./models:/workspace/models:rw
       - ./bentoml_configuration.yaml:${BENTOML_CONFIG}:rw
       - ./bento:${BENTOML_HOME}
+      - ./outdir:/workspace/outdir
     command: >
       mlflow server --serve-artifacts --host 0.0.0.0 --port 5000  
       --backend-store-uri "${MLFLOW_BACKEND_STORE_URI}" 

outdir/.gitignore

@@ -0,0 +1,2 @@
+*
+!.gitignore

pymc_mlflow/1_train_model.py

@@ -15,6 +15,7 @@
 import pymc as pm
 import xarray as xr
 from pymc_experimental.model_builder import ModelBuilder
+import os.path as osp
 
 from numpy.random import RandomState
 
@@ -95,66 +96,98 @@ def _generate_and_preprocess_model_data(
         self.X = X
         self.y = y
 
+class BayesModel(mlflow.pyfunc.PythonModel):
+    def __init__(self, module):
+        self.module = module
+        self.model = None
+
+    def load_context(self, context):
+        model_instance = self.module()
+        idata = context.artifacts["idata"]
+        self.model = model_instance.load(idata)
+
+    def predict(self, context, model_input, params={}):
+        if self.model is None:
+            raise ValueError(
+                "The model has not been loaded. "
+            )
+
+        prediction_data = pd.DataFrame({"input": model_input})
+        predict_proba = params.get("predict_proba")
+        return self.model.predict(prediction_data)
+
+        # if predict_proba:
+        #     return self.model.predict_proba(prediction_data)
+        # else:
+        #     return self.model.predict(prediction_data)
+
 @hydra.main(version_base=None, config_path="../conf", config_name="config")
 def main(config: DictConfig):
     RANDOM_SEED = 8927
-
     rng = np.random.default_rng(RANDOM_SEED)
     az.style.use("arviz-darkgrid")
 
     EXPERIMENT_CONFIG = instantiate(config["experiment"])
-
     mlflow.set_tracking_uri(EXPERIMENT_CONFIG.tracking_uri)
     experiment_name = EXPERIMENT_CONFIG.name
 
     existing_exp = mlflow.get_experiment_by_name(experiment_name)
     if not existing_exp:
         mlflow.create_experiment(experiment_name)
     mlflow.set_experiment(experiment_name)
-
     mlflow.set_tag("task", "pymc_mlflow_model")
 
     x = np.linspace(start=0, stop=1, num=100)
     X = pd.DataFrame(data=x, columns=["input"])
     y = 0.3 * x + 0.5 + rng.normal(0, 1, len(x))
 
     model = LinearModel()
-
     idata = model.fit(X, y)
-
-    # # Load training data set
-    # iris = datasets.load_iris()
-    # X, y = iris.data, iris.target
-
-    # # Train the model
-    # clf = svm.SVC(gamma='scale')
-    # clf.fit(X, y)
-
-    # signature = infer_signature(
-    #     X, 
-    #     y
-    # )
+    fname = osp.join("outdir", "linear_model_v1.nc")
+    model.save(fname)
+
+    model_2 = LinearModel.load(fname)
+    x_pred = np.random.uniform(low=1, high=2, size=10)
+    prediction_data = pd.DataFrame({"input": x_pred})
+    model_2.predict(prediction_data)
+
+    signature = infer_signature(
+        x, 
+        y,
+        params = { "predict_proba": True }
+    )
 
-    # model_name = "iris_clf"
-    # run_id = mlflow.active_run().info.run_id
-    # logged_model = mlflow.sklearn.log_model(
-    #     clf, artifact_path = model_name, signature = signature
-    # )
-    # model_uri = f"runs:/{run_id}/{model_name}" 
+    model_name = "bayes_reg"
+    run_id = mlflow.active_run().info.run_id
+    print(run_id)
+    bayes_model = BayesModel(module=LinearModel)
+
+    logged_model = mlflow.pyfunc.log_model(
+        python_model = bayes_model, 
+        artifact_path = model_name, 
+        artifacts={"idata": fname},
+        signature = signature
+    )
+    model_uri = logged_model.model_uri
+    mlflow.register_model(model_uri, model_name)
+
+    bento_model = bentoml.mlflow.import_model(
+        'bayes_reg', 
+        model_uri,
+        labels=mlflow.active_run().data.tags,
+        metadata={
+            "metrics": mlflow.active_run().data.metrics,
+            "params": mlflow.active_run().data.params,
+        }
+    )
 
-    # print(model_uri)
-    # print(logged_model.model_uri)
-
-    # mlflow.register_model(model_uri, model_name)
-    # bento_model = bentoml.mlflow.import_model(
-    #     'iris_clf', 
-    #     logged_model.model_uri,
-    #     labels=mlflow.active_run().data.tags,
-    #     metadata={
-    #     "metrics": mlflow.active_run().data.metrics,
-    #     "params": mlflow.active_run().data.params,
-    # })
+    loaded_model = mlflow.pyfunc.load_model(model_uri)
+    pred_mean = loaded_model.predict(x_pred, params = {})
+    # pred_samples = loaded_model.predict(
+    #     x_pred, params = { "predict_proba": True }
+    # )
 
+    print(pred_mean)
     mlflow.end_run()
 
 if __name__ == "__main__":

pymc_mlflow/2_deploy_model.py

@@ -0,0 +1,12 @@
+import numpy as np
+import bentoml
+from bentoml.io import NumpyNdarray
+
+bayes_reg_runner = bentoml.mlflow.get("bayes_reg:latest").to_runner()
+
+svc = bentoml.Service("bayes_reg", runners=[bayes_reg_runner])
+
+@svc.api(input=NumpyNdarray(), output=NumpyNdarray())
+def predict(input_series: np.ndarray) -> np.ndarray:
+    result = bayes_reg_runner.predict.run(input_series)
+    return result

scripts/pymc_mlflow.sh

@@ -4,4 +4,4 @@ CMD="docker compose exec mlflow python -m pymc_mlflow"
 
 ${CMD}.1_train_model
 
-# docker compose exec mlflow bentoml serve --host 0.0.0.0 -p 3000 pymc_mlflow.2_deploy_model:svc
+docker compose exec mlflow bentoml serve --host 0.0.0.0 -p 3000 pymc_mlflow.2_deploy_model:svc