From 596ee7f64662dd94f1a21967f15d451bad969a2e Mon Sep 17 00:00:00 2001 From: schneefux Date: Mon, 6 Mar 2017 20:15:51 +0100 Subject: support acquiring in batches --- joblib.py | 46 ++++++++++++++++++++++++++++++---------------- test_joblib.py | 9 +++++++++ worker.py | 58 +++++++++++++++++++++------------------------------------- 3 files changed, 60 insertions(+), 53 deletions(-) diff --git a/joblib.py b/joblib.py index c384e5f..45f1cbf 100644 --- a/joblib.py +++ b/joblib.py @@ -59,33 +59,47 @@ class JobQueue(object): else: return ids[0] - async def acquire(self, jobtype): + async def acquire(self, jobtype, length=None): """Mark a job as running, return id, payload and priority. Return (None, None, None) if no job is available.""" async with self._pool.acquire() as con: + if length is None: + limit = 1 + else: + limit = length while True: try: # do not allow async access async with con.transaction(isolation="serializable"): - result = await con.fetchrow(""" - SELECT id, payload, priority - FROM jobs WHERE - type=$1 AND status='open' - ORDER BY priority DESC - LIMIT 1 - """, jobtype) - if result is None: + result = await con.fetch(""" + UPDATE jobs SET STATUS='running' + FROM ( + SELECT id FROM jobs + WHERE status='open' AND type=$1 + ORDER BY PRIORITY + LIMIT $2 + ) AS open_jobs + WHERE jobs.id=open_jobs.id + RETURNING jobs.id, jobs.payload, jobs.priority + """, jobtype, limit) + if len(result) == 0 and length is None: # no jobs available + # backwards compatibility return None, None, None - jobid, payload, priority = result - await con.execute(""" - UPDATE jobs - SET status='running' - WHERE id=$1 - """, jobid) - return jobid, json.loads(payload), priority + + jobs = [] + for record in result: + jobs.append((record[0], + json.loads(record[1]), + record[2])) + + if length is None: + return jobs[0] + else: + return jobs except asyncpg.exceptions.SerializationError: # job is being picked up by another worker, try again + print("serialization error") pass async def status(self, jobid): diff --git a/test_joblib.py b/test_joblib.py index 1a760f8..fc9640b 100644 --- a/test_joblib.py +++ b/test_joblib.py @@ -58,6 +58,15 @@ class TestJoblib: # there should not be another job assert None == (await queue.acquire(jobtype="testing"))[1] + @pytest.mark.asyncio + async def test_request_and_acquire_batch_batched(self, queue, payload): + # request should succeed + await queue.request(jobtype="testing", payload=[payload]*5) + # acquire should return same payload + assert payload == (await queue.acquire(jobtype="testing", length=5))[0][1] + # there should not be another job + assert None == (await queue.acquire(jobtype="testing"))[1] + @pytest.mark.asyncio async def test_cleanup(self, queue, payload): await queue.request(jobtype="testing", payload=payload) diff --git a/worker.py b/worker.py index 163ebfd..50642f6 100644 --- a/worker.py +++ b/worker.py @@ -36,48 +36,32 @@ class Worker(object): # override pass - async def _work(self): - """Fetch a job and run it. - Return id.""" - jobid, payload, priority = await self._queue.acquire( - jobtype=self._jobtype) - if jobid is None: - raise LookupError - try: - await self._execute_job(jobid, payload, priority) - except JobFailed as err: - raise JobFailed(err.args[0], jobid) - return jobid - async def run(self, batchlimit=1): """Start jobs forever.""" while True: await self._windup() - jobids = [] - error = None - low_load = False - try: - for _ in range(batchlimit): - try: - jobids.append(await self._work()) - except LookupError: - low_load = True - break - except JobFailed as err: - error = err.args[0] - jobids.append(err.args[1]) - break - finally: - if error is not None: - await self._queue.reset(jobids[:-1]) - logging.debug(jobids) - await self._queue.fail(jobids[-1], error) - logging.warning("batch failed, reset") - else: - await self._queue.finish(jobids) - await self._teardown(failed=error is not None) - if low_load: + jobs = await self._queue.acquire(jobtype=self._jobtype, + length=batchlimit) + + if len(jobs) == 0: await asyncio.sleep(0.1) + # nothing to do + continue + + for jobid, payload, priority in jobs: + try: + self._execute_job(jobid, payload, priority) + except JobFailed as err: + error = err.args[0] + finally: + if error is not None: + await self._queue.reset([j[0] for j in jobs]) + await self._queue.fail(jobid, error) + logging.warning("batch failed, reset") + await self._teardown(failed=True) + else: + await self._queue.finish([j[0] for j in jobs]) + await self._teardown(failed=False) async def start(self, batchlimit=1): """Start in background.""" -- cgit v1.3.1