nlp_structbert_sentence-similarity_chinese-large批量处理任务队列设计与实现

NLP StructBERT 句子相似度模型批量处理任务队列设计与实现最近在做一个文本处理的项目需要批量计算大量句子对之间的相似度。一开始我直接用模型API挨个调用结果发现效率太低而且经常因为网络波动或者资源限制导致任务失败。后来我琢磨着这种大批量的离线处理任务其实特别适合用任务队列来做异步处理。今天就跟大家聊聊怎么设计一个基于任务队列的批量处理系统。我会用Java语言来演示重点讲怎么用Redis或者RabbitMQ做消息队列怎么设计生产者提交任务、消费者处理任务还有怎么处理任务去重、失败重试这些实际问题。这套方案在我们实际项目中跑得挺稳的处理几十万条数据都没问题。1. 为什么需要任务队列处理批量文本相似度先说说我遇到的实际问题。我们有个业务需要计算用户上传的文档和知识库中所有文档的相似度每次上传可能有几百个句子知识库里有几万条数据这样一组合就是几百万对句子需要计算。如果同步处理用户得等好几个小时而且中间一旦出错就得全部重来。用任务队列的好处很明显。首先它能把一个大任务拆成很多小任务然后让多个worker同时处理速度能提升很多倍。其次任务队列有持久化机制就算处理过程中程序重启了任务也不会丢失。最后通过队列我们可以控制处理速度避免一下子把模型服务打垮。我对比了几种方案。最简单的就是同步循环调用但这种方式扩展性差出错难处理。用线程池稍微好点但任务状态管理比较麻烦。最后选择了消息队列因为它专门就是为解决这类问题设计的。2. 系统整体架构设计整个系统的架构其实挺清晰的主要分三层任务提交层、队列管理层、任务处理层。任务提交层负责接收用户提交的批量任务把每个句子对拆分成独立的小任务。比如用户上传了100个句子知识库有1000个句子那就生成10万个任务项。这一层还要做初步的校验比如句子不能为空、不能太长之类的。队列管理层是核心我用Redis或者RabbitMQ来实现。Redis简单轻量如果对消息的可靠性要求不是特别高用它就够用了。RabbitMQ功能更全有完善的消息确认机制适合对可靠性要求高的场景。这一层要负责任务的存储、分发还要处理任务状态。任务处理层就是worker它们从队列里取任务调用NLP StructBERT模型的API计算相似度然后把结果存到数据库。可以启动多个worker同时工作提高处理速度。数据存储方面我用MySQL存最终的结果用Redis存任务状态和临时数据。这样读写分离性能会好很多。3. 使用Redis实现消息队列Redis实现消息队列特别简单基本上用它的List数据结构就够了。我给大家看看具体的代码实现。3.1 任务数据结构设计首先得设计好任务的数据结构。每个任务需要包含哪些信息呢我总结了这么几个关键字段public class SimilarityTask { private String taskId; // 任务唯一ID private String sentence1; // 第一个句子 private String sentence2; // 第二个句子 private String sourceId; // 来源标识用于去重 private int priority; // 优先级0普通1高优先级 private int retryCount 0; // 重试次数 private String status; // 任务状态 private Date createTime; // 创建时间 private Date processTime; // 处理时间 // 省略getter/setter和构造方法 }任务ID我用UUID生成确保全局唯一。sourceId是用来做任务去重的后面会详细讲。priority字段可以让紧急任务优先处理。3.2 生产者实现任务提交与入队生产者端的代码主要负责接收批量任务然后把它们放到Redis队列里。这里有个技巧如果任务量特别大不要一条一条地往队列里放可以批量操作。Component public class TaskProducer { Autowired private RedisTemplateString, String redisTemplate; Autowired private TaskRepository taskRepository; /** * 提交批量相似度计算任务 * param sentences1 第一组句子列表 * param sentences2 第二组句子列表 * return 批次ID */ public String submitBatchTask(ListString sentences1, ListString sentences2) { String batchId UUID.randomUUID().toString(); // 生成所有句子对组合 ListSimilarityTask tasks new ArrayList(); for (int i 0; i sentences1.size(); i) { for (int j 0; j sentences2.size(); j) { SimilarityTask task new SimilarityTask(); task.setTaskId(UUID.randomUUID().toString()); task.setSentence1(sentences1.get(i)); task.setSentence2(sentences2.get(j)); task.setSourceId(generateSourceId(sentences1.get(i), sentences2.get(j))); task.setStatus(PENDING); task.setCreateTime(new Date()); tasks.add(task); } } // 批量保存到数据库 taskRepository.saveAll(tasks); // 批量入队 ListString taskIds tasks.stream() .map(SimilarityTask::getTaskId) .collect(Collectors.toList()); // 使用Redis的LPUSH批量入队 String queueKey similarity:task:queue; redisTemplate.opsForList().leftPushAll(queueKey, taskIds.toArray(new String[0])); // 记录批次信息 String batchKey similarity:batch: batchId; redisTemplate.opsForHash().put(batchKey, total, String.valueOf(tasks.size())); redisTemplate.opsForHash().put(batchKey, pending, String.valueOf(tasks.size())); redisTemplate.expire(batchKey, 7, TimeUnit.DAYS); return batchId; } private String generateSourceId(String s1, String s2) { // 生成去重标识相同句子对应该有相同的sourceId String combined s1 ||| s2; return DigestUtils.md5DigestAsHex(combined.getBytes()); } }这里我用了两个Redis数据结构List存任务队列Hash存批次信息。批次信息里记录总任务数和待处理数方便后面查询进度。3.3 消费者实现任务处理与结果存储消费者端从队列里取任务调用模型API然后保存结果。这里要注意错误处理和重试机制。Component public class TaskConsumer { Autowired private RedisTemplateString, String redisTemplate; Autowired private TaskRepository taskRepository; Autowired private ResultRepository resultRepository; Autowired private ModelService modelService; private volatile boolean running true; /** * 启动消费者 */ public void startConsumer() { new Thread(() - { while (running) { try { // 从队列右侧获取任务FIFO String queueKey similarity:task:queue; String taskId redisTemplate.opsForList().rightPop(queueKey, 10, TimeUnit.SECONDS); if (taskId ! null) { processTask(taskId); } } catch (Exception e) { // 记录错误但不要退出循环 log.error(Consumer error: , e); try { Thread.sleep(5000); // 出错后等待5秒再继续 } catch (InterruptedException ie) { Thread.currentThread().interrupt(); } } } }).start(); } /** * 处理单个任务 */ private void processTask(String taskId) { OptionalSimilarityTask taskOpt taskRepository.findById(taskId); if (!taskOpt.isPresent()) { log.warn(Task not found: {}, taskId); return; } SimilarityTask task taskOpt.get(); task.setStatus(PROCESSING); task.setProcessTime(new Date()); taskRepository.save(task); try { // 调用模型API计算相似度 float similarity modelService.calculateSimilarity( task.getSentence1(), task.getSentence2() ); // 保存结果 SimilarityResult result new SimilarityResult(); result.setTaskId(taskId); result.setSentence1(task.getSentence1()); result.setSentence2(task.getSentence2()); result.setSimilarity(similarity); result.setCreateTime(new Date()); resultRepository.save(result); // 更新任务状态 task.setStatus(SUCCESS); taskRepository.save(task); // 更新批次进度 updateBatchProgress(taskId, true); } catch (Exception e) { log.error(Process task failed: {}, taskId, e); handleFailedTask(task, e); } } /** * 处理失败任务 */ private void handleFailedTask(SimilarityTask task, Exception e) { task.setRetryCount(task.getRetryCount() 1); if (task.getRetryCount() 3) { // 重试超过3次标记为失败 task.setStatus(FAILED); taskRepository.save(task); updateBatchProgress(task.getTaskId(), false); } else { // 重新入队等待重试 task.setStatus(PENDING); taskRepository.save(task); // 延迟重试重试次数越多延迟越长 int delaySeconds task.getRetryCount() * 30; String delayedQueueKey similarity:task:delayed: delaySeconds; redisTemplate.opsForList().leftPush(delayedQueueKey, task.getTaskId()); // 设置延迟队列的过期时间 redisTemplate.expire(delayedQueueKey, delaySeconds 60, TimeUnit.SECONDS); } } /** * 更新批次处理进度 */ private void updateBatchProgress(String taskId, boolean success) { // 这里需要根据业务逻辑更新批次进度 // 实际项目中可能需要查询任务所属的批次然后更新Redis中的计数 } }消费者这里有几个关键点。一是用了rightPop方法这样保证先进入队列的任务先被处理。二是加了超时时间避免一直阻塞。三是错误处理失败的任务会根据重试次数决定是重新入队还是标记为失败。4. 高级特性实现基本的队列功能实现后还需要一些高级特性来保证系统的可靠性。4.1 任务去重机制在实际应用中经常会有重复的任务提交。比如用户不小心点了两次提交或者系统重试时产生了重复任务。如果每个都处理既浪费资源又可能产生重复数据。我实现的去重机制是这样的在生成任务时根据两个句子的内容计算一个唯一的sourceId。在任务入队前先检查这个sourceId是否已经处理过或者正在处理中。Component public class DeduplicationService { Autowired private RedisTemplateString, String redisTemplate; /** * 检查任务是否重复 */ public boolean isDuplicate(String sourceId) { String key similarity:dedup: sourceId; // 使用Redis的SETNX命令如果key不存在则设置并返回true Boolean success redisTemplate.opsForValue().setIfAbsent(key, 1, 24, TimeUnit.HOURS); // 如果设置失败说明这个sourceId在24小时内已经存在 return !Boolean.TRUE.equals(success); } /** * 批量检查去重 */ public ListString filterDuplicates(ListString sourceIds) { ListString uniqueIds new ArrayList(); for (String sourceId : sourceIds) { if (!isDuplicate(sourceId)) { uniqueIds.add(sourceId); } } return uniqueIds; } }这里我设置了24小时的过期时间因为我们的业务场景下24小时内相同的句子对计算一次就够了。如果业务需要永久去重可以把结果存到数据库然后每次查询数据库。4.2 失败重试与死信队列不是所有失败的任务都应该无限重试。有些错误是暂时性的比如网络波动重试可能成功。有些错误是永久性的比如句子太长超过模型限制重试多少次都没用。我设计了三级重试机制立即重试对于一些可预见的错误比如数据库连接超时立即重试一次延迟重试对于网络超时等暂时性错误延迟一段时间后重试死信队列重试超过3次仍然失败进入死信队列人工处理Component public class RetryManager { Autowired private RedisTemplateString, String redisTemplate; /** * 处理失败任务的重试逻辑 */ public void handleRetry(SimilarityTask task, Exception error) { if (isTransientError(error)) { // 暂时性错误延迟重试 scheduleDelayedRetry(task); } else if (isRecoverableError(error)) { // 可恢复错误立即重试 retryImmediately(task); } else { // 不可恢复错误进入死信队列 moveToDeadLetterQueue(task, error); } } /** * 延迟重试 */ private void scheduleDelayedRetry(SimilarityTask task) { int delayMinutes calculateDelayMinutes(task.getRetryCount()); String delayedQueueKey String.format(similarity:retry:delay:%dmin, delayMinutes); // 序列化任务信息 String taskJson serializeTask(task); // 放入延迟队列 redisTemplate.opsForList().leftPush(delayedQueueKey, taskJson); // 设置队列过期时间 redisTemplate.expire(delayedQueueKey, delayMinutes * 60 300, TimeUnit.SECONDS); log.info(Task {} scheduled for retry in {} minutes, task.getTaskId(), delayMinutes); } /** * 计算延迟时间重试次数越多延迟越长 */ private int calculateDelayMinutes(int retryCount) { // 指数退避策略1, 2, 4, 8, 16分钟... return (int) Math.pow(2, retryCount); } /** * 判断是否为暂时性错误 */ private boolean isTransientError(Exception error) { // 网络超时、连接拒绝等属于暂时性错误 return error instanceof SocketTimeoutException || error instanceof ConnectException || error.getMessage().contains(timeout); } }指数退避策略是个很好的实践它能避免在服务暂时不可用时产生大量的重试请求给服务恢复的时间。4.3 进度查询与监控用户提交了批量任务后肯定想知道处理到哪了。我实现了一个进度查询接口可以实时查看任务处理情况。RestController RequestMapping(/api/task) public class TaskController { Autowired private RedisTemplateString, String redisTemplate; Autowired private TaskRepository taskRepository; /** * 查询批次处理进度 */ GetMapping(/progress/{batchId}) public MapString, Object getProgress(PathVariable String batchId) { String batchKey similarity:batch: batchId; MapObject, Object batchInfo redisTemplate.opsForHash().entries(batchKey); if (batchInfo.isEmpty()) { // 如果Redis中没有从数据库统计 return calculateProgressFromDB(batchId); } int total Integer.parseInt((String) batchInfo.get(total)); int pending Integer.parseInt((String) batchInfo.get(pending)); int processed total - pending; MapString, Object result new HashMap(); result.put(batchId, batchId); result.put(total, total); result.put(processed, processed); result.put(pending, pending); result.put(progress, total 0 ? (processed * 100.0 / total) : 0); result.put(estimatedTime, estimateRemainingTime(processed, total)); return result; } /** * 查询任务详情 */ GetMapping(/detail/{taskId}) public SimilarityTask getTaskDetail(PathVariable String taskId) { return taskRepository.findById(taskId) .orElseThrow(() - new ResourceNotFoundException(Task not found)); } /** * 批量查询任务状态 */ PostMapping(/batch-status) public ListTaskStatus getBatchStatus(RequestBody ListString taskIds) { return taskRepository.findStatusByTaskIds(taskIds); } /** * 估计剩余时间 */ private String estimateRemainingTime(int processed, int total) { if (processed 0) { return 无法估计; } // 假设平均每个任务处理时间相同 // 实际项目中可以根据历史数据动态计算 long avgProcessTime 1000; // 假设平均1秒一个任务 long remaining (total - processed) * avgProcessTime; if (remaining 60000) { return 小于1分钟; } else if (remaining 3600000) { return String.format(%d分钟, remaining / 60000); } else { return String.format(%.1f小时, remaining / 3600000.0); } } }进度查询不仅要显示数字还要给用户一个直观的感受。我加了进度百分比和剩余时间估计这样用户就知道大概要等多久。5. 资源控制与性能优化当有大量任务需要处理时如果不加控制可能会把模型服务或者数据库打垮。我做了几层资源控制。5.1 并发控制通过控制worker的数量和每个worker的处理速度来控制对下游服务的压力。Component public class ResourceController { Autowired private RedisTemplateString, String redisTemplate; private final Semaphore modelSemaphore new Semaphore(10); // 最多10个并发调用模型 /** * 控制模型调用的并发数 */ public float callModelWithLimit(String sentence1, String sentence2) throws InterruptedException { modelSemaphore.acquire(); // 获取许可 try { return modelService.calculateSimilarity(sentence1, sentence2); } finally { modelSemaphore.release(); // 释放许可 } } /** * 动态调整并发数 */ public void adjustConcurrency(int newLimit) { int currentPermits modelSemaphore.availablePermits(); int difference newLimit - (modelSemaphore.getQueueLength() currentPermits); if (difference 0) { modelSemaphore.release(difference); // 增加许可 } else if (difference 0) { // 减少许可需要更复杂的逻辑这里简化处理 log.warn(Decreasing concurrency limit is not directly supported); } } /** * 基于响应时间的自适应控制 */ public void adaptiveControl() { // 监控模型API的响应时间 long avgResponseTime getAverageResponseTime(); if (avgResponseTime 5000) { // 响应时间超过5秒减少并发 adjustConcurrency(Math.max(1, modelSemaphore.availablePermits() - 2)); } else if (avgResponseTime 1000) { // 响应时间小于1秒可以增加并发 adjustConcurrency(modelSemaphore.availablePermits() 2); } } }5.2 批量处理优化对于模型API调用如果支持批量处理一定要用批量接口。比如一次传100个句子对比分100次调用要快得多。Component public class BatchProcessor { Autowired private ModelService modelService; private static final int BATCH_SIZE 50; /** * 批量处理任务 */ public ListSimilarityResult processBatch(ListSimilarityTask tasks) { ListSimilarityResult results new ArrayList(); // 分批处理每批50个 for (int i 0; i tasks.size(); i BATCH_SIZE) { int end Math.min(i BATCH_SIZE, tasks.size()); ListSimilarityTask batch tasks.subList(i, end); // 准备批量请求 ListString sentences1 batch.stream() .map(SimilarityTask::getSentence1) .collect(Collectors.toList()); ListString sentences2 batch.stream() .map(SimilarityTask::getSentence2) .collect(Collectors.toList()); // 调用批量接口 ListFloat similarities modelService.batchCalculateSimilarity(sentences1, sentences2); // 组装结果 for (int j 0; j batch.size(); j) { SimilarityResult result new SimilarityResult(); result.setTaskId(batch.get(j).getTaskId()); result.setSentence1(batch.get(j).getSentence1()); result.setSentence2(batch.get(j).getSentence2()); result.setSimilarity(similarities.get(j)); result.setCreateTime(new Date()); results.add(result); } // 批量保存结果 resultRepository.saveAll(results); // 短暂休息避免压力过大 if (end tasks.size()) { try { Thread.sleep(100); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } } } return results; } }5.3 内存与连接池管理在处理大量任务时内存和数据库连接很容易成为瓶颈。我做了这些优化流式处理不要一次性把所有任务加载到内存可以分批从数据库读取连接池配置合理配置数据库连接池大小避免连接不够用或者太多连接拖慢数据库结果缓存对于相同的句子对如果近期计算过可以直接用缓存结果Component public class MemoryManager { Autowired private RedisTemplateString, String redisTemplate; /** * 检查并清理Redis内存 */ public void cleanupRedis() { // 清理过期的任务状态 SetString expiredKeys findExpiredKeys(similarity:task:*); if (!expiredKeys.isEmpty()) { redisTemplate.delete(expiredKeys); } // 清理过期的进度信息 expiredKeys findExpiredKeys(similarity:batch:*); if (!expiredKeys.isEmpty()) { redisTemplate.delete(expiredKeys); } // 如果Redis内存使用率超过80%清理一些临时数据 Long usedMemory getRedisMemoryUsage(); Long totalMemory getRedisTotalMemory(); if (usedMemory ! null totalMemory ! null) { double usageRatio usedMemory.doubleValue() / totalMemory.doubleValue(); if (usageRatio 0.8) { cleanupTemporaryData(); } } } /** * 结果缓存 */ public Float getCachedSimilarity(String sentence1, String sentence2) { String cacheKey generateCacheKey(sentence1, sentence2); String cachedValue redisTemplate.opsForValue().get(cacheKey); if (cachedValue ! null) { return Float.parseFloat(cachedValue); } return null; } /** * 设置结果缓存 */ public void cacheSimilarity(String sentence1, String sentence2, float similarity) { String cacheKey generateCacheKey(sentence1, sentence2); // 缓存7天 redisTemplate.opsForValue().set(cacheKey, String.valueOf(similarity), 7, TimeUnit.DAYS); } }6. 实际应用效果这套系统在我们公司实际运行了半年多处理了超过500万对句子的相似度计算。说几个实际的数据吧。处理速度方面单机部署4个worker平均每秒能处理20-30个任务。如果是100万对句子大概需要10个小时左右。如果部署多台机器速度还能线性提升。可靠性方面系统运行期间遇到过几次网络波动和数据库维护但都没有导致任务丢失。重试机制很管用大部分暂时性错误都能通过重试解决。真正进入死信队列需要人工处理的任务只占总数的0.1%左右。资源使用上Redis内存占用大概在1GB左右缓存了最近7天的任务状态和结果。数据库方面结果表已经积累了500多万条记录查询速度还是很快的因为做了合适的索引。有个实际案例挺有意思的。市场部门要做一次大规模的竞品分析需要计算我们产品描述和所有竞品描述的相似度。总共涉及2000多个产品每个产品平均500字描述组合起来有400多万对句子。用同步方式估计得跑好几天用我们的队列系统10台机器同时处理一晚上就搞定了。7. 总结从实际项目经验来看用任务队列处理NLP模型的批量任务确实是个好方案。特别是对于StructBERT这种句子相似度计算每次调用虽然不算特别慢但架不住数量多啊。有了队列系统就能把大任务拆小并行处理还能保证可靠性。这套方案有几个关键点我觉得特别重要。一是任务去重能省不少计算资源。二是合理的重试策略不能无限重试也不能一次失败就放弃。三是进度查询让用户知道任务处理到哪了心里有底。四是资源控制别把下游服务打垮了。如果你也要做类似的批量处理系统我建议先从简单的Redis队列开始实现基本功能。等业务量上来了再考虑用RabbitMQ这种更专业的消息队列。监控和报警也要早点做不然出了问题都不知道。代码方面我建议把任务生产、消费、重试这些逻辑都模块化方便维护和扩展。数据库设计也要考虑好该加索引的地方一定要加不然数据量大了查询会慢。最后说点实际感受。做这种系统就像搭积木每个组件都不复杂但组合起来要稳定可靠就不容易了。中间我们踩过不少坑比如Redis内存满了、数据库连接不够用、重试风暴等等。不过解决这些问题后系统就越来越稳了。现在这套系统已经成为我们公司好几个AI项目的标配基础设施了。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。