Integrating AIShield in Your MLOps Pipeline: A Guide for AI Developers

Photo by Christina @ wocintechchat.com on Unsplash

TL;DR

In under 30 minutes, follow the easy step-by-step process (including code snippets) presented in this guide to effortlessly integrate AIShield’s vulnerability analysis and defense generation features into your MLOps workflow. Experience these three key benefits by implementing AIShield:

• Benefit 1: Ensure the security, robustness, and resistance to adversarial attacks of your AI models by integrating AIShield in your MLOps pipeline.
• Benefit 2: Receive detailed vulnerability analysis and defense reports to identify and address potential weaknesses in your AI models.
• Benefit 3: Deploy secure, high-quality machine learning models with confidence, protecting your applications and users from adversarial attacks.

Act now to protect your organization’s AI models: Sign up for a FREE trial of AIShield API (AWS, non-AWS) today, and seize the opportunity to bolster security of AI model.

You can also check out Reference implementations here,

As AI developers, we strive to create robust and secure machine learning models. One essential aspect of this process is integrating security checks and vulnerability analysis within your MLOps pipeline. AIShield, an AI security platform, offers a powerful solution to analyze, detect, and mitigate potential vulnerabilities in your AI models. In this blog post, we will guide you through integrating AIShield in your MLOps pipeline to ensure your models are secure before deployment.

Overview of AIShield:

AIShield is a comprehensive AI security platform designed to analyze machine learning models for potential vulnerabilities, generate defenses, and provide detailed reports for developers. By integrating AIShield in your MLOps pipeline, you can ensure your models are secure, robust, and resistant to adversarial attacks.

Prerequisites:

To follow this guide, you should have a basic understanding of MLOps and experience with Python and machine learning libraries such as TensorFlow or PyTorch. You will also need an AIShield subscription to access the API.

Train Your Model:

Before integrating AIShield, train your machine learning model using your preferred framework (e.g., TensorFlow or PyTorch).

For the remaining of the explanation, we assume that you have used MNIST dataset and trained a simple CNN model in tensorflow. Let’s have fun

Steps to Integrate AIShield in Your MLOps Pipeline:

Below are the key steps for integrating AIShield into your MLOps pipeline, along with code snippets to guide you through the process.

1. Install AIShield Package: Install the AIShield package using pip:

pip install aishield

2. Import AIShield Library: In your Python script, import the AIShield library:

import aishield as ais

3. Prepare Your Model, Data, and Label Artifacts: Create zip folders for your model, data, and label artifacts. This is required for uploading these artifacts to the AIShield API.

"""
Description: Zip data
"""

from shutil import make_archive
import os

zip_path = "/path/to/zip/folder"
data_path = "/path/to/data/folder"
label_path = "/path/to/label/folder"
model_path = "/path/to/model/folder"

"""
Description: File paths
"""

make_archive(base_name=os.path.join(zip_path, "data"), root_dir=data_path, zip_format="zip")
make_archive(base_name=os.path.join(zip_path, "label"), root_dir=label_path, zip_format="zip")
make_archive(base_name=os.path.join(zip_path, "model"), root_dir=model_path, zip_format="zip")

4. Set Up AIShield API: Configure the AIShield API by providing the API URL, subscription key, and organization ID obtained from the AIShield developer portal.

"""
Description: AIShield API URL and subscription key
""" 
baseurl = "YOUR_API_URL"  # Fill in the API endpoint URL from AIShield developer portal under API tab
api_key = "YOUR_API_KEY"  # Fill in the subscription key from AIShield developer portal under My Dashboard tab
org_id = "YOUR_ORG_ID"    # Fill in the Org_ID provided in the welcome email
url=baseurl+"/api/ais/v1.5" # Updated with the latest version of API

5. Initialize AIShield API Client and Register the Model: Create an AIShield API client instance, register your model, and upload the input artifacts (data, label, and model zip files).

"""
Description: Initialize the AIShield API
"""
client = ais.AIShieldApi(api_url=url, api_key=api_key, org_id=org_id)


"""
Description: Define the task and analysis type
"""
task_type = ais.get_type("task", "image_classification")
analysis_type = ais.get_type("analysis", "extraction")



"""
Description: Perform model registration and upload the input artifacts
"""

status, job_details = client.register_model(task_type=task_type, analysis_type=analysis_type)
model_id = job_details.model_id

data_path = os.path.join(zip_path, 'data.zip')
label_path = os.path.join(zip_path, 'label.zip')
model_path = os.path.join(zip_path, 'model.zip')

upload_status = client.upload_input_artifacts(
    job_details=job_details,
    data_path=data_path,
    label_path=label_path,
    model_path=model_path,
)
print('Upload status: {}'.format(', '.join(upload_status)))

6. Configure Vulnerability Analysis: Define the vulnerability analysis configuration, such as input dimensions, number of classes, attack type, number of attack queries, and encryption strategy.

"""
Description: Specify the appropriate configs required for vulnerability analysis
"""

input_shape = (28, 28, 1)  # Example input shape for MNIST dataset
num_classes = 10  # Example number of classes for MNIST dataset

vuln_config = ais.VulnConfig(task_type=task_type,
                             analysis_type=analysis_type,
                             defense_generate=True)

vuln_config.input_dimensions = input_shape # input dimension for mnist digit classification
vuln_config.number_of_classes = = num_classes # number of classes for mnist digit classification
vuln_config.attack_type = = "greybox"  # greybox or blackbox depending upon the availability of information about Model, Data and Parameters
vuln_config.attack_queries = 60000  # Number of attack queries to be generated for testing model vulnerability
vuln_config.encryption_strategy = 0  # value 0 (or) 1, if model is unencrypted or encrypted(pyc) respectively

7. Perform Vulnerability Analysis: Run the vulnerability analysis using the AIShield API client. Monitor the progress through the provided URL and fetch the job status using the Job ID.

"""
Description: Run vulnerability analysis
"""
job_status, job_details = client.vuln_analysis(model_id, vuln_config)
job_id = job_details.job_id 


"""
Description: Monitor progress for given Job ID using the Link below
"""
printf("Job URL: {client.get_job_url(job_id)}")

"""
Description: Monitor progress for given Job ID periodically at every 10 seconds
"""
job_status = client.job_status(job_id)
while job_status.state != "success":
    printf("Job status: {job_status.state}")
    time.sleep(10)
    job_status = client.get_job_status(job_id)

8. Save and Download Reports and Artifacts: Download the vulnerability report, defense reports, defense artifacts (e.g., the model), and attack samples. Save these files in a specified output folder for further review and analysis.

"""
Description: Creating a directory to save the defense artifacts
"""
OUTPUT_PATH = os.path.join(os.getcwd(), "Output_Artifacts")
os.makedirs(OUTPUT_PATH, exist_ok=True)



if job_status.state == ais.get_type("job_state", "finished"):
    # Download Vulnerability Report
    output_conf = ais.OutputConf(report_type=ais.get_type("report", "vulnerability"),
                                 file_format=ais.get_type("file_format", "pdf"),
                                 save_folder_path=OUTPUT_PATH)
    vulnerability_report = client.save_job_report(job_id=job_id, output_config=output_conf)
    # Download Defense Report
    output_conf = ais.OutputConf(report_type=ais.get_type("report", "defense"),
                                 file_format=ais.get_type("file_format", "pdf"),
                                 save_folder_path=OUTPUT_PATH)
    defense_report = client.save_job_report(job_id=job_id, output_config=output_conf)
    # Download Defense Model
    output_conf = ais.OutputConf(report_type=ais.get_type("report", "defense_artifact"),
                                 file_format=ais.get_type("file_format", "zip"),
                                 save_folder_path=OUTPUT_PATH)
    defense_artifact = client.save_job_report(job_id=job_id, output_config=output_conf)
    # Download Attack Samples
    output_conf = ais.OutputConf(report_type=ais.get_type("report", "attack_samples"),
                                 save_folder_path=OUTPUT_PATH)
    attack_samples = client.save_job_report(job_id=job_id, output_config=output_conf)

Post using AIShield

Analyzing and Implementing Improvements

Review the vulnerability and defense reports generated by AIShield. Identify potential weaknesses and implement necessary improvements to enhance your model’s security by either hardening the model using supplied attack vectors or integrating the defense model provided.

Deploy Your Secure Model

Once you have addressed the vulnerabilities and enhanced your model, integrate the secure model back into your MLOps pipeline for deployment. This ensures that you are deploying a robust and secure machine learning model for real-world applications.

Conclusion:

Integrating AIShield in your MLOps pipeline is an essential step for AI developers to ensure the security and robustness of their machine learning models. By following the steps outlined in this guide, you can seamlessly incorporate AIShield’s vulnerability analysis and defense generation capabilities into your MLOps workflow. This enables you to deploy secure, high-quality models with confidence, protecting your applications and users from potential adversarial attacks.

Act now to protect your organization’s AI models: Sign up for a FREE trial of AIShield API (AWS, non-AWS) today, and seize the opportunity to bolster security of AI model. You can also check out Reference implementations here.