Parallel execution uses thread pool to avoid leaking resources. (#1367)

Co-authored-by: Michael R. Crusoe <[email protected]>

Author: tetron

cwltool/command_line_tool.py

@@ -1145,6 +1145,7 @@ def collect_output(
         r = []  # type: List[CWLOutputType]
         empty_and_optional = False
         debug = _logger.isEnabledFor(logging.DEBUG)
+        result: Optional[CWLOutputType] = None
         if "outputBinding" in schema:
             binding = cast(
                 MutableMapping[str, Union[bool, str, List[str]]],

cwltool/executors.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 """Single and multi-threaded executors."""
 import datetime
+import functools
 import logging
 import math
 import os
@@ -34,6 +35,8 @@
 from .utils import CWLObjectType, JobsType
 from .workflow import Workflow
 from .workflow_job import WorkflowJob, WorkflowJobStep
+from .task_queue import TaskQueue
+
 
 TMPDIR_LOCK = Lock()
 
@@ -277,7 +280,6 @@ class MultithreadedJobExecutor(JobExecutor):
     def __init__(self) -> None:
         """Initialize."""
         super(MultithreadedJobExecutor, self).__init__()
-        self.threads = set()  # type: Set[threading.Thread]
         self.exceptions = []  # type: List[WorkflowException]
         self.pending_jobs = []  # type: List[JobsType]
         self.pending_jobs_lock = threading.Lock()
@@ -339,7 +341,6 @@ def _runner(self, job, runtime_context, TMPDIR_LOCK):
         finally:
             if runtime_context.workflow_eval_lock:
                 with runtime_context.workflow_eval_lock:
-                    self.threads.remove(threading.current_thread())
                     if isinstance(job, JobBase):
                         ram = job.builder.resources["ram"]
                         if not isinstance(ram, str):
@@ -362,6 +363,10 @@ def run_job(
         with self.pending_jobs_lock:
             n = 0
             while (n + 1) <= len(self.pending_jobs):
+                # Simple greedy resource allocation strategy.  Go
+                # through pending jobs in the order they were
+                # generated and add them to the queue only if there
+                # are resources available.
                 job = self.pending_jobs[n]
                 if isinstance(job, JobBase):
                     ram = job.builder.resources["ram"]
@@ -403,26 +408,24 @@ def run_job(
                         n += 1
                         continue
 
-                thread = threading.Thread(
-                    target=self._runner, args=(job, runtime_context, TMPDIR_LOCK)
-                )
-                thread.daemon = True
-                self.threads.add(thread)
                 if isinstance(job, JobBase):
                     ram = job.builder.resources["ram"]
                     if not isinstance(ram, str):
                         self.allocated_ram += ram
                     cores = job.builder.resources["cores"]
                     if not isinstance(cores, str):
                         self.allocated_cores += cores
-                thread.start()
+                self.taskqueue.add(
+                    functools.partial(self._runner, job, runtime_context, TMPDIR_LOCK),
+                    runtime_context.workflow_eval_lock,
+                )
                 self.pending_jobs.remove(job)
 
     def wait_for_next_completion(self, runtime_context):
         # type: (RuntimeContext) -> None
         """Wait for jobs to finish."""
         if runtime_context.workflow_eval_lock is not None:
-            runtime_context.workflow_eval_lock.wait()
+            runtime_context.workflow_eval_lock.wait(timeout=3)
         if self.exceptions:
             raise self.exceptions[0]
 
@@ -434,36 +437,46 @@ def run_jobs(
         runtime_context: RuntimeContext,
     ) -> None:
 
-        jobiter = process.job(job_order_object, self.output_callback, runtime_context)
+        self.taskqueue = TaskQueue(
+            threading.Lock(), psutil.cpu_count()
+        )  # type: TaskQueue
+        try:
 
-        if runtime_context.workflow_eval_lock is None:
-            raise WorkflowException(
-                "runtimeContext.workflow_eval_lock must not be None"
+            jobiter = process.job(
+                job_order_object, self.output_callback, runtime_context
             )
 
-        runtime_context.workflow_eval_lock.acquire()
-        for job in jobiter:
-            if job is not None:
-                if isinstance(job, JobBase):
-                    job.builder = runtime_context.builder or job.builder
-                    if job.outdir is not None:
-                        self.output_dirs.add(job.outdir)
+            if runtime_context.workflow_eval_lock is None:
+                raise WorkflowException(
+                    "runtimeContext.workflow_eval_lock must not be None"
+                )
 
-            self.run_job(job, runtime_context)
+            runtime_context.workflow_eval_lock.acquire()
+            for job in jobiter:
+                if job is not None:
+                    if isinstance(job, JobBase):
+                        job.builder = runtime_context.builder or job.builder
+                        if job.outdir is not None:
+                            self.output_dirs.add(job.outdir)
 
-            if job is None:
-                if self.threads:
-                    self.wait_for_next_completion(runtime_context)
-                else:
-                    logger.error("Workflow cannot make any more progress.")
-                    break
+                self.run_job(job, runtime_context)
+
+                if job is None:
+                    if self.taskqueue.in_flight > 0:
+                        self.wait_for_next_completion(runtime_context)
+                    else:
+                        logger.error("Workflow cannot make any more progress.")
+                        break
 
-        self.run_job(None, runtime_context)
-        while self.threads:
-            self.wait_for_next_completion(runtime_context)
             self.run_job(None, runtime_context)
+            while self.taskqueue.in_flight > 0:
+                self.wait_for_next_completion(runtime_context)
+                self.run_job(None, runtime_context)
 
-        runtime_context.workflow_eval_lock.release()
+            runtime_context.workflow_eval_lock.release()
+        finally:
+            self.taskqueue.drain()
+            self.taskqueue.join()
 
 
 class NoopJobExecutor(JobExecutor):

cwltool/task_queue.py

@@ -0,0 +1,126 @@
+# Copyright (C) The Arvados Authors. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+"""TaskQueue."""
+
+import queue
+import threading
+
+from typing import Callable, Optional
+
+from .loghandler import _logger
+
+
+class TaskQueue(object):
+    """A TaskQueue class.
+
+    Uses a first-in, first-out queue of tasks executed on a fixed number of
+    threads.
+
+    New tasks enter the queue and are started in the order received,
+    as worker threads become available.
+
+    If thread_count == 0 then tasks will be synchronously executed
+    when add() is called (this makes the actual task queue behavior a
+    no-op, but may be a useful configuration knob).
+
+    The thread_count is also used as the maximum size of the queue.
+
+    The threads are created during TaskQueue initialization.  Call
+    join() when you're done with the TaskQueue and want the threads to
+    stop.
+
+
+    Attributes
+    ----------
+    in_flight
+        the number of tasks in the queue
+
+    """
+
+    def __init__(self, lock: threading.Lock, thread_count: int):
+        """Create a new task queue using the specified lock and number of threads."""
+        self.thread_count = thread_count
+        self.task_queue: queue.Queue[Optional[Callable[[], None]]] = queue.Queue(
+            maxsize=self.thread_count
+        )
+        self.task_queue_threads = []
+        self.lock = lock
+        self.in_flight = 0
+        self.error: Optional[BaseException] = None
+
+        for _r in range(0, self.thread_count):
+            t = threading.Thread(target=self._task_queue_func)
+            self.task_queue_threads.append(t)
+            t.start()
+
+    def _task_queue_func(self) -> None:
+        while True:
+            task = self.task_queue.get()
+            if task is None:
+                return
+            try:
+                task()
+            except BaseException as e:
+                _logger.exception("Unhandled exception running task")
+                self.error = e
+            finally:
+                with self.lock:
+                    self.in_flight -= 1
+
+    def add(
+        self,
+        task: Callable[[], None],
+        unlock: Optional[threading.Condition] = None,
+        check_done: Optional[threading.Event] = None,
+    ) -> None:
+        """
+        Add your task to the queue.
+
+        The optional unlock will be released prior to attempting to add the
+        task to the queue.
+
+        If the optional "check_done" threading.Event's flag is set, then we
+        will skip adding this task to the queue.
+
+        If the TaskQueue was created with thread_count == 0 then your task will
+        be synchronously executed.
+
+        """
+        if self.thread_count == 0:
+            task()
+            return
+
+        with self.lock:
+            self.in_flight += 1
+
+        while True:
+            try:
+                if unlock is not None:
+                    unlock.release()
+                if check_done is not None and check_done.is_set():
+                    with self.lock:
+                        self.in_flight -= 1
+                    return
+                self.task_queue.put(task, block=True, timeout=3)
+                return
+            except queue.Full:
+                pass
+            finally:
+                if unlock is not None:
+                    unlock.acquire()
+
+    def drain(self) -> None:
+        """Drain the queue."""
+        try:
+            while not self.task_queue.empty():
+                self.task_queue.get(True, 0.1)
+        except queue.Empty:
+            pass
+
+    def join(self) -> None:
+        """Wait for all threads to complete."""
+        for _t in self.task_queue_threads:
+            self.task_queue.put(None)
+        for t in self.task_queue_threads:
+            t.join()