Create mlserve.py

mlserve.py

@@ -0,0 +1,82 @@
+import mlflow.sklearn
+from flask import Flask, request, jsonify
+from joblib import load
+import requests
+import sys
+
+def download_model(model_url, local_path):
+    """Downloads a model from the specified URL and saves it locally."""
+    try:
+        response = requests.get(model_url)
+        with open(local_path, 'wb') as f:
+            f.write(response.content)
+    except Exception as e:
+        print(f"Error downloading model: {e}")
+        raise  # Re-raise the exception for proper handling
+
+def load_model_and_log(model_path):
+    """Loads a pre-trained Scikit-learn model from the specified path and logs it with MLflow."""
+    try:
+        # Load the model using joblib
+        model = load(model_path)
+
+        # Log the loaded model with MLflow
+        mlflow.sklearn.log_model(model, "loaded_model")
+
+        # Get the active MLflow run (the one you just logged)
+        active_run = mlflow.active_run()
+
+        # Get the run ID
+        run_id = active_run.info.run_id
+
+        print("Run ID:", run_id)
+
+        return model, run_id
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        raise  # Re-raise the exception for proper handling
+
+def load_model_from_mlflow(run_id):
+    """Loads a model from MLflow based on the provided run ID."""
+    try:
+        # Load the model from MLflow
+        loaded_model = mlflow.sklearn.load_model(f"runs:/{run_id}/loaded_model")
+
+        return loaded_model
+    except Exception as e:
+        print(f"Error loading model from MLflow: {e}")
+        raise  # Re-raise the exception for proper handling
+
+def predict(model):
+    """Makes predictions using the provided model."""
+    try:
+        input_data = request.get_json()
+        prediction = model.predict(input_data)
+        return jsonify({'prediction': prediction.tolist()})
+    except Exception as e:
+        print(f"Error making prediction: {e}")
+        return jsonify({'error': str(e)}), 500  # Internal Server Error
+
+app = Flask(__name__)
+
+# Get model URL from command line argument
+model_url = sys.argv[1]
+
+# Path to save the downloaded model locally
+local_model_path = "/tmp/model.joblib"
+
+# Download the model from the specified URL
+download_model(model_url, local_model_path)
+
+# Load the model on app startup and log it with MLflow
+model, run_id = load_model_and_log(local_model_path)
+
+# Load the model from MLflow based on the run ID
+loaded_model = load_model_from_mlflow(run_id)
+
+@app.route('/predict', methods=['POST'])
+def inference():
+    return predict(loaded_model)
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=5000)