Skip to content

Commit b2c6528

Browse files
william-silversmithwilliam-silversmith
authored
feat: Multiprocess LocalTaskQueue (seung-lab#6)
1 parent 086d856 commit b2c6528

7 files changed

Lines changed: 60 additions & 8 deletions

File tree

‎AUTHORS‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
1+
Jingpeng Wu <jingpeng.wu@gmail.com>
12
William Silversmith <william.silversmith@gmail.com>

‎ChangeLog‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,12 @@
11
CHANGES
22
=======
33

4+
0.4.1
5+
-----
6+
7+
* fix: updated ptq to use updated cloud-volume semantics
8+
* Update README.md
9+
410
0.4.0
511
-----
612

‎requirements.txt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@ cloud-volume>=0.20.2
22
google-api-python-client
33
google-cloud>=0.32
44
google-auth-httplib2==0.0.2
5-
numpy>=1.14.3
5+
numpy>=1.13.1
66
pytest>=3.3.1
77
pbr
8-
tqdm==4.23.2
8+
-e git+https://github.com/seung-lab/tqdm.git#egg=tqdm
99
requests>2,<3

‎taskqueue/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
from .registered_task import RegisteredTask, MockTask, PrintTask
2-
from .taskqueue import TaskQueue, MockTaskQueue
2+
from .taskqueue import TaskQueue, MockTaskQueue, LocalTaskQueue
33
from .secrets import (
44
QUEUE_NAME, TEST_QUEUE_NAME, QUEUE_TYPE,
55
PROJECT_NAME, APPENGINE_QUEUE_URL

‎taskqueue/registered_task.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -77,13 +77,17 @@ def __repr__(self):
7777
return string + ")"
7878

7979
class PrintTask(RegisteredTask):
80-
def __init__(self):
80+
def __init__(self, txt=''):
8181
super(PrintTask, self).__init__()
82+
self.txt = txt
8283
def execute(self):
83-
print(self)
84+
if self.txt:
85+
print(str(self) + ": " + str(self.txt))
86+
else:
87+
print(self)
8488

8589
class MockTask(RegisteredTask):
86-
def __init__(self):
90+
def __init__(self, *args):
8791
super(MockTask, self).__init__()
8892
def execute(self):
8993
pass

‎taskqueue/taskqueue.py‎

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,11 @@
44
import json
55
from functools import partial
66

7+
import concurrent.futures
8+
import multiprocessing as mp
79
import googleapiclient.errors
810
import numpy as np
11+
from tqdm import tqdm
912

1013
from cloudvolume.threaded_queue import ThreadedQueue
1114

@@ -16,7 +19,6 @@
1619
from .secrets import PROJECT_NAME, QUEUE_NAME, QUEUE_TYPE
1720

1821
def totask(task):
19-
print(task)
2022
taskobj = payloadBase64Decode(task['payloadBase64'])
2123
taskobj._id = task['id']
2224
return taskobj
@@ -231,3 +233,36 @@ def __enter__(self):
231233

232234
def __exit__(self, exception_type, exception_value, traceback):
233235
pass
236+
237+
class LocalTaskQueue(object):
238+
def __init__(self, parallel=1, queue_name='', queue_server='', progress=True):
239+
if parallel and type(parallel) == bool:
240+
parallel = mp.cpu_count()
241+
242+
self.parallel = parallel
243+
self.queue = []
244+
self.progress = progress
245+
246+
def insert(self, task):
247+
self.queue.append(task)
248+
249+
def wait(self, progress=None):
250+
return self
251+
252+
def kill_threads(self):
253+
return self
254+
255+
def __enter__(self):
256+
return self
257+
258+
def __exit__(self, exception_type, exception_value, traceback):
259+
with tqdm(total=len(self.queue), desc="Tasks") as pbar:
260+
with concurrent.futures.ProcessPoolExecutor(max_workers=self.parallel) as executor:
261+
for _ in executor.map(task_execute, self.queue):
262+
pbar.update()
263+
self.queue = []
264+
265+
# Necessary to define here to make the
266+
# function picklable
267+
def task_execute(task):
268+
task.execute()

‎test/test_taskqueue.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from six.moves import range
55
import pytest
66

7-
from taskqueue import RegisteredTask, TaskQueue, PrintTask
7+
from taskqueue import RegisteredTask, TaskQueue, MockTask, PrintTask, LocalTaskQueue
88
from taskqueue import QUEUE_NAME
99

1010
TRAVIS_BRANCH = None if 'TRAVIS_BRANCH' not in os.environ else os.environ['TRAVIS_BRANCH']
@@ -98,3 +98,9 @@ def test_400_errors():
9898
with TaskQueue(n_threads=1, queue_name=QUEUE_NAME, queue_server=qtype, qurl=QURL) as tq:
9999
tq.delete('nonexistent')
100100

101+
def test_local_taskqueue():
102+
with LocalTaskQueue(parallel=True, progress=False) as tq:
103+
for i in range(20000):
104+
tq.insert(
105+
MockTask(i)
106+
)

0 commit comments

Comments
 (0)