ML Infrastructure Security Intermediate

The infrastructure supporting AI systems - model registries, training pipelines, GPU clusters, and serving platforms - presents traditional IT security risks amplified by the high value of ML assets. This lesson covers security testing of ML infrastructure components, supply chain assessment, and container security for AI deployments.

ML Infrastructure Components

Model Registry Security

Model registries store trained models, metadata, and deployment configurations. Test for:

• Access control - Can unauthorized users download model artifacts or weights?
• Model integrity - Are model checksums verified during deployment?
• Version control - Can an attacker replace a model version with a backdoored copy?
• Metadata exposure - Do model cards or metadata reveal sensitive training details?
• API authentication - Are registry APIs protected against unauthorized access?

Training Pipeline Security

Training pipelines often run with elevated privileges and access sensitive data:

• Pipeline injection - Can an attacker modify pipeline definitions to inject malicious training steps?
• Credential management - Are API keys, database credentials, or cloud tokens properly secured?
• Data access controls - Can training jobs access data outside their intended scope?
• Compute hijacking - Can GPU resources be diverted for cryptocurrency mining?
• Artifact integrity - Are intermediate artifacts (checkpoints, features) verified?

Container and Serving Security

ML serving often uses containers with ML frameworks that have large attack surfaces:

CONTAINER SECURITY CHECKS:
  [ ] Base images scanned for CVEs
  [ ] ML framework versions checked for known vulnerabilities
  [ ] No unnecessary packages installed (minimize attack surface)
  [ ] Containers run as non-root user
  [ ] Read-only file system where possible
  [ ] Network policies restrict container communication
  [ ] Secrets not baked into images
  [ ] Resource limits set (CPU, memory, GPU)

DESERIALIZATION RISKS:
  [ ] Pickle files not loaded from untrusted sources
  [ ] SafeTensors used instead of pickle for model weights
  [ ] Custom deserialization code reviewed for vulnerabilities
  [ ] Model loading sandboxed or restricted

Supply Chain Assessment

• Audit all ML framework dependencies for known CVEs
• Verify checksums and signatures for pretrained models downloaded from model hubs
• Scan pip/conda packages for typosquatting (e.g., tensorfow vs tensorflow)
• Review third-party dataset provenance and integrity
• Assess CI/CD pipeline security for training and deployment workflows