Training Cluster Network Security

from typing import Dict, List
import subprocess
import socket
 
class TrainingNetworkDiscovery:
    """Discover training cluster network configuration."""
 
    def discover_nccl_config(self) -> Dict:
        """Discover NCCL network configuration from environment."""
        import os
 
        nccl_vars = {}
        for key, value in os.environ.items():
            if key.startswith("NCCL_"):
                nccl_vars[key] = value
 
        return {
            "nccl_socket_ifname": nccl_vars.get("NCCL_SOCKET_IFNAME", "not set"),
            "nccl_ib_disable": nccl_vars.get("NCCL_IB_DISABLE", "0"),
            "nccl_net_gdr_level": nccl_vars.get("NCCL_NET_GDR_LEVEL", "not set"),
            "nccl_socket_nthreads": nccl_vars.get("NCCL_SOCKET_NTHREADS", "not set"),
            "nccl_debug": nccl_vars.get("NCCL_DEBUG", "not set"),
            "all_nccl_vars": nccl_vars,
        }
 
    def discover_rdma_interfaces(self) -> List[Dict]:
        """Discover RDMA-capable network interfaces."""
        interfaces = []
        try:
            result = subprocess.run(
                ["ibstat"], capture_output=True, text=True, timeout=10,
            )
            if result.returncode == 0:
                # Parse ibstat output
                current_port = {}
                for line in result.stdout.split("\n"):
                    line = line.strip()
                    if line.startswith("CA"):
                        if current_port:
                            interfaces.append(current_port)
                        current_port = {"device": line}
                    elif ":" in line:
                        key, _, value = line.partition(":")
                        current_port[key.strip()] = value.strip()
                if current_port:
                    interfaces.append(current_port)
        except (subprocess.TimeoutExpired, FileNotFoundError):
            pass
 
        return interfaces
 
    def discover_training_endpoints(self) -> Dict:
        """Discover distributed training endpoints from environment."""
        import os
 
        return {
            "master_addr": os.environ.get("MASTER_ADDR", "not set"),
            "master_port": os.environ.get("MASTER_PORT", "not set"),
            "world_size": os.environ.get("WORLD_SIZE", "not set"),
            "rank": os.environ.get("RANK", "not set"),
            "local_rank": os.environ.get("LOCAL_RANK", "not set"),
        }
 
    def scan_training_ports(
        self, target_host: str, port_range: tuple = (29400, 29510)
    ) -> List[Dict]:
        """Scan for open NCCL/Gloo training ports."""
        open_ports = []
 
        for port in range(port_range[0], port_range[1]):
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(0.5)
            try:
                result = sock.connect_ex((target_host, port))
                if result == 0:
                    open_ports.append({
                        "host": target_host,
                        "port": port,
                        "status": "open",
                        "likely_service": self._identify_service(port),
                    })
            except socket.error:
                pass
            finally:
                sock.close()
 
        return open_ports
 
    @staticmethod
    def _identify_service(port: int) -> str:
        """Identify the likely service based on port number."""
        if 29400 <= port <= 29500:
            return "PyTorch distributed (NCCL/Gloo)"
        elif port == 6379:
            return "Redis (possibly torch.distributed store)"
        elif port == 2379:
            return "etcd (possibly torch.distributed store)"
        return "unknown"

import socket
import struct
from typing import Dict, Optional
 
class NCCLNetworkAssessment:
    """Assess NCCL network security in a training cluster."""
 
    def __init__(self, master_addr: str, master_port: int = 29500):
        self.master_addr = master_addr
        self.master_port = master_port
 
    def test_nccl_accessibility(self) -> Dict:
        """Test if NCCL master port is accessible without authentication."""
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(5)
            result = sock.connect_ex((self.master_addr, self.master_port))
 
            if result == 0:
                # Try to receive initial handshake data
                try:
                    data = sock.recv(1024)
                    return {
                        "accessible": True,
                        "severity": "critical",
                        "finding": "NCCL master port is accessible without authentication",
                        "received_bytes": len(data),
                        "data_preview": data[:64].hex(),
                    }
                except socket.timeout:
                    return {
                        "accessible": True,
                        "severity": "high",
                        "finding": "NCCL master port is open but did not send data",
                    }
            else:
                return {"accessible": False, "severity": "info"}
        except socket.error as e:
            return {"accessible": False, "error": str(e)}
        finally:
            sock.close()
 
    def test_gradient_interception(self) -> Dict:
        """
        Assess the feasibility of gradient interception.
        Checks if training traffic is visible on the network.
        """
        # Check if NCCL is using TCP (interceptable) vs RDMA (harder to intercept)
        import os
        ib_disable = os.environ.get("NCCL_IB_DISABLE", "0")
        net_plugin = os.environ.get("NCCL_NET", "")
 
        if ib_disable == "1" or net_plugin == "Socket":
            return {
                "transport": "TCP/Socket",
                "interceptable": True,
                "severity": "high",
                "finding": "NCCL using TCP transport — gradient traffic is visible to network sniffers",
                "remediation": "Use IPsec or WireGuard to encrypt inter-node traffic",
            }
        else:
            return {
                "transport": "RDMA/IB",
                "interceptable": False,
                "severity": "medium",
                "finding": "NCCL using RDMA — traffic bypasses OS network stack but is not encrypted",
                "remediation": "Use InfiniBand security features or physically isolate the training network",
            }

import torch
import torch.distributed as dist
from typing import Optional, Callable
 
class GradientPoisoningDefense:
    """
    Detect and mitigate gradient poisoning in distributed training.
    Implements gradient integrity checking between training nodes.
    """
 
    def __init__(
        self,
        world_size: int,
        rank: int,
        clip_norm: float = 1.0,
        anomaly_threshold: float = 3.0,
    ):
        self.world_size = world_size
        self.rank = rank
        self.clip_norm = clip_norm
        self.anomaly_threshold = anomaly_threshold
        self._gradient_history: list = []
 
    def validate_gradient(self, gradient: torch.Tensor) -> dict:
        """Validate a gradient for anomalies before aggregation."""
        grad_norm = gradient.norm().item()
        findings = {"valid": True, "checks": []}
 
        # Check 1: Gradient norm within expected range
        if self._gradient_history:
            mean_norm = sum(self._gradient_history) / len(self._gradient_history)
            if grad_norm > mean_norm * self.anomaly_threshold:
                findings["valid"] = False
                findings["checks"].append({
                    "check": "norm_anomaly",
                    "value": grad_norm,
                    "expected_range": f"<{mean_norm * self.anomaly_threshold:.4f}",
                })
 
        # Check 2: NaN/Inf detection
        if torch.isnan(gradient).any() or torch.isinf(gradient).any():
            findings["valid"] = False
            findings["checks"].append({
                "check": "nan_inf_detected",
                "nan_count": torch.isnan(gradient).sum().item(),
                "inf_count": torch.isinf(gradient).sum().item(),
            })
 
        # Check 3: Gradient direction consistency (cosine similarity with recent history)
        # Dramatic direction changes may indicate injection
 
        self._gradient_history.append(grad_norm)
        if len(self._gradient_history) > 100:
            self._gradient_history.pop(0)
 
        return findings
 
    def secure_allreduce(
        self,
        tensor: torch.Tensor,
        op: dist.ReduceOp = dist.ReduceOp.SUM,
    ) -> torch.Tensor:
        """
        AllReduce with gradient validation and norm clipping.
        """
        # Clip gradient norm before sending
        grad_norm = tensor.norm()
        if grad_norm > self.clip_norm:
            tensor = tensor * (self.clip_norm / grad_norm)
 
        # Validate before aggregation
        validation = self.validate_gradient(tensor)
        if not validation["valid"]:
            # Replace with zero tensor if validation fails
            print(f"[Rank {self.rank}] Gradient validation failed: {validation['checks']}")
            tensor = torch.zeros_like(tensor)
 
        # Standard allreduce
        dist.all_reduce(tensor, op=op)
        tensor /= self.world_size
 
        return tensor

#!/bin/bash
# InfiniBand security assessment for training clusters
 
echo "=== InfiniBand Security Assessment ==="
 
# Check InfiniBand devices
echo "[*] InfiniBand devices:"
ibstat 2>/dev/null || echo "  ibstat not available (no IB support?)"
 
# Check subnet manager
echo ""
echo "[*] Subnet Manager:"
sminfo 2>/dev/null || echo "  sminfo not available"
 
# Check partition keys (access control)
echo ""
echo "[*] Partition Keys (P_Keys):"
cat /sys/class/infiniband/*/ports/*/pkeys/* 2>/dev/null | sort -u || echo "  Cannot read P_Keys"
 
# Check for RDMA interfaces
echo ""
echo "[*] RDMA interfaces:"
rdma link show 2>/dev/null || echo "  rdma tool not available"
 
# Check for GPUDirect RDMA
echo ""
echo "[*] GPUDirect RDMA:"
if [ -d /dev/infiniband ]; then
    ls -la /dev/infiniband/
else
    echo "  No InfiniBand devices in /dev"
fi
 
# Check NCCL network configuration
echo ""
echo "[*] NCCL environment:"
env | grep NCCL_ | sort
 
# Check for network namespaces (isolation)
echo ""
echo "[*] Network namespaces:"
ip netns list 2>/dev/null || echo "  No network namespaces configured"
 
# Check firewall rules on training interfaces
echo ""
echo "[*] Firewall rules for training ports:"
iptables -L -n 2>/dev/null | grep -E "2950[0-9]|NCCL" || echo "  No specific rules for training ports"

from typing import Dict, List
 
class RDMASecurityAssessment:
    """Assess RDMA security configuration for training clusters."""
 
    def check_rdma_permissions(self) -> Dict:
        """Check RDMA device permissions and access controls."""
        import os
        from pathlib import Path
 
        findings = []
 
        # Check RDMA device files
        rdma_devices = list(Path("/dev/infiniband").glob("*")) if Path("/dev/infiniband").exists() else []
 
        for device in rdma_devices:
            stat = os.stat(device)
            if stat.st_mode & 0o006:  # World readable/writable
                findings.append({
                    "severity": "high",
                    "device": str(device),
                    "finding": "RDMA device is world-accessible",
                    "permissions": oct(stat.st_mode)[-3:],
                })
 
        # Check for memory registration limits
        try:
            with open("/proc/sys/vm/max_map_count") as f:
                max_map = int(f.read().strip())
                if max_map > 1000000:
                    findings.append({
                        "severity": "medium",
                        "finding": f"High max_map_count ({max_map}) — allows large RDMA memory registrations",
                    })
        except (FileNotFoundError, ValueError):
            pass
 
        return {"findings": findings}
 
    def check_partition_isolation(self) -> Dict:
        """Check if InfiniBand partition keys are configured for isolation."""
        from pathlib import Path
 
        pkey_paths = list(Path("/sys/class/infiniband").glob("*/ports/*/pkeys/*"))
        pkeys = set()
 
        for pkey_path in pkey_paths:
            try:
                pkey = pkey_path.read_text().strip()
                pkeys.add(pkey)
            except (PermissionError, OSError):
                pass
 
        # Default P_Key (0x7fff or 0xffff) means no partitioning
        has_non_default = any(
            k not in ("0x7fff", "0xffff", "0x0000") for k in pkeys
        )
 
        return {
            "partition_keys": list(pkeys),
            "custom_partitions": has_non_default,
            "severity": "high" if not has_non_default else "info",
            "finding": (
                "Only default P_Key found — no InfiniBand partition isolation"
                if not has_non_default
                else "Custom partition keys configured"
            ),
        }

#!/bin/bash
# Configure WireGuard tunnel for encrypted training traffic
 
# Generate keys on each training node
wg genkey | tee /etc/wireguard/privatekey | wg pubkey > /etc/wireguard/publickey
 
# Configuration for training node 0
cat > /etc/wireguard/wg-train.conf << 'WGEOF'
[Interface]
Address = 10.100.0.1/24
ListenPort = 51820
PrivateKey = <node0_private_key>
# Set MTU to accommodate overhead without fragmenting
MTU = 1420
 
# Training node 1
[Peer]
PublicKey = <node1_public_key>
AllowedIPs = 10.100.0.2/32
Endpoint = 10.1.0.2:51820
PersistentKeepalive = 25
 
# Training node 2
[Peer]
PublicKey = <node2_public_key>
AllowedIPs = 10.100.0.3/32
Endpoint = 10.1.0.3:51820
PersistentKeepalive = 25
WGEOF
 
# Bring up the tunnel
wg-quick up wg-train
 
# Configure NCCL to use the WireGuard interface
export NCCL_SOCKET_IFNAME=wg-train
export MASTER_ADDR=10.100.0.1

# NetworkPolicy: Restrict training pod communication
apiVersion: networking.k8s.io/v1
kind: NetworkPolicy
metadata:
  name: training-job-network-policy
  namespace: ml-training
spec:
  podSelector:
    matchLabels:
      job-type: distributed-training
  policyTypes:
  - Ingress
  - Egress
  ingress:
  # Allow NCCL traffic only from other training pods in the same job
  - from:
    - podSelector:
        matchLabels:
          job-type: distributed-training
    ports:
    - protocol: TCP
      port: 29500  # PyTorch distributed master
    - protocol: TCP
      port: 29400  # NCCL
  egress:
  # Allow NCCL traffic to other training pods
  - to:
    - podSelector:
        matchLabels:
          job-type: distributed-training
    ports:
    - protocol: TCP
      port: 29500
    - protocol: TCP
      port: 29400
  # Allow DNS
  - to:
    - namespaceSelector: {}
    ports:
    - protocol: UDP
      port: 53
  # Allow access to storage (model checkpoints, datasets)
  - to:
    - ipBlock:
        cidr: 10.2.0.0/24  # Storage network
    ports:
    - protocol: TCP
      port: 2049  # NFS
    - protocol: TCP
      port: 443   # S3