Kevin Hu commited on
Commit
b08b226
·
1 Parent(s): 6f1cd9e

add retry count to task (#2152)

Browse files

### What problem does this PR solve?


### Type of change

- [x] New Feature (non-breaking change which adds functionality)

api/db/db_models.py CHANGED
@@ -788,6 +788,7 @@ class Task(DataBaseModel):
788
  null=True,
789
  help_text="process message",
790
  default="")
 
791
 
792
 
793
  class Dialog(DataBaseModel):
@@ -982,3 +983,10 @@ def migrate_db():
982
  DB.execute_sql('ALTER TABLE llm ADD PRIMARY KEY (llm_name,fid);')
983
  except Exception as e:
984
  pass
 
 
 
 
 
 
 
 
788
  null=True,
789
  help_text="process message",
790
  default="")
791
+ retry_count = IntegerField(default=0)
792
 
793
 
794
  class Dialog(DataBaseModel):
 
983
  DB.execute_sql('ALTER TABLE llm ADD PRIMARY KEY (llm_name,fid);')
984
  except Exception as e:
985
  pass
986
+ try:
987
+ migrate(
988
+ migrator.add_column('task', 'retry_count', IntegerField(default=0))
989
+ )
990
+ except Exception as e:
991
+ pass
992
+
api/db/services/task_service.py CHANGED
@@ -64,9 +64,20 @@ class TaskService(CommonService):
64
  docs = list(docs.dicts())
65
  if not docs: return []
66
 
67
- cls.model.update(progress_msg=cls.model.progress_msg + "\n" + "Task has been received.",
68
- progress=random.random() / 10.).where(
 
 
 
 
 
 
 
 
69
  cls.model.id == docs[0]["id"]).execute()
 
 
 
70
  return docs
71
 
72
  @classmethod
 
64
  docs = list(docs.dicts())
65
  if not docs: return []
66
 
67
+ msg = "\nTask has been received."
68
+ prog = random.random() / 10.
69
+ if docs[0]["retry_count"] >= 3:
70
+ msg = "\nERROR: Task is abandoned after 3 times attempts."
71
+ prog = -1
72
+
73
+ cls.model.update(progress_msg=cls.model.progress_msg + msg,
74
+ progress=prog,
75
+ retry_count=docs[0]["retry_count"]+1
76
+ ).where(
77
  cls.model.id == docs[0]["id"]).execute()
78
+
79
+ if docs[0]["retry_count"] >= 3: return []
80
+
81
  return docs
82
 
83
  @classmethod