MLFlow
MLflow is an open source platform for managing the end-to-end machine learning lifecycl. see MLflow docs 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.
To install packages, run the following command:
pip install -r requirements.txt
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.
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%.
Save a sample of data, its label, and model as a .zip file. The .zip file is needed during the AIShield API call.
Requirements: Get AIShield API endpoint, x-API-key, and ord_id from the AIShield team.
- 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.
- Upload the data, model, and label to the given URL.
- 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.
- 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
- Load the AIShield provided defense model and pass a sample of original data and attack data to get a prediction from the defense model.
- To integrate with SIEM solutions such as Microsoft Sentinel and Splunk, please follow the instructions below:
- For the Microsoft Sentinel connector, provide azure_log_customer_id and azure_log_shared_key while creating the AISDefenseModel object.
- For Splunk connector, provide splunk_url and splunk authorization while creating the AISDefenseModel object.
To learn more see the reference implementatoin notebook for FLflow