如何优化处理Elasticsearch中需NLP运算的百万级记录

1. ES数据拉取侧优化 #

• 禁用全量一次性拉取逻辑，改用scroll或search_after做分页批量拉取，单批次数据量控制在1000~5000条，可根据单文档体积灵活调整，避免内存被全量数据占满
• 拉取时通过_source参数仅返回NLP任务需要的字段，过滤无关内容，大幅降低单批次数据体积
• 适当调大ES查询队列配置：将thread_pool.search.queue_size调整为1000~2000，避免批量查询请求被ES直接拒绝
• 拉取任务尽量避开ES高峰读写时段，拉取过程中不要同时执行其他高负载写操作，避免ES节点压力过高宕机

2. NLP处理侧优化 #

• 匹配ES拉取的批大小做NLP批量推理，不要单条过模型，在最大化硬件利用率的同时避免显存/内存溢出
• 限制单进程资源上限：CPU环境下，进程占用核心数不超过物理核心的70%；GPU环境下，通过CUDA_VISIBLE_DEVICES指定使用的显卡，同时开启显存动态分配，避免模型预占全部显存
• 处理完成的文档及时从内存中销毁，不要留存不必要的中间结果，每处理完10个批次主动触发一次垃圾回收
• 对大体积NLP模型做量化/剪枝优化，将FP32精度降低至FP16甚至INT8，可将推理速度提升2~4倍，同时资源占用降低50%以上

3. 任务调度侧优化 #

• 按ES分片、时间等维度将全量任务拆分为多个子任务，单个子任务仅处理10万条以内数据，子任务执行完成后完全释放资源再启动下一个
• 增加失败重试机制：单批次处理失败后自动重试2~3次，重试仍失败则将对应文档ID写入失败日志，全量任务跑完后单独补处理，避免单条数据异常导致整个进程崩溃
• 采用生产者消费者模式解耦拉取和处理逻辑：单独开线程拉取ES数据存入缓冲队列，NLP处理线程从队列取数据，对两边速率做削峰，避免短时间内大量数据涌入占满内存

4. 结果回写侧优化 #

• 处理完成的结果不要单条回写ES，通过bulk接口批量提交，批量大小和拉取批次保持一致
• 回写阶段临时将refresh_interval设置为-1关闭实时刷新，全量回写完成后再恢复原有配置，减少ES段合并压力
• 回写阶段临时将索引副本数设为0，全量完成后再恢复副本，可提升写入速度3倍以上

快速排查定位技巧 #

• 任务运行时通过top/nvidia-smi实时监控内存、显存、CPU占用，确定崩溃发生在拉取、推理还是回写阶段，再针对性调整对应参数
• 先使用1万条以内的小批量数据跑通全流程，测试出最优的批大小、资源配置后再启动全量任务，避免反复崩溃浪费时间

内容的提问来源于stack exchange，提问作者Davinder Singh