Ron Au commited on
Commit
3750ff9
1 Parent(s): 2d460b1

fix(eta): Base ETA on original place in queue

Browse files
Files changed (1) hide show
  1. app.py +27 -0
app.py CHANGED
@@ -16,6 +16,29 @@ def index():
16
  tasks = {}
17
 
18
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19
  @app.route('/task/create')
20
  def create_task():
21
  prompt = request.args.get('prompt') or "покемон"
@@ -28,10 +51,14 @@ def create_task():
28
  "task_id": task_id,
29
  "created_at": created_at,
30
  "prompt": prompt,
 
31
  "status": "pending",
32
  "poll_count": 0,
33
  }
34
 
 
 
 
35
  return jsonify(tasks[task_id])
36
 
37
 
 
16
  tasks = {}
17
 
18
 
19
+ def place_in_queue(task_id):
20
+
21
+ pending_tasks = list(task for task in tasks.values()
22
+ if task["status"] == "pending")
23
+
24
+ try:
25
+ return pending_tasks.index(task_id) + 1
26
+ except:
27
+ return 0
28
+
29
+
30
+ def calculate_eta(task_id):
31
+ total_durations = list(task["completed_at"] - task["created_at"]
32
+ for task in tasks.values() if "completed_at" in task)
33
+
34
+ place = tasks[task_id]["initial_place_in_queue"] or 1
35
+
36
+ if len(total_durations):
37
+ return sum(total_durations) / len(total_durations) * place
38
+ else:
39
+ return 40 * place
40
+
41
+
42
  @app.route('/task/create')
43
  def create_task():
44
  prompt = request.args.get('prompt') or "покемон"
 
51
  "task_id": task_id,
52
  "created_at": created_at,
53
  "prompt": prompt,
54
+ "initial_place_in_queue": place_in_queue(task_id),
55
  "status": "pending",
56
  "poll_count": 0,
57
  }
58
 
59
+ print("Place in queue: ", place_in_queue(task_id))
60
+ print("ETA: ", calculate_eta(task_id))
61
+
62
  return jsonify(tasks[task_id])
63
 
64