1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
|
from os import makedirs, path, listdir, rename
from dataclasses import dataclass
from uuid import uuid4
from time import time
@dataclass
class FemtoTask:
id: str
data: bytes
class FemtoQueue:
def __init__(
self,
data_dir: str,
node_name: str,
timeout_stale_ms: int = 30_000,
):
assert node_name != "pending" \
and node_name != "done" \
and node_name != "failed"
self.node_name = node_name
assert timeout_stale_ms > 0
self.timeout_stale_ms = timeout_stale_ms
self.latest_stale_check_ts: float | None = None
self.data_dir = data_dir
self.dir_pending = path.join(data_dir, "pending")
self.dir_in_progress = path.join(data_dir, node_name)
self.dir_done = path.join(data_dir, "done")
self.dir_failed = path.join(data_dir, "failed")
makedirs(self.data_dir, exist_ok=True)
makedirs(self.dir_pending, exist_ok=True)
makedirs(self.dir_in_progress, exist_ok=True)
makedirs(self.dir_done, exist_ok=True)
makedirs(self.dir_failed, exist_ok=True)
def push(self, data: bytes) -> str:
id = f"{uuid4().hex[-12:]}_{str(int(time()))}"
pending_path = path.join(self.dir_pending, id)
with open(pending_path, "wb") as f:
f.write(data)
return id
def _release_stale_tasks(self):
now = time()
# Only run this every `timeout_stale_ms` milliseconds because iterating
# through all tasks is slow
timeout_sec = self.timeout_stale_ms / 1000.0
if self.latest_stale_check_ts is not None and now - self.latest_stale_check_ts < timeout_sec:
return
self.latest_stale_check_ts = now
for dir_name in listdir(self.data_dir):
full_dir_path = path.join(self.data_dir, dir_name)
# Skip non-directories and reserved names
if not path.isdir(full_dir_path):
continue
if dir_name in ("pending", "done", "failed", self.node_name):
continue
# Check tasks in this node's in-progress directory
for task_file in listdir(full_dir_path):
task_path = path.join(full_dir_path, task_file)
modified_time = int(task_file[13:])
if now - modified_time < timeout_sec:
continue
try:
pending_path = path.join(self.dir_pending, task_file)
rename(task_path, pending_path)
except FileNotFoundError:
continue # Task may have been moved by another node
def _pop_task_path(self) -> str | None:
# First check assigned tasks in progress
tasks = listdir(self.dir_in_progress)
if len(tasks) > 0:
id = min(tasks, key=lambda x: x[13:])
return path.join(self.dir_in_progress, id)
# Then check pending tasks
tasks = listdir(self.dir_pending)
if len(tasks) > 0:
id = min(tasks, key=lambda x: x[13:])
return path.join(self.dir_pending, id)
return None
def pop(self) -> FemtoTask | None:
self._release_stale_tasks()
while True:
task = self._pop_task_path()
if task is None: return None
id = path.basename(task)
in_progress_path = path.join(self.dir_in_progress, id)
try:
rename(task, in_progress_path)
except FileNotFoundError:
# If another node grabbed the task, just get another one
continue
with open(in_progress_path, "rb") as f:
content = f.read()
return FemtoTask(id = id, data = content)
def done(self, task: FemtoTask):
in_progress_path = path.join(self.dir_in_progress, task.id)
done_path = path.join(self.dir_done, task.id)
try:
rename(in_progress_path, done_path)
except FileNotFoundError as e:
raise Exception(f"Tried to complete a task that is not in progress, id={task.id}") from e
def fail(self, task: FemtoTask):
in_progress_path = path.join(self.dir_in_progress, task.id)
failed_path = path.join(self.dir_failed, task.id)
try:
rename(in_progress_path, failed_path)
except FileNotFoundError as e:
raise Exception(f"Tried to fail a task that is not in progress, id={task.id}") from e
|
