使用 SkyPilot Sky Batch 构建跨云分布式批处理从 JSONL 到 GPU 图像生成的完整实战指南【免费下载链接】skypilotThe AI Compute Platform for frontier teams. SkyPilot turns fragmented AI compute into one AI supercomputer, so frontier AI teams build custom intelligence faster.项目地址: https://gitcode.com/GitHub_Trending/sk/skypilotSky Batch 是 SkyPilot 内置的分布式批处理框架它把处理大规模数据集抽象为三个简单步骤——指向云端数据、定义 worker 池、映射一个纯 Python 函数从而让数百万条 prompt 的推理、文本向量化、批量图像生成等任务在几行代码内跑在跨云 GPU/CPU 集群上。读完本文你将掌握 Sky Batch 的完整 APIDataset、JsonReader、remote_function、load()、save_results()、ds.map()、worker 池的配置与生命周期管理以及故障容错与部分结果恢复的底层机制并可直接复用仓库中的可运行示例。为什么选择 Sky Batch在 GPU 和计算集群上处理大型数据集对数百万条 prompt 做推理、从文本生成图像、计算 embedding通常需要大量的基础设施工作准备机器、分发数据、处理失败、收集结果。Sky Batch 将这一切自动化提供四大核心能力简单 Python API处理逻辑就是一个普通 Python 函数Sky Batch 自动把它分发到各 worker 上执行自动数据分发数据集被切分为 batch 并分配给 worker无需手动分片或协调复用运行环境worker 跨作业复用昂贵的初始化安装依赖包、下载模型权重、把模型加载到 GPU只需发生一次多云存储从 Amazon S3s3://或 Google Cloud Storagegs://读取输入并写回结果。Sky Batch 的工作原理Sky Batch 只包含三个核心步骤步骤操作说明1创建 Dataset指向云端存储中的输入数据2启动 worker 池定义每个 worker 所需的硬件与软件3对数据集 map 一个函数编写处理逻辑交给 Sky Batch 分发从源码结构看这三步分别对应 sky/batch/dataset.py 中的Dataset类、sky/jobs/client/sdk.py 的pool_apply池管理接口以及 sky/batch/coordinator.py 中的BatchCoordinator编排器。整体架构如下图所示输入数据被切分后分发到一组常驻warmworker 上worker 池可以横跨不同云厂商和 GPU 型号并且池的生命周期长于单个作业——下一次ds.map()直接跳过环境初始化。快速上手一个可运行的端到端示例本示例读取一个 JSONL 文本文件把每条字符串加倍后写回云存储。仓库中完整的可运行版本位于 examples/batch/simple对应脚本 double_text.py 与 run.sh。第一步准备输入数据创建 JSONL 文件每行一个 JSON 对象并上传到云存储$ cat /tmp/data.jsonl EOF {text: hello} {text: world} {text: sky} {text: batch} EOF $ aws s3 cp /tmp/data.jsonl s3://my-bucket/data.jsonl第二步创建 worker 池创建一个 YAML 文件描述每个 worker 所需的硬件和软件# pool.yaml pool: workers: 2 resources: cpus: 1 setup: echo Setup complete!启动 worker 池$ sky jobs pool apply pool.yaml --pool my-pool -y第三步处理数据集编写 Python 脚本定义处理函数并在数据集上运行import sky # 1. 指向输入数据 ds sky.batch.Dataset(sky.batch.JsonReader(s3://my-bucket/data.jsonl)) # 2. 定义处理函数 sky.batch.remote_function def double_text(): for batch in sky.batch.load(): results [{output: item[text] * 2} for item in batch] sky.batch.save_results(results) # 3. 运行 ds.map( double_text, pool_namemy-pool, batch_size2, outputsky.batch.JsonWriter(s3://my-bucket/output.jsonl), )运行脚本后会看到一个进度条显示批次完成情况。任务结束后结果已写入云存储$ python my_script.py $ aws s3 cp s3://my-bucket/output.jsonl /tmp/output.jsonl $ cat /tmp/output.jsonl {output: hellohello} {output: worldworld} {output: skysky} {output: batchbatch}创建 Dataset指向云端输入Dataset包装一个输入读取器InputReader告诉 Sky Batch 数据在哪里以及如何读取。内置的JsonReader从云存储读取 JSONL 文件ds sky.batch.Dataset(sky.batch.JsonReader(s3://my-bucket/prompts.jsonl))JSONL 文件应每行一个 JSON 对象且所有对象遵循同一 schema{prompt: Summarize the theory of relativity} {prompt: Write a haiku about mountains}路径约束源码级验证从 sky/batch/io_formats.py 的JsonReader.__post_init__可以看到读取器会严格校验路径必须满足以下条件否则直接抛出ValueError前缀必须是s3://或gs://当前仅支持这两类云存储必须以.jsonl结尾。此外utils.count_jsonl_lines_from_cloud 通过下载文件并执行wc -l来统计总行数S3/GCS API 不暴露行数随后按start_idx/end_idx区间切分批次并以 md5(path) 做本地缓存文件名避免同一数据集被重复下载。配置 worker 池硬件、软件与生命周期worker 池定义了每个 worker 的硬件和软件配置。创建一个 YAML 文件# pool.yaml pool: workers: 3 # 并行 worker 数量 resources: accelerators: L4:1 # 每个 worker 的 GPU 类型和数量 setup: | pip install vllm # 安装依赖核心字段说明字段说明示例pool.workers并行 worker 数量1、10、100resources.accelerators每个 worker 的 GPU 类型和数量L4:1、A100:2、H100:8setup安装依赖的命令pip install vllm、uv pip install vllm启动池的命令$ sky jobs pool apply pool.yaml --pool my-pool -y关键特性是池在作业结束后保持运行后续作业完全跳过 setup 阶段。关于池的管理、扩缩容与自动伸缩可以查阅 examples/batch 下的相关配置与 README。从 dataset.py 的实现可以看到ds.map()在执行前还会校验池是否存在、以及池中的 replica 信息版本是否满足MIN_BATCH_REPLICA_INFO_VERSION——用旧版 SkyPilot 创建的池不支持 Sky Batch需要sky jobs pool down pool -y后重新 apply。编写 mapper 函数处理逻辑与执行生命周期mapper 函数包含你的处理逻辑它运行在池中的每个 worker 上。了解什么在 worker 上何时运行是写出高效处理代码的关键sky.batch.remote_function def my_mapper(): # 可选一次性初始化加载模型、打开连接等。 # 每个 worker 只运行一次并在所有批次间复用。 for batch in sky.batch.load(): # batch 是来自输入的 dict 列表例如 # [{prompt: hello}, {prompt: world}] results [process(item) for item in batch] # 保存结果每个输入项一个 dict顺序一致 sky.batch.save_results(results)三个关键 APIsky.batch.remote_function将函数标记为在 worker 上远程执行。装饰器在 sky/batch/remote.py 中实现注册阶段会通过_is_remote_function标记并立即做静态校验。sky.batch.load()一个生成器逐批产出输入数据每批是一个 dict 列表。sky.batch.save_results(results)保存本批次处理结果每个批次调用一次结果顺序必须与输入一致。mapper 函数的硬性约束由于 mapper 函数运行在远程 worker 的全新 Python 环境中通过 utils.serialize_function 把函数源码 base64 序列化后传输采用源码而非字节码序列化因此跨 Python 版本兼容所有 import 必须放在函数体内部sky.batch.remote_function def my_mapper(): import torch # 在函数内 import from transformers import pipeline model pipeline(sentiment-analysis) for batch in sky.batch.load(): results model([item[text] for item in batch]) sky.batch.save_results(results)从 remote.py 的源码可以看到remote_function在装饰阶段就通过 AST 静态分析做三重校验不满足会立即抛出ValueError禁止闭包不能捕获外层作用域的变量co_freevars非空即报错禁止模块级全局引用函数引用的名字必须是内置名、局部定义名或skyworker 命名空间中显式可用等白名单中的名字import 必须在函数内模块级 import 会被 AST 识别为全局引用而拒绝。另外在 worker.py 中load()在每次yield后会校验save_results()是否被调用漏调会触发RuntimeErrorsave_results 还会校验结果长度必须与批次长度一致保证一对一的映射关系。写入结果列、输出格式与多输出使用sky.batch.save_results()保存结果。每次调用为每个输入项写入一行结果dict 的键即成为列名类似表格的行与列。例如sky.batch.save_results([ {prompt: hello, output: Hello! How can I help?, score: 0.95}, {prompt: world, output: The world is vast., score: 0.87}, ])这会写入两行promptoutputscorehelloHello! How can I help?0.95worldThe world is vast.0.87内置输出格式Sky Batch 提供内置的输出 writer实现在 sky/batch/io_formats.pyWriter说明JsonWriter(path)以 JSONL 文件写出结果每行一个 JSON 对象JsonWriter(path, columnoutput)只写出指定列到 JSONLImageWriter(path, columnimage)将 PIL Image 逐个保存为 PNG 文件到目录底层细节JsonWriter的路径必须以s3:///gs://开头且以.jsonl结尾ImageWriter的路径必须以/结尾默认提取名为image的列生成的 PNG 文件按全局索引命名00000000.png、00000001.png……索引即输入数据的原始序号。更详细的输入/输出格式自定义指南见 custom_formats 示例。列过滤类似 SQL SELECTcolumn参数选择输出中包含的列用法类似 SQL 的SELECT# 写全部列 sky.batch.JsonWriter(s3://bucket/full.jsonl) # 只写 output 列 sky.batch.JsonWriter(s3://bucket/outputs.jsonl, columnoutput) # 写 prompt 和 score 列 sky.batch.JsonWriter(s3://bucket/meta.jsonl, column[prompt, score])多输出一次运行写多个目的地可以传入一个 writer 列表一次性写入多个目的地。每个 writer 从同一份结果中独立选择自己需要的列ds.map( my_mapper, pool_namemy-pool, batch_size32, output[ sky.batch.ImageWriter(s3://bucket/images/, columnimage), sky.batch.JsonWriter(s3://bucket/meta.jsonl, column[prompt, score]), ], )运行批处理作业ds.map() 参数详解调用ds.map()将你的函数分发到 worker 池ds.map( my_mapper, # 处理函数 pool_namemy-pool, # worker 池名称 batch_size32, # 每批数据条数 outputsky.batch.JsonWriter(s3://bucket/out.jsonl), # 结果写入位置 )参数说明参数说明mapper_fn用sky.batch.remote_function装饰的函数pool_nameworker 池名称用sky jobs pool apply创建batch_size每个 worker 每次处理的条目数output一个输出 writer 或 writer 列表activate_env可选激活环境的 shell 命令如source .venv/bin/activatestream可选是否阻塞并流式显示进度默认True设为False提交后立即返回ds.map()默认阻塞直到所有批次处理完毕并显示进度条进度同时可以通过sky jobs queue查看。执行链路源码级从 dataset.py 看ds.map()会依次执行API 版本门禁检查 → 池存在性与版本校验 → mapper 是否为远程函数、batch_size是否为正 → 输出路径非空校验 →输出路径已存在时交互式确认是否覆盖输入y/yes才继续→ 序列化函数与格式描述 → 构造带batch_coordinator元数据标记的sky.Task以 managed job 形式提交到指定池。任务的stream逻辑在 dataset.py 中实现以 2 秒为间隔轮询 job 状态用 tqdm 渲染批次级进度条。实战案例基于 Stable Diffusion 的 GPU 图像生成本示例使用 Stable Diffusion 从文本 prompt 生成图像并将任务分发到 GPU worker 上。完整可运行版本位于 examples/batch/diffusion。池配置pool.yamlpool: workers: 3 resources: accelerators: L4:1 setup: | uv venv .venv source .venv/bin/activate uv pip install torch torchvision diffusers transformers accelerate safetensors处理脚本import sky sky.batch.remote_function def generate_images(): from diffusers import StableDiffusionPipeline import torch pipe StableDiffusionPipeline.from_pretrained( stable-diffusion-v1-5/stable-diffusion-v1-5, torch_dtypetorch.float16, ) pipe pipe.to(cuda) for batch in sky.batch.load(): prompts [item[prompt] for item in batch] result pipe(prompts) results [] for item, img in zip(batch, result.images): results.append({prompt: item[prompt], image: img}) sky.batch.save_results(results) ds sky.batch.Dataset(sky.batch.JsonReader(s3://my-bucket/prompts.jsonl)) ds.map( generate_images, pool_namediffusion-pool, batch_size3, output[ sky.batch.ImageWriter(s3://my-bucket/images/, columnimage), sky.batch.JsonWriter(s3://my-bucket/manifest.jsonl, column[prompt]), ], activate_envsource .venv/bin/activate, )这个示例演示了 Sky Batch 的三个关键优势Stable Diffusion 只在每个 worker 上加载一次for 循环之前初始化开销被分摊到所有批次上prompt 以每批 3 条分发到 3 个 GPU worker生成图片保存为 PNG同时写一份 manifest 记录 prompt 与文件名的对应关系利用多输出 列过滤。运行与取回结果$ sky jobs pool apply pool.yaml --pool diffusion-pool -y $ python generate_images.py $ aws s3 cp s3://my-bucket/images/ ./images/ --recursive监控与管理作业查看作业状态$ sky jobs queue从 dataset.py 看队列查询会携带batch_total_batches与batch_completed_batches字段用于渲染进度条。取消正在运行的批处理作业$ sky jobs cancel job-id取消批处理作业会停止处理但保留 worker 池运行。池是共享资源可以被其他作业复用。从 coordinator.py 看取消时协调器收到 SIGTERM 后调用cancel()设置取消标志使分发循环提前退出并向各活跃 worker 发送 shutdown 信号、关闭其上的 worker 服务HTTP 进程。用完拆除池$ sky jobs pool down my-pool -y故障容错与部分结果恢复Sky Batch 自动处理失败情况相关机制在 coordinator.py 与 constants.py 中定义worker 失败worker 崩溃或被抢占后其正在处理的批次会被重新分配给其他 worker。协调器会周期性地重新发现池中的 worker见_dispatch_all新扩容的 replica 会被自动纳入分发个别 worker 线程的失败只要其他 worker 能接续处理即可容忍自动重试失败的批次最多重试 3 次MAX_RETRIES 3采用指数退避RETRY_BACKOFF_BASE 2重试计数跨重启持久化防止无限重试部分结果恢复如果作业中途失败已完成的批次结果会保留在云存储中。协调器异常时打印恢复指引你可以手动合并writer sky.batch.JsonWriter(s3://bucket/output.jsonl) writer.reduce_results(job_id42) writer.cleanup(job_id42) # 删除临时批次文件作业 ID 会打印在失败信息中可以直接复制粘贴。可以省略cleanup()调用以保留临时文件用于调试。底层机制批处理状态PENDING / DISPATCHED / COMPLETED / FAILED持久化在数据库中。如果控制器崩溃重启后会执行_resume_from_db()coordinator.py把所有 DISPATCHED在途批次重置为 PENDING 并重新入队同时恢复各批次的重试计数。批次结果以临时文件形式写入云存储命名如batch_00000000-00000031.jsonl存放在输出路径下的.sky_batch_tmp/job_id/目录reduce_results按起始索引排序拼接回最终输出文件保证结果与输入顺序一致。自定义输入与输出格式Sky Batch 支持通过子类化InputReader/OutputWriter并注册到格式注册表来定义自己的格式。仓库中的 examples/batch/custom_formats/process_range.py 展示了完整的自定义示例RangeReader直接从 Pythonrange生成数据无需文件 I/OTextWriter按条目写出.txt文件YamlWriter把各批次结果合并为单个 YAML 文件。自定义 reader 只需实现两个方法框架在 io_formats.py 中以抽象方法约束__len__()返回数据集总条目数download_batch(start_idx, end_idx, cache_dir)下载指定索引区间的数据返回 dict 列表。自定义 writer 需实现三个方法upload_batch()上传单批结果、reduce_results()合并所有批次到最终输出、cleanup()清理临时文件。自定义类定义在 SkyPilot 源码之外时会以源码形式随作业序列化并在 worker 端通过exec重建见to_dict/from_dict因此可以直接在用户脚本中定义新格式。模块结构与进一步探索Sky Batch 的完整实现位于仓库的 sky/batch 目录文件职责sky/batch/init.py公共 API 导出Dataset、JsonReader/Writer、ImageWriter、remote_function、load、save_resultssky/batch/dataset.pyDataset类与ds.map()提交逻辑sky/batch/io_formats.py输入/输出格式基类与内置实现sky/batch/remote.pyremote_function装饰器与 AST 校验sky/batch/worker.pyworker 端常驻服务与load()/save_results()sky/batch/coordinator.py批次编排、分发、HA 恢复与结果合并sky/batch/constants.py端口、超时、重试等常量配置可运行的端到端示例分别位于 examples/batch/simpleCPU 文本加倍、examples/batch/diffusionGPU 图像生成、多输出与 examples/batch/custom_formats自定义格式。执行前请确保本机已安装 SkyPilot并已通过sky check完成云厂商配置worker 端的sky.batch与boto3由框架自动准备。在本地按bash examples/batch/simple/run.sh即可体验完整的建池 → 造数据 → 跑批 → 取结果流程。【免费下载链接】skypilotThe AI Compute Platform for frontier teams. SkyPilot turns fragmented AI compute into one AI supercomputer, so frontier AI teams build custom intelligence faster.项目地址: https://gitcode.com/GitHub_Trending/sk/skypilot创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考