From 96e4c03d119652db5bf134f6f79367a8b6a0466c Mon Sep 17 00:00:00 2001
From: Wojciech Lidwin <59863785+Halal37@users.noreply.github.com>
Date: Fri, 12 May 2023 15:17:15 +0200
Subject: [PATCH] mlflow

---
 Dockerfile        |  1 +
 Jenkinsfile_train |  2 +-
 ium_train.py      | 55 ++++++++++++++++++++++++++++++++---------------
 3 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index a3dc3bc..48bdbe6 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -17,6 +17,7 @@ RUN pip3 install matplotlib
 RUN pip3 install sacred
 RUN pip3 install pymongo
 RUN pip3 install dvc
+RUN pip3 install mlflow
 
 
 WORKDIR /app
diff --git a/Jenkinsfile_train b/Jenkinsfile_train
index 5efcdaa..b943ed8 100644
--- a/Jenkinsfile_train
+++ b/Jenkinsfile_train
@@ -25,7 +25,7 @@ node {
      checkout([$class: 'GitSCM', branches: [[name: '*/master']], extensions: [], userRemoteConfigs: [[credentialsId: 's487197', url: 'https://git.wmi.amu.edu.pl/s487197/ium_487197']]])
  }
  stage('Dockerfile'){
-    def testImage = docker.image('s487197/ium:52')
+    def testImage = docker.image('s487197/ium:55')
 	testImage.inside{ 
          copyArtifacts filter: 'baltimore_train.csv', projectName: 's487197-create-dataset'
 		 sh "python3 ium_sacred.py -epochs $EPOCHS -lr $LR -validation_split $VALIDATION_SPLIT"
diff --git a/ium_train.py b/ium_train.py
index ff960f0..15605b4 100644
--- a/ium_train.py
+++ b/ium_train.py
@@ -11,6 +11,13 @@ import numpy as np
 from sklearn.preprocessing import LabelEncoder
 import argparse
 import shutil
+import mlflow
+import logging
+logging.basicConfig(level=logging.WARN)
+logger = logging.getLogger(__name__)
+
+mlflow.set_tracking_uri("http://localhost:5000")
+mlflow.set_experiment("s487197")
 
 def get_x_y(data):
 
@@ -40,24 +47,38 @@ def train_model():
     data_train, x_train, y_train = get_x_y(train)
     normalizer = tf.keras.layers.Normalization(axis=1)
     normalizer.adapt(np.array(x_train))
-    model = Sequential(normalizer)
-    model.add(Dense(64, activation="relu"))
-    model.add(Dense(10, activation='relu'))
-    model.add(Dense(10, activation='relu'))
-    model.add(Dense(10, activation='relu'))
-    model.add(Dense(5, activation="softmax"))
-    model.compile(Adam(learning_rate=args.lr), loss='sparse_categorical_crossentropy', metrics = ['accuracy'] )
-    model.summary()
+    with mlflow.start_run() as run:
+        print("MLflow run experiment_id: {0}".format(run.info.experiment_id))
+        print("MLflow run artifact_uri: {0}".format(run.info.artifact_uri))
+        mlflow.log_param("epochs", args.epochs)
+        mlflow.log_param("lr", args.lr)
+        mlflow.log_param("validation_split", args.validation_split)
+        model = Sequential(normalizer)
+        model.add(Dense(64, activation="relu"))
+        model.add(Dense(10, activation='relu'))
+        model.add(Dense(10, activation='relu'))
+        model.add(Dense(10, activation='relu'))
+        model.add(Dense(5, activation="softmax"))
+        model.compile(Adam(learning_rate=args.lr), loss='sparse_categorical_crossentropy', metrics = ['accuracy'] )
+        model.summary()
 
-    history = model.fit(
-        x_train,
-        y_train,
-        epochs=args.epochs,
-        validation_split=args.validation_split)
-    hist = pd.DataFrame(history.history)
-    hist['epoch'] = history.epoch
-    model.save('baltimore_model')
-    shutil.make_archive('baltimore', 'zip', 'baltimore_model')
+        history = model.fit(
+           x_train,
+           y_train,
+           epochs=args.epochs,
+           validation_split=args.validation_split)
+        mlflow.log_metric("loss", float(, hist['loss']))
+        mlflow.log_metric('accuracy', float(hist['accuracy']))
+        signature = mlflow.models.signature.infer_signature(train_x, model.predict(x_test))
+        if tracking_url_type_store != "file":
+
+            mlflow.sklearn.log_model(model, "wines-model", registered_model_name="ElasticnetWineModel", signature=signature)
+        else:
+            mlflow.sklearn.log_model(model, "model", signature=signature)
+        hist = pd.DataFrame(history.history)
+        hist['epoch'] = history.epoch
+        model.save('baltimore_model')
+        shutil.make_archive('baltimore', 'zip', 'baltimore_model')
 
 
 train_model()