---
title: MLFlow
slug: mlflow
docTags: 
createdAt: 2023-05-16T09:23:08.000Z
---

MLflow is an open source platform for managing the end-to-end machine learning lifecycl. see [MLflow docs](https://mlflow.org/docs/latest/index.html) to learn more.

This document provides an overview of the use of the AIShield API with MLflow and logs artifact received from the API. In the following sections, you can learn how to install packages, load datasets, train the model, prepare an artifact for an AIShield API call, and perform the AIShield API call.

### Installation

To install packages, run the following command:

`pip install -r requirements.txt`

### Load Dataset

Download the MNIST dataset from TensorFlow, preprocess it, and split the data set into train, test, and validation in a ratio of 54k:6k:10K. Visualize sample data.

### Model Training

Create model architecture and train it on training data. After training, validate on the validation dataset. Model accuracy is expected to be around 99.3%.

### Prepare Artifact for the AIShield API Call

Save a sample of data, its label, and model as a .zip file. The .zip file is needed during the AIShield API call.

### AIShield API Call

Requirements: Get AIShield API endpoint, x-API-key, and ord\_id from the AIShield team.

1. Call the model registration API. This API will return a unique model id and path to upload data, model, and label. It will also return a sample request payload for a model analysis API.
2. Upload the data, model, and label to the given URL.
3. After successful upload, call the model analysis API and pass the payload. In the response, you will get a unique job\_id and monitor link. The monitor link can be used to track the progress of the triggered job.
4. Monitor the progress of the job ID using the get API, and after successful completion, download the artifacts and log them to MLflow as an artifact. The artifacts will contain vulnerability and defense reports, attack samples, defense artifact
5. Load the AIShield provided defense model and pass a sample of original data and attack data to get a prediction from the defense model.
6. To integrate with SIEM solutions such as Microsoft Sentinel and Splunk, please follow the instructions below:
   1.  For the Microsoft Sentinel connector, provide azure\_log\_customer\_id and azure\_log\_shared\_key while creating the AISDefenseModel object.
   2. For Splunk connector, provide splunk\_url and splunk authorization while creating the AISDefenseModel object.

:::hint{type="info"}
To learn more see the [reference implementatoin](https://github.com/bosch-aisecurity-aishield/Reference-Implementations/tree/main/Working_with_ML_Pipelines_and_Partner_Tools/MLflow) notebook for FLflow
:::