Building a Privacy-Preserving AI Infrastructure Stack

As AI systems become deeply embedded in healthcare, finance, and enterprise workflows, the data feeding these models has become both the most valuable and most sensitive asset in the stack. A privacy-preserving AI infrastructure is the set of tools, patterns, and services that allow you to train, serve, and reason over models without exposing raw user data to the system, the operators, or even the model provider. This tutorial walks through what such a stack looks like, why it matters, and how to build one with practical, working code.

What Is a Privacy-Preserving AI Infrastructure Stack?

A privacy-preserving AI stack layers cryptographic and statistical guarantees on top of the usual ML pipeline. Instead of a single monolith, it is composed of several cooperating components:

Data layer: Encrypted storage, tokenization, and access controls (e.g., envelope encryption with a KMS).
Compute layer: Confidential computing enclaves (Intel SGX, AMD SEV-SNP, Nitro Enclaves) where data is decrypted only inside CPU-protected memory.
Training layer: Federated learning and differential privacy to train without centralizing raw data.
Inference layer: Secure multi-party computation (SMPC) or homomorphic encryption for blind inference.
Audit layer: Immutable logs, policy engines, and consent management.

The goal is not a single product but an architecture where each layer adds a guarantee. Even if one component is compromised, the others prevent raw data from leaking.

Why It Matters

Regulations like GDPR, HIPAA, and the EU AI Act impose strict obligations on how personal data is processed. Beyond compliance, there are hard engineering reasons to adopt this stack:

Data minimization: You cannot leak what you never collect in plaintext.
Cross-org collaboration: Hospitals, banks, or competitors can jointly train models without sharing raw records.
Verifiable trust: Cryptographic attestation lets users verify that their data was handled by a specific, audited code path.
Breach containment: Encrypted data at rest and in use dramatically reduces the blast radius of a compromise.

Architecture Overview

Before writing code, it helps to sketch the end-to-end flow. A typical request travels through these stages:

Client
  │  (1) Encrypts payload with enclave public key
  ▼
API Gateway ──► Policy Engine (OPA) ──► Consent Ledger
  │
  ▼
Confidential Inference Enclave (attested)
  │  (2) Decrypts inside protected memory
  │  (3) Runs model
  │  (4) Returns encrypted result
  ▼
Audit Log (append-only, hash-chained)

The client never sends plaintext to the server. The server never sees plaintext outside the enclave. The audit log records that an inference happened, but not the data itself.

Step 1: Setting Up the Confidential Compute Layer

We will use AWS Nitro Enclaves as the confidential compute substrate because it is broadly accessible. The same pattern applies to AMD SEV-SNP or Intel TDX with minor changes. First, install the Nitro CLI on an EC2 instance launched from an enclave-enabled AMI:

sudo amazon-linux-extras install aws-nitro-enclaves-cli
sudo systemctl start nitro-enclaves-allocator
sudo usermod -aG nitro-enclaves ec2-user

Next, define a Dockerfile for the inference server. The enclave image (EIF) is built from this Dockerfile:

FROM python:3.11-slim

WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY model.py server.py ./

# The vsock-proxy will forward TCP to the parent instance
EXPOSE 5000

CMD ["python", "server.py"]

Build and run the enclave:

nitro-cli build-enclave \
  --docker-uri inference-server:latest \
  --output-file inference.eif

nitro-cli run-enclave \
  --cpu-count 2 --memory 4096 \
  --enclave-cid 16 \
  --eif-path inference.eif

The enclave boots with a cryptographic attestation document you can verify from the client side before sending any data.

Step 2: Attestation and Key Exchange

Attestation proves to the client that the enclave is running the exact EIF you expect, on genuine Nitro hardware. The attestation document contains a public key generated inside the enclave. The client uses it to encrypt the payload. Here is a Python implementation using the nitro-enclaves-sdk:

import base64
import json
import requests
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import hashes, serialization

def get_attestation_document(nonce: bytes) -> bytes:
    # The enclave exposes an endpoint that returns its attestation doc
    resp = requests.get("https://enclave.example/attestation",
                        params={"nonce": base64.b64encode(nonce).decode()})
    return base64.b64decode(resp.json()["document"])

def verify_and_extract_pubkey(document: bytes, expected_pcrs: dict) -> bytes:
    # In production, use the AWS Nitro Enclaves SDK COSE parser
    from nitro_enclaves_sdk.attestation import verify_attestation_document
    doc = verify_attestation_document(document, expected_pcrs)
    return doc["public_key"]

def encrypt_payload(pubkey_pem: bytes, payload: dict) -> str:
    pub = serialization.load_pem_public_key(pubkey_pem)
    ct = pub.encrypt(
        json.dumps(payload).encode(),
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        ),
    )
    return base64.b64encode(ct).decode()

The expected_pcrs are SHA-384 hashes of the enclave image, boot code, and certificate. They are produced when you build the EIF and should be pinned in your client configuration. If an attacker swaps the enclave image, the PCR values change and verification fails.

Step 3: The Inference Server Inside the Enclave

Inside the enclave, the server decrypts the payload, runs the model, and returns an encrypted response. Here is a minimal Flask server that does this:

import os
import json
import base64
from flask import Flask, request, jsonify
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from model import load_model, predict

app = Flask(__name__)

# Generate the enclave keypair at startup
private_key = rsa.generate_private_key(public_exponent=65537, key_size=4096)
public_key = private_key.public_key()

model = load_model("/app/weights.bin")

@app.route("/attestation")
def attestation():
    # In production, request the attestation document from the NSM device
    # and embed the public key in the user data field.
    return jsonify({
        "document": base64.b64encode(b"<attestation_doc>").decode(),
        "public_key": public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        ).decode(),
    })

@app.route("/predict", methods=["POST"])
def predict_route():
    encrypted = base64.b64decode(request.json["payload"])
    plaintext = private_key.decrypt(
        encrypted,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        ),
    )
    features = json.loads(plaintext)["features"]

    # Differential privacy: clip and add noise to the input
    import numpy as np
    features = np.clip(features, -1.0, 1.0)
    noise = np.random.normal(0, 0.01, size=len(features))
    features = (np.array(features) + noise).tolist()

    result = predict(model, features)
    return jsonify({"result": result.tolist()})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)

Note the differential privacy step before inference. Even though the data is decrypted inside the enclave, adding calibrated noise ensures that the output cannot be reverse-engineered to reconstruct the input exactly. This is a defense-in-depth measure against model inversion attacks.

Step 4: Federated Training with Differential Privacy

For training, centralizing data is often the bigger risk than inference. Federated learning keeps data on the client and only shares model updates. We will use Flower as the federated orchestrator and Opacus for differential privacy on the client side.

First, install the dependencies:

pip install flwr torch opacus

Define a client that trains locally and applies DP gradient noise:

import flwr as fl
import torch
from torch import nn
from opacus import PrivacyEngine

class DPClient(fl.client.NumPyClient):
    def __init__(self, model, train_loader, epochs=1, target_epsilon=8.0):
        self.model = model
        self.train_loader = train_loader
        self.optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
        self.criterion = nn.CrossEntropyLoss()

        self.privacy_engine = PrivacyEngine()
        self.model, self.optimizer, self.train_loader = (
            self.privacy_engine.make_private_with_epsilon(
                module=self.model,
                optimizer=self.optimizer,
                data_loader=self.train_loader,
                epochs=epochs,
                target_epsilon=target_epsilon,
                target_delta=1e-5,
                max_grad_norm=1.0,
            )
        )

    def get_parameters(self, config):
        return [val.cpu().numpy() for val in self.model.parameters()]

    def fit(self, parameters, config):
        fl.common.parameters_to_ndarrays(parameters)
        self.model.train()
        for _ in range(config.get("epochs", 1)):
            for x, y in self.train_loader:
                self.optimizer.zero_grad()
                out = self.model(x)
                loss = self.criterion(out, y)
                loss.backward()
                self.optimizer.step()
        epsilon = self.privacy_engine.get_epsilon(1e-5)
        return self.get_parameters(config), len(self.train_loader.dataset), {"epsilon": epsilon}

    def evaluate(self, parameters, config):
        # Local evaluation omitted for brevity
        return 0.0, len(self.train_loader.dataset), {}

On the server side, configure secure aggregation so the server never sees individual updates, only the aggregated sum:

import flwr as fl

strategy = fl.server.strategy.FedAvg(
    min_fit_clients=10,
    min_available_clients=10,
    fit_metrics_aggregation_fn=lambda m: {
        "epsilon": sum(v["epsilon"] for v in m.values()) / len(m)
    },
)

fl.server.start_server(
    server_address="0.0.0.0:8080",
    config=fl.server.ServerConfig(num_rounds=50),
    strategy=strategy,
)

Each client reports its spent epsilon, and the server logs the average. This gives you a concrete, auditable privacy budget rather than a vague promise.

Step 5: Policy and Consent Management

Even the best cryptography cannot protect data if the policy layer allows the wrong access. Use Open Policy Agent (OPA) to enforce rules like "inference requests for user X require active consent for purpose Y." Define a Rego policy:

package ai.privacy

default allow = false

allow {
    input.consent[input.purpose].active
    input.consent[input.purpose].expires_at > time.now_ns()
    input.data_classification == "deidentified"
}

deny[msg] {
    not input.consent[input.purpose].active
    msg := sprintf("no active consent for purpose %q", [input.purpose])
}

Call OPA from your API gateway before forwarding to the enclave:

import requests

def authorize(request_payload: dict) -> bool:
    resp = requests.post(
        "http://opa:8181/v1/data/ai/privacy/allow",
        json={"input": request_payload},
    )
    return resp.json().get("result", False)

Pair this with an append-only consent ledger. Every grant, revocation, and inference is recorded with a hash chain so auditors can reconstruct the full history without accessing the underlying data.

Best Practices

Pin PCR values in clients. Attestation is meaningless if clients accept any enclave. Hard-code expected PCRs and fail closed on mismatch.
Track the privacy budget end to end. Use a ledger that accumulates epsilon across training rounds and inference calls. Halt when the budget is exhausted.
Separate keys from compute. Store enclave wrapping keys in a hardware KMS (AWS KMS, GCP Cloud HSM) and require attestation documents before release.
Minimize data in the prompt. For LLM stacks, run a redaction proxy that strips PII before the request ever reaches the model, even inside the enclave.
Log metadata, not payloads. Audit logs should record who, when, and which model, never the input or output. Hash-chain logs to detect tampering.
Test with adversarial inversion. Periodically run model inversion and membership inference attacks against your own stack to validate that DP noise is sufficient.
Version your enclaves. Every model and code change produces new PCRs. Maintain a registry mapping versions to PCRs so clients can choose which to trust.

Conclusion

Building a privacy-preserving AI infrastructure stack is less about a single silver-bullet technology and more about composing layers that each contribute a guarantee: confidential computing protects data in use, differential privacy bounds what can be inferred from outputs, federated learning avoids centralizing raw data, and policy engines enforce consent and access rules. By combining attested enclaves, DP-aware training, secure aggregation, and an append-only audit trail, you create a system where privacy is a verifiable property of the architecture rather than a promise. Start with one layer, measure its overhead, and add the next only when the threat model justifies it. The result is an AI platform that regulators, partners, and users can trust by construction.