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 + 队列驱动的方式按需开关机，能有效控制成本。