Parallelizing cross validation

I have a nested loop that I’m trying to parallelize that looks like this

def est_model(data, params):
   return model

def score_data(data, model):
   return score

for dataset in Datasets:
    raw_obj = BigData[dataset] ## WANT TO WRAP THIS IN DELAYED
    processed_obj = raw_obj.process() ## TO DELAY THIS
   
    X = list(range(processed_obj.n_chunks))
    scores = []
    kf = KFold(n_splits=5)

    for training, validation in kf.split(X):     
        train = processed_obj[training]
        test = processed_obj[validation]
        model = delayed(est_model)(delayed(train, traverse=False), params)
        scores.append(delayed(score_data)(delayed(test, traverse=False), model))
    
    scores_by_dataset.append(scores)

sbs = dask.compute(*scores_by_datsaset)

It currently works fine, but I’d like to further parallelize the initial raw data processing. The problem is that knowing the number of chunks requires this pre-processing step. As far as I can see, this will mean that there is a for loop which depends on a delayed object. Is there a pattern for cross-validation that I’m just missing?

[TomAugspurger]

https://github.com/dask/dask-ml/compare/master...TomAugspurger:train-test-split?expand=1 has train_test_split mostly finished I think. I don’t recall why I didn’t make a PR with it yet.

Haven’t had time to work on the rest, so if anyone is interested in finishing off that PR or picking up the follow-on things like cross_val_score and the various CV strategies I’d appreciate it.

[TomAugspurger]

Thanks @kdubovikov

I’m hoping to implement some CV strategies (and utilities like cross_val_score) in dask-ml this week. I’ll update this issue when that’s done.