GH-63: Add an option to force the number of workers for concurrent processing (GH-64)

Author: ArtyomVancyan

src/thumbnails/__init__.py

@@ -22,7 +22,7 @@
 from .thumbnail import ThumbnailVTT
 from .thumbnail import register_thumbnail
 
-__version__ = "0.1.8"
+__version__ = "0.1.9"
 __all__ = (
     "Generator",
     "Thumbnail",

src/thumbnails/cli.py

@@ -10,6 +10,7 @@
 from .constants import DEFAULT_INTERVAL
 from .constants import DEFAULT_OUTPUT
 from .constants import DEFAULT_SKIP
+from .constants import DEFAULT_WORKERS
 
 # Help messages of the particular option of the CLI.
 HELP_BASE = "The prefix of the thumbnails path can be customized."
@@ -18,6 +19,7 @@
 HELP_FORMAT = "Output format. Default is %s." % DEFAULT_FORMAT
 HELP_COMPRESS = "The image scale coefficient. A number from 0 to 1."
 HELP_INTERVAL = "The interval between neighbor thumbnails in seconds."
+HELP_WORKERS = "Workers number for concurrent processing. Default is calculated automatically."
 
 # This defines a choice of supported values for the '--format' option of the CLI.
 format_choice = click.Choice(ThumbnailFactory.thumbnails.keys(), case_sensitive=False)
@@ -27,6 +29,7 @@ def cli(func):
     @click.command()
     @click.option("--compress", "-C", default=DEFAULT_COMPRESS, help=HELP_COMPRESS)
     @click.option("--interval", "-I", default=DEFAULT_INTERVAL, help=HELP_INTERVAL)
+    @click.option("--workers", "-W", default=DEFAULT_WORKERS, help=HELP_WORKERS)
     @click.option("--base", "-B", default=DEFAULT_BASE, help=HELP_BASE)
     @click.option("--skip", "-S", default=DEFAULT_SKIP, help=HELP_SKIP, is_flag=True)
     @click.option("--output", "-O", default=DEFAULT_OUTPUT, type=click.Path(), help=HELP_OUTPUT)

src/thumbnails/constants.py

@@ -4,3 +4,4 @@
 DEFAULT_FORMAT = "vtt"
 DEFAULT_COMPRESS = 1.0
 DEFAULT_INTERVAL = 1.0
+DEFAULT_WORKERS = 0  # 0 is a flag for calculating the number of workers automatically.

src/thumbnails/generator.py

@@ -10,6 +10,7 @@
 from .constants import DEFAULT_INTERVAL
 from .constants import DEFAULT_OUTPUT
 from .constants import DEFAULT_SKIP
+from .constants import DEFAULT_WORKERS
 from .pathtools import listdir
 from .progress import use_progress
 from .thumbnail import ThumbnailExistsError
@@ -34,10 +35,24 @@ def __init__(self, inputs):
         self.format = DEFAULT_FORMAT
         self.compress = DEFAULT_COMPRESS
         self.interval = DEFAULT_INTERVAL
+        self._workers = DEFAULT_WORKERS
 
         # Remove non-video files in case of input directory already contains other generated files.
         self.inputs = [file for file in self.inputs if re.match(r"^.*\.(?:(?!png|vtt|json).)+$", file)]
 
+    @property
+    def workers(self):
+        """Returns the number of workers for concurrent processing."""
+        if self._workers > 0:
+            return self._workers
+        # Limit the auto-calculated workers for super-multicore machines.
+        return min(32, len(self.inputs), (os.cpu_count() or 1) + 4)
+
+    @workers.setter
+    def workers(self, value):
+        """Sets the number of workers for concurrent processing."""
+        self._workers = value
+
     @staticmethod
     def worker(video, fmt, base, skip, output):
         """Executes the required workflows for generating a thumbnail."""
@@ -60,7 +75,7 @@ def __next__(self):
 
     @use_progress
     def generate(self):
-        with concurrent.futures.ThreadPoolExecutor() as executor:
+        with concurrent.futures.ThreadPoolExecutor(max_workers=self.workers) as executor:
             executor.map(
                 functools.partial(
                     self.worker,