ML Security Challenges
ML infrastructure faces unique security challenges:
- Expensive GPUs: Attackers target GPU instances for crypto mining
- Model theft: Proprietary models are valuable IP
- Data exposure: Training data may contain PII
- Prompt injection: Malicious inputs can manipulate outputs
π‘ 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
- β Network: Firewall configured, private subnets for GPU instances
- β Authentication: JWT or API keys for all endpoints
- β Rate Limiting: Per-user and per-IP limits
- β Input Validation: Prompt injection protection
- β Output Filtering: PII and harmful content removal
- β Encryption: TLS for transit, encrypted storage
- β Secrets: No hardcoded credentials, use secrets manager
- β Model Protection: Read-only mounts, file permissions
- β Audit Logging: All requests logged for compliance
- β Monitoring: Alerts for suspicious activity
β οΈ 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:
- Network: Isolate GPU instances, use private networks
- Application: Authenticate, rate limit, validate inputs
- Data: Encrypt at rest and in transit
- Monitoring: Log everything, alert on anomalies
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.