ML Security Challenges

ML infrastructure faces unique security challenges:

πŸ’‘ Defense in Depth

Secure at every layer: network, authentication, application, and data. No single control is sufficient.

Network Security

Firewall Configuration

# UFW basic setup for GPU server
sudo ufw default deny incoming
sudo ufw default allow outgoing

# Allow SSH (consider restricting to specific IPs)
sudo ufw allow from 10.0.0.0/8 to any port 22

# Allow inference API (through load balancer only)
sudo ufw allow from 10.0.0.0/8 to any port 8000

# Allow metrics scraping (internal only)
sudo ufw allow from 10.0.0.0/8 to any port 9090

# Block all other incoming
sudo ufw enable

VPC and Private Networks

# Keep GPU instances in private subnet
# Only expose through load balancer

β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
β”‚                   VPC (10.0.0.0/16)              β”‚
β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€
β”‚  Public Subnet (10.0.1.0/24)                    β”‚
β”‚  β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”  β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”         β”‚
β”‚  β”‚ Load Balancer β”‚  β”‚  NAT Gateway   β”‚         β”‚
β”‚  β”‚   (HTTPS)     β”‚  β”‚                β”‚         β”‚
β”‚  β””β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”˜  β””β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜         β”‚
β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€
β”‚  Private Subnet (10.0.2.0/24)                   β”‚
β”‚  β”Œβ”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β”  β”Œβ”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β”€β”         β”‚
β”‚  β”‚  GPU Server 1 β”‚  β”‚  GPU Server 2  β”‚         β”‚
β”‚  β”‚  (No public   β”‚  β”‚  (No public    β”‚         β”‚
β”‚  β”‚   IP)         β”‚  β”‚   IP)          β”‚         β”‚
β”‚  β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜  β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜         β”‚
β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜

Kubernetes Network Policy

# k8s/network-policy.yaml
apiVersion: networking.k8s.io/v1
kind: NetworkPolicy
metadata:
  name: inference-server-policy
  namespace: ml
spec:
  podSelector:
    matchLabels:
      app: inference-server
  policyTypes:
    - Ingress
    - Egress
  ingress:
    # Only allow from ingress controller
    - from:
        - namespaceSelector:
            matchLabels:
              name: ingress-nginx
      ports:
        - protocol: TCP
          port: 8000
    # Allow Prometheus scraping
    - from:
        - namespaceSelector:
            matchLabels:
              name: monitoring
      ports:
        - protocol: TCP
          port: 9090
  egress:
    # Allow DNS
    - to:
        - namespaceSelector: {}
      ports:
        - protocol: UDP
          port: 53
    # Allow HTTPS for model downloads
    - to:
        - ipBlock:
            cidr: 0.0.0.0/0
      ports:
        - protocol: TCP
          port: 443

API Authentication

JWT Authentication

from fastapi import FastAPI, Depends, HTTPException, Security
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
import jwt
from datetime import datetime, timedelta

app = FastAPI()
security = HTTPBearer()

SECRET_KEY = "your-secret-key"  # Use env var in production
ALGORITHM = "HS256"

def create_token(user_id: str, expires_hours: int = 24) -> str:
    payload = {
        "sub": user_id,
        "exp": datetime.utcnow() + timedelta(hours=expires_hours),
        "iat": datetime.utcnow()
    }
    return jwt.encode(payload, SECRET_KEY, algorithm=ALGORITHM)

def verify_token(credentials: HTTPAuthorizationCredentials = Security(security)):
    try:
        payload = jwt.decode(
            credentials.credentials, 
            SECRET_KEY, 
            algorithms=[ALGORITHM]
        )
        return payload["sub"]
    except jwt.ExpiredSignatureError:
        raise HTTPException(status_code=401, detail="Token expired")
    except jwt.InvalidTokenError:
        raise HTTPException(status_code=401, detail="Invalid token")

@app.post("/v1/generate")
async def generate(
    prompt: str,
    user_id: str = Depends(verify_token)
):
    # User is authenticated
    return {"result": "..."}

API Key Authentication

from fastapi import Security
from fastapi.security import APIKeyHeader
import hashlib

api_key_header = APIKeyHeader(name="X-API-Key")

# Store hashed API keys in database
API_KEYS = {
    hashlib.sha256(b"sk-live-xxx").hexdigest(): "user-123",
    hashlib.sha256(b"sk-live-yyy").hexdigest(): "user-456",
}

def verify_api_key(api_key: str = Security(api_key_header)):
    key_hash = hashlib.sha256(api_key.encode()).hexdigest()
    user_id = API_KEYS.get(key_hash)
    if not user_id:
        raise HTTPException(status_code=401, detail="Invalid API key")
    return user_id

Rate Limiting

from fastapi import Request
from slowapi import Limiter
from slowapi.util import get_remote_address

limiter = Limiter(key_func=get_remote_address)

# Rate limit by IP
@app.post("/v1/generate")
@limiter.limit("100/minute")
async def generate(request: Request, prompt: str):
    ...

# Rate limit by API key
def get_api_key(request: Request):
    return request.headers.get("X-API-Key", get_remote_address(request))

limiter_by_key = Limiter(key_func=get_api_key)

@app.post("/v1/generate")
@limiter_by_key.limit("1000/hour")
async def generate_with_key(request: Request, prompt: str):
    ...

Secure GPU Infrastructure

GPUBrazil provides isolated instances with private networking and firewall controls built-in.

Get Started Securely β†’

Prompt Injection Protection

import re
from typing import Optional

class PromptGuard:
    """Detect and prevent prompt injection attacks"""
    
    INJECTION_PATTERNS = [
        r"ignore (all |previous |above )?instructions",
        r"disregard (all |previous |above )?instructions", 
        r"forget (all |previous |above )?instructions",
        r"you are now",
        r"pretend (to be|you are)",
        r"act as",
        r"new instructions:",
        r"system prompt:",
    ]
    
    def __init__(self):
        self.patterns = [
            re.compile(p, re.IGNORECASE) 
            for p in self.INJECTION_PATTERNS
        ]
    
    def check(self, prompt: str) -> Optional[str]:
        """Return detected pattern or None if safe"""
        for pattern in self.patterns:
            if pattern.search(prompt):
                return pattern.pattern
        return None
    
    def sanitize(self, prompt: str) -> str:
        """Remove potentially dangerous patterns"""
        sanitized = prompt
        for pattern in self.patterns:
            sanitized = pattern.sub("[FILTERED]", sanitized)
        return sanitized

# Usage
guard = PromptGuard()

@app.post("/v1/generate")
async def generate(prompt: str):
    injection = guard.check(prompt)
    if injection:
        # Log the attempt
        logger.warning(f"Prompt injection attempt: {injection}")
        raise HTTPException(
            status_code=400, 
            detail="Invalid prompt content"
        )
    
    # Process safe prompt
    return generate_response(prompt)

Output Filtering

import re
from presidio_analyzer import AnalyzerEngine
from presidio_anonymizer import AnonymizerEngine

# Initialize PII detection
analyzer = AnalyzerEngine()
anonymizer = AnonymizerEngine()

def filter_pii(text: str) -> str:
    """Remove personally identifiable information from output"""
    results = analyzer.analyze(
        text=text,
        entities=[
            "EMAIL_ADDRESS",
            "PHONE_NUMBER", 
            "CREDIT_CARD",
            "US_SSN",
            "PERSON",
            "LOCATION"
        ],
        language="en"
    )
    
    anonymized = anonymizer.anonymize(
        text=text,
        analyzer_results=results
    )
    
    return anonymized.text

def filter_harmful_content(text: str) -> str:
    """Basic harmful content filter"""
    # Use a content moderation API for production
    # This is a simple example
    harmful_patterns = [
        r"how to (make|build|create) (a )?(bomb|weapon|explosive)",
        r"hack(ing)? (into|someone)",
    ]
    
    for pattern in harmful_patterns:
        text = re.sub(pattern, "[CONTENT REMOVED]", text, flags=re.IGNORECASE)
    
    return text

@app.post("/v1/generate")
async def generate(prompt: str):
    response = model.generate(prompt)
    
    # Filter output
    response = filter_pii(response)
    response = filter_harmful_content(response)
    
    return {"result": response}

Data Protection

Encryption at Rest

# Encrypt model storage
# Option 1: LUKS encryption for disk
sudo cryptsetup luksFormat /dev/sdb
sudo cryptsetup luksOpen /dev/sdb encrypted_models
sudo mkfs.ext4 /dev/mapper/encrypted_models
sudo mount /dev/mapper/encrypted_models /models

# Option 2: Use encrypted EBS/persistent volumes
# AWS: Enable EBS encryption
# GCP: Enable default encryption
# Azure: Enable Azure Disk Encryption

Encryption in Transit

# Always use HTTPS for API endpoints
# nginx.conf
server {
    listen 443 ssl http2;
    server_name api.example.com;
    
    ssl_certificate /etc/letsencrypt/live/api.example.com/fullchain.pem;
    ssl_certificate_key /etc/letsencrypt/live/api.example.com/privkey.pem;
    
    # Strong SSL configuration
    ssl_protocols TLSv1.2 TLSv1.3;
    ssl_prefer_server_ciphers on;
    ssl_ciphers ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256;
    
    # HSTS
    add_header Strict-Transport-Security "max-age=31536000" always;
    
    location / {
        proxy_pass http://localhost:8000;
        proxy_set_header Host $host;
        proxy_set_header X-Real-IP $remote_addr;
    }
}

Secrets Management

# Never hardcode secrets!

# Bad ❌
API_KEY = "sk-live-xxxx"

# Good βœ… - Use environment variables
import os
API_KEY = os.environ["API_KEY"]

# Better βœ… - Use secrets manager
from google.cloud import secretmanager

def get_secret(secret_id: str) -> str:
    client = secretmanager.SecretManagerServiceClient()
    name = f"projects/my-project/secrets/{secret_id}/versions/latest"
    response = client.access_secret_version(request={"name": name})
    return response.payload.data.decode("UTF-8")

API_KEY = get_secret("inference-api-key")

Model Security

Protect Model Weights

# Restrict model file access
chmod 600 /models/*.safetensors
chown ml-service:ml-service /models/

# Use read-only mounts in containers
docker run -v /models:/models:ro inference-server

# Kubernetes - read only
volumes:
  - name: models
    persistentVolumeClaim:
      claimName: model-pvc
      readOnly: true

Model Signing

import hashlib
import json
from pathlib import Path

def sign_model(model_path: str, private_key: str) -> dict:
    """Create signed manifest for model files"""
    manifest = {
        "model_path": model_path,
        "files": {}
    }
    
    for file_path in Path(model_path).rglob("*"):
        if file_path.is_file():
            with open(file_path, "rb") as f:
                file_hash = hashlib.sha256(f.read()).hexdigest()
            manifest["files"][str(file_path)] = file_hash
    
    # Sign with private key (simplified)
    manifest_str = json.dumps(manifest, sort_keys=True)
    manifest["signature"] = sign_with_key(manifest_str, private_key)
    
    return manifest

def verify_model(model_path: str, manifest: dict, public_key: str) -> bool:
    """Verify model hasn't been tampered with"""
    for file_path, expected_hash in manifest["files"].items():
        with open(file_path, "rb") as f:
            actual_hash = hashlib.sha256(f.read()).hexdigest()
        if actual_hash != expected_hash:
            return False
    return True

Audit Logging

import logging
from datetime import datetime
import json

# Configure audit logger
audit_logger = logging.getLogger("audit")
audit_logger.setLevel(logging.INFO)
handler = logging.FileHandler("/var/log/ml-audit.log")
handler.setFormatter(logging.Formatter('%(message)s'))
audit_logger.addHandler(handler)

def log_inference_request(
    user_id: str,
    prompt: str,
    response: str,
    latency_ms: float,
    ip_address: str
):
    """Log inference request for audit trail"""
    audit_logger.info(json.dumps({
        "timestamp": datetime.utcnow().isoformat(),
        "event": "inference_request",
        "user_id": user_id,
        "ip_address": ip_address,
        "prompt_length": len(prompt),
        "response_length": len(response),
        "latency_ms": latency_ms,
        # Don't log actual prompt/response for privacy
        # unless required for compliance
    }))

def log_security_event(
    event_type: str,
    user_id: str,
    details: dict
):
    """Log security-relevant events"""
    audit_logger.warning(json.dumps({
        "timestamp": datetime.utcnow().isoformat(),
        "event": event_type,
        "user_id": user_id,
        "details": details
    }))

Security Checklist

⚠️ Regular Security Audits

Security is not set-and-forget. Schedule regular penetration tests and security reviews, especially after major changes.

Conclusion

ML infrastructure security requires attention at every layer:

Start with the basicsβ€”authentication and network isolationβ€”then progressively enhance your security posture.

Deploy securely on GPUBrazil with isolated instances, private networking, and built-in security controls.