Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Avoid local dir creation, ensure dense array ordering during UMAP save() #823

Merged
merged 8 commits into from
Jan 14, 2025
Merged

Avoid local dir creation, ensure dense array ordering during UMAP save() #823

Changes from 2 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
ed0360d
Avoid local dir creation, fix dense ordering
rishic3 Jan 13, 2025
ba6220a
signoff
rishic3 Jan 13, 2025
297c486
add listdir test, use topandas
rishic3 Jan 14, 2025
2a6e886
license header
rishic3 Jan 14, 2025
3cf20c5
fix overwriting and add test
rishic3 Jan 14, 2025
7b3048c
formatting
rishic3 Jan 14, 2025
3bf1f51
use list()
rishic3 Jan 14, 2025
c0e8b5e
remove unneeded import
rishic3 Jan 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 16 additions & 9 deletions python/src/spark_rapids_ml/umap.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -790,8 +790,9 @@ class UMAP(UMAPClass, _CumlEstimatorSupervised, _UMAPCumlParams):

sample_fraction : float (optional, default=1.0)
The fraction of the dataset to be used for fitting the model. Since fitting is done on a single node, very large
datasets must be subsampled to fit within the node's memory and execute in a reasonable time. Smaller fractions
will result in faster training, but may result in sub-optimal embeddings.
datasets must be subsampled to fit within the node's memory. Smaller fractions will result in faster training, but
may decrease embedding quality. Note: this is not guaranteed to provide exactly the fraction specified of the total
count of the given DataFrame.

featuresCol: str or List[str]
The feature column names, spark-rapids-ml supports vector, array and columnar as the input.\n
Expand Down Expand Up @@ -1473,11 +1474,19 @@ def write_sparse_array(array: scipy.sparse.spmatrix, df_dir: str) -> None:
def write_dense_array(array: np.ndarray, df_path: str) -> None:
schema = StructType(
[
StructField(f"_{i}", FloatType(), False)
for i in range(1, array.shape[1] + 1)
StructField("row_id", LongType(), False),
StructField("data", ArrayType(FloatType(), False), False),
]
)
data_df = spark.createDataFrame(pd.DataFrame(array), schema=schema)
data_df = spark.createDataFrame(
pd.DataFrame(
{
"row_id": range(array.shape[0]),
"data": array.tolist(),
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this better than list(array) ?

Pretty sure we create pandas dfs array columns from np arrays and vice versa elsewhere in our code and would be good to be consistent and/or use best way through out.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While list(array) is more efficient since tolist() does a deep conversion of each row to python lists, Spark will throw an error with list(array) if spark.sql.execution.arrow.pyspark.enabled=false, since pyarrow would no longer handle the numpy -> arrow array conversion.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We pretty much require that to be enabled to get good data transfer from jvm to python workers.

}
),
schema=schema,
)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Below and elsewhere in this class is it correct to use overwrite when writing? This might be counter to the overwrite MLWriter api. If that is not invoked, a user would not expect overwrite to be allowed.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done, thx

data_df.write.parquet(df_path, mode="overwrite")

DefaultParamsWriter.saveMetadata(
Expand All @@ -1495,8 +1504,6 @@ def write_dense_array(array: np.ndarray, df_path: str) -> None:
assert model_attributes is not None

data_path = os.path.join(path, "data")
if not os.path.exists(data_path):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would be good to have a test that checks for expected files and directories?

os.makedirs(data_path)

for key in ["embedding_", "raw_data_"]:
array = model_attributes[key]
Expand Down Expand Up @@ -1547,8 +1554,8 @@ def read_sparse_array(
return scipy.sparse.csr_matrix((data, indices, indptr), shape=csr_shape)

def read_dense_array(df_path: str) -> np.ndarray:
data_df = spark.read.parquet(df_path)
return np.array(data_df.collect(), dtype=np.float32)
data_df = spark.read.parquet(df_path).orderBy("row_id")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if there is test for the order, one that would fail if orderby was omitted.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A multi-gpu env (e.g., DGX) where Spark's default parallelism is >1 would have caught it and I should have tested there with the last PR.
Forcing >1 parallelism would require changing CleanSparkSession to allow a new conf to override the default conf - not sure if that's worth it

return np.array([row.data for row in data_df.collect()], dtype=np.float32)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

toPandas might be better here followed by np.array(list(data_pandas_df.data))


metadata = DefaultParamsReader.loadMetadata(path, self.sc)
data_path = os.path.join(path, "data")
Expand Down
Loading