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 管任务怎么分配。三层各司其职。三、落地方案TerraformGPU 实例声明式管理# 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 }AnsibleCUDA 环境一键配置# 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: 安装 MinicondaPython 环境隔离 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: - transformers4.44.0 # 版本锁定 - datasets2.21.0 - accelerate0.33.0 - wandb0.17.0 # 实验追踪 - tensorboard2.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(fGPU: {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 函数检测空闲 GPU15分钟利用率 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 队列驱动的方式按需开关机能有效控制成本。
