hacnho/mlflow-tensorspec-model-load-amplification-poc

MLFlow TensorSpec signature amplification PoC

This repository contains a benign security research proof of concept for a malformed MLmodel signature that causes MLflow to materialize a very large TensorSpec list during model loading.

Files

• MLmodel - the malicious PoC artifact with 600000 repeated TensorSpec entries
• build-meta.json - artifact size and entry count
• requirements.txt - minimal Python dependencies for the local replay
• build_mlflow_tensorspec_signature_amp.py - PoC builder
• reproduce_mlflow_tensorspec_signature_amp.py - replay helper
• repro.sh - one-command local reproduction

What it demonstrates

The artifact is only about 42.92 MB, but MLflow's signature parsing path materializes the full TensorSpec object graph and drives close to or above 1 GB RSS, with noticeable load delay.

Quick repro

bash repro.sh 600000

Expected output:

[signature path]
{
  "mode": "signature",
  "entries": 600000,
  "elapsed_s": 4.83,
  "rss_mb": 900.4,
  "mlmodel_mb": 42.92
}

[model load path]
{
  "mode": "load",
  "entries": 600000,
  "elapsed_s": 20.824,
  "rss_mb": 900.8,
  "mlmodel_mb": 42.92
}

This PoC is intended for defensive validation and bug bounty triage only.