Ray data migration

[ShreyaR]

Code Pull Requests

This PR adds support for Ray Datasets for distributed prediction.

Documentation Pull Requests

Note that the documentation HTML files are in docs/ while the Markdown sources are in mkdocs/docs.

If you are proposing a modification to the documentation you should change only the Markdown files.

api.md is automatically generated from the docstrings in the code, so if you want to change something in that file, first modify ludwig/api.py docstring, then run mkdocs/code_docs_autogen.py, which will create mkdocs/docs/api.md .

[tgaddair]

This will be more efficient if we turn it into a (stateful) actor:

class BatchInferModel:
    def __init__(self):
        self.model = remote_model.load()
        self.predictor = Predictor(**predictor_kwargs)
    
    def __call__(self, df: pd.DataFrame) -> pd.DataFrame:
        pd_ds = PandasDataset(df, dataset.features, dataset.data_hdf5_fp)
        predictions = self.predictor.batch_predict(self.model, pd_ds, *args, **kwargs)
        ordered_predictions = predictions[output_columns]
        return ordered_predictions

return dataset.ds.map_batchest(BatchInferModel, ...)

This allows us to cache the model and predictor in memory between calls.

[ShreyaR]

Thanks for the suggestion, done!

[tgaddair]

We may need to do a conversion to Dask DataFrame here:

dataset.ds.map_batches(...).to_dask()

The other option is to carry through the RayDataset and then only call the methods to save to a particular format in the end. This might be a better longterm solution, but might be too involved for this PR, since the code already assumes Dask DF will be used downstream. What do you think?

[ShreyaR]

Yeah I agree, I really like the idea of carrying Ray Datasets throughout the prediction and I can do that in a new PR.

[tgaddair]

We need to define this class in-line because we don't want to load the remote model until we're running in the worker context.

[tgaddair]

Probably we want to remove PartitionedDataset here and just raise an exception like:

raise ValueError('Batch inference not supported with TFRecord format at this time')

[tgaddair]

Change isinstance check to RayDataset

[tgaddair]

Nice!

[tgaddair]

Nice work figuring out the flatten stuff. Just a couple small things, then we should be good to land.

[tgaddair]

Last thing we need to double check is the optional dependency imports. I added a few comments about how we can workaround that. Let me know if it makes sense.

[tgaddair]

LGTM! Awesome work.