Perform inference with Dask when using the Ray backend

[tgaddair]

This PR introduces a new PartitionedDataset constructed by the backend when calling create_inference_dataset within preprocess_for_prediction. Unlike other datasets, the PartitionedDataset runs the batch prediction in parallel on each partition of the distributed DataFrame.

A consequence of this rearchitecture is that instead of Predictor.batch_predict returning a dict, it now returns a DataFrame (as the dataset could be large enough to not fit into memory). Therefore, all the postprocessing now happens on the DataFrame with a flat structure, instead of the nested dict structure.

There a few API limitations to the current implementation:

1. No activation collection for Ray or Dask (usually you wouldn't need to run this in a massive dataset, anyway). We can add this in the future.
2. No batch evaluation when using Dask backend. The evaluation requires some tricky state management to collect the metrics across partitions, which we can do using Ray actors (though it is not ideal). We should rearchitect this in the future so that metrics collection can occur after we run the prediction.
3. No overall stats collection for partitioned datasets. It's not impossible, but will require more work, and is not a top priority at this time.

Another limitation we should address soon is that this implementation does not optimize for GPU inference. Because it runs a separate model replica on each partition, we have no way of preventing multiple model replicas from landing on the same GPU. With standalone Dask, this is possible when configuring your workers (see here). However, this approach doesn't work for Ray at the moment. One possibility is to create workers for each GPU and then use ray.remote to send partitions to GPU workers. If we go this route, we need be smart with how we route partitions to model replicas to ensure minimal network overhead. We can do this in a follow-up PR.

[ifokeev]

so cool feature

[w4nderlust]

why do we need this two times? Can't we do it only once at the end?
Will likely also remove the need for saved_keys

[tgaddair]

At least in the original implementation, existing columns were written to numpy before postprocessing, while new columns were written after postprocessing. This change is meant to preserve that behavior, but we can definitely remove the earlier lines if that change is acceptable.

[w4nderlust]

not 100% sure of how the content of the top k column looks like, but before each was a list of values, so in line 431 there was a loop for each value, while here it sems like there's only a loop over preds and not inside each pred

[tgaddair]

Good catch, every row here should be a pred_top_k. Fixed.

[w4nderlust]

This still does not convince me 100%, as most of the metrics summed / averaged across partitions do no make much sense. maybe we can deactivate evaluation also in the ray case, think it through and then reactivate it.

[tgaddair]

I agree, though this implies it is also currently broken in the normal Horovod case as well:

ludwig/ludwig/models/predictor.py

Line 289 in 51e357d

merged_output_metrics = sum_dicts(

I believe once we adopt the Metrics package for PyTorch that natively supports distributed aggregation this problem will be addressed. Until then, I agree that disabling makes sense.

[w4nderlust]

I think it depends on what we are aggregating, it it's accuracies for instance, this wouldn't be correct, it it's aggregation by summing the counts of correct predictions and all predictions and then later diving correct prediction by all predictions, this would be correct. But I believe we are doing hte formaer.