use thread for extract_info

This commit is contained in:
ArabCoders 2024-02-23 18:40:52 +03:00
parent 9dd4063ff4
commit 966e847bbd
2 changed files with 58 additions and 44 deletions

View file

@ -16,21 +16,23 @@ class AsyncPool:
number of items of work is in the queue). `worker_co` will be called
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 job_accept_duration: maximum number of seconds from first push to last push before a TimeoutError will be thrown.
: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 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.
@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
: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
@ -117,25 +119,26 @@ class AsyncPool:
await self.join()
async def push(self, *args, **kwargs) -> asyncio.Future:
""" Method to push work to `worker_co` passed to `__init__`.
"""
Method to push work to `worker_co` passed initially to `__init__`.
:param args: position arguments to be passed to `worker_co`
:param kwargs: keyword arguments to be passed to `worker_co`
: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("Maximum lifetime of {} seconds of AsyncWorkerPool: {} exceeded".format(
self._job_accept_duration, self._name))
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))
self._total_queued += 1
if self._log_every_n is not None and (self._total_queued % self._log_every_n) == 0:
self._logger.info("pushed {}/{} items to {} AsyncWorkerPool".format(
self._total_queued, self._expected_total, self._name))
self._logger.info(f"pushed {self._total_queued}/{self._expected_total} items to {self._name} pool.")
return future
@ -144,15 +147,22 @@ class AsyncPool:
assert self._workers is None
self._exceptions = False
self._workers = [asyncio.ensure_future(
self._worker_loop(), loop=self._loop) for _ in range(self._num_workers)]
self._workers = []
for _ in range(self._num_workers):
self._workers.append(
asyncio.ensure_future(
coro_or_future=self._worker_loop(),
loop=self._loop
)
)
async def join(self):
# no-op if workers aren't running
if not self._workers:
return
self._logger.info('Joining {}'.format(self._name))
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
for _ in range(self._num_workers):
@ -162,13 +172,13 @@ class AsyncPool:
await asyncio.gather(*self._workers)
self._workers = None
except:
self._logger.exception('Exception joining {}'.format(self._name))
self._logger.exception(f'Exception joining {self._name}')
raise
finally:
self._logger.info('Completed {}'.format(self._name))
self._logger.info(f'Completed {self._name}')
if self._exceptions and self._raise_on_join:
raise Exception("Exception occurred in pool {}".format(self._name))
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>

View file

@ -11,7 +11,7 @@ from ItemDTO import ItemDTO
from DataStore import DataStore
from Utils import Notifier, ObjectSerializer, calcDownloadPath, ExtractInfo, isDownloaded, mergeConfig
from AsyncPool import AsyncPool
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
LOG = logging.getLogger('DownloadQueue')
TYPE_DONE: str = 'done'
TYPE_QUEUE: str = 'queue'
@ -217,18 +217,26 @@ class DownloadQueue:
else:
already.add(url)
try:
entry = await asyncio.get_running_loop().run_in_executor(
None,
ExtractInfo,
mergeConfig(self.config.ytdl_options, ytdlp_config),
url,
bool(self.config.ytdl_debug)
)
with ThreadPoolExecutor(thread_name_prefix='extract_info') as pool:
LOG.debug(f'extracting info from {url=}')
entry = await asyncio.get_running_loop().run_in_executor(
pool,
ExtractInfo,
mergeConfig(self.config.ytdl_options, ytdlp_config),
url,
bool(self.config.ytdl_debug)
)
if not entry:
if self.config.keep_archive:
if not entry:
if not self.config.keep_archive:
return {
'status': 'error',
'msg': 'No metadata, most likely video has been downloaded before.' if self.config.keep_archive else 'Unable to extract info check logs.'
}
LOG.debug(f'No metadata, Rechecking with archive disabled. {url=}')
entry = await asyncio.get_running_loop().run_in_executor(
None,
pool,
ExtractInfo,
mergeConfig(self.config.ytdl_options, ytdlp_config),
url,
@ -246,11 +254,6 @@ class DownloadQueue:
'msg': f'[{entry.get("id")}: {entry.get("title")}]: has been downloaded already.'
}
return {
'status': 'error',
'msg': 'No metadata, most likely video has been downloaded before.' if self.config.keep_archive else 'Unable to extract info check logs.'
}
if self.isDownloaded(entry):
raise yt_dlp.utils.ExistingVideoReached()
@ -331,12 +334,13 @@ class DownloadQueue:
loop=asyncio.get_running_loop(),
num_workers=self.config.max_workers,
worker_co=self.__downloadFile,
name='WorkerPool',
name='download_pool',
logger=logging.getLogger('WorkerPool'),
) as executor:
while True:
while True:
if executor.has_open_workers() is False:
LOG.debug(f'Waiting for workers to be available.')
await asyncio.sleep(1)
else:
break