I have a data science application that involves training tens of thousands of small individual Gaussian models. By "small", I mean that any individual model can easily be trained on one of our worker servers. In fact, we should be able to train several at once on each node.

I have been exploring using Spark with Yarn, but Spark seem to be really intended for training big models across multiple machines, not many small models contained on individual machines.

I am imagining a workflow that looks like:

  1. group data by key
  2. send whole groups of data to individual worker machines,
  3. train models for groups on the same worker machine.
  4. Report back or save trained models.

I could use some guidance on how to implement this. The model training is embarrassingly parallelizable.

  • 1
    Not an answer: there was a talk on this at Strata EU this year. Maybe if you contact the speaker nicely they'll give you their slides, or fire up a free 10 day Safari subscription and it's available there. Sep 18, 2018 at 8:18

2 Answers 2


There's no particular right or wrong way to do this because this depends on your projects, and on whether you can exploit the structure of your data for efficiency.

E.g. for a one-off project, you could install the necessary software on all servers, prepare packets of work, prepare a SSH login for all servers, and then use the GNU Parallel to keep all servers busy processing work packets. This is especially well suited for an ad-hoc approach where the input data and output models are stored as plain files, and if you are comfortable with the command line.

If you want to regularly and automatically train new models, it might be better to create a queue of work items, i.e. a shared database that contains all work items and results. You then use some management software to deploy custom-written work server software on all nodes of your cluster. This work server waits on the queue for work packets and writes the results back into the database. This could even be combined with some clever auto-scaling to adapt the number of workers to the amount of pending work, but that might be overkill for a simple project.

In either case:

  1. Start by writing a simple worker software that can train the models locally.
  2. Extend the software so that multiple workers on your local computer can train in parallel – don't use multiple threads. This may involve a database or software like GNU Parallel to synchronize the workers.
  3. Find a way to run those workers distributed over multiple computers. Your worker software is already capable of this, this step is mostly a sysadmin (“ops”) problem.
  • For now it is a one-off, but if results look good it will become recurring. For step (1) I can run locally using Python and scikit-learn.
    – andrew
    Sep 18, 2018 at 18:04
  • For step (2), why no multithreading? Is a for ... & done loop in a shell script sufficient for extending "the software so that multiple workers on your local computer can train in parallel"?
    – andrew
    Sep 18, 2018 at 18:10
  • For step (3), we have a large hadoop cluster. I have seen some threads along the lines of of "create a UDF to use with Pig or Hive". Does this approach sound reasonable?
    – andrew
    Sep 18, 2018 at 18:13
  • 1
    @andrew For (2), because multithreading isn't a suitable step on the way to run distributed instances of your program. This involves problems like making sure the instances don't overwrite each other's output files. Also because of Python-specific restrictions that render threads mostly useless (the “GIL”). Numpy can automatically use real multithreading if you install a suitable linear algebra library, e.g. OpenBLAS. For (3), I'm not familiar with the Hadoop ecosystem. That would probably take care of your problems. But I would first check that the Python/Numpy code works on your cluster.
    – amon
    Sep 18, 2018 at 19:17
  • 1
    @andrew: If you use pandas you might take a look at dask.
    – Giorgio
    Sep 18, 2018 at 19:17

Here is a MCVE example of how I ended up implementing this. This runs entirely in PySpark V2.1.1 with the aid of scikit-learn under some strict assumptions (see Requirement 2).


  1. scikit-learn is installed on every worker machine
  2. for an single model, all data and training overhead can fit on a single worker machine

The general workflow is:

  1. Transform the DataFrame into a RDD[(K,V)] where the keys are group IDs and values are individual data observations
  2. Use groupByKey to shuffle all the data for a single key to a single worker machine
  3. Train models for each key on worker machines
  4. Collect the trained models to serialize for later retrieval

For my application, this approach runs in minutes on a large Spark cluster. I am training tens of thousands of models.

import numpy as np
from pyspark.sql import SparkSession
from sklearn.mixture import GaussianMixture

# Make the example PySpark DataFrame

# Generate example data
nsamps = 500
cv1_1 = np.array([[1,0],[0,1]])
cv1_2 = np.array([[2,0],[0,0.5]])
cv2_1 = np.array([[2, -1.5,],[-1.5, 2]])
cv2_2 = np.array([[2, 1.5,],[1.5, 2]])

mu1_1 = np.array([0,0])
mu1_2 = np.array([0,3])

mu2_1 = mu1_1 + np.array([5,5])
mu2_2 = mu2_1 + np.array([5,5])

# Group 1 data
x1_1 = np.matmul(np.random.randn(nsamps,2), cv1_1) + mu1_1
x1_2 = np.matmul(np.random.randn(nsamps,2), cv1_2) + mu1_2
X1 = np.concatenate([x1_1, x1_2])

# Group 2 data
# X2 = np.matmul(np.random.randn(nsamps,2), cv2_1) + mu2_1
x2_1 = np.matmul(np.random.randn(nsamps,2), cv2_1) + mu2_1
x2_2 = np.matmul(np.random.randn(nsamps,2), cv2_2) + mu2_2
X2 = np.concatenate([x2_1, x2_2])

# Group lables
labs = 2*nsamps*["a"] + 2*nsamps*["b"]

# Create the data frame
X = np.concatenate([X1, X2]).tolist()
dat = [(i, x[0], x[1]) for (i, x) in zip(labs, X)]
cols = ["id", "x", "y"]
df = spark.createDataFrame(dat, cols)

# Shuffle groups to individual workers and train models

# group by ids
kv = df.rdd.map(lambda r: (r.id, [r.x, r.y]))
# create a distrributed RDD where each group is localized on a single worker node
groups = kv.groupByKey()
# a single group is a tuple of id and an iterable with the data
# e.g. (u'a', <pyspark.resultiterable.ResultIterable at 0x7effd7debb90>)

# helper function to train GMMs on the data iterables
def trainGMM(data_itr):
    # Returns a trained GMM
    X = np.array(data_itr.data).astype(np.float64)
    gmm = GaussianMixture(n_components=2, covariance_type='full', tol=0.001, reg_covar=1e-06, max_iter=100,
                          n_init=1, init_params='kmeans', weights_init=None, means_init=None, precisions_init=None,
                          random_state=None, warm_start=False, verbose=0, verbose_interval=10)
    return gmm

# Train GMMs
gmms = groups.mapValues(trainGMM)  # still just a transformation

# Collect and Serialize GMMs

# the trained models are small, so we can collect to a single machine
collected_gmms = gmms.collect()

# pickle models for restoring later
outRoot = "local/output/dir/"

for tup in collected_gmms:
    id = tup[0]
    gmm = tup[1]
    with open("%s/%s_gmm.pkl" % (outRoot, id), 'w') as fout:
        pickle.dump(gmm, fout)

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.