Move subpool_map from Job to context

Author: sylvinus

mrq/context.py

@@ -4,6 +4,7 @@
 import re
 import time
 from .utils import LazyObject
+from itertools import count as itertools_count
 
 
 # greenletid => Job object
@@ -19,11 +20,16 @@
 
 def set_current_job(job):
     current = gevent.getcurrent()
-    _GREENLET_JOBS_REGISTRY[id(current)] = job
 
     current.__dict__["_trace_time"] = 0
     current.__dict__["_trace_switches"] = 0
 
+    if job is None:
+        if id(current) in _GREENLET_JOBS_REGISTRY:
+            del _GREENLET_JOBS_REGISTRY[id(current)]
+    else:
+        _GREENLET_JOBS_REGISTRY[id(current)] = job
+
 
 def get_current_job(greenlet_id=None):
     if greenlet_id is None:
@@ -199,3 +205,33 @@ def metric(name, incr=1, **kwargs):
     cfg = get_current_config()
     if cfg.get("metric_hook"):
         return cfg.get("metric_hook")(name, incr=incr, **kwargs)
+
+
+def subpool_map(pool_size, func, iterable):
+    """ Starts a Gevent pool and run a map. Takes care of setting current_job and cleaning up. """
+
+    if not pool_size:
+        return [func(*args) for args in iterable]
+
+    counter = itertools_count()
+
+    current_job = get_current_job()
+
+    def inner_func(*args):
+        next(counter)
+        if current_job:
+            set_current_job(current_job)
+        ret = func(*args)
+        if current_job:
+            set_current_job(None)
+        return ret
+
+    start_time = time.time()
+    pool = gevent.pool.Pool(size=pool_size)
+    ret = pool.map(inner_func, iterable)
+    pool.join(raise_error=True)
+    total_time = time.time() - start_time
+
+    log.debug("SubPool ran %s greenlets in %0.6fs" % (counter, total_time))
+
+    return ret

mrq/job.py

@@ -13,7 +13,6 @@
 from collections import defaultdict
 import traceback
 import sys
-from itertools import count as itertools_count
 
 
 class Job(object):
@@ -179,31 +178,6 @@ def insert(self, jobs_data, queue=None, return_jobs=True, w=1):
         else:
             return inserted
 
-    def subpool_map(self, pool_size, func, iterable):
-        """ Starts a Gevent pool and run a map. Takes care of setting current_job and cleaning up. """
-
-        if not pool_size:
-            return [func(*args) for args in iterable]
-
-        counter = itertools_count()
-
-        def inner_func(*args):
-            next(counter)
-            set_current_job(self)
-            ret = func(*args)
-            set_current_job(None)
-            return ret
-
-        start_time = time.time()
-        pool = gevent.pool.Pool(size=pool_size)
-        ret = pool.map(inner_func, iterable)
-        pool.join(raise_error=True)
-        total_time = time.time() - start_time
-
-        log.debug("SubPool ran %s greenlets in %0.6fs" % (counter, total_time))
-
-        return ret
-
     def save_status(
             self,
             status,
@@ -403,6 +377,11 @@ def trace_memory_start(self):
         gc.collect()
         self._memory_start = self.worker.get_memory()
 
+    def subpool_map(self, *args, **kwargs):
+        """ Deprecated! Use from mrq.context """
+        from mrq.context import subpool_map
+        return subpool_map(*args, **kwargs)
+
     def trace_memory_stop(self):
         """ Stops measuring memory consumption """
 

mrq/version.py

@@ -1 +1 @@
-VERSION = "0.1.3"
+VERSION = "0.1.4"

tests/tasks/general.py

@@ -1,6 +1,6 @@
 from time import sleep
 from mrq.task import Task
-from mrq.context import log, retry_current_job, connections, get_current_config, get_current_job, progress
+from mrq.context import log, retry_current_job, connections, get_current_config, get_current_job, progress, subpool_map
 from mrq.queue import send_task
 import urllib2
 import json
@@ -142,10 +142,9 @@ def inner(self, x):
         return x
 
     def run(self, params):
-
         self.job = get_current_job()
 
-        return self.job.subpool_map(params["pool_size"], self.inner, params["inner_params"])
+        return subpool_map(params["pool_size"], self.inner, params["inner_params"])
 
 
 class GetMetrics(Task):