diff --git a/AUTHORS b/AUTHORS index 25dda54..75e6363 100644 --- a/AUTHORS +++ b/AUTHORS @@ -1 +1,2 @@ +Jingpeng Wu William Silversmith diff --git a/ChangeLog b/ChangeLog index c0452d1..4045061 100644 --- a/ChangeLog +++ b/ChangeLog @@ -1,6 +1,12 @@ CHANGES ======= +0.4.1 +----- + +* fix: updated ptq to use updated cloud-volume semantics +* Update README.md + 0.4.0 ----- diff --git a/requirements.txt b/requirements.txt index 4d8f038..13e52a0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,8 +2,8 @@ cloud-volume>=0.20.2 google-api-python-client google-cloud>=0.32 google-auth-httplib2==0.0.2 -numpy>=1.14.3 +numpy>=1.13.1 pytest>=3.3.1 pbr -tqdm==4.23.2 +-e git+https://github.com/seung-lab/tqdm.git#egg=tqdm requests>2,<3 diff --git a/taskqueue/__init__.py b/taskqueue/__init__.py index 23dc1c5..2b08fd3 100644 --- a/taskqueue/__init__.py +++ b/taskqueue/__init__.py @@ -1,5 +1,5 @@ from .registered_task import RegisteredTask, MockTask, PrintTask -from .taskqueue import TaskQueue, MockTaskQueue +from .taskqueue import TaskQueue, MockTaskQueue, LocalTaskQueue from .secrets import ( QUEUE_NAME, TEST_QUEUE_NAME, QUEUE_TYPE, PROJECT_NAME, APPENGINE_QUEUE_URL diff --git a/taskqueue/registered_task.py b/taskqueue/registered_task.py index 0b4ad92..864df56 100644 --- a/taskqueue/registered_task.py +++ b/taskqueue/registered_task.py @@ -77,13 +77,17 @@ def __repr__(self): return string + ")" class PrintTask(RegisteredTask): - def __init__(self): + def __init__(self, txt=''): super(PrintTask, self).__init__() + self.txt = txt def execute(self): - print(self) + if self.txt: + print(str(self) + ": " + str(self.txt)) + else: + print(self) class MockTask(RegisteredTask): - def __init__(self): + def __init__(self, *args): super(MockTask, self).__init__() def execute(self): pass \ No newline at end of file diff --git a/taskqueue/taskqueue.py b/taskqueue/taskqueue.py index cb88c8f..ed371a0 100644 --- a/taskqueue/taskqueue.py +++ b/taskqueue/taskqueue.py @@ -4,8 +4,11 @@ import json from functools import partial +import concurrent.futures +import multiprocessing as mp import googleapiclient.errors import numpy as np +from tqdm import tqdm from cloudvolume.threaded_queue import ThreadedQueue @@ -16,7 +19,6 @@ from .secrets import PROJECT_NAME, QUEUE_NAME, QUEUE_TYPE def totask(task): - print(task) taskobj = payloadBase64Decode(task['payloadBase64']) taskobj._id = task['id'] return taskobj @@ -231,3 +233,36 @@ def __enter__(self): def __exit__(self, exception_type, exception_value, traceback): pass + +class LocalTaskQueue(object): + def __init__(self, parallel=1, queue_name='', queue_server='', progress=True): + if parallel and type(parallel) == bool: + parallel = mp.cpu_count() + + self.parallel = parallel + self.queue = [] + self.progress = progress + + def insert(self, task): + self.queue.append(task) + + def wait(self, progress=None): + return self + + def kill_threads(self): + return self + + def __enter__(self): + return self + + def __exit__(self, exception_type, exception_value, traceback): + with tqdm(total=len(self.queue), desc="Tasks") as pbar: + with concurrent.futures.ProcessPoolExecutor(max_workers=self.parallel) as executor: + for _ in executor.map(task_execute, self.queue): + pbar.update() + self.queue = [] + +# Necessary to define here to make the +# function picklable +def task_execute(task): + task.execute() diff --git a/test/test_taskqueue.py b/test/test_taskqueue.py index 693beaf..0bd9e8d 100644 --- a/test/test_taskqueue.py +++ b/test/test_taskqueue.py @@ -4,7 +4,7 @@ from six.moves import range import pytest -from taskqueue import RegisteredTask, TaskQueue, PrintTask +from taskqueue import RegisteredTask, TaskQueue, MockTask, PrintTask, LocalTaskQueue from taskqueue import QUEUE_NAME TRAVIS_BRANCH = None if 'TRAVIS_BRANCH' not in os.environ else os.environ['TRAVIS_BRANCH'] @@ -98,3 +98,9 @@ def test_400_errors(): with TaskQueue(n_threads=1, queue_name=QUEUE_NAME, queue_server=qtype, qurl=QURL) as tq: tq.delete('nonexistent') +def test_local_taskqueue(): + with LocalTaskQueue(parallel=True, progress=False) as tq: + for i in range(20000): + tq.insert( + MockTask(i) + ) \ No newline at end of file