Various changes.

1 files changed, 69 insertions, 68 deletions

diff --git a/executors.py b/executors.py
index 2f4cf83..63efd81 100644
--- a/executors.py
+++ b/executors.py
@@ -23,7 +23,7 @@ import argparse_utils
 import config
 from exec_utils import run_silently, cmd_in_background
 from decorator_utils import singleton
-import histogram
+import histogram as hist
 
 logger = logging.getLogger(__name__)
 
@@ -61,33 +61,6 @@ parser.add_argument(
 
 RSYNC = 'rsync -q --no-motd -W --ignore-existing --timeout=60 --size-only -z'
 SSH = 'ssh -oForwardX11=no'
-HIST = histogram.SimpleHistogram(
-    histogram.SimpleHistogram.n_evenly_spaced_buckets(
-        int(0), int(500), 25
-    )
-)
-
-
-def run_local_bundle(fun, *args, **kwargs):
-    logger.debug(f"Running local bundle at {fun.__name__}")
-    start = time.time()
-    result = fun(*args, **kwargs)
-    end = time.time()
-    duration = end - start
-    logger.debug(f"{fun.__name__} finished; used {duration:.1f}s")
-    HIST.add_item(duration)
-    return result
-
-
-def run_cloud_pickle(pickle):
-    fun, args, kwargs = cloudpickle.loads(pickle)
-    logger.debug(f"Running pickled bundle at {fun.__name__}")
-    start = time.time()
-    result = fun(*args, **kwargs)
-    end = time.time()
-    duration = end - start
-    logger.debug(f"{fun.__name__} finished; used {duration:.1f}s")
-    return result
 
 
 def make_cloud_pickle(fun, *args, **kwargs):
@@ -96,8 +69,14 @@ def make_cloud_pickle(fun, *args, **kwargs):
 
 
 class BaseExecutor(ABC):
-    def __init__(self):
-        pass
+    def __init__(self, *, title=''):
+        self.title = title
+        self.task_count = 0
+        self.histogram = hist.SimpleHistogram(
+            hist.SimpleHistogram.n_evenly_spaced_buckets(
+                int(0), int(500), 50
+            )
+        )
 
     @abstractmethod
     def submit(self,
@@ -111,6 +90,10 @@ class BaseExecutor(ABC):
                  wait: bool = True) -> None:
         pass
 
+    def adjust_task_count(self, delta: int) -> None:
+        self.task_count += delta
+        logger.debug(f'Executor current task count is {self.task_count}')
+
 
 class ThreadExecutor(BaseExecutor):
     def __init__(self,
@@ -126,29 +109,36 @@ class ThreadExecutor(BaseExecutor):
             max_workers=workers,
             thread_name_prefix="thread_executor_helper"
         )
-        self.job_count = 0
+
+    def run_local_bundle(self, fun, *args, **kwargs):
+        logger.debug(f"Running local bundle at {fun.__name__}")
+        start = time.time()
+        result = fun(*args, **kwargs)
+        end = time.time()
+        self.adjust_task_count(-1)
+        duration = end - start
+        logger.debug(f"{fun.__name__} finished; used {duration:.1f}s")
+        self.histogram.add_item(duration)
+        return result
 
     def submit(self,
                function: Callable,
                *args,
                **kwargs) -> fut.Future:
-        self.job_count += 1
-        logger.debug(
-            f'Submitted work to threadpool; there are now {self.job_count} items.'
-        )
+        self.adjust_task_count(+1)
         newargs = []
         newargs.append(function)
         for arg in args:
             newargs.append(arg)
         return self._thread_pool_executor.submit(
-            run_local_bundle,
+            self.run_local_bundle,
             *newargs,
             **kwargs)
 
     def shutdown(self,
                  wait = True) -> None:
-        logger.debug("Shutting down threadpool executor.")
-        print(HIST)
+        logger.debug(f'Shutting down threadpool executor {self.title}')
+        print(self.histogram)
         self._thread_pool_executor.shutdown(wait)
 
 
@@ -165,24 +155,41 @@ class ProcessExecutor(BaseExecutor):
         self._process_executor = fut.ProcessPoolExecutor(
             max_workers=workers,
         )
-        self.job_count = 0
+
+    def run_cloud_pickle(self, pickle):
+        fun, args, kwargs = cloudpickle.loads(pickle)
+        logger.debug(f"Running pickled bundle at {fun.__name__}")
+        result = fun(*args, **kwargs)
+        self.adjust_task_count(-1)
+        return result
 
     def submit(self,
                function: Callable,
                *args,
                **kwargs) -> fut.Future:
-        # Bundle it up before submitting because pickle sucks.
+        start = time.time()
+        self.adjust_task_count(+1)
         pickle = make_cloud_pickle(function, *args, **kwargs)
-        self.job_count += 1
-        logger.debug(
-            f'Submitting work to processpool executor; there are now {self.job_count} items.'
+        result = self._process_executor.submit(
+            self.run_cloud_pickle,
+            pickle
+        )
+        result.add_done_callback(
+            lambda _: self.histogram.add_item(
+                time.time() - start
+            )
         )
-        return self._process_executor.submit(run_cloud_pickle, pickle)
+        return result
 
     def shutdown(self, wait=True) -> None:
-        logger.debug('Shutting down processpool executor')
-        print(HIST)
+        logger.debug(f'Shutting down processpool executor {self.title}')
         self._process_executor.shutdown(wait)
+        print(self.histogram)
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state['_process_executor'] = None
+        return state
 
 
 @dataclass
@@ -565,6 +572,7 @@ class RemoteExecutor(BaseExecutor):
 
     def launch(self, bundle: BundleDetails) -> Any:
         """Find a worker for bundle or block until one is available."""
+        self.adjust_task_count(+1)
         uuid = bundle.uuid
         hostname = bundle.hostname
         avoid_machine = None
@@ -648,6 +656,7 @@ class RemoteExecutor(BaseExecutor):
             # Whether original or backup, if we finished first we must
             # fetch the results if the computation happened on a
             # remote machine.
+            bundle.end_ts = time.time()
             if not was_cancelled:
                 assert bundle.machine is not None
                 if bundle.hostname not in bundle.machine:
@@ -658,31 +667,24 @@ class RemoteExecutor(BaseExecutor):
                     try:
                         run_silently(cmd)
                     except subprocess.CalledProcessError:
-                        pass
+                        logger.critical(f'Failed to copy {username}@{machine}:{result_file}!')
                     run_silently(f'{SSH} {username}@{machine}'
                                  f' "/bin/rm -f {code_file} {result_file}"')
-            bundle.end_ts = time.time()
+                dur = bundle.end_ts - bundle.start_ts
+                self.histogram.add_item(dur)
             assert bundle.worker is not None
             self.status.record_release_worker_already_locked(
                 bundle.worker,
                 bundle.uuid,
                 was_cancelled
             )
-            if not was_cancelled:
-                dur = bundle.end_ts - bundle.start_ts
-                HIST.add_item(dur)
-
-        # Original or not, the results should be back on the local
-        # machine.  Are they?
-        if not os.path.exists(result_file):
-            msg = f'{result_file} unexpectedly missing, wtf?!'
-            logger.critical(msg)
-            bundle.failure_count += 1
-            self.release_worker(bundle.worker)
-            raise Exception(msg)
 
         # Only the original worker should unpickle the file contents
-        # though since it's the only one whose result matters.
+        # though since it's the only one whose result matters.  The
+        # original is also the only job that may delete result_file
+        # from disk.  Note that the original may have been cancelled
+        # if one of the backups finished first; it still must read the
+        # result from disk.
         if is_original:
             logger.debug(f"Unpickling {result_file}.")
             try:
@@ -709,11 +711,11 @@ class RemoteExecutor(BaseExecutor):
                     )
                     backup.is_cancelled.set()
 
-        # This is a backup.
+        # This is a backup job.
         else:
             # Backup results don't matter, they just need to leave the
-            # result file in the right place for their original to
-            # read later.
+            # result file in the right place for their originals to
+            # read/unpickle later.
             result = None
 
             # Tell the original to stop if we finished first.
@@ -725,6 +727,7 @@ class RemoteExecutor(BaseExecutor):
 
         assert bundle.worker is not None
         self.release_worker(bundle.worker)
+        self.adjust_task_count(-1)
         return result
 
     def create_original_bundle(self, pickle):
@@ -811,14 +814,12 @@ class RemoteExecutor(BaseExecutor):
         pickle = make_cloud_pickle(function, *args, **kwargs)
         bundle = self.create_original_bundle(pickle)
         self.total_bundles_submitted += 1
-        logger.debug(
-            f'Submitted work to remote executor; {self.total_bundles_submitted} items now submitted'
-        )
         return self._helper_executor.submit(self.launch, bundle)
 
     def shutdown(self, wait=True) -> None:
         self._helper_executor.shutdown(wait)
-        print(HIST)
+        logging.debug(f'Shutting down RemoteExecutor {self.title}')
+        print(self.histogram)
 
 
 @singleton