1. 项目概述构建可复现的机器学习训练流水线在机器学习项目的实际落地过程中我们常常会遇到一个令人头疼的循环几周前还能完美复现的模型今天再跑一遍结果却天差地别。是数据变了环境依赖升级了还是某个关键的随机种子没保存这种“不可复现”的幽灵不仅消耗着团队大量的调试时间更让模型迭代、A/B测试和线上部署充满了不确定性。今天我想和你分享一套我近年来在多个生产项目中验证过的解决方案它核心围绕两个工具展开Weights Biases和dstack。这套组合拳的目标非常明确——打造一个从数据准备到模型训练每一步都清晰、可追踪、可一键复现的自动化流水线。简单来说WandB 负责解决实验本身的“黑盒”问题它能像一位细心的实验记录员自动记下你每一次训练的所有参数、指标和系统状态。而 dstack 则扮演着“自动化工程师”的角色它用声明式的配置文件帮你把数据预处理、模型训练、验证等一系列任务串联起来并确保每次运行都在一个干净、可控的计算环境中进行。两者结合意味着你不仅能知道模型“为什么”表现好还能确保任何人、在任何时间都能用完全相同的“配方”重新“烘焙”出这个模型。这对于需要长期维护、多人协作或接受严格审计的AI项目来说是至关重要的工程实践。2. 核心工具选型与设计思路拆解在深入实操之前我们先来拆解一下为什么是 WandB dstack 这个组合以及它们各自解决了哪些关键痛点。理解这个设计思路能帮助你在未来面对其他工具时做出更合适的选择。2.1 为什么实验追踪必须独立于本地环境很多团队初期的做法是把训练日志、TensorBoard 文件存在本地或者共享服务器上。这带来了几个致命问题首先环境依赖混乱。你的本地 Python 环境、CUDA 版本和同事的稍有不同就可能导致结果不一致。其次信息孤岛。日志散落在各个工程师的机器上项目负责人很难全局查看进度和对比实验。最后难以回溯。三个月后想回顾某个关键实验的详细参数可能连当时的代码版本都找不到了。WandB 的核心价值在于提供了一个“中心化的实验事实源”。它通过一个轻量的客户端库将每次实验的以下要素自动同步到云端超参数与配置不仅仅是batch_size、learning_rate还包括数据路径、模型结构名称等所有影响结果的配置项。指标时序数据训练/验证损失、准确率等并实时绘制成交互式图表。系统资源GPU 利用率、内存消耗、温度等这对于排查性能瓶颈和成本优化至关重要。输出物可以自动保存模型检查点、预测结果样本甚至整个训练脚本的快照。注意将实验追踪云端化并非只是为了方便查看。它本质上是在建立团队的“机器学习资产目录”。当你的实验记录达到成千上万次时强大的查询、对比和分组功能能让你从“历史数据”中挖掘出有价值的模式比如哪些类型的模型结构对当前数据集更敏感。2.2 为什么需要 dstack 这样的工作流编排工具有了好的实验记录下一步就是保证产生这些记录的“过程”本身是可复现的。手动执行一系列脚本python prepare_data.py,python train.py --config xxx的弊端显而易见顺序可能出错环境可能被污染中间产物管理混乱。dstack 采用了一种“基础设施即代码”和“工作流即代码”的思路。你把整个训练流程准备数据、训练、评估定义在一个 YAML 配置文件里。这个文件明确指定了任务间的依赖关系训练任务必须等数据准备好才能开始。计算资源需求这个任务需要 2 块 A100 GPU、64GB 内存。代码与环境使用哪个 Python 版本依赖哪些requirements.txt。输入与输出从哪里读取数据把生成的模型或数据存到哪里即“产物”。当你执行dstack run train时dstack 会解析工作流依赖按正确顺序执行任务。根据配置在云端AWS/GCP/Azure或本地自动创建一台全新的、指定配置的虚拟机或容器实例。将你的代码仓库和指定的数据同步到这个干净的环境中。执行任务并自动将声明的输出产物如训练好的模型文件保存到云存储中。任务结束后自动清理计算资源按需付费节省成本。这个过程的妙处在于任何持有该 YAML 文件和代码权限的人执行同一个命令都能得到完全一致的环境和运行过程彻底消除了“在我机器上好好的”这类问题。2.3 WandB 与 dstack 的协同分工理解了各自角色后它们的协作关系就清晰了dstack 管“舞台”和“流程”它搭建一个标准、干净的舞台计算环境并严格按照剧本工作流定义指挥演员各个任务脚本依次上场、退场并管理道具数据、模型等产物的流转。WandB 管“现场记录”在“演员”表演脚本运行的过程中WandB 实时、详细地记录下每一个动作细节参数、表现、状态并生成可供回放和分析的报告。这样当你需要复现某个实验时你只需要两个入口1在 dstack 中找到那次历史运行记录一键复现整个流程和环境2在 WandB 中查看那次运行的所有内部细节和结果图表。两者通过唯一的运行名称Run Name相互关联形成了完整的可观测性闭环。3. 环境配置与工具集成实操理论讲完了我们动手搭建。假设我们有一个经典的图像分类项目目录结构如下my_ml_project/ ├── data/ ├── src/ │ ├── prepare_data.py │ └── train.py ├── requirements.txt └── dstack/ └── workflows.yaml3.1 WandB 的初始化与集成首先在 WandB 官网注册账号并创建一个新项目比如my-image-classifier。接下来在代码中集成 WandB。安装与基础配置 在你的requirements.txt中确保包含wandb0.15.0 pytorch-lightning2.0.0 # 如果你用PyTorch Lightning torch torchvision在train.py中进行如下集成。我强烈建议将配置参数集中管理而不是散落在代码各处# train.py import os import torch from pytorch_lightning import Trainer from pytorch_lightning.loggers import WandbLogger from pytorch_lightning.callbacks import ModelCheckpoint import wandb def main(): # 1. 解析配置可以从配置文件、命令行等读取 config { learning_rate: 1e-3, batch_size: 32, max_epochs: 10, model_arch: resnet34, data_path: ./data/processed } # 2. 关键步骤获取 dstack 运行名用于关联 WandB 和 dstack # dstack 会为每次运行注入一个 RUN_NAME 环境变量 run_name os.environ.get(RUN_NAME, flocal-run-{os.getpid()}) # 3. 初始化 WandB # 注意这里我们不在代码中硬编码 API Key而是通过环境变量传递 wandb_logger WandbLogger( namerun_name, # 使用 dstack 运行名实现双向关联 projectmy-image-classifier, configconfig, # 将配置字典记录到 WandB save_dir./logs # 日志本地缓存目录可选 ) # 此时WandB 运行已开始。config 中的参数会自动同步到云端界面。 # 4. 可选记录更复杂的配置或自定义指标 wandb.log({custom_metric: 0.5}) # 可以在训练循环中实时记录 # 5. 将 logger 传递给 PyTorch Lightning Trainer checkpoint_callback ModelCheckpoint(dirpathf./checkpoints/{run_name}) trainer Trainer( max_epochsconfig[max_epochs], loggerwandb_logger, callbacks[checkpoint_callback], acceleratorgpu if torch.cuda.is_available() else cpu, devices1 ) # ... 此处是你的数据模块和模型定义 ... # model MyLightningModel(config) # datamodule MyDataModule(config) # 6. 开始训练WandB 会自动记录损失、准确率等指标 # trainer.fit(model, datamodule) # 7. 训练结束后可以记录测试集指标或上传模型 # test_metrics trainer.test(model, datamodule) # wandb.log({test_accuracy: test_metrics[0][test_acc]}) # 可选将训练好的模型作为 WandB Artifact 保存实现版本化管理 # artifact wandb.Artifact(nameftrained-model-{run_name}, typemodel) # artifact.add_file(local_path_to_model) # wandb_logger.experiment.log_artifact(artifact) wandb.finish() # 结束 WandB 运行 if __name__ __main__: main()实操心得将run_name设置为 dstack 的环境变量是打通两个平台的关键。这样在 WandB 的界面上你一眼就能看出哪个实验对应 dstack 的哪次流水线运行。此外建议将所有可调参数都放入config字典并通过wandb_logger初始化时传入而不是在代码中零散地使用wandb.config.xxx这有利于配置的集中管理和版本控制。3.2 dstack 的安装与云账户配置接下来配置 dstack 来自动化执行这个训练脚本。安装 dstack CLIpip install dstack配置云供应商以 AWS 为例在 AWS IAM 控制台创建一个具有 EC2、S3 和 IAM 相关权限的用户生成访问密钥Access Key ID 和 Secret Access Key。在本地终端初始化 dstack 并配置后端# 初始化 dstack 配置 dstack init # 配置使用 AWS 作为后端并设置区域 dstack config --backend aws dstack config --region us-east-1 # 选择离你近的区域执行dstack config命令时它会引导你输入 AWS 的访问密钥。这些信息会被安全地存储在本地。注意事项关于云成本。dstack 默认会使用按需实例On-Demand对于长期运行的训练任务成本可能较高。在workflows.yaml中你可以通过resources字段指定使用 Spot 实例竞价实例价格可能低至按需实例的 70%-90%。但需注意 Spot 实例可能被中断。对于可容错、支持断点续训的任务这是一个极大的成本优化点。3.3 定义 dstack 工作流配置文件这是整个可复现流水线的“蓝图”定义在./dstack/workflows.yaml中。# ./dstack/workflows.yaml workflows: - name: prepare-data provider: python version: 3.9 # 指定 Python 版本确保环境一致性 script: src/prepare_data.py working-dir: . # 工作目录根路径脚本中的相对路径基于此 requirements: requirements.txt artifacts: - path: ./data/processed # 将处理后的数据保存为产物 name: processed-data # 为产物命名便于引用 - name: train-model provider: python version: 3.9 script: src/train.py working-dir: . requirements: requirements.txt depends-on: - prepare-data # 声明依赖确保先执行数据准备 artifacts: - path: ./checkpoints # 保存模型检查点 name: model-checkpoints - path: ./logs # 保存 WandB 本地日志可选 name: training-logs resources: gpu: 1 # 申请 1 块 GPU memory: 32GB # 申请 32GB 内存 shm-size: 8GB # 共享内存大小对某些数据加载器很重要 env: - WANDB_API_KEY${secrets.WANDB_API_KEY} # 从 dstack 密钥管理注入 API Key - RUN_NAME${run.name} # dstack 会自动注入本次运行的名字配置文件关键点解析depends-on这是定义流水线顺序的核心。train-model依赖于prepare-datadstack 会保证数据准备任务成功完成后才会启动训练任务并将prepare-data任务生成的processed-data产物自动挂载到train-model任务环境的对应路径下。artifacts定义了需要持久化保存的输出目录。任务完成后dstack 会自动将这些目录上传到你配置的云存储如 AWS S3中。后续任务或用户可以通过 dstack CLI 下载这些产物。resources明确指定计算资源。这不仅保证了性能更是可复现性的关键——因为不同的 GPU 型号或数量可能影响浮点精度和训练速度。env环境变量设置。这里做了两件重要的事WANDB_API_KEY${secrets.WANDB_API_KEY}安全地将 WandB API Key 传递给任务环境。你需要在 dstack 的 UI 或通过 CLI (dstack secrets add WANDB_API_KEY your-key) 添加这个密钥避免在代码或配置文件中硬编码敏感信息。RUN_NAME${run.name}将 dstack 自动生成的唯一运行名如lively-dragon-1传递给训练脚本从而在 WandB 中使用相同的名称。4. 运行、监控与产物管理全流程配置完成后我们就可以运行这个可复现的流水线了。4.1 启动流水线与实时监控在项目根目录下执行dstack run train-modeldstack CLI 会做以下几件事解析workflows.yaml发现train-model依赖prepare-data。首先在云端启动一个满足prepare-data任务资源要求默认无GPU的实例。将你的整个代码仓库或指定目录同步到该实例。安装requirements.txt中的依赖。执行python src/prepare_data.py。任务完成后将./data/processed目录作为产物上传到云存储并销毁该实例。接着启动一个拥有 1 块 GPU 和 32GB 内存的新实例用于train-model。同步代码并自动将上一步的processed-data产物下载到实例中的./data/processed路径。设置环境变量安装依赖然后执行python src/train.py。训练过程中你的脚本通过 WandB SDK 将指标实时发送到 WandB 云端。训练结束后将./checkpoints和./logs作为新产物保存并销毁计算实例。如何监控dstack 监控在终端运行dstack ps查看运行状态或使用dstack logs run-name查看实时日志。更直观的方式是访问dstack.ai的 Web UI你可以看到所有历史运行、它们的依赖关系图、状态、日志以及每个任务的资源消耗。WandB 监控同时打开 WandB 项目页面 (wandb.ai/your-username/my-image-classifier)你会看到一个以 dstack 运行名如lively-dragon-1命名的新的实验正在实时更新损失曲线、准确率、GPU 利用率等图表。4.2 版本化数据与模型产物流水线的一次运行很好但如何基于同一份处理好的数据进行多次不同的训练实验呢这就需要用到 dstack 的标签功能。假设我们对prepare-data任务的某次运行结果非常满意想固定这份数据用于后续多轮模型调优。为数据运行打标签 首先找到那次prepare-data的运行 ID在 dstack UI 或通过dstack ps查看。然后为其添加标签dstack tags add run-name-of-prepare v1.0-processed-data这相当于给那份处理好的数据产物拍了个快照并命名为v1.0-processed-data。在训练工作流中引用带标签的数据 修改workflows.yaml中的train-model任务让其依赖特定的数据版本而不是每次都重新运行prepare-data。- name: train-model-v2 provider: python script: src/train.py requirements: requirements.txt depends-on: - prepare-data:v1.0-processed-data # 依赖特定标签的产物 artifacts: - path: ./checkpoints name: model-checkpoints resources: gpu: 1 env: - WANDB_API_KEY${secrets.WANDB_API_KEY} - RUN_NAME${run.name}现在每次运行dstack run train-model-v2它都会直接使用被打上v1.0-processed-data标签的那份数据而不会重新运行数据准备脚本。这完美实现了数据版本化是模型可复现性的基石。下载与复用产物 训练完成后生成的模型产物model-checkpoints也保存在 dstack 中。你可以通过 CLI 下载到本地进行分析或部署dstack artifacts download run-name-of-train ./local-output-dir或者你也可以为某个优秀的模型检查点运行打上标签如best-model-acc-95然后在后续的评估或部署流水线中引用它。4.3 在 WandB 中对比与分析实验当你在固定数据 (v1.0-processed-data) 上用不同的超参数通过修改train.py的config或传递不同参数运行了多次train-model-v2后所有实验都记录在 WandB 中。打开 WandB 项目你可以平行坐标图快速查看哪些超参数组合导致了更高的准确率。结果表格根据验证损失、训练时间等指标进行排序和筛选。分组功能如果你在config中记录了数据版本如data_version: v1.0你可以按数据版本分组排除数据差异对实验结果的干扰。关联 dstack通过相同的run_name你可以直接从 WandB 的实验页面跳转到 dstack 中查看这次训练任务的具体日志、资源消耗和完整的运行配置。5. 高级技巧与常见问题排查在实际生产中你会遇到比示例更复杂的情况。下面分享一些进阶技巧和踩坑经验。5.1 处理大型数据集与中间产物我们的例子中数据产物是直接保存在 dstack 管理的存储中。对于 TB 级的大型数据集每次上传下载可能成为瓶颈。优化策略使用外部版本化数据存储对于原始或预处理后的大型数据集建议使用专门的版本化数据存储系统如DVC或LakeFS。在prepare-data任务中不直接输出数据文件而是输出一个包含数据版本哈希的“指针文件”如data_version.yaml。训练任务读取这个指针文件再通过 DVC 拉取对应版本的数据到本地缓存。利用云原生高速存储在workflows.yaml中可以为任务配置高速临时存储如 AWS 上的 NVMe 实例存储。将数据集预先缓存到这种存储上可以极大加速训练时的数据读取。但需注意实例存储是临时的重要产物仍需通过artifacts字段声明以持久化。分层产物管理区分“中间产物”和“最终产物”。例如预处理中的 tokenized 数据可能是中间产物只在单次流水线运行中有用可以不持久化。而训练好的模型是最终产物必须持久化。通过精细设计artifacts路径可以节省存储成本和传输时间。5.2 超参数扫描与并行实验如何利用这套体系进行超参数优化方案一在单个训练脚本内集成 WandB Sweep。 WandB 提供了强大的超参数扫描功能。你可以在train.py中适配 WandB Sweep 的 API然后在 dstack 的工作流中启动一个执行超参数扫描的任务。这个任务会在内部并行启动多个训练子进程。这种方式的好处是一次 dstack 运行管理一次完整的扫描逻辑集中。缺点是单个任务需要的资源更多需要能同时运行多个实验实例且资源管理不够灵活。方案二使用 dstack 矩阵策略推荐。 dstack 支持在workflows.yaml中定义参数矩阵为每个参数组合创建独立的任务运行。这更适合利用云端的弹性进行大规模并行实验。- name: hyperparam-sweep provider: python script: src/train.py requirements: requirements.txt depends-on: - prepare-data:v1.0-processed-data artifacts: - path: ./checkpoints name: model-checkpoints-${batch_size}-${lr} # 产物名包含超参数便于区分 resources: gpu: 1 env: - WANDB_API_KEY${secrets.WANDB_API_KEY} - RUN_NAME${run.name} # 定义超参数矩阵 params: batch_size lr batch_size: [32, 64, 128] lr: [0.001, 0.0005, 0.0001]运行dstack run hyperparam-sweepdstack 会为batch_size和lr的 9 种组合3x3创建 9 个独立的训练任务每个任务都有自己的运行环境和 WandB 记录。你可以在 WandB 的 Sweep 面板或项目页面直观地对比这 9 个实验。5.3 常见问题与排查清单即使配置得当运行时也可能出错。以下是一个快速排查清单问题现象可能原因排查步骤dstack run失败提示配置错误workflows.yaml语法错误或字段不合法使用dstack check命令验证配置文件。检查缩进、冒号后的空格以及字段名拼写。任务在Submitted状态长时间不开始云账户配额不足、所选区域资源售罄、或 Spot 实例请求未满足1. 检查 dstack UI 或云控制台的实例启动日志。2. 尝试切换可用区或使用按需实例。3. 检查云账户的 vCPU/GPU 限额。任务失败日志显示ModuleNotFoundErrorrequirements.txt依赖缺失或版本冲突1. 在本地使用相同的 Python 版本和pip install -r requirements.txt测试。2. 考虑在任务中增加pip install调试步骤或使用 Docker 镜像确保环境一致性。WandB 无法连接报AuthenticationErrorWANDB_API_KEY环境变量未正确设置1. 确认已在 dstack secrets 中添加了正确的密钥。2. 在任务脚本开头添加print(os.environ.get(WANDB_API_KEY))进行调试注意安全仅限调试。3. 检查网络某些计算环境可能需要配置代理。训练任务找不到依赖任务的数据depends-on路径映射错误或产物未成功生成1. 检查上游任务如prepare-data的日志确认其artifacts声明的路径下确实生成了文件。2. 在下游任务中打印当前工作目录并列出文件检查依赖的产物是否被正确下载到了预期路径。WandB 中的run_name与 dstack 不一致RUN_NAME环境变量未成功传递或脚本中使用了默认值1. 在train.py中打印os.environ.get(RUN_NAME)进行确认。2. 确保workflows.yaml中env部分正确设置了RUN_NAME${run.name}。模型训练结果无法复现随机种子未固定、数据加载顺序随机、或使用了非确定性算法1. 在代码开头固定所有随机种子Python, NumPy, PyTorch等。2. 设置 DataLoader 的worker_init_fn和generator。3. 对于 PyTorch设置torch.backends.cudnn.deterministic True和torch.backends.cudnn.benchmark False。5.4 成本控制与优化建议云上训练成本意识很重要。设置预算告警在云服务商控制台设置每月预算和支出告警。善用 Spot 实例对于可中断的训练任务支持从检查点恢复务必在resources中指定使用 Spot 实例例如resources: {gpu: 1, spot: true}。这通常能节省 60-90% 的成本。自动清理资源确保任务结束后dstack 能正确销毁实例。定期检查 dstack UI确认没有“僵尸”实例在运行。产物生命周期管理定期清理 dstack 中旧的、不再需要的运行记录和产物以节省对象存储费用。可以设置自动保留策略或手动清理。