local_task_manager.py 14.8 KB
Newer Older
LiangLiu's avatar
LiangLiu committed
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
import asyncio
import json
import os

from lightx2v.deploy.common.utils import class_try_catch_async, current_time, str2time, time2str
from lightx2v.deploy.task_manager import ActiveStatus, BaseTaskManager, FinishedStatus, TaskStatus


class LocalTaskManager(BaseTaskManager):
    def __init__(self, local_dir, metrics_monitor=None):
        self.local_dir = local_dir
        if not os.path.exists(self.local_dir):
            os.makedirs(self.local_dir)
        self.metrics_monitor = metrics_monitor

    def get_task_filename(self, task_id):
        return os.path.join(self.local_dir, f"task_{task_id}.json")

    def get_user_filename(self, user_id):
        return os.path.join(self.local_dir, f"user_{user_id}.json")

    def fmt_dict(self, data):
        super().fmt_dict(data)
        for k in ["create_t", "update_t", "ping_t"]:
            if k in data:
                data[k] = time2str(data[k])

    def parse_dict(self, data):
        super().parse_dict(data)
        for k in ["create_t", "update_t", "ping_t"]:
            if k in data:
                data[k] = str2time(data[k])

    def save(self, task, subtasks, with_fmt=True):
        info = {"task": task, "subtasks": subtasks}
        if with_fmt:
            self.fmt_dict(info["task"])
            [self.fmt_dict(x) for x in info["subtasks"]]
        out_name = self.get_task_filename(task["task_id"])
        with open(out_name, "w") as fout:
            fout.write(json.dumps(info, indent=4, ensure_ascii=False))

    def load(self, task_id, user_id=None, only_task=False):
        fpath = self.get_task_filename(task_id)
        info = json.load(open(fpath))
        task, subtasks = info["task"], info["subtasks"]
        if user_id is not None and task["user_id"] != user_id:
            raise Exception(f"Task {task_id} is not belong to user {user_id}")
        self.parse_dict(task)
        if only_task:
            return task
        for sub in subtasks:
            self.parse_dict(sub)
        return task, subtasks

    @class_try_catch_async
    async def insert_task(self, task, subtasks):
        self.save(task, subtasks)
        return True

    @class_try_catch_async
    async def list_tasks(self, **kwargs):
        tasks = []
        fs = [os.path.join(self.local_dir, f) for f in os.listdir(self.local_dir)]
        for f in os.listdir(self.local_dir):
            if not f.startswith("task_"):
                continue
            fpath = os.path.join(self.local_dir, f)
            info = json.load(open(fpath))
            if kwargs.get("subtasks", False):
                items = info["subtasks"]
                assert "user_id" not in kwargs, "user_id is not allowed when subtasks is True"
            else:
                items = [info["task"]]
            for task in items:
                self.parse_dict(task)
                if "user_id" in kwargs and task["user_id"] != kwargs["user_id"]:
                    continue
                if "status" in kwargs:
                    if isinstance(kwargs["status"], list) and task["status"] not in kwargs["status"]:
                        continue
                    elif kwargs["status"] != task["status"]:
                        continue
                if "start_created_t" in kwargs and kwargs["start_created_t"] > task["create_t"]:
                    continue
                if "end_created_t" in kwargs and kwargs["end_created_t"] < task["create_t"]:
                    continue
                if "start_updated_t" in kwargs and kwargs["start_updated_t"] > task["update_t"]:
                    continue
                if "end_updated_t" in kwargs and kwargs["end_updated_t"] < task["update_t"]:
                    continue
                if "start_ping_t" in kwargs and kwargs["start_ping_t"] > task["ping_t"]:
                    continue
                if "end_ping_t" in kwargs and kwargs["end_ping_t"] < task["ping_t"]:
                    continue
LiangLiu's avatar
LiangLiu committed
96
97
98
99
100

                # 如果不是查询子任务,则添加子任务信息到任务中
                if not kwargs.get("subtasks", False):
                    task["subtasks"] = info.get("subtasks", [])

LiangLiu's avatar
LiangLiu committed
101
102
103
                tasks.append(task)
        if "count" in kwargs:
            return len(tasks)
LiangLiu's avatar
LiangLiu committed
104
105
106
107

        sort_key = "update_t" if kwargs.get("sort_by_update_t", False) else "create_t"
        tasks = sorted(tasks, key=lambda x: x[sort_key], reverse=True)

LiangLiu's avatar
LiangLiu committed
108
109
110
111
112
113
114
115
116
117
118
119
        if "offset" in kwargs:
            tasks = tasks[kwargs["offset"] :]
        if "limit" in kwargs:
            tasks = tasks[: kwargs["limit"]]
        return tasks

    @class_try_catch_async
    async def query_task(self, task_id, user_id=None, only_task=True):
        return self.load(task_id, user_id, only_task)

    @class_try_catch_async
    async def next_subtasks(self, task_id):
LiangLiu's avatar
LiangLiu committed
120
        records = []
LiangLiu's avatar
LiangLiu committed
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
        task, subtasks = self.load(task_id)
        if task["status"] not in ActiveStatus:
            return []
        succeeds = set()
        for sub in subtasks:
            if sub["status"] == TaskStatus.SUCCEED:
                succeeds.add(sub["worker_name"])
        nexts = []
        for sub in subtasks:
            if sub["status"] == TaskStatus.CREATED:
                dep_ok = True
                for prev in sub["previous"]:
                    if prev not in succeeds:
                        dep_ok = False
                        break
                if dep_ok:
LiangLiu's avatar
LiangLiu committed
137
                    self.mark_subtask_change(records, sub, sub["status"], TaskStatus.PENDING)
LiangLiu's avatar
LiangLiu committed
138
139
140
141
142
143
144
145
                    sub["params"] = task["params"]
                    sub["status"] = TaskStatus.PENDING
                    sub["update_t"] = current_time()
                    nexts.append(sub)
        if len(nexts) > 0:
            task["status"] = TaskStatus.PENDING
            task["update_t"] = current_time()
            self.save(task, subtasks)
LiangLiu's avatar
LiangLiu committed
146
        self.metrics_commit(records)
LiangLiu's avatar
LiangLiu committed
147
148
149
150
        return nexts

    @class_try_catch_async
    async def run_subtasks(self, cands, worker_identity):
LiangLiu's avatar
LiangLiu committed
151
        records = []
LiangLiu's avatar
LiangLiu committed
152
153
154
155
156
157
158
159
160
        valids = []
        for cand in cands:
            task_id = cand["task_id"]
            worker_name = cand["worker_name"]
            task, subtasks = self.load(task_id)
            if task["status"] in [TaskStatus.SUCCEED, TaskStatus.FAILED, TaskStatus.CANCEL]:
                continue
            for sub in subtasks:
                if sub["worker_name"] == worker_name:
LiangLiu's avatar
LiangLiu committed
161
                    self.mark_subtask_change(records, sub, sub["status"], TaskStatus.RUNNING)
LiangLiu's avatar
LiangLiu committed
162
163
164
165
166
167
168
169
170
                    sub["status"] = TaskStatus.RUNNING
                    sub["worker_identity"] = worker_identity
                    sub["update_t"] = current_time()
                    task["status"] = TaskStatus.RUNNING
                    task["update_t"] = current_time()
                    task["ping_t"] = current_time()
                    self.save(task, subtasks)
                    valids.append(cand)
                    break
LiangLiu's avatar
LiangLiu committed
171
        self.metrics_commit(records)
LiangLiu's avatar
LiangLiu committed
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
        return valids

    @class_try_catch_async
    async def ping_subtask(self, task_id, worker_name, worker_identity):
        task, subtasks = self.load(task_id)
        for sub in subtasks:
            if sub["worker_name"] == worker_name:
                pre = sub["worker_identity"]
                assert pre == worker_identity, f"worker identity not matched: {pre} vs {worker_identity}"
                sub["ping_t"] = current_time()
                self.save(task, subtasks)
                return True
        return False

    @class_try_catch_async
    async def finish_subtasks(self, task_id, status, worker_identity=None, worker_name=None, fail_msg=None, should_running=False):
LiangLiu's avatar
LiangLiu committed
188
        records = []
LiangLiu's avatar
LiangLiu committed
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
        task, subtasks = self.load(task_id)
        subs = subtasks

        if worker_name:
            subs = [sub for sub in subtasks if sub["worker_name"] == worker_name]
        assert len(subs) >= 1, f"no worker task_id={task_id}, name={worker_name}"

        if worker_identity:
            pre = subs[0]["worker_identity"]
            assert pre == worker_identity, f"worker identity not matched: {pre} vs {worker_identity}"

        assert status in [TaskStatus.SUCCEED, TaskStatus.FAILED], f"invalid finish status: {status}"
        for sub in subs:
            if sub["status"] not in FinishedStatus:
                if should_running and sub["status"] != TaskStatus.RUNNING:
                    print(f"task {task_id} is not running, skip finish subtask: {sub}")
                    continue
LiangLiu's avatar
LiangLiu committed
206
                self.mark_subtask_change(records, sub, sub["status"], status, fail_msg=fail_msg)
LiangLiu's avatar
LiangLiu committed
207
208
209
210
211
                sub["status"] = status
                sub["update_t"] = current_time()

        if task["status"] == TaskStatus.CANCEL:
            self.save(task, subtasks)
LiangLiu's avatar
LiangLiu committed
212
            self.metrics_commit(records)
LiangLiu's avatar
LiangLiu committed
213
214
215
216
217
218
219
220
221
222
223
224
225
            return TaskStatus.CANCEL

        running_subs = []
        failed_sub = False
        for sub in subtasks:
            if sub["status"] not in FinishedStatus:
                running_subs.append(sub)
            if sub["status"] == TaskStatus.FAILED:
                failed_sub = True

        # some subtask failed, we should fail all other subtasks
        if failed_sub:
            if task["status"] != TaskStatus.FAILED:
LiangLiu's avatar
LiangLiu committed
226
                self.mark_task_end(records, task, TaskStatus.FAILED)
LiangLiu's avatar
LiangLiu committed
227
228
229
                task["status"] = TaskStatus.FAILED
                task["update_t"] = current_time()
            for sub in running_subs:
LiangLiu's avatar
LiangLiu committed
230
                self.mark_subtask_change(records, sub, sub["status"], TaskStatus.FAILED, fail_msg="other subtask failed")
LiangLiu's avatar
LiangLiu committed
231
232
233
                sub["status"] = TaskStatus.FAILED
                sub["update_t"] = current_time()
            self.save(task, subtasks)
LiangLiu's avatar
LiangLiu committed
234
            self.metrics_commit(records)
LiangLiu's avatar
LiangLiu committed
235
236
237
238
239
            return TaskStatus.FAILED

        # all subtasks finished and all succeed
        elif len(running_subs) == 0:
            if task["status"] != TaskStatus.SUCCEED:
LiangLiu's avatar
LiangLiu committed
240
                self.mark_task_end(records, task, TaskStatus.SUCCEED)
LiangLiu's avatar
LiangLiu committed
241
242
243
                task["status"] = TaskStatus.SUCCEED
                task["update_t"] = current_time()
            self.save(task, subtasks)
LiangLiu's avatar
LiangLiu committed
244
            self.metrics_commit(records)
LiangLiu's avatar
LiangLiu committed
245
246
247
            return TaskStatus.SUCCEED

        self.save(task, subtasks)
LiangLiu's avatar
LiangLiu committed
248
        self.metrics_commit(records)
LiangLiu's avatar
LiangLiu committed
249
250
251
252
        return None

    @class_try_catch_async
    async def cancel_task(self, task_id, user_id=None):
LiangLiu's avatar
LiangLiu committed
253
        records = []
LiangLiu's avatar
LiangLiu committed
254
255
256
257
258
259
        task, subtasks = self.load(task_id, user_id)
        if task["status"] not in ActiveStatus:
            return f"Task {task_id} is not in active status (current status: {task['status']}). Only tasks with status CREATED, PENDING, or RUNNING can be cancelled."

        for sub in subtasks:
            if sub["status"] not in FinishedStatus:
LiangLiu's avatar
LiangLiu committed
260
                self.mark_subtask_change(records, sub, sub["status"], TaskStatus.CANCEL)
LiangLiu's avatar
LiangLiu committed
261
262
                sub["status"] = TaskStatus.CANCEL
                sub["update_t"] = current_time()
LiangLiu's avatar
LiangLiu committed
263
        self.mark_task_end(records, task, TaskStatus.CANCEL)
LiangLiu's avatar
LiangLiu committed
264
265
266
        task["status"] = TaskStatus.CANCEL
        task["update_t"] = current_time()
        self.save(task, subtasks)
LiangLiu's avatar
LiangLiu committed
267
        self.metrics_commit(records)
LiangLiu's avatar
LiangLiu committed
268
269
270
271
        return True

    @class_try_catch_async
    async def resume_task(self, task_id, all_subtask=False, user_id=None):
LiangLiu's avatar
LiangLiu committed
272
        records = []
LiangLiu's avatar
LiangLiu committed
273
274
275
        task, subtasks = self.load(task_id, user_id)
        # the task is not finished
        if task["status"] not in FinishedStatus:
LiangLiu's avatar
LiangLiu committed
276
            return "Active task cannot be resumed"
LiangLiu's avatar
LiangLiu committed
277
278
        # the task is no need to resume
        if not all_subtask and task["status"] == TaskStatus.SUCCEED:
LiangLiu's avatar
LiangLiu committed
279
            return "Succeed task cannot be resumed"
LiangLiu's avatar
LiangLiu committed
280
281
        for sub in subtasks:
            if all_subtask or sub["status"] != TaskStatus.SUCCEED:
LiangLiu's avatar
LiangLiu committed
282
                self.mark_subtask_change(records, sub, None, TaskStatus.CREATED)
LiangLiu's avatar
LiangLiu committed
283
284
285
                sub["status"] = TaskStatus.CREATED
                sub["update_t"] = current_time()
                sub["ping_t"] = 0.0
LiangLiu's avatar
LiangLiu committed
286
        self.mark_task_start(records, task)
LiangLiu's avatar
LiangLiu committed
287
288
289
        task["status"] = TaskStatus.CREATED
        task["update_t"] = current_time()
        self.save(task, subtasks)
LiangLiu's avatar
LiangLiu committed
290
291
292
293
294
295
296
297
298
299
300
301
302
        self.metrics_commit(records)
        return True

    @class_try_catch_async
    async def delete_task(self, task_id, user_id=None):
        task = self.load(task_id, user_id, only_task=True)
        # only allow to delete finished tasks
        if task["status"] not in FinishedStatus:
            return False
        # delete task file
        task_file = self.get_task_filename(task_id)
        if os.path.exists(task_file):
            os.remove(task_file)
LiangLiu's avatar
LiangLiu committed
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
        return True

    @class_try_catch_async
    async def insert_user_if_not_exists(self, user_info):
        fpath = self.get_user_filename(user_info["user_id"])
        if os.path.exists(fpath):
            return True
        self.fmt_dict(user_info)
        with open(fpath, "w") as fout:
            fout.write(json.dumps(user_info, indent=4, ensure_ascii=False))
        return True

    @class_try_catch_async
    async def query_user(self, user_id):
        fpath = self.get_user_filename(user_id)
        if not os.path.exists(fpath):
            return None
        data = json.load(open(fpath))
        self.parse_dict(data)
        return data


async def test():
    from lightx2v.deploy.common.pipeline import Pipeline

    p = Pipeline("/data/nvme1/liuliang1/lightx2v/configs/model_pipeline.json")
    m = LocalTaskManager("/data/nvme1/liuliang1/lightx2v/local_task")
    await m.init()

    keys = ["t2v", "wan2.1", "multi_stage"]
    workers = p.get_workers(keys)
    inputs = p.get_inputs(keys)
    outputs = p.get_outputs(keys)
    params = {
        "prompt": "fake input prompts",
        "resolution": {
            "height": 233,
            "width": 456,
        },
    }

    user_info = {
        "source": "github",
        "id": "test-id-233",
        "username": "test-username-233",
        "email": "test-email-233@test.com",
        "homepage": "https://test.com",
        "avatar_url": "https://test.com/avatar.png",
    }
    user_id = await m.create_user(user_info)
    print(" - create_user:", user_id)

    user = await m.query_user(user_id)
    print(" - query_user:", user)

    task_id = await m.create_task(keys, workers, params, inputs, outputs, user_id)
    print(" - create_task:", task_id)

    tasks = await m.list_tasks()
    print(" - list_tasks:", tasks)

    task = await m.query_task(task_id)
    print(" - query_task:", task)

    subtasks = await m.next_subtasks(task_id)
    print(" - next_subtasks:", subtasks)

    await m.run_subtasks(subtasks, "fake-worker")
    await m.finish_subtasks(task_id, TaskStatus.FAILED)
    await m.cancel_task(task_id)
    await m.resume_task(task_id)
    for sub in subtasks:
        await m.finish_subtasks(sub["task_id"], TaskStatus.SUCCEED, worker_name=sub["worker_name"], worker_identity="fake-worker")

    subtasks = await m.next_subtasks(task_id)
    print(" - final next_subtasks:", subtasks)

    task = await m.query_task(task_id)
    print(" - final task:", task)

    await m.close()


if __name__ == "__main__":
    asyncio.run(test())