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

Added endpoint to expose worker status

Browse source
This commit is contained in:
ArabCoders 2024-03-20 23:13:05 +03:00
parent cccb8efb68
commit 5bba32cd06
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
from datetime import datetime, timezone
import logging
class Terminator:
@ -7,9 +8,13 @@ class Terminator:
class AsyncPool:
def __init__(self, loop, num_workers: int, name: str, logger, worker_co, load_factor: int = 1,
job_accept_duration: int = None, max_task_time: int = None, return_futures: bool = False,
raise_on_join: bool = False, log_every_n: int = None, expected_total=None):
def __init__(self,
num_workers: int, worker_co,
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
`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 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 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.
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
: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
"""
loop = loop if loop else asyncio.get_event_loop()
self._loop = loop
self._loop = loop if loop else asyncio.get_event_loop()
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._workers = None
self._exceptions = False
self._job_accept_duration = job_accept_duration
self._first_push_dt = None
self._max_task_time = max_task_time
self._return_futures = return_futures
self._raise_on_join = raise_on_join
self._name = name
self._worker_co = worker_co
self._total_queued = 0
self._log_every_n = log_every_n
self._expected_total = expected_total
self._active_workers = 0
self._active_worker: dict[str, dict | None] = {}
async def _worker_loop(self):
async def _worker_loop(self, worker_id: str):
while True:
got_obj = False
future = None
@ -65,11 +60,16 @@ class AsyncPool:
if item.__class__ is Terminator:
break
self._active_workers += 1
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
# 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:
future.set_result(result)
@ -85,9 +85,10 @@ class AsyncPool:
# don't log the failure when the client is receiving the future
future.set_exception(e)
else:
self._logger.exception('Worker call failed')
self._logger.exception(f'Worker call failed. {str(e)}')
finally:
self._active_workers -= 1
self._active_worker[worker_id] = None
if got_obj:
self._queue.task_done()
@ -95,21 +96,23 @@ class AsyncPool:
def exceptions(self):
return self._exceptions
@property
def total_queued(self):
return self._total_queued
def has_open_workers(self) -> bool:
"""
: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:
"""
: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):
self.start()
@ -127,18 +130,8 @@ class AsyncPool:
: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
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}")
@ -150,13 +143,15 @@ class AsyncPool:
self._exceptions = False
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(
asyncio.ensure_future(
coro_or_future=self._worker_loop(),
coro_or_future=self._worker_loop(worker_id=workerName),
loop=self._loop
)
)
self._active_worker[workerName] = None
async def join(self):
# no-op if workers aren't running
@ -166,7 +161,7 @@ class AsyncPool:
self._logger.info(f'Joining {self._name}')
# 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):
await self._queue.put(Terminator())
@ -183,5 +178,4 @@ class AsyncPool:
raise Exception(f"Exception occurred in {self._name} pool")
def _time(self):
# utcnow returns a naive datetime, so we have to set the timezone manually <sigh>
return datetime.utcnow().replace(tzinfo=timezone.utc)
return datetime.now(tz=timezone.utc)

59
app/DownloadQueue.py
View file

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

4
app/Utils.py
View file

@ -293,7 +293,9 @@ class ObjectSerializer(json.JSONEncoder):
"""
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:

26
app/main.py
View file

@ -11,6 +11,7 @@ from Utils import ObjectSerializer, Notifier
from aiohttp import web, client
from aiohttp.web import Request, Response
from Webhooks import Webhooks
from Download import Download
from player.M3u8 import M3u8
from player.Segments import Segments
import socketio
@ -342,6 +343,31 @@ class Main:
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:.*}')
async def m3u8(request: Request) -> Response:
file: str = request.match_info.get('file')