Daft 文档嵌入基准测试在 8 个 GPU 节点上为 10,000 份 PDF 生成 Embedding 的完整流水线与实现解析【免费下载链接】DaftHigh-performance data engine for AI and multimodal workloads. Process images, audio, video, and structured data at any scale项目地址: https://gitcode.com/GitHub_Trending/da/Daft本文基于仓库中benchmarking/ai/document_embedding/目录下的基准测试文档与配套实现讲解 Daft 如何完成PDF 下载 → 文本抽取 → 分块 → GPU 嵌入的分布式文档嵌入流水线从 10,000 份 Digital Corpora PDF 中抽取文本、按 2048 字符切块使用sentence-transformers/all-MiniLM-L6-v2生成 384 维向量并写回 Parquet。读完后你将掌握用daft.clsdaft.method.batch构建 GPU 推理 UDF 的关键手法并理解 Daft 与 Ray Data、Spark 在同一硬件上跑出 1m 54s / 14m 32s / 8m 4s 成绩背后的执行机制差异。基准背景任务规模与测试环境该基准定义在 README.md 中任务要点如下输入Digital Corpora PDF 元数据S3 上的 Parquet 表约 10,000 份 PDF处理提取文本、切分内容块使用sentence-transformers/all-MiniLM-L6-v2生成384 维嵌入向量输出包含 embedding、文本块与元数据的 Parquet 文件集群8 个g6.xlargeGPU worker 节点每节点 1 张 L4 GPU基准日期2024 年 9 月 22 日框架版本Daft 0.6.2、Ray Data 2.49.2、AWS EMR Spark 7.10.0版本钉在 pyproject.toml 中可确认daft0.6.2与ray[default]2.49.2。这是 Daft AI 基准系列的四个负载之一音频转写、文档嵌入、图像分类、视频目标检测系列总览见 benchmarking/ai/README.md。四个负载的设计目标都是压测多模态数据处理中云端 I/O CPU 预处理 GPU 推理的混合调度能力文档嵌入负载恰好同时覆盖了三段。运行环境Ray 集群与依赖集群定义cluster.yaml 用 Ray 集群定义文件cluster config描述了测试环境head 节点g6.xlargeresources: {CPU: 0, GPU: 0}仅作为提交节点worker 节点min_workers: 8、max_workers: 8均为g6.xlarge各带 1 GPU即 README 中所述8 worker nodes的来源setup_commands中预装了与基准对齐的依赖版本ray[default]2.49.2、transformers4.56.2、sentence-transformers5.1.1、langchain0.3.27、pymupdf1.26.4并通过pip install daft --pre安装当日预发布版 Daft。提交方式基准通过 Ray Job Submission 提交到集群脚本 run_ai_benchmark.py 的关键逻辑是submission_id client.submit_job( entrypointDAFT_RUNNERray DAFT_PROGRESS_BAR0 python daft_main.py, runtime_env{ working_dir: f./benchmarking/ai/{benchmark_name}, uv: daft_uv_runtime_env(), }, )注意DAFT_RUNNERray环境变量它让 Daft 以 Ray 作为分布式执行后端。计时口径是先做一轮 warmup 跑预热模型下载与 GPU 缓存再正式跑 2 轮取平均这解释了为何官方成绩是稳定值而非单次抖动值。Daft 实现逐段解析 daft_main.py完整实现见 daft_main.py其流水线由 6 个阶段构成。1. 配置与集群预热EMBED_MODEL_ID sentence-transformers/all-MiniLM-L6-v2 EMBEDDING_DIM 384 NUM_GPU_NODES 8 INPUT_PATH s3://daft-oss-public-datasets/digitalcorpora_metadata MAX_PDF_PAGES 100 CHUNK_SIZE 2048 CHUNK_OVERLAP 200 EMBEDDING_BATCH_SIZE 10 daft.set_runner_ray()参数语义单份 PDF 超过 100 页即跳过控制长尾分块大小 2048 字符、重叠 200嵌入批大小 10。脚本开头还有一段值得注意的预热ray.remote def warmup(): pass ray.get([warmup.remote() for _ in range(64)])先派发 64 个空任务确保所有 worker 节点就绪并分配了资源避免把节点拉起时间计入基准耗时。2. 读取元数据并下载 PDFdf daft.read_parquet(INPUT_PATH) df df.where(daft.col(file_name).endswith(.pdf)) df df.with_column(pdf_bytes, df[uploaded_pdf_path].download())uploaded_pdf_path列上的.download()是 Daft 的分布式文件读取表达式引擎会把每个 S3 路径的下载任务并行调度到集群各 worker 上字节以二进制列物化在数据中而不是让单个节点串行requests下载。这一步正是云端 I/O 密集型部分的主体。3. PDF 文本抽取CPU UDFdef extract_text_from_parsed_pdf(pdf_bytes): try: doc pymupdf.Document(streampdf_bytes, filetypepdf) if len(doc) MAX_PDF_PAGES: return None page_texts [{text: page.get_text(), page_number: page.number} for page in doc] return page_texts except Exception as e: print(fError extracting text from PDF {e}) return None df df.with_column( pages, df[pdf_bytes].apply( extract_text_from_parsed_pdf, return_dtypedaft.DataType.list( daft.DataType.struct({text: daft.DataType.string(), page_number: daft.DataType.int64()}) ), ), )这是一个普通 Python UDF对pdf_bytes列逐行调用pymupdf解析返回liststructtext: string, page_number: int64即每页一条记录。注意异常被捕获后返回None——解析失败的 PDF 不会中断整个作业后续由where(...not_null())过滤掉。4. 展开页面并分块df df.explode(pages) df df.with_columns({page_text: col(pages)[text], page_number: col(pages)[page_number]}) df df.where(daft.col(page_text).not_null()) df df.with_column( chunks, df[page_text].apply( chunk, return_dtypedaft.DataType.list( daft.DataType.struct({text: daft.DataType.string(), chunk_id: daft.DataType.int64()}) ), ), ) df df.explode(chunks) df df.with_columns({chunk: col(chunks)[text], chunk_id: col(chunks)[chunk_id]}) df df.where(daft.col(chunk).not_null())explode把一份 PDF 一页一条记录的行转置成每页一行再用 LangChain 的RecursiveCharacterTextSplitter(chunk_size2048, chunk_overlap200)把每页文本切成带chunk_id的块再次explode展开。分块逻辑与 Ray Data、Spark 版本保持完全一致对照 ray_data_main.py 与 spark.ipynb 中的RecursiveCharacterTextSplitter参数保证了三个引擎跑的是同一份逻辑。5. GPU 嵌入daft.cls 与批量方法这是整个基准的核心也是 Daft 与另外两个引擎写法差异最大的地方daft.cls(max_concurrencyNUM_GPU_NODES, gpus1) class Embedder: def __init__(self): from sentence_transformers import SentenceTransformer device cuda if torch.cuda.is_available() else cpu self.model SentenceTransformer(EMBED_MODEL_ID, devicedevice) self.model.compile() daft.method.batch( return_dtypedaft.DataType.fixed_size_list(daft.DataType.float32(), EMBEDDING_DIM), batch_sizeEMBEDDING_BATCH_SIZE, ) def __call__(self, text_col): if len(text_col) 0: return [] embeddings self.model.encode( text_col.to_pylist(), ) return embeddings df df.with_column(embedding, Embedder()(df[chunk]))从 daft/udf/init.py 中cls装饰器的源码文档可以看出其语义daft.cls把 Python 类变成 Daft 用户自定义类初始化一次、多行复用——模型加载这类昂贵初始化被摊销到整个作业上且类实例是懒加载的只在查询执行时才在各 worker 上调用__init__gpus1每个Embedder实例向调度器声明需要 1 张 GPU支持 0~1 的小数共享1.0以上的小数不支持max_concurrency8同步方法下该值即actor 池进程数也就是同时存活的Embedder实例上限。配合 8 个单 GPU 节点正好做到每个节点驻留一个模型实例、8 路并行推理且每个节点只下载/加载一次模型daft.method.batch(batch_size10)方法以 10 行一个daft.Series批次喂入to_pylist()转成 Python 列表后整批model.encode(...)。批式编码避免了逐行推理的 Python 调用与 GPU kernel 启动开销return_dtypefixed_size_list(float32, 384)向引擎声明输出是定长 384 维 float32 向量引擎可据此直接构建 Arrow 列、零拷贝写 Parquet。对照 Ray Data 版本 ray_data_main.py 中对应的写法ds ds.map_batches(Embedder, concurrencyNUM_GPU_NODES, num_gpus1.0, batch_sizeEMBEDDING_BATCH_SIZE)两者在 GPU 并发度与批大小上参数完全对齐比较是公平的。差异体现在流水线其余部分见下文三个引擎的结构性差异。6. 写回与 S3 配置df df.select(uploaded_pdf_path, page_number, chunk_id, chunk, embedding) df.write_parquet(OUTPUT_PATH)输出 Schema 为uploaded_pdf_path / page_number / chunk_id / chunk / embedding与 README 所述Parquet with embeddings, text chunks, and metadata一致。此外脚本中有一行容易被忽略的 I/O 配置daft.set_planning_config( default_io_configdaft.io.IOConfig(s3daft.io.S3Config.from_env().replace(requester_paysTrue)) )输入桶启用了 S3requester_pays计费模式若不显式开启部分请求会被拒绝——这是使用公共数据集桶时的必要配置迁移到自己的 S3 时可删除。三个引擎的结构性差异对比三份实现daft_main.py、ray_data_main.py、spark.ipynb可以在源码层面解释性能差异的来源文件获取方式不同Ray Data 版本先read_parquet(...).take_all()把全部 PDF 路径收拢到驱动端再read_binary_files(file_paths)Spark 版本同样先collect()路径、再binaryFile加载。Daft 版本则直接在 DataFrame 上以.download()表达式分布式下载元数据扫描与文件下载可以融合进同一执行计划。GPU 模型生命周期管理Daft 的daft.cls让模型实例成为 actor 池中的一等公民max_concurrency直接等于实例数Spark 版本用pandas_udf配合模块级_model_cache字典做每 executor 缓存一个模型属于手写等价物且依赖TORCH_HOME等缓存路径环境变量见 spark.ipynb 中get_model()Ray Data 依赖map_batches的 actor pool 行为。类型与列的声明Daft 版本为每步 UDF 显式声明return_dtypeliststruct、fixed_size_listfloat32, 384引擎可在计划构建期确定 Schema 并按 Arrow 布局排布数据Spark 版本则显式写ArrayType/StructTypeUDF SchemaRay Data 版本基本靠自动推断。官方基准结果README.md 记录的 2024-09-22 基准成绩如下EngineRuntimeDaft1m 54sRay Data14m 32sSpark8m 4s适用前提需要说清楚该结果绑定特定硬件8 个g6.xlargeL4 GPU、特定数据约 10,000 份 Digital Corpora PDF与特定版本Daft 0.6.2 / Ray 2.49.2 / EMR Spark 7.10.0是相对参考而非通用性能承诺。复现入口是 run_ai_benchmark.py用法为python run_ai_benchmark.py document_embedding它会自动完成 warmup 与两轮取平均。小结文档嵌入基准完整展示了 Daft 处理多模态 AI 推理混合负载的推荐姿势用daft.read_parquet 列级.download()表达式完成分布式文件获取用普通 UDFapply 显式return_dtype承载 CPU 侧的解析与分块explode完成一对多展开用daft.cls(max_concurrencyN, gpus1)daft.method.batch把 GPU 模型包装成带资源声明的常驻 actor 池模型只加载一次、按批推理、以fixed_size_list声明定长向量列全程单条声明式 DataFrame 链完成最终write_parquet直出 S3。同一逻辑在 Ray Data 与 Spark 上分别需要更多样板代码路径收拢、模型缓存字典、UDF Schema 声明而 Daft 版本约 120 行即可表达完整流水线——这一实现模式同样适用于仓库中 image_classification 与 video_object_detection 等其它 AI 基准。【免费下载链接】DaftHigh-performance data engine for AI and multimodal workloads. Process images, audio, video, and structured data at any scale项目地址: https://gitcode.com/GitHub_Trending/da/Daft创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Daft 文档嵌入基准测试：在 8 个 GPU 节点上为 10,000 份 PDF 生成 Embedding 的完整流水线与实现解析
Daft 文档嵌入基准测试在 8 个 GPU 节点上为 10,000 份 PDF 生成 Embedding 的完整流水线与实现解析【免费下载链接】DaftHigh-performance data engine for AI and multimodal workloads. Process images, audio, video, and structured data at any scale项目地址: https://gitcode.com/GitHub_Trending/da/Daft本文基于仓库中benchmarking/ai/document_embedding/目录下的基准测试文档与配套实现讲解 Daft 如何完成PDF 下载 → 文本抽取 → 分块 → GPU 嵌入的分布式文档嵌入流水线从 10,000 份 Digital Corpora PDF 中抽取文本、按 2048 字符切块使用sentence-transformers/all-MiniLM-L6-v2生成 384 维向量并写回 Parquet。读完后你将掌握用daft.clsdaft.method.batch构建 GPU 推理 UDF 的关键手法并理解 Daft 与 Ray Data、Spark 在同一硬件上跑出 1m 54s / 14m 32s / 8m 4s 成绩背后的执行机制差异。基准背景任务规模与测试环境该基准定义在 README.md 中任务要点如下输入Digital Corpora PDF 元数据S3 上的 Parquet 表约 10,000 份 PDF处理提取文本、切分内容块使用sentence-transformers/all-MiniLM-L6-v2生成384 维嵌入向量输出包含 embedding、文本块与元数据的 Parquet 文件集群8 个g6.xlargeGPU worker 节点每节点 1 张 L4 GPU基准日期2024 年 9 月 22 日框架版本Daft 0.6.2、Ray Data 2.49.2、AWS EMR Spark 7.10.0版本钉在 pyproject.toml 中可确认daft0.6.2与ray[default]2.49.2。这是 Daft AI 基准系列的四个负载之一音频转写、文档嵌入、图像分类、视频目标检测系列总览见 benchmarking/ai/README.md。四个负载的设计目标都是压测多模态数据处理中云端 I/O CPU 预处理 GPU 推理的混合调度能力文档嵌入负载恰好同时覆盖了三段。运行环境Ray 集群与依赖集群定义cluster.yaml 用 Ray 集群定义文件cluster config描述了测试环境head 节点g6.xlargeresources: {CPU: 0, GPU: 0}仅作为提交节点worker 节点min_workers: 8、max_workers: 8均为g6.xlarge各带 1 GPU即 README 中所述8 worker nodes的来源setup_commands中预装了与基准对齐的依赖版本ray[default]2.49.2、transformers4.56.2、sentence-transformers5.1.1、langchain0.3.27、pymupdf1.26.4并通过pip install daft --pre安装当日预发布版 Daft。提交方式基准通过 Ray Job Submission 提交到集群脚本 run_ai_benchmark.py 的关键逻辑是submission_id client.submit_job( entrypointDAFT_RUNNERray DAFT_PROGRESS_BAR0 python daft_main.py, runtime_env{ working_dir: f./benchmarking/ai/{benchmark_name}, uv: daft_uv_runtime_env(), }, )注意DAFT_RUNNERray环境变量它让 Daft 以 Ray 作为分布式执行后端。计时口径是先做一轮 warmup 跑预热模型下载与 GPU 缓存再正式跑 2 轮取平均这解释了为何官方成绩是稳定值而非单次抖动值。Daft 实现逐段解析 daft_main.py完整实现见 daft_main.py其流水线由 6 个阶段构成。1. 配置与集群预热EMBED_MODEL_ID sentence-transformers/all-MiniLM-L6-v2 EMBEDDING_DIM 384 NUM_GPU_NODES 8 INPUT_PATH s3://daft-oss-public-datasets/digitalcorpora_metadata MAX_PDF_PAGES 100 CHUNK_SIZE 2048 CHUNK_OVERLAP 200 EMBEDDING_BATCH_SIZE 10 daft.set_runner_ray()参数语义单份 PDF 超过 100 页即跳过控制长尾分块大小 2048 字符、重叠 200嵌入批大小 10。脚本开头还有一段值得注意的预热ray.remote def warmup(): pass ray.get([warmup.remote() for _ in range(64)])先派发 64 个空任务确保所有 worker 节点就绪并分配了资源避免把节点拉起时间计入基准耗时。2. 读取元数据并下载 PDFdf daft.read_parquet(INPUT_PATH) df df.where(daft.col(file_name).endswith(.pdf)) df df.with_column(pdf_bytes, df[uploaded_pdf_path].download())uploaded_pdf_path列上的.download()是 Daft 的分布式文件读取表达式引擎会把每个 S3 路径的下载任务并行调度到集群各 worker 上字节以二进制列物化在数据中而不是让单个节点串行requests下载。这一步正是云端 I/O 密集型部分的主体。3. PDF 文本抽取CPU UDFdef extract_text_from_parsed_pdf(pdf_bytes): try: doc pymupdf.Document(streampdf_bytes, filetypepdf) if len(doc) MAX_PDF_PAGES: return None page_texts [{text: page.get_text(), page_number: page.number} for page in doc] return page_texts except Exception as e: print(fError extracting text from PDF {e}) return None df df.with_column( pages, df[pdf_bytes].apply( extract_text_from_parsed_pdf, return_dtypedaft.DataType.list( daft.DataType.struct({text: daft.DataType.string(), page_number: daft.DataType.int64()}) ), ), )这是一个普通 Python UDF对pdf_bytes列逐行调用pymupdf解析返回liststructtext: string, page_number: int64即每页一条记录。注意异常被捕获后返回None——解析失败的 PDF 不会中断整个作业后续由where(...not_null())过滤掉。4. 展开页面并分块df df.explode(pages) df df.with_columns({page_text: col(pages)[text], page_number: col(pages)[page_number]}) df df.where(daft.col(page_text).not_null()) df df.with_column( chunks, df[page_text].apply( chunk, return_dtypedaft.DataType.list( daft.DataType.struct({text: daft.DataType.string(), chunk_id: daft.DataType.int64()}) ), ), ) df df.explode(chunks) df df.with_columns({chunk: col(chunks)[text], chunk_id: col(chunks)[chunk_id]}) df df.where(daft.col(chunk).not_null())explode把一份 PDF 一页一条记录的行转置成每页一行再用 LangChain 的RecursiveCharacterTextSplitter(chunk_size2048, chunk_overlap200)把每页文本切成带chunk_id的块再次explode展开。分块逻辑与 Ray Data、Spark 版本保持完全一致对照 ray_data_main.py 与 spark.ipynb 中的RecursiveCharacterTextSplitter参数保证了三个引擎跑的是同一份逻辑。5. GPU 嵌入daft.cls 与批量方法这是整个基准的核心也是 Daft 与另外两个引擎写法差异最大的地方daft.cls(max_concurrencyNUM_GPU_NODES, gpus1) class Embedder: def __init__(self): from sentence_transformers import SentenceTransformer device cuda if torch.cuda.is_available() else cpu self.model SentenceTransformer(EMBED_MODEL_ID, devicedevice) self.model.compile() daft.method.batch( return_dtypedaft.DataType.fixed_size_list(daft.DataType.float32(), EMBEDDING_DIM), batch_sizeEMBEDDING_BATCH_SIZE, ) def __call__(self, text_col): if len(text_col) 0: return [] embeddings self.model.encode( text_col.to_pylist(), ) return embeddings df df.with_column(embedding, Embedder()(df[chunk]))从 daft/udf/init.py 中cls装饰器的源码文档可以看出其语义daft.cls把 Python 类变成 Daft 用户自定义类初始化一次、多行复用——模型加载这类昂贵初始化被摊销到整个作业上且类实例是懒加载的只在查询执行时才在各 worker 上调用__init__gpus1每个Embedder实例向调度器声明需要 1 张 GPU支持 0~1 的小数共享1.0以上的小数不支持max_concurrency8同步方法下该值即actor 池进程数也就是同时存活的Embedder实例上限。配合 8 个单 GPU 节点正好做到每个节点驻留一个模型实例、8 路并行推理且每个节点只下载/加载一次模型daft.method.batch(batch_size10)方法以 10 行一个daft.Series批次喂入to_pylist()转成 Python 列表后整批model.encode(...)。批式编码避免了逐行推理的 Python 调用与 GPU kernel 启动开销return_dtypefixed_size_list(float32, 384)向引擎声明输出是定长 384 维 float32 向量引擎可据此直接构建 Arrow 列、零拷贝写 Parquet。对照 Ray Data 版本 ray_data_main.py 中对应的写法ds ds.map_batches(Embedder, concurrencyNUM_GPU_NODES, num_gpus1.0, batch_sizeEMBEDDING_BATCH_SIZE)两者在 GPU 并发度与批大小上参数完全对齐比较是公平的。差异体现在流水线其余部分见下文三个引擎的结构性差异。6. 写回与 S3 配置df df.select(uploaded_pdf_path, page_number, chunk_id, chunk, embedding) df.write_parquet(OUTPUT_PATH)输出 Schema 为uploaded_pdf_path / page_number / chunk_id / chunk / embedding与 README 所述Parquet with embeddings, text chunks, and metadata一致。此外脚本中有一行容易被忽略的 I/O 配置daft.set_planning_config( default_io_configdaft.io.IOConfig(s3daft.io.S3Config.from_env().replace(requester_paysTrue)) )输入桶启用了 S3requester_pays计费模式若不显式开启部分请求会被拒绝——这是使用公共数据集桶时的必要配置迁移到自己的 S3 时可删除。三个引擎的结构性差异对比三份实现daft_main.py、ray_data_main.py、spark.ipynb可以在源码层面解释性能差异的来源文件获取方式不同Ray Data 版本先read_parquet(...).take_all()把全部 PDF 路径收拢到驱动端再read_binary_files(file_paths)Spark 版本同样先collect()路径、再binaryFile加载。Daft 版本则直接在 DataFrame 上以.download()表达式分布式下载元数据扫描与文件下载可以融合进同一执行计划。GPU 模型生命周期管理Daft 的daft.cls让模型实例成为 actor 池中的一等公民max_concurrency直接等于实例数Spark 版本用pandas_udf配合模块级_model_cache字典做每 executor 缓存一个模型属于手写等价物且依赖TORCH_HOME等缓存路径环境变量见 spark.ipynb 中get_model()Ray Data 依赖map_batches的 actor pool 行为。类型与列的声明Daft 版本为每步 UDF 显式声明return_dtypeliststruct、fixed_size_listfloat32, 384引擎可在计划构建期确定 Schema 并按 Arrow 布局排布数据Spark 版本则显式写ArrayType/StructTypeUDF SchemaRay Data 版本基本靠自动推断。官方基准结果README.md 记录的 2024-09-22 基准成绩如下EngineRuntimeDaft1m 54sRay Data14m 32sSpark8m 4s适用前提需要说清楚该结果绑定特定硬件8 个g6.xlargeL4 GPU、特定数据约 10,000 份 Digital Corpora PDF与特定版本Daft 0.6.2 / Ray 2.49.2 / EMR Spark 7.10.0是相对参考而非通用性能承诺。复现入口是 run_ai_benchmark.py用法为python run_ai_benchmark.py document_embedding它会自动完成 warmup 与两轮取平均。小结文档嵌入基准完整展示了 Daft 处理多模态 AI 推理混合负载的推荐姿势用daft.read_parquet 列级.download()表达式完成分布式文件获取用普通 UDFapply 显式return_dtype承载 CPU 侧的解析与分块explode完成一对多展开用daft.cls(max_concurrencyN, gpus1)daft.method.batch把 GPU 模型包装成带资源声明的常驻 actor 池模型只加载一次、按批推理、以fixed_size_list声明定长向量列全程单条声明式 DataFrame 链完成最终write_parquet直出 S3。同一逻辑在 Ray Data 与 Spark 上分别需要更多样板代码路径收拢、模型缓存字典、UDF Schema 声明而 Daft 版本约 120 行即可表达完整流水线——这一实现模式同样适用于仓库中 image_classification 与 video_object_detection 等其它 AI 基准。【免费下载链接】DaftHigh-performance data engine for AI and multimodal workloads. Process images, audio, video, and structured data at any scale项目地址: https://gitcode.com/GitHub_Trending/da/Daft创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考