Override call_cuml_fit_func to use Dataframe, model saving+loading as numpy

[rishic3]

override call_cuml_fit_func

• removed barrier stages, returns dataframe rather than rdd
• allows for:
  • coalesce() rather than repartition
  • toPandas() rather than collect
  • storing raw_data / embeddings as numpy rather than list
  • saving attributes as float32 rather than python float (float64)
• speedups below

dataset: 50,000 x 3000, float32, parquet

	fit runtime	transform runtime
rdd + repartition + collect (no overrides)	58.5s	29.2s
df + coalesce + collect (python lists)	46.7s	28.9s
df + coalesce + toPandas (numpy arrays)	24.9s	8.7s

dataset: 100,000 x 3000, float32, parquet

	fit runtime	transform runtime
rdd + repartition + collect (no overrides)	113.3s	65.4s
df + coalesce + toPandas (numpy arrays)	48.5s	23.7s

override modelwriter / modelreader

• subclassed cumlmodelreader and cumlmodelwriter to handle numpy saving + loading
  • saves arrays with np.save, creates subdirectory for other model attributes ("metadata")
• allows for continuous use of numpy arrays between fit and transform phases
• saves memory and preserves float32 dtype

[leewyang]

Looks nice! Just some minor comments and some questions for rest of team.

[leewyang]

If this is mostly duplicated code, wondering if it can be refactored into the existing API.

Also, is the fit_multiple_params API (from @wbo4958) explicitly unsupported then? If so, maybe we should document this, especially if it's removed for specific reasons.

[rishic3]

Took out fit_multiple_params since fit_multiple isn't supported in UMAP - basically just trimmed out everything that wasn't relevant to UMAP specifically since this func only lives in UMAP atm.

As for refactoring this into the existing API, don't think there's a clean way without overriding or creating a new call_fit_func in core due to the RDD return signature and the barrier stuff. If we're interested in using dataframes for future algos that don't require NCCL during fit, we could have a second call_cuml_fit_func within core like the one in this PR for those use cases which future algos (and this algo) could inherit from. Not sure if this is preferred,

[leewyang]

Probably should log a warning that we're auto-converting the type (but only if it's not too chatty).

[rishic3]

At the moment I'm not supporting user-control over the "convert_dtype" param from cuml (determines whether the internal computations are float64); currently just defaulting to float32. I figured we could keep it like that for now for perf reasons and add float64 support in a future pr (and like Erik mentioned, maybe include a default conversion to float32 much earlier, long before we get to the cuml side, if desired).

[wbo4958]

@rishic3 Could you run

rdd + no-repartition + collect (no overrides) by adding

--conf spark.sql.files.minPartitionNum=$gpu_workers --conf spark.sql.files.maxPartitionBytes=50000000000

[leewyang]

build

[rishic3]

@rishic3 Could you run

rdd + no-repartition + collect (no overrides) by adding

--conf spark.sql.files.minPartitionNum=$gpu_workers --conf spark.sql.files.maxPartitionBytes=50000000000

Both tests were run with these settings. I used the benchmark spark config

[rishic3]

build