umap fit() implementation and tests

[rishic3]

• future additions:
  • OOM detection of data subsample (requires running a mini spark job to query GPU memory on node)
  • "convert_dtype" fit() param support

[wbo4958]

It would be great if another test can be added to test the UMAP estimator persistence.

[rishic3]

Idea to avoid overriding _call_cuml_fit_func():

Currently, override of _call_cuml_fit_func is only necessitated by the yield statement.

Alternatively, add class method def _per_row_fit_return(self) -> bool (similar to _fit_array_order or _require_nccl_ucx). Add conditional within core._call_cuml_fit_func() to optionally return data row by row as a bonafide RDD based on this class flag. (may be useful if future algos need this too.)

Open to other (cleaner?) suggestions.

[wbo4958]

Idea to avoid overriding _call_cuml_fit_func():

Currently, override of _call_cuml_fit_func is only necessitated by the yield statement.

Alternatively, add class method def _per_row_fit_return(self) -> bool (similar to _fit_array_order or _require_nccl_ucx). Add conditional within core._call_cuml_fit_func() to optionally return data row by row as a bonafide RDD based on this class flag. (may be useful if future algos need this too.)

Open to other (cleaner?) suggestions.

I would suggest the 2nd way which can re-use code.

[lijinf2]

Idea to avoid overriding _call_cuml_fit_func():

Currently, override of _call_cuml_fit_func is only necessitated by the yield statement.

Alternatively, add class method def _per_row_fit_return(self) -> bool (similar to _fit_array_order or _require_nccl_ucx). Add conditional within core._call_cuml_fit_func() to optionally return data row by row as a bonafide RDD based on this class flag. (may be useful if future algos need this too.)

Open to other (cleaner?) suggestions.

Is it possible to set pyspark configuration to increase this serialization limit? One way is to check the expected serialization limit and hint user to increase the pyspark configuration parameter.

another option is to return row by row in core.py for all algorithms. (any downside?)

[rishic3]

Idea to avoid overriding _call_cuml_fit_func():

Currently, override of _call_cuml_fit_func is only necessitated by the yield statement.
Alternatively, add class method def _per_row_fit_return(self) -> bool (similar to _fit_array_order or _require_nccl_ucx). Add conditional within core._call_cuml_fit_func() to optionally return data row by row as a bonafide RDD based on this class flag. (may be useful if future algos need this too.)
Open to other (cleaner?) suggestions.

Is it possible to set pyspark configuration to increase this serialization limit? One way is to check the expected serialization limit and hint user to increase the pyspark configuration parameter.

another option is to return row by row in core.py for all algorithms. (any downside?)

There is the driver.maxResultSize that can be set to unlimited, but I don't think there's a way to increase the JVM 2GB byte array limit.

As for returning row-by-row for all algos, I don't think this would work - UMAP is not implementing the typical get_cuml_fit_func and instead has its own generator function, so we would have to also make this change for other algos.

[eordentlich]

some additional comments. looks good overall.

[rishic3]

build

[rishic3]

build

[eordentlich]

LGTM