From ceea54d460812545866505683a8a1a365aa79214 Mon Sep 17 00:00:00 2001 From: schneefux Date: Mon, 13 Mar 2017 20:52:53 +0100 Subject: rewrite using notifications on one connection --- joblib.py | 303 ++++++++++++++++++++++++++++++++------------------------------ worker.py | 106 ++++++++++++++-------- 2 files changed, 225 insertions(+), 184 deletions(-) diff --git a/joblib.py b/joblib.py index 8faf71f..ab107f8 100644 --- a/joblib.py +++ b/joblib.py @@ -8,15 +8,28 @@ import logging class JobQueue(object): def __init__(self): - self._pool = None + self._con = None + self._listens = {} + + def _listener(self, con, pid, + channel, payload): + """Fire the registered callback. Must not be async.""" + self._listens[channel](payload) + + async def _hook_listener(self): + """Register callbacks.""" + if len(self._listens.keys()) > 0: + for channel in self._listens.keys(): + await self._con.add_listener(channel, + self._listener) async def connect(self, **args): """Connect the database.""" logging.info("connecting to queue database") while True: try: - self._pool = await asyncpg.create_pool( - min_size=1, **args) + self._con = await asyncpg.connect(**args) + await self._hook_listener() break except asyncpg.exceptions.CannotConnectNowError: logging.warning( @@ -29,166 +42,162 @@ class JobQueue(object): async def setup(self): """Initialize the database.""" - async with self._pool.acquire() as con: - await con.execute(""" - CREATE TABLE IF NOT EXISTS - jobs ( - id SERIAL PRIMARY KEY, - priority INT DEFAULT 0, - status TEXT DEFAULT 'open', - type TEXT, - payload JSONB - ) - """) - await con.execute(""" - CREATE UNIQUE INDEX ON jobs(priority, id) - """) + await self._con.execute(""" + CREATE TABLE IF NOT EXISTS + jobs ( + id SERIAL PRIMARY KEY, + priority INT DEFAULT 0, + status TEXT DEFAULT 'open', + type TEXT, + payload JSONB + ) + """) + await self._con.execute(""" + CREATE UNIQUE INDEX ON jobs(priority, id) + """) + + async def listen(self, channel, callback): + """Hook a notification listener.""" + self._listens[channel] = callback async def request(self, jobtype, payload, priority=0): """Create a new job and return its id.""" - async with self._pool.acquire() as con: - insert = await con.prepare(""" - INSERT INTO jobs(type, payload, priority) - VALUES($1, $2, $3) - RETURNING id - """) - if isinstance(payload, list): - payloads = payload - else: - payloads = [payload] - if isinstance(priority, list): - priorities = priority - else: - priorities = [priority] * len(payloads) - ids = [] - async with con.transaction(): - for pl, pr in zip(payloads, priorities): - ids.append(await insert.fetchval(jobtype, - json.dumps(pl), - pr)) - if isinstance(payload, list): - return ids - else: - return ids[0] + if isinstance(payload, list): + payloads = [json.dumps(p) + for p in payload] + else: + payloads = [json.dumps(payload)] + if isinstance(priority, list): + priorities = priority + else: + priorities = [priority] * len(payloads) + insert = await self._con.prepare(""" + INSERT INTO jobs(type, payload, priority) + VALUES($1, $2, $3) + RETURNING id + """) + ids = [] + async with self._con.transaction(): + for pl, pr in zip(payloads, priorities): + ids.append(await insert.fetch(jobtype, pl, pr)) + await self._con.execute("SELECT pg_notify($1 || '_open', '')", + jobtype) + + if isinstance(payload, list): + return ids + else: + return ids[0] 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: - update = await con.prepare(""" - UPDATE jobs SET STATUS='running' - FROM ( - SELECT id FROM jobs - WHERE status='open' AND type=$1 - ORDER BY priority, id - LIMIT $2 - ) AS open_jobs - WHERE jobs.id=open_jobs.id - RETURNING jobs.id, jobs.payload, jobs.priority - """) - 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 update.fetch(jobtype, limit) - if len(result) == 0 and length is None: - # no jobs available - # backwards compatibility - return None, None, None - - 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 - logging.debug("serialization error, retrying") - pass + if length is None: + limit = 1 + else: + limit = length + while True: + try: + # do not allow async access + async with self._con.transaction(isolation="serializable"): + result = await self._con.fetch(""" + UPDATE jobs SET STATUS='running' + FROM ( + SELECT id FROM jobs + WHERE status='open' AND type=$1 + ORDER BY priority, id + 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 + + jobs = [(r[0], json.loads(r[1]), r[2]) for r in result] + + await self._con.execute("SELECT pg_notify($1 || '_running', '')", + jobtype) + + if length is None: + return jobs[0] + else: + return jobs + except asyncpg.exceptions.SerializationError: + # job is being picked up by another worker, try again + pass async def status(self, jobid): """Return the status of a job.""" - async with self._pool.acquire() as con: - return await con.fetchval(""" - SELECT status - FROM jobs WHERE - id=$1 - """, jobid) - - async def finish(self, jobid): + return await self._con.fetchval(""" + SELECT status + FROM jobs WHERE + id=$1 + """, jobid) + + async def finish(self, jobid, jobtype): """Mark jobs as completed.""" - async with self._pool.acquire() as con: - update = await con.prepare(""" - UPDATE jobs - SET status='finished' + if not isinstance(jobid, list): + jobids = [(jobid,)] + else: + jobids = [(jid,) for jid in jobid] + async with self._con.transaction(): + await self._con.executemany(""" + UPDATE jobs SET status='finished' WHERE id=$1 - """) - if not isinstance(jobid, list): - jobids = [jobid] - else: - jobids = jobid - async with con.transaction(): - for jid in jobids: - await update.fetch(jid) - - async def fail(self, jobid, reason): + """, jobids) + await self._con.execute(""" + SELECT pg_notify($1 || '_finished', '') + """, jobtype) + + async def fail(self, jobid, jobtype, reason): """Mark a job as failed.""" - async with self._pool.acquire() as con: - update = await con.prepare(""" - UPDATE jobs - SET status='failed', payload=payload||$2::jsonb + if not isinstance(jobid, list): + jobids = [jobid] + else: + jobids = jobid + if not isinstance(reason, list): + reasons = [json.dumps({"error": reason})] + else: + reasons = [json.dumps({"error": r}) + for r in reason] + assert len(jobids) == len(reasons) + async with self._con.transaction(): + await self._con.executemany(""" + UPDATE jobs SET status='failed', + payload=payload||$2::jsonb WHERE id=$1 - """) - if not isinstance(jobid, list): - jobids = [jobid] - else: - jobids = jobid - if not isinstance(reason, list): - reasons = [reason] - else: - reasons = reason - assert len(jobids) == len(reasons) - async with con.transaction(): - for jid, rsn in zip(jobids, reasons): - await update.fetch(jid, - json.dumps({"error": rsn})) - - async def reset(self, jobid): + """, zip(jobids, reasons)) + await self._con.execute(""" + SELECT pg_notify($1 || '_failed', '') + """, jobtype) + + async def reset(self, jobid, jobtype): """Mark a job as open.""" - async with self._pool.acquire() as con: - update = await con.prepare(""" - UPDATE jobs - SET status='open' + if not isinstance(jobid, list): + jobids = [(jobid,)] + else: + jobids = [(jid,) for jid in jobid] + async with self._con.transaction(): + await self._con.executemany(""" + UPDATE jobs SET status='open' WHERE id=$1 - """) - if not isinstance(jobid, list): - jobids = [jobid] - else: - jobids = jobid - async with con.transaction(): - for jid in jobids: - await update.fetch(jid) + """, jobids) + await self._con.execute(""" + SELECT pg_notify($1 || '_open', '') + """, jobtype) async def cleanup(self): """Reopen all unfinished jobs.""" - async with self._pool.acquire() as con: - while True: - try: - async with con.transaction(isolation="serializable"): - await con.execute(""" - UPDATE jobs - SET status='open' - WHERE status='running' - """) - return - except asyncpg.exceptions.SerializationError: - pass + while True: + try: + async with self._con.transaction(isolation="serializable"): + await self._con.execute(""" + UPDATE jobs + SET status='open' + WHERE status='running' + """) + return + except asyncpg.exceptions.SerializationError: + pass diff --git a/worker.py b/worker.py index e5f0b25..d67223a 100644 --- a/worker.py +++ b/worker.py @@ -14,9 +14,15 @@ class Worker(object): def __init__(self, jobtype): self._queue = None self._jobtype = jobtype + self._need_update = True + self._children = [] async def connect(self, **queuedb): self._queue = joblib.joblib.JobQueue() + # register event handler on new open job + await self._queue.listen( + self._jobtype + "_open", + self._pushed_new) await self._queue.connect(**queuedb) await self._queue.setup() @@ -36,44 +42,70 @@ class Worker(object): # override pass - async def run(self, batchlimit=1): - """Start jobs forever.""" + async def request(self, jobtype, payload, priority=None): + """Queue child jobs to be executed after teardown.""" + self._children.append((jobtype, payload, priority)) + + async def poll(self, batchlimit=1): + """Start a batch of jobs.""" + jobs = await self._queue.acquire(jobtype=self._jobtype, + length=batchlimit) + + if len(jobs) == 0: + logging.debug("no jobs") + return + + await self._windup() + critical_error = False + failed = [] + succeeded = [] + for jobid, payload, priority in jobs: + try: + await self._execute_job(jobid, payload, priority) + succeeded.append(jobid) + except JobFailed as err: + failed.append((jobid, err.args[0])) + if len(err.args) > 1: + # rollback + critical_error = err.args[1] + if critical_error: + break + else: + critical_error = True # default + if critical_error: + await self._queue.reset(succeeded, + self._jobtype) + logging.warning("batch failed, reset") + else: + await self._queue.finish(succeeded, + self._jobtype) + + for jobid, err in failed: + await self._queue.fail(jobid, + self._jobtype, + err) + + await self._teardown(failed=critical_error) + + # spawn child jobs + for jobt, payloads, priorities in self._children: + await self._queue.request( + jobt, payloads, priorities) + self._children = [] + + def _pushed_new(self, _): + """New job event handler.""" + logging.debug("received a notification") + self._need_update = True + + async def run(self, batchlimit): + """Request a poll after a push or a timer.""" while True: - jobs = await self._queue.acquire(jobtype=self._jobtype, - length=batchlimit) - - if len(jobs) == 0: - await asyncio.sleep(1) - # nothing to do - continue - - await self._windup() - critical_error = False - failed = [] - succeeded = [] - for jobid, payload, priority in jobs: - try: - await self._execute_job(jobid, payload, priority) - succeeded.append(jobid) - except JobFailed as err: - failed.append((jobid, err.args[0])) - if len(err.args) > 1: - # rollback - critical_error = err.args[1] - if critical_error: - break - else: - critical_error = True # default - if critical_error: - await self._queue.reset(succeeded) - logging.warning("batch failed, reset") - else: - await self._queue.finish(succeeded) - - for jobid, err in failed: - await self._queue.fail(jobid, err) - - await self._teardown(failed=critical_error) + while not self._need_update: + await asyncio.sleep(0) + logging.debug("polling") + self._need_update = False + await self.poll(batchlimit) async def start(self, batchlimit=1): """Start in background.""" -- cgit v1.3.1