X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=parallelize.py;h=98f883c5716ae8fa3d67101b1d38d20dece7a0b6;hb=36fea7f15ed17150691b5b3ead75450e575229ef;hp=08220951a000e3ee9c24dcc908f54af1067ee265;hpb=11eeb8574b7b4620ac6fd440cb251f8aa2458f5b;p=python_utils.git

diff --git a/parallelize.py b/parallelize.py
index 0822095..98f883c 100644
--- a/parallelize.py
+++ b/parallelize.py
@@ -2,14 +2,11 @@
 
 """A decorator to help with dead simple parallelization."""
 
+import atexit
 from enum import Enum
 import functools
 import typing
 
-ps_count = 0
-thread_count = 0
-remote_count = 0
-
 
 class Method(Enum):
     THREAD = 1
@@ -18,23 +15,23 @@ class Method(Enum):
 
 
 def parallelize(
-        _funct: typing.Optional[typing.Callable] = None,
-        *,
-        method: Method = Method.THREAD
+    _funct: typing.Optional[typing.Callable] = None,
+    *,
+    method: Method = Method.THREAD
 ) -> typing.Callable:
     """Usage:
 
-    @parallelize    # defaults to thread-mode
-    def my_function(a, b, c) -> int:
-        ...do some slow / expensive work, e.g., an http request
+        @parallelize    # defaults to thread-mode
+        def my_function(a, b, c) -> int:
+            ...do some slow / expensive work, e.g., an http request
 
-    @parallelize(method=Method.PROCESS)
-    def my_other_function(d, e, f) -> str:
-        ...do more really expensice work, e.g., a network read
+        @parallelize(method=Method.PROCESS)
+        def my_other_function(d, e, f) -> str:
+            ...do more really expensice work, e.g., a network read
 
-    @parallelize(method=Method.REMOTE)
-    def my_other_other_function(g, h) -> int:
-        ...this work will be distributed to a remote machine pool
+        @parallelize(method=Method.REMOTE)
+        def my_other_other_function(g, h) -> int:
+            ...this work will be distributed to a remote machine pool
 
     This decorator will invoke the wrapped function on:
 
@@ -56,9 +53,14 @@ def parallelize(
     your machine with processes or your network with remote processes
     beyond the control mechanisms built into one instance of the pool.
     Be careful.
+
+    Also note: there is a non trivial overhead of pickling code and
+    scp'ing it over the network when you use Method.REMOTE.  There's
+    a smaller but still considerable cost of creating a new process
+    and passing code to/from it when you use Method.PROCESS.
     """
-    def wrapper(funct: typing.Callable):
 
+    def wrapper(funct: typing.Callable):
         @functools.wraps(funct)
         def inner_wrapper(*args, **kwargs):
             import executors
@@ -71,9 +73,7 @@ def parallelize(
                 newargs.append(smart_future.SmartFuture.resolve(arg))
             newkwargs = {}
             for kw in kwargs:
-                newkwargs[kw] = smart_future.SmartFuture.resolve(
-                    kwargs[kw]
-                )
+                newkwargs[kw] = smart_future.SmartFuture.resolve(kwargs[kw])
 
             executor = None
             if method == Method.PROCESS:
@@ -83,6 +83,7 @@ def parallelize(
             elif method == Method.REMOTE:
                 executor = executors.DefaultExecutors().remote_pool()
             assert executor is not None
+            atexit.register(executors.DefaultExecutors().shutdown)
 
             future = executor.submit(funct, *newargs, **newkwargs)