Joblib: Running Python functions as pipeline jobs
joblib.readthedocs.io
joblib.readthedocs.io
I use joblib at work (it handles edge cases that pickle does not) to package up trained ML models for deployment by a colleague. I train the models outside of a Python environment, pull the models into the Python version of the ML library I use, and then serialize the models to a compressed joblib file. The joblib object takes up less space than the text based model file and deploys quickly for almost instant predictions.
To ensure that the input fields for training line up with the input fields for classification/regression, we add a hash to the joblib object, which is trivially easy. Dump the object on s3 using the hash as the key and it's trivial to have a library of ML models ready to deploy very quickly.
The file format was a zipfile with:
- a joblib dump of the python class wrapping the model and its data processing pipeline code
- a small csv (~100 rows) of test X data, test Y data and the expected outputs of the model on that data
- a little json with important library versions (pandas, numpy, etc.) so that the process loading the code could check it ran the same as the one that produced the model blob
Then the loader would try to load the file and verify it, and if failed would keep the existing model in memory and raise an alert to slack, so engineers could fix it and reupload a new model without downtime.
The whole thing was inspired by the talk "Alex Gaynor: Pickles are for Delis, not Software - PyCon 2014" [1]
At my current job we don't need such infrastructure because lambda/kubeflow gateways take care of it for us.
It's also insecure in the sense that you can pickle and unpickle more or less arbitrary Python objects and object structures, so there's no way to trust a pickle from a source you don't control, unless you want to start down the rabbit hole of encrypting and signing your pickles. But that's actually never come up for me in practice.
The incompatibility issue has always been the bigger deal for me. You don't want to upgrade Python and then have to deal with converting all your model weights to a new serialization format, too. An unstable data interchange format is absolutely an oxymoron, and all but rules out the idea of using pickle in production for much of anything. The best use I can think of is for RPC between clients you control and trust, but I just can't bring myself to have that level of trust, even for a service or host that I'm sure is fully under my control.
Joblib is also used internally for models to be run in parallel, e.g. https://github.com/scikit-learn/scikit-learn/blob/1834cd6b76...
(End users typically just pass n_jobs as a non-zero/one value to use it).
A new safer alternative for scikit-learn model persistence is skops:
- https://skops.readthedocs.io/en/stable/persistence.html
It makes it possible to trust a list of types of Python objects that are safe to load and refuse to load skops files with untrusted types.
import pickle
with open("model.pkl", mode="wb") as f:
pickle.dump(trained_model, f, protocol=pickle.HIGHEST_PROTOCOL)
with open("model.pkl", mode="rb") as f:
trained_model = pickle.load(f)
- https://peps.python.org/pep-0574/
and also provides extra API to handle large data buffers externally ("out-of-band") via custom callbacks. This is in addition to the no-copy semantics memory optimization when loading/storing such arrays "in-band" without providing custom callbacks.
ThreadPoolExecutor if IO-bound
ProcessPoolExecutor if CPU-bound
for example
from concurrent.futures import ThreadPoolExecutor
with ThreadPoolExecutor(max_workers=4) as e:
e.submit(shutil.copy, 'src1.txt', 'dest1.txt')
e.submit(shutil.copy, 'src2.txt', 'dest2.txt')
print("everything done!")PS: The synchronous backend switch is great for getting better backtraces, like in a debugger.
I've used Luigi in a project before; I wasn't too crazy with the class-based syntax at first, but once you get used to it, it's very expressive, and since it's effectively plain python, it means you can customise your pipelines in a very flexible manner.
A tool in similar vein is pydoit. Here you need to be a bit more explicit about declaring the dag, in return it parallelizes well and can tell you what steps are available.