Dilling instead of pickling to allow for functions and modules

airflow/models.py

@@ -5,7 +5,7 @@ import imp
 import jinja2
 import logging
 import os
-import pickle
+import dill as pickle
 import re
 import signal
 import socket
@@ -14,7 +14,6 @@ from sqlalchemy import (
     Column, Integer, String, DateTime, Text, Boolean, ForeignKey)
 from sqlalchemy import func
 from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.ext.serializer import loads, dumps
 from sqlalchemy.dialects.mysql import LONGTEXT
 from sqlalchemy.orm import relationship
 
@@ -230,10 +229,10 @@ class DagPickle(Base):
         self.job = job
         for t in dag.tasks:
             t.materialize_files()
-        self.pickle = dumps(dag)
+        self.pickle = pickle.dumps(dag)
 
     def get_object(self):
-        return loads(self.pickle)
+        return pickle.loads(self.pickle)
 
 
 class TaskInstance(Base):

requirements.txt

@@ -1,5 +1,6 @@
 celery
 chartkick
+dill
 flake8
 flask
 flask-admin

setup.py

@@ -14,6 +14,7 @@ setup(
     install_requires=[
         'celery>=3.1.17',
         'chartkick>=0.4.2',
+        'dill>=0.2.2',
         'flask>=0.10.1',
         'flask-admin>=1.0.9',
         'flask-bootstrap>=3.3.0.1',