Python 基础设施即代码：用 Terraform + Ansible 管 AI 训练环境- 惊觉

Python 基础设施即代码：用 Terraform + Ansible 管 AI 训练环境

一、"我的 GPU 训练任务跑了 3 天，被同事的 Jupyter Notebook 把显存吃光了"

这是一个 AI 团队的真实笑话——也是痛点。团队有 4 块 A100 GPU，但没有任何资源管理和调度机制。大家 SSH 上去手动跑训练脚本，经常一个人占满显存，其他人干等着。更糟的是，每次有新同事入职，需要手动配环境（CUDA、cuDNN、PyTorch 版本），经常出现"在我机器上能跑啊"的问题。

基础设施即代码（IaC）不只是 DevOps 的事，AI 团队的训练环境更需要 IaC——因为环境配置错了，不是服务宕机 5 分钟的问题，而是一个训练任务 3 天的结果作废。

二、AI 训练环境的 IaC 架构

核心思路：Terraform 管"有哪些机器"，Ansible 管"机器上装了什么"，Slurm/K8s 管"任务怎么分配"。三层各司其职。

三、落地方案

Terraform：GPU 实例声明式管理

# main.tf - AI 训练集群基础设施
terraform {
  required_version = ">= 1.5.0"
  backend "s3" {
    bucket = "ai-infra-terraform-state"
    key    = "training-cluster/terraform.tfstate"
    region = "ap-southeast-1"
  }
}

# GPU 实例
resource "aws_instance" "gpu_node" {
  count         = var.gpu_node_count
  ami           = var.gpu_ami
  instance_type = "p4d.24xlarge" # 8x A100 40GB

  vpc_security_group_ids = [
    aws_security_group.gpu_cluster.id
  ]

  # EBS 用于系统盘
  root_block_device {
    volume_size = 200
    volume_type = "gp3"
    encrypted   = true
  }

  # 本地 NVMe SSD 用于临时数据（训练数据集缓存）
  ephemeral_block_device {
    device_name = "/dev/sdb"
  }

  tags = {
    Name        = "gpu-training-node-${count.index}"
    Environment = var.environment
    Team        = "ai-training"
    CostCenter  = "ai-research"
  }

  # 自动加入 Ansible 管理
  provisioner "local-exec" {
    command = <<-EOT
      echo "${self.private_ip}  gpu-node-${count.index}" >> inventory.ini
    EOT
  }
}

# 共享 NFS 存储
resource "aws_efs_file_system" "training_data" {
  creation_token = "ai-training-data"
  encrypted      = true

  lifecycle_policy {
    transition_to_ia = "AFTER_30_DAYS"
  }

  tags = {
    Name = "ai-training-datasets"
  }
}

resource "aws_efs_mount_target" "training_data" {
  count           = length(var.subnet_ids)
  file_system_id  = aws_efs_file_system.training_data.id
  subnet_id       = var.subnet_ids[count.index]
  security_groups = [aws_security_group.efs.id]
}

# 成本控制：非工作时间自动关机
resource "aws_autoscaling_schedule" "night_shutdown" {
  scheduled_action_name  = "night-shutdown"
  autoscaling_group_name = aws_autoscaling_group.gpu_nodes.name
  recurrence            = "0 22 * * *"  # 每晚22点
  min_size              = 0
  desired_capacity      = 0
  max_size              = 0
}

Ansible：CUDA 环境一键配置

# playbook.yml - AI 训练环境标准化部署
- name: 配置 GPU 训练节点
  hosts: gpu_nodes
  become: yes
  vars:
    cuda_version: "12.4"
    cudnn_version: "9.1.0"
    python_version: "3.11"
    pytorch_version: "2.4.0"
    # 关键版本锁定，避免"在我机器上能跑"问题

  tasks:
    - name: 安装 NVIDIA 驱动
      apt:
        name: nvidia-driver-550
        state: present
      register: driver_install

    - name: 重启（驱动安装后需要）
      reboot:
        reboot_timeout: 300
      when: driver_install.changed

    - name: 安装 CUDA Toolkit
      shell: |
        wget https://developer.download.nvidia.com/compute/cuda/{{ cuda_version }}/local_installers/cuda_{{ cuda_version }}_linux.run
        sh cuda_{{ cuda_version }}_linux.run --silent --toolkit
      args:
        creates: /usr/local/cuda-{{ cuda_version }}

    - name: 设置 CUDA 环境变量
      lineinfile:
        path: /etc/profile.d/cuda.sh
        line: "{{ item }}"
        create: yes
      loop:
        - 'export CUDA_HOME=/usr/local/cuda-{{ cuda_version }}'
        - 'export PATH=$CUDA_HOME/bin:$PATH'
        - 'export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH'

    - name: 安装 cuDNN
      unarchive:
        src: /tmp/cudnn-linux-x86_64-{{ cudnn_version }}.tar.xz
        dest: /usr/local/
        remote_src: no
        creates: /usr/local/cuda/include/cudnn.h

    - name: 安装 Miniconda（Python 环境隔离）
      shell: |
        wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh
        bash /tmp/miniconda.sh -b -p /opt/miniconda
      args:
        creates: /opt/miniconda/bin/conda

    - name: 创建 PyTorch 虚拟环境
      shell: |
        source /opt/miniconda/bin/activate
        conda create -n pytorch-{{ pytorch_version }} python={{ python_version }} -y
        conda activate pytorch-{{ pytorch_version }}
        pip install torch=={{ pytorch_version }} torchvision torchaudio \
            --index-url https://download.pytorch.org/whl/cu{{ cuda_version.split('.')[0] }}{{ cuda_version.split('.')[1] }}
      args:
        executable: /bin/bash
        creates: /opt/miniconda/envs/pytorch-{{ pytorch_version }}

    - name: 安装训练依赖
      pip:
        name:
          - transformers==4.44.0     # 版本锁定
          - datasets==2.21.0
          - accelerate==0.33.0
          - wandb==0.17.0            # 实验追踪
          - tensorboard==2.17.0
        virtualenv: /opt/miniconda/envs/pytorch-{{ pytorch_version }}
        state: present  # present 而不是 latest——避免自动升级
    
    - name: 验证 GPU 可用
      shell: |
        source /opt/miniconda/bin/activate pytorch-{{ pytorch_version }}
        python -c "import torch; assert torch.cuda.is_available(), 'GPU不可用!'; print(f'GPU: {torch.cuda.get_device_name(0)}, 显存: {torch.cuda.get_device_properties(0).total_mem/1e9:.0f}GB')"
      register: gpu_check
    - debug:
        msg: "{{ gpu_check.stdout }}"
      when: gpu_check.rc == 0

四、成本控制的 IaC 实践

AI 训练环境最大的隐性成本是"空闲的 GPU"。一块 A100 每小时约 $3，如果 24/7 开机，月成本 $2160。4 块就是 $8640。而实际的训练时间每天可能只有 6-8 小时。

# 成本控制：自动扩缩容
resource "aws_autoscaling_group" "gpu_nodes" {
  min_size = 0          # 允许缩到 0（非工作时间）
  max_size = var.gpu_node_count
  
  tag {
    key   = "AutoShutdown"
    value = "true"
  }
}

# Lambda 函数：检测空闲 GPU（15分钟利用率 < 10% 时自动 shutdown）
# 配合 Slurm 的任务队列：
# - 当有任务排队时：自动启动 GPU 节点
# - 当任务队列为空 + GPU 空闲 > 15 分钟：自动关机

实际节省：GPU 使用时长从 720 小时/月（24/7）降到约 400 小时/月（按需使用），月成本从 $8640 降到约 $4800，节省了 45%。

五、总结

AI 训练环境的 IaC 三件套：Terraform（管资源）、Ansible（管环境）、Slurm/K8s（管调度）。核心价值不是"自动化"（那是最基本的），而是"可复现"——任何一个新节点拉起来，环境完全一致。版本锁定的重要性被低估——CUDA 12.4 和 12.5 的 PyTorch 行为可能有细微差异，训练结果不能复现时排查这个问题能浪费数天。最后，GPU 空闲成本是最大的浪费——用 Auto Scaling + 队列驱动的方式按需开关机，能有效控制成本。

转载自 CSDN-专业IT技术社区

原文链接：https://blog.csdn.net/baronbool/article/details/163219247

Python 基础设施即代码：用 Terraform + Ansible 管 AI 训练环境