Skip to content

Add fit_data_fraction and cache_path to KMeansStage - #1927

Merged
sarahyurick merged 15 commits into
NVIDIA-NeMo:mainfrom
sarahyurick:climb_kmeans
May 7, 2026
Merged

sarahyurick merged 15 commits into
NVIDIA-NeMo:mainfrom
sarahyurick:climb_kmeans

Conversation

@sarahyurick

@sarahyurick sarahyurick commented May 5, 2026 •

Copy link
Copy Markdown
Contributor

TODO:

Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Comment on lines +208 to +210
if self.cache_path is not None:
cp.save(f"{self.cache_path}/kmeans_centroids.npy", self.kmeans.cluster_centers_)
logger.info(f"Saved {self.n_clusters} KMeans centroids to {self.cache_path}/kmeans_centroids.npy")

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Could there be a race condition here?

Comment on lines +303 to +305
if self.cache_path is not None:
cp.save(f"{self.cache_path}/kmeans_centroids.npy", self.kmeans.cluster_centers_)
logger.info(f"Saved {self.n_clusters} KMeans centroids to {self.cache_path}/kmeans_centroids.npy")

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Same comment as above.

Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
sarahyurick and others added 4 commits May 7, 2026 13:57
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>

This branch was previously deployed

2 inactive deployments
nemo-ci — 0223b7fa Deployed May 7, 2026 by copy-pr-bot[bot] via L0_Unit_Test_GPU-sdg #4355
test — 0223b7fa Deployed May 7, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #4355
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants