Securing Model Supply Chain Risk with File Scanning

What is model supply chain risk?

Pickle is a widely used serialization format in Machine Learning, emerging as the most common way to share model and model weights. However, when you load and de-serialize a pickle file, it is susceptible to dangerous arbitrary code execution attacks specified by the original author of the file. To learn more about pickle serialization based attacks: check out our blog post on the topic.

Once a model has been registered with Robust Intelligence, you can run a file scan on the model artifacts, which will identify potentially dangerous imports within the model which pose risks during de-serialization. This is a critical step in securing your model supply chain. All pickle derived model formats, including PyTorch, TensorFlow, Keras, and Scikit-Learn, are supported for file scanning.

Note: File Scanning is currently only enabled on registered Hugging Face models.

Prepare a Hugging Face model for scanning

To prepare a Hugging Face model for scanning:

Create a project

Register the model. If your model is a gated model, additional steps are needed. See “Prepare a gated model for scanning,” below.

model_id = project.register_model(
    name="google/gemma-7b",
    model_config={
        "hugging_face": {
            "model_uri": "google/gemma-7b",
        }
    }
)

Retrieve the model and project IDs and proceed to “Run a file scan,” below.

Prepare a gated model for scanning

To scan a private or gated Hugging Face model, you must first set up a custom integration for the model. The integration stores sensitive tokens/secrets that your artifacts (models/datasets) might need.

To prepare a gated Hugging Face model for scanning:

Create a project
Create the custom integration for the model. When you create the integration:
- Include a key HUGGINGFACE_API_TOKEN with its value set to your Hugging Face API token. Note: You can find your Huggingface API token on your Huggingface profile page.
- You must must manually accept the Hugging Face terms and conditions with the account that owns your API key.
Register the model. When you register the model:
- Specify the integration_id as an argument to the register_model method of the SDK’s Project class. See Project.register_model. Use the integration_id parameter and NOT the model_endpoint_integration_id.
- Copy the integration id and paste it into the integration_id. See the sample value, INTEGRATION_ID in the example below. This is where you register the model to be scanned:
```
model_id = project.register_model(
    name="google/gemma-7b",
    model_config={
        "hugging_face": {
            "model_uri": "google/gemma-7b",
        }
    },
    integration_id=INTEGRATION_ID
)
```
Retrieve the model and project IDs and proceed to “Run a file scan,” below.

Run a file scan on a Hugging Face model

Before you can scan a model, you must have created a project, registered the model, and retrieved the model ID and project ID.

Note: For large generative models, consider specifying the ram_request_megabytes field to a higher value such as 30000 in the start_file_scan method to ensure that the file scan job has enough memory to run.

Use the SDK’s Client.start_file_scan method to start the scan:

from rime_sdk import Client

client = Client(domain=<CLUSTER_URL>, api_key=<API_KEY>)
job = client.start_file_scan(model_id=<MODEL_ID>, project_id=<PROJECT_ID>)
job.get_status(verbose=True, wait_until_finish=True)

Get the results of a file scan

Retrieve via the SDK

Once the file scan job has completed, you can view the results of the file scan via the SDK:

client.get_file_scan_result(file_scan_id=job.job_id)

Severity is scored from 1-3, with 3 corresponding to an Alert level of risk, and 1 being a pass. If vulnerabilities were found during the file scan, they will be present under the file_security_reports field under unsafe_dependencies. The file_security_reports field will be empty if no vulnerabilities were found in any of the model files.

View in the UI

You can also view the latest file scan result for a specific model on its model registry page.

Sign in to a Robust Intelligence instance.

The Workspaces page appears.
Click a workspace.

The Workspaces summary page appears.
Select the project in which the model is registered.

The project overview page appears.
Click the Model Registry tab on the left navigation bar under Artifacts.

The model registry page appears.
Click the model for which you have run a file scan.

The model info page for the model appears.
The severity of the latest file scan result is displayed under the File Scan tab in the lower right.

There are three possible severity levels: Alert, Warning, Pass.

Understanding your file scan results

The result or severity of a file scan result is determined by the most dangerous import amongst all dependencies found across all files that comprise the model. Certain imports, such as any methods from the os or sys module, are considered more dangerous than others, such as numpy or pandas which are more commonly found in model artifacts. Some examples of dangerous imports include dependencies that enable arbitrary code execution, such as the builtin exec function, or libraries that create HTTP connections, such as httplib. Robust Intelligence maintains a list of dangerous dependencies, updating whenever a new CVE is identified to ensure that your model is secured against the latest vulnerabilities.