Pickle dicts (#123)

* dicts when pickling and better error catching

* returning JobResult and JobDef

* fix linting, tests for unpicklable

* more tests for pickling

* correct docs

Author: samuelcolvin

.gitignore

@@ -16,3 +16,4 @@ __pycache__/
 /demo/tmp/
 .vscode/
 .venv/
+/.auto-format

HISTORY.rst

@@ -3,6 +3,10 @@
 History
 -------
 
+v0.16.0b1 (2019-04-23)
+......................
+* use dicts for pickling not tuples, better handling of pickling errors, #123
+
 v0.16.0a5 (2019-04-22)
 ......................
 * use ``pipeline`` in ``enqueue_job``

arq/connections.py

@@ -1,17 +1,16 @@
 import asyncio
 import logging
-import pickle
 from dataclasses import dataclass
 from datetime import datetime, timedelta
-from operator import itemgetter
-from typing import Any, Dict, List, Optional, Union
+from operator import attrgetter
+from typing import Any, List, Optional, Union
 from uuid import uuid4
 
 import aioredis
 from aioredis import MultiExecError, Redis
 
 from .constants import job_key_prefix, queue_name, result_key_prefix
-from .jobs import Job
+from .jobs import Job, JobResult, pickle_job
 from .utils import timestamp_ms, to_ms, to_unix_ms
 
 logger = logging.getLogger('arq.connections')
@@ -96,7 +95,7 @@ async def enqueue_job(
 
             expires_ms = expires_ms or score - enqueue_time_ms + expires_extra_ms
 
-            job = pickle.dumps((enqueue_time_ms, _job_try, function, args, kwargs))
+            job = pickle_job(function, args, kwargs, _job_try, enqueue_time_ms)
             tr = conn.multi_exec()
             tr.psetex(job_key, expires_ms, job)
             tr.zadd(queue_name, score, job_id)
@@ -111,16 +110,16 @@ async def _get_job_result(self, key):
         job_id = key[len(result_key_prefix) :]
         job = Job(job_id, self)
         r = await job.result_info()
-        r['job_id'] = job_id
+        r.job_id = job_id
         return r
 
-    async def all_job_results(self) -> List[Dict]:
+    async def all_job_results(self) -> List[JobResult]:
         """
         Get results for all jobs in redis.
         """
         keys = await self.keys(result_key_prefix + '*')
         results = await asyncio.gather(*[self._get_job_result(k) for k in keys])
-        return sorted(results, key=itemgetter('enqueue_time'))
+        return sorted(results, key=attrgetter('enqueue_time'))
 
 
 async def create_pool(settings: RedisSettings = None, *, _retry: int = 0) -> ArqRedis:

arq/jobs.py

@@ -1,11 +1,16 @@
 import asyncio
+import logging
 import pickle
+from dataclasses import dataclass
+from datetime import datetime
 from enum import Enum
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from .constants import in_progress_key_prefix, job_key_prefix, queue_name, result_key_prefix
 from .utils import ms_to_datetime, poll, timestamp_ms
 
+logger = logging.getLogger('arq.jobs')
+
 
 class JobStatus(str, Enum):
     """
@@ -24,6 +29,25 @@ class JobStatus(str, Enum):
     not_found = 'not_found'
 
 
+@dataclass
+class JobDef:
+    function: str
+    args: tuple
+    kwargs: dict
+    job_try: int
+    enqueue_time: datetime
+    score: Optional[int]
+
+
+@dataclass
+class JobResult(JobDef):
+    success: bool
+    result: Any
+    start_time: datetime
+    finish_time: datetime
+    job_id: Optional[str] = None
+
+
 class Job:
     """
     Holds data a reference to a job.
@@ -46,53 +70,35 @@ async def result(self, timeout: Optional[float] = None, *, pole_delay: float = 0
         async for delay in poll(pole_delay):
             info = await self.result_info()
             if info:
-                result = info['result']
-                if info['success']:
+                result = info.result
+                if info.success:
                     return result
                 else:
                     raise result
             if timeout is not None and delay > timeout:
                 raise asyncio.TimeoutError()
 
-    async def info(self) -> Optional[Dict[str, Any]]:
+    async def info(self) -> Optional[JobDef]:
         """
         All information on a job, including its result if it's available, does not wait for the result.
         """
         info = await self.result_info()
         if not info:
             v = await self._redis.get(job_key_prefix + self.job_id, encoding=None)
             if v:
-                enqueue_time_ms, job_try, function, args, kwargs = pickle.loads(v)
-                info = dict(
-                    enqueue_time=ms_to_datetime(enqueue_time_ms),
-                    job_try=job_try,
-                    function=function,
-                    args=args,
-                    kwargs=kwargs,
-                )
+                info = unpickle_job(v)
         if info:
-            info['score'] = await self._redis.zscore(queue_name, self.job_id)
+            info.score = await self._redis.zscore(queue_name, self.job_id)
         return info
 
-    async def result_info(self) -> Optional[Dict[str, Any]]:
+    async def result_info(self) -> Optional[JobResult]:
         """
         Information about the job result if available, does not wait for the result. Does not raise an exception
         even if the job raised one.
         """
         v = await self._redis.get(result_key_prefix + self.job_id, encoding=None)
         if v:
-            enqueue_time_ms, job_try, function, args, kwargs, s, r, start_time_ms, finish_time_ms = pickle.loads(v)
-            return dict(
-                enqueue_time=ms_to_datetime(enqueue_time_ms),
-                job_try=job_try,
-                function=function,
-                args=args,
-                kwargs=kwargs,
-                result=r,
-                success=s,
-                start_time=ms_to_datetime(start_time_ms),
-                finish_time=ms_to_datetime(finish_time_ms),
-            )
+            return unpickle_result(v)
 
     async def status(self) -> JobStatus:
         """
@@ -110,3 +116,78 @@ async def status(self) -> JobStatus:
 
     def __repr__(self):
         return f'<arq job {self.job_id}>'
+
+
+class PickleError(RuntimeError):
+    pass
+
+
+def pickle_job(function_name: str, args: tuple, kwargs: dict, job_try: int, enqueue_time_ms: int):
+    data = {'t': job_try, 'f': function_name, 'a': args, 'k': kwargs, 'et': enqueue_time_ms}
+    try:
+        return pickle.dumps(data)
+    except Exception as e:
+        raise PickleError(f'unable to pickle job "{function_name}"') from e
+
+
+def pickle_result(
+    function: str,
+    args: tuple,
+    kwargs: dict,
+    job_try: int,
+    enqueue_time_ms: int,
+    success: bool,
+    result: Any,
+    start_ms: int,
+    finished_ms: int,
+    ref: str,
+) -> Optional[bytes]:
+    data = {
+        't': job_try,
+        'f': function,
+        'a': args,
+        'k': kwargs,
+        'et': enqueue_time_ms,
+        's': success,
+        'r': result,
+        'st': start_ms,
+        'ft': finished_ms,
+    }
+    try:
+        return pickle.dumps(data)
+    except Exception:
+        logger.warning('error pickling result of %s', ref, exc_info=True)
+
+    data.update(r=PickleError('unable to pickle result'), s=False)
+    try:
+        return pickle.dumps(data)
+    except Exception:
+        logger.critical('error pickling result of %s even after replacing result', ref, exc_info=True)
+
+
+def unpickle_job(r: bytes) -> JobDef:
+    d = pickle.loads(r)
+    return JobDef(
+        function=d['f'], args=d['a'], kwargs=d['k'], job_try=d['t'], enqueue_time=ms_to_datetime(d['et']), score=None
+    )
+
+
+def unpickle_job_raw(r: bytes) -> tuple:
+    d = pickle.loads(r)
+    return d['f'], d['a'], d['k'], d['t'], d['et']
+
+
+def unpickle_result(r: bytes) -> JobResult:
+    d = pickle.loads(r)
+    return JobResult(
+        job_try=d['t'],
+        function=d['f'],
+        args=d['a'],
+        kwargs=d['k'],
+        enqueue_time=ms_to_datetime(d['et']),
+        score=None,
+        success=d['s'],
+        result=d['r'],
+        start_time=ms_to_datetime(d['st']),
+        finish_time=ms_to_datetime(d['ft']),
+    )

arq/version.py

@@ -2,4 +2,4 @@
 
 __all__ = ['VERSION']
 
-VERSION = StrictVersion('0.16a5')
+VERSION = StrictVersion('0.16b1')

arq/worker.py

@@ -1,7 +1,6 @@
 import asyncio
 import inspect
 import logging
-import pickle
 import signal
 from dataclasses import dataclass
 from datetime import datetime
@@ -15,6 +14,7 @@
 from pydantic.utils import import_string
 
 from arq.cron import CronJob
+from arq.jobs import pickle_result, unpickle_job_raw
 
 from .connections import ArqRedis, RedisSettings, create_pool, log_redis_info
 from .constants import (
@@ -107,7 +107,7 @@ def __init__(self, count, job_results):
 
     def __str__(self):
         if self.count == 1 and self.job_results:
-            exc = self.job_results[0]['result']
+            exc = self.job_results[0].result
             return f'1 job failed "{exc.__class__.__name__}: {exc}"'
         else:
             return f'{self.count} jobs failed'
@@ -219,7 +219,7 @@ async def run_check(self) -> int:
         """
         await self.async_run()
         if self.jobs_failed:
-            failed_job_results = [r for r in await self.pool.all_job_results() if not r['success']]
+            failed_job_results = [r for r in await self.pool.all_job_results() if not r.success]
             raise FailedJobs(self.jobs_failed, failed_job_results)
         else:
             return self.jobs_complete
@@ -290,7 +290,7 @@ async def run_job(self, job_id, score):  # noqa: C901
             self.jobs_failed += 1
             return await asyncio.shield(self.abort_job(job_id))
 
-        enqueue_time_ms, enqueue_job_try, function_name, args, kwargs = pickle.loads(v)
+        function_name, args, kwargs, enqueue_job_try, enqueue_time_ms = unpickle_job_raw(v)
 
         try:
             function: Union[Function, CronJob] = self.functions[function_name]
@@ -376,11 +376,9 @@ async def run_job(self, job_id, score):  # noqa: C901
         result_timeout_s = self.keep_result_s if function.keep_result_s is None else function.keep_result_s
         result_data = None
         if result is not no_result and result_timeout_s > 0:
-            d = enqueue_time_ms, job_try, function_name, args, kwargs, success, result, start_ms, finished_ms
-            try:
-                result_data = pickle.dumps(d)
-            except Exception:
-                logger.critical('error pickling result of %s', ref, exc_info=True)
+            result_data = pickle_result(
+                function_name, args, kwargs, job_try, enqueue_time_ms, success, result, start_ms, finished_ms, ref
+            )
 
         await asyncio.shield(self.finish_job(job_id, finish, result_data, result_timeout_s, incr_score))
 

docs/examples/job_results.py

@@ -25,13 +25,14 @@ async def main():
     debug(await job.info())
     """
     >   docs/examples/job_results.py:23 main
-    {
-        'enqueue_time': datetime.datetime(2019, 3, 3, 12, 32, 19, 975000),
-        'function': 'the_task',
-        'args': (),
-        'kwargs': {},
-        'score': 1551616339975,
-    } (dict) len=5
+    JobDef(
+        function='the_task',
+        args=(),
+        kwargs={},
+        job_try=None,
+        enqueue_time=datetime.datetime(2019, 4, 23, 13, 58, 56, 781000),
+        score=1556027936781
+    ) (JobDef)
     """
 
     # get the Job's status