A secure AI evaluation framework for running AI safety benchmarks within AWS Nitro Enclaves, combining Rust and Python components.
The following repository contains all artifacts for this MPhil project, which is a mix of rust and python (and too many scripts to glue stuff together). A good starting point is to look at root Makefile and the README.md files in all the subdirectories:
analysis- code and artefacts to reproduce plots in my thesis based on benchmark runsenclave- first prototype for running the enclave and the host, sending files, running llama in the enclave, generating attestation documents. Does not include AI Safety benchmarks.llama_runner- running llama et al. using llama-cpp2 through the actor modelbert_runner- running bert models using rust-bert through the actor model (compiles against mock implementation by default or links with libpytdorch if compiled through-F use_rust_bert)evaluation- scaffolding (messages, file_transfer, dataset handling) to run AI safety benchmarks in haven. Configurations and prompts for the audit code for different graphs is included inevaluation/tasks. Recreates huggingfaces dataset library in dataset.rs. Includes code to run the evaluation withoutevaluation_enclave- protocol for the enclave side. uses type-state pattern:InitializedState -> LlamaLoadedState -> BertLoadedState -> DatasetLoadedState -> EvaluatedState -> AttestedStateevaluation_host- protocol for the host side. uses type-state-pattern:Disconnected -> Connected -> LlamaSent -> BertSent -> DatasetSent -> EvaluationComplete -> AttestationReceivedquantization- quantize models in AWS Nitro Enclaves (through llama.cpp and pytorch)scripts- the ducttape (basically)vsock- abstracts away (most) of the vsock trouble, takes a server/client handle that interact with the sockgpu_baseline- essentiallyevaluationbut in Python to run on GPUs through vLLM
Make sure you have AWS Nitro CLI and SDK installed here and sufficient memory allocated in the nitro allocator.
From the project directory:
cd enclave && cargo build
make build-docker && make build-eif && make run-enclaveRun make terminate-enclave if you wish to stop the enclave.
cd enclave && cargo build && cargo runMake sure to have uv installed.
All the raw data is in quantization_ablation_model for local analysis and in the remote_experiments for AWS. Make sure input datasets contain classification pairs for postprocessing.
cd analysis
uv run local_analysis.py
uv run aws_analysis.py

