训练任务添加并发控制

This commit is contained in:
wanxiaoli 2026-08-05 16:41:09 +08:00
parent 34d259feb4
commit 43e820fd3b
3 changed files with 32 additions and 2 deletions

View File

@ -6,6 +6,7 @@ import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.autoconfigure.data.redis.RedisAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.web.servlet.ServletComponentScan;
import org.springframework.scheduling.annotation.EnableScheduling;
@SpringBootApplication(
scanBasePackages = {
@ -27,6 +28,7 @@ import org.springframework.boot.web.servlet.ServletComponentScan;
"com.yfd.platform.**.mapper",
"com.yfd.business.css.**.mapper"
})
@EnableScheduling
public class CriticalScenarioApplication {
public static void main(String[] args) {
SpringApplication.run(CriticalScenarioApplication.class, args);

View File

@ -45,6 +45,9 @@ public class CriticalDataController {
@PostMapping
public boolean create(@RequestBody CriticalData data) {
data.setModifier(currentUsername());
if (data.getExtraFeatures() != null && data.getExtraFeatures().isBlank()) {
data.setExtraFeatures(null);
}
return criticalDataService.save(data);
}
@ -82,6 +85,9 @@ public class CriticalDataController {
public boolean update(@RequestBody CriticalData data) {
data.setModifier(currentUsername());
data.setUpdatedAt(LocalDateTime.now());
if (data.getExtraFeatures() != null && data.getExtraFeatures().isBlank()) {
data.setExtraFeatures(null);
}
return criticalDataService.updateById(data);
}

View File

@ -1,6 +1,7 @@
package com.yfd.business.css.service.impl;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.conditions.update.UpdateWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
@ -19,6 +20,7 @@ import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.scheduling.annotation.Async;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.client.HttpStatusCodeException;
@ -78,6 +80,7 @@ public class ModelTrainServiceImpl extends ServiceImpl<ModelTrainTaskMapper, Mod
private static final Pattern DERIVED_EXPR_ALLOWED_PATTERN = Pattern.compile("^[A-Za-z0-9_+\\-*/()\\s]+$");
private static final int MAX_TRAINING = 2;
private static final int MAX_PENDING = 5;
private static final int TRAINING_TIMEOUT_MINUTES = 120;
@Override
public String uploadDataset(MultipartFile file) {
@ -188,7 +191,7 @@ public class ModelTrainServiceImpl extends ServiceImpl<ModelTrainTaskMapper, Mod
@Override
public Map<String, Object> checkCapacity() {
long training = countByStatus(List.of("Training", "TRAINING"));
long training = countActiveTraining();
long pending = countByStatus(List.of("Pending", "PENDING"));
boolean canSubmit = !(training >= MAX_TRAINING && pending >= MAX_PENDING);
@ -212,7 +215,7 @@ public class ModelTrainServiceImpl extends ServiceImpl<ModelTrainTaskMapper, Mod
}
private void assertCapacityForSubmit() {
long training = countByStatus(List.of("Training", "TRAINING"));
long training = countActiveTraining();
long pending = countByStatus(List.of("Pending", "PENDING"));
if (training >= MAX_TRAINING && pending >= MAX_PENDING) {
throw new BizException(buildCapacityBlockReason(training, pending));
@ -225,9 +228,28 @@ public class ModelTrainServiceImpl extends ServiceImpl<ModelTrainTaskMapper, Mod
return this.count(q);
}
private long countActiveTraining() {
LocalDateTime threshold = LocalDateTime.now().minusMinutes(TRAINING_TIMEOUT_MINUTES);
QueryWrapper<ModelTrainTask> q = new QueryWrapper<>();
q.in("status", List.of("Training", "TRAINING"));
q.ge("updated_at", threshold);
return this.count(q);
}
private String buildCapacityBlockReason(long training, long pending) {
return "当前训练中(" + training + ")且等待中(" + pending + ")已达上限,请稍后再试";
}
@Scheduled(fixedDelay = 300000)
public void recycleTimeoutTrainingTasks() {
LocalDateTime threshold = LocalDateTime.now().minusMinutes(TRAINING_TIMEOUT_MINUTES);
UpdateWrapper<ModelTrainTask> uw = new UpdateWrapper<>();
uw.in("status", List.of("Training", "TRAINING"));
uw.lt("updated_at", threshold);
uw.set("status", "Failed");
uw.setSql("error_log = CONCAT(IFNULL(error_log,''), '\\n[auto-timeout] training timeout at ', NOW())");
this.update(uw);
}
@Async
public void asyncCallTrain(ModelTrainTask task) {