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

Merge pull request #80 from arabcoders/dev

Browse source 
Added endpoint to expose worker status
This commit is contained in:
Abdulmohsen 2024-03-20 23:16:03 +03:00 committed by GitHub
commit 6475820679
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 98 additions and 69 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

78
app/AsyncPool.py
View file

@ -1,5 +1,6 @@
import asyncio import asyncio
from datetime import datetime, timezone from datetime import datetime, timezone
import logging
class Terminator: class Terminator:
@ -7,9 +8,13 @@ class Terminator:
class AsyncPool: class AsyncPool:
def __init__(self, loop, num_workers: int, name: str, logger, worker_co, load_factor: int = 1, def __init__(self,
job_accept_duration: int = None, max_task_time: int = None, return_futures: bool = False, num_workers: int, worker_co,
raise_on_join: bool = False, log_every_n: int = None, expected_total=None): name: str, logger: logging.Logger,
loop: asyncio.AbstractEventLoop = None,
load_factor: int = 1, max_task_time: int = None,
return_futures: bool = False, raise_on_join: bool = False
):
""" """
This class will create `num_workers` asyncio tasks to work against a queue of This class will create `num_workers` asyncio tasks to work against a queue of
`num_workers * load_factor` items of back-pressure (IOW we will block after such `num_workers * load_factor` items of back-pressure (IOW we will block after such
@ -23,37 +28,27 @@ class AsyncPool:
:param logger: logger to use :param logger: logger to use
:param worker_co: async coroutine to call when an item is retrieved from the queue :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 load_factor: multiplier used for number of items in queue
:param job_accept_duration: maximum number of seconds from first push to last push before a TimeoutError will be thrown.
Set to None for no limit. Note this does not get reset on aenter/aexit.
:param max_task_time: maximum time allowed for each task before a CancelledError is raised in the task. :param max_task_time: maximum time allowed for each task before a CancelledError is raised in the task.
Set to None for no limit. Set to None for no limit.
:param return_futures: set to reture to return a future for each `push` (imposes CPU overhead) :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 :param raise_on_join: raise on join if any exceptions have occurred, default is False
:param log_every_n: (optional) set to number of `push`s each time a log statement should be printed (default does not print every-n pushes)
:param expected_total: (optional) expected total number of jobs (used for `log_event_n` logging)
:return: instance of AsyncWorkerPool :return: instance of AsyncWorkerPool
""" """
loop = loop if loop else asyncio.get_event_loop() self._loop = loop if loop else asyncio.get_event_loop()
self._loop = loop
self._num_workers = num_workers self._num_workers = num_workers
self._logger = logger self._logger = logger if logger else logging.getLogger(__name__)
self._queue = asyncio.Queue(num_workers * load_factor) self._queue = asyncio.Queue(num_workers * load_factor)
self._workers = None self._workers = None
self._exceptions = False self._exceptions = False
self._job_accept_duration = job_accept_duration
self._first_push_dt = None
self._max_task_time = max_task_time self._max_task_time = max_task_time
self._return_futures = return_futures self._return_futures = return_futures
self._raise_on_join = raise_on_join self._raise_on_join = raise_on_join
self._name = name self._name = name
self._worker_co = worker_co self._worker_co = worker_co
self._total_queued = 0 self._active_worker: dict[str, dict | None] = {}
self._log_every_n = log_every_n
self._expected_total = expected_total
self._active_workers = 0
async def _worker_loop(self): async def _worker_loop(self, worker_id: str):
while True: while True:
got_obj = False got_obj = False
future = None future = None
@ -65,11 +60,16 @@ class AsyncPool:
if item.__class__ is Terminator: if item.__class__ is Terminator:
break break
self._active_workers += 1
future, args, kwargs = item future, args, kwargs = item
self._active_worker[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 # 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 # so be wary of catching TimeoutErrors in this loop
result = await asyncio.wait_for(self._worker_co(*args, **kwargs), self._max_task_time) result = await asyncio.wait_for(self._worker_co(*args, **kwargs), timeout=self._max_task_time)
if future: if future:
future.set_result(result) future.set_result(result)
@ -85,9 +85,10 @@ class AsyncPool:
# don't log the failure when the client is receiving the future # don't log the failure when the client is receiving the future
future.set_exception(e) future.set_exception(e)
else: else:
self._logger.exception('Worker call failed') self._logger.exception(f'Worker call failed. {str(e)}')
finally: finally:
self._active_workers -= 1 self._active_worker[worker_id] = None
if got_obj: if got_obj:
self._queue.task_done() self._queue.task_done()
@ -95,21 +96,23 @@ class AsyncPool:
def exceptions(self): def exceptions(self):
return self._exceptions return self._exceptions
@property
def total_queued(self):
return self._total_queued
def has_open_workers(self) -> bool: def has_open_workers(self) -> bool:
""" """
:return: True if there are open workers. :return: True if there are open workers.
""" """
return self._active_workers < self._num_workers return self.get_available_workers() > 0
def get_available_workers(self) -> int: def get_available_workers(self) -> int:
""" """
:return: number of available workers. :return: number of available workers.
""" """
return self._num_workers - self._active_workers return sum(1 for worker_status in self._active_worker.values() if worker_status is None)
def get_workers_status(self) -> dict[str, datetime | None]:
"""
:return: dictionary of worker status.
"""
return self._active_worker
async def __aenter__(self): async def __aenter__(self):
self.start() self.start()
@ -127,18 +130,8 @@ class AsyncPool:
:return: future of result. :return: future of result.
""" """
if self._first_push_dt is None:
self._first_push_dt = self._time()
if self._job_accept_duration is not None and (self._time() - self._first_push_dt) > self._job_accept_duration:
raise TimeoutError(f"Max life time of {self._job_accept_duration}s exceeded for {self._name} pool.")
future = asyncio.futures.Future(loop=self._loop) if self._return_futures else None future = asyncio.futures.Future(loop=self._loop) if self._return_futures else None
await self._queue.put((future, args, kwargs)) await self._queue.put((future, args, kwargs))
self._total_queued += 1
if self._log_every_n is not None and (self._total_queued % self._log_every_n) == 0:
self._logger.info(f"pushed {self._total_queued}/{self._expected_total} items to {self._name} pool.")
self._logger.debug(f"'{self._name}' pool has received a new job. {args} {kwargs}") self._logger.debug(f"'{self._name}' pool has received a new job. {args} {kwargs}")
@ -150,13 +143,15 @@ class AsyncPool:
self._exceptions = False self._exceptions = False
self._workers = [] self._workers = []
for _ in range(self._num_workers): for worker_number in range(self._num_workers):
workerName = f"worker_{worker_number+1}"
self._workers.append( self._workers.append(
asyncio.ensure_future( asyncio.ensure_future(
coro_or_future=self._worker_loop(), coro_or_future=self._worker_loop(worker_id=workerName),
loop=self._loop loop=self._loop
) )
) )
self._active_worker[workerName] = None
async def join(self): async def join(self):
# no-op if workers aren't running # no-op if workers aren't running
@ -166,7 +161,7 @@ class AsyncPool:
self._logger.info(f'Joining {self._name}') self._logger.info(f'Joining {self._name}')
# The Terminators will kick each worker from being blocked against the _queue.get() and allow # The Terminators will kick each worker from being blocked against the _queue.get() and allow
# each one to exit # each one to exit.
for _ in range(self._num_workers): for _ in range(self._num_workers):
await self._queue.put(Terminator()) await self._queue.put(Terminator())
@ -183,5 +178,4 @@ class AsyncPool:
raise Exception(f"Exception occurred in {self._name} pool") raise Exception(f"Exception occurred in {self._name} pool")
def _time(self): def _time(self):
# utcnow returns a naive datetime, so we have to set the timezone manually <sigh> return datetime.now(tz=timezone.utc)
return datetime.utcnow().replace(tzinfo=timezone.utc)

59
app/DownloadQueue.py
View file

@ -25,6 +25,7 @@ class DownloadQueue:
queue: DataStore = None queue: DataStore = None
done: DataStore = None done: DataStore = None
event: asyncio.Event = None event: asyncio.Event = None
pool: AsyncPool = None
def __init__(self, notifier: Notifier, connection: Connection, serializer: ObjectSerializer): def __init__(self, notifier: Notifier, connection: Connection, serializer: ObjectSerializer):
self.config = Config.get_instance() self.config = Config.get_instance()
@ -122,6 +123,9 @@ class DownloadQueue:
is_manifestless = entry.get('is_manifestless', False) is_manifestless = entry.get('is_manifestless', False)
options.update({'is_manifestless': is_manifestless}) options.update({'is_manifestless': is_manifestless})
live_status: list = ['is_live', 'is_upcoming']
is_live = entry.get('is_live', None) or live_in or entry.get('live_status', None) in live_status
dl = ItemDTO( dl = ItemDTO(
id=entry.get('id'), id=entry.get('id'),
title=entry.get('title'), title=entry.get('title'),
@ -134,7 +138,7 @@ class DownloadQueue:
output_template=output_template if output_template else self.config.output_template, output_template=output_template if output_template else self.config.output_template,
datetime=formatdate(time.time()), datetime=formatdate(time.time()),
error=error, error=error,
is_live=entry.get('is_live', None) or entry.get('live_status', None) in ['is_live', 'is_upcoming'] or live_in, is_live=is_live,
live_in=live_in, live_in=live_in,
options=options options=options
) )
@ -161,7 +165,7 @@ class DownloadQueue:
debug=bool(self.config.ytdl_debug) debug=bool(self.config.ytdl_debug)
) )
if dlInfo.info.live_in or 'is_upcoming' == entry.get('live_status',None): if dlInfo.info.live_in or 'is_upcoming' == entry.get('live_status', None):
dlInfo.info.status = 'not_live' dlInfo.info.status = 'not_live'
itemDownload = self.done.put(dlInfo) itemDownload = self.done.put(dlInfo)
NotifyEvent = 'completed' NotifyEvent = 'completed'
@ -328,42 +332,45 @@ class DownloadQueue:
return items return items
async def __download_pool(self): async def __download_pool(self):
async with AsyncPool( self.pool = AsyncPool(
loop=asyncio.get_running_loop(), loop=asyncio.get_running_loop(),
num_workers=self.config.max_workers, num_workers=self.config.max_workers,
worker_co=self.__downloadFile, worker_co=self.__downloadFile,
name='download_pool', name='download_pool',
logger=logging.getLogger('WorkerPool'), logger=logging.getLogger('WorkerPool'),
) as executor: )
lastLog = time.time()
self.pool.start()
lastLog = time.time()
while True:
while True: while True:
while True: if self.pool.has_open_workers() is True:
if executor.has_open_workers() is True: break
break if time.time() - lastLog > 120:
if time.time() - lastLog > 120: lastLog = time.time()
lastLog = time.time() LOG.info(f'Waiting for worker to be free. {self.pool.get_workers_status()}')
LOG.info(f'Waiting for worker to be free.') await asyncio.sleep(1)
await asyncio.sleep(1)
while not self.queue.hasDownloads(): while not self.queue.hasDownloads():
LOG.info(f"Waiting for item to download. '{executor.get_available_workers()}' free workers.") LOG.info(f"Waiting for item to download. '{self.pool.get_available_workers()}' free workers.")
await self.event.wait() await self.event.wait()
self.event.clear() self.event.clear()
LOG.debug(f"Cleared wait event.") LOG.debug(f"Cleared wait event.")
entry = self.queue.getNextDownload() entry = self.queue.getNextDownload()
await asyncio.sleep(0.2) await asyncio.sleep(0.2)
if entry is None: if entry is None:
continue continue
LOG.debug(f"Pushing {entry=} to executor.") LOG.debug(f"Pushing {entry=} to executor.")
if entry.started() is False and entry.is_canceled() is False: if entry.started() is False and entry.is_canceled() is False:
await executor.push(id=entry.info._id, entry=entry) await self.pool.push(id=entry.info._id, entry=entry)
LOG.debug(f"Pushed {entry=} to executor.") LOG.debug(f"Pushed {entry=} to executor.")
await asyncio.sleep(1) await asyncio.sleep(1)
async def __download(self): async def __download(self):
while True: while True:

4
app/Utils.py
View file

@ -293,7 +293,9 @@ class ObjectSerializer(json.JSONEncoder):
""" """
def default(self, obj): def default(self, obj):
return obj.__dict__ if isinstance(obj, object) else json.JSONEncoder.default(self, obj) if isinstance(obj, object) and hasattr(obj, '__dict__'):
return obj.__dict__
return json.JSONEncoder.default(self, obj)
class Notifier: class Notifier:

26
app/main.py
View file

@ -11,6 +11,7 @@ from Utils import ObjectSerializer, Notifier
from aiohttp import web, client from aiohttp import web, client
from aiohttp.web import Request, Response from aiohttp.web import Request, Response
from Webhooks import Webhooks from Webhooks import Webhooks
from Download import Download
from player.M3u8 import M3u8 from player.M3u8 import M3u8
from player.Segments import Segments from player.Segments import Segments
import socketio import socketio
@ -342,6 +343,31 @@ class Main:
return web.Response(text=self.serializer.encode(history)) return web.Response(text=self.serializer.encode(history))
@self.routes.get(self.config.url_prefix + 'workers')
async def workers(_) -> Response:
if self.queue.pool is None:
return web.json_response({"error": "Worker pool not initialized."}, status=404)
status = self.queue.pool.get_workers_status()
data = []
for worker in status:
worker_status = status.get(worker)
data.append({
'id': worker,
'data': {"status": 'Waiting for download.'} if worker_status is None else worker_status,
})
def safe_serialize(obj):
def default(o): return f"<<non-serializable: {type(o).__qualname__}>>"
return json.dumps(obj, default=default)
return web.Response(text=safe_serialize(data), headers={
'Content-Type': 'application/json',
})
@self.routes.get(self.config.url_prefix + 'm3u8/{file:.*}') @self.routes.get(self.config.url_prefix + 'm3u8/{file:.*}')
async def m3u8(request: Request) -> Response: async def m3u8(request: Request) -> Response:
file: str = request.match_info.get('file') file: str = request.match_info.get('file')