easyvvuq.actions.action_statuses

  1"""Implements ActionPool - a thin wrapper around the Python Executor interface
  2that is meant to simplify the execution of actions and retrieval of results.
  3This object is instantiated by the Campaign. The user would never instantiate it
  4manually. The user does interact with it to track the progress of execution.
  5"""
  6import concurrent
  7from concurrent.futures import ThreadPoolExecutor
  8from dask.distributed import Client
  9from tqdm import tqdm
 10import copy
 11
 12from . import QCGPJPool
 13
 14__copyright__ = """
 15
 16    Copyright 2020 Vytautas Jancauskas
 17
 18    This file is part of EasyVVUQ
 19
 20    EasyVVUQ is free software: you can redistribute it and/or modify
 21    it under the terms of the Lesser GNU General Public License as published by
 22    the Free Software Foundation, either version 3 of the License, or
 23    (at your option) any later version.
 24
 25    EasyVVUQ is distributed in the hope that it will be useful,
 26    but WITHOUT ANY WARRANTY; without even the implied warranty of
 27    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 28    Lesser GNU General Public License for more details.
 29
 30    You should have received a copy of the Lesser GNU General Public License
 31    along with this program.  If not, see <https://www.gnu.org/licenses/>.
 32
 33"""
 34__license__ = "LGPL"
 35
 36
 37class ActionPool:
 38    """A class that handles the execution of Actions.
 39
 40    Parameters
 41    ----------
 42    campaign: Campaign
 43        An instance of an EasyVVUQ campaign.
 44    actions: Actions
 45        An instance of `Actions` containing things to be done as part of the simulation.
 46    inits: iterable
 47        Initial inputs to be passed to each `Actions` representing a sample. Will usually contain
 48        dictionaries with the following information: {'run_id': ..., 'campaign_dir': ...,
 49        'run_info': ...}.
 50    sequential: bool
 51        Will run the actions sequentially.
 52    """
 53
 54    def __init__(self, campaign, actions, inits, sequential=False):
 55        self.campaign = campaign
 56        self.actions = actions
 57        self.inits = inits
 58        self.sequential = sequential
 59        self.futures = []
 60        self.results = []
 61        self._collate_callback = lambda previous: previous
 62
 63    def start(self, pool=None):
 64        """Start the actions.
 65
 66        Parameters
 67        ----------
 68        pool: An Executor instance (e.g. ThreadPoolExecutor)
 69
 70        Returns
 71        -------
 72        ActionPool
 73            Starts execution and returns a reference to itself for tracking progress
 74            and for collation.
 75        """
 76        if pool is None:
 77            pool = ThreadPoolExecutor()
 78        self.pool = pool
 79        for previous in self.inits:
 80            previous = copy.copy(previous)
 81            if self.sequential:
 82                result = self.actions.start(previous)
 83                self.results.append(result)
 84            else:
 85                future = self.pool.submit(self.actions.start, previous)
 86                self.futures.append(future)
 87        return self
 88
 89    def progress(self):
 90        """Some basic stats about the action statuses status.
 91
 92        Returns
 93        -------
 94        dict
 95            A dictionary with four keys - 'ready', 'active' and 'finished', 'failed'.
 96            Values under "ready" correspond to `Actions` waiting for execution, "active"
 97            corresponds to the number of currently running tasks.
 98        """
 99        ready = 0
100        running = 0
101        done = 0
102        failed = 0
103        for future in self.futures:
104            if future.running():
105                running += 1
106            elif future.done():
107                if not future.result():
108                    failed += 1
109                else:
110                    done += 1
111            else:
112                ready += 1
113        return {'ready': ready, 'active': running, 'finished': done, 'failed': failed}
114
115    def add_collate_callback(self, fn):
116        """Adds a callback to be called after collation is done.
117
118        Parameters
119        ----------
120        fn - A callable that takes previous as it's only input.
121        """
122        self._collate_callback = fn
123
124    def collate(self, progress_bar=False):
125        """A command that will block until all Futures in the pool have finished.
126        It will also store the results gather from `Actions` in the database.
127
128        Parameters
129        ----------
130        progress_bar: bool
131           Whether to show progress bar
132        """
133        if not progress_bar:
134            def tqdm_(x, total=None): return x
135        else:
136            tqdm_ = tqdm
137        if isinstance(self.pool, Client):
138            self.results = self.pool.gather(self.futures)
139        if self.sequential or isinstance(self.pool, Client):
140            for result in tqdm_(self.results, total=len(self.results)):
141                result = self._collate_callback(result)
142                self.campaign.campaign_db.store_result(
143                    result['run_id'], result, change_status=result['collated'])
144        else:
145            if isinstance(self.pool, QCGPJPool):
146                as_completed_fn = self.pool.as_completed
147                self.add_collate_callback(self.pool.convert_results)
148            else:
149                as_completed_fn = concurrent.futures.as_completed
150
151            for future in tqdm_(as_completed_fn(self.futures), total=len(self.futures)):
152                result = self._collate_callback(future.result())
153                self.campaign.campaign_db.store_result(
154                    result['run_id'], result, change_status=result['collated'])
155        self.campaign.campaign_db.session.commit()