forked from harry0703/MoneyPrinterTurbo
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: add support for maximum concurrency of /api/v1/videos
- Loading branch information
kevin.zhang
committed
Apr 16, 2024
1 parent
414bcb0
commit abe12ab
Showing
7 changed files
with
170 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,57 @@ | ||
import threading | ||
from typing import Callable, Any, Dict | ||
|
||
|
||
class TaskManager: | ||
def __init__(self, max_concurrent_tasks: int): | ||
self.max_concurrent_tasks = max_concurrent_tasks | ||
self.current_tasks = 0 | ||
self.lock = threading.Lock() | ||
self.queue = self.create_queue() | ||
|
||
def create_queue(self): | ||
raise NotImplementedError() | ||
|
||
def add_task(self, func: Callable, *args: Any, **kwargs: Any): | ||
with self.lock: | ||
if self.current_tasks < self.max_concurrent_tasks: | ||
print(f"add task: {func.__name__}, current_tasks: {self.current_tasks}") | ||
self.execute_task(func, *args, **kwargs) | ||
else: | ||
print(f"enqueue task: {func.__name__}, current_tasks: {self.current_tasks}") | ||
self.enqueue({"func": func, "args": args, "kwargs": kwargs}) | ||
|
||
def execute_task(self, func: Callable, *args: Any, **kwargs: Any): | ||
thread = threading.Thread(target=self.run_task, args=(func, *args), kwargs=kwargs) | ||
thread.start() | ||
|
||
def run_task(self, func: Callable, *args: Any, **kwargs: Any): | ||
try: | ||
with self.lock: | ||
self.current_tasks += 1 | ||
func(*args, **kwargs) # 在这里调用函数,传递*args和**kwargs | ||
finally: | ||
self.task_done() | ||
|
||
def check_queue(self): | ||
with self.lock: | ||
if self.current_tasks < self.max_concurrent_tasks and not self.is_queue_empty(): | ||
task_info = self.dequeue() | ||
func = task_info['func'] | ||
args = task_info.get('args', ()) | ||
kwargs = task_info.get('kwargs', {}) | ||
self.execute_task(func, *args, **kwargs) | ||
|
||
def task_done(self): | ||
with self.lock: | ||
self.current_tasks -= 1 | ||
self.check_queue() | ||
|
||
def enqueue(self, task: Dict): | ||
raise NotImplementedError() | ||
|
||
def dequeue(self): | ||
raise NotImplementedError() | ||
|
||
def is_queue_empty(self): | ||
raise NotImplementedError() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,18 @@ | ||
from queue import Queue | ||
from typing import Dict | ||
|
||
from app.controllers.manager.base_manager import TaskManager | ||
|
||
|
||
class InMemoryTaskManager(TaskManager): | ||
def create_queue(self): | ||
return Queue() | ||
|
||
def enqueue(self, task: Dict): | ||
self.queue.put(task) | ||
|
||
def dequeue(self): | ||
return self.queue.get() | ||
|
||
def is_queue_empty(self): | ||
return self.queue.empty() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,48 @@ | ||
import json | ||
from typing import Dict | ||
|
||
import redis | ||
|
||
from app.controllers.manager.base_manager import TaskManager | ||
from app.models.schema import VideoParams | ||
from app.services import task as tm | ||
|
||
FUNC_MAP = { | ||
'start': tm.start, | ||
# 'start_test': tm.start_test | ||
} | ||
|
||
|
||
class RedisTaskManager(TaskManager): | ||
def __init__(self, max_concurrent_tasks: int, redis_url: str): | ||
self.redis_client = redis.Redis.from_url(redis_url) | ||
super().__init__(max_concurrent_tasks) | ||
|
||
def create_queue(self): | ||
return "task_queue" | ||
|
||
def enqueue(self, task: Dict): | ||
task_with_serializable_params = task.copy() | ||
|
||
if 'params' in task['kwargs'] and isinstance(task['kwargs']['params'], VideoParams): | ||
task_with_serializable_params['kwargs']['params'] = task['kwargs']['params'].dict() | ||
|
||
# 将函数对象转换为其名称 | ||
task_with_serializable_params['func'] = task['func'].__name__ | ||
self.redis_client.rpush(self.queue, json.dumps(task_with_serializable_params)) | ||
|
||
def dequeue(self): | ||
task_json = self.redis_client.lpop(self.queue) | ||
if task_json: | ||
task_info = json.loads(task_json) | ||
# 将函数名称转换回函数对象 | ||
task_info['func'] = FUNC_MAP[task_info['func']] | ||
|
||
if 'params' in task_info['kwargs'] and isinstance(task_info['kwargs']['params'], dict): | ||
task_info['kwargs']['params'] = VideoParams(**task_info['kwargs']['params']) | ||
|
||
return task_info | ||
return None | ||
|
||
def is_queue_empty(self): | ||
return self.redis_client.llen(self.queue) == 0 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -73,7 +73,7 @@ class MaterialInfo: | |
# ] | ||
|
||
|
||
class VideoParams: | ||
class VideoParams(BaseModel): | ||
""" | ||
{ | ||
"video_subject": "", | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters