diff options
Diffstat (limited to 'joblib.py')
| -rw-r--r-- | joblib.py | 85 |
1 files changed, 56 insertions, 29 deletions
@@ -43,11 +43,21 @@ class JobQueue(object): async def request(self, jobtype, payload, priority=0): """Create a new job and return its id.""" async with self._pool.acquire() as con: - return await con.fetchval(""" - INSERT INTO jobs(type, payload, priority) - VALUES($1, $2, $3) - RETURNING id - """, jobtype, json.dumps(payload), priority) + if isinstance(payload, list): + payloads = payload + else: + payloads = [payload] + ids = [] + for pl in payloads: + ids.append(await con.fetchval(""" + INSERT INTO jobs(type, payload, priority) + VALUES($1, $2, $3) + RETURNING id + """, jobtype, json.dumps(pl), priority)) + if isinstance(payload, list): + return ids + else: + return ids[0] async def acquire(self, jobtype): """Mark a job as running, return id, payload and priority. @@ -76,7 +86,7 @@ class JobQueue(object): 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: @@ -87,34 +97,51 @@ class JobQueue(object): """, jobid) async def finish(self, jobid): - """Mark a job as completed.""" + """Mark jobs as completed.""" async with self._pool.acquire() as con: - while True: - try: - async with con.transaction(isolation="serializable"): - await con.execute(""" - UPDATE jobs - SET status='finished' - WHERE id=$1 - """, jobid) - return - except asyncpg.exceptions.SerializationError: - pass + if not isinstance(jobid, list): + jobids = [jobid] + else: + jobids = jobid + for jid in jobids: + await con.execute(""" + UPDATE jobs + SET status='finished' + WHERE id=$1 + """, jid) async def fail(self, jobid, reason): """Mark a job as failed.""" async with self._pool.acquire() as con: - while True: - try: - async with con.transaction(isolation="serializable"): - await con.execute(""" - UPDATE jobs - SET status='failed', payload=payload||$2::jsonb - WHERE id=$1 - """, jobid, json.dumps({"error": reason})) - return - except asyncpg.exceptions.SerializationError: - pass + 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) + for jid, rsn in zip(jobids, reasons): + await con.execute(""" + UPDATE jobs + SET status='failed', payload=payload||$2::jsonb + WHERE id=$1 + """, jid, json.dumps({"error": rsn})) + + async def reset(self, jobid): + """Mark a job as open.""" + async with self._pool.acquire() as con: + if not isinstance(jobid, list): + jobids = [jobid] + else: + jobids = jobid + for jid in jobids: + await con.execute(""" + UPDATE jobs + SET status='open' + WHERE id=$1 + """, jid) async def cleanup(self): """Reopen all unfinished jobs.""" |