support job batching (#12)

* support job batching

* worker: pull more frequently

1 files changed, 56 insertions, 29 deletions

diff --git a/joblib.py b/joblib.py
index 7ecde6a..9e8cf84 100644
--- a/joblib.py
+++ b/joblib.py
@@ -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."""