AI Workload Isolation

import subprocess
import json
from typing import Dict, List
 
class ContainerSecurityAuditor:
    """Audit container security configuration for AI workloads."""
 
    def __init__(self, container_id: str):
        self.container_id = container_id
 
    def inspect_container(self) -> Dict:
        """Retrieve container configuration."""
        result = subprocess.run(
            ["docker", "inspect", self.container_id],
            capture_output=True, text=True, check=True,
        )
        return json.loads(result.stdout)[0]
 
    def check_gpu_isolation(self) -> List[Dict]:
        """Check for GPU isolation weaknesses in the container configuration."""
        config = self.inspect_container()
        findings = []
 
        # Check for privileged mode
        if config.get("HostConfig", {}).get("Privileged", False):
            findings.append({
                "severity": "critical",
                "finding": "Container running in privileged mode",
                "detail": "Privileged mode disables all container isolation. "
                         "GPU workloads should use --gpus flag instead.",
                "remediation": "Use 'docker run --gpus device=0' instead of '--privileged'",
            })
 
        # Check for device mappings
        devices = config.get("HostConfig", {}).get("Devices", []) or []
        nvidia_devices = [d for d in devices if "nvidia" in str(d)]
        if nvidia_devices:
            findings.append({
                "severity": "medium",
                "finding": f"Direct NVIDIA device access: {len(nvidia_devices)} devices",
                "detail": "Direct device mapping bypasses NVIDIA container toolkit isolation",
                "remediation": "Use NVIDIA container toolkit with --gpus flag",
            })
 
        # Check for host PID namespace
        if config.get("HostConfig", {}).get("PidMode") == "host":
            findings.append({
                "severity": "high",
                "finding": "Container shares host PID namespace",
                "detail": "Host PID namespace allows observing GPU processes from other containers",
                "remediation": "Remove --pid=host flag",
            })
 
        # Check seccomp profile
        seccomp = config.get("HostConfig", {}).get("SecurityOpt", [])
        if any("seccomp=unconfined" in str(s) for s in seccomp):
            findings.append({
                "severity": "high",
                "finding": "Seccomp profile disabled",
                "detail": "Unconfined seccomp allows all syscalls including ioctl for GPU manipulation",
                "remediation": "Apply a custom seccomp profile that allows required GPU ioctls only",
            })
 
        # Check for SYS_ADMIN capability
        caps = config.get("HostConfig", {}).get("CapAdd", []) or []
        if "SYS_ADMIN" in caps:
            findings.append({
                "severity": "high",
                "finding": "SYS_ADMIN capability granted",
                "detail": "SYS_ADMIN enables mount namespace escapes and device access",
                "remediation": "Remove SYS_ADMIN and use specific capabilities instead",
            })
 
        return findings

#!/bin/bash
# Audit NVIDIA container toolkit configuration
# Run on the host system
 
echo "=== NVIDIA Container Toolkit Security Audit ==="
 
# Check toolkit version
echo "[*] NVIDIA Container Toolkit version:"
nvidia-container-cli --version 2>/dev/null || echo "  Not found"
 
# Check if CDI (Container Device Interface) is enabled
echo "[*] CDI configuration:"
if [ -f /etc/nvidia-container-runtime/config.toml ]; then
    grep -i "cdi" /etc/nvidia-container-runtime/config.toml
else
    echo "  Config file not found"
fi
 
# Check GPU isolation mode
echo "[*] GPU isolation mode:"
nvidia-smi -q | grep -i "MIG\|isolation\|accounting" || echo "  Could not query GPU"
 
# Check if MIG (Multi-Instance GPU) is available and enabled
echo "[*] MIG (Multi-Instance GPU) status:"
nvidia-smi mig -lgi 2>/dev/null || echo "  MIG not available on this GPU"
 
# Check for IOMMU (required for DMA protection)
echo "[*] IOMMU status:"
if dmesg | grep -qi "IOMMU enabled"; then
    echo "  IOMMU is enabled (good)"
else
    echo "  WARNING: IOMMU may not be enabled"
fi
 
# List all NVIDIA device files
echo "[*] NVIDIA device files:"
ls -la /dev/nvidia* 2>/dev/null || echo "  No NVIDIA devices found"

import subprocess
import json
from typing import List, Dict
 
class MIGSecurityManager:
    """Manage and audit MIG instances for workload isolation."""
 
    def list_mig_instances(self) -> List[Dict]:
        """List all MIG GPU instances and their configuration."""
        result = subprocess.run(
            ["nvidia-smi", "mig", "-lgi"],
            capture_output=True, text=True,
        )
        # Parse nvidia-smi output
        instances = []
        for line in result.stdout.strip().split("\n"):
            if "MIG" in line or "GPU" in line:
                instances.append({"raw": line.strip()})
        return instances
 
    def verify_memory_isolation(self, instance_id: int) -> Dict:
        """
        Verify that a MIG instance cannot access memory from other instances.
        Returns test results.
        """
        # This test allocates memory on one MIG instance and verifies
        # it cannot be read from another instance
        test_script = f"""
import torch
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "MIG-{instance_id}"
# Allocate and fill a tensor with known pattern
t = torch.full((1024, 1024), 42.0, device="cuda")
print(f"Allocated on MIG instance {instance_id}: {{t.sum().item()}}")
# Free the tensor
del t
torch.cuda.empty_cache()
"""
        result = subprocess.run(
            ["python3", "-c", test_script],
            capture_output=True, text=True, timeout=30,
        )
        return {
            "instance_id": instance_id,
            "success": result.returncode == 0,
            "output": result.stdout.strip(),
            "error": result.stderr.strip() if result.returncode != 0 else None,
        }

#!/bin/bash
# Configure SR-IOV for NVIDIA GPU (A100 example)
# Must be run on the bare-metal host with root privileges
 
# Enable SR-IOV on the GPU PCI device
GPU_PCI_ID="0000:3b:00.0"  # Example — find with lspci | grep NVIDIA
 
echo "=== Configuring GPU SR-IOV ==="
 
# Check if SR-IOV is supported
echo "[*] Checking SR-IOV capability:"
lspci -vvv -s "$GPU_PCI_ID" | grep -i "single root"
 
# Enable virtual functions (e.g., 4 VFs)
echo 4 > "/sys/bus/pci/devices/$GPU_PCI_ID/sriov_numvfs"
 
# Verify VFs were created
echo "[*] Virtual functions:"
lspci | grep -i nvidia | grep "Virtual Function"
 
# Assign VFs to VMs using VFIO
echo "[*] Binding VF to VFIO driver for VM assignment:"
VF_PCI_ID="0000:3b:00.1"  # First VF
echo "vfio-pci" > "/sys/bus/pci/devices/$VF_PCI_ID/driver_override"
echo "$VF_PCI_ID" > /sys/bus/pci/drivers/vfio-pci/bind

from typing import Dict, Any
import hashlib
 
class TEEAttestationVerifier:
    """
    Verify TEE attestation for AI workloads.
    Ensures the inference environment is genuine and unmodified.
    """
 
    def __init__(self, expected_measurements: Dict[str, str]):
        """
        Args:
            expected_measurements: Expected hash values for TEE measurements.
                Keys are measurement names (e.g., 'mrenclave', 'mrsigner'),
                values are expected hex-encoded hash values.
        """
        self.expected = expected_measurements
 
    def verify_sgx_attestation(self, attestation_report: Dict[str, Any]) -> Dict:
        """Verify an Intel SGX attestation report."""
        results = {"verified": True, "checks": []}
 
        # Verify MRENCLAVE (hash of the enclave code and initial data)
        mrenclave = attestation_report.get("mrenclave", "")
        if mrenclave != self.expected.get("mrenclave", ""):
            results["verified"] = False
            results["checks"].append({
                "check": "mrenclave",
                "status": "FAIL",
                "detail": "Enclave measurement does not match expected value. "
                         "The inference code may have been modified.",
            })
        else:
            results["checks"].append({
                "check": "mrenclave",
                "status": "PASS",
            })
 
        # Verify MRSIGNER (hash of the enclave signing key)
        mrsigner = attestation_report.get("mrsigner", "")
        if mrsigner != self.expected.get("mrsigner", ""):
            results["verified"] = False
            results["checks"].append({
                "check": "mrsigner",
                "status": "FAIL",
                "detail": "Signer identity does not match. "
                         "The enclave was not signed by the expected party.",
            })
        else:
            results["checks"].append({
                "check": "mrsigner",
                "status": "PASS",
            })
 
        # Check for debug mode (should be disabled in production)
        if attestation_report.get("attributes", {}).get("debug", False):
            results["verified"] = False
            results["checks"].append({
                "check": "debug_mode",
                "status": "FAIL",
                "detail": "Enclave is running in debug mode. "
                         "Memory can be inspected by the host.",
            })
 
        return results
 
    def verify_model_integrity_in_tee(
        self, model_bytes: bytes, expected_hash: str
    ) -> bool:
        """Verify model integrity inside the TEE before loading."""
        actual_hash = hashlib.sha256(model_bytes).hexdigest()
        return actual_hash == expected_hash

#!/bin/bash
# Discover the isolation architecture of the target AI infrastructure
 
echo "=== AI Workload Isolation Assessment ==="
 
# Determine if running in a container
echo "[*] Container detection:"
if [ -f /.dockerenv ]; then
    echo "  Running inside Docker container"
elif grep -q "containerd" /proc/1/cgroup 2>/dev/null; then
    echo "  Running inside containerd"
elif [ -f /run/secrets/kubernetes.io ]; then
    echo "  Running inside Kubernetes pod"
else
    echo "  Not detected as container"
fi
 
# Check GPU access method
echo "[*] GPU access method:"
if [ -e /dev/nvidia0 ]; then
    echo "  Direct GPU device access"
    ls -la /dev/nvidia*
fi
 
# Check for MIG
echo "[*] MIG configuration:"
nvidia-smi mig -lgi 2>/dev/null || echo "  MIG not available/enabled"
 
# Check capabilities
echo "[*] Container capabilities:"
cat /proc/self/status | grep -i cap
 
# Check mount namespace
echo "[*] Mount namespace:"
ls -la /proc/self/ns/mnt
 
# Check for shared GPU memory accessible from this context
echo "[*] GPU memory status:"
nvidia-smi --query-gpu=memory.used,memory.free --format=csv 2>/dev/null