Datatailr joblib Backend for sklearn

On Datatailr you can create and run complex data processing jobs as DAGs (Directed Acyclic Graphs). Each node on the graph is a task whose code is executed in an isolated container, the environment of which you define and control.

Datatailr takes care of the orchestration, as well as scaling and optimizing the compute resources based on the requirements of the workload.

Our clients use Datatailr to build and run complex ad-hoc workflows onto the cloud (we support AWS/GCP/Azure/Hetzner). Each one is designed for a specific use case.

Recently, we used this mechanism to build a new Python joblib backend. This backend extends sklearn’s existing parallelism onto containers, distributed across multiple servers in your cloud account.

Here's a small example showing how it works:

import joblib 
import numpy as np 
from sklearn.datasets import load_digits 
from sklearn.model_selection import RandomizedSearchCV 
from sklearn.svm import SVC 
 
from datatailr.joblib import register_datatailr 
 
digits = load_digits() 
param_space = { 
    "C": np.logspace(-6, 6, 30), 
    "gamma": np.logspace(-8, 8, 30), 
} 
# 300 candidates x 5-fold CV = 1,500 fits 
search = RandomizedSearchCV( 
    SVC(kernel="rbf"), param_space, cv=5, n_iter=300 
) 
 
register_datatailr() 
with joblib.parallel_backend("datatailr", n_jobs=4): 
    search.fit(digits.data, digits.target)

import joblib 
import numpy as np 
from sklearn.datasets import load_digits 
from sklearn.model_selection import RandomizedSearchCV 
from sklearn.svm import SVC 
 
from datatailr.joblib import register_datatailr 
 
digits = load_digits() 
param_space = { 
    "C": np.logspace(-6, 6, 30), 
    "gamma": np.logspace(-8, 8, 30), 
} 
# 300 candidates x 5-fold CV = 1,500 fits 
search = RandomizedSearchCV( 
    SVC(kernel="rbf"), param_space, cv=5, n_iter=300 
) 
 
register_datatailr() 
with joblib.parallel_backend("datatailr", n_jobs=4): 
    search.fit(digits.data, digits.target)

import joblib 
import numpy as np 
from sklearn.datasets import load_digits 
from sklearn.model_selection import RandomizedSearchCV 
from sklearn.svm import SVC 
 
from datatailr.joblib import register_datatailr 
 
digits = load_digits() 
param_space = { 
    "C": np.logspace(-6, 6, 30), 
    "gamma": np.logspace(-8, 8, 30), 
} 
# 300 candidates x 5-fold CV = 1,500 fits 
search = RandomizedSearchCV( 
    SVC(kernel="rbf"), param_space, cv=5, n_iter=300 
) 
 
register_datatailr() 
with joblib.parallel_backend("datatailr", n_jobs=4): 
    search.fit(digits.data, digits.target)

import joblib 
import numpy as np 
from sklearn.datasets import load_digits 
from sklearn.model_selection import RandomizedSearchCV 
from sklearn.svm import SVC 
 
from datatailr.joblib import register_datatailr 
 
digits = load_digits() 
param_space = { 
    "C": np.logspace(-6, 6, 30), 
    "gamma": np.logspace(-8, 8, 30), 
} 
# 300 candidates x 5-fold CV = 1,500 fits 
search = RandomizedSearchCV( 
    SVC(kernel="rbf"), param_space, cv=5, n_iter=300 
) 
 
register_datatailr() 
with joblib.parallel_backend("datatailr", n_jobs=4): 
    search.fit(digits.data, digits.target)

Here we leave n_jobs off RandomizedSearchCV. The default is None, which means scikit-learn takes the worker count from the with block. Here we used 4. 

Now, instead of spawning 4 local processes, which will be limited by the resources of your local machine, datatailr deploys 4 workers and orchestrates the workload as well as scaling the cluster. 

When the search finishes, the workers are disposed and the results are retrieved into your search object. All worker nodes are automatically created when a new fit is requested and terminated when their resources are no longer needed. 

You can run this from a workstation hosted on datatailr, or from your local environment. The workers use the workstation image, so the execution environment matches your setup. 

Conclusion

Initial container set up takes just a few seconds, but the speedup scales linearly with the number of workers.

Scaling happens automatically and doesn’t require changing your code or dealing with infrastructure management.