From 4d0123c8b0705c02c1e6cdac25a8a2a7d0178c3e Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Fri, 5 Jan 2024 13:04:01 -0800 Subject: [PATCH] [BACKEND] Optimize memory usage (#872) Reduce memory usage significantly when pulling large channels all at once by - Not loading entry script until it is being processed - Deleting the entry memory once fully processed - Disable yt-dlp from keeping entry info in memory since ytdl-sub reads it from file --- src/ytdl_sub/config/overrides.py | 2 +- src/ytdl_sub/downloaders/url/downloader.py | 40 ++++++++++------- src/ytdl_sub/entries/base_entry.py | 17 ------- src/ytdl_sub/entries/entry.py | 18 ++++---- .../subscription_ytdl_options.py | 1 + src/ytdl_sub/utils/scriptable.py | 45 +++++++++++++++---- tools/docgen/plugins.py | 2 + 7 files changed, 74 insertions(+), 51 deletions(-) diff --git a/src/ytdl_sub/config/overrides.py b/src/ytdl_sub/config/overrides.py index f68ce727..3372e936 100644 --- a/src/ytdl_sub/config/overrides.py +++ b/src/ytdl_sub/config/overrides.py @@ -56,7 +56,7 @@ class Overrides(DictFormatterValidator, Scriptable): def __init__(self, name, value): DictFormatterValidator.__init__(self, name, value) - Scriptable.__init__(self) + Scriptable.__init__(self, initialize_base_script=True) for key in self._keys: self.ensure_variable_name_valid(key) diff --git a/src/ytdl_sub/downloaders/url/downloader.py b/src/ytdl_sub/downloaders/url/downloader.py index 67feabd0..b1be016b 100644 --- a/src/ytdl_sub/downloaders/url/downloader.py +++ b/src/ytdl_sub/downloaders/url/downloader.py @@ -349,39 +349,43 @@ class MultiUrlDownloader(SourcePlugin[MultiUrlValidator]): ) def _iterate_child_entries( - self, url_validator: UrlValidator, entries: List[Entry] + self, entries: List[Entry], download_reversed: bool ) -> Iterator[Entry]: - entries_to_iterate = entries - if url_validator.download_reverse: - entries_to_iterate = reversed(entries) + # Iterate a list of entries, and delete the entries after yielding + indices = list(range(len(entries))) + if download_reversed: + indices = reversed(indices) - for entry in entries_to_iterate: + for idx in indices: self._url_state.entries_downloaded += 1 - if self._is_downloaded(entry): + if self._is_downloaded(entries[idx]): download_logger.info( "Already downloaded entry %d/%d: %s", self._url_state.entries_downloaded, self._url_state.entries_total, - entry.title, + entries[idx].title, ) + del entries[idx] continue - yield entry - self._mark_downloaded(entry) + yield entries[idx] + self._mark_downloaded(entries[idx]) + + del entries[idx] def _iterate_parent_entry( - self, url_validator: UrlValidator, parent: EntryParent + self, parent: EntryParent, download_reversed: bool ) -> Iterator[Entry]: for entry_child in self._iterate_child_entries( - url_validator=url_validator, entries=parent.entry_children() + entries=parent.entry_children(), download_reversed=download_reversed ): yield entry_child # Recursion the parent's parent entries for parent_child in reversed(parent.parent_children()): for entry_child in self._iterate_parent_entry( - url_validator=url_validator, parent=parent_child + parent=parent_child, download_reversed=download_reversed ): yield entry_child @@ -415,9 +419,9 @@ class MultiUrlDownloader(SourcePlugin[MultiUrlValidator]): def _iterate_entries( self, - url_validator: UrlValidator, parents: List[EntryParent], orphans: List[Entry], + download_reversed: bool, ) -> Iterator[Entry]: """ Downloads the leaf entries from EntryParent trees @@ -426,11 +430,13 @@ class MultiUrlDownloader(SourcePlugin[MultiUrlValidator]): with self._separate_download_archives(clear_info_json_files=True): for parent in parents: for entry_child in self._iterate_parent_entry( - url_validator=url_validator, parent=parent + parent=parent, download_reversed=download_reversed ): yield entry_child - for orphan in self._iterate_child_entries(url_validator=url_validator, entries=orphans): + for orphan in self._iterate_child_entries( + entries=orphans, download_reversed=download_reversed + ): yield orphan def download_metadata(self) -> Iterable[Entry]: @@ -454,7 +460,9 @@ class MultiUrlDownloader(SourcePlugin[MultiUrlValidator]): "Beginning downloads for %s", self.overrides.apply_formatter(collection_url.url) ) for entry in self._iterate_entries( - url_validator=collection_url, parents=parents, orphans=orphan_entries + parents=parents, + orphans=orphan_entries, + download_reversed=collection_url.download_reverse, ): entry.initialize_script(self.overrides).add( {v.ytdl_sub_input_url: self.overrides.apply_formatter(collection_url.url)} diff --git a/src/ytdl_sub/entries/base_entry.py b/src/ytdl_sub/entries/base_entry.py index 6b056900..1d8acd3e 100644 --- a/src/ytdl_sub/entries/base_entry.py +++ b/src/ytdl_sub/entries/base_entry.py @@ -116,23 +116,6 @@ class BaseEntry(ABC): """ return self._working_directory - def add_kwargs(self, variables_to_add: Dict[str, Any]) -> "BaseEntry": - """ - Adds variables to kwargs. Use with caution since yt-dlp data can be overwritten. - Plugins should use ``add_variables``. - - Parameters - ---------- - variables_to_add - Variables to add to kwargs - - Returns - ------- - self - """ - self._kwargs = dict(self._kwargs, **variables_to_add) - return self - def get_download_info_json_name(self) -> str: """ Returns diff --git a/src/ytdl_sub/entries/entry.py b/src/ytdl_sub/entries/entry.py index 1d684c5b..831bda6c 100644 --- a/src/ytdl_sub/entries/entry.py +++ b/src/ytdl_sub/entries/entry.py @@ -44,12 +44,6 @@ class Entry(BaseEntry, Scriptable): BaseEntry.__init__(self, entry_dict=entry_dict, working_directory=working_directory) Scriptable.__init__(self) - def _add_entry_kwargs_to_script(self) -> None: - # Add entry metadata, but avoid the `.add()` helper since it also adds sanitized - self.unresolvable.remove(v.entry_metadata.variable_name) - self.script.add({v.entry_metadata.variable_name: ScriptUtils.to_script(self._kwargs)}) - self.update_script() - def initialize_script(self, other: Optional[Scriptable] = None) -> "Entry": """ Initializes the entry script using the Overrides script, then adding @@ -57,12 +51,20 @@ class Entry(BaseEntry, Scriptable): """ # Overrides contains added variables that are unresolvable, add them here if other: - self.script = copy.deepcopy(other.script) - self.unresolvable = copy.deepcopy(other.unresolvable) + self._script = copy.deepcopy(other.script) + self._unresolvable = copy.deepcopy(other.unresolvable) + else: + self.initialize_base_script() self._add_entry_kwargs_to_script() return self + def _add_entry_kwargs_to_script(self) -> None: + # Add entry metadata, but avoid the `.add()` helper since it also adds sanitized + self.unresolvable.remove(v.entry_metadata.variable_name) + self.script.add({v.entry_metadata.variable_name: ScriptUtils.to_script(self._kwargs)}) + self.update_script() + def get(self, variable: Variable, expected_type: Type[TypeT]) -> TypeT: """ Gets a variable of an expected type. Will error if it does not exist or is not resolved. diff --git a/src/ytdl_sub/subscriptions/subscription_ytdl_options.py b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py index c92ea028..77ec1c6b 100644 --- a/src/ytdl_sub/subscriptions/subscription_ytdl_options.py +++ b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py @@ -81,6 +81,7 @@ class SubscriptionYTDLOptions: "skip_download": True, "writethumbnail": False, "writeinfojson": True, + "extract_flat": "discard", # do not store info.json in mem since its in file } @property diff --git a/src/ytdl_sub/utils/scriptable.py b/src/ytdl_sub/utils/scriptable.py index aba9d1fc..b1dc20e7 100644 --- a/src/ytdl_sub/utils/scriptable.py +++ b/src/ytdl_sub/utils/scriptable.py @@ -2,6 +2,7 @@ import copy from abc import ABC from typing import Any from typing import Dict +from typing import Optional from typing import Set from ytdl_sub.entries.script.function_scripts import CUSTOM_FUNCTION_SCRIPTS @@ -13,21 +14,47 @@ from ytdl_sub.script.utils.exceptions import RuntimeException from ytdl_sub.utils.exceptions import StringFormattingException from ytdl_sub.utils.script import ScriptUtils +_BASE_SCRIPT: Script = Script( + ScriptUtils.add_sanitized_variables( + dict(copy.deepcopy(VARIABLE_SCRIPTS), **copy.deepcopy(CUSTOM_FUNCTION_SCRIPTS)) + ) +) + class Scriptable(ABC): """ Shared class between Entry and Overrides to manage their underlying Script. """ - _BASE_SCRIPT: Script = Script( - ScriptUtils.add_sanitized_variables( - dict(copy.deepcopy(VARIABLE_SCRIPTS), **copy.deepcopy(CUSTOM_FUNCTION_SCRIPTS)) - ) - ) + def __init__(self, initialize_base_script: bool = False): + self._script: Optional[Script] = None + self._unresolvable: Optional[Set[str]] = None - def __init__(self): - self.script = copy.deepcopy(Scriptable._BASE_SCRIPT) - self.unresolvable: Set[str] = copy.deepcopy(UNRESOLVED_VARIABLES) + if initialize_base_script: + self.initialize_base_script() + + def initialize_base_script(self): + """ + Initializes with base values + """ + self._script = copy.deepcopy(_BASE_SCRIPT) + self._unresolvable = copy.deepcopy(UNRESOLVED_VARIABLES) + + @property + def script(self) -> Script: + """ + Initialized script + """ + assert self._script is not None, "Not initialized" + return self._script + + @property + def unresolvable(self) -> Set[str]: + """ + Initialized unresolvable variables + """ + assert self._unresolvable is not None, "Not initialized" + return self._unresolvable def update_script(self) -> None: """ @@ -45,7 +72,7 @@ class Scriptable(ABC): for var, definition in values.items() } - self.unresolvable -= set(list(values_as_str.keys())) + self._unresolvable -= set(list(values_as_str.keys())) self.script.add( ScriptUtils.add_sanitized_variables( { diff --git a/tools/docgen/plugins.py b/tools/docgen/plugins.py index 10323066..8080c297 100644 --- a/tools/docgen/plugins.py +++ b/tools/docgen/plugins.py @@ -39,6 +39,8 @@ def should_filter_property(property_name: str) -> bool: "dict_with_format_strings", "subscription_name", "list", + "script", + "unresolvable", )