danvics/ytptube
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Updated the worker pool to be dynamic, and make it so that live stream downloading doesn't count towards the worker pool limit.

Browse source
This commit is contained in:
ArabCoders 2025-03-08 19:53:15 +03:00
parent fbdeec6a99
commit e4794eda7c
3 changed files with 146 additions and 85 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.vscode/launch.json vendored
View file

@ -35,7 +35,6 @@
"YTP_DOWNLOAD_PATH": "${workspaceFolder}/var/downloads",
"YTP_TEMP_PATH": "${workspaceFolder}/var/tmp",
"PYDEVD_DISABLE_FILE_VALIDATION": "1",
"YTP_MAX_WORKERS": "2",
"YTP_IGNORE_UI": "true",
"YTP_PIP_IGNORE_UPDATES": "true",
"YTP_LOG_LEVEL": "DEBUG",

228
app/library/AsyncPool.py
View file

@ -1,5 +1,6 @@
import asyncio
import logging
import uuid
from datetime import UTC, datetime
@ -27,18 +28,20 @@ class AsyncPool:
against each item retrieved from the queue. If any exceptions are raised out of
worker_co, self.exceptions will be set to True.
:param loop: asyncio loop to use
:param num_workers: number of async tasks which will pull from the internal queue
:param name: name of the worker pool (used for logging)
:param logger: logger to use
:param worker_co: async coroutine to call when an item is retrieved from the queue
:param load_factor: multiplier used for number of items in queue
:param max_task_time: maximum time allowed for each task before a CancelledError is raised in the task.
Set to None for no limit.
:param return_futures: set to reture to return a future for each `push` (imposes CPU overhead)
:param raise_on_join: raise on join if any exceptions have occurred, default is False
Args:
num_workers (int): number of async tasks which will pull from the internal queue
worker_co (coroutine): async coroutine to call when an item is retrieved from the queue
name (str): name of the worker pool (used for logging)
logger (logging.Logger): logger to use
loop (asyncio.AbstractEventLoop|None): asyncio loop to use
load_factor (int): multiplier used for number of items in queue
max_task_time (int|None): maximum time allowed for each task before a CancelledError is raised in the task.
return_futures (bool): set to reture to return a future for each `push` (imposes CPU overhead)
raise_on_join (bool): raise on join if any exceptions have occurred, default is False
Returns:
AsyncPool: instance of AsyncPool
:return: instance of AsyncWorkerPool
"""
self._loop = loop if loop else asyncio.get_event_loop()
self._num_workers = num_workers
@ -53,77 +56,100 @@ class AsyncPool:
self._worker_co = worker_co
self._status: dict[str, dict | None] = {}
async def _worker_loop(self, worker_id: str):
"""
This is the main loop for each worker. It will pull from the queue and call the worker_co
coroutine with the arguments passed to `push`.
:param worker_id: id of the worker
:return: None
"""
while True:
got_obj = False
future = None
try:
item = await self._queue.get()
got_obj = True
if item.__class__ is Terminator:
break
future, args, kwargs = item
self._status[worker_id] = {
"started": self._time().isoformat(),
"data": kwargs.get("entry", {"info": {}}).info.__dict__,
}
# the wait_for will cancel the task (task sees CancelledError) and raises a TimeoutError from here
# so be wary of catching TimeoutErrors in this loop
result = await asyncio.wait_for(self._worker_co(*args, **kwargs), timeout=self._max_task_time)
if future:
future.set_result(result)
except (KeyboardInterrupt, MemoryError, SystemExit) as e:
if future:
future.set_exception(e)
self._exceptions = True
raise
except BaseException as e:
if isinstance(e, asyncio.exceptions.CancelledError):
raise
self._exceptions = True
if future:
# don't log the failure when the client is receiving the future
future.set_exception(e)
else:
self._logger.exception(f"Worker call failed. {e!s}")
finally:
self._status[worker_id] = None
if got_obj:
self._queue.task_done()
@property
def exceptions(self):
return self._exceptions
async def _worker_loop(self, worker_id: str):
"""
The main persistent worker loop that continuously processes jobs from the shared queue.
Args:
worker_id (str): the id of the worker processing this job.
"""
while True:
item = await self._queue.get()
should_continue = await self._process_item(worker_id, item, from_queue=True)
if not should_continue:
break
async def _process_item(self, worker_id: str, item: tuple | Terminator, from_queue: bool = True) -> bool:
"""
Processes a single job item.
Args:
worker_id (str): the id of the worker processing this job.
item (tuple|Terminator): the job item (tuple of (future, args, kwargs)) or a Terminator.
from_queue (bool): indicates whether the job came from the shared queue; if True, task_done() is called.
Returns:
bool: False if the item is a Terminator (indicating termination), True otherwise.
"""
future = None
is_terminator = isinstance(item, Terminator)
try:
if is_terminator:
return False
future, args, kwargs = item
self._status[worker_id] = {
"started": self._time().isoformat(),
"data": kwargs.get("entry", {"info": {}}).info.__dict__,
}
result = await asyncio.wait_for(self._worker_co(*args, **kwargs), timeout=self._max_task_time)
if future:
future.set_result(result)
except (KeyboardInterrupt, MemoryError, SystemExit) as e:
if future:
future.set_exception(e)
self._exceptions = True
raise
except Exception as e:
self._exceptions = True
if future:
future.set_exception(e)
else:
self._logger.exception(f"Worker call failed. {e!s}")
finally:
self._status[worker_id] = None
if not is_terminator and from_queue:
self._queue.task_done()
return True
def has_open_workers(self) -> bool:
"""
:return: True if there are open workers.
Check if there are open workers.
Returns:
bool: True if there are open workers, False otherwise.
"""
return self.get_available_workers() > 0
def get_available_workers(self) -> int:
"""
:return: number of available workers.
Get the number of available workers.
Returns:
int: number of available workers.
"""
return sum(1 for worker_status in self._status.values() if worker_status is None)
def get_workers_status(self) -> dict[str, datetime | None]:
"""
:return: dictionary of worker status.
Get the status of all workers.
Returns:
dict: dictionary of worker status.
"""
return self._status
@ -134,36 +160,65 @@ class AsyncPool:
async def __aexit__(self, exc_type, exc_val, exc_tb):
await self.join()
async def push(self, *args, **kwargs) -> asyncio.Future:
async def push(self, *args, is_temp: bool = False, **kwargs) -> asyncio.Future:
"""
Method to push work to `worker_co` passed initially to `__init__`.
Push work to the worker_co. If is_temp is True, a temporary worker is spawned
:param args: position arguments to be passed to `worker_co`
:param kwargs: keyword arguments to be passed to `worker_co`
Args:
args: position arguments to be passed to `worker_co`
is_temp (bool): flag to indicate if the job is temporary, default is False
kwargs (dict): keyword arguments to be passed to `worker_co`
Returns:
asyncio.Future: future of result.
:return: future of result.
"""
future = asyncio.futures.Future(loop=self._loop) if self._return_futures else None
if is_temp:
temp_worker_id = f"temp_worker_{uuid.uuid4()!s}"
self._logger.info(f"Creating temporary worker '{temp_worker_id}'.")
# Spawn a temporary worker that processes this one job.
task = asyncio.ensure_future(
self._process_item(temp_worker_id, (future, args, kwargs), from_queue=False), loop=self._loop
)
self._workers[temp_worker_id] = task
# Attach a callback to clean up temporary worker entries when done.
def cleanup_temp_worker(_):
self._workers.pop(temp_worker_id, None)
self._status.pop(temp_worker_id, None)
self._logger.info(f"Temporary worker '{temp_worker_id}' has terminated.")
task.add_done_callback(cleanup_temp_worker)
return future
await self._queue.put((future, args, kwargs))
self._logger.debug(f"'{self._name}' pool has received a new job. {args} {kwargs}")
return future
def start(self):
"""Will start up worker pool and reset exception state"""
"""
Will start up worker pool and reset exception state
"""
self._exceptions = False
for worker_number in range(self._num_workers):
worker_id = f"worker_{worker_number+1}"
self._create_worker(worker_id)
async def restart(self, worker_id: str, msg: str = None) -> bool:
async def restart(self, worker_id: str, msg: str | None = None) -> bool:
"""
Will restart the worker pool
:param worker_id: worker id to restart
:param msg: message to send to the worker
:return: True if worker was restarted
Args:
worker_id (str): worker id to restart
msg (str|None): message to send to the worker, default is None
Returns:
bool: True if worker was restarted.
"""
if worker_id not in self._workers:
self._logger.warning(f"Worker {worker_id} does not exist.")
@ -185,18 +240,25 @@ class AsyncPool:
return True
async def on_shutdown(self, _) -> None:
"""Will shutdown the worker pool"""
"""
Will shutdown the worker pool
"""
try:
await asyncio.wait_for(asyncio.gather(*[self.stop(worker_id) for worker_id in self._workers]), timeout=10)
except Exception:
self._logger.error(f"Exception shutting down {self._name}")
async def stop(self, worker_id: str, msg: str = None) -> bool:
async def stop(self, worker_id: str, msg: str | None = None) -> bool:
"""
Will stop the worker
:param worker_id: worker id to stop
:param msg: message to send to the worker
:return: True if worker was stopped
Args:
worker_id (str): worker id to stop
msg (str|None): message to send to the worker, default is None
Returns:
bool: True if worker was stopped.
"""
if worker_id not in self._workers:
self._logger.warning(f"Worker {worker_id} does not exist.")

2
app/library/DownloadQueue.py
View file

@ -640,7 +640,7 @@ class DownloadQueue(metaclass=Singleton):
LOG.debug(f"Pushing {entry=} to executor.")
if entry.started() is False and entry.is_cancelled() is False:
await self.pool.push(id=entry.info._id, entry=entry)
await self.pool.push(is_temp=entry.is_live, id=entry.info._id, entry=entry)
LOG.debug(f"Pushed {entry=} to executor.")
await asyncio.sleep(1)